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

llg3d/post/process.py

@@ -1,105 +1,112 @@
-#!/usr/bin/env python3
 """
-Post-processes a set of runs grouped into a `run.json` file or
-into a set of SLURM job arrays:
+Post-processes a set of runs.
 
 1. Extracts result data,
 2. Plots the computed average magnetization against temperature,
-3. Interpolates the computed points using cubic splines,
-4. Determines the Curie temperature as the value corresponding to the minimal (negative) slope of the interpolated curve.
+3. Interpolates the computed points using a PCHIP interpolator,
+4. Determines the Curie temperature as the value below which the magnetization
+   drops under 0.1.
 """
 
-import json
 from pathlib import Path
 
+
 import numpy as np
-from scipy.interpolate import interp1d
+from scipy.interpolate import PchipInterpolator
+from scipy.optimize import brentq
 
+from llg3d.io import load_results
 
-class MagData:
-    """
-    Class to handle magnetization data and interpolation according to temperature
-    """
 
-    n_interp = 200
+class MagTempData:
+    """
+    Handle magnetization data using the npz format.
 
-    def __init__(self, job_dir: Path = None, run_file: Path = Path("run.json")) -> None:
+    - Extracts result data,
+    - Interpolates the computed points using a PCHIP interpolator,
+    - Determines the Curie temperature as the value below which the magnetization
+      drops under 0.1.
 
-        if job_dir:
-            self.parentpath = job_dir
-            data, self.run = self.process_slurm_jobs()
-        elif run_file:
-            self.parentpath = run_file.parent
-            data, self.run = self.process_json(run_file)
+    Args:
+        *files: Paths to the result .npz files
+    """
 
-        self.temperature = data[:, 0]
-        self.m1_mean = data[:, 1]
-        self.interp = interp1d(self.temperature, self.m1_mean, kind="cubic")
-        self.T = np.linspace(
+    n_interp = 200  #: number of interpolation points
+
+    def __init__(self, *files: Path | str) -> None:
+        #: list of result files
+        self.files: list[Path] = [Path(file) for file in files]
+        self.params: dict = {}  #: common parameters of the runs
+        # Extract data from the runs
+        data = self._process_jobs()
+        self.temperature = data[:, 0]  #: temperatures from the runs
+        self.m1_mean = data[:, 1]  #: mean magnetization
+        # Use PCHIP interpolator which preserves positivity and monotonicity
+        #: interpolated magnetization function
+        self.interp = PchipInterpolator(self.temperature, self.m1_mean)
+        self.temperature_interp = np.linspace(
             self.temperature.min(), self.temperature.max(), self.n_interp
-        )
+        )  #: finer temperature grid for interpolation
 
-    def process_slurm_jobs(self) -> tuple[np.array, dict]:
+    @property
+    def T_Curie(self) -> float:
         """
-        Iterates through calculation directories to assemble data.
+        Return the Curie temperature.
 
-        Args:
-            parentdir (str): path to the directory containing the runs
+        It is defined as the temperature at which the magnetization equals 0.1,
+        found using a root-finding algorithm for precision.
 
         Returns:
-            tuple: (data, run) where data is a numpy array (T, <m>) and run
-            is a descriptive dictionary of the run
-        """
-        json_filename = "run.json"
-
-        # List of run directories
-        jobdirs = [f for f in self.parentpath.iterdir() if f.is_dir()]
-        if len(jobdirs) == 0:
-            exit(f"No job directories found in {self.parentpath}")
-        data = []
-        # Iterating through run directories
-        for jobdir in jobdirs:
-            try:
-                # Reading the JSON file
-                with open(jobdir / json_filename) as f:
-                    run = json.load(f)
-                    # Adding temperature and averaging value to the data list
-                    data.extend(
-                        [[float(T), res["m1_mean"]] for T, res in run["results"].items()]
-                    )
-            except FileNotFoundError:
-                print(f"Warning: {json_filename} file not found " f"in {jobdir.as_posix()}")
-
-        data.sort()  # Sorting by increasing temperatures
+            float: Curie temperature
 
