First plugin fluidfft-mpi4pyfft (untested) + fft3d/testing.py

plugins/README.md

@@ -5,7 +5,7 @@ Directory containing the plugins, i.e. Python packages declaring the
 
 We should have
 
-- [ ] fluidfft-mpi4pyfft
+- [x] fluidfft-mpi4pyfft
 - [ ] fluidfft-fftw
 - [ ] fluidfft-mpi_with_fftw
 - [ ] fluidfft-fftwmpi

plugins/fluidfft-mpi4pyfft/LICENSE

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2024 Pierre Augier
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

plugins/fluidfft-mpi4pyfft/pyproject.toml

@@ -0,0 +1,20 @@
+[build-system]
+requires = ["flit_core >=3.2,<4"]
+build-backend = "flit_core.buildapi"
+
+[project]
+name = "fluidfft_mpi4pyfft"
+version = "0.0.1"
+description = "Fluidfft plugin using mpi4pyfft"
+authors = [{name = "Pierre Augier", email = "pierre.augier@univ-grenoble-alpes.fr"}]
+license = {file = "LICENSE"}
+classifiers = ["License :: OSI Approved :: MIT License"]
+dependencies = ["fluidfft", "mpi4py-fft"]
+
+[project.urls]
+Home = "https://fluidfft.readthedocs.io"
+
+[project.entry-points."fluidfft.plugins"]
+
+"fft3d.mpi_with_mpi4pyfft" = "fluidfft_mpi4pyfft.mpi_with_mpi4pyfft"
+"fft3d.mpi_with_mpi4pyfft_slab" = "fluidfft_mpi4pyfft.fft3d.mpi_with_mpi4pyfft_slab"

src/fluidfft/fft3d/mpi_with_mpi4pyfft.py → plugins/fluidfft-mpi4pyfft/src/fluidfft_mpi4pyfft/mpi_with_mpi4pyfft.py

@@ -3,7 +3,7 @@
 from mpi4py import MPI
 from mpi4py_fft import PFFT, newDistArray
 
-from .base import BaseFFTMPI
+from fluidfft.fft3d.base import BaseFFTMPI
 
 
 class FFT3DMPIWithMPI4PYFFT(BaseFFTMPI):

plugins/fluidfft-mpi4pyfft/tests/test_with_mpi4pyfft.py

@@ -0,0 +1,22 @@
+from fluidfft.fft3d.testing import make_testop_functions
+
+
+from fluidfft import import_fft_class
+
+
+class Tests3D(unittest.TestCase):
+    pass
+
+
+def complete_class(name, cls):
+    tests = make_testop_functions(name, cls)
+
+    for key, test in tests.items():
+        setattr(Tests3D, "test_operator3d_{}_{}".format(name, key), test)
+
+
+methods = ["fft3d.mpi_with_mpi4pyfft", "fft3d.mpi_with_mpi4pyfft_slab"]
+for method in methods:
+    name = method.split(".")[1]
+    cls = import_fft_class(method)
+    complete_class(name, cls)

pyproject.toml

@@ -84,17 +84,17 @@ fluidfft-bench-analysis = "fluidfft.bench_analysis:run"
 # should be in fluidfft-pfft
 "fft3d.mpi_with_pfft" = "fluidfft.fft3d.mpi_with_pfft"
 
-# should be in fluidfft-mpi4pyfft (a pure Python package)
-"fft3d.mpi_with_mpi4pyfft" = "fluidfft.fft3d.mpi_with_mpi4pyfft"
-"fft3d.mpi_with_mpi4pyfft_slab" = "fluidfft.fft3d.mpi_with_mpi4pyfft_slab"
-
 [tool.pdm]
 distribution = true
 package-dir = "src"
 
 [tool.pdm.dev-dependencies]
 build = ["setuptools", "transonic", "pythran", "wheel", "jinja2", "cython"]
 
