llg3d 1.4.1__py3-none-any.whl → 2.0.1__py3-none-any.whl

llg3d/post/temperature.py

@@ -1,30 +1,24 @@
 #!/usr/bin/env python3
-"""
-Plot the magnetization vs temperature and determine the Curie temperature.
-"""
+"""Plot the magnetization vs temperature and determine the Curie temperature."""
 
 import argparse
 from pathlib import Path
 
 import matplotlib.pyplot as plt
-import numpy as np
 
 from .process import MagData
 
 
 def plot_m_vs_T(m: MagData, show: bool):
     """
-    Plots the data (T, <m>), interpolates the values,
-    calculates the Curie temperature.
-    Exports to PNG.
+    Plots the data (T, <m>).
+
+    Interpolates the values, calculates the Curie temperature, exports to PNG.
 
     Args:
-        data: numpy array (T, <m>)
-        parentdir: path to the directory containing the runs
-        run: descriptive dictionary of the run
+        m: Magnetization data object
         show: display the graph in a graphical window
     """
-
     print(f"T_Curie = {m.T_Curie:.0f} K")
 
     fig, ax = plt.subplots()
@@ -57,9 +51,7 @@ def plot_m_vs_T(m: MagData, show: bool):
 
 
 def main():
-    """
-    Parses the command line to execute processing functions
-    """
+    """Parses the command line to execute processing functions."""
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--job_dir", type=Path, help="Slurm main job directory")
     parser.add_argument(

llg3d/simulation.py

@@ -0,0 +1,95 @@
+"""
+Define the Simulation class.
+
+Example usage:
+
+    >>> from llg3d.main import Simulation
+    >>> from llg3d.parameters import parameters
+    >>> run_parameters = {name: value["default"] for name, value in parameters.items()}
+    >>> sim = Simulation(run_parameters)
+    >>> sim.run()
+    >>> sim.save()
+
+"""
+
+import inspect
+
+from . import rank, size
+from .element import get_element_class
+from .parameters import Parameter
+from .output import write_json
+
+
+class Simulation:
+    """
+    Class to encapsulate the simulation logic.
+
+    Args:
+        params: Dictionary of simulation parameters.
+    """
+
+    json_file = "run.json"  #: JSON file to store the results
+
+    def __init__(self, params: dict[str, Parameter]):
+        self.params: dict[str, Parameter] = params.copy()  #: simulation parameters
+        self.simulate: callable = self._get_simulate_function_from_name(
+            self.params["solver"]
+        )  #: simulation function imported from the solver module
+        self.total_time: None | float = None  #: total simulation time
+        self.filenames: list[str] = []  #: list of output filenames
+        self.m1_mean: None | float = None  #: space and time average of m1
+        self.params["np"] = size  # Add a parameter for the number of processes
+        # Reference the element class from the element string
+        self.params["element_class"] = get_element_class(params["element"])
+
+    def run(self):
+        """Runs the simulation and store the results."""
+        self.total_time, self.filenames, self.m1_mean = self.simulate(**self.params)
+
+    def _get_simulate_function_from_name(self, name: str) -> callable:
+        """
+        Retrieves the simulation function for a given solver name.
+
+        Args:
+            name: Name of the solver
+
+        Returns:
+            callable: The simulation function
+
+        Example:
+            >>> simulate = self.get_simulate_function_from_name("mpi")
+
+            Will return the `simulate` function from the `llg3d.solver.mpi` module.
+        """
+        module = __import__(f"llg3d.solver.{name}", fromlist=["simulate"])
+        return inspect.getattr_static(module, "simulate")
+
+    def save(self):
+        """Saves the results of the simulation to a JSON file."""
+        params = self.params.copy()  # save the parameters
+        del params["element_class"]  # remove class object before serialization
+        if rank == 0:
+            results = {"total_time": self.total_time}
+            # Export the integral of m1
+            if len(self.filenames) > 0:
+                results["integral_file"] = self.filenames[0]
+                print(f"Integral of m1 in {self.filenames[0]}")
+            # Export the x-profiles of m1, m2 and m3
+            for i, filename in enumerate(self.filenames[1:]):
+                results[f"xprofile_m{i}"] = filename
+                print(f"x-profile of m{i} in {filename}")
+
+            print(
+                f"""\
+N iterations      = {params["N"]}
+total_time [s]    = {self.total_time:.03f}
+time/ite [s/iter] = {self.total_time / params["N"]:.03e}\
+"""
+            )
+            # Export the mean of m1
+            if params["N"] > params["start_averaging"]:
+                print(f"m1_mean           = {self.m1_mean:e}")
+                results["m1_mean"] = float(self.m1_mean)
+
+            write_json(self.json_file, {"params": params, "results": results})
+            print(f"Summary in {self.json_file}")

llg3d/solver/__init__.py

@@ -0,0 +1,45 @@
+"""
+Solver module for LLG3D.
+
+This module contains different solver implementations.
+"""
+
+import importlib.util
+
+
+__all__ = ["numpy", "solver", "rank", "size", "comm", "status"]
+
+LIB_AVAILABLE: dict[str, bool] = {}
+
+# Check for other solver availability
+for lib in "opencl", "jax", "mpi4py":
+    if importlib.util.find_spec(lib, package=__package__) is not None:
+        LIB_AVAILABLE[lib] = True
+        __all__.append(lib)
+    else:
+        LIB_AVAILABLE[lib] = False
+
+
+# MPI utilities
+if LIB_AVAILABLE["mpi4py"]:
+    from mpi4py import MPI
+
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    size = comm.Get_size()
+    status = MPI.Status()
+else:
+    # MPI library is not available: use dummy values
+    class DummyComm:
+        pass
+
+    comm = DummyComm()
+    rank = 0
+    size = 1
+
+    class DummyStatus:
+        pass
+
+    status = DummyStatus()
+
+from . import numpy, solver  # noqa: E402

llg3d/solver/jax.py

@@ -0,0 +1,387 @@
+"""LLG3D solver using XLA compilation."""
+
+import os
+import time
+
+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
+
+
+# JIT compile individual components for better performance and modularity
+@jax.jit
+def compute_H_anisotropy(
+    m: jnp.ndarray, coeff_2: float, anisotropy: int
+) -> jnp.ndarray:
+    """
+    Compute anisotropy field (JIT compiled).
+
+    Args:
+        m: Magnetization array (shape (3, nx, ny, nz))
+        coeff_2: Coefficient for anisotropy
+        anisotropy: Anisotropy type (0: uniaxial, 1: cubic)
+
+    Returns:
+        Anisotropy field array (shape (3, nx, ny, nz))
+    """
+    m1, m2, m3 = m
+
+    m1m1 = m1 * m1
+    m2m2 = m2 * m2
+    m3m3 = m3 * m3
+
+    # Uniaxial anisotropy
+    aniso_1_uniaxial = m1
+    aniso_2_uniaxial = jnp.zeros_like(m1)
+    aniso_3_uniaxial = jnp.zeros_like(m1)
+
+    # Cubic anisotropy
+    aniso_1_cubic = -(1 - m1m1 + m2m2 * m3m3) * m1
+    aniso_2_cubic = -(1 - m2m2 + m1m1 * m3m3) * m2
+    aniso_3_cubic = -(1 - m3m3 + m1m1 * m2m2) * m3
+
+    # Select based on anisotropy type
+    aniso_1 = jnp.where(
+        anisotropy == 0,
+        aniso_1_uniaxial,
+        jnp.where(anisotropy == 1, aniso_1_cubic, jnp.zeros_like(m1)),
+    )
+    aniso_2 = jnp.where(
+        anisotropy == 0,
+        aniso_2_uniaxial,
+        jnp.where(anisotropy == 1, aniso_2_cubic, jnp.zeros_like(m1)),
+    )
+    aniso_3 = jnp.where(
+        anisotropy == 0,
+        aniso_3_uniaxial,
+        jnp.where(anisotropy == 1, aniso_3_cubic, jnp.zeros_like(m1)),
+    )
+
+    return coeff_2 * jnp.stack([aniso_1, aniso_2, aniso_3], axis=0)
+
+
+@jax.jit
+def laplacian3D(
+    m_i: jnp.ndarray,
+    dx2_inv: float,
+    dy2_inv: float,
+    dz2_inv: float,
+    center_coeff: float,
+) -> jnp.ndarray:
+    """
+    Compute Laplacian for a single component with Neumann boundary conditions.
+
+    (JIT compiled)
+
+    Args:
+        m_i: Single component of magnetization (shape (nx, ny, nz))
+        dx2_inv: Inverse of squared grid spacing in x direction
+        dy2_inv: Inverse of squared grid spacing in y direction
+        dz2_inv: Inverse of squared grid spacing in z direction
+        center_coeff: Coefficient for the center point
+
+    Returns:
+        Laplacian of m_i (shape (nx, ny, nz))
+    """
+    m_i_padded = jnp.pad(m_i, ((1, 1), (1, 1), (1, 1)), mode="reflect")
+    return (
+        dx2_inv * (m_i_padded[2:, 1:-1, 1:-1] + m_i_padded[:-2, 1:-1, 1:-1])
+        + dy2_inv * (m_i_padded[1:-1, 2:, 1:-1] + m_i_padded[1:-1, :-2, 1:-1])
+        + dz2_inv * (m_i_padded[1:-1, 1:-1, 2:] + m_i_padded[1:-1, 1:-1, :-2])
+        + center_coeff * m_i
+    )
+
+
+@jax.jit
+def compute_laplacian(m: jnp.ndarray, dx: float, dy: float, dz: float) -> jnp.ndarray:
+    """
+    Compute 3D Laplacian with Neumann boundary conditions (JIT compiled).
+
+    Args:
+        m: Magnetization array (shape (3, nx, ny, nz))
+        dx: Grid spacing in x direction
+        dy: Grid spacing in y direction
+        dz: Grid spacing in z direction
+
+    Returns:
+        Laplacian of m (shape (3, nx, ny, nz))
+    """
+    dx2_inv, dy2_inv, dz2_inv = 1 / dx**2, 1 / dy**2, 1 / dz**2
+    center_coeff = -2 * (dx2_inv + dy2_inv + dz2_inv)
+
+    return jnp.stack(
+        [
+            laplacian3D(m[0], dx2_inv, dy2_inv, dz2_inv, center_coeff),
+            laplacian3D(m[1], dx2_inv, dy2_inv, dz2_inv, center_coeff),
+            laplacian3D(m[2], dx2_inv, dy2_inv, dz2_inv, center_coeff),
+        ],
+        axis=0,
+    )
+
+
+@jax.jit
+def compute_space_average_jax(m1: jnp.ndarray) -> float:
+    """
+    Compute space average using midpoint method on GPU (JIT compiled).
+
+    Args:
+        m1: First component of magnetization (shape (nx, ny, nz))
+
+    Returns:
+        Space average of m1
+    """
+    # Get dimensions directly from the array shape
+    Jx, Jy, Jz = m1.shape
+
+    # Create 3D coordinate grids using the shape
+    i_coords = jnp.arange(Jx)
+    j_coords = jnp.arange(Jy)
+    k_coords = jnp.arange(Jz)
+
+    # Create 3D coordinate grids
+    ii, jj, kk = jnp.meshgrid(i_coords, j_coords, k_coords, indexing="ij")
+
+    # Apply midpoint weights (0.5 on edges, 1.0 elsewhere)
+    weights = jnp.ones_like(m1)
+    weights = jnp.where((ii == 0) | (ii == Jx - 1), weights * 0.5, weights)
+    weights = jnp.where((jj == 0) | (jj == Jy - 1), weights * 0.5, weights)
+    weights = jnp.where((kk == 0) | (kk == Jz - 1), weights * 0.5, weights)
+
+    # Compute weighted sum and normalize
+    weighted_sum = jnp.sum(weights * m1)
+
+    # Compute ncell from the weights (this is the effective cell count)
+    ncell = jnp.sum(weights)
+
+    return weighted_sum / ncell
+
+
+@jax.jit
+def cross_product(a: jnp.ndarray, b: jnp.ndarray) -> jnp.ndarray:
+    r"""
+    Compute cross product :math:`a \times b` (JIT compiled).
+
+    Args:
+        a: First vector (shape (3, nx, ny, nz))
+        b: Second vector (shape (3, nx, ny, nz))
+
+    Returns:
+        Cross product :math:`a \times b` (shape (3, nx, ny, nz))
+    """
+    # Use JAX's optimized cross product function directly on axis 0
+    return jnp.cross(a, b, axis=0)
+
+
+# JIT compile the slope computation for performance
+@jax.jit
+def compute_slope(
+    g_params: dict, e_params: dict, m: jnp.ndarray, R_random: jnp.ndarray
+) -> jnp.ndarray:
+    """
+    JIT-compiled version of compute_slope_jax using modular sub-functions.
+
+    Args:
+        g_params: Grid parameters dict (dx, dy, dz)
+        e_params: Element parameters dict (coeff_1, coeff_2, coeff_3, lambda_G,
+            anisotropy)
+        m: Magnetization array (shape (3, nx, ny, nz))
+        R_random: Random field array (shape (3, nx, ny, nz))
+
+    Returns:
+        Slope array (shape (3, nx, ny, nz))
+    """
+    # Extract parameters
+    dx, dy, dz = g_params["dx"], g_params["dy"], g_params["dz"]
+    coeff_1 = e_params["coeff_1"]
+    coeff_2 = e_params["coeff_2"]
+    coeff_3 = e_params["coeff_3"]
+    lambda_G = e_params["lambda_G"]
+    anisotropy = e_params["anisotropy"]
+
+    # Compute components using modular sub-functions
+    H_aniso = compute_H_anisotropy(m, coeff_2, anisotropy)
+    laplacian_m = compute_laplacian(m, dx, dy, dz)
+
+    # Effective field
+    R_eff = coeff_1 * laplacian_m + R_random + H_aniso
+    R_eff = R_eff.at[0].add(coeff_3)
+
+    # Cross products using modular functions
+    m_cross_R_eff = cross_product(m, R_eff)
+    m_cross_m_cross_R_eff = cross_product(m, m_cross_R_eff)
+
+    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.
+
+    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')
+
+    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}"
+            )
+
+    # 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
+
+    print(f"Available JAX devices: {jax.devices()}")
+    print(f"Using JAX on device: {jax.devices()[0]}")
+    print(f"Precision: {precision} ({jnp_float})")
+
+    # Initialize random key for JAX
+    key = random.PRNGKey(seed)
+
+    g = Grid(Jx, Jy, Jz, dx)
+    dims = g.dims
+
+    e = element_class(T, H_ext, g, dt)
+    print(f"CFL = {e.get_CFL()}")
+
+    # Prepare parameters for JIT compilation using to_dict methods
+    g_params = g.to_dict()
+    e_params = e.to_dict()
+
+    # --- Initialization ---
+    def theta_init(shape):
+        """Initialization of theta."""
+        return jnp.zeros(shape, dtype=jnp_float)
+
+    def phi_init(t, shape):
+        """Initialization of phi."""
+        return jnp.zeros(shape, dtype=jnp_float) + e.gamma_0 * H_ext * t
+
+    m_n = jnp.zeros((3,) + dims, dtype=jnp_float)
+
+    theta = theta_init(dims)
+    phi = phi_init(0, dims)
+
+    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))
+
+    f_mean, f_profiles, output_filenames = get_output_files(g, T, n_mean, n_profile)
+
+    t = 0.0
+    m1_average = 0.0
+
+    # === JIT WARMUP: Pre-compile all functions to exclude compilation time ===
+    print("Warming up JIT compilation...")
+
+    # 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)
+
+    # 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])
+
+    # Force compilation and execution to complete
+    jax.block_until_ready(m_n)
+    print("JIT warmup completed.")
+
+    start_time = time.perf_counter()
+
+    for n in progress_bar(range(1, N + 1), "Iteration : ", 40):
+        t += dt
+
+        # Generate random field for temperature effect
+        key, subkey = random.split(key)
+        R_random = e.coeff_4 * random.normal(subkey, (3,) + dims, dtype=jnp_float)
+
+        # 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)
+
+        # Update magnetization
+        m_n = m_n + dt * 0.5 * (s_pre + s_cor)
+
+        # Renormalize to unit sphere
+        norm = jnp.sqrt(m_n[0] ** 2 + m_n[1] ** 2 + m_n[2] ** 2)
+        m_n = m_n / norm
+
+        # 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")
+
+    total_time = time.perf_counter() - start_time
+
+    close_output_files(f_mean, f_profiles)
+
+    if n > start_averaging:
+        m1_average /= N - start_averaging
+
+    return total_time, output_filenames, m1_average