feectools 0.3.1__tar.gz → 0.5.0__tar.gz

{feectools-0.3.1/feectools.egg-info → feectools-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.3.1
+Version: 0.5.0
 Summary: Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies.
 Author-email: Psydac development team <psydac@googlegroups.com>
 Maintainer-email: Stefan Possanner <stefan.possanner@ipp.mpg.de>, Max Lindqvist <max.lindqvist@ipp.mpg.de>, Yaman Güçlü <yaman.guclu@gmail.com>, Martin Campos Pinto <martin.campos-pinto@ipp.mpg.de>, Ahmed Ratnani <ratnaniahmed@gmail.com>
@@ -42,7 +42,8 @@ Requires-Dist: matplotlib
 Requires-Dist: pyyaml>=5.1
 Requires-Dist: packaging
 Requires-Dist: pyevtk
-Requires-Dist: cunumpy<0.6,>=0.5.0
+Requires-Dist: maybempi>=0.1.2
+Requires-Dist: cunumpy>=0.6.1
 Requires-Dist: pyccel>=2.1.0
 Requires-Dist: h5py
 Requires-Dist: tblib

feectools-0.5.0/feectools/__init__.py

@@ -0,0 +1,7 @@
+# Bind this rank to its own GPU (by its rank within the node) and create the CUDA context before MPI starts, as
+# CUDA-aware MPI requires. Every feectools module imports this package first, and MPI starts only when a module
+# imports `from maybempi import MPI` (maybempi uses mpi4py under an MPI launcher and a serial stand-in otherwise;
+# MAYBEMPI=0/1 overrides the detection). A no-op on the NumPy backend.
+from cunumpy.cuda import bind_local_device
+
+bind_local_device()

{feectools-0.3.1 → feectools-0.5.0}/feectools/ddm/blocking_data_exchanger.py

@@ -3,8 +3,7 @@
 import cunumpy as xp
 from cunumpy.mpi import synchronize_for_mpi
 import numpy as np
-from feectools.ddm.mpi import mpi as MPI
-
+from maybempi import MPI
 from .cart import CartDecomposition, find_mpi_type
 from .basic import CartDataExchanger
 

{feectools-0.3.1 → feectools-0.5.0}/feectools/ddm/cart.py

@@ -10,16 +10,8 @@ from itertools import product
 
 from cunumpy.xp import array_backend, to_numpy
 
-# Bind this rank to its own GPU (by its rank within the node) and create the
-# CUDA context before MPI is initialized, as CUDA-aware MPI requires. Must stay
-# above the feectools.ddm.mpi import, which initializes MPI as a side effect.
-# A no-op on the NumPy backend.
-from cunumpy.cuda import bind_local_device
-
-bind_local_device()
-
-from feectools.ddm.mpi import mpi as MPI
-from feectools.ddm.mpi import MockMPI
+# The GPU of this rank is bound in feectools/__init__.py, before MPI starts.
+from maybempi import MPI, is_serial
 from feectools.ddm.partition import compute_dims, partition_procs_per_patch
 
 
@@ -46,7 +38,7 @@ def find_mpi_type( dtype ):
         MPI datatype to be used for communication.
 
     """
-    if not isinstance(MPI, MockMPI):
+    if not is_serial(MPI):
         if isinstance( dtype, MPI.Datatype ):
             mpi_type = dtype
         else:
@@ -92,7 +84,7 @@ class MultiPatchDomainDecomposition:
     def __init__(self, ncells, periods, comm=None, num_threads=None):
 
         assert len( ncells ) == len( periods )
-        if not isinstance(MPI, MockMPI) and comm is not None:
+        if not is_serial(MPI) and comm is not None:
             assert isinstance( comm, MPI.Comm )
         num_threads = num_threads if num_threads else int(os.environ.get('OMP_NUM_THREADS', 1))
 
@@ -236,7 +228,7 @@ class DomainDecomposition:
         assert len( ncells ) == len( periods )
         assert all( n >=1 for n in ncells )
         assert all( isinstance( period, bool ) for period in periods )
-        if isinstance(MPI, MockMPI): 
+        if is_serial(MPI):
             comm = None
         else:
             if comm is not None:

{feectools-0.3.1 → feectools-0.5.0}/feectools/ddm/interface_data_exchanger.py

@@ -1,8 +1,7 @@
 # coding: utf-8
 
 from cunumpy.mpi import synchronize_for_mpi
-from feectools.ddm.mpi import mpi as MPI
-
+from maybempi import MPI
 from .cart import InterfaceCartDecomposition, find_mpi_type
 
 __all__ = ('InterfaceCartDataExchanger',)

{feectools-0.3.1 → feectools-0.5.0}/feectools/ddm/nonblocking_data_exchanger.py

@@ -5,7 +5,7 @@ from cunumpy.mpi import synchronize_for_mpi
 import numpy as np
 from itertools import product
 
-from feectools.ddm.mpi import mpi as MPI
+from maybempi import MPI
 from .cart import CartDecomposition, find_mpi_type
 from .basic import CartDataExchanger
 

{feectools-0.3.1 → feectools-0.5.0}/feectools/ddm/tests/test_cart_1d.py

@@ -14,7 +14,7 @@ from feectools.ddm.nonblocking_data_exchanger import NonBlockingCartDataExchange
 def run_cart_1d( data_exchanger_type, verbose=False ):
 
     import cunumpy as xp
-    from feectools.ddm.mpi import mpi as MPI
+    from maybempi import MPI
     from feectools.ddm.cart import DomainDecomposition, CartDecomposition
 
     #---------------------------------------------------------------------------

{feectools-0.3.1 → feectools-0.5.0}/feectools/ddm/tests/test_cart_2d.py

@@ -9,7 +9,7 @@ from feectools.ddm.nonblocking_data_exchanger import NonBlockingCartDataExchange
 def run_cart_2d( data_exchanger_type, verbose=False , nprocs=None, reverse_axis=None):
 
     import cunumpy as xp
-    from feectools.ddm.mpi import mpi as MPI  
+    from maybempi import MPI
     from feectools.ddm.cart import DomainDecomposition, CartDecomposition
 
     #---------------------------------------------------------------------------

{feectools-0.3.1 → feectools-0.5.0}/feectools/ddm/tests/test_cart_3d.py

@@ -9,7 +9,7 @@ from feectools.ddm.nonblocking_data_exchanger import NonBlockingCartDataExchange
 def run_cart_3d( data_exchanger_type, verbose=False ):
 
     import cunumpy as xp
-    from feectools.ddm.mpi import mpi as MPI 
+    from maybempi import MPI
     from feectools.ddm.cart import DomainDecomposition, CartDecomposition
 
     #---------------------------------------------------------------------------

{feectools-0.3.1 → feectools-0.5.0}/feectools/ddm/tests/test_coarsen.py

@@ -2,9 +2,7 @@ import numpy as np
 import pytest
 
 from feectools.ddm.cart import DomainDecomposition
-from feectools.ddm.mpi import mpi as MPI
-
-
+from maybempi import MPI
 def _comm(parallel):
     return MPI.COMM_WORLD if parallel else None
 

{feectools-0.3.1 → feectools-0.5.0}/feectools/ddm/tests/test_multicart_2d.py

@@ -41,7 +41,7 @@ def get_plus_starts_ends(minus_starts, minus_ends, minus_npts, plus_npts, minus_
 def run_carts_2d():
     import cunumpy as xp
     
-    from feectools.ddm.mpi import mpi as MPI
+    from maybempi import MPI
     from feectools.ddm.cart import MultiPatchDomainDecomposition, CartDecomposition, create_interfaces_cart
     from feectools.ddm.blocking_data_exchanger  import BlockingCartDataExchanger
     from feectools.ddm.interface_data_exchanger import InterfaceCartDataExchanger

{feectools-0.3.1 → feectools-0.5.0}/feectools/fem/tensor.py

@@ -5,8 +5,7 @@ We assume here that a tensor space is the product of fem spaces whom basis are
 of compact support
 
 """
-from feectools.ddm.mpi import mpi as MPI
-    
+from maybempi import MPI
 import cunumpy as xp
 import numpy as np
 import itertools

{feectools-0.3.1 → feectools-0.5.0}/feectools/fem/tests/test_spline_interpolation.py

@@ -5,8 +5,7 @@
 #---------------------------------------------------------------------------#
 import time
 
-from feectools.ddm.mpi import mpi as MPI
-    
+from maybempi import MPI
 import cunumpy as xp
 import pytest
 

{feectools-0.3.1 → feectools-0.5.0}/feectools/fem/tests/test_splines_par.py

@@ -7,8 +7,7 @@ from feectools.fem.splines import SplineSpace
 from feectools.fem.tensor  import TensorFemSpace
 from feectools.fem.vector  import VectorFemSpace
 from feectools.ddm.cart    import DomainDecomposition
-from feectools.ddm.mpi import mpi as MPI
-
+from maybempi import MPI
 import cunumpy as xp
 
 def test_2d_1():

{feectools-0.3.1 → feectools-0.5.0}/feectools/fem/tests/test_tensor.py

@@ -19,7 +19,7 @@ import numpy as np
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 from PIL import Image
-from mpi4py import MPI
+from maybempi import MPI
 
 from sympde.topology.domain import Square
 from sympde.topology.space import ScalarFunctionSpace

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/basic.py

@@ -385,7 +385,7 @@ class LinearOperator(ABC):
 
         if not is_sparse:
             if comm is not None:
-                from feectools.ddm.mpi import mpi as MPI
+                from maybempi import MPI
                 comm.Allreduce(MPI.IN_PLACE, out, op=MPI.SUM)
             return out
 

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/block.py

@@ -1076,7 +1076,7 @@ class BlockLinearOperator(LinearOperator):
         if not self.codomain.parallel:
             return blocks, blocks_T
 
-        from feectools.ddm.mpi import mpi as MPI
+        from maybempi import MPI
         from feectools.linalg.stencil import StencilInterfaceMatrix
 
         if not isinstance(self.codomain, BlockVectorSpace):

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/solvers.py

@@ -1649,7 +1649,7 @@ class LSMR(InverseLinearOperator):
 
             test1 = normr / normb
             if (normA * normr) != 0:test2 = normar / (normA * normr)
-            else:test2 = xp.infty
+            else:test2 = xp.inf
             test3 = 1 / condA
             t1    = test1 / (1 + normA * normx / normb)
             rtol  = btol + atol * normA * normx / normb

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/stencil.py

@@ -15,7 +15,7 @@ from cunumpy.mpi import synchronize_for_mpi
 from cunumpy.xp import array_backend
 from scipy.sparse import coo_matrix, diags as sp_diags
 
-from feectools.ddm.mpi import mpi as MPI
+from maybempi import MPI
 from feectools.linalg.basic  import VectorSpace, Vector, LinearOperator
 from feectools.linalg.memory import stencil_matrix_memory
 from feectools.ddm.cart      import find_mpi_type, CartDecomposition, InterfaceCartDecomposition

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/tests/test_block.py

@@ -5,7 +5,7 @@ import cunumpy as xp
 from scipy.sparse import csr_matrix
 from random import random, seed
 
-from feectools.ddm.mpi import mpi as MPI
+from maybempi import MPI
 from feectools.linalg.direct_solvers import SparseSolver
 from feectools.linalg.stencil        import StencilVectorSpace, StencilVector, StencilMatrix
 from feectools.linalg.block          import BlockVectorSpace, BlockVector

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/tests/test_fft.py

@@ -7,7 +7,7 @@ import pytest
 import scipy.fft as scifft
 import cunumpy as xp
 
-from feectools.ddm.mpi import mpi as MPI
+from maybempi import MPI
 from feectools.linalg.fft import *
 from feectools.ddm.cart               import DomainDecomposition, CartDecomposition
 from feectools.linalg.stencil import StencilVector

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/tests/test_mpi_device.py

@@ -23,7 +23,7 @@ import numpy as np
 import pytest
 import cunumpy as xp
 
-from feectools.ddm.mpi import mpi as MPI
+from maybempi import MPI
 from feectools.ddm.cart import DomainDecomposition, CartDecomposition
 from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix
 

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/tests/test_stencil_interface_matrix.py

@@ -231,7 +231,7 @@ def test_stencil_interface_matrix_3d_serial_init(dtype, n1, n2, n3, p1, p2, p3,
 @pytest.mark.mpi
 def test_stencil_interface_matrix_2d_parallel_dot(n1, n2, p1, p2, expected):
 
-    from feectools.ddm.mpi import mpi as MPI
+    from maybempi import MPI
     from feectools.ddm.cart     import MultiPatchDomainDecomposition, CartDecomposition, create_interfaces_cart
     from feectools.linalg.block import BlockVectorSpace, BlockVector, BlockLinearOperator
 

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/tests/test_stencil_vector.py

@@ -6,7 +6,7 @@
 import pytest
 import cunumpy as xp
 
-from feectools.ddm.mpi import mpi as MPI
+from maybempi import MPI
 from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix
 from feectools.linalg.utilities import array_to_psydac, petsc_to_psydac
 from feectools.ddm.cart import DomainDecomposition, CartDecomposition

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/tests/test_stencil_vector_space.py

@@ -6,7 +6,7 @@
 import pytest
 import cunumpy as xp
 
-from feectools.ddm.mpi import mpi as MPI
+from maybempi import MPI
 from feectools.linalg.stencil import StencilVectorSpace, StencilVector
 from feectools.ddm.cart import DomainDecomposition, CartDecomposition, find_mpi_type
 

{feectools-0.3.1 → feectools-0.5.0}/feectools/linalg/tests/test_toarray.py

@@ -6,7 +6,7 @@
 import pytest
 import cunumpy as xp
 
-from feectools.ddm.mpi import mpi as MPI
+from maybempi import MPI
 from feectools.ddm.cart import DomainDecomposition, CartDecomposition
 from feectools.linalg.basic import LinearOperator, MatrixFreeLinearOperator, IdentityOperator
 from feectools.linalg.stencil import StencilVectorSpace, StencilMatrix

{feectools-0.3.1 → feectools-0.5.0/feectools.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.3.1
+Version: 0.5.0
 Summary: Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies.
 Author-email: Psydac development team <psydac@googlegroups.com>
 Maintainer-email: Stefan Possanner <stefan.possanner@ipp.mpg.de>, Max Lindqvist <max.lindqvist@ipp.mpg.de>, Yaman Güçlü <yaman.guclu@gmail.com>, Martin Campos Pinto <martin.campos-pinto@ipp.mpg.de>, Ahmed Ratnani <ratnaniahmed@gmail.com>
@@ -42,7 +42,8 @@ Requires-Dist: matplotlib
 Requires-Dist: pyyaml>=5.1
 Requires-Dist: packaging
 Requires-Dist: pyevtk
-Requires-Dist: cunumpy<0.6,>=0.5.0
+Requires-Dist: maybempi>=0.1.2
+Requires-Dist: cunumpy>=0.6.1
 Requires-Dist: pyccel>=2.1.0
 Requires-Dist: h5py
 Requires-Dist: tblib

{feectools-0.3.1 → feectools-0.5.0}/feectools.egg-info/SOURCES.txt

@@ -32,7 +32,6 @@ feectools/ddm/basic.py
 feectools/ddm/blocking_data_exchanger.py
 feectools/ddm/cart.py
 feectools/ddm/interface_data_exchanger.py
-feectools/ddm/mpi.py
 feectools/ddm/nonblocking_data_exchanger.py
 feectools/ddm/partition.py
 feectools/ddm/petsc.py

{feectools-0.3.1 → feectools-0.5.0}/feectools.egg-info/requires.txt

@@ -5,7 +5,8 @@ matplotlib
 pyyaml>=5.1
 packaging
 pyevtk
-cunumpy<0.6,>=0.5.0
+maybempi>=0.1.2
+cunumpy>=0.6.1
 pyccel>=2.1.0
 h5py
 tblib

{feectools-0.3.1 → feectools-0.5.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name            = "feectools"
-version         = "0.3.1"
+version         = "0.5.0"
 description     = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies."
 readme          = "README.md"
 requires-python = ">= 3.10"
@@ -30,7 +30,8 @@ dependencies = [
     'pyyaml >= 5.1',
     'packaging',
     'pyevtk',
-    'cunumpy >= 0.5.0, < 0.6',
+    'maybempi >= 0.1.2',
+    'cunumpy >= 0.6.1',
 
     # Our packages from PyPi
     'pyccel >= 2.1.0',

feectools-0.3.1/feectools/__init__.py

@@ -1,4 +0,0 @@
-# MPI switch, read once when feectools.ddm.mpi is first imported:
-#   None  -> use mpi4py if it is available (default)
-#   False -> do not import mpi4py, use MockMPI (serial runs; saves ~1 s of import time)
-use_mpi = None

feectools-0.3.1/feectools/ddm/mpi.py

@@ -1,223 +0,0 @@
-"""Detection of whether the process was launched by an MPI launcher.
-
-Importing ``mpi4py.MPI`` calls ``MPI_Init``, and any collective (``bcast``,
-``Barrier``, ...) issued afterwards costs something even on a single process.
-A plain ``python script.py`` run should therefore never touch MPI at all, even
-when mpi4py happens to be installed. This module answers the only question
-that decides it: was this process started by ``mpirun``/``mpiexec``/``srun``
-(or an equivalent launcher)?
-
-The answer is read from the environment the launcher sets up, so it is
-available before mpi4py is imported.
-"""
-
-import os
-import sys
-
-from dataclasses import dataclass
-from time import time
-from typing import TYPE_CHECKING
-
-
-# Per-process variables exported by the process managers behind the common
-# launchers. Each is set only for processes started *by* the launcher, so the
-# presence of any one of them means "this rank belongs to an MPI job".
-# SLURM_PROCID is deliberately absent: it is also set for the script of a
-# plain `sbatch` job, which is not an MPI launch. `srun` is covered by the
-# PMI/PMIX variables its MPI plugin exports.
-_LAUNCHER_ENV_VARS = (
-    "OMPI_COMM_WORLD_RANK",  # Open MPI (and derivatives: Spectrum, ...)
-    "PMI_RANK",  # MPICH, Intel MPI, MS-MPI, Cray, srun (pmi2)
-    "PMIX_RANK",  # PMIx, used by srun --mpi=pmix and Open MPI 5
-    "MV2_COMM_WORLD_RANK",  # MVAPICH2
-    "MPI_LOCALRANKID",  # Hydra (mpiexec.hydra)
-    "ALPS_APP_PE",  # Cray ALPS aprun
-    "PALS_RANKID",  # Cray PALS palsrun
-)
-
-# Escape hatch: force the decision either way without touching code, e.g. for
-# a launcher whose variables are not listed above.
-_OVERRIDE_ENV_VAR = "STRUPHY_MPI"
-
-_TRUE_VALUES = ("1", "true", "yes", "on")
-_FALSE_VALUES = ("0", "false", "no", "off")
-
-
-def _override() -> bool | None:
-    """Value of ``STRUPHY_MPI``, or None if unset/unrecognized."""
-    value = os.environ.get(_OVERRIDE_ENV_VAR)
-    if value is None:
-        return None
-    value = value.strip().lower()
-    if value in _TRUE_VALUES:
-        return True
-    if value in _FALSE_VALUES:
-        return False
-    return None
-
-
-def launched_under_mpi() -> bool:
-    """Whether this process was started by an MPI launcher.
-
-    Returns
-    -------
-    bool
-        True if a launcher's per-rank environment variable is present, or if
-        the application itself already initialized MPI (in which case using
-        the communicator is free). ``STRUPHY_MPI=0``/``1`` overrides
-        the detection.
-    """
-    override = _override()
-    if override is not None:
-        return override
-
-    if any(var in os.environ for var in _LAUNCHER_ENV_VARS):
-        return True
-
-    # The application may have initialized MPI itself (embedded interpreter,
-    # or an explicit `from mpi4py import MPI`). Only inspect mpi4py if it is
-    # already imported: importing it here is exactly what must be avoided.
-    mpi_module = sys.modules.get("mpi4py.MPI")
-    if mpi_module is not None:
-        try:
-            return bool(mpi_module.Is_initialized())
-        except AttributeError:
-            return False
-
-    return False
-
-
-# Might not be needed
-class MPICommWrapper:
-    def __init__(self, use_mpi=True):
-        self.use_mpi = use_mpi
-        if use_mpi:
-            from mpi4py import MPI
-
-            self.comm = MPI.COMM_WORLD
-        else:
-            self.comm = MockComm()
-
-    def __getattr__(self, name):
-        return getattr(self.comm, name)
-
-
-class MockComm:
-    def __getattr__(self, name):
-        # Return a function that does nothing and returns None
-        def dummy(*args, **kwargs):
-            return None
-
-        return dummy
-
-    # Override some functions
-    def Get_rank(self):
-        return 0
-
-    def Get_size(self):
-        return 1
-
-    def Barrier(self):
-        return
-
-
-class MPIwrapper:
-    def __init__(
-        self,
-        use_mpi: bool = False,
-        verbose: bool = False,
-    ):
-        self.use_mpi = use_mpi
-        if use_mpi:
-            from mpi4py import MPI
-
-            self._MPI = MPI
-            if verbose:
-                print("MPI is enabled")
-        else:
-            self._MPI = MockMPI()
-            if verbose:
-                print("MPI is NOT enabled")
-
-    @property
-    def MPI(self):
-        return self._MPI
-
-
-class MockMPI:
-    def __getattr__(self, name):
-        # Return a function that does nothing and returns None
-        def dummy(*args, **kwargs):
-            return None
-
-        return dummy
-
-    # Override some functions
-    @property
-    def COMM_WORLD(self):
-        return MockComm()
-
-    # def comm_Get_rank(self):
-    #     return 0
-
-    # def comm_Get_size(self):
-    #     return 1
-
-
-def _mpi_disabled():
-    """True if the user or the host application opted out of MPI.
-
-    Importing mpi4py initializes MPI, which takes close to a second. Serial runs that
-    never use MPI can skip it entirely, in two ways:
-
-    * export ``FEECTOOLS_MPI=0`` before starting Python, or
-    * set ``feectools.use_mpi = False`` before this module is first imported
-      (in-process, so it is not inherited by subprocesses such as ``mpirun``).
-
-    The MockMPI wrapper below is then used, exactly as if mpi4py were not installed.
-    """
-    import os
-
-    import feectools
-
-    if getattr(feectools, 'use_mpi', None) is False:
-        return True
-    return os.environ.get('FEECTOOLS_MPI', '').strip().lower() in ('0', 'false', 'no', 'off')
-
-
-if launched_under_mpi():
-    try:
-        # MPI with the CuPy backend needs a CUDA-aware MPI library: device
-        # buffers are passed to MPI directly (after synchronize_for_mpi, see
-        # the data exchangers). A non-CUDA-aware MPI segfaults on them.
-        if _mpi_disabled():
-            raise ImportError("MPI disabled (feectools.use_mpi = False or FEECTOOLS_MPI=0)")
-
-        from mpi4py import MPI
-
-        _comm = MPI.COMM_WORLD
-        # rank = _comm.Get_rank()
-        # size = _comm.Get_size()
-        mpi_enabled = True
-    except ImportError:
-        # mpi4py not installed
-        mpi_enabled = False
-    except Exception:
-        # mpi4py installed but not running under mpirun
-        mpi_enabled = False
-else:
-    mpi_enabled = False
-
-# TODO: add environment variable for mpi use
-mpi_wrapper = MPIwrapper(
-    use_mpi=mpi_enabled,
-    verbose=False,
-)
-
-# TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime.
-if TYPE_CHECKING:
-    from mpi4py import MPI
-
-    mpi = MPI
-else:
-    mpi = mpi_wrapper.MPI