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

llg3d/solvers/mpi.py

@@ -0,0 +1,370 @@
+"""
+LLG3D solver using MPI.
+
+The parallelization is done in the x direction.
+"""
+
+import argparse
+import sys
+import time
+
+import mpi4py
+import numpy as np
+
+from .. import solvers
+from .base import BaseSolver
+from .math_utils import cross_product
+from .profiling import timeit
+
+mpi4py.rc.initialize = False
+mpi4py.rc.finalize = True
+
+from mpi4py import MPI  # noqa: E402
+
+comm: MPI.Comm
+rank: int
+size: int
+status: MPI.Status
+
+Boundaries = tuple[np.ndarray, np.ndarray, MPI.Request | None, MPI.Request | None]
+Boundaries_3d = tuple[Boundaries, Boundaries, Boundaries]
+
+
+def initialize_mpi():
+    """Initialize MPI variables."""
+    if not MPI.Is_initialized():
+        MPI.Init()
+
+    global comm, rank, size, status
+    # Set MPI variables with the actual MPI values
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    size = comm.Get_size()
+    status = MPI.Status()
+    # Set these variables in the solver module as well
+    solvers.comm = comm
+    solvers.rank = rank
+    solvers.size = size
+    solvers.status = status
+    solvers.mpi_initialized = True
+
+
+class MPISolver(BaseSolver):
+    """MPI LLG3D solver."""
+
+    solver_type = "mpi"
+
+    def __post_init__(self):
+        """Initialize MPI and check parameters."""
+        initialize_mpi()
+        super().__post_init__()
+        if self.Jx % size != 0:
+            raise ValueError(
+                f"Jx ({self.Jx}) must be divisible by the number of processes ({size})"
+            )
+
+    def _get_boundaries_x(self, m: np.ndarray) -> Boundaries:
+        """
+        Returns the boundaries asynchronously.
+
+        Allows overlapping communication time of boundaries with calculations
+        if non-blocking communication is used.
+
+        Args:
+            m: Magnetization array (shape (3, nx, ny, nz))
+
+        Returns:
+            - m_i_x_start: Start boundary in x direction
+            - m_i_x_end: End boundary in x direction
+            - request_start: Request for start boundary or None
+            - request_end: Request for end boundary or None
+        """
+        # Extract slices for Neumann boundary conditions
+
+        # Create boundary buffers with the same dtype as `m` to avoid
+        # MPI packing/unpacking issues when precision is `single`.
+        m_i_x_start = np.empty((1, self.grid.Jy, self.grid.Jz), dtype=self.np_float)
+        m_i_x_end = np.empty_like(m_i_x_start)
+
+        # Prepare ring communication:
+        # Even if procs 0 and size - 1 shouldn't receive anything from left
+        # and right respectively, it's simpler to express it like this
+        right = (rank + 1) % size
+        left = (rank - 1 + size) % size
+
+        if self.blocking:
+            # Wait for boundaries to be available
+            comm.Sendrecv(
+                m[:1, :, :], dest=left, sendtag=0, recvbuf=m_i_x_end, source=right
+            )
+            comm.Sendrecv(
+                m[-1:, :, :], dest=right, sendtag=1, recvbuf=m_i_x_start, source=left
+            )
+            return m_i_x_start, m_i_x_end, None, None
+        else:
+            request_start = comm.Irecv(m_i_x_start, source=left, tag=201)
+            request_end = comm.Irecv(m_i_x_end, source=right, tag=202)
+            comm.Isend(m[-1:, :, :], dest=right, tag=201)
+            comm.Isend(m[:1, :, :], dest=left, tag=202)
+
+            return m_i_x_start, m_i_x_end, request_start, request_end
+
+    @timeit
+    def _get_boundaries_3d(self, m: np.ndarray) -> Boundaries_3d:
+        """
+        Returns the boundaries for all components.
+
+        Args:
+            m: Magnetization array (shape (3, nx, ny, nz))
+
+        Returns:
+            Boundaries for all three components
+        """
+        return (
+            self._get_boundaries_x(m[0]),
+            self._get_boundaries_x(m[1]),
+            self._get_boundaries_x(m[2]),
+        )
+
+    def laplacian_3d(
+        self,
+        m_i: np.ndarray,
+        m_i_x_start: np.ndarray,
+        m_i_x_end: np.ndarray,
+        request_start: MPI.Request | None,
+        request_end: MPI.Request | None,
+    ) -> np.ndarray:
+        """
+        Returns the Laplacian of m_i in 3D.
+
+        We start by calculating contributions in y and z, to wait
+        for the end of communications in x.
+
+        Args:
+            m_i: i-component of the magnetization array (shape (nx, ny, nz))
+            m_i_x_start: start boundary in x direction
+            m_i_x_end: end boundary in x direction
+            request_start: request for start boundary
+            request_end: request for end boundary
+
+        Returns:
+            Laplacian of m_i (shape (nx, ny, nz))
+        """
+        # Extract slices for Neumann boundary conditions
+        m_i_y_start = m_i[:, 1:2, :]
+        m_i_y_end = m_i[:, -2:-1, :]
+
+        m_i_z_start = m_i[:, :, 1:2]
+        m_i_z_end = m_i[:, :, -2:-1]
+
+        m_i_y_plus = np.concatenate((m_i[:, 1:, :], m_i_y_end), axis=1)
+        m_i_y_minus = np.concatenate((m_i_y_start, m_i[:, :-1, :]), axis=1)
+        m_i_z_plus = np.concatenate((m_i[:, :, 1:], m_i_z_end), axis=2)
+        m_i_z_minus = np.concatenate((m_i_z_start, m_i[:, :, :-1]), axis=2)
+
+        if self.grid.uniform:  # Uniform grid spacing: limit truncature errors
+            laplacian = self.grid.inv_dx2 * (
+                m_i_y_plus + m_i_y_minus + m_i_z_plus + m_i_z_minus - 6 * m_i
+            )
+        else:
+            laplacian = (
+                self.grid.inv_dy2 * (m_i_y_plus + m_i_y_minus)
+                + self.grid.inv_dz2 * (m_i_z_plus + m_i_z_minus)
+                + self.grid.center_coeff * m_i
+            )
+
+        # Wait for x-boundaries to be available (communications completed)
+        if request_start is not None:  # non blocking communication
+            request_start.Wait(status)
+        if request_end is not None:  # non blocking communication
+            request_end.Wait(status)
+
+        # For extreme procs, apply Neumann boundary conditions in x
+        if rank == size - 1:
+            m_i_x_end = m_i[-2:-1, :, :]
+        if rank == 0:
+            m_i_x_start = m_i[1:2, :, :]
+
+        m_i_x_plus = np.concatenate((m_i[1:, :, :], m_i_x_end), axis=0)
+        m_i_x_minus = np.concatenate((m_i_x_start, m_i[:-1, :, :]), axis=0)
+
+        laplacian += self.grid.inv_dx2 * (m_i_x_plus + m_i_x_minus)
+
+        return laplacian
+
+    @timeit
+    def compute_laplacian(self, m: np.ndarray, boundaries: Boundaries_3d) -> np.ndarray:
+        """
+        Compute the laplacian of m in 3D.
+
+        Args:
+            m: Magnetization array (shape (3, nx, ny, nz))
+            boundaries: Boundaries for x direction
+
+        Returns:
+            Laplacian of m (shape (3, nx, ny, nz))
+        """
+        return np.stack(
+            [
+                self.laplacian_3d(m[0], *boundaries[0]),
+                self.laplacian_3d(m[1], *boundaries[1]),
+                self.laplacian_3d(m[2], *boundaries[2]),
+            ],
+            axis=0,
+        )
+
+    @timeit
+    def compute_slope(
+        self,
+        m: np.ndarray,
+        R_random: np.ndarray,
+        H_aniso: np.ndarray,
+        boundaries: Boundaries_3d,
+    ) -> np.ndarray:
+        """
+        Compute the slope of the LLG equation.
+
+        Args:
+            m: Magnetization array (shape (3, nx, ny, nz))
+            R_random: Random field array (shape (3, nx, ny, nz))
+            H_aniso: Pre-allocated buffer for anisotropy field.
+            boundaries: Boundaries for x direction
+
+        Returns:
+            Slope array (shape (3, nx, ny, nz))
+        """
+        # Precalculate terms used multiple times
+
+        self.elem.compute_H_anisotropy(m, H_aniso)
+
+        laplacian_m = self.compute_laplacian(m, boundaries)
+        R_eff = self.elem.coeff_1 * laplacian_m + R_random + H_aniso
+        R_eff[0] += self.elem.coeff_3
+
+        m_cross_R_eff = cross_product(m, R_eff)
+        m_cross_m_cross_R_eff = cross_product(m, m_cross_R_eff)
+
+        s = -(m_cross_R_eff + self.elem.lambda_G * m_cross_m_cross_R_eff)
+
+        return s
+
+    @timeit
+    def _xyz_average(self, m: np.ndarray) -> float:
+        """
+        Returns the spatial average of m with shape (g.dims) using the midpoint method.
+
+        Performs the local sum on each process and then reduces it to process 0.
+
+        Args:
+            m: Array to be integrated (shape (x, y, z))
+
+        Returns:
+            Spatial average of m
+        """
+        # Make a copy of m to avoid modifying its value
+        mm = m.copy()
+
+        # On y and z edges, divide the contribution by 2
+        mm[:, 0, :] /= 2
+        mm[:, -1, :] /= 2
+        mm[:, :, 0] /= 2
+        mm[:, :, -1] /= 2
+
+        # On x edges (only on extreme procs), divide the contribution by 2
+        if rank == 0:
+            mm[0] /= 2
+        if rank == size - 1:
+            mm[-1] /= 2
+        local_sum = mm.sum()
+
+        # Sum across all processes gathered by process 0
+        global_sum = comm.reduce(local_sum)
+
+        # Spatial average is the global sum divided by the number of cells
+        return float(global_sum) / self.grid.ncell if global_sum is not None else 0.0
+
+    @timeit
+    def _yz_average(self, m_i: np.ndarray) -> np.ndarray:
+        """
+        Compute the integral of m_i over y and z directions.
+
+        Gather the local profiles from all processes to form the global profile.
+
+        Args:
+            m_i: Magnetization component array (shape (nx, ny, nz))
+
+        Returns:
+            1D array of the integral over y and z (shape (nx,))
+        """
+        local_yz_average = super()._yz_average(m_i)
+        # global coordinates
+        global_yz_average = np.empty(self.grid.Jx, dtype=self.np_float)
+        comm.Gather(local_yz_average, global_yz_average)
+        return global_yz_average
+
+    def _simulate(self) -> float:
+        """
+        Simulates the system for N iterations.
+
+        Returns:
+            The time taken for the simulation
+        """
+        m_n = self._init_m_n()
+        H_aniso = np.empty_like(m_n)
+        t = 0.0
+        self._record(m_n, t, 0)  # Record the initial solution
+
+        start_time = time.perf_counter()
+
+        for n in self._progress_bar():
+            t += self.dt
+
+            # Prediction phase
+            x_boundaries: Boundaries_3d = self._get_boundaries_3d(m_n)
+
+            # Adding randomness: effect of temperature
+            R_random = self._get_R_random()
+
+            s_pre = self.compute_slope(m_n, R_random, H_aniso, x_boundaries)
+            m_pre = m_n + self.dt * s_pre
+
+            # Correction phase
+            x_boundaries = self._get_boundaries_3d(m_pre)
+
+            s_cor = self.compute_slope(m_pre, R_random, H_aniso, x_boundaries)
+            m_n += self.dt * 0.5 * (s_pre + s_cor)
+
+            # renormalize to verify the constraint of being on the sphere
+            self._normalize(m_n)
+
+            self._record(m_n, t, n)
+
+        total_time = time.perf_counter() - start_time
+
+        self._finalize()
+
+        return total_time
+
+
+class ArgumentParser(argparse.ArgumentParser):
+    """An argument parser compatible with MPI."""
+
+    def _print_message(self, message, file=None):
+        if rank == 0 and message:
+            if file is None:
+                file = sys.stderr
+            file.write(message)
+
+    def exit(self, status: int = 0, message: str | None = None):
+        """
+        Exit the program using MPI finalize.
+
+        Args:
+            status: Exit status code
+            message: Optional exit message
+        """
+        if message:
+            self._print_message(message, sys.stderr)
+        comm.barrier()
+        MPI.Finalize()
+        exit(status)

