Fix all MPI ranks should use the same sigma max

The CFS parameter sigma max values for PMLs should be the same for all
PMLs. Therefore it needs to be broadcast to all MPI ranks.

box_half_model.in

@@ -1,11 +0,0 @@
-#title: Hertzian dipole over a half-space
-#domain: 0.010 0.010 0.001
-#dx_dy_dz: 0.001 0.001 0.001
-#time_window: 3e-9
-
-#waveform: gaussiandot 1 1e9 myWave
-#hertzian_dipole: z 0.005 0.005 0.001 myWave
-#rx: 0.007 0.007 0.001
-
-#material: 8 0 1 0 half_space
-#box: 0 0 0 0.010 0.005 0.001 half_space

gprMax/grid/cuda_grid.py

@@ -21,6 +21,7 @@ from importlib import import_module
 import numpy as np
 
 from gprMax.grid.fdtd_grid import FDTDGrid
+from gprMax.pml import CUDAPML
 
 
 class CUDAGrid(FDTDGrid):
@@ -36,12 +37,19 @@ class CUDAGrid(FDTDGrid):
         # Blocks per grid - used for main electric/magnetic field updates
         self.bpg = None
 
+    def _construct_pml(self, pml_ID: str, thickness: int) -> CUDAPML:
+        return super()._construct_pml(pml_ID, thickness, CUDAPML)
+
     def set_blocks_per_grid(self):
         """Set the blocks per grid size used for updating the electric and
         magnetic field arrays on a GPU.
         """
 
-        self.bpg = (int(np.ceil(((self.nx + 1) * (self.ny + 1) * (self.nz + 1)) / self.tpb[0])), 1, 1)
+        self.bpg = (
+            int(np.ceil(((self.nx + 1) * (self.ny + 1) * (self.nz + 1)) / self.tpb[0])),
+            1,
+            1,
+        )
 
     def htod_geometry_arrays(self):
         """Initialise an array for cell edge IDs (ID) on compute device."""

gprMax/grid/fdtd_grid.py

@@ -21,17 +21,19 @@ import itertools
 import logging
 import sys
 from collections import OrderedDict
-from typing import Any, Iterable, List, Union
+from typing import Any, Iterable, List, Tuple, Union
 
 import numpy as np
 from terminaltables import SingleTable
 from tqdm import tqdm
+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, build_pml, print_pml_info
+from gprMax.pml import CFS, PML, print_pml_info
 from gprMax.receivers import Rx
 from gprMax.snapshots import Snapshot
 from gprMax.sources import HertzianDipole, MagneticDipole, Source, TransmissionLine, VoltageSource
@@ -152,7 +154,7 @@ class FDTDGrid:
             self.initialise_dispersive_update_coeff_array()
         self._build_materials()
 
-    def _build_pmls(self, build_pml_func=build_pml) -> None:
+    def _build_pmls(self) -> None:
         pbar = tqdm(
             total=sum(1 for value in self.pmls["thickness"].values() if value > 0),
             desc=f"Building PML boundaries [{self.name}]",
@@ -162,12 +164,127 @@ class FDTDGrid:
         )
         for pml_id, thickness in self.pmls["thickness"].items():
             if thickness > 0:
-                # TODO: Consider making this a method of FDTDGrid (that
-                # can be overriden in MPIGrid)
-                build_pml_func(self, pml_id, thickness)
+                pml = self._construct_pml(pml_id, thickness)
+                averageer, averagemr = self._calculate_average_pml_material_properties(pml)
+                pml.calculate_update_coeffs(averageer, averagemr)
+                self.pmls["slabs"].append(pml)
                 pbar.update()
         pbar.close()
 
