diff --git a/devito/data/decomposition.py b/devito/data/decomposition.py index 2bb3faf797..c9cc01be9d 100644 --- a/devito/data/decomposition.py +++ b/devito/data/decomposition.py @@ -324,12 +324,16 @@ def index_glb_to_loc(self, *args, rel=True): rel_ofs = self.glb_min + abs_ofs - base if abs_ofs >= base and abs_ofs <= top: return rel_ofs + elif abs_ofs > top: + return top + 1 else: return None else: rel_ofs = abs_ofs - (self.glb_max - top) if abs_ofs >= self.glb_max - top and abs_ofs <= self.glb_max - base: return rel_ofs + elif abs_ofs > self.glb_max - base: + return self.glb_max - base + 1 else: return None else: diff --git a/devito/mpi/distributed.py b/devito/mpi/distributed.py index 387695b7b1..055b77e3da 100644 --- a/devito/mpi/distributed.py +++ b/devito/mpi/distributed.py @@ -14,6 +14,7 @@ from devito.parameters import configuration from devito.tools import EnrichedTuple, as_tuple, ctypes_to_cstr, filter_ordered from devito.types import CompositeObject, Object +from devito.types.utils import DimensionTuple # Do not prematurely initialize MPI @@ -126,7 +127,8 @@ def glb_shape(self): @property def shape(self): """The calling MPI rank's local shape.""" - return tuple(len(i) for i in self.glb_numb) + return DimensionTuple(*[len(i) for i in self.glb_numb], + getters=self.dimensions) @property def dimensions(self): diff --git a/devito/types/grid.py b/devito/types/grid.py index de91284441..4f25c5b10f 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -13,6 +13,7 @@ from devito.types.args import ArgProvider from devito.types.basic import Scalar from devito.types.dense import Function +from devito.types.utils import DimensionTuple from devito.types.dimension import (Dimension, SpaceDimension, TimeDimension, Spacing, SteppingDimension, SubDimension) @@ -37,7 +38,7 @@ def __init__(self, shape=None, dimensions=None, dtype=None): @property def shape(self): """Shape of the physical domain.""" - return self._shape + return DimensionTuple(*self._shape, getters=self.dimensions) @property def dimensions(self): @@ -226,12 +227,13 @@ def origin_ioffset(self): """Offset index of the local (per-process) origin from the domain origin.""" grid_origin = [min(i) for i in self.distributor.glb_numb] assert len(grid_origin) == len(self.spacing) - return tuple(grid_origin) + return DimensionTuple(*grid_origin, getters=self.dimensions) @property def origin_offset(self): """Physical offset of the local (per-process) origin from the domain origin.""" - return tuple(i*h for i, h in zip(self.origin_ioffset, self.spacing)) + return DimensionTuple(*[i*h for i, h in zip(self.origin_ioffset, self.spacing)], + getters=self.dimensions) @property def time_dim(self): diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index 984f2c29ef..9c1e5815e6 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -162,6 +162,47 @@ def define(self, dimensions): assert u0.data.all() == u1.data.all() == u2.data.all() == u3.data.all() + sd_specs = [('middle', 1, 7), ('middle', 2, 3), ('middle', 7, 1), + ('middle', 5, 5), ('right', 3), ('right', 7), ('left', 3), + ('left', 7)] + + @pytest.mark.parametrize('spec', sd_specs) + @pytest.mark.parallel(mode=[2, 3]) + def test_subdomains_mpi(self, spec): + + class sd0(SubDomain): + name = 'd0' + + def __init__(self, spec): + super().__init__() + self.spec = spec + + def define(self, dimensions): + x = dimensions[0] + return {x: self.spec} + s_d0 = sd0(spec) + + grid = Grid(shape=(11,), extent=(10.,), subdomains=(s_d0,)) + x = grid.dimensions[0] + xd0 = grid.subdomains['d0'].dimensions[0] + f = Function(name='f', grid=grid) + + Operator(Eq(f, f+1, subdomain=grid.subdomains['d0']))() + + # Sets ones on a global array according to the subdomains specified + # then slices this according to the indices on this rank and compares + # to the operator output. + check = np.zeros(grid.shape) + + mM_map = {x.symbolic_min: 0, x.symbolic_max: grid.shape[0]-1} + t_map = {k: v for k, v in xd0._thickness_map.items() if v is not None} + start = int(xd0.symbolic_min.subs({**mM_map, **t_map})) + stop = int(xd0.symbolic_max.subs({**mM_map, **t_map})+1) + + check[slice(start, stop)] = 1 + + assert np.all(check[grid.distributor.glb_slices[x]] == f.data) + class TestMultiSubDomain(object):