feectools 0.1.10__tar.gz → 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (111) hide show
  1. {feectools-0.1.10/feectools.egg-info → feectools-0.2.0}/PKG-INFO +1 -1
  2. feectools-0.2.0/feectools/__init__.py +4 -0
  3. {feectools-0.1.10 → feectools-0.2.0}/feectools/api/settings.py +8 -5
  4. feectools-0.2.0/feectools/ddm/mpi.py +225 -0
  5. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/partition.py +41 -2
  6. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/basic.py +137 -27
  7. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/direct_solvers.py +5 -2
  8. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/fft.py +0 -6
  9. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/kron.py +0 -6
  10. feectools-0.2.0/feectools/linalg/tests/test_toarray.py +162 -0
  11. {feectools-0.1.10 → feectools-0.2.0/feectools.egg-info}/PKG-INFO +1 -1
  12. {feectools-0.1.10 → feectools-0.2.0}/feectools.egg-info/SOURCES.txt +1 -0
  13. {feectools-0.1.10 → feectools-0.2.0}/pyproject.toml +1 -1
  14. feectools-0.1.10/feectools/ddm/mpi.py +0 -114
  15. feectools-0.1.10/feectools/utilities/__init__.py +0 -0
  16. {feectools-0.1.10 → feectools-0.2.0}/AUTHORS +0 -0
  17. {feectools-0.1.10 → feectools-0.2.0}/LICENSE +0 -0
  18. {feectools-0.1.10 → feectools-0.2.0}/README.md +0 -0
  19. {feectools-0.1.10/feectools → feectools-0.2.0/feectools/accelerate}/__init__.py +0 -0
  20. {feectools-0.1.10 → feectools-0.2.0}/feectools/accelerate/accelerate.py +0 -0
  21. {feectools-0.1.10 → feectools-0.2.0}/feectools/accelerate/compile_psydac.mk +0 -0
  22. {feectools-0.1.10/feectools/accelerate → feectools-0.2.0/feectools/api}/__init__.py +0 -0
  23. {feectools-0.1.10 → feectools-0.2.0}/feectools/api/essential_bc.py +0 -0
  24. {feectools-0.1.10 → feectools-0.2.0}/feectools/api/fem_bilinear_form.py +0 -0
  25. {feectools-0.1.10 → feectools-0.2.0}/feectools/api/fem_common.py +0 -0
  26. {feectools-0.1.10 → feectools-0.2.0}/feectools/api/fem_sum_form.py +0 -0
  27. {feectools-0.1.10 → feectools-0.2.0}/feectools/core/__init__.py +0 -0
  28. {feectools-0.1.10 → feectools-0.2.0}/feectools/core/bsplines.py +0 -0
  29. {feectools-0.1.10 → feectools-0.2.0}/feectools/core/bsplines_kernels.py +0 -0
  30. {feectools-0.1.10 → feectools-0.2.0}/feectools/core/field_evaluation_kernels.py +0 -0
  31. {feectools-0.1.10/feectools/api → feectools-0.2.0/feectools/core/tests}/__init__.py +0 -0
  32. {feectools-0.1.10 → feectools-0.2.0}/feectools/core/tests/test_bsplines.py +0 -0
  33. {feectools-0.1.10 → feectools-0.2.0}/feectools/core/tests/test_bsplines_kernel.py +0 -0
  34. {feectools-0.1.10 → feectools-0.2.0}/feectools/core/tests/test_bsplines_pyccel.py +0 -0
  35. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/__init__.py +0 -0
  36. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/basic.py +0 -0
  37. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/blocking_data_exchanger.py +0 -0
  38. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/cart.py +0 -0
  39. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/interface_data_exchanger.py +0 -0
  40. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/nonblocking_data_exchanger.py +0 -0
  41. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/petsc.py +0 -0
  42. {feectools-0.1.10/feectools/core → feectools-0.2.0/feectools/ddm}/tests/__init__.py +0 -0
  43. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/tests/test_cart_1d.py +0 -0
  44. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/tests/test_cart_2d.py +0 -0
  45. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/tests/test_cart_3d.py +0 -0
  46. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/tests/test_multicart_2d.py +0 -0
  47. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/tests/test_partition.py +0 -0
  48. {feectools-0.1.10 → feectools-0.2.0}/feectools/ddm/utilities.py +0 -0
  49. {feectools-0.1.10/feectools/ddm/tests → feectools-0.2.0/feectools/feec}/__init__.py +0 -0
  50. {feectools-0.1.10 → feectools-0.2.0}/feectools/feec/derivatives.py +0 -0
  51. {feectools-0.1.10 → feectools-0.2.0}/feectools/feec/dof_kernels.py +0 -0
  52. {feectools-0.1.10 → feectools-0.2.0}/feectools/feec/global_geometric_projectors.py +0 -0
  53. {feectools-0.1.10 → feectools-0.2.0}/feectools/feec/hodge.py +0 -0
  54. {feectools-0.1.10/feectools/feec → feectools-0.2.0/feectools/fem}/__init__.py +0 -0
  55. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/basic.py +0 -0
  56. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/grid.py +0 -0
  57. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/lst_preconditioner.py +0 -0
  58. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/partitioning.py +0 -0
  59. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/projectors.py +0 -0
  60. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/splines.py +0 -0
  61. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tensor.py +0 -0
  62. {feectools-0.1.10/feectools/fem → feectools-0.2.0/feectools/fem/tests}/__init__.py +0 -0
  63. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/analytical_profiles_1d.py +0 -0
  64. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/analytical_profiles_base.py +0 -0
  65. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/splines_error_bounds.py +0 -0
  66. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/test_dirichlet_projectors.py +0 -0
  67. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/test_spline_histopolation.py +0 -0
  68. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/test_spline_interpolation.py +0 -0
  69. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/test_splines.py +0 -0
  70. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/test_splines_par.py +0 -0
  71. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/test_tensor.py +0 -0
  72. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/test_vector_spaces.py +0 -0
  73. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/tests/utilities.py +0 -0
  74. {feectools-0.1.10 → feectools-0.2.0}/feectools/fem/vector.py +0 -0
  75. {feectools-0.1.10/feectools/fem/tests → feectools-0.2.0/feectools/linalg}/__init__.py +0 -0
  76. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/block.py +0 -0
  77. {feectools-0.1.10/feectools/linalg → feectools-0.2.0/feectools/linalg/kernels}/__init__.py +0 -0
  78. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/kernels/axpy_kernels.py +0 -0
  79. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/kernels/inner_kernels.py +0 -0
  80. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/kernels/matvec_kernels.py +0 -0
  81. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/kernels/stencil2IJV_kernels.py +0 -0
  82. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/kernels/stencil2coo_kernels.py +0 -0
  83. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/kernels/transpose_kernels.py +0 -0
  84. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/memory.py +0 -0
  85. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/solvers.py +0 -0
  86. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/sparse.py +0 -0
  87. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/stencil.py +0 -0
  88. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/stencil_dot_kernels.py +0 -0
  89. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/stencil_transpose_kernels.py +0 -0
  90. {feectools-0.1.10/feectools/linalg/kernels → feectools-0.2.0/feectools/linalg/tests}/__init__.py +0 -0
  91. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/tests/test_block.py +0 -0
  92. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/tests/test_fft.py +0 -0
  93. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/tests/test_kron_stencil_matrix.py +0 -0
  94. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/tests/test_linalg.py +0 -0
  95. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/tests/test_matrix_free.py +0 -0
  96. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/tests/test_solvers.py +0 -0
  97. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/tests/test_stencil_interface_matrix.py +0 -0
  98. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/tests/test_stencil_vector.py +0 -0
  99. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/tests/test_stencil_vector_space.py +0 -0
  100. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/tests/utilities.py +0 -0
  101. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/topetsc.py +0 -0
  102. {feectools-0.1.10 → feectools-0.2.0}/feectools/linalg/utilities.py +0 -0
  103. {feectools-0.1.10/feectools/linalg/tests → feectools-0.2.0/feectools/utilities}/__init__.py +0 -0
  104. {feectools-0.1.10 → feectools-0.2.0}/feectools/utilities/quadratures.py +0 -0
  105. {feectools-0.1.10 → feectools-0.2.0}/feectools/utilities/utils.py +0 -0
  106. {feectools-0.1.10 → feectools-0.2.0}/feectools/version.py +0 -0
  107. {feectools-0.1.10 → feectools-0.2.0}/feectools.egg-info/dependency_links.txt +0 -0
  108. {feectools-0.1.10 → feectools-0.2.0}/feectools.egg-info/entry_points.txt +0 -0
  109. {feectools-0.1.10 → feectools-0.2.0}/feectools.egg-info/requires.txt +0 -0
  110. {feectools-0.1.10 → feectools-0.2.0}/feectools.egg-info/top_level.txt +0 -0
  111. {feectools-0.1.10 → feectools-0.2.0}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: feectools