+    PmlType = TypeVar("PmlType", bound=PML)
+
+    def _construct_pml(self, pml_ID: str, thickness: int, pml_type: type[PmlType] = PML) -> PmlType:
+        """Builds instances of the PML and calculates the initial parameters and
+            coefficients including setting profile (based on underlying material
+            er and mr from solid array).
+
+        Args:
+            G: FDTDGrid class describing a grid in a model.
+            pml_ID: string identifier of PML slab.
+            thickness: int with thickness of PML slab in cells.
+        """
+        if pml_ID == "x0":
+            pml = pml_type(
+                self,
+                ID=pml_ID,
+                direction="xminus",
+                xs=0,
+                xf=thickness,
+                ys=0,
+                yf=self.ny,
+                zs=0,
+                zf=self.nz,
+            )
+        elif pml_ID == "xmax":
+            pml = pml_type(
+                self,
+                ID=pml_ID,
+                direction="xplus",
+                xs=self.nx - thickness,
+                xf=self.nx,
+                ys=0,
+                yf=self.ny,
+                zs=0,
+                zf=self.nz,
+            )
+        elif pml_ID == "y0":
+            pml = pml_type(
+                self,
+                ID=pml_ID,
+                direction="yminus",
+                xs=0,
+                xf=self.nx,
+                ys=0,
+                yf=thickness,
+                zs=0,
+                zf=self.nz,
+            )
+        elif pml_ID == "ymax":
+            pml = pml_type(
+                self,
+                ID=pml_ID,
+                direction="yplus",
+                xs=0,
+                xf=self.nx,
+                ys=self.ny - thickness,
+                yf=self.ny,
+                zs=0,
+                zf=self.nz,
+            )
+        elif pml_ID == "z0":
+            pml = pml_type(
+                self,
+                ID=pml_ID,
+                direction="zminus",
+                xs=0,
+                xf=self.nx,
+                ys=0,
+                yf=self.ny,
+                zs=0,
+                zf=thickness,
+            )
+        elif pml_ID == "zmax":
+            pml = pml_type(
+                self,
+                ID=pml_ID,
+                direction="zplus",
+                xs=0,
+                xf=self.nx,
+                ys=0,
+                yf=self.ny,
+                zs=self.nz - thickness,
+                zf=self.nz,
+            )
+        else:
+            raise ValueError(f"Unknown PML ID '{pml_ID}'")
+
+        return pml
+
+    def _calculate_average_pml_material_properties(self, pml: PML) -> Tuple[float, float]:
+        # Arrays to hold values of permittivity and permeability (avoids accessing
+        # Material class in Cython.)
+        ers = np.zeros(len(self.materials))
+        mrs = np.zeros(len(self.materials))
+
+        for i, m in enumerate(self.materials):
+            ers[i] = m.er
+            mrs[i] = m.mr
+
+        if pml.ID[0] == "x":
+            n1 = self.ny
+            n2 = self.nz
+            solid = self.solid[pml.xs, :, :]
+        elif pml.ID[0] == "y":
+            n1 = self.nx
+            n2 = self.nz
+            solid = self.solid[:, pml.ys, :]
+        elif pml.ID[0] == "z":
+            n1 = self.nx
+            n2 = self.ny
+            solid = self.solid[:, :, pml.zs]
+
+        return pml_average_er_mr(n1, n2, config.get_model_config().ompthreads, solid, ers, mrs)
+
     def _build_components(self) -> None:
         # Build the model, i.e. set the material properties (ID) for every edge
         # of every Yee cell

gprMax/grid/mpi_grid.py

@@ -19,15 +19,17 @@
 import itertools
 import logging
 from enum import IntEnum, unique
-from typing import List, Optional, TypeVar, Union
+from typing import List, Optional, Tuple, TypeVar, Union
 
 import numpy as np
 import numpy.typing as npt
 from mpi4py import MPI
 from numpy import ndarray
 
+from gprMax import config
+from gprMax.cython.pml_build import pml_sum_er_mr
 from gprMax.grid.fdtd_grid import FDTDGrid
-from gprMax.pml import build_pml_mpi
+from gprMax.pml import MPIPML, PML
 from gprMax.receivers import Rx
 from gprMax.sources import Source
 
@@ -62,12 +64,13 @@ class MPIGrid(FDTDGrid):
         self.x_comm = comm.Sub([False, True, True])
         self.y_comm = comm.Sub([True, False, True])
         self.z_comm = comm.Sub([True, True, False])
