diff --git a/src/pg_rad/main.py b/src/pg_rad/main.py index f30f478..f83be0f 100644 --- a/src/pg_rad/main.py +++ b/src/pg_rad/main.py @@ -98,7 +98,7 @@ def main(): if args.showplots: plotter.plot() - if not (args.save and args.showplots): + if not (args.save or args.showplots): logger.warning( "No output produced. Use --save flag to save outputs and/or " "--showplots to display interactive plots." diff --git a/src/pg_rad/simulator/engine.py b/src/pg_rad/simulator/engine.py index 7b8c950..f22305c 100644 --- a/src/pg_rad/simulator/engine.py +++ b/src/pg_rad/simulator/engine.py @@ -34,6 +34,7 @@ class SimulationEngine: return SimulationOutput( name=self.landscape.name, + size=self.landscape.size, count_rate=count_rate_results, sources=source_results ) diff --git a/src/pg_rad/simulator/outputs.py b/src/pg_rad/simulator/outputs.py index 9aca77e..404e742 100644 --- a/src/pg_rad/simulator/outputs.py +++ b/src/pg_rad/simulator/outputs.py @@ -27,5 +27,6 @@ class SourceOutput: @dataclass class SimulationOutput: name: str + size: tuple count_rate: CountRateOutput sources: List[SourceOutput] diff --git a/src/pg_rad/utils/export.py b/src/pg_rad/utils/export.py index 0d9fffd..a4c38ca 100644 --- a/src/pg_rad/utils/export.py +++ b/src/pg_rad/utils/export.py @@ -1,9 +1,11 @@ +from dataclasses import asdict from datetime import datetime as dt +import json import os import logging import re -from numpy import array, full_like +from numpy import array, full_like, ndarray from pandas import DataFrame from pg_rad.simulator.outputs import SimulationOutput @@ -12,6 +14,13 @@ from pg_rad.simulator.outputs import SimulationOutput logger = logging.getLogger(__name__) +class NumpyEncoder(json.JSONEncoder): + def default(self, obj): + if isinstance(obj, ndarray): + return obj.tolist() + return super().default(obj) + + def generate_folder_name(sim: SimulationOutput) -> str: formatted_sim_name = re.sub(r"\s+", '_', sim.name.lower()) folder_name = ( @@ -35,9 +44,18 @@ def save_results(sim: SimulationOutput, folder_name: str) -> None: df = generate_df(sim) csv_name = generate_csv_name(sim) df.to_csv(f"{folder_name}/{csv_name}.csv", index=False) + with open(f"{folder_name}/parameters.json", 'w') as f: + json.dump(generate_sim_param_dict(sim), 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.""" @@ -62,6 +80,7 @@ def generate_df(sim: SimulationOutput) -> DataFrame: 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 = [ [ @@ -81,5 +100,6 @@ def generate_csv_name(sim: SimulationOutput) -> str: src_str = "_".join(src_str_array.flat) - csv_name = f"{num_src}_src_{bkg_cps}_cps_bkg_{src_str}" + src_ids_str = "_".join(src_ids) + csv_name = f"{src_ids_str}_src_{bkg_cps}_cps_bkg_{src_str}" return csv_name