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

nmag/simulation/dynamics/integrator.py

@@ -0,0 +1,263 @@
+from __future__ import annotations
+
+import sys
+import time
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Protocol, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+from si.physical import SI
+
+from ...backends import _selected_llg_backend
+from ...dynamics import (
+    DEFAULT_STOPPING_DM_DT,
+    ConvergenceTracker,
+    IntegratorConfig,
+    IntegratorStats,
+)
+from ..support import _si_unit, _simulation_compatibility_binding
+from .advance import SimulationTimeAdvanceMixin
+
+FloatArray = NDArray[np.float64]
+Action = Callable[[Any], Any]
+
+
+class _DenseOutput(Protocol):
+    def __call__(self, time_seconds: float) -> FloatArray: ...
+
+
+class _OdeIntegrator(Protocol):
+    t: float
+    y: FloatArray
+    status: str
+
+    def step(self) -> str | None: ...
+
+    def dense_output(self) -> _DenseOutput: ...
+
+
+def _dop853_class() -> Callable[..., _OdeIntegrator]:
+    from scipy.integrate import DOP853
+
+    return cast(Callable[..., _OdeIntegrator], DOP853)
+
+
+def _dop853_constructor() -> Callable[..., _OdeIntegrator]:
+    """Honor the historical module-level test and extension override."""
+    dynamics_module = sys.modules.get("nmag.simulation.dynamics")
+    return getattr(dynamics_module, "_dop853_class", _dop853_class)()
+
+
+class SimulationIntegratorSetupMixin:
+    if TYPE_CHECKING:
+        _integrator: _OdeIntegrator | None
+        _integrator_is_stale: bool
+        _integrator_rhs_evaluations: int
+        _integrator_config: IntegratorConfig
+        _integrator_effective_max_step_seconds: float
+        _last_integrator_stats: IntegratorStats
+        _stage_wall_started: float
+
+        def __getattr__(self, name: str) -> Any: ...
+
+    def _initialise_dynamics(self) -> None:
+        self._integrator = None
+        self._integrator_is_stale = True
+        self._integrator_rhs_evaluations = 0
+        self._integrator_config = IntegratorConfig()
+        self._integrator_effective_max_step_seconds = self._integrator_config.maximum_step_seconds
+        self._last_integrator_stats = IntegratorStats()
+        self._stage_wall_started = time.perf_counter()
+        self.stopping_dm_dt = DEFAULT_STOPPING_DM_DT
+        self.max_dm_dt: float | None = None
+        self.convergence = ConvergenceTracker()
+        self.max_time_reached = SI(1.0, "s")
+        self._restarting = False
+        self.action_abbreviations: dict[str, Action] = {
+            # Legacy relaxation emits separate average and field rows at stage end.
+            "save_averages": lambda sim: sim.save_data(),
+            "save_fields": lambda sim: sim.save_data(fields="all"),
+            "save_restart": lambda sim: sim.save_restart_file(),
+            "do_next_stage": self.hysteresis_next_stage,
+            "do_exit": self.hysteresis_exit,
+        }
+
+    @property
+    def last_integrator_stats(self) -> IntegratorStats:
+        return self._last_integrator_stats
+
+    @property
+    def integrator_config(self) -> IntegratorConfig:
+        config = self._integrator_config
+        return IntegratorConfig(
+            relative_tolerance=config.relative_tolerance,
+            absolute_tolerance=config.absolute_tolerance,
+            maximum_step_seconds=config.maximum_step_seconds,
+            exact_tstop=config.exact_tstop,
+        )
+
+    @property
+    def effective_integrator_max_step(self) -> SI:
+        return SI(self._integrator_effective_max_step_seconds, "s")
+
+    def _invalidate_integrator(self) -> None:
+        self._integrator_is_stale = True
+
+    def set_params(
+        self,
+        stopping_dm_dt: SI | float | None = None,
+        ts_rel_tol: float | None = None,
+        ts_abs_tol: float | None = None,
+        exact_tstop: bool | None = None,
+        ts_max_step: SI | float | None = None,
+    ) -> None:
+        """Update convergence and DOP853 integration controls.
+
+        Args:
+            stopping_dm_dt: Positive convergence rate in 1/s.
+            ts_rel_tol: Positive relative error tolerance.
+            ts_abs_tol: Positive absolute error tolerance.
+            exact_tstop: Whether :meth:`advance_time` normally reconstructs
+                state exactly at its requested target.
+            ts_max_step: Positive configured maximum step in seconds. Exchange
+                stability can impose a smaller effective maximum.
+
+        Raises:
+            ValueError: If a supplied rate, tolerance, or step is not finite
+                and positive.
+        """
+        if stopping_dm_dt is not None:
+            value = (
+                stopping_dm_dt.in_units_of(_si_unit("1/s"))
+                if isinstance(stopping_dm_dt, SI)
+                else float(stopping_dm_dt)
+            )
+            if not np.isfinite(value) or value <= 0.0:
+                raise ValueError("stopping_dm_dt must be a positive finite rate.")
+            self.stopping_dm_dt = value
+        if ts_rel_tol is not None:
+            self._integrator_config.relative_tolerance = float(ts_rel_tol)
+        if ts_abs_tol is not None:
+            self._integrator_config.absolute_tolerance = float(ts_abs_tol)
+        if exact_tstop is not None:
+            self._integrator_config.exact_tstop = bool(exact_tstop)
+        if ts_max_step is not None:
+            self._integrator_config.maximum_step_seconds = (
+                ts_max_step.in_units_of(_si_unit("s"))
+                if isinstance(ts_max_step, SI)
+                else float(ts_max_step)
+            )
+        self._integrator_config.validate()
+        self._invalidate_integrator()
+
+    def _llg_time_derivative(self, _time_seconds: float, state: FloatArray) -> FloatArray:
+        self._integrator_rhs_evaluations += 1
+        m: FloatArray = np.reshape(np.asarray(state, dtype=np.float64), (-1, 3))
+        if not np.all(np.isfinite(m)):
+            raise FloatingPointError("The integrator supplied a non-finite magnetisation state.")
+
+        self._fields["m"] = m
+        self._invalidate_demag()
+        h_total = self._subfield_array("H_total")
+        pin = self._subfield_array("pin")
+        unit_scale = np.ones(len(m), dtype=float)
+        coefficients = self._nodal_material_coefficients()
+        dm_dcurrent = self._get_dm_dcurrent() if "current_density" in self._fields else None
+        backend = _simulation_compatibility_binding("_selected_llg_backend", _selected_llg_backend)(
+            len(m), getattr(self, "config", None)
+        )
+        with np.errstate(over="ignore", invalid="ignore"):
+            if backend == "rust":
+                derivative = self._llg_rhs_rust(
+                    m,
+                    h_total,
+                    pin,
+                    unit_scale,
+                    coefficients.precession,
+                    coefficients.damping,
+                    coefficients.normalisation,
+                    dm_dcurrent,
+                    coefficients.stt_adiabatic,
+                    coefficients.stt_nonadiabatic,
+                )
+            else:
+                derivative = self._llg_rhs_python(
+                    m,
+                    h_total,
+                    pin,
+                    unit_scale,
+                    coefficients.precession,
+                    coefficients.damping,
+                    coefficients.normalisation,
+                    dm_dcurrent,
+                    coefficients.stt_adiabatic,
+                    coefficients.stt_nonadiabatic,
+                )
+        if not np.all(np.isfinite(derivative)):
+            raise FloatingPointError("The LLG right-hand side produced non-finite values.")
+        return np.ravel(np.asarray(derivative, dtype=np.float64))
+
+    def reinitialise(
+        self,
+        rel_tolerance: float | None = None,
+        abs_tolerance: float | None = None,
+        initial_time: SI | float | None = None,
+    ) -> None:
+        """Rebuild the adaptive integrator from the current physical state.
+
+        Args:
+            rel_tolerance: Optional new relative tolerance.
+            abs_tolerance: Optional new absolute tolerance.
+            initial_time: Non-negative physical start time. Omit it to retain
+                the simulation clock.
+
+        Raises:
+            RuntimeError: If mesh or magnetization has not been set.
+            ValueError: If the time or tolerances are invalid.
+        """
+        if rel_tolerance is not None or abs_tolerance is not None:
+            self.set_params(ts_rel_tol=rel_tolerance, ts_abs_tol=abs_tolerance)
+        if self.mesh is None or "m" not in self._fields:
+            raise RuntimeError("A mesh and magnetisation are required before time integration.")
+
+        if initial_time is None:
+            initial_seconds = self.clock.time_reached_si.in_units_of(_si_unit("s"))
+        elif isinstance(initial_time, SI):
+            initial_seconds = initial_time.in_units_of(_si_unit("s"))
+        else:
+            initial_seconds = float(initial_time)
+        if not np.isfinite(initial_seconds) or initial_seconds < 0.0:
+            raise ValueError("initial_time must be a finite non-negative time.")
+
+        self._integrator_config.validate()
+        state: FloatArray = np.array(
+            np.ravel(np.asarray(self._fields["m"], dtype=np.float64)),
+            dtype=np.float64,
+            copy=True,
+        )
+        upper_bound = self.max_time_reached.in_units_of(_si_unit("s"))
+        if upper_bound <= initial_seconds:
+            upper_bound = float(np.nextafter(initial_seconds, np.inf))
+        self._integrator_effective_max_step_seconds = min(
+            self._integrator_config.maximum_step_seconds,
+            self._exchange_explicit_step_limit_seconds(),
+        )
+        self._integrator = _dop853_constructor()(
+            self._llg_time_derivative,
+            initial_seconds,
+            state,
+            upper_bound,
+            rtol=self._integrator_config.relative_tolerance,
+            atol=self._integrator_config.absolute_tolerance,
+            max_step=self._integrator_effective_max_step_seconds,
+        )
+        self.clock.time_reached_su = initial_seconds
+        self.clock.time_reached_si = SI(initial_seconds, "s")
+        self.clock.stage_time = SI(initial_seconds, "s")
+        self._integrator_is_stale = False
+
+
+class SimulationIntegratorMixin(SimulationIntegratorSetupMixin, SimulationTimeAdvanceMixin):
+    """Compose adaptive-integrator setup with its stepping lifecycle."""