3
- Version: 0.1.10
3
+ Version: 0.2.0
4
4
  Summary: Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies.
5
5
  Author-email: Psydac development team <psydac@googlegroups.com>
6
6
  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>
@@ -0,0 +1,4 @@
1
+ # MPI switch, read once when feectools.ddm.mpi is first imported:
2
+ # None -> use mpi4py if it is available (default)
3
+ # False -> do not import mpi4py, use MockMPI (serial runs; saves ~1 s of import time)
4
+ use_mpi = None
@@ -43,13 +43,16 @@ PSYDAC_BACKEND_NVPYCCEL = {'name': 'pyccel',
43
43
  'openmp' : False}
44
44
  # ...
45
45
 
46
- # Get gfortran version
47
- gfortran_version_output = subprocess.check_output(['gfortran', '--version']).decode('utf-8') # nosec B603, B607
48
- gfortran_version_string = re.search(r"(\d+\.\d+\.\d+)", gfortran_version_output).group()
49
- gfortran_version = Version(gfortran_version_string)
46
+ def _get_gfortran_version():
47
+ """Version of the installed gfortran. Only queried where it is needed (see below), since it spawns a subprocess."""
48
+ gfortran_version_output = subprocess.check_output(['gfortran', '--version']).decode('utf-8') # nosec B603, B607
49
+ gfortran_version_string = re.search(r"(\d+\.\d+\.\d+)", gfortran_version_output).group()
50
+ return Version(gfortran_version_string)
50
51
 
51
52
  # Platform-dependent flags
52
- if platform.system() == "Darwin" and platform.machine() == 'arm64' and gfortran_version >= Version("14"):
53
+ # (the gfortran version is only relevant on Apple silicon; evaluating it lazily avoids running
54
+ # a subprocess on every import and lets feectools be imported without gfortran on other platforms)
55
+ if platform.system() == "Darwin" and platform.machine() == 'arm64' and _get_gfortran_version() >= Version("14"):
53
56
 
54
57
  # Apple silicon requires architecture-specific flags (see https://github.com/pyccel/psydac/pull/411)
55
58
  # which are only available on GCC version >= 14
