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

llg3d/solvers/opencl.py

@@ -0,0 +1,439 @@
+"""LLG3D solver using OpenCL."""
+
+from __future__ import annotations
+
+import os
+import time
+from pathlib import Path
+from typing import Any, ClassVar
+
+import numpy as np
+import pyopencl as cl
+from pyopencl import array as clarray
+from pyopencl import clrandom
+from pyopencl import mem_flags as mf
+
+from .base import BaseSolver
+from .profiling import timeit
+
+
+class TimedKernel:
+    """
+    Wrapper for OpenCL kernels with automatic timing.
+
+    Args:
+        kernel: The underlying OpenCL kernel
+        name: Name of the kernel for logging
+        solver: Reference to the solver for profiling storage
+    """
+
+    def __init__(self, kernel: cl.Kernel, name: str, solver: "OpenCLSolver"):
+        self.kernel = kernel
+        self.name = name
+        self.solver = solver
+
+    def __call__(self, *args: Any, **kwargs: Any) -> None:
+        """
+        Execute the kernel and wait for completion with optional timing.
+
+        Args:
+            *args: Positional arguments for the kernel
+            **kwargs: Keyword arguments for the kernel
+        """
+        # Record profiling stats if enabled (use wall-clock for reliability)
+        start = time.perf_counter() if self.solver.profiling else None
+
+        # Execute kernel
+        event: cl.Event = self.kernel(*args, **kwargs)
+        event.wait()  # Always wait for kernel to complete
+
+        # Record elapsed time if profiling enabled
+        if start is not None:
+            elapsed = time.perf_counter() - start
+            stats = self.solver.profiling_stats[f"kernel_{self.name}"]
+            stats["time"] += elapsed
+            stats["calls"] += 1
+
+
+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
+
+    Raises:
+        RuntimeError: If no suitable device is found
+        ValueError: If the device selection string is invalid
+    """
+    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.dtype(np.float64 if precision == "double" else np.float32)
+
+
+class Program:
+    """Class to manage the OpenCL kernels for the LLG3D simulation."""
+
+    def __init__(self, solver: OpenCLSolver):
+        self.solver = solver  #: The OpenCLSolver instance
+        self.cl_program: cl.Program = self._get_built_program()  #: The OpenCL program
+
+    def _get_built_program(self) -> cl.Program:
+        """
+        Return the OpenCL program built from the source code.
+
+        Returns:
+            The OpenCL program object
+        """
+        grid = self.solver.grid
+        opencl_code: str = (Path(__file__).parent / "llg3d.cl").read_text()
+        build_options = (
+            "-D USE_DOUBLE_PRECISION" if self.solver.np_float == np.float64 else ""
+        )
+        build_options += f" -D NX={grid.Jx} -D NY={grid.Jy} -D NZ={grid.Jz}"
+        build_options += " -cl-fp32-correctly-rounded-divide-sqrt"
+
+        # Optional error detection in kernels, opt-in via env var
+        if self.solver.error_check_enabled:
+            build_options += " -D ENABLE_ERROR_CHECK"
+
+        # Add anisotropy type directive
+        if self.solver.elem.anisotropy == "uniaxial":
+            build_options += " -D ANISOTROPY_UNIAXIAL"
+        # else: cubic is the default
+
+        return cl.Program(self.solver.context, opencl_code).build(options=build_options)
+
+    def get_kernel(self, kernel_name: str, arg_types: list = [None]) -> TimedKernel:
+        """
+        Returns the specified kernel by name, wrapped with timing.
+
+        Args:
+            kernel_name: Name of the kernel to retrieve
+            arg_types: List of argument types for the kernel
+
+        Returns:
+            The TimedKernel wrapper around the OpenCL kernel
+        """
+        kernel: cl.Kernel = getattr(self.cl_program, kernel_name)
+        kernel.set_arg_types(arg_types)
+        return TimedKernel(kernel, kernel_name, self.solver)
+
+
+class OpenCLSolver(BaseSolver):
+    """OpenCL-based LLG3D solver."""
+
+    solver_type: ClassVar[str] = "opencl"  #: Solver type name
+
+    def __post_init__(self) -> None:
+        """Initialize OpenCL solver and check parameters."""
+        print("Initializing OpenCL solver...")
+        print("Initializing context...")
+        self.context, opencl_device = get_context_and_device(self.device)
+        self.queue = cl.CommandQueue(self.context)  #: OpenCL command queue
+        #: Numpy float type (np.dtype(np.float32) or np.dtype(np.float64))
+        self.np_float: np.dtype = get_precision(opencl_device, self.precision)
+
+        # Optional error detection in kernels, opt-in via env var
+        error_check_env = os.getenv("LLG3D_ENABLE_ERROR_CHECK", "0").lower()
+        #: Whether error checking is enabled in kernels
+        self.error_check_enabled: bool = error_check_env in {"1", "true", "on", "yes"}
+
+        super().__post_init__()
+        # Check that the grid is uniform
+        if not self.grid.uniform:
+            raise ValueError("OpenCLSolver only supports uniform grids.")
+        # Create OpenCL kernels
+        program = Program(self)
+        kernel_arg_types = [None] * 4 + [self.np_float] * 7
+        # Fused kernels combining slope computation and updates
+        self.predict_kernel = program.get_kernel(
+            "predict", kernel_arg_types
+        )  #: Prediction kernel with slope storage
+        # Correction kernel with error codes
+        correct_arg_types = kernel_arg_types + [None]
+        self.correct_and_normalize_kernel = program.get_kernel(
+            "correct_and_normalize", correct_arg_types
+        )  #: Correction kernel (slope + update_2 + normalize) with error codes
+        #: Weighted reduction kernel for m1 averaging
+        self.sum_m1_weighted_kernel = program.get_kernel(
+            "sum_m1_weighted", [None, None, None]
+        )
+
+        # Buffer for error codes (one int per work-item, thread-safe). Always present to simplify call.
+        if self.error_check_enabled:
+            err_buf_bytes = self.grid.ntot * 4
+            host_err_shape = self.grid.ntot
+        else:
+            err_buf_bytes = 4  # minimal placeholder when disabled
+            host_err_shape = 1
+
+        self.d_error_codes = cl.Buffer(
+            self.context, cl.mem_flags.READ_WRITE, size=err_buf_bytes
+        )
+        self.h_error_codes = np.zeros(host_err_shape, dtype=np.int32)
+
+        # Pre-allocate buffer for reduction (used in _xyz_average)
+        self.wgroup_size = 256
+        self.num_groups = (self.grid.ntot + self.wgroup_size - 1) // self.wgroup_size
+        self.d_partial_sums = cl.Buffer(
+            self.context,
+            cl.mem_flags.READ_WRITE,
+            size=self.num_groups * self.np_float.itemsize,
+        )
+        self.h_partial_sums = np.empty(self.num_groups, dtype=self.np_float)
+
+    def _init_rng(self) -> Any:
+        """
+        Initialize a random number generator for temperature fluctuations.
+
+        Returns:
+            An OpenCL random number generator
+        """
+        return clrandom.PhiloxGenerator(self.context, seed=self.seed)
+
+    @timeit
+    def _compute_R_random(self, d_R_random: clarray.Array):
+        """
+        Compute random number array for thermal fluctuations.
+
+        Args:
+            d_R_random: Device array to fill with random numbers
+        """
+        self.rng.fill_normal(d_R_random)  # type: ignore[attr-defined]
+        self.queue.finish()  # ensure the array is filled
+
+    @timeit
+    def _xyz_average(self, m: Any) -> float:
+        """
+        Compute the space average of m_n using a GPU reduction kernel.
+
+        Args:
+            m: Current magnetization device array
+
+        Returns:
+            The space average of m1
+        """
+        # Launch reduction kernel (buffers pre-allocated in __post_init__)
+        local_mem_size = self.wgroup_size * self.np_float.itemsize
+        self.sum_m1_weighted_kernel(
+            self.queue,
+            (self.num_groups * self.wgroup_size,),  # Global work size
+            (self.wgroup_size,),  # Local work size
+            m.data,
+            self.d_partial_sums,
+            cl.LocalMemory(local_mem_size),
+        )
+
+        # Read back partial sums and finish reduction on host
+        cl.enqueue_copy(self.queue, self.h_partial_sums, self.d_partial_sums).wait()
+        weighted_sum = float(np.sum(self.h_partial_sums))
+
+        return weighted_sum / self.grid.ncell
+
+    @timeit
+    def _check_errors(self, iteration: int) -> None:
+        """
+        Check for non-finite values detected by the kernel.
+
+        Args:
+            iteration: Current iteration number (for error message)
+
+        Raises:
+            RuntimeError: If non-finite values are detected
+        """
+        if not self.error_check_enabled:
+            return
+
+        cl.enqueue_copy(self.queue, self.h_error_codes, self.d_error_codes).wait()
+        if np.any(self.h_error_codes):
+            error_gids = np.where(self.h_error_codes)[0]
+            first_gid = error_gids[0]
+            # Convert gid to (i, j, k) coordinates
+            k = first_gid % self.grid.Jz
+            j = (first_gid // self.grid.Jz) % self.grid.Jy
+            i = first_gid // (self.grid.Jy * self.grid.Jz)
+
+            raise RuntimeError(
+                f"Non-finite value detected at iteration n={iteration}\n"
+                f"Location: gid={first_gid}, i={i}, j={j}, k={k}"
+            )
+
+    @timeit
+    def _update_x_profiles(self, m_n: Any, t: float):
+        """
+        Update the x profiles of m_n at time t.
+
+        Reads the device buffer back to host and computes the y-z averaged
+        profiles for each magnetization component, then stores them in
+        `self.x_profiles`.
+        """
+        # Initialize x_profiles on first use
+        if "x_profiles" not in self.records:
+            self.records["x_profiles"] = {
+                "t": [],
+                "m1": [],
+                "m2": [],
+                "m3": [],
+            }
+
+        # allocate host array and copy from device
+        h_m_n = m_n.map_to_host(is_blocking=True)
+
+        # save profiles
+        self.records["x_profiles"]["t"].append(t)
+        self.records["x_profiles"]["m1"].append(self._yz_average(h_m_n[0]))
+        self.records["x_profiles"]["m2"].append(self._yz_average(h_m_n[1]))
+        self.records["x_profiles"]["m3"].append(self._yz_average(h_m_n[2]))
+
+    def _simulate(self) -> float:
+        """
+        Simulates the system over N iterations.
+
+        Returns:
+            The time taken for the simulation
+        """
+        elem = self.elem
+        grid = self.grid
+        queue = self.queue
+        h_m_n = self._init_m_n()  # Create initial magnetization array on host
+        # Create device buffers
+        # Use clarray
+        d_m_n = clarray.to_device(queue, h_m_n)
+        d_R_random = clarray.Array(queue, (3, *grid.dims), self.np_float)
+        d_m_np1 = clarray.empty_like(d_m_n)
+        d_s_pre = cl.Buffer(self.context, mf.READ_WRITE, size=h_m_n.nbytes)
+
+        t = 0.0
+        self._record(d_m_n, t, 0)  # Record the initial solution
+        start_time = time.perf_counter()
+
+        for n in self._progress_bar():
+            t += self.dt
+
+            self._compute_R_random(d_R_random)
+
+            # Prediction phase: compute slope, store it, and apply first Euler update
+            self.predict_kernel(
+                queue,
+                (grid.ntot,),
+                None,
+                d_m_n.data,
+                d_m_np1.data,
+                d_R_random.data,
+                d_s_pre,
+                grid.inv_dx2,
+                elem.coeff_1,
+                elem.coeff_2,
+                elem.coeff_3,
+                elem.coeff_4,
+                elem.lambda_G,
+                self.dt,
+            )
+
+            # Correction phase: compute slope, apply midpoint update and normalize
+            self.correct_and_normalize_kernel(
+                queue,
+                (grid.ntot,),
+                None,
+                d_m_n.data,
+                d_m_np1.data,
+                d_R_random.data,
+                d_s_pre,
+                grid.inv_dx2,
+                elem.coeff_1,
+                elem.coeff_2,
+                elem.coeff_3,
+                elem.coeff_4,
+                elem.lambda_G,
+                self.dt,
+                self.d_error_codes,
+            )
+
+            # Check for errors if enabled
+            self._check_errors(n)
+
+            # Swap the buffers for the next iteration
+            d_m_n, d_m_np1 = d_m_np1, d_m_n
+
+            self._record(d_m_n, t, n)
+
+        total_time = time.perf_counter() - start_time
+
+        self._finalize()
+
+        return total_time

llg3d/solvers/profiling.py

@@ -0,0 +1,38 @@
+"""Profiling utilities for solvers."""
+
+import time
+from functools import wraps
+from typing import TYPE_CHECKING, Callable
+
+if TYPE_CHECKING:
+    from .base import BaseSolver
+
+ProfilingStats = dict[str, dict[str, float]]
+
+
+def timeit(func: Callable) -> Callable:
+    """
+    Decorator to time functions only if profiling is enabled.
+
+    Args:
+        func: The function to time.
+
+    Returns:
+        A wrapped function that records timing stats to solver.profiling_stats.
+    """
+
+    @wraps(func)
+    def timeit_wrapper(solver: "BaseSolver", *args, **kwargs):
+        # Skip timing overhead if profiling is disabled
+        if not solver.profiling:
+            return func(solver, *args, **kwargs)
+
+        start = time.perf_counter()
+        try:
+            return func(solver, *args, **kwargs)
+        finally:
+            stats = solver.profiling_stats[func.__name__]
+            stats["time"] += time.perf_counter() - start
+            stats["calls"] += 1
+
+    return timeit_wrapper

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

@@ -1,19 +1,20 @@
 Metadata-Version: 2.4
 Name: llg3d
-Version: 2.0.1
+Version: 3.0.0
 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.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: AUTHORS
 Requires-Dist: numpy
 Requires-Dist: matplotlib
 Requires-Dist: scipy
+Requires-Dist: tqdm
 Provides-Extra: mpi
 Requires-Dist: mpi4py; extra == "mpi"
 Provides-Extra: opencl
@@ -22,6 +23,8 @@ 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"
+Provides-Extra: git
+Requires-Dist: GitPython; extra == "git"
 Dynamic: license-file
 
 # LLG3D: A solver for the stochastic Landau-Lifshitz-Gilbert equation in 3D

llg3d-3.0.0.dist-info/RECORD

@@ -0,0 +1,36 @@
+llg3d/__init__.py,sha256=MssuCTRN-5px2s7MXHPvkp4Mo7HxaGQ_Hsy2HW9bg8Q,75
+llg3d/__main__.py,sha256=3D1q7AG5vU6gr-V0iuo5oNYl-Og2SvJ4YaBTdqOVVaw,115
+llg3d/element.py,sha256=GnafEd8vrDr59ixi05zyMesUCUAFhFccdZ-I1y-NAlA,7259
+llg3d/grid.py,sha256=iaT8QUZPAjt1fD39H64-c4EF4SY-qt4406pbHlO8l0s,4158
+llg3d/io.py,sha256=mJdPv65r8u5ODrxJB39u_dVnZi-PGBwC_mos4SXkrEc,12366
+llg3d/main.py,sha256=uUgT4PGBr255OZqYXS2x-gvaFKVf_RKo0JliDQlQcT4,1729
+llg3d/parameters.py,sha256=Ox8-HL2AKw-Ogw-xkCFXiEm6ponYpf6I8C6BHWpgWf8,6064
+llg3d/py.typed,sha256=iMysHiY_j0PVKAdpVSJynuRB8D7puNCzeIDwHpWZ4xE,58
+llg3d/benchmarks/__init__.py,sha256=zw3npEgBEO_1rr7MRq9YPO3nOGRKwIhe6kFnpxI7ZJQ,26
+llg3d/benchmarks/compare_commits.py,sha256=9fY6RMXlEfDKKvz4D5NO3jtuf4JB26zIF5L9mMKrsa4,9904
+llg3d/benchmarks/efficiency.py,sha256=PrDE3KxVAtWVRotcod0CiegQ_2PZa6qI9ByRRadHq7s,15461
+llg3d/benchmarks/utils.py,sha256=XSECFgoquWu7LcgCsPYx5YG3CtRlPHXUewMPGthP5wY,708
+llg3d/post/__init__.py,sha256=S4UiquBtCm_iHQ0vaU0E4T69TZfhiYq7i26AOz7Vs40,29
+llg3d/post/extract.py,sha256=73FWYE1-HYLzU1tjjAOvhHrZFe6SVLMLa78jJxgz0pM,3301
+llg3d/post/info.py,sha256=7dAWsfO9IkCkiWgBVxVm36Y3rK78PBzL_rb6LuajANU,5232
+llg3d/post/m1_vs_T.py,sha256=t7HAPlcH_v1iLmbFplhZRmHRVoQLSA9t287ITZ47Fhw,2715
+llg3d/post/m1_vs_time.py,sha256=6ynl4OK4lPPDLnhY_k-qSBovldO4Wa95-Z8LfiohIr0,1507
+llg3d/post/process.py,sha256=usvRLnvHDqaV3EKiP6il3se7ha_G_PHVS0LP_8c87KM,3845
+llg3d/post/utils.py,sha256=oZBYbPhJJYzHsRrnENyxmZotylLlYVOOFePr3zc-Fo0,1114
+llg3d/post/x_profiles.py,sha256=pexMn2hUPPNvdTOsc5ilXm6JN3zlV5xu4xU2LXXOp4I,3993
+llg3d/solvers/__init__.py,sha256=aNn3LjP1nDZj7PqJ_KGxY3hj-BDa-RrOcE00DnjLgDk,3436
+llg3d/solvers/base.py,sha256=cucNTqBSbl9MgArGbmlYORKPWrsJFuZRQCZYgACuj8U,11879
+llg3d/solvers/math_utils.py,sha256=F1uwrfNCg0tCLlkXKoMFHcOaSFsIrrvzO768IpfvZx4,1014
+llg3d/solvers/mpi.py,sha256=8zBWU1k1Ns4nM2-Dz-gmoGeYnCWuIh8jN-b6USWBV2Q,11530
+llg3d/solvers/numpy.py,sha256=pjNQJzTyqOwe5XjGKwK_Vjp55ZfmAtC1WOPaSQ6OpdM,3616
+llg3d/solvers/opencl.py,sha256=i7y50uQ8bV5EcV3PDcEvJoPo3Qt6pvAq_awy7mgHDGk,15190
+llg3d/solvers/profiling.py,sha256=B5PgxtE3xCcv7Qi51VVYAToBzR5HwjS4j9LPKg1ajyc,991
+llg3d/solvers/experimental/__init__.py,sha256=F4smHsF9_hR5KXIsh6Xd46sSzJE6N7yFuH0vmj0ygGA,147
+llg3d/solvers/experimental/jax.py,sha256=FrcdQYMVroPbYBbqvZMPIH9gv8n-PkT2vdUO5u5aYHg,11249
+llg3d-3.0.0.dist-info/licenses/AUTHORS,sha256=vhJ88HikYvOrGiB_l1xH2X6hyn9ZJafx6mpoMYNhh1I,297
+llg3d-3.0.0.dist-info/licenses/LICENSE,sha256=aFxTGAyyve8nM9T2jWTarJzQhdSSC3MbbN1heNAev9c,1062
+llg3d-3.0.0.dist-info/METADATA,sha256=dD3kJMCX07AFNIHAud5mEILEzBzdNZLfvaK065sH9FM,2030
+llg3d-3.0.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+llg3d-3.0.0.dist-info/entry_points.txt,sha256=28japzG2aS2e_J42LXw7SAK4ufbG22LGHcI2D6zVz9Y,374
+llg3d-3.0.0.dist-info/top_level.txt,sha256=cBZ0roaXt3CAXqYojuO84lGPCtWuLlXxLGLYRKmHZy0,6
+llg3d-3.0.0.dist-info/RECORD,,

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

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

llg3d-3.0.0.dist-info/entry_points.txt

@@ -0,0 +1,9 @@
+[console_scripts]
+llg3d = llg3d.main:main
+llg3d.bench.compare_commits = llg3d.benchmarks.compare_commits:main
+llg3d.bench.efficiency = llg3d.benchmarks.efficiency:main
+llg3d.extract = llg3d.post.extract:main
+llg3d.info = llg3d.post.info:main
+llg3d.m1_vs_T = llg3d.post.m1_vs_T:main
+llg3d.m1_vs_time = llg3d.post.m1_vs_time:main
+llg3d.x_profiles = llg3d.post.x_profiles:main

llg3d/output.py

@@ -1,107 +0,0 @@
-"""Utility functions for LLG3D."""
-
-import json
-import sys
-from typing import Iterable, TextIO
-
-from .solver import rank
-from .grid import Grid
-
-
-def progress_bar(
-    it: Iterable, prefix: str = "", size: int = 60, out: TextIO = sys.stdout
-):
-    """
-    Displays a progress bar.
-
-    (Source: https://stackoverflow.com/a/34482761/16593179)
-
-    Args:
-        it: Iterable object to iterate over
-        prefix: Prefix string for the progress bar
-        size: Size of the progress bar (number of characters)
-        out: Output stream (default is sys.stdout)
-    """
-    count = len(it)
-
-    def show(j):
-        x = int(size * j / count)
-        if rank == 0:
-            print(
-                f"{prefix}[{'█' * x}{('.' * (size - x))}] {j}/{count}",
-                end="\r",
-                file=out,
-                flush=True,
-            )
-
-    show(0)
-    for i, item in enumerate(it):
-        yield item
-        # To avoid slowing down the computation, we do not display at every iteration
-        if i % 5 == 0:
-            show(i + 1)
-    show(i + 1)
-    if rank == 0:
-        print("\n", flush=True, file=out)
-
-
-def write_json(json_file: str, run: dict):
-    """
-    Writes the run dictionary to a JSON file.
-
-    Args:
-        json_file: Name of the JSON file
-        run: Dictionary containing the run information
-    """
-    with open(json_file, "w") as f:
-        json.dump(run, f, indent=4)
-
-
-def get_output_files(g: Grid, T: float, n_mean: int, n_profile: int) -> tuple:
-    """
-    Open files and list them.
-
-    Args:
-        g: Grid object
-        T: temperature
-        n_mean: Number of iterations for integral output
-        n_profile: Number of iterations for profile output
-
-    Returns:
-        - a file handler for storing m space integral over time
-        - a file handler for storing x-profiles of m_i
-        - a list of output filenames
-    """
-    f_mean = None
-    f_profiles = None
-    output_filenames = []
-    if n_mean != 0:
-        output_filenames.append(g.get_filename(T, extension="txt"))
-    if n_profile != 0:
-        output_filenames.extend(
-            [g.get_filename(T, name=f"m{i + 1}", extension="npy") for i in range(3)]
-        )
-    if rank == 0:
-        if n_mean != 0:
-            f_mean = open(output_filenames[0], "w")  # integral of m1
-        if n_profile != 0:
-            f_profiles = [
-                open(output_filename, "wb") for output_filename in output_filenames[1:]
-            ]  # x profiles of m_i
-
-    return f_mean, f_profiles, output_filenames
-
-
-def close_output_files(f_mean: TextIO, f_profiles: list[TextIO] = None):
-    """
-    Close all output files.
-
-    Args:
-        f_mean: file handler for storing m space integral over time
-        f_profiles: file handlers for storing x-profiles of m_i
-    """
-    if f_mean is not None:
-        f_mean.close()
-    if f_profiles is not None:
-        for f_profile in f_profiles:
-            f_profile.close()