Initial MPIFractalBox implementation

gprMax/cython/fractals_generate.pyx

@@ -70,9 +70,14 @@ cpdef void generate_fractal3D(
     int nx,
     int ny,
     int nz,
+    int ox,
+    int oy,
+    int oz,
+    int gx,
+    int gy,
+    int gz,
     int nthreads,
     float D,
-    np.int32_t[:] global_size,
     np.float64_t[:] weighting,
     np.float64_t[:] v1,
     np.complex128_t[:, :, ::1] A,
@@ -97,17 +102,17 @@ cpdef void generate_fractal3D(
     cdef double v2x, v2y, v2z, rr, B
     cdef int sx, sy, sz
 
-    sx = global_size[0] // 2
-    sy = global_size[1] // 2
-    sz = global_size[2] // 2
+    sx = gx // 2
+    sy = gy // 2
+    sz = gz // 2
 
     for i in prange(nx, nogil=True, schedule='static', num_threads=nthreads):
         for j in range(ny):
             for k in range(nz):
                 # Positional vector for current position
-                v2x = ((weighting[0] * i) + sx) % global_size[0]
-                v2y = ((weighting[1] * j) + sy) % global_size[1]
-                v2z = ((weighting[2] * k) + sz) % global_size[2]
+                v2x = ((weighting[0] * (i + ox)) + sx) % gx
+                v2y = ((weighting[1] * (j + oy)) + sy) % gy
+                v2z = ((weighting[2] * (k + oz)) + sz) % gz
 
                 # Calulate norm of v2 - v1
                 rr = ((v2x - v1[0])**2 + (v2y - v1[1])**2 + (v2z - v1[2])**2)**(1/2)

gprMax/fractals.py

@@ -19,9 +19,14 @@
 from typing import Optional
 
 import numpy as np
+from mpi4py import MPI
+from mpi4py_fft import PFFT, DistArray, newDistArray
+from mpi4py_fft.pencil import Subcomm
 from scipy import fftpack
 
 import gprMax.config as config
+from gprMax.grid.mpi_grid import MPIGrid
+from gprMax.output_controllers.grid_view import MPIGridView
 
 from .cython.fractals_generate import generate_fractal2D, generate_fractal3D
 from .utilities.utilities import round_value
@@ -158,6 +163,11 @@ class FractalVolume:
 
         self.ID = None
         self.operatingonID = None
+        self.start = np.array([xs, ys, zs], dtype=np.int32)
+        self.stop = np.array([xf, yf, zf], dtype=np.int32)
+        self.size = self.stop - self.start
+        self.original_start = self.start.copy()
+        self.original_stop = self.stop.copy()
         self.xs = xs
         self.xf = xf
         self.ys = ys
@@ -183,7 +193,7 @@ class FractalVolume:
         self.nbins = 0
         self.fractalsurfaces = []
 
-    def generate_fractal_volume(self):
+    def generate_fractal_volume(self) -> bool:
         """Generate a 3D volume with a fractal distribution."""
 
         # Scale filter according to size of fractal volume
@@ -219,13 +229,21 @@ class FractalVolume:
         rng = np.random.default_rng(seed=self.seed)
         A = rng.standard_normal(size=(self.nx, self.ny, self.nz))
 
+        print(A)
+
         # 3D FFT
         A = fftpack.fftn(A)
-        # Shift the zero frequency component to the centre of the array
-        A = fftpack.fftshift(A)
+
+        print(A)
 
         # Generate fractal
         generate_fractal3D(
+            self.nx,
+            self.ny,
+            self.nz,
+            0,
+            0,
+            0,
             self.nx,
             self.ny,
             self.nz,
@@ -237,8 +255,8 @@ class FractalVolume:
             self.fractalvolume,
         )
 
-        # Shift the zero frequency component to the start of the array
-        self.fractalvolume = fftpack.ifftshift(self.fractalvolume)
+        print(self.fractalvolume)
+
         # Set DC component of FFT to zero
         self.fractalvolume[0, 0, 0] = 0
         # Take the real part (numerical errors can give rise to an imaginary part)
@@ -248,6 +266,11 @@ class FractalVolume:
             config.sim_config.dtypes["float_or_double"], copy=False
         )
 
+        print(self.fractalvolume)
+
+        print("min:", np.amin(self.fractalvolume))
+        print("max:", np.amax(self.fractalvolume))
+
         # Bin fractal values
         bins = np.linspace(np.amin(self.fractalvolume), np.amax(self.fractalvolume), self.nbins)
         for j in range(self.ny):
@@ -256,6 +279,10 @@ class FractalVolume:
                     self.fractalvolume[:, j, k], bins, right=True
                 )
 