+# plugins = [
+#     "-e fluidfft-mpi4pyfft @ file:///${PROJECT_ROOT}/plugins/fluidfft-mpi4pyfft",
+# ]
+
 test = [
     "pytest",
     "coverage",
@@ -118,7 +118,7 @@ doc = [
 lint = ["black", "pylint"]
 
 [tool.pdm.scripts]
-black = 'black -l 82 src doc src_cy tests --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/"'
-lint = {shell="pylint -rn --rcfile=pylintrc --jobs=$(nproc) src doc tests --exit-zero"}
-black_check = 'black --check -l 82 src doc src_cy tests --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/"'
+black = 'black -l 82 src doc src_cy tests plugins --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/"'
+lint = {shell="pylint -rn --rcfile=pylintrc --jobs=$(nproc) src doc tests plugins --exit-zero"}
+black_check = 'black --check -l 82 src doc src_cy tests plugins --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/"'
 validate_code = {composite = ["black_check", "lint"]}

src/fluidfft/fft3d/__init__.py

@@ -35,6 +35,8 @@ class :class:`fluidfft.fft3d.operators.OperatorsPseudoSpectral3D` defined in
 
 """
 
+import sys
+
 from .. import import_fft_class
 
 __all__ = [
@@ -58,8 +60,6 @@ class :class:`fluidfft.fft3d.operators.OperatorsPseudoSpectral3D` defined in
     "fftwmpi3d",
     "p3dfft",
     "pfft",
-    "mpi4pyfft",
-    "mpi4pyfft_slab",
 ]
 methods_mpi = ["fft3d.mpi_with_" + method for method in methods_mpi]
 
@@ -78,3 +78,9 @@ def get_classes_mpi():
         method: import_fft_class(method, raise_import_error=False)
         for method in methods_mpi
     }
+
+
+if any("pytest" in part for part in sys.argv):
+    import pytest
+
+    pytest.register_assert_rewrite("fluidfft.fft3d.testing")

src/fluidfft/fft3d/testing.py

@@ -0,0 +1,141 @@
+import numpy as np
+
+from .operators import OperatorsPseudoSpectral3D, vector_product
+
+from fluiddyn.util import mpi
+
+rank = mpi.rank
+nb_proc = mpi.nb_proc
+
+
+def make_testop_functions(name, cls):
+    tests = {}
+    shapes = {"even": (4, 8, 12)}
+    if nb_proc == 1:
+        shapes["odd"] = (5, 3, 7)
+
+    for key, (n0, n1, n2) in shapes.items():
+
+        def test(self, n0=n0, n1=n1, n2=n2):
+            try:
+                op = OperatorsPseudoSpectral3D(n2, n1, n0, 12, 8, 4, fft=cls)
+            except ValueError:
+                print(
+                    "ValueError while instantiating OperatorsPseudoSpectral3D"
+                    " for {}".format(cls)
+                )
+                return
+
+            op_fft = op._op_fft
+
+            op_fft.run_tests()
+
+            a = np.random.random(op_fft.get_local_size_X()).reshape(
+                op_fft.get_shapeX_loc()
+            )
+            a0 = a.copy()
+            afft = op.fft3d(a)
+            self.assertTrue(np.allclose(a, a0))
+            afft0 = afft.copy()
+            a = op.ifft3d(afft)
+            self.assertTrue(np.allclose(afft, afft0))
+            afft = op.fft3d(a)
+
+            nrja = op.compute_energy_from_X(a)
+            nrjafft = op.compute_energy_from_K(afft)
+            self.assertAlmostEqual(nrja, nrjafft)
+
+            energy_fft = 0.5 * abs(afft) ** 2
+            nrj = op.sum_wavenumbers(energy_fft)
+            self.assertAlmostEqual(nrjafft, nrj)
+
+            try:
+                nrj_versatile = op.sum_wavenumbers_versatile(energy_fft)
+            except NotImplementedError:
+                pass
+            else:
+                self.assertAlmostEqual(nrj_versatile, nrj)
+
+            try:
+                E_kx, E_ky, E_kz = op.compute_1dspectra(energy_fft)
+            except NotImplementedError:
+                pass
+            else:
+                self.assertAlmostEqual(nrj, E_kx.sum() * op.deltakx)
+                self.assertAlmostEqual(nrj, E_ky.sum() * op.deltaky)
+                self.assertAlmostEqual(nrj, E_kz.sum() * op.deltakz)
+
+                self.assertEqual(E_kx.shape[0], op.nkx_spectra)
+                self.assertEqual(E_ky.shape[0], op.nky_spectra)
+                self.assertEqual(E_kz.shape[0], op.nkz_spectra)
+
+            E_k = op.compute_3dspectrum(energy_fft)
+            self.assertAlmostEqual(nrja, E_k.sum() * op.deltak_spectra3d)
+
+            E_kz_kh = op.compute_spectrum_kzkh(energy_fft)
+            self.assertAlmostEqual(nrja, E_kz_kh.sum() * op.deltakh * op.deltakz)
+
+            try:
+                E_kx_kyz, E_ky_kzx, E_kz_kxy = op.compute_spectra_2vars(
+                    energy_fft
+                )
+            except NotImplementedError:
+                pass
+            else:
+                self.assertAlmostEqual(
+                    E_kx_kyz.sum() * op.deltakx, E_ky_kzx.sum() * op.deltaky
+                )
+                self.assertAlmostEqual(E_kz_kxy.sum() * op.deltakz, nrja)
+
+            op.produce_str_describing_grid()
+            op.produce_str_describing_oper()
+            op.produce_long_str_describing_oper()
+            op.create_arrayX(value=None, shape="loc")
+            op.create_arrayX(value=None, shape="seq")
+            op.create_arrayX(value=0.0)
+            op.create_arrayK(value=1.0)
+            op.create_arrayX_random(max_val=2)
+            op.create_arrayK_random(min_val=-1, max_val=1, shape="seq")
+
+            op.project_perpk3d(afft, afft, afft)
+            op.divfft_from_vecfft(afft, afft, afft)
+            op.rotfft_from_vecfft(afft, afft, afft)
+            op.rotfft_from_vecfft_outin(afft, afft, afft, afft, afft, afft)
+            op.rotzfft_from_vxvyfft(afft, afft)
+
+            # depreciated...
+            # op.vgradv_from_v(a, a, a)
+            # op.vgradv_from_v2(a, a, a)
+            # op.div_vv_fft_from_v(a, a, a)
+            op.div_vb_fft_from_vb(a, a, a, a)
+            vector_product(a, a, a, a, a, a)
+
+            X, Y, Z = op.get_XYZ_loc()
+            self.assertEqual(X.shape, op.shapeX_loc)
+            self.assertEqual(Y.shape, op.shapeX_loc)
+            self.assertEqual(Z.shape, op.shapeX_loc)
+
+            X = np.ascontiguousarray(X)
+            Y = np.ascontiguousarray(Y)
+            Z = np.ascontiguousarray(Z)
+            root = 0
+            X_seq = op.gather_Xspace(X, root=root)
+            Y_seq = op.gather_Xspace(Y, root=root)
+            Z_seq = op.gather_Xspace(Z, root=root)
+
+            if rank == root:
+                self.assertTrue(np.allclose(X_seq[0, 0, :], op.x_seq))
+                self.assertTrue(np.allclose(Y_seq[0, :, 0], op.y_seq))
+                self.assertTrue(np.allclose(Z_seq[:, 0, 0], op.z_seq))
+
+            X_scatter = op.scatter_Xspace(X_seq, root=root)
+            Y_scatter = op.scatter_Xspace(Y_seq, root=root)
+            Z_scatter = op.scatter_Xspace(Z_seq, root=root)
+
+            self.assertTrue(np.allclose(X, X_scatter))
+            self.assertTrue(np.allclose(Y, Y_scatter))
+            self.assertTrue(np.allclose(Z, Z_scatter))
+
+        tests[key] = test
+
+    return tests