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

llg3d/grid.py

@@ -1,9 +1,10 @@
-"""Module to define the grid for the simulation."""
+"""Define the computational grid for the simulation."""
 
-from dataclasses import dataclass
+from typing import ClassVar
+from dataclasses import dataclass, field, asdict
 import numpy as np
 
-from .solver import rank, size
+from . import solvers
 
 
 @dataclass
@@ -15,109 +16,98 @@ class Grid:
     Jy: int  #: number of points in y direction
     Jz: int  #: number of points in z direction
     dx: float  #: grid spacing in x direction
+    dy: float = field(init=False)  #: grid spacing in y direction
+    dz: float = field(init=False)  #: grid spacing in z direction
+    dV: float = field(init=False)  #: elemental volume
+    Lx: float = field(init=False)  #: physical length in x direction
+    Ly: float = field(init=False)  #: physical length in y direction
+    Lz: float = field(init=False)  #: physical length in z direction
+    dims: tuple[int, int, int] = field(init=False)  #: local grid dimensions
+    V: float = field(init=False)  #: total volume
+    ntot: int = field(init=False)  #: total number of points
+    ncell: int = field(init=False)  #: total number of cells
+    inv_dx2: float = field(init=False)  #: :math:`1/dx^2`
+    inv_dy2: float = field(init=False)  #: :math:`1/dy^2`
+    inv_dz2: float = field(init=False)  #: :math:`1/dz^2`
+    center_coeff: float = field(init=False)  #: center coefficient for Laplacian
+    uniform: ClassVar[bool] = True  #: whether the grid is uniform
 
     def __post_init__(self) -> None:
         """Compute grid characteristics."""
-        self.dy = self.dz = self.dx  # Setting dx = dy = dz
+        self.dy = self.dz = self.dx  # Enforce dx = dy = dz
         self.Lx = (self.Jx - 1) * self.dx
         self.Ly = (self.Jy - 1) * self.dy
         self.Lz = (self.Jz - 1) * self.dz
-        # shape of the local array to the process
-        self.dims = self.Jx // size, self.Jy, self.Jz
-        # elemental volume of a cell
+        self.dims = self.Jx // solvers.size, self.Jy, self.Jz
         self.dV = self.dx * self.dy * self.dz
-        # total volume
         self.V = self.Lx * self.Ly * self.Lz
-        # total number of points
         self.ntot = self.Jx * self.Jy * self.Jz
         self.ncell = (self.Jx - 1) * (self.Jy - 1) * (self.Jz - 1)
+        # precompute the Laplacian coefficients (uniform grid spacing)
+        self.inv_dx2 = self.inv_dy2 = self.inv_dz2 = 1 / self.dx**2
+        self.center_coeff = -6.0 * self.inv_dx2
 
     def __str__(self) -> str:
-        """Print grid information."""
+        """Return grid information."""
         header = "\t\t".join(("x", "y", "z"))
         s = f"""\
+---
 \t{header}
 J\t{self.Jx}\t\t{self.Jy}\t\t{self.Jz}
 L\t{self.Lx:.08e}\t{self.Ly:.08e}\t{self.Lz:.08e}
 d\t{self.dx:.08e}\t{self.dy:.08e}\t{self.dz:.08e}
-
+---
 dV    = {self.dV:.08e}
 V     = {self.V:.08e}
 ntot  = {self.ntot:d}
 ncell = {self.ncell:d}
-"""
+---"""
         return s
 
-    def get_filename(
-        self, T: float, name: str = "m1_mean", extension: str = "txt"
-    ) -> str:
+    def get_x_coords(
+        self, local: bool = True, dtype: np.dtype = np.dtype(np.float64)
+    ) -> np.ndarray:
         """
-        Returns the output file name for a given temperature.
+        Returns the x coordinates.
 
         Args:
-            T: temperature
-            name: file name
-            extension: file extension
+            local: if True, returns the local coordinates,
+                   otherwise the global coordinates
+            dtype: data type of the coordinates
 
         Returns:
-            file name
-
-        >>> g = Grid(Jx=300, Jy=21, Jz=21, dx=1.e-9)
-        >>> g.get_filename(1100)
-        'm1_mean_T1100_300x21x21.txt'
+            1D array with the x coordinates
         """
-        suffix = f"T{int(T)}_{self.Jx}x{self.Jy}x{self.Jz}"
-        return f"{name}_{suffix}.{extension}"
+        x_global = np.linspace(0, self.Lx, self.Jx, dtype=dtype)  # global coordinates
+        # Split x into local parts if needed
+        return x_global if not local else np.split(x_global, solvers.size)[solvers.rank]
 
     def get_mesh(
-        self, local: bool = True, dtype=np.float64
+        self, local: bool = True, dtype: np.dtype = np.dtype(np.float64)
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         """
         Returns a meshgrid of the coordinates.
 