-        return np.array(data), run
-
-
-    def process_json(json_filepath: Path) -> tuple[np.array, dict]:
+        Raises:
+            ValueError: If the magnetization never crosses 0.1 in the dataset
         """
-        Reads the run.json file and extracts result data.
-
-        Args:
-            json_filepath: path to the run.json file
+        # Check if magnetization ever crosses 0.1
+        T_min = self.temperature.min()
+        T_max = self.temperature.max()
+        m1_min = self.interp(T_min)
+        m1_max = self.interp(T_max)
+
+        # If 0.1 is never reached
+        if (m1_min < 0.1 and m1_max < 0.1) or (m1_min > 0.1 and m1_max > 0.1):
+            raise ValueError(
+                f"Magnetization never crosses 0.1 in the dataset. "
+                f"Range: [{m1_min:.4f}, {m1_max:.4f}]"
+            )
+
+        # Find the exact temperature where m = 0.1 using Brent's method
+        T_curie = brentq(lambda T: float(self.interp(T)) - 0.1, T_min, T_max)
+        return float(T_curie)
+
+    def _process_jobs(self) -> np.ndarray:
+        """
+        Iterates through calculation directories to assemble data.
 
         Returns:
-            tuple: (data, run) where data is a numpy array (T, <m>) and run
-            is a descriptive dictionary of the run
+            data a numpy array (T, <m>)
+
+        Raises:
+            ValueError: If any file does not end with .npz
         """
-        with open(json_filepath) as f:
-            run = json.load(f)
+        for file in self.files:
+            if not file.name.endswith(".npz"):
+                raise ValueError(f"File {file} should end with .npz")
+        # Get parameters from the first run file
 
-        data = [[int(T), res["m1_mean"]] for T, res in run["results"].items()]
+        first_results = load_results(self.files[0])
+        self.params = first_results.params.as_dict()  # Store common parameters
+        data = []
+        # Iterating through run directories
+        for file in self.files:
+            print(f"Processing file: {file}")
+            run_results = load_results(file)
+            params = run_results.params.as_dict()
+            m1_mean = np.nan
+            if "observables" in run_results.results:
+                m1_mean = run_results.results["observables"].get("m1_mean", np.nan)
+            data.append([params["T"], m1_mean])
 
         data.sort()  # Sorting by increasing temperatures
 
-        return np.array(data), run
-
-    @property
-    def T_Curie(self) -> float:
-        """
-        Return the Curie temperature defined as the temperature at
-        which the magnetization is below 0.1
-        """
-        i_max = np.where(0.1 - self.interp(self.T) > 0)[0].min()
-        return self.T[i_max]
+        return np.array(data)

llg3d/post/utils.py

@@ -0,0 +1,38 @@
+"""Command-line argument parsing for post-processing scripts."""
+
+import argparse
+from pathlib import Path
+
+
+def get_cli_args(
+    description: str | None, default_image_filepath: Path
+) -> argparse.Namespace:
+    """
+    Parse command-line arguments for post-processing scripts.
+
+    Args:
+        description: Description of the script for the help message.
+        default_image_filepath: Default path to save the output image.
+
+    Returns:
+        Parsed command-line arguments.
+    """
+    parser = argparse.ArgumentParser(
+        description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument("files", nargs="+", type=Path, help="Path to the result files.")
+    parser.add_argument(
+        "-i",
+        "--image_filepath",
+        type=Path,
+        default=default_image_filepath,
+        help="Path to save the image",
+    )
+    parser.add_argument(
+        "-s",
+        "--show",
+        action="store_true",
+        default=False,
+        help="Display the plot (omit to disable display in non-interactive runs).",
+    )
+    return parser.parse_args()

llg3d/post/x_profiles.py

@@ -0,0 +1,141 @@
+"""Plot x-profiles of magnetization from a result file."""
+
+import argparse
+from pathlib import Path
+
+from matplotlib import pyplot as plt
+
+from ..grid import Grid
+from ..io import RunResults, load_results
+from .extract import extract_values
+
+
+def parse_slice(slice_str: str) -> slice:
+    """
+    Parse a string representation of a slice into a slice object.
+
+    Handles formats like 'start:stop:step', 'start:stop', ':stop', '::step', etc.
+
+    Args:
+        slice_str: String representation of the slice.
+
+    Returns:
+        A slice object corresponding to the input string.
+    """
+    parts = slice_str.split(":")
+    # Convert parts to int or None
+    args = [int(p) if p else None for p in parts]
+    return slice(*args)
+
+
+def plot_x_profiles(
+    file: Path,
+    image_filepath: Path | str | None = None,
+    show: bool = False,
+    time_slice: str = ":",
+    m_index: int = 0,
+):
+    """
+    Plot the results from the given files.
+
+    Args:
+        file: Path to the result file.
+        image_filepath: Path to the output image file.
+            if None, the image will not be saved.
+        show: display the graph in a graphical window.
+        time_slice: String representing the time slice to plot (e.g. ":-5", "::10").
+        m_index: Index of the magnetization component to plot (0, 1, or 2).
+    """
+    fig, ax = plt.subplots()
+    result: RunResults = load_results(file)
+    grid_params = extract_values(
+        file, "params/Jx", "params/Jy", "params/Jz", "params/dx"
+    )
+    grid = Grid(*grid_params)  # type: ignore
+    x_coords = grid.get_x_coords(local=False)
+    x_profiles = result.get_record("x_profiles")
+
+    # Parse and apply slice
+    sl = parse_slice(time_slice)
+    times = x_profiles["t"][sl]
+    m_profiles = x_profiles[f"m{m_index}"][sl]
+
+    for i, time in enumerate(times):
+        ax.plot(x_coords, m_profiles[i], label=f"t = {time:.3e} s")
+    ax.set_xlabel("x")
+    ax.set_ylabel(rf"$m_{{{m_index}}}$")
+    ax.grid()
+    ax.legend()
+    ax.set_title(rf"Longitudinal profiles of magnetization $m_{{{m_index + 1}}}$")
+
+    if show:
+        plt.show()
+
+    if image_filepath is not None:
+        fig.savefig(image_filepath, dpi=300)
+        print(f"Written to {image_filepath}")
+
+
+def get_cli_args(
+    description: str | None, default_image_filepath: Path
+) -> argparse.Namespace:
+    """
+    Parse command-line arguments for post-processing scripts.
+
+    Args:
+        description: Description of the script for the help message.
+        default_image_filepath: Default path to save the output image.
+
+    Returns:
+        Parsed command-line arguments.
+    """
+    parser = argparse.ArgumentParser(
+        description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument("file", type=Path, help="Path to the result file.")
+    parser.add_argument(
+        "-i",
+        "--image_filepath",
+        type=Path,
+        default=default_image_filepath,
+        help="Path to save the image",
+    )
+    parser.add_argument(
+        "-s",
+        "--show",
+        action="store_true",
+        default=False,
+        help="Display the plot (omit to disable display in non-interactive runs).",
+    )
+    parser.add_argument(
+        "-t",
+        "--time",
+        type=str,
+        default=":",
+        help="Slice for time steps (e.g., '-5:' for last 5, '::10' for every 10th).",
+    )
+    parser.add_argument(
+        "-m",
+        "--m_index",
+        type=int,
+        choices=[1, 2, 3],
+        default=1,
+        help="Index of the magnetization component to plot (1, 2, or 3).",
+    )
+    return parser.parse_args()
+
+
+def main():  # pragma: no cover
+    """Parse CLI arguments and call the plot function."""
+    args = get_cli_args(
+        description=__doc__,
+        default_image_filepath=Path("x_profile.png"),
+    )
+
+    plot_x_profiles(
+        args.file,
+        image_filepath=args.image_filepath,
+        show=args.show,
+        time_slice=args.time,
+        m_index=args.m_index,
+    )

llg3d/py.typed

@@ -0,0 +1 @@
+# Marker file to indicate llg3d is PEP 561 type-annotated

llg3d/solvers/__init__.py

@@ -0,0 +1,153 @@
+"""
+Define various types of solvers.
+
+Example:
+    To initialize one of the solver classes:
+
+    >>> from llg3d.parameters import RunParameters
+    >>> from llg3d.solvers.numpy import NumpySolver
+    >>> solver = NumpySolver(**RunParameters(solver="numpy").as_dict())
+"""
+
+from typing import TYPE_CHECKING
+
+import os
+import importlib.util
+
+if TYPE_CHECKING:
+    from .base import BaseSolver
+
+
+def get_size() -> int:
+    """
+    Return the number of parallel MPI processes.
+
+    Use environment variables to avoid initializing MPI unnecessarily.
+
+    Returns:
+        Number of MPI processes if in an MPI environment, else 1.
+    """
+    # Open MPI
+    if "OMPI_COMM_WORLD_SIZE" in os.environ:
+        return int(os.environ["OMPI_COMM_WORLD_SIZE"])
+
+    # MPICH/Intel MPI
+    if "PMI_SIZE" in os.environ:
+        return int(os.environ["PMI_SIZE"])
+
+    # SLURM
+    if "SLURM_NTASKS" in os.environ:
+        return int(os.environ["SLURM_NTASKS"])
+
+    return 1
+
+
+def get_rank() -> int:
+    """
+    Return the rank of the current MPI process.
+
+    Use environment variables to avoid initializing MPI unnecessarily.
+
+    Returns:
+        Rank of the current MPI process if in an MPI environment, else 0.
+    """
+    # PMIx
+    if "PMIX_RANK" in os.environ:
+        return int(os.environ["PMIX_RANK"])
+
+    # Open MPI
+    if "OMPI_COMM_WORLD_RANK" in os.environ:
+        return int(os.environ["OMPI_COMM_WORLD_RANK"])
+
+    # MPICH/Intel MPI
+    if "PMI_RANK" in os.environ:
+        return int(os.environ["PMI_RANK"])
+
+    # SLURM
+    if "SLURM_PROCID" in os.environ:
+        return int(os.environ["SLURM_PROCID"])
+
+    return 0
+
+
+__all__ = [
+    "rank",
+    "size",
+    "comm",
+    "status",
+    "mpi_initialized",
+    "get_size",
+    "get_rank",
+    "get_solver_class",
+    "LIB_AVAILABLE",
+]
+
+LIB_AVAILABLE: dict[str, bool] = {}
+
+# Check for other solver availability
+for lib in "pyopencl", "jax", "mpi4py":
+    LIB_AVAILABLE[lib] = importlib.util.find_spec(lib, package=__package__) is not None
+
+
+# MPI library: initialize dummy variables at first:
+# it prevents from initializing the MPI communicator if not needed
+class _DummyComm:
+    pass
+
+
+class _DummyStatus:
+    pass
+
+
+comm = _DummyComm()
+rank = get_rank()
+size = get_size()
+status = _DummyStatus()
+mpi_initialized = False
+
+
+def get_solver_class(solver_name: str) -> "type[BaseSolver]":
+    """
+    Get the solver class based on the solver name.
+
+    Args:
+        solver_name: Name of the solver ("mpi", "numpy", "opencl", "jax")
+
+    Returns:
+        The solver class
+
+    Raises:
+        ValueError: If the selected solver is not compatible with MPI
+            or if the solver name is unknown
+
+    Example:
+        >>> Solver = get_solver_class("numpy")
+        >>> Solver.__name__
+        "NumpySolver"
+    """
+    if size > 1 and solver_name != "mpi":
+        raise ValueError(f"Solver method '{solver_name}' is not compatible with MPI.")
+
+    Solver: type[BaseSolver]
+    if solver_name == "mpi":
+        from .mpi import MPISolver
+
+        Solver = MPISolver
+    elif solver_name == "numpy":
+        from .numpy import NumpySolver
+
+        Solver = NumpySolver
+
+    elif solver_name == "opencl":
+        from .opencl import OpenCLSolver
+
+        Solver = OpenCLSolver
+
+    elif solver_name == "jax":
+        from .experimental.jax import JaxSolver
+
+        Solver = JaxSolver
+    else:
+        raise ValueError(f"Unknown solver method '{solver_name}'.")
+
+    return Solver