diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 7760bc286..f0485ba32 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -1,4 +1,5 @@ import inspect +import json import logging import os import pickle @@ -15,11 +16,9 @@ from pyevtk.hl import gridToVTK from struphy.feec.psydac_derham import Derham, SplineFunction -from struphy.fields_background import equils from struphy.fields_background.base import FluidEquilibrium -from struphy.geometry import domains from struphy.geometry.base import Domain -from struphy.io.options import BaseUnits, EnvironmentOptions, Time +from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, Time from struphy.io.setup import import_parameters_py from struphy.kinetic_background import maxwellians from struphy.kinetic_background.base import KineticBackground @@ -125,10 +124,10 @@ def __init__( self, path: str, ): - logger.info(f"\nReading in paramters from {path} ... ") + logger.info(f"\nReading in parameters from {path} ... ") params_path = os.path.join(path, "parameters.py") - bin_path = os.path.join(path, "env.bin") + json_path = os.path.join(path, "config.json") if os.path.exists(params_path): params_in = import_parameters_py(params_path) @@ -141,33 +140,45 @@ def __init__( model = params_in.model sim = params_in.sim - elif os.path.exists(bin_path): - with open(os.path.join(path, "env.bin"), "rb") as f: - env = pickle.load(f) - with open(os.path.join(path, "time_opts.bin"), "rb") as f: - time_opts = pickle.load(f) - with open(os.path.join(path, "domain.bin"), "rb") as f: - # WORKAROUND: cannot pickle pyccelized classes at the moment - domain_dct = pickle.load(f) - domain: Domain = getattr(domains, domain_dct["name"])(**domain_dct["params"]) - with open(os.path.join(path, "equil.bin"), "rb") as f: - # WORKAROUND: cannot pickle pyccelized classes at the moment - equil_dct = pickle.load(f) - if equil_dct: - equil: FluidEquilibrium = getattr(equils, equil_dct["name"])(**equil_dct["params"]) - else: - equil = None - with open(os.path.join(path, "grid.bin"), "rb") as f: - grid = pickle.load(f) - with open(os.path.join(path, "derham_opts.bin"), "rb") as f: - derham_opts = pickle.load(f) - with open(os.path.join(path, "model_class.bin"), "rb") as f: - model_class: StruphyModel = pickle.load(f) - model = model_class() + elif os.path.exists(json_path): + with open(json_path, "r") as f: + dct = json.load(f) + env = EnvironmentOptions.from_dict(dct["env"]) + time_opts = Time.from_dict(dct["time_opts"]) + domain: Domain = Domain.from_dict(dct["domain"]) + equil = FluidEquilibrium.from_dict(dct.get("equil")) + + grid_dct = dct.get("grid") + if grid_dct is not None: + grid_dct = dict(grid_dct) + if "num_elements" in grid_dct and grid_dct["num_elements"] is not None: + grid_dct["num_elements"] = tuple(grid_dct["num_elements"]) + if "mpi_dims_mask" in grid_dct and grid_dct["mpi_dims_mask"] is not None: + grid_dct["mpi_dims_mask"] = tuple(grid_dct["mpi_dims_mask"]) + grid = TensorProductGrid.from_dict(grid_dct) + else: + grid = None + + derham_dct = dct.get("derham_opts") + if derham_dct is not None: + derham_dct = dict(derham_dct) + if "degree" in derham_dct and derham_dct["degree"] is not None: + derham_dct["degree"] = tuple(derham_dct["degree"]) + if "bcs" in derham_dct and derham_dct["bcs"] is not None: + derham_dct["bcs"] = tuple(None if bc is None else tuple(bc) for bc in derham_dct["bcs"]) + if "nquads" in derham_dct and derham_dct["nquads"] is not None: + derham_dct["nquads"] = tuple(derham_dct["nquads"]) + if "nquads_proj" in derham_dct and derham_dct["nquads_proj"] is not None: + derham_dct["nquads_proj"] = tuple(derham_dct["nquads_proj"]) + derham_opts = DerhamOptions.from_dict(derham_dct) + else: + derham_opts = None + + model: StruphyModel = StruphyModel.from_dict(dct["model"]) sim = None else: - raise FileNotFoundError(f"Neither of the paths {params_path} or {bin_path} exists.") + raise FileNotFoundError(f"Neither of the paths {params_path} or {json_path} exists.") logger.info("... Done.") diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 74967f6c6..2c683fab3 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -3,7 +3,6 @@ import json import logging import os -import pickle import shutil import sysconfig import time @@ -185,29 +184,9 @@ def __init__( ) except shutil.SameFileError: pass - # pickle struphy objects + # save simulation configuration as JSON else: - with open(os.path.join(path_out, "env.bin"), "wb") as f: - pickle.dump(env, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "time_opts.bin"), "wb") as f: - pickle.dump(time_opts, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "domain.bin"), "wb") as f: - # WORKAROUND: cannot pickle pyccelized classes at the moment - tmp_dct = {"name": domain.__class__.__name__, "params": domain.params} - pickle.dump(tmp_dct, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "equil.bin"), "wb") as f: - # WORKAROUND: cannot pickle pyccelized classes at the moment - if equil is not None: - tmp_dct = {"name": equil.__class__.__name__, "params": equil.params} - else: - tmp_dct = {} - pickle.dump(tmp_dct, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "grid.bin"), "wb") as f: - pickle.dump(grid, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "derham_opts.bin"), "wb") as f: - pickle.dump(derham_opts, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "model_class.bin"), "wb") as f: - pickle.dump(model.__class__, f, pickle.HIGHEST_PROTOCOL) + self.export(os.path.join(path_out, "config.json")) # config clones if self.comm is None: @@ -1300,11 +1279,11 @@ def _initialize_hdf5_datasets(self, data: DataContainer, size: int): return save_keys_all, save_keys_end def _write_run_metadata(self, one_time_step: bool = False): - """Write run-specific JSON metadata for each sim.run() event, reusing to_json().""" + """Write run-specific JSON metadata for each sim.run() event, reusing to_run_metadata().""" if self.rank != 0: return - self.to_json( + self.to_run_metadata( file_path=os.path.join(self.env.path_out, "run_metadata.json"), started_at_epoch_s=self.start_time, one_time_step=one_time_step, @@ -1385,8 +1364,13 @@ def _collect_particle_metadata(self) -> dict: particle_metadata[species_name] = species_metadata return particle_metadata - def to_json(self, file_path: str = None, **extra_data) -> str: - """Assemble the run's data and metadata by hand and serialize to a JSON string. + def to_run_metadata(self, file_path: str = None, **extra_data) -> str: + """Snapshot of the reconstructible config (see :meth:`to_dict`) plus run-specific, + non-reconstructible facts (MPI layout, live particle counts, caller-supplied + timestamps, ...), serialized to a JSON string. + + This is metadata for humans/logging, not a serialization meant to be fed back + into :meth:`from_dict` — use :meth:`to_dict`/:meth:`export` for that. Parameters ---------- @@ -1394,30 +1378,24 @@ def to_json(self, file_path: str = None, **extra_data) -> str: If given, also write the JSON string to this file. **extra_data - Additional key/value pairs merged into the "data" section, + Additional key/value pairs merged into the config, e.g. call-specific facts like a start timestamp. Returns ------- str - The JSON-encoded simulation configuration. + The JSON-encoded simulation metadata. """ - config = { - "name": self.name, - "description": self.description, - "model_name": self.model_name, - "parameter_file": self.params_path, - "mpi_ranks": self.comm_size, - "use_mpi_comm_world": self.comm is not None, - "env": self.env.to_dict(), - "time_opts": self.time_opts.to_dict(), - "domain": self.domain.to_dict(), - "equil": self.equil.to_dict() if self.equil is not None else None, - "grid": self.grid.to_dict() if self.grid is not None else None, - "derham_opts": self.derham_opts.to_dict() if self.derham_opts is not None else None, - "particle_species": self._collect_particle_metadata(), - **extra_data, - } + config = self.to_dict() + config.update( + { + "model_name": self.model_name, + "mpi_ranks": self.comm_size, + "use_mpi_comm_world": self.comm is not None, + "particle_species": self._collect_particle_metadata(), + **extra_data, + }, + ) json_str = json.dumps(config, indent=4) if file_path is not None: