tests: Added additional tests for substitution and unexpansion of cus…

…tom coefficients with nested expressions
devitocodes · Oct 27, 2023 · 4f43bc9 · 4f43bc9
1 parent 213a51b
commit 4f43bc9
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 0 deletions.
diff --git a/tests/test_symbolic_coefficients.py b/tests/test_symbolic_coefficients.py
@@ -459,3 +459,34 @@ def test_compound_subs(self):
 
         # `str` simply because some objects are of type EvalDerivative
         assert str(eq.evaluate.rhs) == str(term0 + term1 + term2)
+
+    def test_compound_nested_subs(self):
+        grid = Grid(shape=(11, 11))
+        x, y = grid.dimensions
+        hx, hy = grid.spacing_symbols
+
+        f = Function(name='f', grid=grid, space_order=2)
+        p = TimeFunction(name='p', grid=grid, space_order=2,
+                         coefficients='symbolic')
+
+        coeffs0 = np.array([100, 100, 100])
+        coeffs1 = np.array([200, 200, 200])
+
+        subs = Substitutions(Coefficient(1, p, x, coeffs0),
+                             Coefficient(1, p, y, coeffs1))
+
+        eq = Eq(p.forward, (f*p.dx).dy, coefficients=subs)
+
+        mul = lambda e, i: sp.Mul(f.subs(y, y+i*hy), e, 200, evaluate=False)
+        term0 = mul(p*100 +
+                    p.subs(x, x-hx)*100 +
+                    p.subs(x, x+hx)*100, 0)
+        term1 = mul(p.subs(y, y-hy)*100 +
+                    p.subs({x: x-hx, y: y-hy})*100 +
+                    p.subs({x: x+hx, y: y-hy})*100, -1)
+        term2 = mul(p.subs(y, y+hy)*100 +
+                    p.subs({x: x-hx, y: y+hy})*100 +
+                    p.subs({x: x+hx, y: y+hy})*100, 1)
+
+        # `str` simply because some objects are of type EvalDerivative
+        assert str(eq.evaluate.rhs) == str(term0 + term1 + term2)
diff --git a/tests/test_unexpansion.py b/tests/test_unexpansion.py
@@ -57,6 +57,25 @@ def test_numeric_coeffs(self):
         # Non-perfect mixed derivative
         Operator(Eq(u, (u.dx + v.dx).dx, coefficients=coeffs), opt=opt).cfunction
 
+    def test_compound_nested_subs(self):
+        grid = Grid(shape=(11, 11), extent=(10., 10.))
+        x, y = grid.dimensions
+        hx, hy = grid.spacing_symbols
+
+        f = Function(name='f', grid=grid, space_order=2)
+        p = TimeFunction(name='p', grid=grid, space_order=2,
+                         coefficients='symbolic')
+
+        coeffs0 = np.full(3, 1)
+        coeffs1 = np.full(3, 2)
+
+        subs = Substitutions(Coefficient(1, p, x, coeffs0),
+                             Coefficient(1, p, y, coeffs1))
+
+        opt = ('advanced', {'expand': False})
+
+        Operator(Eq(p.forward, (f*p.dx).dy, coefficients=subs), opt=opt).cfunction
+
 
 class Test1Pass(object):