llg3d 2.0.0__py3-none-any.whl → 3.0.0__py3-none-any.whl

llg3d/solvers/base.py

@@ -0,0 +1,345 @@
+"""Define the base solver class."""
+
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, ClassVar
+
+import numpy as np
+from tqdm import tqdm
+
+from ..element import Element, get_element_class
+from ..grid import Grid
+from ..io import (
+    Metrics,
+    Observables,
+    RecordsBuffer,
+    format_profiling_table,
+    get_tqdm_file,
+    save_results,
+)
+from ..parameters import InitType, RunParameters
+from . import rank, size
+from .math_utils import normalize
+from .profiling import ProfilingStats, timeit
+
+
+@dataclass
+class BaseSolver(ABC, RunParameters):
+    """Abstract data base class for LLG3D solvers."""
+
+    solver_type: ClassVar[str] = "base"  #: Solver type name
+
+    def __post_init__(self) -> None:
+        """Initialize the solver after dataclass creation."""
+        # Ensure solver name matches class solver_type
+        if self.solver != self.solver_type:
+            raise ValueError(
+                f"Trying to initialize a {self.__class__.__name__}, but solver name  "
+                f'mismatch: expected "{self.solver_type}", got "{self.solver}"'
+            )
+
+        # Initialize results structure aligned with RunResults
+        self.metrics: Metrics = {
+            "total_time": 0.0,
+            "time_per_ite": 0.0,
+            "efficiency": 0.0,
+            "CFL": 0.0,
+        }
+        # Physical observables
+        self.observables: Observables = {}
+        # Records with optional x_profiles and xyz_average
+        self.records: RecordsBuffer = {}
+
+        self.np = size  # Add a parameter for the number of processes
+
+        self.np_float: np.dtype = np.dtype(
+            np.float64 if self.precision == "double" else np.float32
+        )
+        self.grid: Grid = Grid(self.Jx, self.Jy, self.Jz, self.dx)
+        if rank == 0:
+            print(self.grid)
+        # Reference the element class from the element string
+        ElementClass: type[Element] = get_element_class(self.element)
+        # Pass dtype to Element so its scalar coefficients have the correct
+        # numpy dtype (`self.np_float`) and avoid implicit promotion.
+        self.elem: Element = ElementClass(
+            self.T, self.H_ext, self.grid, self.dt, dtype=self.np_float
+        )
+
+        self.rng: np.random.Generator = self._init_rng()
+        self.profiling_stats: ProfilingStats = defaultdict(
+            lambda: {"time": 0.0, "calls": 0}
+        )
+        self._tqdm_file = get_tqdm_file()  #: File object for tqdm output
+
+    def theta_init_0(self, t: float) -> np.ndarray:
+        """Initialization of theta with 0."""
+        return np.zeros(self.grid.dims, dtype=self.np_float)
+
+    def theta_init_dw(self, t: float) -> np.ndarray:
+        """Initialization of theta with a domain wall profile."""
+        x, _, _ = self.grid.get_mesh(local=size > 1, dtype=self.np_float)
+        return 2.0 * np.arctan(
+            np.exp(
+                -(
+                    x
+                    - self.grid.Lx / 2
+                    + self.elem.d_0 * self.elem.coeff_3 * self.elem.lambda_G * t
+                )
+                / self.elem.d_0
+            )
+        )
+
+    def phi_init(self, t: float) -> np.ndarray:
+        """Initialization of phi."""
+        return (
+            np.zeros(self.grid.dims, dtype=self.np_float)
+            + self.elem.gamma_0 * self.elem.H_ext * t
+        )
+
+    def _init_m_n(self) -> np.ndarray:
+        """Initialize the magnetization array at time step n."""
+        m_n = np.zeros((3,) + self.grid.dims, dtype=self.np_float)
+
+        if self.init_type == "0":
+            theta = self.theta_init_0(0)
+        elif self.init_type == "dw":
+            theta = self.theta_init_dw(0)
+        else:
+            raise ValueError(
+                f"Unknown initialization type: {self.init_type}, "
+                f"should be in {InitType}"
+            )
+
+        phi = self.phi_init(0)
+
+        m_n[0] = np.cos(theta)
+        m_n[1] = np.sin(theta) * np.cos(phi)
+        m_n[2] = np.sin(theta) * np.sin(phi)
+        # renormalize to verify the constraint of being on the sphere
+        normalize(m_n)
+        return m_n
+
+    def _init_rng(self) -> np.random.Generator:
+        """
+        Initialize a random number generator for temperature fluctuations.
+
+        Returns:
+            A numpy random number generator
+        """
+        # Initialize a sequence of random seeds
+        # See: https://numpy.org/doc/stable/reference/random/parallel.html
+        ss = np.random.SeedSequence(self.seed)
+
+        # Deploy size x SeedSequence to pass to child processes
+        child_seeds = ss.spawn(size)
+        streams = [np.random.default_rng(s) for s in child_seeds]
+        rng = streams[rank]
+        return rng
+
+    @timeit
+    def _get_R_random(self) -> np.ndarray:
+        """
+        Generate the random field for temperature fluctuations.
+
+        Returns:
+            Random field array (shape (3, nx, ny, nz))
+        """
+        R_random = self.elem.coeff_4 * self.rng.standard_normal(
+            (3, *self.grid.dims), dtype=self.np_float
+        )
+        return R_random
+
+    @timeit
+    def _normalize(self, m_n: np.ndarray) -> None:
+        r"""
+        Normalize the magnetization array (in place).
+
+        .. math::
+
+            \mathbf{m}_n = \frac{\mathbf{m}_n}{|\mathbf{m}_n|}
+
+        Args:
+            m_n: Magnetization array at time step n (shape (3, nx, ny, nz)).
+        """
+        normalize(m_n)
+
+    @timeit
+    def _xyz_average(self, m: np.ndarray) -> float:
+        """
+        Returns the spatial average of m with shape (g.dims) using the midpoint method.
+
+        Args:
+            m: Array to be integrated
+
+        Returns:
+            Spatial average of m
+        """
+        mm = m.copy()  # copy m to avoid modifying its value
+
+        # on the edges, we divide the contribution by 2
+        # x
+        mm[0, :, :] /= 2
+        mm[-1, :, :] /= 2
+        # y
+        mm[:, 0, :] /= 2
+        mm[:, -1, :] /= 2
+        # z
+        mm[:, :, 0] /= 2
+        mm[:, :, -1] /= 2
+
+        average = mm.sum() / self.grid.ncell
+        return float(average)
+
+    def _record_xyz_average(self, m_n: np.ndarray, t: float, n: int) -> None:
+        """Update the time average of m1."""
+        xyz_average = self._xyz_average(m_n[0])
+        if rank == 0:
+            # Ensure xyz_average list exists
+            if "xyz_average" not in self.records:
+                self.records["xyz_average"] = []
+            # Record the mean value at time t
+            self.records["xyz_average"].append((t, xyz_average))
+            # Update time average of m1
+            if n >= self.start_averaging:
+                # Initialize m1_mean on first use
+                if "m1_mean" not in self.observables:
+                    self.observables["m1_mean"] = 0.0
+                # Accumulate time average (each sample contributes equally)
+                self.observables["m1_mean"] += xyz_average
+
+    def _finalize(self) -> None:
+        """Normalize m1_mean by the actual number of samples accumulated."""
+        if rank == 0:
+            if "m1_mean" in self.observables:
+                # Divide by actual number of samples
+                # (accounting for n_mean sampling interval)
+                num_samples = (self.N - self.start_averaging) // self.n_mean
+                self.observables["m1_mean"] /= num_samples
+
+    @timeit
+    def _yz_average(self, m_i: np.ndarray) -> np.ndarray:
+        """
+        Returns the spatial average of m using the midpoint method along y and z.
+
+        Args:
+            m_i: Array to be integrated
+
+        Returns:
+            Spatial average of m in y and z of shape (g.dims[0],)
+        """
+        # Make a copy of m to avoid modifying its value
+        mm = m_i.copy()
+
+        # On y and z edges, divide the contribution by 2
+        mm[:, 0, :] /= 2
+        mm[:, -1, :] /= 2
+        mm[:, :, 0] /= 2
+        mm[:, :, -1] /= 2
+
+        n_cell_yz = (mm.shape[1] - 1) * (mm.shape[2] - 1)
+        return mm.sum(axis=(1, 2)) / n_cell_yz
+
+    @timeit
+    def _update_x_profiles(self, m_n: np.ndarray, t: float) -> None:
+        """Update x profiles of the averaged m_i in y and z."""
+        # Initialize x_profiles on first use
+        if "x_profiles" not in self.records:
+            self.records["x_profiles"] = {
+                "t": [],
+                "m1": [],
+                "m2": [],
+                "m3": [],
+            }
+        x_prof = self.records["x_profiles"]
+        x_prof["t"].append(t)
+        x_prof["m1"].append(self._yz_average(m_n[0]))
+        x_prof["m2"].append(self._yz_average(m_n[1]))
+        x_prof["m3"].append(self._yz_average(m_n[2]))
+
+    def _record(self, m_n: Any, t: float, n: int) -> None:
+        """Record simulation data."""
+        # Record the average of m1
+        if self.n_mean != 0 and n % self.n_mean == 0:
+            self._record_xyz_average(m_n, t, n)
+        # Record the x profiles of the averaged m_i in y and z
+        if self.n_profile != 0 and n % self.n_profile == 0:
+            self._update_x_profiles(m_n, t)
+
+    def _progress_bar(self):
+        """Return a progress bar for the given range using tqdm."""
+        if self._tqdm_file is None:
+            # corresponds to rank != 0 in MPI
+            return range(self.N)
+
+        return tqdm(
+            range(self.N),
+            file=self._tqdm_file,
+            dynamic_ncols=True,
+            leave=True,
+        )
+
+    @abstractmethod
+    def _simulate(self) -> float:
+        """
+        Simulates the system for N iterations.
+
+        Returns:
+            total_time: Total simulation time
+        """
+
+    def run(self) -> None:
+        """Runs the simulation and store the results."""
+        total_time = self._simulate()
+        self.metrics["total_time"] = total_time
+        time_per_ite = total_time / self.N if self.N > 0 else 0.0
+        self.metrics["time_per_ite"] = time_per_ite
+        self.metrics["efficiency"] = time_per_ite / self.grid.ntot
+        self.metrics["CFL"] = float(self.elem.get_CFL())
+        if rank == 0:
+            # Store only profiling stats for functions/kernels that were actually called
+            profiling_filtered = {
+                k: v for k, v in self.profiling_stats.items() if v.get("calls", 0) > 0
+            }
+            if profiling_filtered:
+                self.metrics["profiling_stats"] = profiling_filtered
+
+    def _format_profiling(self) -> str:
+        """Format the profiling information for display."""
+        return format_profiling_table(self.profiling_stats, self.metrics["total_time"])
+
+    def save(self, dir_path: str | Path = ".") -> None:
+        """
+        Saves the results of the simulation to a .npz file.
+
+        Args:
+            dir_path: Directory path to save the results
+        """
+        if rank == 0:
+            if self.metrics["total_time"] > 0:
+                s = f"""\
+N iterations          = {self.N}
+total_time [s]        = {self.metrics["total_time"]:.03f}
+time/ite [s/ite]      = {self.metrics["time_per_ite"]:.03e}
+efficiency [s/ite/pt] = {self.metrics["efficiency"]:.03e}
+CFL                   = {self.metrics["CFL"]:.03e}"""
+                print(s)
+
+                # Print profiling info only if enabled
+                if self.profiling:
+                    print(f"Profiling info:\n{self._format_profiling()}")
+
+            # Export the mean of m1 over space and time
+            if "m1_mean" in self.observables:
+                print(f"m1_mean           = {self.observables['m1_mean']:e}")
+
+            print(f"Saving {self.result_file}")
+            save_results(
+                Path(dir_path) / self.result_file,
+                self,
+                self.metrics,
+                observables=self.observables if self.observables else None,
+                records_buffer=self.records if self.records else None,
+            )

llg3d/solvers/experimental/__init__.py

@@ -0,0 +1,9 @@
+"""
+Experimental solvers live here.
+
+Currently contains JAX-based solver implementations.
+"""
+
+from .jax import JaxSolver
+
+__all__ = ["JaxSolver"]