mirror of
https://github.com/pim-n/pg-rad
synced 2026-06-17 15:59:35 +02:00
128 lines
3.6 KiB
Python
128 lines
3.6 KiB
Python
from dataclasses import asdict
|
|
from datetime import datetime as dt
|
|
import json
|
|
import os
|
|
import logging
|
|
import re
|
|
|
|
from numpy import array, full_like, ndarray, bool_
|
|
from pandas import DataFrame
|
|
|
|
from pg_rad.simulator.outputs import SimulationOutput
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logging.getLogger("PIL").setLevel(logging.WARNING)
|
|
|
|
|
|
class NumpyEncoder(json.JSONEncoder):
|
|
def default(self, obj):
|
|
if isinstance(obj, ndarray):
|
|
return obj.tolist()
|
|
elif isinstance(obj, bool_):
|
|
return bool(obj)
|
|
return super().default(obj)
|
|
|
|
|
|
def generate_folder_name(sim: SimulationOutput) -> str:
|
|
formatted_sim_name = re.sub(r"\s+", '_', sim.name.lower())
|
|
folder_name = (
|
|
formatted_sim_name +
|
|
'_result_' +
|
|
dt.today().strftime('%Y%m%d_%H%M')
|
|
)
|
|
return folder_name
|
|
|
|
|
|
def save_results(sim: SimulationOutput, folder_name: str) -> None:
|
|
"""Parse all simulation output and save to a folder."""
|
|
if not os.path.exists(folder_name):
|
|
os.makedirs(folder_name)
|
|
else:
|
|
logger.warning("Folder already exists. Overwrite?")
|
|
ans = input("[type 'n' to cancel overwrite] ")
|
|
if ans.lower() == 'n':
|
|
return
|
|
|
|
logger.debug(
|
|
f"Integrated counts: {list(sim.count_rate.integrated_counts)}"
|
|
)
|
|
logger.debug(
|
|
f"Distances: {list(sim.count_rate.distance)}"
|
|
)
|
|
|
|
df = generate_df(sim)
|
|
csv_name = generate_csv_name(sim)
|
|
df.to_csv(f"{folder_name}/{csv_name}.csv", index=False)
|
|
param_dict = generate_sim_param_dict(sim)
|
|
with open(f"{folder_name}/parameters.json", 'w') as f:
|
|
json.dump(param_dict, f, cls=NumpyEncoder)
|
|
logger.info(f"Simulation output saved to {folder_name}!")
|
|
|
|
|
|
def generate_sim_param_dict(sim: SimulationOutput) -> dict:
|
|
"""Parse simulation parameters and hyperparameters to dictionary."""
|
|
d = asdict(sim)
|
|
d.pop('count_rate')
|
|
return d
|
|
|
|
|
|
def generate_df(sim: SimulationOutput) -> DataFrame:
|
|
"""Parse simulation output to CSV format and the name of CSV."""
|
|
|
|
br_array = full_like(
|
|
sim.count_rate.integrated_counts,
|
|
sim.count_rate.mean_bkg_cps
|
|
)
|
|
|
|
east_coords = sim.count_rate.x[1:]
|
|
north_coords = sim.count_rate.y[1:]
|
|
|
|
if len(east_coords) != sim.count_rate.integrated_counts.shape[0]:
|
|
east_coords = None
|
|
north_coords = None
|
|
logger.warning(
|
|
"PG-RAD currently does not support this experimental path."
|
|
" Only ROI_P, ROI_BR and Dist will be saved."
|
|
)
|
|
|
|
result_df = DataFrame(
|
|
{
|
|
"East": east_coords,
|
|
"North": north_coords,
|
|
"ROI_P": sim.count_rate.integrated_counts,
|
|
"ROI_BR": br_array,
|
|
"Dist": sim.count_rate.distance
|
|
}
|
|
)
|
|
|
|
return result_df
|
|
|
|
|
|
def generate_csv_name(sim: SimulationOutput) -> str:
|
|
"""Generate CSV name according to Alex' specification"""
|
|
num_src = len(sim.sources)
|
|
src_ids = [str(i+1) for i in range(num_src)]
|
|
bkg_cps = round(sim.count_rate.mean_bkg_cps)
|
|
source_param_strings = [
|
|
[
|
|
str(round(s.activity))+"MBq",
|
|
str(round(s.dist_from_path))+"m",
|
|
str(round(s.position[0]))+'_'+str(round(s.position[1]))
|
|
]
|
|
for s in sim.sources
|
|
]
|
|
|
|
if num_src == 1:
|
|
src_str = "_".join(source_param_strings[0])
|
|
else:
|
|
src_str_array = array(
|
|
[list(item) for item in zip(*source_param_strings)]
|
|
)
|
|
|
|
src_str = "_".join(src_str_array.flat)
|
|
|
|
src_ids_str = "_".join(src_ids)
|
|
csv_name = f"{src_ids_str}_src_{bkg_cps}_cps_bkg_{src_str}"
|
|
return csv_name
|