Merge pull request #3081 from devitocodes/fix-tensor-name-underscore

types: Fix TensorFunction/VectorFunction name with underscores

Author: FabioLuporini

devito/ir/support/guards.py

@@ -214,7 +214,8 @@ class BaseGuardOverflow(Guard):
     A guard for buffer overflow.
     """
 
-    pass
+    def __new__(cls, p0, p1, **kwargs):
+        return super().__new__(cls, p0, p1, evaluate=False)
 
 
 class GuardOverflowGe(BaseGuardOverflow, Ge):

devito/types/basic.py

@@ -1640,9 +1640,15 @@ def grid(self):
 
     @property
     def name(self):
+        # `_name` may be unset if `__init_finalize__` was skipped in `_fromrep`
+        name = getattr(self, '_name', None)
+        if name is not None:
+            return name
+        # Components are named `{name}_{suffix}`, with the suffix built from
+        # the Dimension names, so the base name is everything before the last `_`
         for c in self.values():
             try:
-                return c.name.split('_')[0]
+                return c.name.rsplit('_', 1)[0]
             except AttributeError:
                 # `c` is not a devito object
                 pass
@@ -1660,7 +1666,7 @@ def _rebuild(self, *args, **kwargs):
         # We need to rebuild the components with the new name then
         # rebuild the matrix
         newname = kwargs.pop('name', self.name)
-        comps = [f.func(*args, name=f.name.replace(self.name, newname), **kwargs)
+        comps = [f.func(*args, name=f"{newname}{f.name[len(self.name):]}", **kwargs)
                  for f in self.flat()]
         # Rebuild the matrix with the new components
         return self._new(*self.shape, comps)

tests/test_pickle.py

@@ -13,14 +13,16 @@
 )
 from devito.data import LEFT, OWNED
 from devito.finite_differences.tools import centered, direct, left, right, transpose
-from devito.ir import Backward, Forward, GuardBound, GuardBoundNext, GuardFactor
+from devito.ir import (
+    Backward, Forward, GuardBound, GuardBoundNext, GuardFactor, GuardOverflow
+)
 from devito.mpi.halo_scheme import Halo
 from devito.mpi.routines import (
     MPIMsgEnriched, MPIRegion, MPIRequestObject, MPIStatusObject
 )
 from devito.symbolics import (
     CallFromPointer, Cast, DefFunction, FieldFromPointer, IntDiv, ListInitializer, SizeOf,
-    indexify, pow_to_mul
+    as_long, indexify, pow_to_mul
 )
 from devito.symbolics.unevaluation import Mod
 from devito.tools import EnrichedTuple
@@ -566,6 +568,22 @@ def test_guard_bound_next(self, pickle, direction):
 
             assert str(gbn) == str(new_gbn)
 
+    @pytest.mark.parametrize('negated', [False, True])
+    def test_guard_overflow(self, pickle, negated):
+        d = Dimension(name='d')
+        freespace = Scalar(name='freespace')
+        size = SizeOf(np.float32)*as_long((d + 4)*(d + 8))
+
+        guard = GuardOverflow(freespace, size)
+        if negated:
+            guard = guard.negated
+
+        new_guard = pickle.loads(pickle.dumps(guard))
+
+        assert type(new_guard) is type(guard)
+        assert str(new_guard) == str(guard)
+        assert new_guard.canonical is new_guard
+
     def test_temp_function(self, pickle):
         grid = Grid(shape=(3, 3))
         d = Dimension(name='d')

tests/test_tensors.py

@@ -487,6 +487,22 @@ def test_rebuild(func1):
         assert j.dimensions == tuple(new_dims)
 
 
+@pytest.mark.parametrize('func1', [TensorFunction, TensorTimeFunction,
+                                   VectorFunction, VectorTimeFunction])
+@pytest.mark.parametrize('name', ['tau_s', 'my_tau_x', 'x'])
+def test_name_with_underscore(func1, name):
+    grid = Grid(tuple([5]*3))
+    f1 = func1(name=name, grid=grid)
+    assert f1.name == name
+    assert (2 * f1).name == name
+    assert (f1 + f1).name == name
+
+    f2 = f1.func(name='f2')
+    assert f2.name == 'f2'
+    for (i, j) in zip(f1.flat(), f2.flat(), strict=True):
+        assert j.name == f"f2{i.name[len(name):]}"
+
+
 @pytest.mark.parametrize('func1', [Function, TimeFunction,
                                    TensorFunction, TensorTimeFunction,
                                    VectorFunction, VectorTimeFunction])