@@ -0,0 +1,225 @@
1
+ """Detection of whether the process was launched by an MPI launcher.
2
+
3
+ Importing ``mpi4py.MPI`` calls ``MPI_Init``, and any collective (``bcast``,
4
+ ``Barrier``, ...) issued afterwards costs something even on a single process.
5
+ A plain ``python script.py`` run should therefore never touch MPI at all, even
6
+ when mpi4py happens to be installed. This module answers the only question
7
+ that decides it: was this process started by ``mpirun``/``mpiexec``/``srun``
8
+ (or an equivalent launcher)?
9
+
10
+ The answer is read from the environment the launcher sets up, so it is
11
+ available before mpi4py is imported.
12
+ """
13
+
14
+ import os
15
+ import sys
16
+
17
+ from dataclasses import dataclass
18
+ from time import time
19
+ from typing import TYPE_CHECKING
20
+
21
+
22
+ # Per-process variables exported by the process managers behind the common
23
+ # launchers. Each is set only for processes started *by* the launcher, so the
24
+ # presence of any one of them means "this rank belongs to an MPI job".
25
+ # SLURM_PROCID is deliberately absent: it is also set for the script of a
26
+ # plain `sbatch` job, which is not an MPI launch. `srun` is covered by the
27
+ # PMI/PMIX variables its MPI plugin exports.
28
+ _LAUNCHER_ENV_VARS = (
29
+ "OMPI_COMM_WORLD_RANK", # Open MPI (and derivatives: Spectrum, ...)
30
+ "PMI_RANK", # MPICH, Intel MPI, MS-MPI, Cray, srun (pmi2)
31
+ "PMIX_RANK", # PMIx, used by srun --mpi=pmix and Open MPI 5
32
+ "MV2_COMM_WORLD_RANK", # MVAPICH2
33
+ "MPI_LOCALRANKID", # Hydra (mpiexec.hydra)
34
+ "ALPS_APP_PE", # Cray ALPS aprun
35
+ "PALS_RANKID", # Cray PALS palsrun
36
+ )
37
+
38
+ # Escape hatch: force the decision either way without touching code, e.g. for
39
+ # a launcher whose variables are not listed above.
40
+ _OVERRIDE_ENV_VAR = "STRUPHY_MPI"
41
+
42
+ _TRUE_VALUES = ("1", "true", "yes", "on")
43
+ _FALSE_VALUES = ("0", "false", "no", "off")
44
+
45
+
46
+ def _override() -> bool | None:
47
+ """Value of ``STRUPHY_MPI``, or None if unset/unrecognized."""
48
+ value = os.environ.get(_OVERRIDE_ENV_VAR)
49
+ if value is None:
50
+ return None
51
+ value = value.strip().lower()
52
+ if value in _TRUE_VALUES:
53
+ return True
54
+ if value in _FALSE_VALUES:
55
+ return False
56
+ return None
57
+
58
+
59
+ def launched_under_mpi() -> bool:
60
+ """Whether this process was started by an MPI launcher.
61
+
62
+ Returns
63
+ -------
64
+ bool
65
+ True if a launcher's per-rank environment variable is present, or if
66
+ the application itself already initialized MPI (in which case using
67
+ the communicator is free). ``STRUPHY_MPI=0``/``1`` overrides
68
+ the detection.
69
+ """
70
+ override = _override()
71
+ if override is not None:
72
+ return override
73
+
74
+ if any(var in os.environ for var in _LAUNCHER_ENV_VARS):
75
+ return True
76
+
77
+ # The application may have initialized MPI itself (embedded interpreter,
78
+ # or an explicit `from mpi4py import MPI`). Only inspect mpi4py if it is
79
+ # already imported: importing it here is exactly what must be avoided.
80
+ mpi_module = sys.modules.get("mpi4py.MPI")
81
+ if mpi_module is not None:
82
+ try:
83
+ return bool(mpi_module.Is_initialized())
84
+ except AttributeError:
85
+ return False
86
+
87
+ return False
88
+
89
+
90
+ # Might not be needed
91
+ class MPICommWrapper:
92
+ def __init__(self, use_mpi=True):
93
+ self.use_mpi = use_mpi
94
+ if use_mpi:
95
+ from mpi4py import MPI
96
+
97
+ self.comm = MPI.COMM_WORLD
98
+ else:
99
+ self.comm = MockComm()
100
+
101
+ def __getattr__(self, name):
102
+ return getattr(self.comm, name)
103
+
104
+
105
+ class MockComm:
106
+ def __getattr__(self, name):
107
+ # Return a function that does nothing and returns None
108
+ def dummy(*args, **kwargs):
109
+ return None
110
+
111
+ return dummy
112
+
113
+ # Override some functions
114
+ def Get_rank(self):
115
+ return 0
116
+
117
+ def Get_size(self):
118
+ return 1
119
+
120
+ def Barrier(self):
121
+ return
122
+
123
+
124
+ class MPIwrapper:
125
+ def __init__(
126
+ self,
127
+ use_mpi: bool = False,
128
+ verbose: bool = False,
129
+ ):
130
+ self.use_mpi = use_mpi
131
+ if use_mpi:
132
+ from mpi4py import MPI
133
+
134
+ self._MPI = MPI
135
+ if verbose:
136
+ print("MPI is enabled")
137
+ else:
138
+ self._MPI = MockMPI()
139
+ if verbose:
140
+ print("MPI is NOT enabled")
141
+
142
+ @property
143
+ def MPI(self):
144
+ return self._MPI
145
+
146
+
147
+ class MockMPI:
148
+ def __getattr__(self, name):
149
+ # Return a function that does nothing and returns None
150
+ def dummy(*args, **kwargs):
151
+ return None
152
+
153
+ return dummy
154
+
155
+ # Override some functions
156
+ @property
157
+ def COMM_WORLD(self):
158
+ return MockComm()
159
+
160
+ # def comm_Get_rank(self):
161
+ # return 0
162
+
163
+ # def comm_Get_size(self):
164
+ # return 1
165
+
166
+
167
+ def _mpi_disabled():
168
+ """True if the user or the host application opted out of MPI.
169
+
170
+ Importing mpi4py initializes MPI, which takes close to a second. Serial runs that
171
+ never use MPI can skip it entirely, in two ways:
172
+
173
+ * export ``FEECTOOLS_MPI=0`` before starting Python, or
174
+ * set ``feectools.use_mpi = False`` before this module is first imported
175
+ (in-process, so it is not inherited by subprocesses such as ``mpirun``).
176
+
177
+ The MockMPI wrapper below is then used, exactly as if mpi4py were not installed.
178
+ """
179
+ import os
180
+
181
+ import feectools
182
+
183
+ if getattr(feectools, 'use_mpi', None) is False:
184
+ return True
185
+ return os.environ.get('FEECTOOLS_MPI', '').strip().lower() in ('0', 'false', 'no', 'off')
186
+
187
+
188
+ if launched_under_mpi():
189
+ try:
190
+ # Disable MPI when using CuPy due to known segfault issues with OpenMPI + CUDA
191
+ import os
192
+ if os.environ.get('ARRAY_BACKEND') == 'cupy':
193
+ raise ImportError("MPI disabled when using CuPy backend")
194
+
195
+ if _mpi_disabled():
196
+ raise ImportError("MPI disabled (feectools.use_mpi = False or FEECTOOLS_MPI=0)")
197
+
198
+ from mpi4py import MPI
199
+
200
+ _comm = MPI.COMM_WORLD
201
+ # rank = _comm.Get_rank()
202
+ # size = _comm.Get_size()
203
+ mpi_enabled = True
204
+ except ImportError:
205
+ # mpi4py not installed
206
+ mpi_enabled = False
207
+ except Exception:
208
+ # mpi4py installed but not running under mpirun
209
+ mpi_enabled = False
210
+ else:
211
+ mpi_enabled = False
212
+
213
+ # TODO: add environment variable for mpi use
214
+ mpi_wrapper = MPIwrapper(
215
+ use_mpi=mpi_enabled,
216
+ verbose=False,
217
+ )
218
+
219
+ # TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime.
220
+ if TYPE_CHECKING:
221
+ from mpi4py import MPI
222
+
223
+ mpi = MPI
224
+ else:
225
+ mpi = mpi_wrapper.MPI
@@ -2,10 +2,49 @@ import cunumpy as xp
2
2
  import numpy as np