+        Use ij indexing.
+
         Args:
             local: if True, returns the local coordinates,
                    otherwise the global coordinates
-            dtype: data type of the coordinates)
+            dtype: data type of the coordinates
 
         Returns:
-            tuple of 3D arrays with the coordinates
+            Tuple of 3D arrays with the coordinates
         """
-        x_global = np.linspace(0, self.Lx, self.Jx, dtype=dtype)  # global coordinates
+        x = self.get_x_coords(local=local, dtype=dtype)
         y = np.linspace(0, self.Ly, self.Jy, dtype=dtype)
         z = np.linspace(0, self.Lz, self.Jz, dtype=dtype)
-        if local:
-            x_local = np.split(x_global, size)[rank]  # local coordinates
-            return np.meshgrid(x_local, y, z, indexing="ij")
-        else:
-            return np.meshgrid(x_global, y, z, indexing="ij")
+        return np.meshgrid(x, y, z, indexing="ij")
 
-    def to_dict(self) -> dict:
+    def as_dict(self) -> dict:
         """
         Export grid parameters to a dictionary for JAX JIT compatibility.
 
         Returns:
             Dictionary containing grid parameters needed for computations
         """
-        return {
-            "dx": self.dx,
-            "dy": self.dy,
-            "dz": self.dz,
-            "Jx": self.Jx,
-            "Jy": self.Jy,
-            "Jz": self.Jz,
-            "dV": self.dV,
-        }
-
-    def get_laplacian_coeff(self) -> tuple[float, float, float, float]:
-        """
-        Returns the coefficients for the laplacian computation.
-
-        Returns:
-            Tuple of coefficients (dx2_inv, dy2_inv, dz2_inv, center_coeff)
-        """
-        dx2_inv = 1 / self.dx**2
-        dy2_inv = 1 / self.dy**2
-        dz2_inv = 1 / self.dz**2
-        center_coeff = -2 * (dx2_inv + dy2_inv + dz2_inv)
-        return dx2_inv, dy2_inv, dz2_inv, center_coeff
+        return asdict(self)

llg3d/io.py

@@ -0,0 +1,395 @@
+"""
+Input/Output functions.
+
+Extensibility Guide
+===================
+
+The I/O system is designed to be fully extensible. You can add arbitrary fields
+to `metrics` and `records` without modifying this module.
+
+Adding custom metrics
+---------------------
+
+    >>> solver.metrics['custom_value'] = 42.0
+    >>> solver.metrics['convergence_rate'] = 0.95
+
+Adding custom records
+---------------------
+
+    >>> # Simple array
+    >>> solver.records['energy_evolution'] = np.array([...])
+
+    >>> # Nested structure (automatically flattened in .npz)
+    >>> solver.records['field_snapshots'] = {
+    ...     'times': np.array([0, 1, 2]),
+    ...     'data': np.array([...])
+    ... }
+
+    >>> # Deeply nested (unlimited depth)
+    >>> solver.records['analysis'] = {
+    ...     'spectra': {
+    ...         'fourier': np.array([...]),
+    ...         'wavelets': np.array([...])
+    ...     }
+    ... }
+
+TypedDicts (`Metrics`, `Records`, etc.) are documentation only. Runtime accepts
+any dict[str, Any] structure.
+"""
+
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import (
+    Any,
+    NotRequired,
+    TextIO,
+    TypedDict,
+    cast,
+)
+
+import numpy as np
+
+from .parameters import RunParameters
+from .solvers import rank, size
+from .solvers.profiling import ProfilingStats
+
+
+def get_tqdm_file() -> TextIO | None:
+    """Get a TQDM-compatible file for progress bar output in MPI."""
+    if size == 1:
+        return sys.stdout  # single process uses default
+    if rank != 0:
+        return None  # other ranks disable progress bar
+    try:
+        # For MPI, try to open /dev/tty for direct terminal output
+        return open("/dev/tty", "w")
+    except OSError:
+        return sys.stdout  # fallback
+
+
+class MetricsRequired(TypedDict):
+    """Required fields for simulation metrics."""
+
+    total_time: float  #: Total wall-clock time of the simulation
+    time_per_ite: float  #: Average time per iteration
+    efficiency: float  #: Time per iteration per cell
+    CFL: float  #: CFL condition value
+
+
+class Metrics(MetricsRequired, total=False):
+    """
+    Structure for simulation metrics (extensible).
+
+    Required fields: total_time, time_per_ite, efficiency, CFL
+    All other fields are optional and can be added dynamically.
+    """
+
+    #: Profiling statistics
+    profiling_stats: NotRequired[ProfilingStats]
+
+
+class Observables(TypedDict, total=False):
+    """
+    Physical observables from the simulation (extensible).
+
+    All fields are optional and can be added dynamically.
+    """
+
+    m1_mean: float  #: Time-averaged magnetization in x direction
+
+
+class XProfiles(TypedDict):
+    """Structure for cross-sectional profiles (arrays in final records)."""
+
+    t: np.ndarray  #: Time points for profiles
+    m1: np.ndarray  #: m1 component profiles
+    m2: np.ndarray  #: m2 component profiles
+    m3: np.ndarray  #: m3 component profiles
+
+
+class XProfilesBuffer(TypedDict):
+    """Structure for cross-sectional profiles during accumulation (lists)."""
+
+    t: list[float]  #: Time points for profiles
+    m1: list[np.ndarray]  #: m1 component profiles
+    m2: list[np.ndarray]  #: m2 component profiles
+    m3: list[np.ndarray]  #: m3 component profiles
+
+
+class RecordsBuffer(TypedDict, total=False):
+    """
+    Records during simulation (accumulation phase with lists).
+
+    BaseSolver.records uses this structure during simulation, accumulating
+    data as lists. Before saving, BaseSolver.save() converts lists to arrays.
+    """
+
+    xyz_average: list[tuple[float, float]]  #: Accumulated (t, value) pairs
+    x_profiles: XProfilesBuffer  #: Profiles during accumulation
+
+
+class Records(TypedDict, total=False):
+    """
+    Records after saving (finalized with numpy arrays).
+
+    Returned by load_results() with all data as read-only numpy arrays.
+    """
+
+    xyz_average: np.ndarray  #: Space-averaged magnetization over time (shape: (2, N))
+    x_profiles: XProfiles  #: Cross-sectional profiles in yz plane
+
+
+class SimulationResults(TypedDict):
+    """Structure for simulation results."""
+
+    metrics: Metrics  #: Simulation metrics (performance, numerical quality)
+    observables: NotRequired[Observables]  #: Physical observables (results)
+    records: NotRequired[Records]  #: Optional time-series records
+
+
+@dataclass
+class RunResults:
+    """Structure for loaded simulation results."""
+
+    params: RunParameters
+    results: SimulationResults
+    file: str | Path
+
+    def get_record(self, record_name: str) -> np.ndarray | dict[str, np.ndarray]:
+        """Return a record by name from results or raise a descriptive error."""
+        try:
+            if "records" not in self.results:
+                raise KeyError(f"RunResults from '{self.file}' has no records.")
+
+            records = self.results["records"]
+            # Cast to dict to allow dynamic access with a variable key for mypy
+            record = cast(dict, records)[record_name]
+            if isinstance(record, (np.ndarray, dict)):
+                return record
+            else:
+                raise TypeError(f"'{record_name}' is not in the expected format.")
+        except KeyError as e:
+            msg = (
+                f"RunResults from '{self.file}' does not contain the required record "
+                f"'{record_name}'."
+            )
+            raise KeyError(msg) from e
+
+
+def save_results(
+    output_file: str | Path,
+    params: RunParameters,
+    metrics: Metrics,
+    observables: Observables | None = None,
+    records_buffer: RecordsBuffer | None = None,
+) -> None:
+    """
+    Saves simulation results to a .npz file with hierarchical structure.
+
+    Args:
+        output_file: Path to the output .npz file.
+        params: Dataclass of simulation parameters.
+        metrics: Dictionary of simulation metrics. Must contain required fields:
+            total_time, time_per_ite, CFL. Additional fields are allowed.
+        observables: Dictionary of physical observables (results).
+        records_buffer: RecordsBuffer or custom dict of records
+            (accumulation phase with lists/dicts).
+
+    Note:
+        Both metrics and records accept any key-value pairs, allowing easy
+        extension without modifying this function. See Metrics and Records
+        TypedDict for recommended structure.
+
+    Example:
+        >>> records_buffer = {
+        ...     'xyz_average': arr1,
+        ...     'x_profiles': {'t': t, 'm1': m1, ...},
+        ...     'custom_data': arr2  # Any new field works automatically
+        ... }
+        >>> save_results('run.npz', params, metrics, records_buffer=records_buffer)
+    """
+    data_to_save: dict[str, Any] = {}
+
+    if records_buffer is None:
+        records_buffer = {}
+    if observables is None:
+        observables = {}
+
+    # Helper to convert lists to numpy arrays recursively
+    def _convert_lists_to_arrays(d: Any) -> Any:
+        """Recursively convert lists to numpy arrays."""
+        if isinstance(d, list):
+            return np.array(d)
+        elif isinstance(d, dict):
+            return {k: _convert_lists_to_arrays(v) for k, v in d.items()}
+        else:
+            return d
+
+    # Convert lists in records_buffer to arrays
+    records_buffer = _convert_lists_to_arrays(records_buffer)
+
+    # Helper to flatten nested dicts into NPZ keys
+    def _flatten_dict(d: dict | Any, prefix: str):
+        """Recursively flatten nested dicts into NPZ keys."""
+        if isinstance(d, dict):
+            for key, value in d.items():
+                full_key = f"{prefix}/{key}"
+                _flatten_dict(value, full_key)
+        else:
+            # Leaf value - save as array
+            data_to_save[prefix] = d
+
+    # Flatten params
+    _flatten_dict(params.as_dict(), "params")
+
+    # Flatten metrics
+    _flatten_dict(metrics, "results/metrics")
+
+    # Flatten observables
+    if observables:
+        _flatten_dict(observables, "results/observables")
+
+    # Flatten records
+    _flatten_dict(records_buffer, "results/records")
+
+    np.savez(output_file, **data_to_save, allow_pickle=False)
+
+
+def load_results(result_file: str | Path) -> RunResults:
+    """
+    Loads simulation results from a .npz file with hierarchical structure.
+
+    Numpy arrays are set to read-only to prevent accidental modification.
+
+    Example:
+        >>> run_results = load_results("run.npz")
+        >>> params = run_results["params"]
+        >>> metrics = run_results["results"]["metrics"]
+        >>> xyz_average = run_results["results"]["records"]["xyz_average"]
+        >>> m1_prof = run_results["results"]["records"]["x_profiles"]["m1"]
+
+    Args:
+        result_file: path to the .npz file.
+
+    Returns:
+        Dictionary containing the loaded data with hierarchical structure.
+    """
+    npz_data: np.lib.npyio.NpzFile = np.load(result_file, allow_pickle=False)
+
+    def _unflatten_dict(prefix: str) -> dict[str, Any]:
+        """Reconstruct nested dict from flat NPZ keys with given prefix."""
+        result: dict[str, Any] = {}
+
+        for key in npz_data.files:
+            if not key.startswith(prefix + "/"):
+                continue
+
+            # Extract relative path after prefix
+            rel_path = key[len(prefix) + 1 :]
+            parts = rel_path.split("/")
+
+            # Navigate/create nested structure
+            current = result
+            for part in parts[:-1]:
+                if part not in current:
+                    current[part] = {}
+                current = current[part]
+
+            # Store the value
+            arr = npz_data[key]
+            arr.flags.writeable = False
+            current[parts[-1]] = arr
+
+        return result
+
+    # Load params
+    params_dict = _unflatten_dict("params")
+    params = RunParameters(**params_dict)
+
+    # Load metrics
+    metrics: Metrics = cast(Metrics, _unflatten_dict("results/metrics"))
+
+    # Load observables
+    observables: Observables | None = None
+    observables_dict = _unflatten_dict("results/observables")
+    if observables_dict:
+        observables = cast(Observables, observables_dict)
+
+    # Load records
+    records: Records = cast(Records, _unflatten_dict("results/records"))
+
+    # Build the results dict
+    simulation_results: SimulationResults = {"metrics": metrics}
+    if observables:
+        simulation_results["observables"] = observables
+    if records:
+        simulation_results["records"] = records
+
+    run_results = RunResults(
+        params=params, results=simulation_results, file=result_file
+    )
+
+    return run_results
+
+
+def format_profiling_table(
+    profiling_dict: dict[str, dict[str, float | int]], total_time: float | None = None
+) -> str:
+    """
+    Format profiling statistics as a table.
+
+    Args:
+        profiling_dict: Dictionary with profiling stats, where each entry has
+            'time' (float) and 'calls' (int) keys
+        total_time: Total simulation time for percentage calculation.
+            If None, the percentage column is omitted.
+
+    Returns:
+        A formatted table string
+    """
+    if not profiling_dict:
+        return "(empty)"
+
+    # Determine the maximum width for alignment
+    col1_width = max(len(name) for name in profiling_dict.keys())
+
+    # Header
+    if total_time is not None:
+        s = (
+            f"{'Function':<{col1_width}} | {'Calls':>5} | {'total_time (s)':>14} "
+            f"| {'%':>6} | {'Avg Time (s)':>13}\n"
+        )
+    else:
+        s = (
+            f"{'Function':<{col1_width}} | {'Calls':>5} | {'total_time (s)':>14} "
+            f"| {'Avg Time (s)':>13}\n"
+        )
+    s += "-" * (len(s) - 1) + "\n"
+
+    # Sort by total time descending
+    sorted_profiling = dict(
+        sorted(profiling_dict.items(), key=lambda item: item[1]["time"], reverse=True)
+    )
+
+    # Data rows
+    for name, stats in sorted_profiling.items():
+        func_time = stats["time"]
+        calls = stats["calls"]
+        avg_time = func_time / calls if calls > 0 else 0.0
+
+        if total_time is not None:
+            percent = (func_time / total_time * 100) if total_time > 0 else 0.0
+            s += (
+                f"{name:<{col1_width}} | {calls:>5} | {func_time:>14.6f} "
+                f"| {percent:>5.1f}% | {avg_time:>13.6f}\n"
+            )
+        else:
+            s += (
+                f"{name:<{col1_width}} | {calls:>5} | {func_time:>14.6f} "
+                f"| {avg_time:>13.6f}\n"
+            )
+
+    return s

llg3d/main.py

@@ -1,67 +1,64 @@
-"""Define a CLI for the llg3d package."""
+"""Define a CLI for running LLG3D simulations."""
 
 import argparse
 
-
-from . import rank, size, LIB_AVAILABLE
-from .parameters import parameters, get_parameter_list
-from .simulation import Simulation
-
-if LIB_AVAILABLE["mpi4py"]:
-    # Use the MPI version of the ArgumentParser
-    from .solver.mpi import ArgumentParser
-else:
-    # Use the original version of the ArgumentParser
-    from argparse import ArgumentParser
+from .parameters import arg_parameters
+from . import solvers
 
 
 def parse_args(args: list[str] | None) -> argparse.Namespace:
     """
