llg3d 2.0.1__py3-none-any.whl → 3.1.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"]

llg3d/{solver → solvers/experimental}/jax.py

@@ -1,4 +1,12 @@
-"""LLG3D solver using XLA compilation."""
+"""
+LLG3D solver using XLA compilation.
+
+.. warning::
+
+    It is experimental and not maintained.
+"""
+
+from typing import ClassVar
 
 import os
 import time
@@ -7,9 +15,7 @@ import jax
 import jax.numpy as jnp
 from jax import random
 
-from ..output import progress_bar, get_output_files, close_output_files
-from ..grid import Grid
-from ..element import Element
+from ..base import BaseSolver
 
 
 # JIT compile individual components for better performance and modularity
@@ -157,7 +163,7 @@ def compute_space_average_jax(m1: jnp.ndarray) -> float:
     # Compute ncell from the weights (this is the effective cell count)
     ncell = jnp.sum(weights)
 
-    return weighted_sum / ncell
+    return weighted_sum / ncell  # type: ignore
 
 
 @jax.jit
@@ -217,171 +223,139 @@ def compute_slope(
     return -(m_cross_R_eff + lambda_G * m_cross_m_cross_R_eff)
 
 
-def simulate(
-    N: int,
-    Jx: int,
-    Jy: int,
-    Jz: int,
-    dx: float,
-    T: float,
-    H_ext: float,
-    dt: float,
-    start_averaging: int,
-    n_mean: int,
-    n_profile: int,
-    element_class: Element,
-    precision: str,
-    seed: int,
-    device: str = "auto",
-    **_,
-) -> tuple[float, str, float]:
-    """
-    Simulates the system for N iterations using JAX.
+class JaxSolver(BaseSolver):
+    """JAX-based LLG3D solver."""
 
-    Args:
-        N: Number of iterations
-        Jx: Number of grid points in x direction
-        Jy: Number of grid points in y direction
-        Jz: Number of grid points in z direction
-        dx: Grid spacing
-        T: Temperature in Kelvin
-        H_ext: External magnetic field strength
-        dt: Time step for the simulation
-        start_averaging: Number of iterations for averaging
-        n_mean: Number of iterations for integral output
-        n_profile: Number of iterations for profile output
-        element_class: Element of the sample (default: Cobalt)
-        precision: Precision of the simulation (single or double)
-        seed: Random seed for temperature fluctuations
-        device: Device to use ('cpu', 'gpu', 'gpu:0', 'gpu:1', etc., or 'auto')
+    solver_type: ClassVar[str] = "jax"  #: Solver type name
 
-    Returns:
-        - The time taken for the simulation
-        - The output filenames
-        - The average magnetization
-    """
-    # Configure JAX
-    if device == "auto":
-        # Let JAX choose the best available device
-        pass
-    elif device == "cpu":
-        jax.config.update("jax_platform_name", "cpu")
-    elif device == "gpu":
-        jax.config.update("jax_platform_name", "gpu")
-    elif device.startswith("gpu:"):
-        # Select specific GPU using environment variable
-        jax.config.update("jax_platform_name", "gpu")
-        gpu_id = device.split(":")[1]
-        # Check if CUDA_VISIBLE_DEVICES is already set externally
-        if "CUDA_VISIBLE_DEVICES" not in os.environ:
-            os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
-            print(f"Set CUDA_VISIBLE_DEVICES={gpu_id}")
-        else:
-            cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
-            print(
-                f"Using external CUDA_VISIBLE_DEVICES={cuda_visible_devices}"
-            )
+    def _xyz_average(self, m1: jnp.ndarray) -> float:
+        """Compute the space average of m1 using JAX."""
+        return compute_space_average_jax(m1)
 
-    # Set precision
-    if precision == "double":
-        jax.config.update("jax_enable_x64", True)
-        jnp_float = jnp.float64
-    else:
-        jax.config.update("jax_enable_x64", False)
-        jnp_float = jnp.float32
+    def _simulate(self) -> float:
+        """
+        Simulates the system for N iterations using JAX.
 
-    print(f"Available JAX devices: {jax.devices()}")
-    print(f"Using JAX on device: {jax.devices()[0]}")
-    print(f"Precision: {precision} ({jnp_float})")
+        Attributes:
+            device: Device to use ('cpu', 'gpu', 'gpu:0', 'gpu:1', etc., or 'auto')
 
-    # Initialize random key for JAX
-    key = random.PRNGKey(seed)
+        Returns:
+            The time taken for the simulation
 
-    g = Grid(Jx, Jy, Jz, dx)
-    dims = g.dims
+        Raises:
+            NotImplementedError: If n_profile is not zero
+        """
+        if self.n_profile != 0:
+            raise NotImplementedError(
+                "Saving x-profiles is not implemented for the JAX solver."
+            )
 
-    e = element_class(T, H_ext, g, dt)
-    print(f"CFL = {e.get_CFL()}")
+        # Configure JAX
+        if self.device == "auto":
+            # Let JAX choose the best available device
+            pass
+        elif self.device == "cpu":
+            jax.config.update("jax_platform_name", "cpu")
+        elif self.device == "gpu":
+            jax.config.update("jax_platform_name", "gpu")
+        elif self.device.startswith("gpu:"):
+            # Select specific GPU using environment variable
+            jax.config.update("jax_platform_name", "gpu")
+            gpu_id = self.device.split(":")[1]
+            # Check if CUDA_VISIBLE_DEVICES is already set externally
+            if "CUDA_VISIBLE_DEVICES" not in os.environ:
+                os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+                print(f"Set CUDA_VISIBLE_DEVICES={gpu_id}")
+            else:
+                cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+                print(f"Using external CUDA_VISIBLE_DEVICES={cuda_visible_devices}")
+
+        # Set precision
+        if self.precision == "double":
+            jax.config.update("jax_enable_x64", True)
+            jnp_float = jnp.float64
+        else:
+            jax.config.update("jax_enable_x64", False)
+            jnp_float = jnp.float32
 
-    # Prepare parameters for JIT compilation using to_dict methods
-    g_params = g.to_dict()
-    e_params = e.to_dict()
+        print(f"Available JAX devices: {jax.devices()}")
+        print(f"Using JAX on device: {jax.devices()[0]}")
+        print(f"Precision: {self.precision} ({jnp_float})")
 
-    # --- Initialization ---
-    def theta_init(shape):
-        """Initialization of theta."""
-        return jnp.zeros(shape, dtype=jnp_float)
+        # Initialize random key for JAX
+        key = random.PRNGKey(self.seed)
 
-    def phi_init(t, shape):
-        """Initialization of phi."""
-        return jnp.zeros(shape, dtype=jnp_float) + e.gamma_0 * H_ext * t
+        # Prepare parameters for JIT compilation using to_dict methods
+        g_params = self.grid.as_dict()
+        e_params = self.elem.to_dict()
 
-    m_n = jnp.zeros((3,) + dims, dtype=jnp_float)
+        # --- Initialization ---
+        def theta_init(shape):
+            """Initialization of theta."""
+            return jnp.zeros(shape, dtype=jnp_float)
 
-    theta = theta_init(dims)
-    phi = phi_init(0, dims)
+        def phi_init(t, shape):
+            """Initialization of phi."""
+            return (
+                jnp.zeros(shape, dtype=jnp_float) + self.elem.gamma_0 * self.H_ext * t
+            )
 
-    m_n = m_n.at[0].set(jnp.cos(theta))
-    m_n = m_n.at[1].set(jnp.sin(theta) * jnp.cos(phi))
-    m_n = m_n.at[2].set(jnp.sin(theta) * jnp.sin(phi))
+        m_n = jnp.zeros((3, *self.grid.dims), dtype=jnp_float)
 
-    f_mean, f_profiles, output_filenames = get_output_files(g, T, n_mean, n_profile)
+        theta = theta_init(self.grid.dims)
+        phi = phi_init(0, self.grid.dims)
 
-    t = 0.0
-    m1_average = 0.0
+        m_n = m_n.at[0].set(jnp.cos(theta))
+        m_n = m_n.at[1].set(jnp.sin(theta) * jnp.cos(phi))
+        m_n = m_n.at[2].set(jnp.sin(theta) * jnp.sin(phi))
 
-    # === JIT WARMUP: Pre-compile all functions to exclude compilation time ===
-    print("Warming up JIT compilation...")
+        t = 0.0
 
-    # Generate dummy random field for warmup
-    warmup_key = random.PRNGKey(42)
-    R_warmup = e.coeff_4 * random.normal(warmup_key, (3,) + dims, dtype=jnp_float)
+        # === JIT WARMUP: Pre-compile all functions to exclude compilation time ===
+        print("Warming up JIT compilation...")
 
-    # Warmup all JIT functions with actual data shapes
-    _ = compute_slope(g_params, e_params, m_n, R_warmup)
-    if n_mean != 0:
-        _ = compute_space_average_jax(m_n[0])
+        # Generate dummy random field for warmup
+        warmup_key = random.PRNGKey(42)
+        R_warmup = self.elem.coeff_4 * random.normal(
+            warmup_key, (3, *self.grid.dims), dtype=jnp_float
+        )
 
-    # Force compilation and execution to complete
-    jax.block_until_ready(m_n)
-    print("JIT warmup completed.")
+        # Warmup all JIT functions with actual data shapes
+        _ = compute_slope(g_params, e_params, m_n, R_warmup)
+        if self.n_mean != 0:
+            _ = compute_space_average_jax(m_n[0])
 
-    start_time = time.perf_counter()
+        # Force compilation and execution to complete
+        jax.block_until_ready(m_n)
+        print("JIT warmup completed.")
 
-    for n in progress_bar(range(1, N + 1), "Iteration : ", 40):
-        t += dt
+        start_time = time.perf_counter()
 
-        # Generate random field for temperature effect
-        key, subkey = random.split(key)
-        R_random = e.coeff_4 * random.normal(subkey, (3,) + dims, dtype=jnp_float)
+        for n in self._progress_bar():
+            t += self.dt
 
-        # Use JIT-compiled version for better performance
-        s_pre = compute_slope(g_params, e_params, m_n, R_random)
-        m_pre = m_n + dt * s_pre
-        s_cor = compute_slope(g_params, e_params, m_pre, R_random)
+            # Generate random field for temperature effect
+            key, subkey = random.split(key)
+            R_random = self.elem.coeff_4 * random.normal(
+                subkey, (3, *self.grid.dims), dtype=jnp_float
+            )
 
-        # Update magnetization
-        m_n = m_n + dt * 0.5 * (s_pre + s_cor)
+            # Use JIT-compiled version for better performance
+            s_pre = compute_slope(g_params, e_params, m_n, R_random)
+            m_pre = m_n + self.dt * s_pre
+            s_cor = compute_slope(g_params, e_params, m_pre, R_random)
 
-        # Renormalize to unit sphere
-        norm = jnp.sqrt(m_n[0] ** 2 + m_n[1] ** 2 + m_n[2] ** 2)
-        m_n = m_n / norm
+            # Update magnetization
+            m_n = m_n + self.dt * 0.5 * (s_pre + s_cor)
 
-        # Export the average of m1 to a file (optimized for GPU)
-        if n_mean != 0 and n % n_mean == 0:
-            # Compute space average directly on GPU without CPU transfer
-            m1_mean = compute_space_average_jax(m_n[0])
-            # Convert to Python float for file writing
-            m1_mean = float(m1_mean)
-            if n >= start_averaging:
-                m1_average += m1_mean * n_mean
-            f_mean.write(f"{t:10.8e} {m1_mean:10.8e}\n")
+            # Renormalize to unit sphere
+            norm = jnp.sqrt(m_n[0] ** 2 + m_n[1] ** 2 + m_n[2] ** 2)
+            m_n = m_n / norm
 
-    total_time = time.perf_counter() - start_time
+            self._record(m_n, t, n)
 
-    close_output_files(f_mean, f_profiles)
+        total_time = time.perf_counter() - start_time
 
-    if n > start_averaging:
-        m1_average /= N - start_averaging
+        self._finalize()
 
-    return total_time, output_filenames, m1_average
+        return total_time