+        self.pml_comm = MPI.COMM_NULL
 
         self.mpi_tasks = np.array(self.comm.dims)
 
         self.lower_extent: npt.NDArray[np.intc] = np.zeros(3, dtype=int)
         self.upper_extent: npt.NDArray[np.intc] = np.zeros(3, dtype=int)
-        self.negative_halo_offset: npt.NDArray[np.intc] = np.zeros(3, dtype=int)
+        self.negative_halo_offset: npt.NDArray[np.bool_] = np.zeros(3, dtype=int)
         self.global_size: npt.NDArray[np.intc] = np.zeros(3, dtype=int)
 
         self.neighbours = np.full((3, 2), -1, dtype=int)
@@ -262,8 +265,58 @@ class MPIGrid(FDTDGrid):
         self._halo_swap_array(self.Hy)
         self._halo_swap_array(self.Hz)
 
-    def _build_pmls(self):
-        super()._build_pmls(build_pml_func=build_pml_mpi)
+    def _construct_pml(self, pml_ID: str, thickness: int) -> MPIPML:
+        pml = super()._construct_pml(pml_ID, thickness, MPIPML)
+        if pml.ID[0] == "x":
+            pml.comm = self.x_comm
+        elif pml.ID[0] == "y":
+            pml.comm = self.y_comm
+        elif pml.ID[0] == "z":
+            pml.comm = self.z_comm
+        pml.global_comm = self.pml_comm
+
+        return pml
+
+    def _calculate_average_pml_material_properties(self, pml: MPIPML) -> Tuple[float, float]:
+        # Arrays to hold values of permittivity and permeability (avoids
+        # accessing Material class in Cython.)
+        ers = np.zeros(len(self.materials))
+        mrs = np.zeros(len(self.materials))
+
+        for i, m in enumerate(self.materials):
+            ers[i] = m.er
+            mrs[i] = m.mr
+
+        # Need to account for the negative halo (remove it) to avoid
+        # double counting. The solid array does not have a positive halo
+        # so we don't need to consider that.
+        if pml.ID[0] == "x":
+            o1 = self.negative_halo_offset[1]
+            o2 = self.negative_halo_offset[2]
+            n1 = self.ny - o1
+            n2 = self.nz - o2
+            solid = self.solid[pml.xs, o1:, o2:]
+        elif pml.ID[0] == "y":
+            o1 = self.negative_halo_offset[0]
+            o2 = self.negative_halo_offset[2]
+            n1 = self.nx - o1
+            n2 = self.nz - o2
+            solid = self.solid[o1:, pml.ys, o2:]
+        elif pml.ID[0] == "z":
+            o1 = self.negative_halo_offset[0]
+            o2 = self.negative_halo_offset[1]
+            n1 = self.nx - o1
+            n2 = self.ny - o2
+            solid = self.solid[o1:, o2:, pml.zs]
+        else:
+            raise ValueError(f"Unknown PML ID '{pml.ID}'")
+
+        sumer, summr = pml_sum_er_mr(n1, n2, config.get_model_config().ompthreads, solid, ers, mrs)
+        n = pml.comm.allreduce(n1 * n2, MPI.SUM)
+        averageer = pml.comm.allreduce(sumer, MPI.SUM) / n
+        averagemr = pml.comm.allreduce(summr, MPI.SUM) / n
+
+        return averageer, averagemr
 
     def build(self):
         if any(self.global_size + 1 < self.mpi_tasks):
@@ -276,6 +329,20 @@ class MPIGrid(FDTDGrid):
         self.set_halo_map()
         self.scatter_grid()
 
+        # TODO: Check PML is not thicker than the grid size
+
+        # Get PMLs present in this grid
+        pmls = [
+            PML.boundaryIDs.index(key) for key, value in self.pmls["thickness"].items() if value > 0
+        ]
+        if len(pmls) > 0:
+            # Use PML ID as the key to ensure rank 0 is always the same
+            # PML. This is needed to ensure the CFS sigma.max parameter
+            # is calculated using the first PML present.
+            self.pml_comm = self.comm.Split(0, pmls[0])
+        else:
+            self.pml_comm = self.comm.Split(MPI.UNDEFINED)
+
         super().build()
 
     def has_neighbour(self, dim: Dim, dir: Dir) -> bool:

