Merge pull request #2736 from devitocodes/conddim-size

api: fix arg processing for conditon+factor

Author: mloubout

devito/types/dimension.py

@@ -14,6 +14,7 @@
 from devito.types.args import ArgProvider
 from devito.types.basic import Symbol, DataSymbol, Scalar
 from devito.types.constant import Constant
+from devito.types.relational import relational_min, relational_max
 
 
 __all__ = ['Dimension', 'SpaceDimension', 'TimeDimension', 'DefaultDimension',
@@ -1015,7 +1016,20 @@ def _arg_defaults(self, _min=None, size=None, alias=None):
         dim = alias or self
         if dim.uses_symbolic_factor:
             factor = defaults[dim.symbolic_factor.name] = self.factor_data
-            defaults[dim.parent.max_name] = range(0, factor*size - 1)
+            if dim.condition is None:
+                d0 = 0
+                d1 = sympy.S.Infinity
+            else:
+                d0 = relational_min(dim.condition, dim.parent)
+                d1 = relational_max(dim.condition, dim.parent)
+            if d1 < sympy.S.Infinity:
+                # We make sure the condition size matches the input size
+                size0 = (d1 - d0 + factor) // factor
+                if size < size0:
+                    raise ValueError(f"Incompatible size for ConditionalDimension "
+                                     f"{self.name}: {size} < {size0}")
+            else:
+                defaults[dim.parent.max_name] = range(d0, d0 + factor*size - 1)
 
         return defaults
 
@@ -1482,7 +1496,7 @@ def _arg_defaults(self, **kwargs):
     def _arg_values(self, *args, **kwargs):
         return {}
 
-    def _arg_check(self, *args):
+    def _arg_check(self, *args, **kwargs):
         """A CustomDimension performs no runtime checks."""
         return
 

devito/types/relational.py

@@ -3,7 +3,7 @@
 
 import sympy
 
-__all__ = ['Le', 'Lt', 'Ge', 'Gt', 'Ne', 'relational_min']
+__all__ = ['Le', 'Lt', 'Ge', 'Gt', 'Ne', 'relational_min', 'relational_max']
 
 
 class AbstractRel:
@@ -250,3 +250,44 @@ def _(expr, s):
         return expr.lts
     else:
         return 0
+
+
+def relational_max(expr, s):
+    """
+    Infer the maximum valid value for symbol `s` in the expression `expr`.
+    For example
+        - if `expr` is `s < 10`, then the maximum valid value for `s` is 9
+        - if `expr` is `s >= 10`, then the maximum valid value for `s` is Inf
+    """
+    if not expr.has(s):
+        return sympy.S.Infinity
+
+    return _relational_max(expr, s)
+
+
+@singledispatch
+def _relational_max(s, expr):
+    return sympy.S.Infinity
+
+
+@_relational_max.register(sympy.And)
+def _(expr, s):
+    return min([_relational_max(e, s) for e in expr.args])
+
+
+@_relational_max.register(Gt)
+@_relational_max.register(Lt)
+def _(expr, s):
+    if s == expr.lts:
+        return expr.gts - 1
+    else:
+        return sympy.S.Infinity
+
+
+@_relational_max.register(Ge)
+@_relational_max.register(Le)
+def _(expr, s):
+    if s == expr.lts:
+        return expr.gts
+    else:
+        return sympy.S.Infinity

tests/test_dimension.py

@@ -1982,6 +1982,31 @@ def test_symbolic_factor_override(self):
 
         assert all(np.all(usave.data[i] == 16 + i*8) for i in range(4))
 
+    def test_factor_and_condition(self):
+        grid = Grid(shape=(10, 10))
+        time = grid.time_dim
+
+        nt = 200
+        bounds = (10, 100)
+        factor = 5
+
+        condition = And(Ge(time, bounds[0]), Le(time, bounds[1]))
+        time_under = ConditionalDimension(name='timeu', parent=time,
+                                          factor=factor, condition=condition)
+        buffer_size = (bounds[1] - bounds[0] + factor) // factor
+
+        rec = SparseTimeFunction(name='rec', grid=grid, npoint=1, nt=nt)
+        u = TimeFunction(name='u', grid=grid, space_order=2)
+        usaved = TimeFunction(name='usaved', grid=grid, space_order=2,
+                              time_dim=time_under, save=buffer_size)
+
+        eq = [Eq(u.forward, u+1), Eq(usaved, u)] + rec.interpolate(u)
+
+        op = Operator(eq)
+        op(time_m=1, time_M=nt-1, usaved=usaved, rec=rec)
+        for t in range(buffer_size):
+            assert np.all(usaved.data[t] == t*factor + bounds[0] - 1)
+
 
 class TestCustomDimension: