api: Add support for subdomain allocation for Functions

devito/types/grid.py

@@ -29,11 +29,20 @@ class CartesianDiscretization(ABC):
     physical domains by congruent parallelotopes (e.g., "tiles" or "bricks").
     """
 
-    def __init__(self, shape=None, dimensions=None, dtype=None):
+    def __init__(self, shape=None, dimensions=None, dtype=None, comm=None, topology=None):
         self._shape = as_tuple(shape)
         self._dimensions = as_tuple(dimensions)
         self._dtype = dtype
 
+        # Create a Distributor, used internally to implement domain decomposition
+        # by all Functions defined on this Grid
+        self._distributor = Distributor(shape, dimensions, comm, topology)
+
+    @property
+    def shape_local(self):
+        """Shape of the local (per-process) physical domain."""
+        return self._distributor.shape
+
     @property
     def shape(self):
         """Shape of the physical domain."""
@@ -56,6 +65,99 @@ def dtype(self):
         """
         return self._dtype
 
+    @property
+    def distributor(self):
+        """The Distributor used for MPI-decomposing the CartesianDiscretization."""
+        return self._distributor
+
+    @property
+    def comm(self):
+        """The MPI communicator inherited from the distributor."""
+        return self._distributor.comm
+
+    def is_distributed(self, dim):
+        """
+        True if `dim` is a distributed Dimension for this CartesianDiscretization,
+        False otherwise.
+        """
+        return any(dim is d for d in self.distributor.dimensions)
+
+    @cached_property
+    def _arg_names(self):
+        ret = []
+        ret.append(self.time_dim.spacing.name)
+        ret.extend([i.name for i in self.origin_map])
+        for i in self.spacing_map:
+            try:
+                ret.append(i.name)
+            except AttributeError:
+                # E.g., {n*h_x: v} (the case of ConditionalDimension)
+                ret.extend([a.name for a in i.free_symbols])
+        return tuple(ret)
+
+    def _arg_defaults(self):
+        """A map of default argument values defined by this Grid."""
+        args = ReducerMap()
+
+        # Dimensions size
+        for k, v in self.dimension_map.items():
+            args.update(k._arg_defaults(_min=0, size=v.loc))
+
+        # Dimensions spacing
+        args.update({k.name: v for k, v in self.spacing_map.items()})
+
+        # Grid origin
+        args.update({k.name: v for k, v in self.origin_map.items()})
+
+        # MPI-related objects
+        if self.distributor.is_parallel:
+            distributor = self.distributor
+            args[distributor._obj_comm.name] = distributor._obj_comm.value
+            args[distributor._obj_neighborhood.name] = distributor._obj_neighborhood.value
+
+        return args
+
+    def _arg_values(self, **kwargs):
+        values = dict(self._arg_defaults())
+
+        # Override spacing and origin if necessary
+        values.update({i: kwargs[i] for i in self._arg_names if i in kwargs})
+
+        return values
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        # A Distributor wraps an MPI communicator, which can't and shouldn't be pickled
+        state.pop('_distributor')
+        return state
+
+    def __setstate__(self, state):
+        for k, v in state.items():
+            setattr(self, k, v)
+        self._distributor = Distributor(self.shape, self.dimensions, MPI.COMM_SELF)
+
+    @property
+    def time_dim(self):
+        """Time dimension associated with this Grid."""
+        return self._time_dim
+
+    @property
+    def stepping_dim(self):
+        """Stepping dimension associated with this Grid."""
+        return self._stepping_dim
+
+    def time_dimension_handle (self, time_dimension, dtype):
+        # Store or create default symbols for time and stepping dimensions
+        if time_dimension is None:
+            spacing = Scalar(name='dt', dtype=dtype, is_const=True)
+            self._time_dim = TimeDimension(name='time', spacing=spacing)
+            self._stepping_dim = SteppingDimension(name='t', parent=self.time_dim)
+        elif isinstance(time_dimension, TimeDimension):
+            self._time_dim = time_dimension
+            self._stepping_dim = SteppingDimension(name='%s_s' % self.time_dim.name,
+                                                   parent=self.time_dim)
+        else:
+            raise ValueError("`time_dimension` must be None or of type TimeDimension")
 
 class Grid(CartesianDiscretization, ArgProvider):
 
@@ -161,11 +263,8 @@ def __init__(self, shape, extent=None, origin=None, dimensions=None,
                                      "of type `%s`" % (d, type(d)))
             dimensions = dimensions
 
-        super().__init__(shape, dimensions, dtype)
 
-        # Create a Distributor, used internally to implement domain decomposition
-        # by all Functions defined on this Grid
-        self._distributor = Distributor(shape, dimensions, comm, topology)
+        super().__init__(shape, dimensions, dtype, comm, topology)
 
         # The physical extent
         self._extent = as_tuple(extent or tuple(1. for _ in self.shape))
@@ -181,20 +280,11 @@ def __init__(self, shape, extent=None, origin=None, dimensions=None,
                                             is_const=True)
                                      for d in self.dimensions)
 
+        self.time_dimension_handle(time_dimension, dtype)
+
         # Sanity check
         assert (self.dim == len(self.origin) == len(self.extent) == len(self.spacing))
 