+        print(self.fractalvolume)
+
+        return True
+
     def generate_volume_mask(self):
         """Generate a 3D volume to use as a mask for adding rough surfaces,
         water and grass/roots. Zero signifies the mask is not set, one
@@ -272,6 +299,218 @@ class FractalVolume:
         self.mask[maskxs:maskxf, maskys:maskyf, maskzs:maskzf] = 1
 
 
+class MPIFractalVolume(FractalVolume):
+    def __init__(
+        self,
+        xs: int,
+        xf: int,
+        ys: int,
+        yf: int,
+        zs: int,
+        zf: int,
+        dimension: float,
+        seed: Optional[int],
+        grid: MPIGrid,
+    ):
+        super().__init__(xs, xf, ys, yf, zs, zf, dimension, seed)
+        self.grid = grid
+
+    def generate_fractal_volume(self) -> bool:
+        """Generate a 3D volume with a fractal distribution."""
+
+        if self.grid.local_bounds_overlap_grid(self.start, self.stop):
+            grid_view = MPIGridView(self.grid, self.xs, self.ys, self.zs, self.xf, self.yf, self.zf)
+        else:
+            # The MPIGridView will create a new communicator using
+            # MPI_Split. Calling this here prevents deadlock if not all
+            # ranks need to generate the fractal volume.
+            self.grid.comm.Split(MPI.UNDEFINED)
+            return False
+
+        # Scale filter according to size of fractal volume
+        sorted_size = np.sort(grid_view.global_size)
+        min_size = sorted_size[1] if sorted_size[0] == 1 else sorted_size[0]
+        filterscaling = np.where(grid_view.global_size == 1, 1, min_size / grid_view.global_size)
+
+        # Adjust weighting to account for filter scaling
+        self.weighting = np.multiply(self.weighting, filterscaling)
+
+        # Positional vector at centre of array, scaled by weighting
+        v1 = self.weighting * grid_view.global_size / 2
+
+        # 3D array of random numbers to be convolved with the fractal function
+        rng = np.random.default_rng(seed=self.seed)
+
+        # print(type(grid_view.comm))
+
+        # fft = PFFT(
+        #     grid_view.comm,
+        #     shape=tuple(grid_view.global_size),
+        #     axes=(0, 1, 2),
+        #     dtype=self.dtype,
+        #     collapse=False,
+        #     backend="fftw",
+        # )
+        # A = newDistArray(fft, False)
+        if all([dim > 1 for dim in grid_view.comm.dims]):
+            raise ValueError(
+                "Fractal volume must be positioned such that its MPI decomposition is 1 in at least"
+                f" 1 dimension. Current decompostion is: {grid_view.comm.dims}"
+            )
+
+        comm = Subcomm(grid_view.comm)
+        A = DistArray(grid_view.global_size, comm, dtype=self.dtype)
+        fft = PFFT(
+            None,
+            axes=tuple(np.argsort(grid_view.comm.dims)[::-1]),
+            darray=A,
+            collapse=False,
+            backend="fftw",
+        )
+
+        print(A.alignment)
+
+        # Decomposition of A may be different to the MPIGrid and GridView
+        start = np.array(A.substart)
+        shape = np.array(A.shape)
+
+        for index in np.ndindex(*A.global_shape):
+            index = np.array(index)
+            if any(index < start) or any(index >= start + shape):
+                rng.standard_normal()
+            else:
+                index -= start
+                A[index[0], index[1], index[2]] = rng.standard_normal()
+
+        print(A.shape)
+        print(A.global_shape)
+        print(A.dtype)
+
+        print(A)
+
+        fft.forward()
+
+        print(A.shape)
+        print(A.global_shape)
+        print(A.dtype)
+
+        print(A)
+
+        A_hat = newDistArray(fft)
+
+        # 3D FFT
+        fft.forward(A, A_hat, normalize=False)
+        A_hat_out = np.zeros_like(A_hat)
+
+        # A = fftpack.fftn(A)
+        # Shift the zero frequency component to the centre of the array
+        # A = fftpack.fftshift(A)
+
+        print(A_hat)
+
+        print("Calling generate_fractal3D")
+        print(A_hat.shape)
+        print(A_hat.dtype)
+
+        # Generate fractal
+        generate_fractal3D(
+            A_hat.shape[0],
+            A_hat.shape[1],
+            A_hat.shape[2],
+            A_hat.substart[0],
+            A_hat.substart[1],
+            A_hat.substart[2],
+            A_hat.global_shape[0],
+            A_hat.global_shape[1],
+            A_hat.global_shape[2],
+            config.get_model_config().ompthreads,
+            self.dimension,
+            self.weighting,
+            v1,
+            A_hat,
+            A_hat_out,
+        )
+
+        print("Returned generate_fractal3D")
+
+        print(A_hat_out)
+
+        print(A_hat_out.shape)
+        print(A_hat_out.dtype)
+
+        # print("A_hat_out -> A_hat")
+
+        # print(A.shape)
+        # print(out.shape)
+        # print(A.dtype)
+        # print(out.dtype)
+        # print(A.size)
+        # print(out.size)
+
+        # A_hat[:] = A_hat_out
+
+        # print(A.shape)
+        # print(A.global_shape)
+
+        # Shift the zero frequency component to the start of the array
+        # self.fractalvolume = fftpack.ifftshift(self.fractalvolume)
+        # Set DC component of FFT to zero
+        if all(start == 0):
+            A_hat_out[0, 0, 0] = 0
+
+        print("Backward transform")
+
+        # Take the real part (numerical errors can give rise to an imaginary part)
+        # of the IFFT, and convert type to floattype. N.B calculation of fractals
+        # must always be carried out at double precision, i.e. float64, complex128
+        Aj = np.real(fft.backward(A_hat_out, normalize=True)).astype(
+            config.sim_config.dtypes["float_or_double"], copy=False
+        )
+
+        print(Aj)
+
+        print(Aj.shape)
+        print(Aj.dtype)
+
+        min_value = np.array(np.amin(Aj), dtype=config.sim_config.dtypes["float_or_double"])
+        max_value = np.array(np.amax(Aj), dtype=config.sim_config.dtypes["float_or_double"])
+        print(f"min: {min_value}")
+        print(f"max: {max_value}")
+
+        grid_view.comm.Allreduce(MPI.IN_PLACE, [min_value, MPI.FLOAT], MPI.MIN)
+        print(f"min: {min_value}")
+        grid_view.comm.Allreduce(MPI.IN_PLACE, [max_value, MPI.FLOAT], MPI.MAX)
+        print(f"max: {max_value}")
+
+        self.fractalvolume = np.zeros_like(Aj)
+
+        # Bin fractal values
+        bins = np.linspace(min_value, max_value, self.nbins)
+        for j in range(shape[1]):
+            for k in range(shape[2]):
+                self.fractalvolume[:, j, k] = np.digitize(Aj[:, j, k], bins, right=True)
+
+        # self.fractalvolume = A.copy()  # np.zeros(grid_view.size, dtype=self.dtype)
+
+        print(self.fractalvolume.shape)
+        print(self.fractalvolume.dtype)
+
+        # self.fractalvolume = A.copy()  # np.zeros(grid_view.size, dtype=self.dtype)
+        print(self.fractalvolume)
+
+        print("Complete")
+
+        self.nx = self.fractalvolume.shape[0]
+        self.ny = self.fractalvolume.shape[1]
+        self.nz = self.fractalvolume.shape[2]
+
+        self.xs = 1 if self.xs < 0 else self.xs
+        self.ys = 1 if self.ys < 0 else self.ys
+        self.zs = 1 if self.zs < 0 else self.zs
+
+        return True
+
+
 class Grass:
     """Geometry information for blades of grass."""
 

gprMax/grid/fdtd_grid.py

@@ -32,7 +32,6 @@ from typing_extensions import TypeVar
 from gprMax import config
 from gprMax.cython.pml_build import pml_average_er_mr
 from gprMax.cython.yee_cell_build import build_electric_components, build_magnetic_components
-from gprMax.fractals import FractalVolume
 from gprMax.materials import ListMaterial, Material, PeplinskiSoil, RangeMaterial, process_materials
 from gprMax.pml import CFS, PML, print_pml_info
 from gprMax.receivers import Rx
@@ -111,7 +110,7 @@ class FDTDGrid:
         # Materials used by this grid
         self.materials: List[Material] = []
         self.mixingmodels: List[Union[PeplinskiSoil, RangeMaterial, ListMaterial]] = []
-        self.fractalvolumes: List[FractalVolume] = []
+        self.fractalvolumes = []  # List[FractalVolume]
 
         # Sources and receivers contained inside this grid
         self.waveforms: List[Waveform] = []

gprMax/user_objects/cmds_geometry/fractal_box.py

@@ -22,8 +22,9 @@ import numpy as np
 
 import gprMax.config as config
 from gprMax.cython.geometry_primitives import build_voxels_from_array, build_voxels_from_array_mask
-from gprMax.fractals import FractalVolume
+from gprMax.fractals import FractalVolume, MPIFractalVolume
 from gprMax.grid.fdtd_grid import FDTDGrid
+from gprMax.grid.mpi_grid import MPIGrid
 from gprMax.materials import ListMaterial
 from gprMax.user_objects.cmds_geometry.cmds_geometry import check_averaging, rotate_2point_object
 from gprMax.user_objects.rotatable import RotatableMixin
@@ -110,12 +111,7 @@ class FractalBox(RotatableMixin, GeometryUserObject):
         p3 = uip.round_to_grid_static_point(p1)
         p4 = uip.round_to_grid_static_point(p2)
 
-        grid_contains_fractal_box, p1, p2 = uip.check_box_points(p1, p2, self.__str__())
-
-        # Exit early if none of the fractal box is in this grid as there
-        # is nothing else to do.
-        if not grid_contains_fractal_box:
-            return
+        p1, p2 = uip.check_output_object_bounds(p1, p2, self.__str__())
 
         xs, ys, zs = p1
         xf, yf, zf = p2
@@ -172,7 +168,10 @@ class FractalBox(RotatableMixin, GeometryUserObject):
             )
             raise ValueError
 
-        self.volume = FractalVolume(xs, xf, ys, yf, zs, zf, frac_dim, seed)
+        if isinstance(grid, MPIGrid):
+            self.volume = MPIFractalVolume(xs, xf, ys, yf, zs, zf, frac_dim, seed, grid)
+        else:
+            self.volume = FractalVolume(xs, xf, ys, yf, zs, zf, frac_dim, seed)
         self.volume.ID = ID
         self.volume.operatingonID = mixing_model_id
         self.volume.nbins = nbins
@@ -685,7 +684,8 @@ class FractalBox(RotatableMixin, GeometryUserObject):
                     )
                     raise ValueError
                 else:
-                    self.volume.generate_fractal_volume()
+                    if not self.volume.generate_fractal_volume():
+                        return
                     for i in range(0, self.volume.nx):
                         for j in range(0, self.volume.ny):
                             for k in range(0, self.volume.nz):

requirements.txt

@@ -9,6 +9,7 @@ jinja2
 jupyter
 matplotlib
 # mpi4py
+mpi4py-fft
 numpy<2
 numpy-stl
 pandas