From 955113cacafda08b5548ff01238e2ca7f24af5ba Mon Sep 17 00:00:00 2001 From: nmannall Date: Fri, 1 Mar 2024 14:55:11 +0000 Subject: [PATCH] Add MPIGrid --- gprMax/grid/mpi_grid.py | 47 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 gprMax/grid/mpi_grid.py diff --git a/gprMax/grid/mpi_grid.py b/gprMax/grid/mpi_grid.py new file mode 100644 index 00000000..6661d1aa --- /dev/null +++ b/gprMax/grid/mpi_grid.py @@ -0,0 +1,47 @@ +from typing import Optional + +from mpi4py import MPI + +from gprMax.grid.fdtd_grid import FDTDGrid + + +class MPIGrid(FDTDGrid): + xmin: int + ymin: int + zmin: int + xmax: int + ymax: int + zmax: int + + def __init__(self, mpi_tasks_x: int, mpi_tasks_y: int, mpi_tasks_z: int, comm: Optional[MPI.Intracomm] = None): + super().__init__() + + if comm is None: + self.comm = MPI.COMM_WORLD + else: + self.comm = comm + + if mpi_tasks_x * mpi_tasks_y * mpi_tasks_z > self.size: + # TODO: Raise expection - insufficient MPI tasks to create the grid as requested + pass + + self.mpi_tasks_x = mpi_tasks_x + self.mpi_tasks_y = mpi_tasks_y + self.mpi_tasks_z = mpi_tasks_z + + self.rank = self.comm.rank + self.size = self.comm.size + + def initialise_field_arrays(self): + super().initialise_field_arrays() + + self.local_grid_size_x = self.nx // self.mpi_tasks_x + self.local_grid_size_y = self.ny // self.mpi_tasks_y + self.local_grid_size_z = self.nz // self.mpi_tasks_z + + self.xmin = (self.rank % self.nx) * self.local_grid_size_x + self.ymin = ((self.mpi_tasks_x * self.rank) % self.ny) * self.local_grid_size_y + self.zmin = ((self.mpi_tasks_y * self.mpi_tasks_x * self.rank) % self.nz) * self.local_grid_size_z + self.xmax = self.xmin + self.local_grid_size_x + self.ymax = self.ymin + self.local_grid_size_y + self.zmax = self.zmin + self.local_grid_size_z