3
3
  import numpy.ma as ma
4
4
 
5
- from sympy.ntheory import factorint
5
+ __all__ = ('compute_dims', 'partition_procs_per_patch')
6
+
7
+ #==============================================================================
8
+ def factorint(n, multiple=False):
9
+ """
10
+ Prime factorization of an integer by trial division.
6
11
 
12
+ Drop-in replacement for the subset of ``sympy.ntheory.factorint`` used here.
13
+ Importing sympy takes ~1 s and is not needed for the small integers
14
+ (process counts, number of grid points) that are factorized in this module.
7
15
 
8
- __all__ = ('compute_dims', 'partition_procs_per_patch')
16
+ Parameters
17
+ ----------
18
+ n : int
19
+ Integer to factorize.
20
+
21
+ multiple : bool
22
+ If False (default), return a dict {prime: multiplicity}.
23
+ If True, return the list of primes in ascending order, repeated
24
+ according to their multiplicity.
25
+ """
26
+ n = int(n)
27
+ factors = {}
28
+
29
+ # same conventions as sympy for non-positive input
30
+ if n == 0:
31
+ factors[0] = 1
32
+ else:
33
+ if n < 0:
34
+ factors[-1] = 1
35
+ n = -n
36
+ p = 2
37
+ while p * p <= n:
38
+ while n % p == 0:
39
+ factors[p] = factors.get(p, 0) + 1
40
+ n //= p
41
+ p += 1 if p == 2 else 2
42
+ if n > 1:
43
+ factors[n] = factors.get(n, 0) + 1
44
+
45
+ if multiple:
46
+ return [p for p in sorted(factors) for _ in range(factors[p])]
47
+ return factors
9
48
 
10
49
  #==============================================================================
11
50
  def partition_procs_per_patch(npts, size):
@@ -7,6 +7,7 @@ provides the fundamental classes for linear algebra operations.
7
7
 
8
8
  """
9
9
 
10
+ import itertools
10
11
  from abc import ABC, abstractmethod
11
12
  from types import LambdaType
12
13
  from inspect import signature
@@ -280,13 +281,143 @@ class LinearOperator(ABC):
280
281
  upon convertion to matrix.
281
282
  """
282
283
 