-        # Store or create default symbols for time and stepping dimensions
-        if time_dimension is None:
-            spacing = Scalar(name='dt', dtype=dtype, is_const=True)
-            self._time_dim = TimeDimension(name='time', spacing=spacing)
-            self._stepping_dim = SteppingDimension(name='t', parent=self.time_dim)
-        elif isinstance(time_dimension, TimeDimension):
-            self._time_dim = time_dimension
-            self._stepping_dim = SteppingDimension(name='%s_s' % self.time_dim.name,
-                                                   parent=self.time_dim)
-        else:
-            raise ValueError("`time_dimension` must be None or of type TimeDimension")
 
     def __repr__(self):
         return "Grid[extent=%s, shape=%s, dimensions=%s]" % (
@@ -228,16 +318,6 @@ def origin_offset(self):
         assert len(grid_origin) == len(self.spacing)
         return tuple(i*h for i, h in zip(grid_origin, self.spacing))
 
-    @property
-    def time_dim(self):
-        """Time dimension associated with this Grid."""
-        return self._time_dim
-
-    @property
-    def stepping_dim(self):
-        """Stepping dimension associated with this Grid."""
-        return self._stepping_dim
-
     @property
     def subdomains(self):
         """The SubDomains defined in this Grid."""
@@ -283,87 +363,13 @@ def spacing_map(self):
 
         return mapper
 
-    @property
-    def shape_local(self):
-        """Shape of the local (per-process) physical domain."""
-        return self._distributor.shape
 
     @property
     def dimension_map(self):
         """Map between SpaceDimensions and their global/local size."""
         return {d: GlobalLocal(g, l)
                 for d, g, l in zip(self.dimensions, self.shape, self.shape_local)}
 
-    @property
-    def distributor(self):
-        """The Distributor used for MPI-decomposing the CartesianDiscretization."""
-        return self._distributor
-
-    @property
-    def comm(self):
-        """The MPI communicator inherited from the distributor."""
-        return self._distributor.comm
-
-    def is_distributed(self, dim):
-        """
-        True if `dim` is a distributed Dimension for this CartesianDiscretization,
-        False otherwise.
-        """
-        return any(dim is d for d in self.distributor.dimensions)
-
-    @cached_property
-    def _arg_names(self):
-        ret = []
-        ret.append(self.time_dim.spacing.name)
-        ret.extend([i.name for i in self.origin_map])
-        for i in self.spacing_map:
-            try:
-                ret.append(i.name)
-            except AttributeError:
-                # E.g., {n*h_x: v} (the case of ConditionalDimension)
-                ret.extend([a.name for a in i.free_symbols])
-        return tuple(ret)
-
-    def _arg_defaults(self):
-        """A map of default argument values defined by this Grid."""
-        args = ReducerMap()
-
-        # Dimensions size
-        for k, v in self.dimension_map.items():
-            args.update(k._arg_defaults(_min=0, size=v.loc))
-
-        # Dimensions spacing
-        args.update({k.name: v for k, v in self.spacing_map.items()})
-
-        # Grid origin
-        args.update({k.name: v for k, v in self.origin_map.items()})
-
-        # MPI-related objects
-        if self.distributor.is_parallel:
-            distributor = self.distributor
-            args[distributor._obj_comm.name] = distributor._obj_comm.value
-            args[distributor._obj_neighborhood.name] = distributor._obj_neighborhood.value
-
-        return args
-
-    def _arg_values(self, **kwargs):
-        values = dict(self._arg_defaults())
-
-        # Override spacing and origin if necessary
-        values.update({i: kwargs[i] for i in self._arg_names if i in kwargs})
-
-        return values
-
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        # A Distributor wraps an MPI communicator, which can't and shouldn't be pickled
-        state.pop('_distributor')
-        return state
-
-    def __setstate__(self, state):
-        for k, v in state.items():
-            setattr(self, k, v)
-        self._distributor = Distributor(self.shape, self.dimensions, MPI.COMM_SELF)
 
 
 class AbstractSubDomain(CartesianDiscretization):
@@ -380,7 +386,6 @@ def __init__(self, *args, **kwargs):
             self.name = self.__class__.__name__
 
         # All other attributes get initialized upon `__subdomain_finalize__`
-        super().__init__()
 
     def __subdomain_finalize__(self, grid, **kwargs):
         """
@@ -477,7 +482,7 @@ class SubDomain(AbstractSubDomain):
     Interior : An example of preset Subdomain.
     """
 
-    def __subdomain_finalize__(self, grid, **kwargs):
+    def __subdomain_finalize__(self, grid, comm=None, topology=None, time_dimension=None, **kwargs):
         # Create the SubDomain's SubDimensions
         sub_dimensions = []
         sdshape = []
@@ -523,6 +528,8 @@ def __subdomain_finalize__(self, grid, **kwargs):
         self._shape = tuple(sdshape)
         self._dimensions = tuple(sub_dimensions)
         self._dtype = grid.dtype
+        self._distributor = Distributor(self._shape, self._dimensions, comm, topology)
+        self.time_dimension_handle(time_dimension, self._dtype)
 
     def define(self, dimensions):
         """

tests/test_subdomains.py

@@ -38,6 +38,23 @@ def define(self, dimensions):
             expr = Operator._lower_exprs([eq0])[0]
         assert expr.rhs == x1 * f[x1 + 1, y1 + 1] + y1
 
+    def test_subfunctions_dimensions (self):
+        class sd0(SubDomain):
+            name = 'd0'
+
+            def define(self, dimensions):
+                x, y = dimensions
+                return {x: ('middle', 3, 4), y: ('middle', 4, 3)}
+        s_d0 = sd0()
+        grid = Grid(shape=(10, 10), subdomains=(s_d0,))
+        dim = s_d0.shape
+
+        f = Function(name='f', grid=grid.subdomains['d0'], dtype=np.int32)
+        g = TimeFunction(name='g', grid=grid.subdomains['d0'], dtype=np.int32)
+
+        assert dim == f.data.shape
+        assert dim == g.data.shape[1:]
+
     def test_multiple_middle(self):
         """
         Test Operator with two basic 'middle' subdomains defined.