Merge pull request #2904 from devitocodes/allow-spacing-as-grid-input

api: allow spacing as grid input

Author: mloubout

.github/workflows/pytest-core-nompi.yaml

@@ -25,6 +25,7 @@ jobs:
       DEVITO_ARCH: "${{ matrix.arch }}"
       DEVITO_LANGUAGE: ${{ matrix.language }}
       OMP_NUM_THREADS: 2
+      DEVITO_DEVELOP: 1
 
     strategy:
       # Prevent cancellation if a single workflow fails

devito/types/grid.py

@@ -150,8 +150,8 @@ class Grid(CartesianDiscretization, ArgProvider):
     _default_dimensions = ('x', 'y', 'z')
 
     def __init__(self, shape, extent=None, origin=None, dimensions=None,
-                 time_dimension=None, dtype=np.float32, subdomains=None,
-                 comm=None, topology=None):
+                 spacing=None, time_dimension=None, dtype=np.float32,
+                 subdomains=None, comm=None, topology=None):
         shape = as_tuple(shape)
 
         # Create or pull the SpaceDimensions
@@ -193,9 +193,16 @@ def __init__(self, shape, extent=None, origin=None, dimensions=None,
             self._topology = None
         self._distributor = Distributor(shape, dimensions, comm, self._topology)
 
-        # The physical extent
-        extent = as_tuple(extent or tuple(1. for _ in self.shape))
-        self._extent = tuple(dtype(e) for e in extent)
+        # The physical extent and grid spacing
+        if spacing is not None:
+            self._spacing = tuple(dtype(s) for s in as_tuple(spacing))
+        else:
+            self._spacing = None
+
+        if extent is not None:
+            self._extent = tuple(dtype(e) for e in as_tuple(extent))
+        else:
+            self._extent = tuple(1. for _ in shape) if spacing is None else None
 
         # The origin of the grid
         origin = as_tuple(origin or tuple(0. for _ in self.shape))
@@ -230,10 +237,13 @@ def __repr__(self):
         return 'Grid' + \
             f'[extent={self.extent}, shape={self.shape}, dimensions={self.dimensions}]'
 
-    @property
+    @cached_property
     def extent(self):
         """Physical extent of the domain in m."""
-        return self._extent
+        if self._extent is not None:
+            return self._extent
+        extent = ((np.array(self.shape) - 1)*np.array(self.spacing)).astype(self.dtype)
+        return as_tuple(extent)
 
     @property
     def origin(self):
@@ -293,9 +303,11 @@ def volume_cell(self):
         """Volume of a single cell e.g  h_x*h_y*h_z in 3D."""
         return prod(d.spacing for d in self.dimensions).subs(self.spacing_map)
 
-    @property
+    @cached_property
     def spacing(self):
         """Spacing between grid points in m."""
+        if self._spacing is not None:
+            return self._spacing
         spacing = (np.array(self.extent) / (np.array(self.shape) - 1)).astype(self.dtype)
         return as_tuple(spacing)
 

tests/test_symbolics.py

@@ -132,6 +132,24 @@ def test_real():
         assert s.is_imaginary is np.iscomplexobj(dtype(0))
 
 
+@pytest.mark.parametrize('spacing, extent, shape, expected, broken', [
+    ((0.5, 0.5), None, (11, 11), ((0.5, 0.5), (5.0, 5.0)), False),
+    (None, (5.0, 5.0), (11, 11), ((0.5, 0.5), (5.0, 5.0)), False),
+    ((0.5, 0.5), (5.0, 5.0), (11, 11), ((0.5, 0.5), (5.0, 5.0)), False),
+    (None, (.3, .3), (151, 146), ((0.002, 0.002), (.3, .3)), 'spacing'),
+    ((.002, .002), (.3, .3), (151, 146), ((0.002, 0.002), (.3, .3)), False),
+    ((.002, .002), None, (151, 146), ((0.002, 0.002), (.3, .3)), 'extent'),
+    (None, None, (11, 11), ((.1, .1), (1.0, 1.0)), False),
+])
+def test_grid_inputs(spacing, extent, shape, expected, broken):
+    grid = Grid(shape=shape, spacing=spacing, extent=extent)
+    sp, ex = expected
+    sp = np.array(sp).astype(grid.dtype)
+    ex = np.array(ex).astype(grid.dtype)
+    assert np.allclose(grid.spacing, sp, atol=0, rtol=0) is (broken != 'spacing')
+    assert np.allclose(grid.extent, ex, atol=0, rtol=0) is (broken != 'extent')
+
+
 def test_constant():
     c = Constant(name='c')