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

nmag/simulation/exchange/llg_rhs.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from ...backends import _load_rust_accelerator
+from ..support import _simulation_compatibility_binding
+
+
+class SimulationLlgRhsMixin:
+    if TYPE_CHECKING:
+
+        def __getattr__(self, name: str) -> Any: ...
+
+    def _llg_rhs_python(
+        self,
+        m: np.ndarray,
+        h_total: np.ndarray,
+        pin: np.ndarray,
+        ms_values: np.ndarray,
+        precession_coeff: float | np.ndarray,
+        damping_coeff: float | np.ndarray,
+        normalisation_coeff: float | np.ndarray,
+        dm_dcurrent: np.ndarray | None = None,
+        stt_adiabatic_coeff: float | np.ndarray = 0.0,
+        stt_nonadiabatic_coeff: float | np.ndarray = 0.0,
+    ) -> np.ndarray:
+        m, h_total, pin, ms_values = self._validated_llg_inputs(m, h_total, pin, ms_values)
+        precession_values, damping_values, normalisation_values = self._nodal_coefficients(
+            len(m),
+            (
+                ("precession", precession_coeff),
+                ("damping", damping_coeff),
+                ("normalisation", normalisation_coeff),
+            ),
+        )
+        rhs, mdotm = self._llg_field_rhs(
+            m,
+            h_total,
+            precession_values,
+            damping_values,
+            normalisation_values,
+        )
+        if dm_dcurrent is not None:
+            rhs += self._llg_stt_rhs(
+                m,
+                mdotm,
+                dm_dcurrent,
+                stt_adiabatic_coeff,
+                stt_nonadiabatic_coeff,
+            )
+        return ms_values[:, np.newaxis] * rhs * pin[:, np.newaxis]
+
+    @staticmethod
+    def _validated_llg_inputs(
+        m: np.ndarray,
+        h_total: np.ndarray,
+        pin: np.ndarray,
+        ms_values: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        m = np.asarray(m, dtype=float)
+        h_total = np.asarray(h_total, dtype=float)
+        pin = np.asarray(pin, dtype=float)
+        ms_values = np.asarray(ms_values, dtype=float)
+        if m.ndim != 2 or m.shape[1] != 3:
+            raise ValueError(f"m must have shape (n, 3), got {m.shape}.")
+        if h_total.shape != m.shape:
+            raise ValueError(f"h_total must have shape {m.shape}, got {h_total.shape}.")
+        if pin.shape != (len(m),):
+            raise ValueError(f"pin must have shape ({len(m)},), got {pin.shape}.")
+        if ms_values.shape != (len(m),):
+            raise ValueError(f"ms_values must have shape ({len(m)},), got {ms_values.shape}.")
+        return m, h_total, pin, ms_values
+
+    @staticmethod
+    def _nodal_coefficients(
+        point_count: int,
+        coefficients: tuple[tuple[str, float | np.ndarray], ...],
+    ) -> tuple[np.ndarray, ...]:
+        arrays: list[np.ndarray] = []
+        for name, coefficient in coefficients:
+            array = np.asarray(coefficient, dtype=float)
+            if array.ndim == 0:
+                array = np.full(point_count, float(array), dtype=float)
+            if array.shape != (point_count,):
+                raise ValueError(
+                    f"{name} coefficient must have shape ({point_count},), got {array.shape}."
+                )
+            arrays.append(array)
+        return tuple(arrays)
+
+    @staticmethod
+    def _llg_field_rhs(
+        m: np.ndarray,
+        h_total: np.ndarray,
+        precession_values: np.ndarray,
+        damping_values: np.ndarray,
+        normalisation_values: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        mxh = np.cross(m, h_total)
+        mdoth = np.einsum("ij,ij->i", m, h_total)
+        mdotm = np.einsum("ij,ij->i", m, m)
+        damping = m * mdoth[:, np.newaxis] - h_total * mdotm[:, np.newaxis]
+        norm_error = (1.0 - mdotm)[:, np.newaxis] * m
+        rhs = (
+            precession_values[:, np.newaxis] * mxh
+            + damping_values[:, np.newaxis] * damping
+            + normalisation_values[:, np.newaxis] * norm_error
+        )
+        return rhs, mdotm
+
+    def _llg_stt_rhs(
+        self,
+        m: np.ndarray,
+        mdotm: np.ndarray,
+        dm_dcurrent: np.ndarray,
+        stt_adiabatic_coeff: float | np.ndarray,
+        stt_nonadiabatic_coeff: float | np.ndarray,
+    ) -> np.ndarray:
+        directional_derivative = np.asarray(dm_dcurrent, dtype=float)
+        if directional_derivative.shape != m.shape:
+            raise ValueError(
+                f"dm_dcurrent must have shape {m.shape}, got {directional_derivative.shape}."
+            )
+        adiabatic_values, nonadiabatic_values = self._nodal_coefficients(
+            len(m),
+            (
+                ("stt_adiabatic", stt_adiabatic_coeff),
+                ("stt_nonadiabatic", stt_nonadiabatic_coeff),
+            ),
+        )
+        mx_directional = np.cross(m, directional_derivative)
+        mdot_directional = np.einsum("ij,ij->i", m, directional_derivative)
+        mxmx_directional = (
+            m * mdot_directional[:, np.newaxis] - directional_derivative * mdotm[:, np.newaxis]
+        )
+        return (
+            adiabatic_values[:, np.newaxis] * mxmx_directional
+            + nonadiabatic_values[:, np.newaxis] * mx_directional
+        )
+
+    def _llg_rhs_rust(
+        self,
+        m: np.ndarray,
+        h_total: np.ndarray,
+        pin: np.ndarray,
+        ms_values: np.ndarray,
+        precession_coeff: float | np.ndarray,
+        damping_coeff: float | np.ndarray,
+        normalisation_coeff: float | np.ndarray,
+        dm_dcurrent: np.ndarray | None = None,
+        stt_adiabatic_coeff: float | np.ndarray = 0.0,
+        stt_nonadiabatic_coeff: float | np.ndarray = 0.0,
+    ) -> np.ndarray:
+        rust_accel = _simulation_compatibility_binding(
+            "_load_rust_accelerator",
+            _load_rust_accelerator,
+        )("NmagConfig.accelerator['llg']")
+        try:
+            llg_rhs = rust_accel.llg_rhs
+        except AttributeError as exc:
+            raise RuntimeError(
+                "NmagConfig.accelerator['llg']='rust' requires an nmag_accel build with llg_rhs. "
+                "Rebuild it with "
+                "`maturin develop --release --manifest-path rust/nmag_accel/Cargo.toml`."
+            ) from exc
+        coefficient_arrays = [
+            np.asarray(value, dtype=float)
+            for value in (
+                precession_coeff,
+                damping_coeff,
+                normalisation_coeff,
+            )
+        ]
+        if dm_dcurrent is None and all(array.ndim == 0 for array in coefficient_arrays):
+            result = llg_rhs(
+                np.asarray(m, dtype=float),
+                np.asarray(h_total, dtype=float),
+                np.asarray(pin, dtype=float),
+                np.asarray(ms_values, dtype=float),
+                *(float(array) for array in coefficient_arrays),
+            )
+        else:
+            try:
+                heterogeneous_rhs = (
+                    rust_accel.llg_rhs_heterogeneous
+                    if dm_dcurrent is None
+                    else rust_accel.llg_rhs_stt_heterogeneous
+                )
+            except AttributeError as exc:
+                raise RuntimeError(
+                    "NmagConfig.accelerator['llg']='rust' requires an nmag_accel build with "
+                    "the heterogeneous LLG kernels. Rebuild the accelerator."
+                ) from exc
+            expanded = [
+                np.full(len(m), float(array), dtype=float) if array.ndim == 0 else array
+                for array in coefficient_arrays
+            ]
+            arguments: list[np.ndarray] = [
+                np.asarray(m, dtype=float),
+                np.asarray(h_total, dtype=float),
+                np.asarray(pin, dtype=float),
+                np.asarray(ms_values, dtype=float),
+                *expanded,
+            ]
+            if dm_dcurrent is not None:
+                stt_arrays = [
+                    np.asarray(value, dtype=float)
+                    for value in (stt_adiabatic_coeff, stt_nonadiabatic_coeff)
+                ]
+                expanded_stt = [
+                    np.full(len(m), float(array), dtype=float) if array.ndim == 0 else array
+                    for array in stt_arrays
+                ]
+                arguments.extend([np.asarray(dm_dcurrent, dtype=float), *expanded_stt])
+            result = heterogeneous_rhs(*arguments)
+        return np.asarray(result, dtype=float)

nmag/simulation/fields/__init__.py

@@ -0,0 +1,19 @@
+"""Focused field-operation mixins exposed through one compatibility façade."""
+
+from .arrays import SimulationFieldArrayMixin
+from .availability import SimulationFieldAvailabilityMixin
+from .averages import SimulationFieldAverageMixin
+from .derived import SimulationFieldDerivedMixin
+from .maxangle import SimulationFieldMaxangleMixin
+from .probes import SimulationFieldProbeMixin
+
+
+class SimulationFieldMixin(
+    SimulationFieldAvailabilityMixin,
+    SimulationFieldAverageMixin,
+    SimulationFieldMaxangleMixin,
+    SimulationFieldProbeMixin,
+    SimulationFieldArrayMixin,
+    SimulationFieldDerivedMixin,
+):
+    """Compose the public field API from focused implementation mixins."""

nmag/simulation/fields/arrays.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import time
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from ..support import (
+    _NO_CONSTANT_AVERAGE,
+    DERIVED_FIELD_NAMES,
+    _copy_average_value,
+)
+
+
+class SimulationFieldArrayMixin:
+    if TYPE_CHECKING:
+        _subfield_array_cache: dict[str, np.ndarray] | None
+        _subfield_average_cache: dict[tuple[str, str | None], object] | None
+        _active_subfield_array_timings: dict[str, float] | None
+
+        def __getattr__(self, name: str) -> Any: ...
+
+    def _subfield_array(self, subfieldname: str) -> np.ndarray:
+        started = time.perf_counter()
+        if self._subfield_array_cache is not None and subfieldname in self._subfield_array_cache:
+            self._record_active_subfield_array_timing(
+                f"{subfieldname}:cache_hit",
+                time.perf_counter() - started,
+            )
+            return self._subfield_array_cache[subfieldname]
+
+        data = self._compute_subfield_array(subfieldname)
+        self._record_active_subfield_array_timing(
+            f"{subfieldname}:compute",
+            time.perf_counter() - started,
+        )
+        if self._subfield_array_cache is not None:
+            self._subfield_array_cache[subfieldname] = data
+        return data
+
+    def _record_active_subfield_array_timing(self, name: str, seconds: float) -> None:
+        timings = self._active_subfield_array_timings
+        if timings is not None:
+            timings[name] = timings.get(name, 0.0) + seconds
+
+    @contextmanager
+    def _record_active_subfield_array_timing_block(self, name: str) -> Generator[None]:
+        started = time.perf_counter()
+        try:
+            yield
+        finally:
+            self._record_active_subfield_array_timing(
+                name,
+                time.perf_counter() - started,
+            )
+
+    @contextmanager
+    def _subfield_array_cache_scope(self) -> Generator[None]:
+        previous_cache = self._subfield_array_cache
+        self._subfield_array_cache = previous_cache if previous_cache is not None else {}
+        try:
+            yield
+        finally:
+            self._subfield_array_cache = previous_cache
+
+    @contextmanager
+    def _subfield_average_cache_scope(self) -> Generator[None]:
+        previous_cache = self._subfield_average_cache
+        self._subfield_average_cache = previous_cache if previous_cache is not None else {}
+        try:
+            yield
+        finally:
+            self._subfield_average_cache = previous_cache
+
+    def _cached_subfield_average(
+        self,
+        key: tuple[str, str | None],
+    ) -> object:
+        cache = self._subfield_average_cache
+        if cache is None or key not in cache:
+            return _NO_CONSTANT_AVERAGE
+        return cache[key]
+
+    def _cache_subfield_average(
+        self,
+        key: tuple[str, str | None],
+        value: object,
+    ) -> None:
+        cache = self._subfield_average_cache
+        if cache is not None:
+            cache[key] = _copy_average_value(value)
+
+    def _compute_subfield_array(self, subfieldname: str) -> np.ndarray:
+        if subfieldname in DERIVED_FIELD_NAMES:
+            return self._derived_subfield_array(subfieldname)
+        if subfieldname == "H_demag":
+            return np.asarray(self._get_demag_nodal_field(), dtype=float)
+        if subfieldname == "H_ext":
+            h_ext = np.asarray(self._fields.get("H_ext", np.zeros(3)), dtype=float)
+            if h_ext.shape != (3,):
+                raise ValueError(f"H_ext must be a homogeneous 3-vector, got shape {h_ext.shape}.")
+            point_count = len(self.mesh.points) if self.mesh is not None else 1
+            return np.tile(h_ext, (point_count, 1))
+        if subfieldname not in self._fields:
+            raise KeyError(f"Unknown or unset subfield '{subfieldname}'.")
+        return np.asarray(self._fields[subfieldname], dtype=float)
+
+    def _is_subfield_available(self, subfieldname: str) -> bool:
+        return self.is_subfield_available(subfieldname)

nmag/simulation/fields/availability.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from si.physical import SI
+
+if TYPE_CHECKING:
+    from simulation.quantity import Quantity
+
+from ..support import (
+    DERIVED_FIELD_NAMES,
+    _si_unit,
+    _simulation_compatibility_binding,
+)
+
+_known_field_quantities_cache: list[Quantity] | None = None
+_known_quantities_by_name_cache: dict[str, Quantity] | None = None
+
+
+def _known_field_quantities() -> list[Quantity]:
+    global _known_field_quantities_cache
+    if _known_field_quantities_cache is None:
+        from simulation.quantity import known_field_quantities
+
+        _known_field_quantities_cache = known_field_quantities
+    return _known_field_quantities_cache
+
+
+def _known_quantities_by_name() -> dict[str, Quantity]:
+    global _known_quantities_by_name_cache
+    if _known_quantities_by_name_cache is None:
+        from simulation.quantity import known_quantities_by_name
+
+        _known_quantities_by_name_cache = known_quantities_by_name
+    return _known_quantities_by_name_cache
+
+
+class SimulationFieldAvailabilityMixin:
+    if TYPE_CHECKING:
+
+        def __getattr__(self, name: str) -> Any: ...
+
+    def get_all_field_names(self) -> list[str]:
+        """Return the field names currently available for saving or access."""
+        has_mesh = self.mesh is not None
+        has_m = "m" in self._fields
+        has_mesh_and_m = has_mesh and has_m
+        has_tetrahedral_mesh = self._has_tetrahedral_mesh_or_empty()
+        return [
+            quantity.name
+            for quantity in _simulation_compatibility_binding(
+                "_known_field_quantities", _known_field_quantities
+            )()
+            if self._is_subfield_available_from_state(
+                quantity.name,
+                has_mesh=has_mesh,
+                has_m=has_m,
+                has_mesh_and_m=has_mesh_and_m,
+                has_tetrahedral_mesh=has_tetrahedral_mesh,
+            )
+        ]
+
+    def is_subfield_available(self, subfieldname: str) -> bool:
+        """Return whether the current model can provide a named field.
+
+        The check does not trigger expensive FEM/BEM work. Actual field access
+        can still raise a numerical error during calculation.
+
+        Args:
+            subfieldname: Candidate field name.
+
+        Returns:
+            True when the field is set or derivable from current state.
+        """
+        if subfieldname == "H_ext":
+            return True
+        if subfieldname not in DERIVED_FIELD_NAMES and subfieldname != "H_demag":
+            return subfieldname in self._fields
+
+        has_mesh = self.mesh is not None
+        has_m = "m" in self._fields
+        has_tetrahedral_mesh = self._has_tetrahedral_mesh_or_empty()
+        return self._is_subfield_available_from_state(
+            subfieldname,
+            has_mesh=has_mesh,
+            has_m=has_m,
+            has_mesh_and_m=has_mesh and has_m,
+            has_tetrahedral_mesh=has_tetrahedral_mesh,
+        )
+
+    def _is_subfield_available_from_state(
+        self,
+        subfieldname: str,
+        *,
+        has_mesh: bool,
+        has_m: bool,
+        has_mesh_and_m: bool,
+        has_tetrahedral_mesh: bool,
+    ) -> bool:
+        has_m_on_tetrahedral_mesh = self._has_m_on_tetrahedral_mesh(
+            has_mesh_and_m,
+            has_tetrahedral_mesh,
+        )
+        state_availability = {
+            "pin": has_mesh,
+            "H_anis": has_mesh_and_m,
+            "E_anis": has_mesh_and_m,
+            "H_exch": has_m_on_tetrahedral_mesh,
+            "E_exch": has_m_on_tetrahedral_mesh,
+            "dm_dcurrent": self._has_current_density(has_m_on_tetrahedral_mesh),
+            "dmdt": has_m_on_tetrahedral_mesh,
+            "H_total": has_m_on_tetrahedral_mesh,
+            "E_total": has_m_on_tetrahedral_mesh,
+            "E_demag": self._has_demag_magnetisation(has_mesh_and_m),
+            "phi": self._has_demag_magnetisation(has_mesh_and_m),
+            "rho": self._has_demag_magnetisation(has_mesh_and_m),
+            "M": has_mesh_and_m,
+            "E_ext": has_mesh_and_m,
+            "H_demag": self._has_demag_nodal_field(has_m, has_tetrahedral_mesh),
+            "H_ext": True,
+        }
+        if subfieldname in state_availability:
+            return state_availability[subfieldname]
+        return subfieldname not in DERIVED_FIELD_NAMES and subfieldname in self._fields
+
+    @staticmethod
+    def _has_m_on_tetrahedral_mesh(has_mesh_and_m: bool, has_tetrahedral_mesh: bool) -> bool:
+        return has_mesh_and_m and has_tetrahedral_mesh
+
+    def _has_current_density(self, has_m_on_tetrahedral_mesh: bool) -> bool:
+        return has_m_on_tetrahedral_mesh and "current_density" in self._fields
+
+    def _has_demag_magnetisation(self, has_mesh_and_m: bool) -> bool:
+        return self.do_demag and has_mesh_and_m
+
+    def _has_demag_nodal_field(self, has_m: bool, has_tetrahedral_mesh: bool) -> bool:
+        return self.do_demag and has_m and has_tetrahedral_mesh
+
+    def get_materials_of_field(self, field_name: str) -> list[Any]:
+        quantity = _simulation_compatibility_binding(
+            "_known_quantities_by_name", _known_quantities_by_name
+        )()[field_name]
+        if "?" in (quantity.signature or ""):
+            return self.materials
+        return []
+
+    def get_subfield(self, subfieldname: str, units: SI | None = None) -> Any:
+        """Return all nodal or cell values for one available field.
+
+        Args:
+            subfieldname: Name returned by :meth:`get_all_field_names`.
+            units: Optional compatible SI unit for returned numeric values.
+
+        Returns:
+            Field values as Python scalars or nested lists.
+
+        Raises:
+            KeyError: If the field is unknown, unset, or disabled.
+        """
+        data, field_units = self._subfield_data_and_units(subfieldname)
+        if units is None or field_units is None:
+            return data.tolist() if isinstance(data, np.ndarray) else data
+        return (np.asarray(data) * field_units.in_units_of(units)).tolist()
+
+    def _subfield_data_and_units(self, subfieldname: str) -> tuple[Any, SI | None]:
+        if subfieldname == "H_demag":
+            if not self.is_subfield_available(subfieldname):
+                raise KeyError("Subfield 'H_demag' is unavailable when demag is disabled.")
+            data = np.asarray(self._get_demag_nodal_field(), dtype=float)
+            return data, _si_unit("A/m")
+        if subfieldname in self._fields:
+            data = self._fields[subfieldname]
+        elif self.is_subfield_available(subfieldname):
+            data = self._subfield_array(subfieldname)
+        else:
+            raise KeyError(f"Unknown or unset subfield '{subfieldname}'.")
+        quantity = _simulation_compatibility_binding(
+            "_known_quantities_by_name", _known_quantities_by_name
+        )().get(subfieldname)
+        return data, None if quantity is None else quantity.units
+
+    def _has_tetrahedral_mesh_or_empty(self) -> bool:
+        if self.mesh is None:
+            return False
+        simplices = np.asarray(self.mesh.simplices)
+        return simplices.size == 0 or (simplices.ndim == 2 and simplices.shape[1] == 4)