nmag-python-3 0.0.2__py3-none-any.whl

si/py.typed

@@ -0,0 +1 @@
+# Marker for PEP 561 inline type information.

simulation/__init__.py

@@ -0,0 +1 @@
+"""Simulation support modules for the nmag Python 3 rewrite."""

simulation/clock.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from math import isfinite
+from typing import Any
+
+from tabulate import tabulate
+
+from si.physical import SI
+
+_zero_seconds_cache: SI | None = None
+
+
+def _zero_seconds() -> SI:
+    global _zero_seconds_cache
+    if _zero_seconds_cache is None:
+        _zero_seconds_cache = SI(0.0, "s")
+    return SI(_zero_seconds_cache._quantity)
+
+
+def fmt_time(t: SI, fmt_ps: str = ".2f", fmt_ns: str = ".2f") -> str:
+    """Formats an SI time object into picoseconds or nanoseconds."""
+    t_ps = float(t / SI(1e-12, "s"))
+
+    ps_str = f"{t_ps:{fmt_ps}}"
+    ns_str = f"{(t_ps / 1000.0):{fmt_ns}}"
+
+    return f"{ps_str} ps" if t_ps < 100.0 else f"{ns_str} ns"
+
+
+@dataclass
+class SimulationClock:
+    """
+    This object specifies all the parameters which define the current time
+    in the simulation, such as the simulation time, step number, ...
+
+    Attributes:
+      id: Unique identifier for data saved. Incremented on save.
+      stage: Stage number. Increments when the external field changes.
+      step: Total number of steps performed (always increases).
+      stage_step: Step number from the beginning of the current stage.
+      zero_stage_step: The value of 'step' at the beginning of the stage.
+      time: Total simulation time (always increases).
+      stage_time: The simulation time from the beginning of the stage.
+      zero_stage_time: The value of 'time' at the beginning of the stage.
+      real_time: The real world time used for advancing time.
+      last_step_dt_si: Last time step's length in SI units.
+      convergence: Flag indicating if convergence is reached.
+      stage_end: Flag indicating the end of a stage.
+      exit_hysteresis: Flag to signal exit from hysteresis loop.
+    """
+
+    id: int = -1
+    stage: int = 1
+    step: int = 0
+    time: SI = field(default_factory=_zero_seconds)
+    stage_step: int = 0
+    stage_time: SI = field(default_factory=_zero_seconds)
+    real_time: SI = field(default_factory=_zero_seconds)
+    stage_end: bool = False
+    convergence: bool = False
+    exit_hysteresis: bool = False
+    zero_stage_time: SI = field(default_factory=_zero_seconds)
+    zero_stage_step: int = 0
+    time_reached_su: float = 0.0
+    time_reached_si: SI = field(default_factory=_zero_seconds)
+    last_step_dt_su: float = 0.0
+    last_step_dt_si: SI = field(default_factory=_zero_seconds)
+
+    def __getitem__(self, key: str, /) -> Any:
+        if key not in self.__dataclass_fields__:
+            raise KeyError(key)
+        return getattr(self, key)
+
+    def get(self, key: str, default: Any = None, /) -> Any:
+        try:
+            return self[key]
+        except KeyError:
+            return default
+
+    def copy(self) -> dict[str, Any]:
+        return asdict(self)
+
+    def checkpoint_state(self) -> dict[str, int | float | bool]:
+        """Return a JSON-safe representation of the simulation clock."""
+        return {
+            "id": self.id,
+            "stage": self.stage,
+            "step": self.step,
+            "stage_step": self.stage_step,
+            "zero_stage_step": self.zero_stage_step,
+            "time_seconds": self.time.in_units_of(SI(1.0, "s")),
+            "stage_time_seconds": self.stage_time.in_units_of(SI(1.0, "s")),
+            "real_time_seconds": self.real_time.in_units_of(SI(1.0, "s")),
+            "zero_stage_time_seconds": self.zero_stage_time.in_units_of(SI(1.0, "s")),
+            "time_reached_seconds": self.time_reached_si.in_units_of(SI(1.0, "s")),
+            "last_step_seconds": self.last_step_dt_si.in_units_of(SI(1.0, "s")),
+            "stage_end": self.stage_end,
+            "convergence": self.convergence,
+            "exit_hysteresis": self.exit_hysteresis,
+        }
+
+    @classmethod
+    def from_checkpoint_state(cls, state: dict[str, object]) -> SimulationClock:
+        """Build a validated clock from :meth:`checkpoint_state` data."""
+        integer_names = ("id", "stage", "step", "stage_step", "zero_stage_step")
+        seconds_names = (
+            "time_seconds",
+            "stage_time_seconds",
+            "real_time_seconds",
+            "zero_stage_time_seconds",
+            "time_reached_seconds",
+            "last_step_seconds",
+        )
+        boolean_names = ("stage_end", "convergence", "exit_hysteresis")
+        required = {*integer_names, *seconds_names, *boolean_names}
+        missing = sorted(required.difference(state))
+        if missing:
+            raise ValueError(f"Checkpoint clock state is missing {', '.join(missing)}.")
+
+        integers: dict[str, int] = {}
+        for name in integer_names:
+            value = state[name]
+            if isinstance(value, bool) or not isinstance(value, int):
+                raise ValueError(f"Checkpoint clock {name} must be an integer.")
+            integers[name] = value
+        if integers["stage"] < 1 or any(
+            integers[name] < 0 for name in ("step", "stage_step", "zero_stage_step")
+        ):
+            raise ValueError("Checkpoint clock counters must be non-negative and stage at least one.")
+
+        seconds: dict[str, float] = {}
+        for name in seconds_names:
+            value = state[name]
+            if isinstance(value, bool) or not isinstance(value, (int, float)):
+                raise ValueError(f"Checkpoint clock {name} must be a finite number.")
+            converted = float(value)
+            if not isfinite(converted) or converted < 0.0:
+                raise ValueError(f"Checkpoint clock {name} must be a non-negative finite number.")
+            seconds[name] = converted
+
+        booleans: dict[str, bool] = {}
+        for name in boolean_names:
+            value = state[name]
+            if not isinstance(value, bool):
+                raise ValueError(f"Checkpoint clock {name} must be a boolean.")
+            booleans[name] = value
+
+        return cls(
+            id=integers["id"],
+            stage=integers["stage"],
+            step=integers["step"],
+            stage_step=integers["stage_step"],
+            zero_stage_step=integers["zero_stage_step"],
+            time=SI(seconds["time_seconds"], "s"),
+            stage_time=SI(seconds["stage_time_seconds"], "s"),
+            real_time=SI(seconds["real_time_seconds"], "s"),
+            zero_stage_time=SI(seconds["zero_stage_time_seconds"], "s"),
+            time_reached_su=seconds["time_reached_seconds"],
+            time_reached_si=SI(seconds["time_reached_seconds"], "s"),
+            last_step_dt_su=seconds["last_step_seconds"],
+            last_step_dt_si=SI(seconds["last_step_seconds"], "s"),
+            stage_end=booleans["stage_end"],
+            convergence=booleans["convergence"],
+            exit_hysteresis=booleans["exit_hysteresis"],
+        )
+
+    # __init__ and __repr__ are GONE (auto-generated)
+
+    def inc_stage(self, stage: int | None = None) -> None:
+        """Advance the clock to the next stage."""
+        if stage is None:
+            self.stage += 1
+        else:
+            self.stage = stage
+        self.stage_step = 0
+        self.stage_time = SI(0.0, "s")
+        self.convergence = False
+        self.zero_stage_step = self.step
+        self.zero_stage_time = self.time
+
+    def record_advance(
+        self,
+        *,
+        stage_time_seconds: float,
+        accepted_steps: int,
+        last_step_seconds: float,
+        wall_seconds: float,
+    ) -> None:
+        """Atomically record an accepted group of integration steps."""
+        previous_stage_seconds = self.stage_time.in_units_of(SI(1.0, "s"))
+        delta_seconds = stage_time_seconds - previous_stage_seconds
+        if delta_seconds < 0.0:
+            raise ValueError("stage time cannot move backwards.")
+        if accepted_steps < 0:
+            raise ValueError("accepted_steps cannot be negative.")
+
+        self.step += accepted_steps
+        self.stage_step += accepted_steps
+        self.time += SI(delta_seconds, "s")
+        self.stage_time = SI(stage_time_seconds, "s")
+        self.time_reached_su = stage_time_seconds
+        self.time_reached_si = SI(stage_time_seconds, "s")
+        self.last_step_dt_su = last_step_seconds
+        self.last_step_dt_si = SI(last_step_seconds, "s")
+        self.real_time += SI(max(0.0, wall_seconds), "s")
+
+    # This method was updated to use tabulate, the format of the data printed out might
+    # look slightly different, but there is a lot less manual formatting code here now.
+    def __str__(self) -> str:
+        ft = fmt_time
+
+        rows = [
+            [
+                f"ID={self.id}",
+                f"Step={self.step}",
+                f"Time={ft(self.time)}",
+                f"Last step size={ft(self.last_step_dt_si)}",
+            ],
+            [
+                "",
+                f"Stage={self.stage}",
+                f"Stage-step={self.stage_step}",
+                f"Stage-time={ft(self.stage_time)}",
+            ],
+            [
+                "",
+                f"Convergence={self.convergence}",
+                f"Stage-end={self.stage_end}",
+                f"Exit hysteresis={self.exit_hysteresis}",
+            ],
+        ]
+
+        table = tabulate(rows, tablefmt="pipe")
+
+        sep_line = "=" * (len(table.splitlines()[0]))
+        return f"{sep_line}\n{table}\n{sep_line}"

