feectools 0.1.9__tar.gz → 0.1.11__tar.gz

{feectools-0.1.9/feectools.egg-info → feectools-0.1.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.1.9
+Version: 0.1.11
 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>

feectools-0.1.11/feectools/__init__.py

@@ -0,0 +1,4 @@
+# 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.1.9 → feectools-0.1.11}/feectools/api/settings.py

@@ -43,13 +43,16 @@ PSYDAC_BACKEND_NVPYCCEL = {'name': 'pyccel',
                        'openmp'  : False}
 # ...
 
-# Get gfortran version
-gfortran_version_output = subprocess.check_output(['gfortran', '--version']).decode('utf-8') # nosec B603, B607
-gfortran_version_string = re.search(r"(\d+\.\d+\.\d+)", gfortran_version_output).group()
-gfortran_version = Version(gfortran_version_string)
+def _get_gfortran_version():
+    """Version of the installed gfortran. Only queried where it is needed (see below), since it spawns a subprocess."""
+    gfortran_version_output = subprocess.check_output(['gfortran', '--version']).decode('utf-8') # nosec B603, B607
+    gfortran_version_string = re.search(r"(\d+\.\d+\.\d+)", gfortran_version_output).group()
+    return Version(gfortran_version_string)
 
 # Platform-dependent flags
-if platform.system() == "Darwin" and platform.machine() == 'arm64' and gfortran_version >= Version("14"):
+# (the gfortran version is only relevant on Apple silicon; evaluating it lazily avoids running
+#  a subprocess on every import and lets feectools be imported without gfortran on other platforms)
+if platform.system() == "Darwin" and platform.machine() == 'arm64' and _get_gfortran_version() >= Version("14"):
 
     # Apple silicon requires architecture-specific flags (see https://github.com/pyccel/psydac/pull/411)
     # which are only available on GCC version >= 14

feectools-0.1.11/feectools/ddm/mpi.py

@@ -0,0 +1,225 @@
+"""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:
+        # Disable MPI when using CuPy due to known segfault issues with OpenMPI + CUDA
+        import os
+        if os.environ.get('ARRAY_BACKEND') == 'cupy':
+            raise ImportError("MPI disabled when using CuPy backend")
+
+        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

{feectools-0.1.9 → feectools-0.1.11}/feectools/ddm/partition.py

@@ -2,10 +2,49 @@ import cunumpy as xp
 import numpy as np
 import numpy.ma as ma
 
-from sympy.ntheory import factorint
+__all__ = ('compute_dims', 'partition_procs_per_patch')
+
+#==============================================================================
+def factorint(n, multiple=False):
+    """
+    Prime factorization of an integer by trial division.
 
+    Drop-in replacement for the subset of ``sympy.ntheory.factorint`` used here.
+    Importing sympy takes ~1 s and is not needed for the small integers
+    (process counts, number of grid points) that are factorized in this module.
 
-__all__ = ('compute_dims', 'partition_procs_per_patch')
+    Parameters
+    ----------
+    n : int
+        Integer to factorize.
+
+    multiple : bool
+        If False (default), return a dict {prime: multiplicity}.
+        If True, return the list of primes in ascending order, repeated
+        according to their multiplicity.
+    """
+    n = int(n)
+    factors = {}
+
+    # same conventions as sympy for non-positive input
+    if n == 0:
+        factors[0] = 1
+    else:
+        if n < 0:
+            factors[-1] = 1
+            n = -n
+        p = 2
+        while p * p <= n:
+            while n % p == 0:
+                factors[p] = factors.get(p, 0) + 1
+                n //= p
+            p += 1 if p == 2 else 2
+        if n > 1:
+            factors[n] = factors.get(n, 0) + 1
+
+    if multiple:
+        return [p for p in sorted(factors) for _ in range(factors[p])]
+    return factors
 
 #==============================================================================
 def partition_procs_per_patch(npts, size):

{feectools-0.1.9 → feectools-0.1.11}/feectools/linalg/basic.py

@@ -756,7 +756,7 @@ class ScaledLinearOperator(LinearOperator):
         self._scalar = c
 
     def toarray(self):
-        return self._scalar * self._operator.toarray() 
+        return self._scalar * self._operator.toarray
 
     def tosparse(self):
         return self._scalar * self._operator.tosparse().tocsr()

{feectools-0.1.9 → feectools-0.1.11}/feectools/linalg/direct_solvers.py

@@ -6,9 +6,7 @@
 from abc                 import abstractmethod
 import cunumpy as xp
 from cunumpy.xp import array_backend
-from scipy.linalg.lapack import dgbtrf, dgbtrs, sgbtrf, sgbtrs, cgbtrf, cgbtrs, zgbtrf, zgbtrs
 from scipy.sparse        import spmatrix, dia_matrix
-from scipy.sparse.linalg import splu
 
 from feectools.linalg.basic    import LinearSolver
 
@@ -53,6 +51,9 @@ class BandedSolver(LinearSolver):
         self._l    = l
         self._transposed = transposed
 
+        # imported here: scipy.linalg costs ~0.3 s to import and is only needed once a solver is built
+        from scipy.linalg.lapack import dgbtrf, dgbtrs, sgbtrf, sgbtrs, cgbtrf, cgbtrs, zgbtrf, zgbtrs
+
         # ... LU factorization
         if bmat.dtype == xp.float32:
             self._factor_function = sgbtrf
@@ -186,6 +187,8 @@ class SparseSolver (LinearSolver):
 
         assert isinstance(spmat, spmatrix)
 
+        from scipy.sparse.linalg import splu  # deferred, see BandedSolver
+
         self._space = xp.ndarray
         self._splu  = splu(spmat.tocsc())
         self._transposed = transposed

{feectools-0.1.9 → feectools-0.1.11}/feectools/linalg/solvers.py

@@ -4,12 +4,14 @@ This module provides iterative solvers and preconditioners.
 
 """
 import cunumpy as xp
-from math import sqrt
+from math import sqrt, inf
 
 from feectools.utilities.utils  import is_real
 from feectools.linalg.utilities import _sym_ortho
 from feectools.linalg.basic     import (Vector, LinearOperator,
         InverseLinearOperator, IdentityOperator, ScaledLinearOperator)
+from feectools.linalg.block import BlockVector, BlockVectorSpace
+
 
 __all__ = (
     'inverse',
@@ -20,7 +22,7 @@ __all__ = (
     'PBiConjugateGradientStabilized',
     'MinimumResidual',
     'LSMR',
-    'GMRES'
+    'GMRES',
 )
 
 #===============================================================================
@@ -64,6 +66,8 @@ def inverse(A, solver, **kwargs):
         'minres'   : MinimumResidual,
         'lsmr'     : LSMR,
         'gmres'    : GMRES,
+        'uzawa'    : UzawaSolver,
+        'schur'    : SchurComplementSolver 
     }
 
     # Check solver input
@@ -1912,3 +1916,219 @@ class GMRES(InverseLinearOperator):
     def dot(self, b, out=None):
         return self.solve(b, out=out)
 
+
+class UzawaSolver(InverseLinearOperator):
+    def __init__(self, A, *, A11, A22, B1, B2,
+                x0=None, tol=1e-6, maxiter=1000, verbose=False, recycle=False,
+                inner_tol=1e-7, inner_maxiter=1000):
+
+
+        self._options = {
+            "x0": x0, "tol": tol, "maxiter": maxiter,
+            "verbose": verbose, "recycle": recycle,
+        }
+        super().__init__(A, **self._options)
+
+        self._inner_tol = inner_tol if inner_tol is not None else tol
+        self._inner_maxiter = inner_maxiter
+
+        self._A11 = A11
+        self._A22 = A22
+        self._B1  = B1
+        self._B2  = B2
+
+        self._A11inv = self._inner_solve(A11, tol=self._inner_tol, maxiter=self._inner_maxiter)
+        self._A22inv = self._inner_solve(A22, tol=self._inner_tol, maxiter=self._inner_maxiter)
+
+        self._info = None
+
+    @staticmethod
+    def _inner_solve(A, tol=1e-10, maxiter=200):
+        """Iterative GMRES inverse — works for non-symmetric positive definite operators."""
+        from feectools.linalg.solvers import inverse
+        return inverse(A, "gmres", tol=tol, maxiter=maxiter, recycle=True)
+
+    def update_A11(self, A11):
+        """Recompute A11 inverse when dt changes."""
+        self._A11 = A11
+        self._A11inv = self._inner_solve(A11, tol=self._inner_tol, maxiter=self._inner_maxiter)
+
+    def solve(self, b, out=None):
+
+        A11    = self._A11
+        A22    = self._A22
+        B1     = self._B1
+        B2     = self._B2
+        A11inv = self._A11inv
+        A22inv = self._A22inv
+
+        tol     = self._options["tol"]
+        maxiter = self._options["maxiter"]
+        verbose = self._options["verbose"]
+        recycle = self._options["recycle"]
+
+        F    = b[0]
+        g    = b[1]
+        f_u  = F[0]
+        f_ue = F[1]
+
+        # initial guess
+        x0 = self._options["x0"]
+        if x0 is not None:
+            u  = x0[0][0].copy()
+            ue = x0[0][1].copy()
+            p  = x0[1].copy()
+        else:
+            u  = A11.domain.zeros()
+            ue = A22.domain.zeros()
+            p  = B1.codomain.zeros()
+
+
+        if verbose:
+            print("Uzawa solver:")
+            print("+---------+---------------------+")
+            print("+ Iter. # | L2-norm of residual |")
+            print("+---------+---------------------+")
+            template = "| {:7d} | {:19.2e} |"
+
+        for iteration in range(1, maxiter + 1):
+
+            # solve A11 * u = f_u - B1^T * p
+            rhs_u = f_u - B1.T.dot(p)
+            u = A11inv.dot(rhs_u)
+
+            # solve A22 * ue = f_ue - B2^T * p
+            rhs_ue = f_ue - B2.T.dot(p)
+            ue = A22inv.dot(rhs_ue)
+
+            # constraint residual: R = B1*u + B2*ue - g
+            R = B1.dot(u) + B2.dot(ue) - g
+            residual_norm = sqrt(R.inner(R).real)
+
+            if verbose:
+                print(template.format(iteration, residual_norm))
+
+            if residual_norm < tol:
+                break
+
+            # pressure update: steepest descent step size
+            S_R   = B1.dot(A11inv.dot(B1.T.dot(R))) + B2.dot(A22inv.dot(B2.T.dot(R)))
+            alpha = R.inner(R).real / R.inner(S_R).real
+            p    += alpha * R
+
+        if verbose:
+            print("+---------+---------------------+")
+
+        self._info = {
+            'niter'   : iteration,
+            'success' : residual_norm < tol,
+            'res_norm': residual_norm,
+        }
+
+        if recycle:
+            block_u = BlockVector(self.domain.spaces[0], blocks=[u, ue])
+            self._options["x0"] = BlockVector(self.domain, blocks=[block_u, p])
+
+        if out is not None:
+            out[0][0] = u
+            out[0][1] = ue
+            out[1]    = p
+            return out
+
+        block_u = BlockVector(BlockVectorSpace(A11.domain, A22.domain), blocks=[u, ue])
+        return BlockVector(self.domain, blocks=[block_u, p])
+
+    def dot(self, b, out=None):
+        return self.solve(b, out=out)
+
+
+class SchurComplementSolver(InverseLinearOperator):
+    def __init__(self, A, *, A11, A22, B1, B2,
+                 x0=None, tol=1e-6, maxiter=1000, verbose=False, recycle=False,
+                 inner_tol=1e-7, inner_maxiter=1000):
+
+        self._options = {
+            "x0": x0, "tol": tol, "maxiter": maxiter,
+            "verbose": verbose, "recycle": recycle,
+        }
+        super().__init__(A, **self._options)
+
+        self._inner_tol     = inner_tol if inner_tol is not None else tol
+        self._inner_maxiter = inner_maxiter
+
+        self._A11 = A11
+        self._A22 = A22
+        self._B1  = B1
+        self._B2  = B2
+
+        self._A11inv = self._make_inner_solver(A11, self._inner_tol, self._inner_maxiter)
+        self._A22inv = self._make_inner_solver(A22, self._inner_tol, self._inner_maxiter)
+
+        self._Sinv = self._build_schur_solver()
+
+        self._info = None
+
+    def _make_inner_solver(self, A, tol, maxiter):
+        from feectools.linalg.solvers import inverse
+        return inverse(A, "gmres", tol=tol, maxiter=maxiter, recycle=True)
+
+    def _build_schur_operator(self):
+        return (self._B1 @ self._A11inv @ self._B1.T
+              + self._B2 @ self._A22inv @ self._B2.T)
+
+    def _build_schur_solver(self):
+        from feectools.linalg.solvers import inverse
+        S = self._build_schur_operator()
+        return inverse(S, "gmres",
+                       tol=self._options["tol"],
+                       maxiter=self._options["maxiter"],
+                       verbose=self._options["verbose"],
+                       recycle=self._options["recycle"])
+
+    def update_A11(self, A11):
+        """Recompute A11 inverse and update Schur operator when dt changes."""
+        self._A11    = A11
+        self._A11inv = self._make_inner_solver(A11, self._inner_tol, self._inner_maxiter)
+        # update operator in-place so Sinv keeps its recycled state
+        self._Sinv.linop = self._build_schur_operator()
+
+    def solve(self, b, out=None):
+
+        A11inv = self._A11inv
+        A22inv = self._A22inv
+        B1     = self._B1
+        B2     = self._B2
+        Sinv   = self._Sinv
+
+        F    = b[0]
+        g    = b[1]
+        f_u  = F[0]
+        f_ue = F[1]
+
+        # Schur complement RHS: B1*u_f + B2*ue_f - g
+        rhs_p = B1.dot(A11inv.dot(f_u)) + B2.dot(A22inv.dot(f_ue)) - g
+
+        # solve S * p = rhs_p
+        p = Sinv.dot(rhs_p)
+
+        # recover velocities
+        u  = A11inv.dot(f_u  - B1.T.dot(p))
+        ue = A22inv.dot(f_ue - B2.T.dot(p))
+
+        self._info = Sinv.get_info()
+
+        if self._options["recycle"]:
+            block_u = BlockVector(self.domain.spaces[0], blocks=[u, ue])
+            self._options["x0"] = BlockVector(self.domain, blocks=[block_u, p])
+
+        if out is not None:
+            out[0][0] = u
+            out[0][1] = ue
+            out[1]    = p
+            return out
+
+        block_u = BlockVector(self.domain.spaces[0], blocks=[u, ue])
+        return BlockVector(self.domain, blocks=[block_u, p])
+
+    def dot(self, b, out=None):
+        return self.solve(b, out=out)

{feectools-0.1.9 → feectools-0.1.11/feectools.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.1.9
+Version: 0.1.11
 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>

{feectools-0.1.9 → feectools-0.1.11}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name            = "feectools"
-version         = "0.1.9"
+version         = "0.1.11"
 description     = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies."
 readme          = "README.md"
 requires-python = ">= 3.10"

feectools-0.1.9/feectools/ddm/mpi.py

@@ -1,114 +0,0 @@
-from dataclasses import dataclass
-from time import time
-from typing import TYPE_CHECKING
-
-
-# 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
-
-
-try:
-    # Disable MPI when using CuPy due to known segfault issues with OpenMPI + CUDA
-    import os
-    if os.environ.get('ARRAY_BACKEND') == 'cupy':
-        raise ImportError("MPI disabled when using CuPy backend")
-    
-    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
-
-# 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