-    Argument parser for llg3d.
+    Argument parser for LLG3D.
 
-    Automatically adds arguments from the parameter dictionary.
+    Automatically adds arguments from :class:`~llg3d.parameters.arg_parameters`.
+
+    Args:
+        args: List of command line arguments
 
     Returns:
         argparse.Namespace: Parsed arguments
     """
+    ArgumentParser: type[argparse.ArgumentParser]
+    if solvers.size > 1:
+        # Use the MPI version of the ArgumentParser
+        from .solvers.mpi import ArgumentParser
+
+    else:
+        # Use the original version of the ArgumentParser
+        from argparse import ArgumentParser
+
     parser = ArgumentParser(
         description=__doc__,
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
 
-    if size > 1:
-        parameters["solver"]["default"] = "mpi"
-
     # Automatically add arguments from the parameter dictionary
-    for name, parameter in parameters.items():
+    for name, parameter in arg_parameters.items():
         if "action" not in parameter:
             parameter["type"] = type(parameter["default"])
-        parser.add_argument(f"--{name}", **parameter)
+        parser.add_argument(f"--{name}", **parameter)  # type: ignore[arg-type]
 
     return parser.parse_args(args)
 
 
-def main(arg_list: list[str] = None):
+def main(arg_list: list[str] | None = None):
     """
     Evaluates the command line and runs the simulation.
 
     Args:
         arg_list: List of command line arguments
     """
+    if solvers.size > 1:  # Ensure MPI global variables are initialized
+        from .solvers.mpi import initialize_mpi
+
+        initialize_mpi()
+
     args = parse_args(arg_list)
 
-    if size > 1 and args.solver != "mpi":
-        raise ValueError(f"Solver method {args.solver} is not compatible with MPI.")
-    if args.solver == "mpi" and not LIB_AVAILABLE["mpi4py"]:
-        raise ValueError(
-            "The MPI solver method requires to install the mpi4py package, "
-            "for example using pip: pip install mpi4py"
-        )
-
-    if rank == 0:
-        # Display parameters as a list
-        print(get_parameter_list(vars(args)))
-
-    simulation = Simulation(vars(args))
-    simulation.run()
-    simulation.save()
+    Solver = solvers.get_solver_class(args.solver)
+    s = Solver(**vars(args))
+    if solvers.rank == 0:
+        # Display the solver parameters
+        print(s)
+    s.run()
+    s.save()