llg3d/solvers/numpy.py

@@ -0,0 +1,126 @@
+"""LLG3D solver using NumPy."""
+
+from typing import ClassVar
+import time
+
+import numpy as np
+
+from .base import BaseSolver, timeit
+from .math_utils import cross_product
+
+
+class NumpySolver(BaseSolver):
+    """NumPy-based LLG3D solver."""
+
+    solver_type: ClassVar[str] = "numpy"  #: Solver type name
+
+    @timeit
+    def laplacian_3d(self, m_i: np.ndarray) -> np.ndarray:
+        """
+        Returns the laplacian of m_i in 3D with Neumann boundary conditions.
+
+        Args:
+            m_i: Magnetization direction (shape (nx, ny, nz))
+
+        Returns:
+            Laplacian of m_i (shape (nx, ny, nz))
+        """
+        m_i_padded = np.pad(m_i, ((1, 1), (1, 1), (1, 1)), mode="reflect")
+
+        laplacian = (
+            self.grid.inv_dx2
+            * (m_i_padded[2:, 1:-1, 1:-1] + m_i_padded[:-2, 1:-1, 1:-1])
+            + self.grid.inv_dy2
+            * (m_i_padded[1:-1, 2:, 1:-1] + m_i_padded[1:-1, :-2, 1:-1])
+            + self.grid.inv_dz2
+            * (m_i_padded[1:-1, 1:-1, 2:] + m_i_padded[1:-1, 1:-1, :-2])
+            + self.grid.center_coeff * m_i
+        )
+        return laplacian
+
+    @timeit
+    def compute_laplacian(self, m: np.ndarray) -> np.ndarray:
+        """
+        Compute the laplacian of m in 3D with Neumann boundary conditions.
+
+        Args:
+            m: Magnetization array (shape (3, nx, ny, nz))
+
+        Returns:
+            Laplacian of m (shape (3, nx, ny, nz))
+        """
+        return np.stack(
+            [
+                self.laplacian_3d(m[0]),
+                self.laplacian_3d(m[1]),
+                self.laplacian_3d(m[2]),
+            ],
+            axis=0,
+        )
+
+    @timeit
+    def _compute_slope(
+        self, m_n: np.ndarray, R_random: np.ndarray, H_aniso: np.ndarray
+    ) -> np.ndarray:
+        """
+        Compute the slope of the LLG equation.
+
+        Args:
+            m_n: Magnetization array at time step n (shape (3, nx, ny, nz)).
+            R_random: Random field array (shape (3, nx, ny, nz)).
+            H_aniso: Optional pre-allocated buffer for anisotropy field.
+
+        Returns:
+            Slope array (shape (3, nx, ny, nz))
+        """
+        self.elem.compute_H_anisotropy(m_n, H_aniso)
+
+        laplacian_m = self.compute_laplacian(m_n)
+        R_eff = self.elem.coeff_1 * laplacian_m + R_random + H_aniso
+        R_eff[0] += self.elem.coeff_3
+
+        m_cross_R_eff = cross_product(m_n, R_eff)
+        m_cross_m_cross_R_eff = cross_product(m_n, m_cross_R_eff)
+
+        s = -(m_cross_R_eff + self.elem.lambda_G * m_cross_m_cross_R_eff)
+
+        return s
+
+    def _simulate(self) -> float:
+        """
+        Simulates the system for N iterations.
+
+        Returns:
+            The time taken for the simulation
+        """
+        m_n = self._init_m_n()
+        H_aniso = np.empty_like(m_n)
+        t = 0.0
+        self._record(m_n, t, 0)  # Record the initial solution
+
+        start_time = time.perf_counter()
+
+        for n in self._progress_bar():
+            t += self.dt
+
+            # Adding randomness: temperature effect
+            R_random = self._get_R_random()
+
+            # Prediction phase
+            s_pre = self._compute_slope(m_n, R_random, H_aniso)
+            m_pre = m_n + self.dt * s_pre
+
+            # Correction phase
+            s_cor = self._compute_slope(m_pre, R_random, H_aniso)
+            m_n += self.dt * 0.5 * (s_pre + s_cor)
+
+            # We renormalize to verify the constraint of being on the sphere
+            self._normalize(m_n)
+
+            self._record(m_n, t, n)
+
+        total_time = time.perf_counter() - start_time
+
+        self._finalize()
+
+        return total_time