283
- @abstractmethod
284
- def tosparse(self):
285
- """ Convert to a sparse matrix in any of the formats supported by scipy.sparse."""
284
+ def toarray(self, out=None, is_sparse=False, format='csr'):
285
+ """
286
+ Assemble the global matrix of the linear operator column by column.
286
287
 
287
- @abstractmethod
288
- def toarray(self):
289
- """ Convert to Numpy 2D array. """
288
+ Column j is computed as ``self.dot(e_j)``, where e_j is the j-th
289
+ canonical basis vector of the domain, in the global numbering of
290
+ ``Vector.toarray()`` (C-ordered within each StencilVector, blocks
291
+ concatenated). Since only ``dot`` is used, this works for any linear
292
+ operator, including matrix-free ones, whose domain is a
293
+ StencilVectorSpace or a (possibly nested) BlockVectorSpace thereof.
294
+
295
+ The cost is one call to ``dot`` per global degree of freedom of the
296
+ domain, hence this default is meant for testing and small problems.
297
+ Subclasses with an explicit matrix representation should override it.
298
+
299
+ In parallel, all ranks call ``dot`` collectively, each rank fills the
300
+ rows it owns, and every rank receives the full matrix.
301
+
302
+ Parameters
303
+ ----------
304
+ out : numpy.ndarray, optional
305
+ Dense array of shape ``self.shape`` into which the result is
306
+ written in place. Must be None if ``is_sparse`` is True.
307
+
308
+ is_sparse : bool
309
+ If True, return a scipy.sparse matrix, otherwise a dense
310
+ numpy.ndarray.
311
+
312
+ format : str
313
+ Sparse format, only used if ``is_sparse`` is True: one of 'csr'
314
+ (default), 'csc', 'bsr', 'lil', 'dok', 'coo' or 'dia'.
315
+
316
+ Returns
317
+ -------
318
+ numpy.ndarray or scipy.sparse matrix
319
+ The global matrix of shape ``self.shape``, identical on all ranks.
320
+ """
321
+ from feectools.linalg.block import BlockVectorSpace
322
+ from feectools.linalg.stencil import StencilVectorSpace
323
+
324
+ # Flatten the (possibly nested) block structure of the domain
325
+ def leaves(w):
326
+ if isinstance(w.space, StencilVectorSpace):
327
+ return [w]
328
+ elif isinstance(w.space, BlockVectorSpace):
329
+ return [lw for b in w.blocks for lw in leaves(b)]
330
+ else:
331
+ raise TypeError(f'{type(self).__name__}.toarray() requires a domain made of '
332
+ f'StencilVectorSpaces, not {type(w.space).__name__}.')
333
+
334
+ e_j = self.domain.zeros()
335
+ Ae_j = self.codomain.zeros()
336
+ e_j_leaves = leaves(e_j)
337
+ offsets = xp.cumsum([0] + [lw.space.dimension for lw in e_j_leaves[:-1]])
338
+
339
+ if is_sparse:
340
+ assert out is None, 'out must be None if is_sparse is True.'
341
+ assert format in ('csr', 'csc', 'bsr', 'lil', 'dok', 'coo', 'dia'), \
342
+ f'Unknown sparse format {format!r}.'
343
+ rows, cols, data = [], [], []
344
+ elif out is None:
345
+ out = xp.zeros(self.shape, dtype=self.dtype)
346
+ else:
347
+ assert isinstance(out, xp.ndarray)
348
+ assert out.shape == self.shape, f'out has shape {out.shape}, expected {self.shape}.'
349
+
350
+ # Index ranges owned by each rank, for every leaf
351
+ bounds = [(lw.starts, lw.ends) for lw in e_j_leaves]
352
+ if e_j_leaves[0].space.parallel:
353
+ comm = e_j_leaves[0].space.cart.comm
354
+ rank = comm.Get_rank()
355
+ all_bounds = comm.allgather(bounds)
356
+ else:
357
+ comm = None
358
+ rank = 0
359
+ all_bounds = [bounds]
360
+
361
+ # All ranks loop over all columns, since dot() is collective;
362
+ # only the owner of index i sets the entry of e_j to one.
363
+ for owner, owner_bounds in enumerate(all_bounds):
364
+ for lw, offset, (starts, ends) in zip(e_j_leaves, offsets, owner_bounds):
365
+ for i in itertools.product(*(range(s, e + 1) for s, e in zip(starts, ends))):
366
+ if rank == owner:
367
+ lw[i] = 1
368
+ lw.update_ghost_regions()
369
+ self.dot(e_j, out=Ae_j)
370
+ if rank == owner:
371
+ lw[i] = 0
372
+
373
+ # Global vector with nonzeros only in the rows owned by this rank
374
+ col_j = Ae_j.toarray()
375
+ j = offset + xp.ravel_multi_index(i, lw.space.npts)
376
+ if is_sparse:
377
+ nz = xp.flatnonzero(col_j)
378
+ rows.append(nz)
379
+ cols.append(xp.full(nz.size, j))
380
+ data.append(col_j[nz])
381
+ else:
382
+ out[:, j] = col_j
383
+ # Clear the ghost regions of the last entry set to one
384
+ lw.update_ghost_regions()
385
+
386
+ if not is_sparse:
387
+ if comm is not None:
388
+ from feectools.ddm.mpi import mpi as MPI
389
+ comm.Allreduce(MPI.IN_PLACE, out, op=MPI.SUM)
390
+ return out
391
+
392
+ rows = xp.concatenate(rows) if rows else xp.zeros(0, dtype=int)
393
+ cols = xp.concatenate(cols) if cols else xp.zeros(0, dtype=int)
394
+ data = xp.concatenate(data) if data else xp.zeros(0, dtype=self.dtype)
395
+ if comm is not None:
396
+ rows = xp.concatenate(comm.allgather(rows))
397
+ cols = xp.concatenate(comm.allgather(cols))
398
+ data = xp.concatenate(comm.allgather(data))
399
+
400
+ return coo_matrix((data, (rows, cols)), shape=self.shape).asformat(format)
401
+
402
+ def tosparse(self, format='csr'):
403
+ """
404
+ Assemble the global matrix of the linear operator as a scipy.sparse matrix.
405
+
406
+ Default implementation calling the generic ``LinearOperator.toarray``
407
+ with ``is_sparse=True``; see there for cost and parallel behavior.
408
+ Subclasses with an explicit matrix representation should override it.
409
+
410
+ Parameters
411
+ ----------
412
+ format : str
413
+ One of 'csr' (default), 'csc', 'bsr', 'lil', 'dok', 'coo' or 'dia'.
414
+
415
+ Returns
416
+ -------
417
+ scipy.sparse matrix
418
+ The global matrix of shape ``self.shape``, identical on all ranks.
419
+ """
420
+ return LinearOperator.toarray(self, is_sparse=True, format=format)
290
421
 
291
422
  @abstractmethod
292
423
  def dot(self, v, out=None):
@@ -977,9 +1108,6 @@ class ComposedLinearOperator(LinearOperator):
977
1108
  def dtype(self):
978
1109
  return None
979
1110
 
980
- def toarray(self):
981
- raise NotImplementedError('toarray() is not defined for ComposedLinearOperators.')
982
-
983
1111
  def tosparse(self):
984
1112
  mats = [M.tosparse() for M in self._multiplicants]
985
1113
  M = mats[0]
