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

llg3d/solver/opencl.py

@@ -0,0 +1,330 @@
+"""LLG3D Solver using OpenCL."""
+
+import time
+from pathlib import Path
+
+import numpy as np
+import pyopencl as cl
+from pyopencl import clrandom
+from pyopencl import array as clarray
+
+from ..output import progress_bar, get_output_files, close_output_files
+from ..grid import Grid
+from ..element import Element, Cobalt
+
+
+def get_context_and_device(
+    device_selection: str = "auto",
+) -> tuple[cl.Context, cl.Device]:
+    """
+    Get the OpenCL context and device.
+
+    Args:
+        device_selection:
+
+            - ``"auto"``: Let OpenCL choose automatically
+            - ``"cpu"``: Select CPU device
+            - ``"gpu"``: Select first available GPU
+            - ``"gpu:N"``: Select specific GPU by index (e.g., ``"gpu:0"``, ``"gpu:1"``)
+
+    Returns:
+        - The OpenCL context
+        - The OpenCL device
+    """
+    if device_selection == "auto":
+        context = cl.create_some_context(interactive=False)
+        device = context.devices[0]
+        return context, device
+
+    # Get all platforms and devices
+    platforms = cl.get_platforms()
+    all_devices = []
+
+    for platform in platforms:
+        all_devices.extend(platform.get_devices())
+
+    if not all_devices:
+        raise RuntimeError("No OpenCL devices found")
+
+    # Filter devices based on selection
+    if device_selection == "cpu":
+        cpu_devices = [d for d in all_devices if d.type & cl.device_type.CPU]
+        if not cpu_devices:
+            raise RuntimeError("No CPU devices found")
+        selected_device = cpu_devices[0]
+    elif device_selection == "gpu":
+        gpu_devices = [d for d in all_devices if d.type & cl.device_type.GPU]
+        if not gpu_devices:
+            raise RuntimeError("No GPU devices found")
+        selected_device = gpu_devices[0]
+    elif device_selection.startswith("gpu:"):
+        gpu_devices = [d for d in all_devices if d.type & cl.device_type.GPU]
+        if not gpu_devices:
+            raise RuntimeError("No GPU devices found")
+
+        gpu_index = int(device_selection.split(":")[1])
+        if gpu_index >= len(gpu_devices):
+            raise RuntimeError(
+                f"GPU index {gpu_index} not available. Found {len(gpu_devices)} GPU(s)"
+            )
+        selected_device = gpu_devices[gpu_index]
+    else:
+        raise ValueError(f"Invalid device selection: {device_selection}")
+
+    # Create context with selected device
+    context = cl.Context([selected_device])
+    print(f"Selected OpenCL device: {selected_device.name} ({selected_device.type})")
+
+    return context, selected_device
+
+
+def get_precision(device: cl.Device, precision: str) -> np.dtype:
+    """
+    Get the numpy float type based on the precision.
+
+    Args:
+        device: OpenCL device
+        precision: Precision of the simulation (single or double)
+
+    Returns:
+        The numpy float type (float32 or float64)
+
+    Raises:
+        RuntimeError: If double precision is asked while the device does not support it
+    """
+    # Check that cl device supports double precision
+    if precision == "double" and not device.double_fp_config:
+        raise RuntimeError("The selected device does not support double precision.")
+
+    return np.float64 if precision == "double" else np.float32
+
+
+class Program:
+    """Class to manage the OpenCL kernels for the LLG3D simulation."""
+
+    def __init__(self, g: Grid, context: cl.Context, np_float: np.dtype):
+        self.grid = g
+        self.context = context
+        self.np_float = np_float
+        self.cl_program = self._get_built_program()
+
+    def _get_built_program(self) -> cl.Program:
+        """
+        Return the OpenCL program built from the source code.
+
+        Returns:
+            The OpenCL program object
+        """
+        opencl_code = (Path(__file__).parent / "llg3d.cl").read_text()
+        build_options = "-D USE_DOUBLE_PRECISION" if self.np_float == np.float64 else ""
+        build_options += (
+            f" -D NX={self.grid.Jx} -D NY={self.grid.Jy} -D NZ={self.grid.Jz}"
+        )
+        return cl.Program(self.context, opencl_code).build(options=build_options)
+
+    def get_kernel(self, kernel_name: str, arg_types: list = [None]) -> cl.Kernel:
+        """
+        Returns the specified kernel by name.
+
+        Args:
+            kernel_name: Name of the kernel to retrieve
+            arg_types: List of argument types for the kernel
+
+        Returns:
+            The OpenCL kernel object
+        """
+        kernel: cl.Kernel = getattr(self.cl_program, kernel_name)
+        kernel.set_arg_types(arg_types)
+        return kernel
+
+
+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 over N iterations.
+
+    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 grid object
+        - The time taken for the simulation
+        - The output filename
+    """
+    context, opencl_device = get_context_and_device(device)
+    np_float = get_precision(opencl_device, precision)
+
+    g = Grid(Jx, Jy, Jz, dx)
+    print(g)
+
+    e = element_class(T, H_ext, g, dt)
+    if not isinstance(e, Cobalt):
+        raise NotImplementedError(
+            f"Element is {type(e)} but only {Cobalt} is supported at the moment."
+        )
+    print(f"CFL = {e.get_CFL()}")
+
+    # --- Initialization ---
+
+    def theta_init(shape):
+        """Initialization of theta."""
+        return np.zeros(shape, dtype=np_float)
+
+    def phi_init(t, shape):
+        """Initialization of phi."""
+        return np.zeros(shape, dtype=np_float) + e.gamma_0 * H_ext * t
+
+    m_n = np.zeros((3,) + g.dims, dtype=np_float)
+
+    theta = theta_init(g.dims)
+    phi = phi_init(0, g.dims)
+
+    m_n[0] = np.cos(theta)
+    m_n[1] = np.sin(theta) * np.cos(phi)
+    m_n[2] = np.sin(theta) * np.sin(phi)
+
+    queue = cl.CommandQueue(context)
+
+    program = Program(g, context, np_float)
+    slope_kernel = program.get_kernel("slope", [None] * 3 + [np_float] * 8)
+    update_1_kernel = program.get_kernel("update_1", [None] * 3 + [np_float])
+    update_2_kernel = program.get_kernel("update_2", [None] * 4 + [np_float])
+    normalize_kernel = program.get_kernel("normalize")
+
+    # Create a CL array for m1 component in order to compute averages
+    d_m1 = clarray.empty(queue, g.ntot, np_float)
+    copy_m1_kernel = program.get_kernel("copy_m1", [None, None])
+
+    mf = cl.mem_flags
+    mem_size = m_n.nbytes
+
+    d_m_n = cl.Buffer(context, mf.READ_WRITE | mf.COPY_HOST_PTR, hostbuf=m_n)
+    d_R_alea = cl.array.Array(queue, (3,) + g.dims, np_float)
+    d_m_np1 = cl.Buffer(context, mf.READ_WRITE, mem_size)
+    d_s_pre = cl.Buffer(context, mf.READ_WRITE, mem_size)
+    d_s_cor = cl.Buffer(context, mf.READ_WRITE, mem_size)
+
+    rng = clrandom.PhiloxGenerator(context, seed=seed)
+
+    f_mean, f_profiles, output_filenames = get_output_files(g, T, n_mean, n_profile)
+
+    t = 0.0
+    m1_average = 0.0
+
+    start_time = time.perf_counter()
+
+    for n in progress_bar(range(1, N + 1), "Iteration : ", 40):
+        t += dt
+
+        rng.fill_normal(d_R_alea)
+        queue.finish()  # ensure the array is filled
+
+        # Prediction phase
+
+        # calculate s_i_pre from m_i^n
+        slope_kernel(
+            queue,
+            g.dims,
+            None,
+            d_m_n,
+            d_R_alea.data,
+            d_s_pre,
+            g.dx,
+            g.dy,
+            g.dz,
+            e.coeff_1,
+            e.coeff_2,
+            e.coeff_3,
+            e.coeff_4,
+            e.lambda_G,
+        )
+
+        # m_i^n+1 = m_i^n + dt * s_i_pre
+        update_1_kernel(queue, (3 * g.ntot,), None, d_m_n, d_m_np1, d_s_pre, dt)
+        queue.finish()
+
+        # # Correction phase
+
+        # calculate s_i_cor from m_i^n+1
+        slope_kernel(
+            queue,
+            g.dims,
+            None,
+            d_m_np1,
+            d_R_alea.data,
+            d_s_cor,
+            g.dx,
+            g.dy,
+            g.dz,
+            e.coeff_1,
+            e.coeff_2,
+            e.coeff_3,
+            e.coeff_4,
+            e.lambda_G,
+        )
+        # m_i^n+1 = m_i^n + dt * (s_i_pre + s_i_cor) / 2
+        # Update using the corrected values
+        update_2_kernel(
+            queue, (3 * g.ntot,), None, d_m_n, d_m_np1, d_s_pre, d_s_cor, dt
+        )
+        queue.finish()
+
+        # Normalization
+        normalize_kernel(queue, (g.ntot,), None, d_m_np1).wait()
+
+        # Swap the buffers for the next iteration
+        d_m_n, d_m_np1 = d_m_np1, d_m_n
+
+        # Space average of m_1 using the midpoint method with OpenCL sum
+        if n_mean != 0 and n % n_mean == 0:
+            # Copy only the first component from d_m_n to d_m1 with weights applied
+            # d_m_n contains [m1, m2, m3] interleaved, we want only m1
+            copy_m1_kernel(queue, g.dims, None, d_m_n, d_m1.data)
+
+            # Use PyOpenCL array sum to compute the weighted sum
+            weighted_sum = clarray.sum(d_m1).get()
+            m1_mean = weighted_sum / g.ncell
+
+            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

llg3d/solver/solver.py

@@ -0,0 +1,89 @@
+"""Common functions for the LLG3D solver."""
+
+import numpy as np
+
+from ..grid import Grid
+from ..element import Element
+
+
+def cross_product(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    r"""
+    Compute cross product :math:`a \times b`.
+
+    This implementation is faster than np.cross for large arrays.
+
+    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))
+    """
+    return np.stack(
+        [
+            a[1] * b[2] - a[2] * b[1],  # x-component
+            a[2] * b[0] - a[0] * b[2],  # y-component
+            a[0] * b[1] - a[1] * b[0],  # z-component
+        ],
+        axis=0,
+    )
+
+
+def compute_H_anisotropy(e: Element, m: np.ndarray) -> np.ndarray:
+    """
+    Compute the anisotropy field.
+
+    Args:
+        e: Element object
+        m: Magnetization array (shape (3, nx, ny, nz)).
+
+    Returns:
+        Anisotropy field array (shape (3, nx, ny, nz))
+    """
+    m1, m2, m3 = m
+
+    m1m1 = m1 * m1
+    m2m2 = m2 * m2
+    m3m3 = m3 * m3
+
+    if e.anisotropy == "uniaxial":
+        aniso_1 = m1
+        aniso_2 = np.zeros_like(m1)
+        aniso_3 = np.zeros_like(m1)
+
+    if e.anisotropy == "cubic":
+        aniso_1 = -(1 - m1m1 + m2m2 * m3m3) * m1
+        aniso_2 = -(1 - m2m2 + m1m1 * m3m3) * m2
+        aniso_3 = -(1 - m3m3 + m1m1 * m2m2) * m3
+
+    return e.coeff_2 * np.stack([aniso_1, aniso_2, aniso_3], axis=0)
+
+
+def space_average(g: Grid, m: np.ndarray, copy: bool = True) -> float:
+    """
+    Returns the spatial average of m with shape (g.dims) using the midpoint method.
+
+    Args:
+        g: Grid object
+        m: Array to be integrated
+        copy: If True, copy m to avoid modifying its value
+
+    Returns:
+        Spatial average of m
+    """
+    # copy m to avoid modifying its value
+    mm = m.copy() if copy else m
+
+    # 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() / g.ncell
+    return average

{llg3d-1.4.1.dist-info → llg3d-2.0.1.dist-info}/METADATA

@@ -1,36 +1,28 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: llg3d
-Version: 1.4.1
-Summary: Solveur pour l'équation de Landau-Lifshitz-Gilbert stochastique en 3D
+Version: 2.0.1
+Summary: A solver for the stochastic Landau-Lifshitz-Gilbert equation in 3D
 Author-email: Clémentine Courtès <clementine.courtes@math.unistra.fr>, Matthieu Boileau <matthieu.boileau@math.unistra.fr>
 Project-URL: Homepage, https://gitlab.math.unistra.fr/llg3d/llg3d
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.6
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: AUTHORS
 Requires-Dist: numpy
-Requires-Dist: mpi4py
 Requires-Dist: matplotlib
 Requires-Dist: scipy
-Provides-Extra: doc
-Requires-Dist: Sphinx >=7.2.2 ; extra == 'doc'
-Requires-Dist: myst-parser ; extra == 'doc'
-Requires-Dist: furo ; extra == 'doc'
-Requires-Dist: nbsphinx ; extra == 'doc'
-Requires-Dist: sphinx-copybutton ; extra == 'doc'
-Requires-Dist: sphinx-autobuild ; extra == 'doc'
-Requires-Dist: sphinx-prompt ; extra == 'doc'
-Requires-Dist: sphinx-last-updated-by-git ; extra == 'doc'
-Requires-Dist: sphinxcontrib-programoutput ; extra == 'doc'
-Requires-Dist: sphinxcontrib-bibtex ; extra == 'doc'
-Requires-Dist: ipython ; extra == 'doc'
-Provides-Extra: test
-Requires-Dist: pytest ; extra == 'test'
-Requires-Dist: pytest-cov ; extra == 'test'
-Requires-Dist: pytest-mpi ; extra == 'test'
+Provides-Extra: mpi
+Requires-Dist: mpi4py; extra == "mpi"
+Provides-Extra: opencl
+Requires-Dist: pyopencl; extra == "opencl"
+Requires-Dist: mako; extra == "opencl"
+Provides-Extra: jax
+Requires-Dist: jax[cuda]; sys_platform != "darwin" and extra == "jax"
+Requires-Dist: jax[cpu]; sys_platform == "darwin" and extra == "jax"
+Dynamic: license-file
 
 # LLG3D: A solver for the stochastic Landau-Lifshitz-Gilbert equation in 3D
 
@@ -40,6 +32,6 @@ Requires-Dist: pytest-mpi ; extra == 'test'
 [![Doc](https://img.shields.io/badge/doc-sphinx-blue)](https://llg3d.pages.math.unistra.fr/llg3d/)
 [![SWH](https://archive.softwareheritage.org/badge/origin/https://gitlab.math.unistra.fr/llg3d/llg3d/)](https://archive.softwareheritage.org/browse/origin/?origin_url=https://gitlab.math.unistra.fr/llg3d/llg3d)
 
-LLG3D is written in Python and utilizes the MPI library for parallelizing computations.
+LLG3D is written in Python and may run in parallel using MPI, OpenCL or JAX.
 
 See the [documentation](https://llg3d.pages.math.unistra.fr/llg3d/).

llg3d-2.0.1.dist-info/RECORD

@@ -0,0 +1,25 @@
+llg3d/__init__.py,sha256=YhFE4dB9NPWdcadzg14xMgOao1PqDgzxU2iFKEp7Vjs,187
+llg3d/__main__.py,sha256=3D1q7AG5vU6gr-V0iuo5oNYl-Og2SvJ4YaBTdqOVVaw,115
+llg3d/element.py,sha256=B2GDv5lTNTgvbIqSThEaA2o3hQRkhQdIgXg-MxzeSuc,4304
+llg3d/grid.py,sha256=goXO-sh9wJ2YDqe3IerqE6NlvUFacp328RfXHHW6bFY,3848
+llg3d/main.py,sha256=8ltKUT9jA3u1E0cEUjxemRfRPhERnIm_i5L6sP0_G0I,1888
+llg3d/output.py,sha256=HCL1V1oehbV2OERwbS0WCG00W1kNeNQpgcP33LIrGWU,2934
+llg3d/parameters.py,sha256=VoAkm8v6MiOkwRaHupWgHVTQaSDSGIKz5jWnRT8PBx0,2442
+llg3d/simulation.py,sha256=8RbAzLhlhCEZZCGxqNEczxGvfrLsnxKbRCkfpy1cCtU,3429
+llg3d/post/__init__.py,sha256=rX0jQIXGQYzFRoXSsPu2oEDB4YWK5IxFs0cJqXnQb8g,39
+llg3d/post/plot_results.py,sha256=fmCxeKmRfpHa_sNHMqczYNWQOCLJzxfRDQL22AeeJcM,1550
+llg3d/post/process.py,sha256=YKOUIgjAOD3q6_qoxzWVSC0WHBofRZJafwWzdMXWQXA,3569
+llg3d/post/temperature.py,sha256=SHf3s1LHIY0SBDKLoSMGSKo1e7EXcKeMXZdlEsChCBk,2150
+llg3d/solver/__init__.py,sha256=fqd1wZkIAi0sMiOQg0tbfLInMDjTCFkRF5NA_FV1wyU,901
+llg3d/solver/jax.py,sha256=VKn7YanidgXhpoIKEJy0sL0HZ1wV3a9F1uWmbAmWOi0,12026
+llg3d/solver/mpi.py,sha256=hJSXbZxpmi3z6PyMxyvuSBHFaqDv4S_C2FnAwAgI_r4,13041
+llg3d/solver/numpy.py,sha256=3trvsn2W-CjGdFF70oQZ6Jw8avX47OJ1DighschMcHE,6040
+llg3d/solver/opencl.py,sha256=EEExOmSUDs0UPt-E4XvPv4NA_S4SuIqw5Uo9p4KghcM,10222
+llg3d/solver/solver.py,sha256=MXscOjO0-RUIqOir1D23_8xCFw8ulDADTFQwQbMkBvc,2142
+llg3d-2.0.1.dist-info/licenses/AUTHORS,sha256=vhJ88HikYvOrGiB_l1xH2X6hyn9ZJafx6mpoMYNhh1I,297
+llg3d-2.0.1.dist-info/licenses/LICENSE,sha256=aFxTGAyyve8nM9T2jWTarJzQhdSSC3MbbN1heNAev9c,1062
+llg3d-2.0.1.dist-info/METADATA,sha256=J7kKUK6YL7cIKhHmEJLMud6rzr93AgCGozbmHIYVwPY,1948
+llg3d-2.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+llg3d-2.0.1.dist-info/entry_points.txt,sha256=kjuf09uDEuROnDeuHkVen3oAS4ThEPS084TpHmQAbJ8,133
+llg3d-2.0.1.dist-info/top_level.txt,sha256=cBZ0roaXt3CAXqYojuO84lGPCtWuLlXxLGLYRKmHZy0,6
+llg3d-2.0.1.dist-info/RECORD,,

{llg3d-1.4.1.dist-info → llg3d-2.0.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (70.1.1)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

llg3d-2.0.1.dist-info/entry_points.txt

@@ -0,0 +1,4 @@
+[console_scripts]
+llg3d = llg3d.main:main
+llg3d.plot_results = llg3d.post.plot_results:main
+llg3d.post = llg3d.post.temperature:main