simulation/data_writer.py

@@ -0,0 +1,273 @@
+from __future__ import annotations
+
+import csv
+import logging
+import time
+from collections.abc import Generator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, cast
+
+from si.physical import SI
+from simulation.data_writer_collection import (
+    DataWriterCollectionMixin,
+    SimulationSource,
+    _SubfieldRequest,
+)
+from simulation.quantity import Quantity, known_quantities
+
+log = logging.getLogger("nmag")
+
+
+class DataWriter(DataWriterCollectionMixin):
+    def __init__(self, ndt_filename: Path, h5_filename: Path, *, append: bool = False) -> None:
+        """
+        :param ndt_filename: Path to the tabular output (.ndt).
+        :param h5_filename: Path to the spatial output (.h5).
+        """
+        self.ndt_filename = ndt_filename
+        self.h5_filename = h5_filename
+        self.quantities = known_quantities
+
+        self.quantities_by_name: dict[str, Quantity] = {q.name: q for q in known_quantities}
+
+        self._header_written: bool = False
+        self._column_names: list[str] | None = None
+        self._column_index_by_name: dict[str, int] | None = None
+        self._column_units: dict[str, str] | None = None
+        self._subfield_requests: list[_SubfieldRequest] = []
+        self._writes_maxangle: bool = False
+        self._last_saved_step: int = -1
+        self.last_save_timings_seconds: dict[str, float] = {}
+        self._active_save_timings: dict[str, float] | None = None
+        self._append_schema_pending = False
+        if append and self.ndt_filename.exists():
+            self._recover_ndt_schema()
+            self._append_schema_pending = True
+
+    def _recover_ndt_schema(self) -> None:
+        """Recover the schema and last step from an NDT file opened for append."""
+
+        with open(self.ndt_filename, newline="", encoding="utf-8") as stream:
+            rows = [row for row in csv.reader(stream, delimiter="\t") if row]
+        if len(rows) < 2 or not rows[0][0].startswith("# Simulation:"):
+            raise ValueError(f"Cannot append to malformed NDT output: {self.ndt_filename}.")
+        column_names = rows[1]
+        if len(column_names) != len(set(column_names)):
+            raise ValueError(f"Cannot append to NDT output with duplicate columns: {self.ndt_filename}.")
+        self._column_names = column_names
+        self._column_index_by_name = {name: index for index, name in enumerate(column_names)}
+        self._column_units = {}
+        self._header_written = True
+        step_index = self._column_index_by_name.get("step")
+        if step_index is not None:
+            for row in reversed(rows[2:]):
+                if len(row) != len(column_names):
+                    continue
+                try:
+                    self._last_saved_step = int(float(row[step_index]))
+                except ValueError:
+                    continue
+                break
+
+    def save(
+        self,
+        source: SimulationSource,
+        fields: str | list[str] | None = None,
+        avoid_same_step: bool = False,
+    ) -> None:
+
+        current_step = source.step
+        if avoid_same_step and current_step == self._last_saved_step:
+            return
+
+        timings: dict[str, float] = {}
+        self._active_save_timings = timings
+        total_started = time.perf_counter()
+        try:
+            with self._record_timing("write_ndt_row"):
+                self._write_ndt_row(source)
+
+            if fields is not None:
+                with self._record_timing("save_spatial_fields"):
+                    self._trigger_field_save(source, fields)
+        finally:
+            timings["total"] = time.perf_counter() - total_started
+            self.last_save_timings_seconds = dict(sorted(timings.items()))
+            self._active_save_timings = None
+
+        self._last_saved_step = current_step
+
+    @contextmanager
+    def _record_timing(self, name: str) -> Generator[None]:
+        started = time.perf_counter()
+        try:
+            yield
+        finally:
+            timings = self._active_save_timings
+            if timings is not None:
+                timings[name] = timings.get(name, 0.0) + time.perf_counter() - started
+
+    def _trigger_field_save(self, source: SimulationSource, fields: str | list[str]) -> None:
+        with self._record_timing("spatial_detail:field_selection"):
+            field_names_to_save: list[str] = []
+
+            if fields == "all":
+                field_names_to_save = source.get_all_field_names()
+            elif isinstance(fields, list):
+                field_names_to_save = fields
+            else:
+                raise ValueError(f"Invalid fields argument: {fields}")
+
+        if field_names_to_save:
+            with self._record_timing("spatial_detail:write_dispatch"):
+                source.save_spatial_fields(
+                    filename=str(self.h5_filename), fieldnames=field_names_to_save
+                )
+            self._record_source_spatial_save_timings(source)
+
+    def _write_ndt_row(self, source: SimulationSource) -> None:
+        quantities: list[Quantity] = []
+        with self._record_timing("gather_ndt_columns"):
+            if self._append_schema_pending:
+                with self._record_timing("gather_ndt_append_schema_columns"):
+                    columns, _quantities = self._gather_data(source)
+            elif self._header_written:
+                with self._record_timing("gather_ndt_existing_schema_columns"):
+                    columns = self._gather_existing_schema_data(source)
+                quantities: list[Quantity] = []
+            else:
+                with self._record_timing("gather_ndt_initial_schema_columns"):
+                    columns, quantities = self._gather_data(source)
+
+        if not self._header_written:
+            with self._record_timing("build_ndt_header_schema"):
+                col_names: list[str] = []
+                col_units: dict[str, str] = {}
+
+                for (name, _), qty in zip(columns, quantities, strict=True):
+                    col_names.append(name)
+                    if qty.units:
+                        col_units[name] = qty.units.dens_str()
+                    else:
+                        col_units[name] = "-"
+
+                column_index_by_name = {name: index for index, name in enumerate(col_names)}
+            with self._record_timing("format_ndt_values"):
+                row_values: list[Any] = [self._ndt_cell_value(value) for _, value in columns]
+
+            with open(self.ndt_filename, "w", newline="", encoding="utf-8") as f:
+                writer = csv.writer(f, delimiter="\t")
+                with self._record_timing("write_ndt_header"):
+                    f.write(f"# Simulation: {source.name}\n")
+                    writer.writerow(col_names)
+
+                with self._record_timing("write_ndt_values"):
+                    writer.writerow(row_values)
+
+            self._column_names = col_names
+            self._column_index_by_name = column_index_by_name
+            self._column_units = col_units
+            self._header_written = True
+            return
+
+        if self._column_names is None:
+            log.error("Column names not initialized.")
+            return
+
+        with self._record_timing("validate_ndt_schema"):
+            actual_names = [name for name, _ in columns]
+            if actual_names != self._column_names:
+                raise ValueError(
+                    "Cannot append output with a different NDT schema. "
+                    f"Expected {self._column_names!r}, got {actual_names!r}."
+                )
+            self._append_schema_pending = False
+
+        with self._record_timing("format_ndt_values"):
+            row_values = self._existing_schema_row_values(columns)
+
+        with self._record_timing("write_ndt_values"):
+            with open(self.ndt_filename, "a", newline="", encoding="utf-8") as f:
+                writer = csv.writer(f, delimiter="\t")
+                writer.writerow(row_values)
+
+    def _existing_schema_row_values(
+        self,
+        columns: list[tuple[str, Any]],
+    ) -> list[Any]:
+        if self._column_names is None:
+            return []
+        if self._column_index_by_name is None:
+            self._column_index_by_name = {
+                name: index for index, name in enumerate(self._column_names)
+            }
+
+        row_values: list[Any] = [None] * len(self._column_names)
+        for name, value in columns:
+            index = self._column_index_by_name.get(name)
+            if index is not None:
+                row_values[index] = self._ndt_cell_value(value)
+        return row_values
+
+    @staticmethod
+    def _ndt_cell_value(value: Any) -> Any:
+        if isinstance(value, SI):
+            return value.magnitude
+        return value
+
+    def _record_source_average_timings(
+        self,
+        source: SimulationSource,
+        prefix: str,
+    ) -> None:
+        timings = self._active_save_timings
+        if timings is None:
+            return
+        source_timings = getattr(
+            source,
+            "last_subfield_average_timings_seconds",
+            None,
+        )
+        if not isinstance(source_timings, dict):
+            return
+        for name, seconds in cast(dict[object, object], source_timings).items():
+            if isinstance(seconds, (int, float)):
+                timing_name = f"average_detail:{prefix}:{name!s}"
+                timings[timing_name] = timings.get(timing_name, 0.0) + float(seconds)
+
+    def _record_source_spatial_save_timings(self, source: SimulationSource) -> None:
+        timings = self._active_save_timings
+        if timings is None:
+            return
+        source_timings = getattr(
+            source,
+            "last_spatial_save_timings_seconds",
+            None,
+        )
+        if not isinstance(source_timings, dict):
+            return
+        for name, seconds in cast(dict[object, object], source_timings).items():
+            if isinstance(seconds, (int, float)):
+                timing_name = f"spatial_detail:{name!s}"
+                timings[timing_name] = timings.get(timing_name, 0.0) + float(seconds)
+
+    def _record_source_maxangle_timings(
+        self,
+        source: SimulationSource,
+        prefix: str,
+    ) -> None:
+        timings = self._active_save_timings
+        if timings is None:
+            return
+        source_timings = getattr(
+            source,
+            "last_maxangle_timings_seconds",
+            None,
+        )
+        if not isinstance(source_timings, dict):
+            return
+        for name, seconds in cast(dict[object, object], source_timings).items():
+            if isinstance(seconds, (int, float)):
+                timing_name = f"average_detail:{prefix}:{name!s}"
+                timings[timing_name] = timings.get(timing_name, 0.0) + float(seconds)