"""
Input/Output functions.
Extensibility Guide
===================
The I/O system is designed to be fully extensible. You can add arbitrary fields
to `metrics` and `records` without modifying this module.
Adding custom metrics
---------------------
>>> solver.metrics['custom_value'] = 42.0
>>> solver.metrics['convergence_rate'] = 0.95
Adding custom records
---------------------
>>> # Simple array
>>> solver.records['energy_evolution'] = np.array([...])
>>> # Nested structure (automatically flattened in .npz)
>>> solver.records['field_snapshots'] = {
... 'times': np.array([0, 1, 2]),
... 'data': np.array([...])
... }
>>> # Deeply nested (unlimited depth)
>>> solver.records['analysis'] = {
... 'spectra': {
... 'fourier': np.array([...]),
... 'wavelets': np.array([...])
... }
... }
TypedDicts (`Metrics`, `Records`, etc.) are documentation only. Runtime accepts
any dict[str, Any] structure.
"""
from __future__ import annotations
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import (
Any,
NotRequired,
TextIO,
TypedDict,
cast,
)
import numpy as np
from .parameters import RunParameters
from .solvers import rank, size
from .solvers.profiling import ProfilingStats
[docs]
def get_tqdm_file() -> TextIO | None:
"""Get a TQDM-compatible file for progress bar output in MPI."""
if size == 1:
return sys.stdout # single process uses default
if rank != 0:
return None # other ranks disable progress bar
try:
# For MPI, try to open /dev/tty for direct terminal output
return open("/dev/tty", "w")
except OSError:
return sys.stdout # fallback
[docs]
class MetricsRequired(TypedDict):
"""Required fields for simulation metrics."""
total_time: float #: Total wall-clock time of the simulation
time_per_ite: float #: Average time per iteration
throughput: float #: Grid point update rate [point/s]
CFL: float #: CFL condition value
[docs]
class Metrics(MetricsRequired, total=False):
"""
Structure for simulation metrics (extensible).
Required fields: total_time, time_per_ite, throughput, CFL
All other fields are optional and can be added dynamically.
"""
#: Profiling statistics
profiling_stats: NotRequired[ProfilingStats]
[docs]
class Observables(TypedDict, total=False):
"""
Physical observables from the simulation (extensible).
All fields are optional and can be added dynamically.
"""
m1_mean: float #: Time-averaged magnetization in x direction
[docs]
class XProfiles(TypedDict):
"""Structure for cross-sectional profiles (arrays in final records)."""
t: np.ndarray #: Time points for profiles
m1: np.ndarray #: m1 component profiles
m2: np.ndarray #: m2 component profiles
m3: np.ndarray #: m3 component profiles
[docs]
class XProfilesBuffer(TypedDict):
"""Structure for cross-sectional profiles during accumulation (lists)."""
t: list[float] #: Time points for profiles
m1: list[np.ndarray] #: m1 component profiles
m2: list[np.ndarray] #: m2 component profiles
m3: list[np.ndarray] #: m3 component profiles
[docs]
class RecordsBuffer(TypedDict, total=False):
"""
Records during simulation (accumulation phase with lists).
BaseSolver.records uses this structure during simulation, accumulating
data as lists. Before saving, BaseSolver.save() converts lists to arrays.
"""
xyz_average: list[tuple[float, float]] #: Accumulated (t, value) pairs
x_profiles: XProfilesBuffer #: Profiles during accumulation
[docs]
class Records(TypedDict, total=False):
"""
Records after saving (finalized with numpy arrays).
Returned by load_results() with all data as read-only numpy arrays.
"""
xyz_average: np.ndarray #: Space-averaged magnetization over time (shape: (2, N))
x_profiles: XProfiles #: Cross-sectional profiles in yz plane
[docs]
class SimulationResults(TypedDict):
"""Structure for simulation results."""
metrics: Metrics #: Simulation metrics (performance, numerical quality)
observables: NotRequired[Observables] #: Physical observables (results)
records: NotRequired[Records] #: Optional time-series records
[docs]
@dataclass
class RunResults:
"""Structure for loaded simulation results."""
params: RunParameters
results: SimulationResults
file: str | Path
[docs]
def get_record(self, record_name: str) -> np.ndarray | dict[str, np.ndarray]:
"""Return a record by name from results or raise a descriptive error."""
try:
if "records" not in self.results:
raise KeyError(f"RunResults from '{self.file}' has no records.")
records = self.results["records"]
# Cast to dict to allow dynamic access with a variable key for mypy
record = cast(dict, records)[record_name]
if isinstance(record, (np.ndarray, dict)):
return record
else:
raise TypeError(f"'{record_name}' is not in the expected format.")
except KeyError as e:
msg = (
f"RunResults from '{self.file}' does not contain the required record "
f"'{record_name}'."
)
raise KeyError(msg) from e
[docs]
def save_results(
output_file: str | Path,
params: RunParameters,
metrics: Metrics,
observables: Observables | None = None,
records_buffer: RecordsBuffer | None = None,
) -> None:
"""
Saves simulation results to a .npz file with hierarchical structure.
Args:
output_file: Path to the output .npz file.
params: Dataclass of simulation parameters.
metrics: Dictionary of simulation metrics. Must contain required fields:
total_time, time_per_ite, CFL. Additional fields are allowed.
observables: Dictionary of physical observables (results).
records_buffer: RecordsBuffer or custom dict of records
(accumulation phase with lists/dicts).
Note:
Both metrics and records accept any key-value pairs, allowing easy
extension without modifying this function. See Metrics and Records
TypedDict for recommended structure.
Example:
>>> records_buffer = {
... 'xyz_average': arr1,
... 'x_profiles': {'t': t, 'm1': m1, ...},
... 'custom_data': arr2 # Any new field works automatically
... }
>>> save_results('run.npz', params, metrics, records_buffer=records_buffer)
"""
data_to_save: dict[str, Any] = {}
if records_buffer is None:
records_buffer = {}
if observables is None:
observables = {}
# Helper to convert lists to numpy arrays recursively
def _convert_lists_to_arrays(d: Any) -> Any:
"""Recursively convert lists to numpy arrays."""
if isinstance(d, list):
return np.array(d)
elif isinstance(d, dict):
return {k: _convert_lists_to_arrays(v) for k, v in d.items()}
else:
return d
# Convert lists in records_buffer to arrays
records_buffer = _convert_lists_to_arrays(records_buffer)
# Helper to flatten nested dicts into NPZ keys
def _flatten_dict(d: dict | Any, prefix: str):
"""Recursively flatten nested dicts into NPZ keys."""
if isinstance(d, dict):
for key, value in d.items():
full_key = f"{prefix}/{key}"
_flatten_dict(value, full_key)
else:
# Leaf value - save as array
data_to_save[prefix] = d
# Flatten params
_flatten_dict(params.as_dict(), "params")
# Flatten metrics
_flatten_dict(metrics, "results/metrics")
# Flatten observables
if observables:
_flatten_dict(observables, "results/observables")
# Flatten records
_flatten_dict(records_buffer, "results/records")
np.savez(output_file, **data_to_save, allow_pickle=False)
[docs]
def load_results(result_file: str | Path) -> RunResults:
"""
Loads simulation results from a .npz file with hierarchical structure.
Numpy arrays are set to read-only to prevent accidental modification.
Example:
>>> run_results = load_results("run.npz")
>>> params = run_results["params"]
>>> metrics = run_results["results"]["metrics"]
>>> xyz_average = run_results["results"]["records"]["xyz_average"]
>>> m1_prof = run_results["results"]["records"]["x_profiles"]["m1"]
Args:
result_file: path to the .npz file.
Returns:
Dictionary containing the loaded data with hierarchical structure.
"""
npz_data: np.lib.npyio.NpzFile = np.load(result_file, allow_pickle=False)
def _unflatten_dict(prefix: str) -> dict[str, Any]:
"""Reconstruct nested dict from flat NPZ keys with given prefix."""
result: dict[str, Any] = {}
for key in npz_data.files:
if not key.startswith(prefix + "/"):
continue
# Extract relative path after prefix
rel_path = key[len(prefix) + 1 :]
parts = rel_path.split("/")
# Navigate/create nested structure
current = result
for part in parts[:-1]:
if part not in current:
current[part] = {}
current = current[part]
# Store the value
arr = npz_data[key]
arr.flags.writeable = False
current[parts[-1]] = arr
return result
# Load params
params_dict = _unflatten_dict("params")
params = RunParameters(**params_dict)
# Load metrics
metrics: Metrics = cast(Metrics, _unflatten_dict("results/metrics"))
# Load observables
observables: Observables | None = None
observables_dict = _unflatten_dict("results/observables")
if observables_dict:
observables = cast(Observables, observables_dict)
# Load records
records: Records = cast(Records, _unflatten_dict("results/records"))
# Build the results dict
simulation_results: SimulationResults = {"metrics": metrics}
if observables:
simulation_results["observables"] = observables
if records:
simulation_results["records"] = records
run_results = RunResults(
params=params, results=simulation_results, file=result_file
)
return run_results