# third party imports
import ast
import copy
import dataclasses
import glob
import hashlib
import importlib
import inspect
import json
import logging
import os
import shutil
import textwrap
import time
from collections.abc import Sequence
from pathlib import Path
from typing import TYPE_CHECKING
import cunumpy as xp
import h5py
import numpy as np
import yaml
from cunumpy.kernels import Kernel
from feectools.linalg.memory import stencil_matrix_memory
from feectools.linalg.stencil import StencilVector
from line_profiler import profile
from maybempi import MPI, SerialMPI
from pyevtk.hl import gridToVTK
from scope_profiler import ProfileManager
if TYPE_CHECKING:
import pyvista as pv
# api imports
from struphy import (
BaseUnits,
DerhamOptions,
EnvironmentOptions,
FieldsBackground,
ProfilingOptions,
Time,
domains,
equils,
grids,
set_logging_level,
)
# core imports
from struphy.feec.basis_projection_ops import BasisProjectionOperators
from struphy.feec.mass import WeightedMassOperators
from struphy.feec.memory import linop_nbytes, vector_nbytes
from struphy.feec.psydac_derham import Derham
from struphy.fields_background.base import (
FluidEquilibrium,
FluidEquilibriumWithB,
MHDequilibrium,
NumericalMHDequilibrium,
)
from struphy.fields_background.projected_equils import (
ProjectedFluidEquilibrium,
ProjectedFluidEquilibriumWithB,
ProjectedMHDequilibrium,
)
from struphy.geometry.base import Domain
from struphy.initial.base import Perturbation
from struphy.io.output_handling import DataContainer
from struphy.models.base import StruphyModel
from struphy.models.species import (
DiagnosticSpecies,
FieldSpecies,
FluidSpecies,
ParticleSpecies,
Species,
)
from struphy.models.variables import FEECVariable, PICVariable, SPHVariable
from struphy.physics.physics import Units
from struphy.pic.base import Particles
from struphy.post_processing.output import Output
from struphy.propagators.base import Propagator
from struphy.simulation.base import SimulationBase
from struphy.utils.clone_config import CloneConfig
from struphy.utils.kernel_compilation import collect_kernels, missing_cuda
from struphy.utils.progress import tqdm
from struphy.utils.utils import ruff_autofix_and_format
logger = logging.getLogger("struphy")
class CuPyJSONEncoder(json.JSONEncoder):
"""JSON encoder that handles CuPy arrays and NumPy arrays."""
def default(self, obj):
# Check if it has a .get() method (CuPy array)
if hasattr(obj, "get"):
return obj.get().tolist() if hasattr(obj.get(), "tolist") else obj.get()
# Handle NumPy arrays and scalars
import numpy as np
if isinstance(obj, np.ndarray):
return obj.tolist()
if isinstance(obj, (np.integer, np.floating)):
return float(obj) if isinstance(obj, np.floating) else int(obj)
return super().default(obj)
[docs]
class Simulation(SimulationBase):
"""Top-level class to configure and run a Struphy simulation.
The `Simulation` class wraps model setup, MPI configuration, output
management, normalization (units), memory allocation and time stepping.
It initializes the model's variables and propagators, prepares runtime
metadata and output folders, and provides the main `run()` entry point
to execute the simulation.
Parameters
----------
model : StruphyModel
Physics model that provides species, propagators and variables.
name : str, optional
Name of the simulation.
description : str, optional
Description of the simulation.
params_path : str, optional
Path to a Python parameter file to save alongside outputs.
env : EnvironmentOptions
Runtime and output environment options.
time_opts : Time
Time-stepping options (dt, Tend, split algorithm, ...).
domain : Domain
Computational domain description.
equil : FluidEquilibrium, optional
Initial fluid equilibrium (may be None).
grid : TensorProductGrid
Spatial grid used for FEEC variables.
derham_opts : DerhamOptions
Options for discrete differential operators.
profiling_opts : ProfilingOptions
Options passed to scope-profiler when profiling is active.
comm: MPI.Intracomm, optional
MPI communicator for parallel execution. If None, uses MPI.COMM_WORLD.
logging_level : int, optional
Logging level (e.g., logging.INFO, logging.DEBUG). If None, uses default.
"""
def __init__(
self,
model: StruphyModel,
name: str = "",
description: str = "",
params_path: str = None,
env: EnvironmentOptions = EnvironmentOptions(),
time_opts: Time = Time(),
domain: Domain = domains.Cuboid(),
equil: FluidEquilibrium = None,
grid: grids.TensorProductGrid = grids.TensorProductGrid(),
derham_opts: DerhamOptions = DerhamOptions(),
profiling_opts: ProfilingOptions = ProfilingOptions(),
comm: MPI.Intracomm = None,
logging_level: int | None = None,
):
if logging_level is not None:
set_logging_level(logging_level)
self.model = model
self.name = name
self.description = description
self.params_path = params_path
self.env = env # Set name first since it's used as a label
self.time_opts = time_opts
self.domain = domain
self.equil = equil
self.grid = grid
self.derham_opts = derham_opts
self.profiling_opts = profiling_opts
self.comm = comm
if logger.level <= logging.INFO and self.rank == 0:
self.show_parameters()
self.model.units.show_units()
self.model.show_equation_params()
# synchronize MPI processes to set same start time of simulation for all processes
self.Barrier()
self.start_time = time.time()
self.clone_config = self._create_clone_config()
self._output = None
self.Barrier()
# ----------------
# Abstract methods
# ----------------
def show_parameters(self):
"""Print the current simulation configuration to stdout.
Only the MPI rank 0 prints to avoid clutter from multiple processes.
"""
print(f"\nNew instance of Simulation from file\n{self.params_path}\n")
print(self.model)
print("")
print(self.env)
print("")
print(self.time_opts)
print("")
print(self.domain)
print("")
print(self.equil)
print("")
print(self.grid)
print("")
print(self.derham_opts)
def show_propagator_options(self):
# Display propagator options and intial conditions:
if self.rank == 0:
print("\nPROPAGATOR OPTIONS:")
for prop in self.model.prop_list:
assert isinstance(prop, Propagator)
print(prop)
def show_initial_conditions(self):
if self.rank == 0:
print("\nINITIAL CONDITIONS:")
for species in self.model.species.values():
assert isinstance(species, Species)
for variable in species.variables.values():
if isinstance(variable, FEECVariable) or isinstance(variable, SPHVariable):
variable.show_backgrounds()
variable.show_perturbations()
elif isinstance(variable, PICVariable):
variable.show_backgrounds()
variable.show_perturbations()
variable.show_initial_condition()
def allocate(self):
"""Allocate FEEC structures, model variables and propagators.
This prepares FEEC operators, allocates variable storage for all
species (fields, fluids, particles) and passes allocation info to
propagators. Prints progress on MPI rank 0.
"""
logger.debug("\nAllocating simulation data ...")
with ProfileManager.profile_region("setup: allocate"):
# feec
with ProfileManager.profile_region("allocate feec", functions=[self._allocate_feec]):
self._allocate_feec(self.grid, self.derham_opts)
# allocate model variables
with ProfileManager.profile_region("allocate variables", functions=[self._allocate_variables]):
self._allocate_variables()
# pass info to propagators
with ProfileManager.profile_region("allocate propagators", functions=[self._allocate_propagators]):
self._allocate_propagators()
# allocate helper fields and perform initial solves if needed
with ProfileManager.profile_region("model post allocate", functions=[self.model.post_allocate]):
self.model.post_allocate()
logger.debug("... Done.")
def kernels(self) -> tuple[Kernel, ...]:
"""The :class:`~cunumpy.kernels.Kernel` objects the time stepping calls, each once.
Collected after :meth:`allocate` from the domain (:meth:`Domain.kernels
<struphy.geometry.base.Domain.kernels>`), the Derham complex (:meth:`Derham.kernels
<struphy.feec.psydac_derham.Derham.kernels>`), the particles of every kinetic variable
(:meth:`Particles.kernels <struphy.pic.base.Particles.kernels>`) and the propagators of the model
(:meth:`Propagator.kernels <struphy.propagators.base.Propagator.kernels>`). Kernels that only
diagnostics or the initial solves call are not included; they run during the setup anyway.
"""
owners = [self.domain, self.derham]
for species in self.model.species.values():
for variable in species.variables.values():
if isinstance(variable, (PICVariable, SPHVariable)):
owners.append(variable.particles)
owners += self.model.prop_list
return collect_kernels(*owners)
def compile_cuda_kernels(self) -> tuple[Kernel, ...]:
"""Compile the CUDA versions of the kernels of :meth:`kernels` now, before the time stepping.
A CUDA kernel is otherwise compiled on its first call, which puts the compile time into the
first time step (and its profiling regions). :meth:`run` calls this method after
:meth:`allocate`, in the profiling region ``"setup: compile cuda kernels"``. CuPy caches the
compiled kernels on disk, so a second run on the same machine mostly loads them.
On the NumPy backend this does nothing (the pyccel kernels are compiled by ``struphy compile``).
Returns
-------
tuple[Kernel, ...]
The kernels whose CUDA version was compiled; empty on the NumPy backend.
Raises
------
NotImplementedError
On the CuPy backend, if a kernel the time stepping calls has no CUDA version. All such
kernels are listed, and nothing is compiled; without this check the run would fail at the
first call of the first one, in the time loop.
"""
if xp.get_backend() != "cupy":
return ()
kernels = self.kernels()
missing = missing_cuda(kernels)
if missing:
lines = "\n".join(
f" - {kernel.name}" + (f" (expected at {kernel.cuda_path})" if kernel.cuda_path else "")
for kernel in missing
)
raise NotImplementedError(
f"Model {self.model_name} cannot run on the CuPy backend: {len(missing)} of the "
f"{len(kernels)} kernels its time stepping calls have no CUDA version:\n{lines}",
)
compiled = []
with ProfileManager.profile_region("setup: compile cuda kernels"):
for kernel in kernels:
if kernel.compile():
compiled.append(kernel)
logger.info(f"Compiled {len(compiled)} CUDA kernels: {', '.join(k.name for k in compiled)}")
return tuple(compiled)
def estimate_mem(self, print_report: bool = False) -> dict:
"""Estimate the memory footprint of all model variables and FEEC matrices, in bytes,
BEFORE calling :meth:`allocate`.
Builds a throwaway Derham sequence (cheap: metadata only, no mass/basis operators) to obtain
the FEEC coefficient space sizes and MPI domain decomposition, then calls ``estimate_mem()`` on
every model variable, mirroring the loop in :meth:`_allocate_variables`. The throwaway Derham is
discarded afterward -- calling :meth:`allocate` still performs the full allocation from scratch.
On top of the variables, the following FEEC matrices are estimated (they are usually much
larger than the spline coefficient vectors):
* the derivative matrices (grad, curl, div) of the Derham sequence (matrix-free, hence
essentially free of charge),
* the standard mass matrices M0, M1, M2, M3 and Mv, created with ``dry_run=True`` so that
only their sizes (including the zero-block structure) are computed, see
:meth:`~struphy.feec.mass.WeightedMassOperators.estimate_mem`.
Note
----
Which matrices a model really allocates is only known once the propagators have been
allocated (they fetch and build their operators in ``allocate()``, partly under names
assembled at run time). The standard mass matrices are therefore used as a baseline:
a model may not need all of them, but it may also build additional weighted mass matrices
(e.g. ``M2n``), basis projection operators and preconditioners which are *not* included here.
Use :meth:`report_mem` after :meth:`allocate` for the exact numbers.
Parameters
----------
print_report : bool
If True, print a breakdown of the estimated memory usage of each variable on
MPI rank 0: both local (this rank) and global (summed over all ranks) values.
Returns
-------
dict
Mapping ``{"species.variable": local_bytes}`` for the model variables and
``{"matrices.<name>": local_bytes}`` for the FEEC matrices (bytes on the *current*
MPI rank), plus a ``"total"`` entry with the local sum over all entries.
"""
logger.debug("\nEstimating memory usage ...")
if self.grid is None or self.derham_opts is None:
raise RuntimeError(
"Simulation.estimate_mem() requires 'grid' and 'derham_opts' to be set (needed to build a Derham sequence for FEEC variable sizing)."
)
if self.clone_config is None:
derham_comm = self.comm
else:
derham_comm = self.clone_config.sub_comm
derham = Derham(
self.grid,
self.derham_opts,
comm=derham_comm,
domain=self.domain,
)
mem = {}
if self.model.field_species:
for species, spec in self.model.field_species.items():
for k, v in spec.variables.items():
assert isinstance(v, FEECVariable)
mem[f"{species}.{k}"] = v.estimate_mem(derham=derham)
if self.model.fluid_species:
for species, spec in self.model.fluid_species.items():
for k, v in spec.variables.items():
assert isinstance(v, FEECVariable)
mem[f"{species}.{k}"] = v.estimate_mem(derham=derham)
if self.model.particle_species:
for species, spec in self.model.particle_species.items():
for k, v in spec.variables.items():
if isinstance(v, PICVariable):
mem[f"{species}.{k}"] = v.estimate_mem(
clone_config=self.clone_config,
derham=derham,
domain=self.domain,
equil=self.equil,
)
elif isinstance(v, SPHVariable):
mem[f"{species}.{k}"] = v.estimate_mem(
derham=derham,
domain=self.domain,
equil=self.equil,
)
if self.model.diagnostic_species:
for species, spec in self.model.diagnostic_species.items():
for k, v in spec.variables.items():
assert isinstance(v, FEECVariable)
mem[f"{species}.{k}"] = v.estimate_mem(derham=derham)
# FEEC matrices: derivative matrices (already allocated by the throwaway Derham) ...
mem["matrices.derivatives"] = sum(linop_nbytes(op) for op in (derham.grad, derham.curl, derham.div))
# ... and the standard mass matrices (dry run, i.e. sized but not allocated)
mass_ops = WeightedMassOperators(derham, self.domain, eq_mhd=self.equil)
for name, nbytes in mass_ops.estimate_mem().items():
mem[f"matrices.{name}"] = nbytes
total_local = sum(mem.values())
mem["total"] = total_local
if print_report:
if self.comm is not None:
total_global = self.comm.allreduce(total_local, op=MPI.SUM)
else:
total_global = total_local
if self.rank == 0:
print("\nESTIMATED MEMORY USAGE (before allocate()):")
for name, nbytes in mem.items():
if name == "total":
continue
print(f" {name}: {nbytes / 1e6:.2f} MB (local, rank 0)")
print(
f" TOTAL: {total_local / 1e6:.2f} MB (local, rank 0), "
f"{total_global / 1e6:.2f} MB (global, summed over {self.comm_size} rank(s))"
)
logger.debug("... Done.")
return mem
def report_mem(self, print_report: bool = True) -> dict:
"""Actual local (per-MPI-rank) memory footprint of the big arrays, in bytes,
AFTER calling :meth:`allocate`.
In contrast to :meth:`estimate_mem`, nothing is estimated here: the spline coefficient
vectors and marker arrays are measured on the allocated objects, and the FEEC matrices are
obtained from the (weak) registry of all allocated stencil matrices,
:data:`~feectools.linalg.memory.stencil_matrix_memory`. The matrix number therefore
includes mass matrices, basis projection operators, preconditioners and any other
stencil matrix built by the propagators.
Parameters
----------
print_report : bool
If True (default), print the breakdown on MPI rank 0.
Returns
-------
dict
Mapping ``{"feec_matrices": ..., "spline_coeffs": ..., "markers": ..., "total": ...}``
with the local number of bytes.
"""
mem = {"feec_matrices": stencil_matrix_memory.nbytes, "spline_coeffs": 0, "markers": 0}
for species in (
self.model.field_species,
self.model.fluid_species,
self.model.particle_species,
self.model.diagnostic_species,
):
if not species:
continue
for spec in species.values():
for v in spec.variables.values():
if isinstance(v, FEECVariable):
mem["spline_coeffs"] += vector_nbytes(v.spline.vector)
elif isinstance(v, (PICVariable, SPHVariable)):
mem["markers"] += v.particles.nbytes_local
if v.n_to_save > 0:
mem["markers"] += v.saved_markers.nbytes
total_local = sum(mem.values())
mem["total"] = total_local
if print_report:
if self.comm is not None:
total_global = self.comm.allreduce(total_local, op=MPI.SUM)
else:
total_global = total_local
if self.rank == 0:
print("\nMEMORY USAGE (after allocate()):")
for name, nbytes in mem.items():
if name == "total":
continue
print(f" {name}: {nbytes / 1e6:.2f} MB (local, rank 0)")
print(
f" TOTAL: {total_local / 1e6:.2f} MB (local, rank 0), "
f"{total_global / 1e6:.2f} MB (global, summed over {self.comm_size} rank(s))"
)
return mem
def save_geometry_and_equil_vtk(self):
"""Write a VTK file with geometry and (projected) equilibrium fields.
Only executed on MPI rank 0. Outputs basic diagnostic fields such as
jacobian determinant, pressure and |B| when available.
"""
# store geometry vtk
if self.rank == 0:
grids_log = [
xp.linspace(1e-6, 1.0, 32),
xp.linspace(0.0, 1.0, 32),
xp.linspace(0.0, 1.0, 32),
]
tmp = self.domain(*grids_log)
grids_phy = [
DataContainer._as_numpy_array(tmp[0]),
DataContainer._as_numpy_array(tmp[1]),
DataContainer._as_numpy_array(tmp[2]),
]
pointData = {}
det_df = self.domain.jacobian_det(*grids_log)
pointData["det_df"] = DataContainer._as_numpy_array(det_df)
if self.equil is not None:
p0 = self.equil.p0(*grids_log)
pointData["p0"] = DataContainer._as_numpy_array(p0)
n0 = self.equil.n0(*grids_log)
pointData["n0"] = DataContainer._as_numpy_array(n0)
if isinstance(self.equil, FluidEquilibriumWithB):
absB0 = self.equil.absB0(*grids_log)
pointData["absB0"] = DataContainer._as_numpy_array(absB0)
gridToVTK(os.path.join(self.env.path_out, "geometry"), *grids_phy, pointData=pointData)
def create_geometry_mesh(
self,
nx: int = 32,
ny: int = 32,
nz: int = 32,
):
"""Create a PyVista mesh with geometry and (projected) equilibrium fields.
Returns a StructuredGrid mesh with basic diagnostic fields such as
jacobian determinant, pressure and |B| when available.
Returns
-------
pyvista.StructuredGrid
Mesh containing geometry and equilibrium field data.
"""
grids_log = [
xp.linspace(1e-6, 1.0, nx),
xp.linspace(0.0, 1.0, ny),
xp.linspace(0.0, 1.0, nz),
]
tmp = self.domain(*grids_log)
grids_phy = [
DataContainer._as_numpy_array(tmp[0]),
DataContainer._as_numpy_array(tmp[1]),
DataContainer._as_numpy_array(tmp[2]),
]
# Create PyVista structured grid
import pyvista as pv
mesh = pv.StructuredGrid(grids_phy[0], grids_phy[1], grids_phy[2])
# Add point data
det_df = self.domain.jacobian_det(*grids_log)
mesh["det_df"] = DataContainer._as_numpy_array(det_df).ravel(order="F")
if self.equil is not None:
p0 = self.equil.p0(*grids_log)
mesh["p0"] = DataContainer._as_numpy_array(p0).ravel(order="F")
if isinstance(self.equil, FluidEquilibriumWithB):
absB0 = self.equil.absB0(*grids_log)
mesh["absB0"] = DataContainer._as_numpy_array(absB0).ravel(order="F")
return mesh
def show_domain(
self,
scalars: list | str | None = None,
nx: int = 32,
ny: int = 32,
nz: int = 32,
window_size: tuple | None = None,
zoom_factor: int = 1.0,
) -> "pv.Plotter":
"""Visualize the geometry and (projected) equilibrium fields using PyVista."""
import pyvista as pv
if self.rank == 0:
mesh = self.create_geometry_mesh(nx=nx, ny=ny, nz=nz)
pv.set_jupyter_backend("static")
if scalars:
if isinstance(scalars, str):
scalars_to_plot = [scalars]
else:
scalars_to_plot = scalars
else:
scalar_names = mesh.array_names
scalars_to_plot = scalar_names[:3] if len(scalar_names) >= 3 else scalar_names
if window_size is None:
window_size = (len(scalars_to_plot) * 500, 250)
# Create a plotter with three subplots side by side
plotter = pv.Plotter(shape=(1, len(scalars_to_plot)), window_size=window_size)
for idx, scalar_name in enumerate(scalars_to_plot):
plotter.subplot(0, idx)
plotter.add_mesh(
mesh,
scalars=scalar_name,
show_edges=False,
cmap="jet",
scalar_bar_args={
"title": scalar_name,
"vertical": True,
"title_font_size": 12,
"label_font_size": 10,
"height": 0.8,
},
)
plotter.view_isometric()
plotter.camera.zoom(zoom_factor)
plotter.show()
return plotter
return None
def initialize_data_storage(self):
"""Create the `DataContainer` and register time datasets.
Initializes `time_state` arrays (normalized and physical time and
index) and registers them with the output `DataContainer` so they
are saved during the run (and on restart).
"""
# data object for saving (will either create new hdf5 files if restart==False or open existing files if restart==True)
# use the simulation communicator when storing the outputs
self.data = DataContainer(self.env.path_out, comm=self.comm)
# time quantities (current time value, value in seconds and index); host bookkeeping of the time loop,
# NumPy on both backends, so that reading them for the loop control and saving them never waits for the device
self.time_state = {}
self.time_state["value"] = np.zeros(1, dtype=float)
self.time_state["value_sec"] = np.zeros(1, dtype=float)
self.time_state["index"] = np.zeros(1, dtype=int)
# add time quantities to data object for saving
for key, val in self.time_state.items():
key_time = "time/" + key
key_time_restart = "restart/time/" + key
self.data.add_data({key_time: val})
self.data.add_data({key_time_restart: val})
def run(self, one_time_step: bool = False, profiling_activated: bool | None = None) -> Output:
"""Main entry point to execute the simulation time loop.
Responsibilities include allocation (when not restarting),
initialization of output storage, handling restarts, running the
main time-stepping loop, saving data at intervals, and finalizing
profiling and metadata. Prints progress on MPI rank 0.
Parameters
----------
one_time_step : bool
If True, only perform one time step (useful for testing).
profiling_activated : bool | None
If True, activate profiling with scope-profiler for this run. If
None, profiling is disabled.
Returns
-------
Output
The output of this run, see :attr:`output`.
"""
if profiling_activated is None:
profiling_activated = False
logger.info(f"\nStarting run for model {self.model_name} on {self.comm_size} ranks ...")
if self.name != "":
logger.info(f"Simulation name: {self.name}")
if self.description != "":
logger.info(f"Description: {self.description}")
self._remove_existing_output_files()
self._setup_folders()
self._copy_parameter_file()
self.Barrier()
self._output = None
with ProfileManager.session(
options=self.profiling_opts,
deactivate_profiling=not profiling_activated,
file_path=self.profiling_opts.file_path or self.profiling_filepath,
):
with ProfileManager.profile_region("setup: total"):
# equation paramters
self.allocate()
# CUDA kernels (CuPy backend only): compile them now rather than in the first time step
self.compile_cuda_kernels()
# time step dependent data of the propagators (e.g. preconditioners), rather than in the first step
with ProfileManager.profile_region("setup: prepare time step", functions=[self._prepare_time_step]):
self._prepare_time_step()
with ProfileManager.profile_region("setup: run metadata", functions=[self._write_run_metadata]):
self._write_run_metadata(
one_time_step=one_time_step,
profiling_activated=profiling_activated,
)
# output
with ProfileManager.profile_region("setup: data storage", functions=[self.initialize_data_storage]):
self.initialize_data_storage()
# peek view into geometry
with ProfileManager.profile_region("setup: geometry vtk", functions=[self.save_geometry_and_equil_vtk]):
self.save_geometry_and_equil_vtk()
# plasma parameters
with ProfileManager.profile_region("setup: plasma params", functions=[self.compute_plasma_params]):
self.compute_plasma_params()
# print info on mpi procs
if self.comm_size < 32:
if self.derham is not None:
logger.info(f"\nderham.domain_array:\n{self.derham.domain_array}")
else:
for _, species in self.model.species.items():
for _, variable in species.variables.items():
if isinstance(variable, (PICVariable, SPHVariable)):
logger.info(f"\nparticle domain_array:\n{variable.particles.domain_array}")
break
if self.rank < 32:
logger.debug("")
logger.debug(f"Rank {self.rank}: executing run() for model {self.model_name} ...")
if self.comm_size > 32 and self.rank == 32:
logger.debug(f"Ranks > 31: executing run() for model {self.model_name} ...")
# retrieve time parameters
dt = self.time_opts.dt
if one_time_step:
Tend = dt
else:
Tend = self.time_opts.Tend
split_algo = self.time_opts.split_algo
# set initial conditions for all variables
if self.env.restart:
with ProfileManager.profile_region("setup: restart", functions=[self._initialize_from_restart]):
self._initialize_from_restart(self.data)
with h5py.File(self.data.file_path, "a") as file:
self.time_state["value"][0] = file["restart/time/value"][-1]
self.time_state["value_sec"][0] = file["restart/time/value_sec"][-1]
self.time_state["index"][0] = file["restart/time/index"][-1]
start_step = file["restart/time/index"][-1]
total_steps = int(round((Tend - float(self.time_state["value"][0])) / dt))
logger.info(f"""\n!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
RESTARTing from:
self.time_state["value"][0]={float(self.time_state["value"][0])}
self.time_state["value_sec"][0]={float(self.time_state["value_sec"][0])}
self.time_state["index"][0]={int(self.time_state["index"][0])}
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
""")
else:
total_steps = int(round(Tend / dt))
start_step = 0
total_steps_str = str(total_steps)
# compute initial scalars and kinetic data, pass time state to all propagators
initial_diagnostics_functions = [
self.model.update_scalar_quantities,
self.model.update_markers_to_be_saved,
self.model.update_distr_functions,
self._add_time_state,
]
with ProfileManager.profile_region(
"setup: initial diagnostics",
functions=initial_diagnostics_functions,
):
self.model.update_scalar_quantities()
self.model.update_markers_to_be_saved()
self.model.update_distr_functions()
self._add_time_state(self.time_state["value"])
# add all variables to be saved to data object
with ProfileManager.profile_region("setup: hdf5 datasets", functions=[self._initialize_hdf5_datasets]):
save_keys_all, save_keys_end = self._initialize_hdf5_datasets(self.data, self.comm_size)
# ======================== main time loop ======================
self.model.update_scalar_quantities()
if logger.level <= logging.INFO and self.rank == 0:
print("\nINITIAL SCALAR QUANTITIES:")
self.model.print_scalar_quantities()
print(f"START TIME STEPPING WITH '{split_algo}' SPLITTING:")
# time loop
run_time_now = 0.0
show_progress_bar = logger.getEffectiveLevel() <= logging.WARNING and self.rank == 0
pbar = tqdm(
total=total_steps,
disable=not show_progress_bar,
desc="Time stepping",
unit="step",
)
while True:
self.Barrier()
# stop time loop?
break_cond_1 = float(self.time_state["value"][0]) >= Tend
break_cond_2 = run_time_now > self.env.max_runtime
if break_cond_1 or break_cond_2:
# save restart data (other data already saved below)
with ProfileManager.profile_region("save data", functions=[self.data.save_data]):
self.data.save_data(keys=save_keys_end)
end_time = time.time()
logger.info(f"\nTime steps done: {int(self.time_state['index'][0])}")
logger.info(f"wall-clock time of simulation [sec]: {end_time - self.start_time}")
logger.info("")
break
if self.env.sort_step and int(self.time_state["index"][0]) % self.env.sort_step == 0:
t0 = time.time()
sort_functions = [val.do_sort for val in self.model.pointer.values() if isinstance(val, Particles)]
with ProfileManager.profile_region("sort particles", functions=sort_functions):
for key, val in self.model.pointer.items():
if isinstance(val, Particles):
val.do_sort()
t1 = time.time()
message = "Particles sorted | wall clock [s]: {0:8.4f} | sorting duration [s]: {1:8.4f}".format(
run_time_now * 60,
t1 - t0,
)
logger.info(message)
logger.info("")
# update time and index (round time to 10 decimals for a clean time grid!)
self.time_state["value"][0] = round(float(self.time_state["value"][0]) + dt, 14)
self.time_state["value_sec"][0] = round(
float(self.time_state["value_sec"][0]) + dt * self.model.units.t, 14
)
self.time_state["index"][0] += 1
# perform one time step dt
t0 = time.time()
with ProfileManager.profile_region("model.integrate", functions=[self.model.integrate]):
self.model.integrate(dt, split_algo)
t1 = time.time()
run_time_now = (time.time() - self.start_time) / 60
# update diagnostics data and save data
if int(self.time_state["index"][0]) % self.env.save_step == 0:
# compute scalars and kinetic data
diagnostics_functions = [
self.model.update_scalar_quantities,
self.model.update_markers_to_be_saved,
self.model.update_distr_functions,
]
with ProfileManager.profile_region("diagnostics", functions=diagnostics_functions):
self.model.update_scalar_quantities()
self.model.update_markers_to_be_saved()
self.model.update_distr_functions()
# extract FEEC coefficients
feec_species = (
self.model.field_species | self.model.fluid_species | self.model.diagnostic_species
)
for species, val in feec_species.items():
assert isinstance(val, Species)
for variable, subval in val.variables.items():
assert isinstance(subval, FEECVariable)
spline = subval.spline
# in-place extraction of FEM coefficients from field.vector --> field.vector_stencil!
spline.extract_coeffs(update_ghost_regions=False)
# save data (everything but restart data)
with ProfileManager.profile_region("save data", functions=[self.data.save_data]):
self.data.save_data(keys=save_keys_all)
# print current time and scalar quantities to screen
step = str(int(self.time_state["index"][0])).zfill(len(total_steps_str))
message = "\n" + "-" * 80 + "\n"
message += "time step:".ljust(25) + f"{step}/{total_steps + start_step}".rjust(25)
message += (
"\n"
+ "normalized time:".ljust(25)
+ "{0:4.2e} / {1:4.2e}".format(float(self.time_state["value"][0]), Tend).rjust(25)
)
message += (
"\n"
+ "physical time [s]:".ljust(25)
+ "{0:4.2e} / {1:4.2e}".format(
float(self.time_state["value_sec"][0]),
Tend * self.model.units.t,
).rjust(25)
)
message += "\n" + "wall clock time [s]:".ljust(25) + "{0:8.4f}".format(run_time_now * 60).rjust(25)
message += "\n" + "last step duration [s]:".ljust(25) + "{0:8.4f}".format(t1 - t0).rjust(25)
logger.info(message)
if logger.level <= logging.INFO and self.rank == 0:
self.model.print_scalar_quantities()
if show_progress_bar:
pbar.update(1)
pbar.close()
# ===================================================================
self.Barrier()
logger.info("Struphy run finished.")
if self.clone_config is not None:
self.clone_config.free()
return self.output
@property
def output(self) -> Output:
"""The output of this simulation in ``env.path_out``, see :class:`~struphy.Output`.
Scalars are available as soon as data is written; fields and particle products are
post-processed on first access, or explicitly with ``sim.output.pproc(...)``.
"""
if self._output is None or self._output.path_out != Path(self.env.path_out).resolve():
self._output = Output(self.env.path_out)
return self._output
# ---------------------
# Code specific methods
# ---------------------
def compute_plasma_params(self):
"""
Compute and print volume averaged plasma parameters for each species of the model.
Global parameters:
- plasma volume
- transit length
- magnetic field
Species dependent parameters:
- mass
- charge
- density
- pressure
- thermal energy kBT
- Alfvén speed v_A
- thermal speed v_th
- thermal frequency Omega_th
- cyclotron frequency Omega_c
- plasma frequency Omega_p
- Alfvèn frequency Omega_A
- thermal Larmor radius rho_th
- MHD length scale v_a/Omega_c
- rho/L
- alpha = Omega_p/Omega_c
- epsilon = 1/(t*Omega_c)
"""
# units affices for printing
units_affix = {}
units_affix["plasma volume"] = " m³"
units_affix["transit length"] = " m"
units_affix["magnetic field"] = " T"
units_affix["mass"] = " kg"
units_affix["charge"] = " C"
units_affix["density"] = " m⁻³"
units_affix["pressure"] = " Pa"
units_affix["kBT"] = " keV"
units_affix["v_A"] = " m/s"
units_affix["v_th"] = " m/s"
units_affix["vth1"] = " m/s"
units_affix["vth2"] = " m/s"
units_affix["vth3"] = " m/s"
units_affix["Omega_th"] = " Mrad/s"
units_affix["Omega_c"] = " Mrad/s"
units_affix["Omega_p"] = " Mrad/s"
units_affix["Omega_A"] = " Mrad/s"
units_affix["rho_th"] = " m"
units_affix["v_A/Omega_c"] = " m"
units_affix["rho_th/L"] = ""
units_affix["alpha"] = ""
units_affix["epsilon"] = ""
h = 1 / 20
eta1 = xp.linspace(h / 2.0, 1.0 - h / 2.0, 20)
eta2 = xp.linspace(h / 2.0, 1.0 - h / 2.0, 20)
eta3 = xp.linspace(h / 2.0, 1.0 - h / 2.0, 20)
# global parameters
# plasma volume (hat x^3)
det_tmp = self.domain.jacobian_det(eta1, eta2, eta3)
vol1 = xp.mean(xp.abs(det_tmp))
# plasma volume (m⁻³)
plasma_volume = vol1 * self.model.units.x**3
# transit length (m)
transit_length = plasma_volume ** (1 / 3)
# magnetic field (T)
if isinstance(self.equil, FluidEquilibriumWithB):
B_tmp = self.equil.absB0(eta1, eta2, eta3)
else:
B_tmp = xp.zeros((eta1.size, eta2.size, eta3.size))
magnetic_field = xp.mean(B_tmp * xp.abs(det_tmp)) / vol1 * self.model.units.B
B_max = xp.max(B_tmp) * self.model.units.B
B_min = xp.min(B_tmp) * self.model.units.B
if magnetic_field < 1e-14:
magnetic_field = xp.nan
# logger.info("\n+++++++ WARNING +++++++ magnetic field is zero - set to nan !!")
logger.info("\nPLASMA PARAMETERS:")
logger.info(
"Plasma volume:".ljust(25) + "{:4.3e}".format(plasma_volume) + units_affix["plasma volume"],
)
logger.info(
"Transit length:".ljust(25) + "{:4.3e}".format(transit_length) + units_affix["transit length"],
)
logger.info(
"Avg. magnetic field:".ljust(25) + "{:4.3e}".format(magnetic_field) + units_affix["magnetic field"],
)
logger.info(
"Max magnetic field:".ljust(25) + "{:4.3e}".format(B_max) + units_affix["magnetic field"],
)
logger.info(
"Min magnetic field:".ljust(25) + "{:4.3e}".format(B_min) + units_affix["magnetic field"],
)
def spawn_sister(
self,
model: StruphyModel = None,
params_path: str = None,
env: EnvironmentOptions = None,
time_opts: Time = None,
domain: Domain = None,
equil: FluidEquilibrium = None,
grid: grids.TensorProductGrid = None,
derham_opts: DerhamOptions = None,
profiling_opts: ProfilingOptions = None,
comm: MPI.Intracomm = None,
):
"""Spawn a sister simulation with parameters that default to the current instance.
This can be used to quickly generate multiple similar simulations."""
if model is None:
model = self.model
if params_path is None:
params_path = self.params_path
if env is None:
env = self.env
if time_opts is None:
time_opts = self.time_opts
if domain is None:
domain = self.domain
if equil is None:
equil = self.equil
if grid is None:
grid = self.grid
if derham_opts is None:
derham_opts = self.derham_opts
if profiling_opts is None:
profiling_opts = self.profiling_opts
if comm is None:
comm = self.comm
sister = Simulation(
model=model,
params_path=params_path,
env=env,
time_opts=time_opts,
domain=domain,
equil=equil,
grid=grid,
derham_opts=derham_opts,
profiling_opts=profiling_opts,
comm=comm,
)
return sister
# ---------------
# Private methods
# ---------------
def _setup_folders(self):
"""
Setup output folders.
"""
if self.rank == 0:
# create output folder if it does not exit
if not os.path.exists(self.env.path_out):
os.makedirs(self.env.path_out, exist_ok=True)
logger.debug("Created folder " + self.env.path_out)
# create data folder in output folder if it does not exist
if not os.path.exists(os.path.join(self.env.path_out, "data/")):
os.mkdir(os.path.join(self.env.path_out, "data/"))
logger.debug("Created folder " + os.path.join(self.env.path_out, "data/"))
def _remove_existing_output_files(self):
"""Removes post_processing/, meta.txt and profile_tmp.
If not restart, also removes existing hdf5 and png files in output folder."""
if self.rank == 0:
# remove post_processing folder
folder = os.path.join(self.env.path_out, "post_processing")
if os.path.exists(folder):
shutil.rmtree(folder)
logger.info("Removed existing folder " + folder)
# remove meta file
file = os.path.join(self.env.path_out, "meta.txt")
if os.path.exists(file):
os.remove(file)
logger.info("Removed existing file " + file)
# remove profiling file
file = os.path.join(self.env.path_out, "profile_tmp")
if os.path.exists(file):
os.remove(file)
logger.info("Removed existing file " + file)
# remove hdf5 and png files (if NOT a restart)
if not self.env.restart:
files = glob.glob(os.path.join(self.env.path_out, "data", "*.hdf5"))
for n, file in enumerate(files):
os.remove(file)
if n < 10: # print only ten statements in case of many processes
logger.info("Removed existing file " + file)
files = glob.glob(os.path.join(self.env.path_out, "*.png"))
for n, file in enumerate(files):
os.remove(file)
if n < 10: # print only ten statements in case of many processes
logger.info("Removed existing file " + file)
def _copy_parameter_file(self):
"""Copy the parameter file to the output folder for reference."""
if self.rank != 0:
return
if self.params_path is not None:
try:
shutil.copy2(
self.params_path,
os.path.join(self.env.path_out, "parameters.py"),
)
except shutil.SameFileError:
pass
def _create_clone_config(self) -> CloneConfig | None:
"""Setup domain cloning communicators, None if there is only one clone (or no MPI).
MPI.COMM_WORLD : comm
within a clone: : sub_comm
between the clones : inter_comm
"""
if self.comm is None or self.env.num_clones == 1:
return None
clone_config = CloneConfig(comm=self.comm, params=None, num_clones=self.env.num_clones)
clone_config.print_clone_config()
if self.model.particle_species:
clone_config.print_particle_config()
return clone_config
@profile
def _allocate_feec(self, grid: grids.TensorProductGrid, derham_opts: DerhamOptions):
"""Create the discrete Derham sequence, mass/basis operators and projected equilibrium.
This sets up the 3D Derham object (unless grid or derham_opts are
None), creates weighted mass and basis projection operators, and
constructs a projected equilibrium appropriate for the chosen
equilibrium type.
"""
# check for polar singularity
if self.domain.pole:
assert (
derham_opts.polar_splines
), """Polar singularity detected in domain but derham_opts.polar_splines is False.
You have two options:
a) set derham_opts.polar_splines = True to use polar splines
b) modify your domain definition to cut a hole around the singularity."""
# create discrete derham sequence
if self.clone_config is None:
derham_comm = self.comm
else:
derham_comm = self.clone_config.sub_comm
if grid is None or derham_opts is None:
logger.debug(f"\n{grid=}, {derham_opts=}: no Derham object set up.")
self._derham = None
else:
with ProfileManager.profile_region("setup: derham", functions=[Derham.__init__]):
self._derham = Derham(
grid,
derham_opts,
comm=derham_comm,
domain=self.domain,
)
# create weighted mass and basis operators
if self.derham is None:
self._mass_ops = None
self._basis_ops = None
else:
with ProfileManager.profile_region("setup: mass ops", functions=[WeightedMassOperators.__init__]):
self._mass_ops = WeightedMassOperators(self.derham, self.domain, eq_mhd=self.equil)
with ProfileManager.profile_region("setup: basis ops", functions=[BasisProjectionOperators.__init__]):
self._basis_ops = BasisProjectionOperators(
self.derham,
self.domain,
eq_mhd=self.equil,
)
# create projected equilibrium
if self.derham is None:
self._projected_equil = None
elif self.equil is None:
self._projected_equil = None
else:
projected_equil_functions = [
ProjectedMHDequilibrium.__init__,
ProjectedFluidEquilibriumWithB.__init__,
ProjectedFluidEquilibrium.__init__,
]
with ProfileManager.profile_region("setup: projected equil", functions=projected_equil_functions):
if isinstance(self.equil, MHDequilibrium):
self._projected_equil = ProjectedMHDequilibrium(
self.equil,
self.derham,
)
elif isinstance(self.equil, FluidEquilibriumWithB):
self._projected_equil = ProjectedFluidEquilibriumWithB(
self.equil,
self.derham,
)
elif isinstance(self.equil, FluidEquilibrium):
self._projected_equil = ProjectedFluidEquilibrium(
self.equil,
self.derham,
)
else:
self._projected_equil = None
@profile
def _allocate_variables(self):
"""
Allocate memory for model variables and set initial conditions.
"""
# allocate memory for FE coeffs of electromagnetic fields/potentials
if self.model.field_species:
for species, spec in self.model.field_species.items():
assert isinstance(spec, FieldSpecies)
for k, v in spec.variables.items():
assert isinstance(v, FEECVariable)
with ProfileManager.profile_region(f"allocate var: {species}.{k}", functions=[v.allocate]):
v.allocate(
derham=self.derham,
domain=self.domain,
equil=self.equil,
)
# allocate memory for FE coeffs of fluid variables
if self.model.fluid_species:
for species, spec in self.model.fluid_species.items():
assert isinstance(spec, FluidSpecies)
for k, v in spec.variables.items():
assert isinstance(v, FEECVariable)
with ProfileManager.profile_region(f"allocate var: {species}.{k}", functions=[v.allocate]):
v.allocate(
derham=self.derham,
domain=self.domain,
equil=self.equil,
)
# allocate memory for marker arrays of kinetic variables
if self.model.particle_species:
for species, spec in self.model.particle_species.items():
assert isinstance(spec, ParticleSpecies)
for k, v in spec.variables.items():
if isinstance(v, PICVariable):
with ProfileManager.profile_region(f"allocate var: {species}.{k}", functions=[v.allocate]):
v.allocate(
clone_config=self.clone_config,
derham=self.derham,
domain=self.domain,
equil=self.equil,
projected_equil=self.projected_equil,
)
if isinstance(v, SPHVariable):
with ProfileManager.profile_region(f"allocate var: {species}.{k}", functions=[v.allocate]):
v.allocate(
derham=self.derham,
domain=self.domain,
equil=self.equil,
projected_equil=self.projected_equil,
)
# allocate memory for FE coeffs of fluid variables
if self.model.diagnostic_species:
for species, spec in self.model.diagnostic_species.items():
assert isinstance(spec, DiagnosticSpecies)
for k, v in spec.variables.items():
assert isinstance(v, FEECVariable)
with ProfileManager.profile_region(f"allocate var: {species}.{k}", functions=[v.allocate]):
v.allocate(
derham=self.derham,
domain=self.domain,
equil=self.equil,
)
# TODO: allocate memory for FE coeffs of diagnostics
# if self.params.diagnostic_fields is not None:
# for key, val in self.diagnostics.items():
# if "params" in key:
# continue
# else:
# val["obj"] = self.derham.create_spline_function(
# key,
# val["space"],
# bckgr_params=None,
# pert_params=None,
# )
# self._pointer[key] = val["obj"].vector
def _prepare_time_step(self):
"""Call :meth:`~struphy.propagators.base.Propagator.prepare_time_step` of each propagator with the time step
it receives from :meth:`~struphy.models.base.StruphyModel.integrate` (``dt / 2`` for all but the last
propagator with Strang splitting)."""
dt = self.time_opts.dt
props = self.model.prop_list
for n, prop in enumerate(props):
half = self.time_opts.split_algo == "Strang" and n < len(props) - 1
prop.prepare_time_step(dt / 2 if half else dt)
@profile
def _allocate_propagators(self):
"""Allocate propagators and bind shared FEEC/domain operators.
Assigns `derham`, `domain`, `mass_ops`, `basis_ops` and
`projected_equil` on the `Propagator` base class so individual
propagator instances can access shared resources, then calls each
propagator's `allocate` method.
"""
# set propagators base class attributes (then available to all propagators)
Propagator.derham = self.derham
Propagator.domain = self.domain
if self.derham is not None:
Propagator.mass_ops = self.mass_ops
Propagator.basis_ops = self.basis_ops
Propagator.projected_equil = self.projected_equil
assert len(self.model.prop_list) > 0, "No propagators in this model, check the model class."
for prop in self.model.prop_list:
assert isinstance(prop, Propagator)
with ProfileManager.profile_region("allocate prop: " + prop.__class__.__name__, functions=[prop.allocate]):
prop.allocate()
logger.debug(f"\nAllocated propagator '{prop.__class__.__name__}'.")
@profile
@staticmethod
def _binned_background(background, bin_plot) -> xp.ndarray:
"""Evaluate a kinetic background on the bin centers of ``bin_plot``.
Directions that are not binned are evaluated at zero; velocity directions that are not
binned are integrated out like the binned data (exact for Maxwellians).
"""
centers = {
dim: edges[:-1] + (edges[1] - edges[0]) / 2
for dim, edges in zip(bin_plot.slice.split("_"), bin_plot.bin_edges)
}
grids = [centers.get(dim, xp.zeros(1)) for dim in ("e1", "e2", "e3")]
factor = 1.0
for component in range(1, background.vdim + 1):
dim = f"v{component}"
if dim in centers:
grids.append(centers[dim])
else:
grids.append(xp.zeros(1))
factor *= xp.sqrt(2 * xp.pi)
return background(*xp.meshgrid(*grids, indexing="ij")).squeeze() * factor
def _initialize_hdf5_datasets(self, data: DataContainer, size: int):
"""
Create datasets in hdf5 files according to model unknowns and diagnostics data.
Parameters
----------
data : struphy.io.output_handling.DataContainer
The data object that links to the hdf5 files.
size : int
Number of MPI processes of the model run.
Returns
-------
save_keys_all : list
Keys of datasets which are saved during the simulation.
save_keys_end : list
Keys of datasets which are saved at the end of a simulation to enable restarts.
"""
# save scalar quantities in group 'scalar/'
# the host copies of the scalars (Scalars.to_host): saving them copies nothing from the device
for key in self.model.scalars.dct:
key_scalar = "scalar/" + key
data.add_data({key_scalar: self.model.scalars.host_value(key)})
with h5py.File(data.file_path, "a") as file:
# store grid_info only for runs with 512 ranks or smaller
if self.model.scalars.dct and self.derham is not None:
if size <= 512:
file["scalar"].attrs["grid_info"] = DataContainer._as_numpy_array(self.derham.domain_array)
else:
file["scalar"].attrs["grid_info"] = DataContainer._as_numpy_array(self.derham.domain_array[0])
else:
pass
# save feec data in group 'feec/'
feec_species = self.model.field_species | self.model.fluid_species | self.model.diagnostic_species
for species, val in feec_species.items():
assert isinstance(val, Species)
species_path = os.path.join("feec", species)
species_path_restart = os.path.join("restart", species)
for variable, subval in val.variables.items():
assert isinstance(subval, FEECVariable)
spline = subval.spline
# in-place extraction of FEM coefficients from field.vector --> field.vector_stencil!
spline.extract_coeffs(update_ghost_regions=False)
# save numpy array to be updated each time step.
if subval.save_data:
key_field = os.path.join(species_path, variable)
if isinstance(spline.vector_stencil, StencilVector):
data.add_data(
{key_field: spline.vector_stencil._data},
)
else:
for n in range(3):
key_component = os.path.join(key_field, str(n + 1))
data.add_data(
{key_component: spline.vector_stencil[n]._data},
)
# save field meta data
file[key_field].attrs["space_id"] = spline.space_id
file[key_field].attrs["starts"] = DataContainer._as_numpy_array(spline.starts)
file[key_field].attrs["ends"] = DataContainer._as_numpy_array(spline.ends)
file[key_field].attrs["pads"] = DataContainer._as_numpy_array(spline.pads)
# save numpy array to be updated only at the end of the simulation for restart.
if self.env.save_restart:
key_field_restart = os.path.join(species_path_restart, variable)
if isinstance(spline.vector_stencil, StencilVector):
data.add_data(
{key_field_restart: spline.vector_stencil._data},
)
else:
for n in range(3):
key_component_restart = os.path.join(key_field_restart, str(n + 1))
data.add_data(
{key_component_restart: spline.vector_stencil[n]._data},
)
# save kinetic data in group 'kinetic/'
for name, species in self.model.particle_species.items():
assert isinstance(species, ParticleSpecies)
assert len(species.variables) == 1, "More than 1 variable per kinetic species is not allowed."
for varname, var in species.variables.items():
assert isinstance(var, PICVariable | SPHVariable)
obj = var.particles
assert isinstance(obj, Particles)
key_spec = os.path.join("kinetic", name)
# restart data
if self.env.save_restart:
key_spec_restart = os.path.join("restart", name)
data.add_data({key_spec_restart: obj.markers})
# marker data
if var.n_to_save > 0:
key_mks = os.path.join(key_spec, "markers")
data.add_data({key_mks: var.saved_markers})
# binning plot data
for bin_plot in species.saving_params.binning_plots:
# define slice name with binning quantity
slice, output_quantity = bin_plot.slice, bin_plot.output_quantity
slice = f"{slice}_{output_quantity}"
key_f = os.path.join(key_spec, "f", slice)
key_df = os.path.join(key_spec, "df", slice)
data.add_data({key_f: bin_plot.f})
data.add_data({key_df: bin_plot.df})
for dim, be in enumerate(bin_plot.bin_edges):
file[key_f].attrs["bin_centers" + "_" + str(dim + 1)] = DataContainer._as_numpy_array(
be[:-1] + (be[1] - be[0]) / 2
)
# the static background of a delta-f species, so that post-processing can
# reconstruct the full f without the simulation's configuration
if var.space == "DeltaFParticles6D":
key_background = os.path.join(key_spec, "f_background", slice)
if key_background in file:
del file[key_background]
file.create_dataset(
key_background,
data=DataContainer._as_numpy_array(self._binned_background(var.backgrounds, bin_plot)),
)
for i, kd_plot in enumerate(species.saving_params.kernel_density_plots):
key_n = os.path.join(key_spec, "n_sph", f"view_{i}")
data.add_data({key_n: kd_plot.n_sph})
# save 1d point values, not meshgrids, because attrs size is limited
eta1 = kd_plot.plot_pts[0][:, 0, 0]
eta2 = kd_plot.plot_pts[1][0, :, 0]
eta3 = kd_plot.plot_pts[2][0, 0, :]
file[key_n].attrs["eta1"] = DataContainer._as_numpy_array(eta1)
file[key_n].attrs["eta2"] = DataContainer._as_numpy_array(eta2)
file[key_n].attrs["eta3"] = DataContainer._as_numpy_array(eta3)
# TODO: maybe add other data
# else:
# data.add_data({key_dat: val1})
# keys to be saved at each time step and only at end (restart)
save_keys_all = []
save_keys_end = []
for key in data.dset_dict:
if "restart" in key:
save_keys_end.append(key)
else:
save_keys_all.append(key)
return save_keys_all, save_keys_end
def _write_run_metadata(self, one_time_step: bool = False, profiling_activated: bool = False):
"""Write run-specific JSON metadata for each sim.run() event, reusing to_run_metadata()."""
if self.rank != 0:
return
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,
profiling_activated=profiling_activated,
)
def _add_time_state(self, time_state):
"""Add a pointer to the time variable of the dynamics ('t')
to the model and to all propagators of the model.
Parameters
----------
time_state : ndarray
Of size 1, holds the current physical time 't'.
"""
assert time_state.size == 1
self._time_state = time_state
for _, prop in self.model.propagators.__dict__.items():
if isinstance(prop, Propagator):
prop.add_time_state(time_state)
def _initialize_from_restart(self, data: DataContainer):
"""
Set initial conditions for FE coefficients (electromagnetic and fluid) and markers from restart group in hdf5 files.
Parameters
----------
data : struphy.io.output_handling.DataContainer
The data object that links to the hdf5 files.
"""
with h5py.File(data.file_path, "a") as file:
for species, val in self.model.species.items():
for variable, subval in val.variables.items():
# initialize feec variables
if isinstance(subval, FEECVariable):
key_restart = os.path.join("restart", species, variable)
subval.spline.initialize_coeffs_from_restart_file(
file,
key=key_restart,
)
# initialize pic and sph variables
elif isinstance(subval, (PICVariable, SPHVariable)):
key_restart = os.path.join("restart", species)
subval.particles._markers[:, :] = file[key_restart][-1, :, :]
# refresh the cached hole/ghost masks for the restored markers
subval.particles.update_holes()
subval.particles._update_ghost_particles()
if self.comm_size > 1:
subval.particles.mpi_sort_markers(do_test=True)
def to_dict(self) -> dict:
"""Serialize the simulation configuration to a dictionary."""
return {
"name": self.name,
"description": self.description,
"model": self.model.to_dict(),
"params_path": self.params_path,
"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,
"profiling_opts": vars(self.profiling_opts).copy(),
}
@staticmethod
def _serialize_function_globals(func, source: str, seen=None) -> dict:
"""Capture the module-level names a user function needs to be re-executed.
Default arguments (``def f(x, r=r_minus)``) and body references (``R0``,
helper functions) are resolved in the defining module, so the source alone
cannot be restored. Plain values, modules and other top-level functions are
recorded; anything else is left out and fails only if it is actually used.
"""
seen = set() if seen is None else seen
seen = seen | {id(func)}
names = sorted(
{
node.id
for node in ast.walk(ast.parse(source))
if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load)
}
)
def plain(item):
if item is None or isinstance(item, (bool, int, float, str)):
return True
return isinstance(item, (list, tuple)) and all(plain(each) for each in item)
referenced = {}
for name in names:
if name == func.__name__ or name not in func.__globals__:
continue
item = func.__globals__[name]
if inspect.ismodule(item):
referenced[name] = {"type": "module", "name": item.__name__}
elif plain(item):
referenced[name] = {"type": "value", "value": Simulation._serialize_initial_condition(item)}
elif inspect.isfunction(item) and id(item) not in seen:
serialized = Simulation._serialize_initial_condition(item, _seen=seen)
if serialized.get("serialization") != "unsupported":
referenced[name] = serialized
return referenced
@staticmethod
def _deserialize_function_globals(referenced: dict) -> dict:
"""Rebuild the namespace recorded by :meth:`_serialize_function_globals`."""
namespace = {}
for name, item in referenced.items():
if item["type"] == "module":
namespace[name] = importlib.import_module(item["name"])
elif item["type"] == "value":
namespace[name] = Simulation._deserialize_initial_condition(item["value"])
else:
namespace[name] = Simulation._deserialize_initial_condition(item)
return namespace
@staticmethod
def _serialize_initial_condition(value, _seen=None):
"""Convert initial-condition definitions into JSON-compatible provenance data.
This deliberately captures constructor parameters rather than evaluated FEEC
coefficients or particle data. It is therefore small and records the setup
that produced the initial state. The representation is not yet used to
reconstruct a simulation from output.
"""
if value is None or isinstance(value, (bool, int, float, str)):
return value
if isinstance(value, dict):
return {str(key): Simulation._serialize_initial_condition(item) for key, item in value.items()}
if isinstance(value, (list, tuple)):
return [Simulation._serialize_initial_condition(item) for item in value]
if (
type(value).__module__
not in {
"struphy.initial.perturbations",
"struphy.kinetic_background.maxwellians",
"struphy.kinetic_background.base",
"struphy.io.options",
}
and not inspect.isfunction(value)
and (isinstance(value, Perturbation) or callable(value))
):
cls = type(value)
if "<locals>" in cls.__qualname__:
return {
"type": "python_class",
"serialization": "unsupported",
"reason": "nested classes are not supported",
}
if cls.__module__ == "struphy" or cls.__module__.startswith("struphy."):
# Struphy's own classes are importable; re-executing their source would
# lose the names their module imports.
data = {"type": "python_class", "name": cls.__qualname__, "module": cls.__module__}
else:
try:
source = textwrap.dedent(inspect.getsource(cls))
except (OSError, TypeError):
return {
"type": "python_class",
"serialization": "unsupported",
"reason": "source code is unavailable",
}
data = {
"type": "python_class",
"name": cls.__name__,
"source": source,
"source_sha256": hashlib.sha256(source.encode()).hexdigest(),
}
if hasattr(value, "params"):
data["params"] = Simulation._serialize_initial_condition(value.params)
elif hasattr(value, "__dict__"):
data["state"] = Simulation._serialize_initial_condition(vars(value))
else:
data.update(
{
"serialization": "unsupported",
"reason": "callable objects without instance state are not supported",
}
)
return data
if dataclasses.is_dataclass(value) and not isinstance(value, type):
return {
"type": type(value).__name__,
"params": Simulation._serialize_initial_condition(
{field.name: getattr(value, field.name) for field in dataclasses.fields(value) if field.init}
),
}
if hasattr(value, "params"):
parameters = value.params
try:
signature = inspect.signature(type(value))
except (TypeError, ValueError):
signature = None
if signature is not None and not any(
parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values()
):
parameters = {
name: parameter_value
for name, parameter_value in parameters.items()
if name in signature.parameters
}
return {
"type": type(value).__name__,
"params": Simulation._serialize_initial_condition(parameters),
}
if inspect.isfunction(value):
# Top-level functions are fully captured in metadata. Their source is
# self-contained: neither the module name nor a source-file reference is
# required to recover it later.
if value.__name__ == "<lambda>" or "<locals>" in value.__qualname__ or value.__closure__ is not None:
return {
"type": "python_function",
"serialization": "unsupported",
"reason": "lambdas, nested functions, and closures are not supported",
}
try:
source = textwrap.dedent(inspect.getsource(value))
except (OSError, TypeError):
return {
"type": "python_function",
"serialization": "unsupported",
"reason": "source code is unavailable",
}
data = {
"type": "python_function",
"name": value.__name__,
"source": source,
"source_sha256": hashlib.sha256(source.encode()).hexdigest(),
}
referenced = Simulation._serialize_function_globals(value, source, _seen)
if referenced:
data["globals"] = referenced
return data
if callable(value):
return {
"type": "callable",
"serialization": "unsupported",
"reason": "only top-level Python functions are currently supported",
}
# CuPyJSONEncoder handles NumPy/CuPy arrays after this traversal. Keep
# other values visible in provenance rather than making metadata writing fail.
return value
@staticmethod
def _deserialize_initial_condition(value):
"""Rebuild one initial-condition definition from metadata."""
if value is None or isinstance(value, (bool, int, float, str)):
return value
if isinstance(value, list):
return tuple(Simulation._deserialize_initial_condition(item) for item in value)
if (
not isinstance(value, dict)
or "type" not in value
or ("params" not in value and value["type"] not in {"python_function", "python_class", "callable"})
):
return {key: Simulation._deserialize_initial_condition(item) for key, item in value.items()}
kind = value["type"]
if kind == "python_function":
if value.get("serialization") == "unsupported":
raise ValueError(f"Cannot restore initial-condition function: {value['reason']}.")
source = value["source"]
if hashlib.sha256(source.encode()).hexdigest() != value["source_sha256"]:
raise ValueError("Initial-condition function source hash does not match its metadata.")
import cunumpy as xp
import numpy as np
namespace = {"np": np, "numpy": np, "xp": xp, "cp": xp, "cupy": xp}
namespace.update(Simulation._deserialize_function_globals(value.get("globals", {})))
exec(source, namespace) # noqa: S102 -- reconstruct saved Python function
return namespace[value["name"]]
if kind == "python_class":
if value.get("serialization") == "unsupported":
raise ValueError(f"Cannot restore initial-condition class: {value['reason']}.")
if "module" in value:
initial_condition_class = importlib.import_module(value["module"])
for part in value["name"].split("."):
initial_condition_class = getattr(initial_condition_class, part)
else:
source = value["source"]
if hashlib.sha256(source.encode()).hexdigest() != value["source_sha256"]:
raise ValueError("Initial-condition class source hash does not match its metadata.")
import cunumpy as xp
import numpy as np
namespace = {
"np": np,
"numpy": np,
"xp": xp,
"cp": xp,
"cupy": xp,
"Perturbation": Perturbation,
"dataclass": dataclasses.dataclass,
}
exec(source, namespace) # noqa: S102 -- reconstruct saved Python class
initial_condition_class = namespace[value["name"]]
if "params" in value:
return initial_condition_class(**Simulation._deserialize_initial_condition(value["params"]))
initial_condition = initial_condition_class.__new__(initial_condition_class)
initial_condition.__dict__.update(Simulation._deserialize_initial_condition(value["state"]))
return initial_condition
if kind == "callable":
raise ValueError(f"Cannot restore initial-condition callable: {value['reason']}.")
if kind == "FieldsBackground":
return FieldsBackground(**Simulation._deserialize_initial_condition(value["params"]))
from struphy.initial import perturbations
from struphy.kinetic_background import base as kinetic_background_base
from struphy.kinetic_background import maxwellians
params = Simulation._deserialize_initial_condition(value["params"])
for module in (equils, perturbations, maxwellians, kinetic_background_base):
initial_condition_class = getattr(module, kind, None)
if initial_condition_class is not None:
return initial_condition_class(**params)
raise ValueError(f"Unknown initial-condition type '{kind}'.")
def _restore_initial_conditions(self, metadata: dict):
"""Attach metadata initial conditions to the reconstructed model variables."""
version = metadata.get("model", {}).get("initial_conditions_schema_version", 1)
if version != 1:
raise ValueError(f"Unsupported initial-conditions metadata schema version: {version}.")
model_species = metadata.get("model", {}).get("species", {})
definitions = {
species_name: {
name: variable["initial_conditions"]
for name, variable in species.get("variables", {}).items()
if "initial_conditions" in variable
}
for species_name, species in model_species.items()
}
for species_name, variables in definitions.items():
species = self.model.species.get(species_name)
if species is None:
continue
for variable_name, entry in variables.items():
variable = species.variables.get(variable_name)
if variable is None:
continue
variable._backgrounds = self._deserialize_initial_condition(entry["backgrounds"])
variable._perturbations = self._deserialize_initial_condition(entry["perturbations"])
if isinstance(variable, PICVariable):
variable._initial_condition = self._deserialize_initial_condition(entry["initial_condition"])
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.
The configuration snapshot can also be restored by :meth:`from_output`;
variable details and propagator options are recorded for inspection,
while run-specific facts do not restore live simulation state.
Parameters
----------
file_path : str, optional
If given, also write the JSON string to this file.
**extra_data
Additional key/value pairs merged into the config,
e.g. call-specific facts like a start timestamp.
Returns
-------
str
The JSON-encoded simulation metadata.
"""
config = self.to_dict()
config["model"] = self.model.to_dict(initial_condition_serializer=self._serialize_initial_condition)
config.update(
{
"mpi_ranks": self.comm_size,
"use_mpi_comm_world": self.comm is not None,
**extra_data,
},
)
json_str = json.dumps(config, indent=4, cls=CuPyJSONEncoder)
if file_path is not None:
with open(file_path, "w") as f:
f.write(json_str)
return json_str
@classmethod
def from_dict(cls, dct) -> "Simulation":
"""Deserialize a simulation configuration from a dictionary."""
return cls(
name=dct["name"],
description=dct["description"],
model=StruphyModel.from_dict(dct["model"]),
params_path=dct["params_path"],
env=EnvironmentOptions.from_dict(dct["env"]),
time_opts=Time.from_dict(dct["time_opts"]),
domain=domains.Cuboid.from_dict(dct["domain"]),
equil=FluidEquilibrium.from_dict(dct["equil"]),
grid=grids.TensorProductGrid.from_dict(dct["grid"]) if dct["grid"] is not None else None,
derham_opts=DerhamOptions.from_dict(dct["derham_opts"]) if dct["derham_opts"] is not None else None,
profiling_opts=ProfilingOptions(
**{
key: value
for key, value in dct.get("profiling_opts", {}).items()
if key in ProfilingOptions.__dataclass_fields__
}
),
)
@classmethod
def from_file(cls, file_path: str) -> "SimulationBase":
"""Deserialize a simulation configuration from a YAML or JSON file.
Initial conditions in run metadata are restored when present. Embedded
Python functions and classes are reconstructed from their saved source.
"""
file_path = os.fspath(file_path)
if file_path.endswith(".yaml") or file_path.endswith(".yml"):
with open(file_path, "r") as f:
dct = yaml.safe_load(f)
elif file_path.endswith(".json"):
with open(file_path, "r") as f:
dct = json.load(f)
else:
raise ValueError("Unsupported file format. Use .yaml, .yml or .json.")
metadata = copy.deepcopy(dct)
# YAML and JSON do not have a native tuple type,
# so when you load them with PyYAML or json,
# sequences are always converted to lists
def convert_lists_to_tuples(obj):
if isinstance(obj, dict):
for k, v in obj.items():
obj[k] = convert_lists_to_tuples(v)
return obj
elif isinstance(obj, list):
return tuple(convert_lists_to_tuples(i) for i in obj)
else:
return obj
# Convert lists to tuples for relevant keys
dct = convert_lists_to_tuples(dct)
sim = cls.from_dict(dct)
sim._restore_initial_conditions(metadata)
return sim
@classmethod
def from_output(cls, path_out: str) -> "Simulation":
"""Restore the simulation that wrote the output folder ``path_out``.
The configuration is read from the ``run_metadata.json`` written by :meth:`run`; a copied
parameter file is never executed. The metadata holds the options objects and the
arguments of the model (and thus its units), which is all that post-processing and
plotting need. Initial conditions are restored when present, including
embedded Python functions and classes from saved source.
Nothing is allocated, and ``env`` points at ``path_out`` even if the folder was moved.
"""
path_out = os.path.abspath(path_out)
config_path = os.path.join(path_out, "run_metadata.json")
if not os.path.exists(config_path):
raise FileNotFoundError(f"run_metadata.json does not exist in {path_out}; is it a Struphy output folder?")
sim = cls.from_file(config_path)
sim.env = dataclasses.replace(
sim.env, out_folders=os.path.dirname(path_out), sim_folder=os.path.basename(path_out)
)
return sim
def generate_script(
self,
include_main_guard: bool = False,
include_defaults: bool = False,
) -> str:
"""Generate a Python script that can be used to reproduce the simulation."""
script = f"""
from struphy import (
BaseUnits,
DerhamOptions,
EnvironmentOptions,
FieldsBackground,
ProfilingOptions,
Simulation,
Time,
domains,
equils,
grids,
perturbations,
)
from struphy.models import {self.model.__class__.__name__}
"""
sim_setup = ""
sim_class_def = "sim = Simulation("
if include_defaults:
sim_setup += f"model = {self.model.__repr__()}\n"
sim_class_def += "model=model,"
sim_setup += f"env = {self.env.__repr__()}\n"
sim_class_def += "env=env,"
sim_setup += f"time_opts = {self.time_opts.__repr__()}\n"
sim_class_def += "time_opts=time_opts,"
sim_setup += f"domain = domains.{self.domain.__repr__()}\n"
sim_class_def += "domain=domain,"
if self.grid is None:
sim_setup += "grid = None\n"
else:
sim_setup += f"grid = grids.{self.grid.__repr__()}\n"
sim_class_def += "grid=grid,"
if self.derham_opts is None:
sim_setup += "derham_opts = None\n"
else:
sim_setup += f"derham_opts = {self.derham_opts.__repr__()}\n"
sim_class_def += "derham_opts=derham_opts,"
sim_setup += f"profiling_opts = {self.profiling_opts.__repr__()}\n"
sim_class_def += "profiling_opts=profiling_opts,"
else:
# Only include parameters that are not default to avoid
# cluttering the script with unnecessary lines
sim_setup += f"model = {self.model.__repr_no_defaults__()}\n"
sim_class_def += "model=model,"
if not self.env.is_default:
sim_setup += f"env = {self.env.__repr_no_defaults__()}\n"
sim_class_def += "env=env,"
if not self.time_opts.is_default:
sim_setup += f"time_opts = {self.time_opts.__repr_no_defaults__()}\n"
sim_class_def += "time_opts=time_opts,"
if not self.domain.is_default:
sim_setup += f"domain = domains.{self.domain.__repr_no_defaults__()}\n"
sim_class_def += "domain=domain,"
# None (no FEEC Derham) is not the constructor default, so pass it explicitly
if self.grid is None:
sim_setup += "grid = None\n"
sim_class_def += "grid=grid,"
elif not self.grid.is_default:
sim_setup += f"grid = grids.{self.grid.__repr_no_defaults__()}\n"
sim_class_def += "grid=grid,"
if self.derham_opts is None:
sim_setup += "derham_opts = None\n"
sim_class_def += "derham_opts=derham_opts,"
elif not self.derham_opts.is_default:
sim_setup += f"derham_opts = {self.derham_opts.__repr_no_defaults__()}\n"
sim_class_def += "derham_opts=derham_opts,"
if self.profiling_opts != ProfilingOptions():
sim_setup += f"profiling_opts = ProfilingOptions(**{self.profiling_opts.to_kwargs()!r})\n"
sim_class_def += "profiling_opts=profiling_opts,"
# This is a bit of a special case since the default is None,
if self.equil is not None:
if include_defaults:
sim_setup += f"equil = equils.{self.equil.__repr__()}\n"
else:
sim_setup += f"equil = equils.{self.equil.__repr_no_defaults__()}\n"
sim_class_def += "equil=equil,"
if self.params_path is not None:
sim_class_def += f"params_path={repr(self.params_path)},\n"
sim_class_def += ")\n"
# Rebuild from the serialized configuration so post-construction model
# settings (particle parameters and propagator options) are preserved.
script += sim_setup + "\n" + sim_class_def
script += f"\n# Restore the complete model configuration\nsim = Simulation.from_dict({self.to_dict()!r})\n"
if include_main_guard:
script += """
if __name__ == "__main__":
sim.run()"""
return ruff_autofix_and_format(script)
def save_script(
self,
file_path: str,
include_main_guard: bool = False,
):
"""Save the generated script to a file."""
script = self.generate_script(include_main_guard=include_main_guard)
with open(file_path, "w") as f:
f.write(script)
def __eq__(self, value: "Simulation") -> bool:
assert isinstance(value, Simulation), "Comparison only implemented between Simulation instances."
return self.to_dict() == value.to_dict()
# ------------------------------------------------------
# Common properties with setters (from input parameters)
# ------------------------------------------------------
@property
def model(self) -> StruphyModel:
"""StruphyModel object containing the PDE of the model."""
return self._model
@model.setter
def model(self, value: StruphyModel):
assert isinstance(value, StruphyModel)
assert hasattr(value, "propagators"), "Attribute 'self.propagators' must be set in model __init__!"
self._model = value
self._model_name = value.__class__.__name__
@property
def name(self) -> str:
"""Name of the simulation."""
return self._name
@name.setter
def name(self, value: str):
assert isinstance(value, str)
self._name = value
@property
def description(self) -> str:
"""Description of the simulation."""
return self._description
@description.setter
def description(self, value: str):
assert isinstance(value, str)
self._description = value
@property
def params_path(self):
"""Path to parameter file used for the run. Can be None if Simulation is instantiated in a notebook environment (no parameter file in this case)."""
return self._params_path
@params_path.setter
def params_path(self, value: str | None):
if value is not None:
assert isinstance(value, str)
assert value[-3:] == ".py", f"Parameter file must be a Python file, got {value}."
self._params_path = value
@property
def env(self) -> EnvironmentOptions:
"""EnvironmentOptions object containing options related to the environment of the run."""
return self._env
@env.setter
def env(self, value: EnvironmentOptions):
"""Update the environment options for the simulation."""
assert isinstance(value, EnvironmentOptions)
self._env = value
@property
def profiling_filepath(self) -> str:
"""Path to the profiling file, if profiling is enabled."""
return os.path.join(
self.env.out_folders,
self.env.sim_folder,
"profiling_data.h5",
)
@property
def time_opts(self):
"""Time object containing time stepping parameters."""
return self._time_opts
@time_opts.setter
def time_opts(self, value: Time):
assert isinstance(value, Time)
self._time_opts = value
@property
def domain(self):
"""Domain object, see :ref:`avail_mappings`."""
return self._domain
@domain.setter
def domain(self, value: Domain):
"""Set the domain. If an equilibrium is already set, the domain is passed on to it
(unless the equilibrium is numerical, in which case it dictates the domain, see :attr:`equil`)."""
assert isinstance(value, Domain)
equil = getattr(self, "_equil", None)
if isinstance(equil, NumericalMHDequilibrium):
self._domain = equil.domain
else:
self._domain = value
if equil is not None:
equil.domain = value
@property
def equil(self):
"""Fluid equilibrium object, see :ref:`fluid_equil`."""
return self._equil
@equil.setter
def equil(self, value: FluidEquilibrium | None):
"""Set the fluid equilibrium and link it to the domain and the model units.
If a numerical equilibrium is used, the domain is taken from this equilibrium."""
assert value is None or isinstance(value, FluidEquilibrium)
self._equil = value
if value is None:
return
if isinstance(value, NumericalMHDequilibrium):
self._domain = value.domain
else:
value.domain = self.domain
if hasattr(value, "units"):
assert isinstance(value.units, Units)
value.units.derive_units(
velocity_scale=self.model.velocity_scale,
A_bulk=self.model.bulk_species.mass_number,
Z_bulk=self.model.bulk_species.charge_number,
)
@property
def grid(self):
"""Grid object, see :ref:`grids`."""
return self._grid
@grid.setter
def grid(self, value: grids.TensorProductGrid | None):
assert value is None or isinstance(value, grids.TensorProductGrid)
self._grid = value
@property
def derham_opts(self):
"""DerhamOptions object containing options for the setup of the 3d Derham sequence."""
return self._derham_opts
@derham_opts.setter
def derham_opts(self, value: DerhamOptions | None):
assert value is None or isinstance(value, DerhamOptions)
self._derham_opts = value
@property
def profiling_opts(self):
"""Options passed to scope-profiler for profiling sessions."""
return self._profiling_opts
@profiling_opts.setter
def profiling_opts(self, value: ProfilingOptions):
assert isinstance(value, ProfilingOptions)
self._profiling_opts = value
@property
def comm(self):
"""MPI communicator of the run, None if MPI is not available."""
return self._comm
@comm.setter
# NOTE: string annotation, MPI.Intracomm is maybempi's stand-in in serial runs
def comm(self, value: "MPI.Intracomm | None"):
"""Set the MPI communicator; this also updates :attr:`rank`, :attr:`comm_size` and :attr:`Barrier`.
If None is passed, MPI.COMM_WORLD is used (unless MPI is not available)."""
if isinstance(MPI, SerialMPI):
self._comm = None
self._rank = 0
self._comm_size = 1
self._barrier = lambda: None
else:
self._comm = MPI.COMM_WORLD if value is None else value
self._rank = self._comm.Get_rank()
self._comm_size = self._comm.Get_size()
self._barrier = self._comm.Barrier
logger.info(f"\nMPI comm: {self._comm}")
logger.info(f"MPI size: {self._comm_size} processes")
logger.info(f"MPI rank: {self._rank}")
@property
def start_time(self) -> float:
"""Wall-clock time (epoch seconds) at which the simulation was set up."""
return self._start_time
@start_time.setter
def start_time(self, value: float):
assert isinstance(value, float)
self._start_time = value
# -----------------------------------------------------------------
# Common properties (derived from the above properties, no setters)
# -----------------------------------------------------------------
@property
def rank(self) -> int:
"""Rank of the current process in :attr:`comm`."""
return self._rank
@property
def comm_size(self) -> int:
"""Number of processes in :attr:`comm`."""
return self._comm_size
@property
def Barrier(self):
"""Barrier of :attr:`comm` (a no-op if MPI is not available)."""
return self._barrier
@property
def model_name(self) -> str:
"""Class name of the model."""
return self._model_name
@property
def derham(self):
"""3d Derham sequence, see :ref:`derham`; None before :meth:`allocate`."""
return getattr(self, "_derham", None)
@property
def mass_ops(self):
"""WeighteMassOperators object, see :ref:`mass_ops`."""
return getattr(self, "_mass_ops", None)
@property
def basis_ops(self):
"""Basis projection operators."""
return getattr(self, "_basis_ops", None)
@property
def projected_equil(self):
"""Fluid equilibrium projected on 3d Derham sequence with commuting projectors."""
return getattr(self, "_projected_equil", None)
@property
def clone_config(self):
"""Config in case domain clones are used."""
return self._clone_config
@clone_config.setter
def clone_config(self, new):
"""Set the clone config, both here and on the model."""
assert isinstance(new, CloneConfig) or new is None
self._clone_config = new
self.model.clone_config = new