gprMax/grid/opencl_grid.py

@@ -19,6 +19,7 @@
 from importlib import import_module
 
 from gprMax.grid.fdtd_grid import FDTDGrid
+from gprMax.pml import PML, OpenCLPML
 
 
 class OpenCLGrid(FDTDGrid):
@@ -29,6 +30,9 @@ class OpenCLGrid(FDTDGrid):
 
         self.clarray = import_module("pyopencl.array")
 
+    def _construct_pml(self, pml_ID: str, thickness: int) -> OpenCLPML:
+        return super()._construct_pml(pml_ID, thickness, OpenCLPML)
+
     def htod_geometry_arrays(self, queue):
         """Initialise an array for cell edge IDs (ID) on compute device.
 

gprMax/pml.py

@@ -91,7 +91,7 @@ class CFS:
         self.kappa = CFSParameter(ID="kappa", scalingprofile="constant", min=1, max=1)
         self.sigma = CFSParameter(ID="sigma", scalingprofile="quartic", min=0, max=None)
 
-    def calculate_sigmamax(self, d, er, mr, G):
+    def calculate_sigmamax(self, d, er, mr):
         """Calculates an optimum value for sigma max based on underlying
             material properties.
 
@@ -205,7 +205,7 @@ class PML:
     #                       x-axis, y-axis, or z-axis
     directions = ["xminus", "yminus", "zminus", "xplus", "yplus", "zplus"]
 
-    def __init__(self, G, ID=None, direction=None, xs=0, xf=0, ys=0, yf=0, zs=0, zf=0):
+    def __init__(self, G, ID: str, direction: str, xs=0, xf=0, ys=0, yf=0, zs=0, zf=0):
         """
         Args:
             G: FDTDGrid class describing a grid in a model.
@@ -345,7 +345,7 @@ class PML:
 
         for x, cfs in enumerate(self.CFS):
             if not cfs.sigma.max:
-                cfs.calculate_sigmamax(self.d, er, mr, self.G)
+                cfs.calculate_sigmamax(self.d, er, mr)
             Ealpha, Halpha = cfs.calculate_values(self.thickness, cfs.alpha)
             Ekappa, Hkappa = cfs.calculate_values(self.thickness, cfs.kappa)
             Esigma, Hsigma = cfs.calculate_values(self.thickness, cfs.sigma)
@@ -682,6 +682,39 @@ class OpenCLPML(PML):
         event.wait()
 
 
+class MPIPML(PML):
+    comm: MPI.Cartcomm
+    global_comm: MPI.Comm
+
+    COORDINATOR_RANK = 0
+
+    def calculate_update_coeffs(self, er: float, mr: float):
+        """Calculates electric and magnetic update coefficients for the PML.
+
+        Args:
+            er: float of average permittivity of underlying material
+            mr: float of average permeability of underlying material
+        """
+        for cfs in self.CFS:
+            if not cfs.sigma.max:
+                if self.global_comm.rank == self.COORDINATOR_RANK:
+                    cfs.calculate_sigmamax(self.d, er, mr)
+                    buffer = np.array([cfs.sigma.max])
+                else:
+                    buffer = np.empty(1)
+
+                # Needs to be non-blocking because some ranks will
+                # contain multiple PMLs, but the material properties for
+                # a PML cannot be calculated until all ranks have
+                # completed that stage. Therefore a blocking broadcast
+                # would wait for ranks that are stuck calculating the
+                # material properties of the PML.
+                self.global_comm.Ibcast(buffer, self.COORDINATOR_RANK).Wait()
+                cfs.sigma.max = buffer[0]
+
+        super().calculate_update_coeffs(er, mr)
+
+
 def print_pml_info(G):
     """Prints information about PMLs.