@@ -1084,12 +1212,6 @@ class PowerLinearOperator(LinearOperator):
1084
1212
  """ Returns the power to which the operator is raised. """
1085
1213
  return self._factorial
1086
1214
 
1087
- def toarray(self):
1088
- raise NotImplementedError('toarray() is not defined for PowerLinearOperators.')
1089
-
1090
- def tosparse(self):
1091
- raise NotImplementedError('tosparse() is not defined for PowerLinearOperators.')
1092
-
1093
1215
  def transpose(self, conjugate=False):
1094
1216
  return PowerLinearOperator(domain=self.codomain, codomain=self.domain, A=self._operator.transpose(conjugate=conjugate), n=self._factorial)
1095
1217
 
@@ -1207,12 +1329,6 @@ class InverseLinearOperator(LinearOperator):
1207
1329
  elif key == 'verbose':
1208
1330
  assert isinstance(value, bool), "verbose must be a bool"
1209
1331
 
1210
- def toarray(self):
1211
- raise NotImplementedError('toarray() is not defined for InverseLinearOperators.')
1212
-
1213
- def tosparse(self):
1214
- raise NotImplementedError('tosparse() is not defined for InverseLinearOperators.')
1215
-
1216
1332
  def get_info(self):
1217
1333
  """ Returns the previous convergence information. """
1218
1334
  return self._info
@@ -1366,12 +1482,6 @@ class MatrixFreeLinearOperator(LinearOperator):
1366
1482
 
1367
1483
  return out
1368
1484
 
1369
- def toarray(self):
1370
- raise NotImplementedError('toarray() is not defined for MatrixFreeLinearOperator.')
1371
-
1372
- def tosparse(self):
1373
- raise NotImplementedError('tosparse() is not defined for MatrixFreeLinearOperator.')
1374
-
1375
1485
  def transpose(self, conjugate=False):
1376
1486
  if self._dot_transpose is None:
1377
1487
  raise NotImplementedError('no transpose dot method was given -- cannot create the transpose operator')
@@ -6,9 +6,7 @@
6
6
  from abc import abstractmethod
7
7
  import cunumpy as xp
8
8
  from cunumpy.xp import array_backend
9
- from scipy.linalg.lapack import dgbtrf, dgbtrs, sgbtrf, sgbtrs, cgbtrf, cgbtrs, zgbtrf, zgbtrs
10
9
  from scipy.sparse import spmatrix, dia_matrix
11
- from scipy.sparse.linalg import splu
12
10
 
13
11
  from feectools.linalg.basic import LinearSolver
14
12
 
@@ -53,6 +51,9 @@ class BandedSolver(LinearSolver):
53
51
  self._l = l
54
52
  self._transposed = transposed
55
53
 
54
+ # imported here: scipy.linalg costs ~0.3 s to import and is only needed once a solver is built
55
+ from scipy.linalg.lapack import dgbtrf, dgbtrs, sgbtrf, sgbtrs, cgbtrf, cgbtrs, zgbtrf, zgbtrs
56
+
56
57
  # ... LU factorization
57
58
  if bmat.dtype == xp.float32:
58
59
  self._factor_function = sgbtrf
@@ -186,6 +187,8 @@ class SparseSolver (LinearSolver):
186
187
 
187
188
  assert isinstance(spmat, spmatrix)
188
189
 
190
+ from scipy.sparse.linalg import splu # deferred, see BandedSolver
191
+
189
192
  self._space = xp.ndarray
190
193
  self._splu = splu(spmat.tocsc())
191
194
  self._transposed = transposed
@@ -21,12 +21,6 @@ class DistributedFFTBase(LinearOperator):
21
21
  The function at position i is applied to the i-th tensor direction.
22
22
  If only a single callable is given, it is used for all directions.
