api: fix arg processing for subsampling factor

Author: mloubout

devito/types/dimension.py

@@ -915,7 +915,7 @@ def __init_finalize__(self, name, parent=None, factor=None, condition=None,
 
         # Process subsampling factor
         fname = f"{name}f"
-        if factor is None:
+        if factor is None or factor == 1:
             self._factor = None
         elif is_number(factor):
             self._factor = int(factor)
@@ -966,9 +966,15 @@ def free_symbols(self):
             pass
         return retval
 
-    def _arg_values(self, interval, grid=None, **kwargs):
-        # Parent dimension define the interval
-        fact = self.factor
+    def _arg_values(self, interval, grid=None, args=None, **kwargs):
+        if self.symbolic_factor is not None:
+            fname = self.symbolic_factor.name
+            args = args or {}
+            fact = kwargs.get(fname, args.get(fname, self.factor))
+        else:
+            # No factor
+            return {}
+
         toint = lambda x: math.ceil(x / fact)
         vals = {}
         try:
@@ -981,6 +987,9 @@ def _arg_values(self, interval, grid=None, **kwargs):
         except (KeyError, TypeError):
             pass
 
+        if self.symbolic_factor is not None:
+            vals[self.symbolic_factor.name] = fact
+
         return vals
 
     def _arg_defaults(self, _min=None, size=None, alias=None):
@@ -990,11 +999,9 @@ def _arg_defaults(self, _min=None, size=None, alias=None):
         # `factor` endpoints are legal, so we return them all. It's then
         # up to the caller to decide which one to pick upon reduction
         dim = alias or self
-        if dim.condition is not None or size is None or dim._factor is None:
-            return defaults
-
-        factor = defaults[dim.symbolic_factor.name] = self.factor
-        defaults[dim.parent.max_name] = range(0, factor*size - 1)
+        if dim.symbolic_factor is not None:
+            factor = defaults[dim.symbolic_factor.name] = self.factor
+            defaults[dim.parent.max_name] = range(0, factor*size - 1)
 
         return defaults
 

tests/test_dimension.py

@@ -17,6 +17,7 @@
 from devito.ir import SymbolRegistry
 from devito.symbolics import indexify, retrieve_functions, IntDiv, INT
 from devito.types import Array, StencilDimension, Symbol
+from devito.types.basic import Scalar
 from devito.types.dimension import AffineIndexAccessFunction, Thickness
 
 
@@ -1012,9 +1013,9 @@ def test_issue_1592(self):
         op = Operator(Eq(v.forward, v.dx))
         op.apply(time=6)
         exprs = FindNodes(Expression).visit(op)
-        assert exprs[-1].expr.lhs.indices[0] == IntDiv(time, time_sub.factor) + 1
-        assert time_sub.factor.data == 2
-        assert time_sub.factor.is_Constant
+        assert exprs[-1].expr.lhs.indices[0] == IntDiv(time, time_sub.symbolic_factor) + 1
+        assert time_sub.factor == 2
+        assert isinstance(time_sub.symbolic_factor, Scalar)
 
     def test_issue_1753(self):
         grid = Grid(shape=(3, 3, 3))

tests/test_pickle.py

@@ -320,6 +320,7 @@ def test_conditional_dimension(self, pickle):
         assert cd.name == new_cd.name
         assert cd.parent.name == new_cd.parent.name
         assert cd.factor == new_cd.factor
+        assert cd.symbolic_factor == new_cd.symbolic_factor
         assert cd.condition == new_cd.condition
 
     def test_incr_dimension(self, pickle):
@@ -603,7 +604,7 @@ def test_equation(self, pickle):
         assert new_eq.lhs.name == f.name
         assert str(new_eq.rhs) == 'f(x) + 1'
         assert new_eq.implicit_dims[0].name == 'xs'
-        assert new_eq.implicit_dims[0].factor.data == 4
+        assert new_eq.implicit_dims[0].factor == 4
 
     @pytest.mark.parametrize('typ', [ctypes.c_float, 'struct truct'])
     def test_Cast(self, pickle, typ):