nmag/simulation/exchange/__init__.py

@@ -0,0 +1,13 @@
+"""Exchange and LLG compatibility façade composed from focused mixins."""
+
+from .coefficients import SimulationExchangeCoefficientMixin
+from .fields import SimulationExchangeFieldMixin
+from .llg_rhs import SimulationLlgRhsMixin
+
+
+class SimulationExchangeMixin(
+    SimulationExchangeCoefficientMixin,
+    SimulationLlgRhsMixin,
+    SimulationExchangeFieldMixin,
+):
+    """Provide material coefficients, LLG kernels, and exchange/STT fields."""

nmag/simulation/exchange/coefficients.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from si.physical import SI
+
+from ...dynamics import DOP853_EXCHANGE_RESOLUTION_RADIUS, NodalMaterialCoefficients
+from ..support import MU0, _si_dimensionless, _si_unit
+
+
+class SimulationExchangeCoefficientMixin:
+    if TYPE_CHECKING:
+        _exchange_spectral_bound_cache: tuple[int, float] | None
+        _nodal_material_coefficients_cache: tuple[int, NodalMaterialCoefficients] | None
+
+        def __getattr__(self, name: str) -> Any: ...
+
+    @staticmethod
+    def _material_dynamic_coefficients(
+        material: Any,
+    ) -> tuple[float, float, float, float, float, float]:
+        exchange = material.exchange_coupling.in_units_of(_si_unit("J/m"))
+        ms = material.Ms.in_units_of(_si_unit("A/m"))
+        exchange_prefactor = 0.0 if ms == 0.0 else 2.0 * exchange / (MU0 * ms)
+        gamma = material.llg_gamma_G.in_units_of(_si_unit("m/A/s"))
+        damping = (
+            material.llg_damping.in_units_of(_si_dimensionless())
+            if isinstance(material.llg_damping, SI)
+            else float(material.llg_damping)
+        )
+        gilbert_to_ll = 1.0 / (1.0 + damping * damping)
+        precession = -gamma * gilbert_to_ll
+        damping_coeff = precession * damping
+        if material.do_precession is False:
+            precession = 0.0
+        normalisation = material.llg_normalisationfactor.in_units_of(_si_unit("1/s"))
+        stt_adiabatic = material.su_llg_stt_adiab.in_units_of(_si_unit("m^3/A/s"))
+        stt_nonadiabatic = material.su_llg_stt_nadiab.in_units_of(_si_unit("m^3/A/s"))
+        return (
+            exchange_prefactor,
+            precession,
+            damping_coeff,
+            normalisation,
+            stt_adiabatic,
+            stt_nonadiabatic,
+        )
+
+    def _nodal_material_coefficients(self) -> NodalMaterialCoefficients:
+        token = self._mesh_geometry_token()
+        cached = self._nodal_material_coefficients_cache
+        if cached is not None and cached[0] == token:
+            return cached[1]
+
+        point_count = len(self._mesh_points())
+        simplices = np.asarray(self._require_mesh().simplices, dtype=int)
+        regions = np.asarray(self._require_mesh().regions or [1] * len(simplices), dtype=int)
+        if len(regions) != len(simplices):
+            raise ValueError("Mesh regions must contain one entry per simplex.")
+
+        region_coefficients = {
+            int(region): self._material_dynamic_coefficients(self._simplex_material(int(region)))
+            for region in np.unique(regions)
+        }
+        if not region_coefficients:
+            defaults = (
+                self._material_dynamic_coefficients(self.materials[0])
+                if self.materials
+                else (
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                )
+            )
+            arrays = tuple(np.full(point_count, value, dtype=float) for value in defaults)
+        else:
+            nodal_regions = np.full(point_count, -1, dtype=int)
+            for region, coefficients in region_coefficients.items():
+                region_nodes = np.unique(simplices[regions == region].ravel())
+                assigned = nodal_regions[region_nodes]
+                for previous_region in np.unique(assigned[assigned >= 0]):
+                    previous = region_coefficients[int(previous_region)]
+                    if not np.allclose(previous, coefficients, rtol=1.0e-12, atol=0.0):
+                        conflict_nodes = region_nodes[assigned == previous_region]
+                        raise NotImplementedError(
+                            "A mesh node shared by regions with different exchange or LLG "
+                            f"coefficients requires material-specific magnetisation DOFs; node "
+                            f"{int(conflict_nodes[0])} belongs to regions {int(previous_region)} "
+                            f"and {region}."
+                        )
+                nodal_regions[region_nodes[assigned < 0]] = region
+
+            first_region = next(iter(region_coefficients))
+            nodal_regions[nodal_regions < 0] = first_region
+            arrays = tuple(
+                np.asarray(
+                    [region_coefficients[int(region)][component] for region in nodal_regions],
+                    dtype=float,
+                )
+                for component in range(6)
+            )
+
+        coefficients = NodalMaterialCoefficients(*arrays)
+        self._nodal_material_coefficients_cache = (token, coefficients)
+        return coefficients
+
+    def _exchange_prefactor(self) -> float:
+        if not self.materials:
+            return 0.0
+
+        prefactors = [
+            self._material_dynamic_coefficients(material)[0] for material in self.materials
+        ]
+
+        first = prefactors[0]
+        if any(
+            abs(prefactor - first) > 1.0e-12 * max(1.0, abs(first)) for prefactor in prefactors[1:]
+        ):
+            raise NotImplementedError("The exchange MVP currently supports one exchange prefactor.")
+        return float(first)
+
+    def _llg_coefficients(self) -> tuple[float, float, float]:
+        if not self.materials:
+            return 0.0, 0.0, 0.0
+
+        coefficients = [
+            self._material_dynamic_coefficients(material)[1:] for material in self.materials
+        ]
+
+        first = coefficients[0]
+        for coeffs in coefficients[1:]:
+            if any(
+                abs(value - reference) > 1.0e-12 * max(1.0, abs(reference))
+                for value, reference in zip(coeffs, first, strict=True)
+            ):
+                raise NotImplementedError(
+                    "The dmdt MVP currently supports one LLG coefficient set."
+                )
+        return (float(first[0]), float(first[1]), float(first[2]))
+
+    def _exchange_laplacian_spectral_bound(self) -> float:
+        """Return a cached upper bound for the lumped FEM Laplacian spectrum."""
+        token = self._mesh_geometry_token()
+        if self._exchange_spectral_bound_cache is not None:
+            cached_token, cached_bound = self._exchange_spectral_bound_cache
+            if cached_token == token:
+                return cached_bound
+
+        points = self._mesh_points()
+        simplices = np.asarray(self._require_mesh().simplices, dtype=int)
+        if simplices.size == 0:
+            bound = 0.0
+        else:
+            stiffness, _gradients_by_cell, volumes = self._demag_fem_geometry_for_mesh(
+                points, simplices
+            )
+            lumped_volumes = self._incident_cell_volume_sums(points, simplices, volumes) / 4.0
+            present = lumped_volumes > 0.0
+            if not np.any(present):
+                bound = 0.0
+            else:
+                absolute_row_sums = np.asarray(
+                    np.abs(stiffness[present]).sum(axis=1),
+                ).ravel()
+                bound = float(np.max(absolute_row_sums / lumped_volumes[present]))
+
+        self._exchange_spectral_bound_cache = (token, bound)
+        return bound
+
+    def _exchange_explicit_step_limit_seconds(self) -> float:
+        """Resolve stiff exchange modes accurately with the explicit integrator."""
+        spectral_bound = self._exchange_laplacian_spectral_bound()
+        coefficients = self._nodal_material_coefficients()
+        dynamic_coefficients = np.hypot(coefficients.precession, coefficients.damping)
+        maximum_rate = (
+            float(np.max(np.abs(coefficients.exchange_prefactor) * dynamic_coefficients))
+            * spectral_bound
+        )
+        if maximum_rate <= 0.0:
+            return float("inf")
+        return DOP853_EXCHANGE_RESOLUTION_RADIUS / maximum_rate

nmag/simulation/exchange/fields.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+
+class SimulationExchangeFieldMixin:
+    if TYPE_CHECKING:
+
+        def __getattr__(self, name: str) -> Any: ...
+
+    def _get_dm_dcurrent(self) -> np.ndarray:
+        if "current_density" not in self._fields:
+            return np.zeros((len(self._mesh_points()), 3), dtype=float)
+        points = self._mesh_points()
+        simplices = np.asarray(self._require_mesh().simplices, dtype=int)
+        if simplices.ndim != 2 or simplices.shape[1] != 4:
+            raise NotImplementedError("Spin-transfer torque currently supports tetrahedral meshes.")
+        stiffness, gradients_by_cell, volumes = self._demag_fem_geometry_for_mesh(
+            points,
+            simplices,
+        )
+        del stiffness
+        m = np.asarray(self._fields["m"], dtype=float)
+        cell_gradients = np.einsum(
+            "cla,clj->caj",
+            m[simplices],
+            gradients_by_cell,
+            optimize=True,
+        )
+        nodal_gradients = np.zeros((len(points), 3, 3), dtype=float)
+        weighted_gradients = cell_gradients * volumes[:, np.newaxis, np.newaxis]
+        for local_index in range(4):
+            np.add.at(
+                nodal_gradients,
+                simplices[:, local_index],
+                weighted_gradients,
+            )
+        weights = self._incident_cell_volume_sums(points, simplices, volumes)
+        present = weights > 0.0
+        nodal_gradients[present] /= weights[present, np.newaxis, np.newaxis]
+        current_density = np.asarray(self._fields["current_density"], dtype=float)
+        if current_density.shape != (len(points), 3):
+            raise ValueError(
+                "current_density must contain one 3-vector per mesh point; "
+                f"got {current_density.shape}."
+            )
+        return np.einsum(
+            "nij,nj->ni",
+            nodal_gradients,
+            current_density,
+            optimize=True,
+        )
+
+    def _get_exchange_nodal_field(self) -> np.ndarray:
+        token = self._demag_token()
+        if self._exchange_cache_token == token:
+            return self._exchange_nodal_cache
+
+        points = self._mesh_points()
+        simplices = np.asarray(self._require_mesh().simplices, dtype=int)
+        exchange = np.zeros((len(points), 3), dtype=float)
+        if simplices.size == 0:
+            self._exchange_cache_token = token
+            self._exchange_nodal_cache = exchange
+            return exchange
+        if simplices.ndim != 2 or simplices.shape[1] != 4:
+            raise NotImplementedError("The exchange MVP currently supports tetrahedral 3D meshes.")
+        if self._exchange_is_zero_by_construction():
+            self._exchange_cache_token = token
+            self._exchange_nodal_cache = exchange
+            return exchange
+
+        stiffness, _gradients_by_cell, volumes = self._demag_fem_geometry_for_mesh(
+            points, simplices
+        )
+        lumped_volumes = (
+            self._incident_cell_volume_sums(
+                points,
+                simplices,
+                volumes,
+            )
+            / 4.0
+        )
+
+        present = lumped_volumes > 0.0
+        if np.any(present):
+            cofield = stiffness @ np.asarray(self._fields["m"], dtype=float)
+            prefactors = self._nodal_material_coefficients().exchange_prefactor
+            exchange[present] = -(
+                prefactors[present, np.newaxis]
+                * cofield[present]
+                / lumped_volumes[present, np.newaxis]
+            )
+
+        self._exchange_cache_token = token
+        self._exchange_nodal_cache = exchange
+        return exchange
+
+    def _demag_token(self) -> tuple[int, int]:
+        if self.mesh is None:
+            raise RuntimeError("A mesh must be loaded before using demag fields.")
+        if "m" not in self._fields:
+            raise RuntimeError("Magnetisation must be set before using demag fields.")
+        return (id(self._require_mesh().raw_mesh), id(self._fields["m"]))