23
23
  """
24
- def toarray(self):
25
- raise NotImplementedError('toarray() is not defined for DistributedFFTBase.')
26
-
27
- def tosparse(self):
28
- raise NotImplementedError('tosparse() is not defined for DistributedFFTBase.')
29
-
30
24
  # Possible additions for the future:
31
25
  # * split off the LinearSolver class when used with the space ndarray (as used in the KroneckerLinearSolver),
32
26
  # and make it state if it works in-place (or if it needs temporary memory), and what its optimal
@@ -528,12 +528,6 @@ class KroneckerLinearSolver(LinearOperator):
528
528
  @property
529
529
  def dtype(self):
530
530
  return None
531
-
532
- def toarray(self):
533
- raise NotImplementedError('toarray() is not defined for KroneckerLinearSolvers.')
534
-
535
- def tosparse(self):
536
- raise NotImplementedError('tosparse() is not defined for KroneckerLinearSolvers.')
537
531
 
538
532
  def transpose(self, conjugate=False):
539
533
  new_domain = self._codomain
@@ -0,0 +1,162 @@
1
+ # -*- coding: UTF-8 -*-
2
+ #
3
+ # Tests for the generic LinearOperator.toarray() and LinearOperator.tosparse(),
4
+ # which assemble the matrix of any linear operator from its dot() method.
5
+ #
6
+ import pytest
7
+ import cunumpy as xp
8
+
9
+ from feectools.ddm.mpi import mpi as MPI
10
+ from feectools.ddm.cart import DomainDecomposition, CartDecomposition
11
+ from feectools.linalg.basic import LinearOperator, MatrixFreeLinearOperator, IdentityOperator
12
+ from feectools.linalg.stencil import StencilVectorSpace, StencilMatrix
13
+ from feectools.linalg.block import BlockVectorSpace, BlockLinearOperator
14
+ from feectools.linalg.solvers import inverse
15
+ from feectools.linalg.tests.test_block import compute_global_starts_ends
16
+
17
+ SPARSE_FORMATS = ['csr', 'csc', 'bsr', 'lil', 'dok', 'coo', 'dia']
18
+
19
+ #===============================================================================
20
+ # HELPERS
21
+ #===============================================================================
22
+ def get_space(npts, pads, periods, comm=None):
23
+ D = DomainDecomposition(npts, periods=periods, comm=comm)
24
+ global_starts, global_ends = compute_global_starts_ends(D, npts)
25
+ C = CartDecomposition(D, npts, global_starts, global_ends, pads=pads, shifts=[1] * len(npts))
26
+ return StencilVectorSpace(C)
27
+
28
+ def get_random_matrix(V, W, seed):
29
+ """ Random StencilMatrix, identical on all ranks in the global numbering. """
30
+ rng = xp.random.default_rng(seed)
31
+ M = StencilMatrix(V, W)
32
+ # Draw the full (global) band and keep the local part, so the matrix
33
+ # does not depend on the domain decomposition.
34
+ shape = tuple(W.npts) + M._data.shape[W.ndim:]
35
+ band = rng.random(shape) - 0.5
36
+ idx = tuple(slice(s, e + 1) for s, e in zip(W.starts, W.ends))
37
+ M[idx] = band[idx]
38
+ M.remove_spurious_entries()
39
+ return M
40
+
41
+ def matrix_free(A):
42
+ """ Hide the explicit matrix of A, so that the generic toarray() is used. """
43
+ return MatrixFreeLinearOperator(A.domain, A.codomain, lambda v, out=None: A.dot(v, out=out))
44
+
45
+ def reference(A):
46
+ """ Global dense matrix of A, from its own (row-local in parallel) tosparse(). """
47
+ local = A.tosparse().toarray()
48
+ comm = A.domain.cart.comm if isinstance(A.domain, StencilVectorSpace) else A.domain.spaces[0].cart.comm
49
+ if A.domain.parallel:
50
+ glob = xp.zeros_like(local)
51
+ comm.Allreduce(local, glob, op=MPI.SUM)
52
+ return glob
53
+ return local
54
+
55
+ def check_all(O, ref):
56
+ """ Check dense, in-place and all sparse outputs of the generic toarray()/tosparse(). """
57
+ assert xp.allclose(O.toarray(), ref)
58
+
59
+ out = xp.full(O.shape, 7.0)
60
+ res = O.toarray(out=out)
61
+ assert res is out
62
+ assert xp.allclose(out, ref)
63
+
64
+ for fmt in SPARSE_FORMATS:
65
+ S = O.toarray(is_sparse=True, format=fmt)
66
+ assert S.format == fmt
67
+ assert S.shape == O.shape
68
+ assert xp.allclose(S.toarray(), ref)
69
+
70
+ S = O.tosparse()
71
+ assert S.format == 'csr'
72
+ assert xp.allclose(S.toarray(), ref)
73
+ assert O.tosparse('csc').format == 'csc'
74
+
75
+ def stencil_case(periods, comm):
76
+ V = get_space([6, 5], [2, 1], periods, comm=comm)
77
+ return get_random_matrix(V, V, seed=0)
78
+
79
+ def block_case(periods, comm):
80
+ """ 2x2 block operator, and a nested block operator [[B, 0], [0, A]]. """
81
+ V = get_space([6, 5], [2, 1], periods, comm=comm)
82
+ A = get_random_matrix(V, V, seed=0)
83
+ A01 = get_random_matrix(V, V, seed=1)
84
+ VV = BlockVectorSpace(V, V)
85
+ B = BlockLinearOperator(VV, VV, blocks=[[A, A01], [None, A]])
86
+ Vn = BlockVectorSpace(VV, V)
87
+ N = BlockLinearOperator(Vn, Vn, blocks=[[B, None], [None, A]])
88
+ return A, B, N
89
+
90
+ #===============================================================================
91
+ # SERIAL TESTS
92
+ #===============================================================================
93
+ @pytest.mark.parametrize('periods', [[False, False], [True, False], [True, True]])
94
+ def test_toarray_stencil(periods):
95
+ A = stencil_case(periods, comm=None)
96
+ check_all(matrix_free(A), reference(A))
97
+
98
+ @pytest.mark.parametrize('periods', [[False, False], [True, True]])
99
+ def test_toarray_block(periods):
100
+ A, B, N = block_case(periods, comm=None)
101
+ refA, refB = reference(A), reference(B)
102
+ check_all(matrix_free(B), refB)
103
+
104
+ Z = xp.zeros((refB.shape[0], refA.shape[1]))
105
+ refN = xp.block([[refB, Z], [Z.T, refA]])
106
+ check_all(matrix_free(N), refN)
107
+
108
+ def test_toarray_composite_operators():
109
+ """ Operators that used to raise NotImplementedError in toarray()/tosparse(). """
110
+ A = stencil_case([True, False], comm=None)
111
+ ref = reference(A)
112
+ V = A.domain
113
+
114
+ assert xp.allclose((A @ A).toarray(), ref @ ref)
115
+ assert xp.allclose((A ** 2).toarray(), ref @ ref)
116
+ assert xp.allclose((A ** 2).tosparse().toarray(), ref @ ref)
117
+
118
+ M = A.T @ A + IdentityOperator(V)
119
+ refM = ref.T @ ref + xp.eye(ref.shape[0])
120
+ Minv = inverse(M, 'cg', tol=1e-13, maxiter=1000)
121
+ assert xp.allclose(Minv.toarray(), xp.linalg.inv(refM), atol=1e-8)
122
+
123
+ def test_toarray_invalid_input():
124
+ O = matrix_free(stencil_case([False, False], comm=None))
125
+ with pytest.raises(AssertionError):
126
+ O.toarray(out=xp.zeros((O.shape[0], O.shape[1] + 1)))
127
+ with pytest.raises(AssertionError):
128
+ O.toarray(out=xp.zeros(O.shape), is_sparse=True)
129
+ with pytest.raises(AssertionError):
130
+ O.toarray(is_sparse=True, format='xyz')
131
+
132
+ #===============================================================================
133
+ # PARALLEL TESTS
134
+ #===============================================================================
135
+ @pytest.mark.parametrize('periods', [[False, False], [True, False], [True, True]])
136
+ @pytest.mark.parallel
137
+ def test_toarray_stencil_parallel(periods):
138
+ A = stencil_case(periods, comm=MPI.COMM_WORLD)
139
+ check_all(matrix_free(A), reference(A))
140
+
141
+ @pytest.mark.parametrize('periods', [[False, False], [True, True]])
142
+ @pytest.mark.parallel
143
+ def test_toarray_block_parallel(periods):
144
+ A, B, N = block_case(periods, comm=MPI.COMM_WORLD)
145
+ refA, refB = reference(A), reference(B)
146
+ check_all(matrix_free(B), refB)
147
+
148
+ Z = xp.zeros((refB.shape[0], refA.shape[1]))
149
+ refN = xp.block([[refB, Z], [Z.T, refA]])
150
+ check_all(matrix_free(N), refN)
151
+
152
+ @pytest.mark.parallel
153
+ def test_toarray_composite_operators_parallel():
154
+ A = stencil_case([True, False], comm=MPI.COMM_WORLD)
155
+ ref = reference(A)
156
+ assert xp.allclose((A @ A).toarray(), ref @ ref)
157
+ assert xp.allclose((A ** 2).tosparse().toarray(), ref @ ref)
158
+
159
+ #===============================================================================
160
+ if __name__ == '__main__':
161
+ import sys
162
+ pytest.main(sys.argv)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: feectools
3
- Version: 0.1.10
3
+ Version: 0.2.0
4
4
  Summary: Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies.
5
5
  Author-email: Psydac development team <psydac@googlegroups.com>
6
6
  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>
@@ -100,6 +100,7 @@ feectools/linalg/tests/test_solvers.py
100
100
  feectools/linalg/tests/test_stencil_interface_matrix.py
101
101
  feectools/linalg/tests/test_stencil_vector.py
102
102
  feectools/linalg/tests/test_stencil_vector_space.py
103
+ feectools/linalg/tests/test_toarray.py
103
104
  feectools/linalg/tests/utilities.py
104
105
  feectools/utilities/__init__.py
105
106
  feectools/utilities/quadratures.py
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "feectools"
7
- version = "0.1.10"
7
+ version = "0.2.0"
8
8
  description = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies."
9
9
  readme = "README.md"
10
10
  requires-python = ">= 3.10"
@@ -1,114 +0,0 @@
1
- from dataclasses import dataclass
2
- from time import time
3
- from typing import TYPE_CHECKING
4
-
5
-
6
- # Might not be needed
7
- class MPICommWrapper:
8
- def __init__(self, use_mpi=True):
9
- self.use_mpi = use_mpi
10
- if use_mpi:
11
- from mpi4py import MPI
12
-
13
- self.comm = MPI.COMM_WORLD
14
- else:
15
- self.comm = MockComm()
16
-
17
- def __getattr__(self, name):
18
- return getattr(self.comm, name)
19
-
20
-
21
- class MockComm:
22
- def __getattr__(self, name):
23
- # Return a function that does nothing and returns None
24
- def dummy(*args, **kwargs):
25
- return None
26
-
27
- return dummy
28
-
29
- # Override some functions
30
- def Get_rank(self):
31
- return 0
32
-
33
- def Get_size(self):
34
- return 1
35
-
36
- def Barrier(self):
37
- return
38
-
39
-
40
- class MPIwrapper:
41
- def __init__(
42
- self,
43
- use_mpi: bool = False,
44
- verbose: bool = False,
45
- ):
46
- self.use_mpi = use_mpi
47
- if use_mpi:
48
- from mpi4py import MPI
49
-
50
- self._MPI = MPI
51
- if verbose:
52
- print("MPI is enabled")
53
- else:
54
- self._MPI = MockMPI()
55
- if verbose:
56
- print("MPI is NOT enabled")
57
-
58
- @property
59
- def MPI(self):
60
- return self._MPI
61
-
62
-
63
- class MockMPI:
64
- def __getattr__(self, name):
65
- # Return a function that does nothing and returns None
66
- def dummy(*args, **kwargs):
67
- return None
68
-
69
- return dummy
70
-
71
- # Override some functions
72
- @property
73
- def COMM_WORLD(self):
74
- return MockComm()
75
-
76
- # def comm_Get_rank(self):
77
- # return 0
78
-
79
- # def comm_Get_size(self):
80
- # return 1
81
-
82
-
83
- try:
84
- # Disable MPI when using CuPy due to known segfault issues with OpenMPI + CUDA
85
- import os
86
- if os.environ.get('ARRAY_BACKEND') == 'cupy':
87
- raise ImportError("MPI disabled when using CuPy backend")
88
-
89
- from mpi4py import MPI
90
-
91
- _comm = MPI.COMM_WORLD
92
- # rank = _comm.Get_rank()
93
- # size = _comm.Get_size()
94
- mpi_enabled = True
95
- except ImportError:
96
- # mpi4py not installed
97
- mpi_enabled = False
98
- except Exception:
99
- # mpi4py installed but not running under mpirun
100
- mpi_enabled = False
101
-
102
- # TODO: add environment variable for mpi use
103
- mpi_wrapper = MPIwrapper(
104
- use_mpi=mpi_enabled,
105
- verbose=False,
106
- )
107
-
108
- # TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime.
109
- if TYPE_CHECKING:
110
- from mpi4py import MPI
111
-
112
- mpi = MPI
113
- else:
114
- mpi = mpi_wrapper.MPI
File without changes
File without changes
File without changes
File without changes
File without changes