Merge pull request #2072 from devitocodes/tens-transp

symbolics: evaluate transpose on elements of tensors in case of deriv…

Author: mloubout

devito/types/basic.py

@@ -677,6 +677,16 @@ def __init_finalize__(self, *args, **kwargs):
     def doit(self, **hint):
         return self
 
+    def transpose(self, inner=True):
+        new = super().transpose()
+        if inner:
+            return new.applyfunc(lambda x: getattr(x, 'T', x))
+        return new
+
+    def adjoint(self, inner=True):
+        # Real valued adjoint is transpose
+        return self.transpose(inner=inner)
+
     def _eval_matrix_mul(self, other):
         """
         Copy paste from sympy to avoid explicit call to sympy.Add

examples/seismic/elastic/operators.py

@@ -61,7 +61,7 @@ def ForwardOperator(model, geometry, space_order=4, save=False, **kwargs):
     # Particle velocity
     eq_v = v.dt - b * div(tau)
     # Stress
-    e = (grad(v.forward) + grad(v.forward).T)
+    e = (grad(v.forward) + grad(v.forward).transpose(inner=False))
     eq_tau = tau.dt - lam * diag(div(v.forward)) - mu * e
 
     u_v = Eq(v.forward, model.damp * solve(eq_v, v.forward))

examples/seismic/tutorials/06_elastic.ipynb

@@ -137,7 +137,7 @@
     "\n",
     "# First order elastic wave equation\n",
     "pde_v = v.dt - ro * div(tau)\n",
-    "pde_tau = tau.dt - l * diag(div(v.forward)) - mu * (grad(v.forward) + grad(v.forward).T)\n",
+    "pde_tau = tau.dt - l * diag(div(v.forward)) - mu * (grad(v.forward) + grad(v.forward).transpose(inner=False))\n",
     "# Time update\n",
     "u_v = Eq(v.forward, solve(pde_v, v.forward))\n",
     "u_t = Eq(tau.forward, solve(pde_tau, tau.forward))\n",
@@ -304,7 +304,7 @@
     "\n",
     "# First order elastic wave equation\n",
     "pde_v = v.dt - ro * div(tau)\n",
-    "pde_tau = tau.dt - l * diag(div(v.forward)) - mu * (grad(v.forward) + grad(v.forward).T)\n",
+    "pde_tau = tau.dt - l * diag(div(v.forward)) - mu * (grad(v.forward) + grad(v.forward).transpose(inner=False))\n",
     "# Time update\n",
     "u_v = Eq(v.forward, solve(pde_v, v.forward))\n",
     "u_t = Eq(tau.forward, solve(pde_tau, tau.forward))\n",
@@ -445,7 +445,7 @@
   },
   "hide_input": false,
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -459,7 +459,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.9"
+   "version": "3.9.16"
   },
   "latex_envs": {
    "LaTeX_envs_menu_present": true,
@@ -485,5 +485,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

examples/seismic/tutorials/06_elastic_varying_parameters.ipynb

@@ -48,7 +48,7 @@ def ForwardOperator(model, geometry, space_order=4, save=False, **kwargs):
     pde_v = v.dt - b * div(tau)
     u_v = Eq(v.forward, model.damp * solve(pde_v, v.forward))
     # Strain
-    e = grad(v.forward) + grad(v.forward).T
+    e = grad(v.forward) + grad(v.forward).transpose(inner=False)
 
     # Stress equations
     pde_tau = tau.dt - r.forward - l * t_ep / t_s * diag(div(v.forward)) - \

examples/seismic/tutorials/09_viscoelastic.ipynb

@@ -149,6 +149,37 @@ def test_vector_transpose(func1):
     assert np.all([f1[i] == f2[i] for i in range(3)])
 
 
+@pytest.mark.parametrize('func1', [VectorFunction, VectorTimeFunction])
+def test_vector_transpose_deriv(func1):
+    grid = Grid(tuple([5]*3))
+    f1 = func1(name="f1", grid=grid)
+    f2 = f1.dx.T
+    assert all([f2[i] == f1[i].dx.T for i in range(3)])
+
+
+@pytest.mark.parametrize('func1', [TensorFunction, TensorTimeFunction])
+def test_tensor_transpose_deriv(func1):
+    grid = Grid(tuple([5]*3))
+    f1 = func1(name="f1", grid=grid)
+    f2 = f1.dx.T
+    assert np.all([f2[i, j] == f1[j, i].dx.T for i in range(3) for j in range(3)])
+
+
+@pytest.mark.parametrize('func1', [TensorFunction, TensorTimeFunction,
+                                   VectorFunction, VectorTimeFunction])
+def test_transpose_vs_T(func1):
+    grid = Grid(tuple([5]*3))
+    f1 = func1(name="f1", grid=grid)
+    f2 = f1.dx.T
+    f3 = f1.dx.transpose(inner=True)
+    f4 = f1.dx.transpose(inner=False)
+    # inner=True is the same as T
+    assert f3 == f2
+    # inner=False doesn't tranpose inner derivatives
+    for f4i, f2i in zip(f4, f2):
+        assert f4i == f2i.T
+
+
 @pytest.mark.parametrize('func1', [TensorFunction, TensorTimeFunction,
                                    VectorFunction, VectorTimeFunction])
 def test_tensor_fd(func1):