Add matrix multiplication helper (#807)

docs/examples/ex04.py

@@ -54,7 +54,7 @@
 from skfem import *
 from skfem.models.elasticity import (linear_elasticity, lame_parameters,
                                      linear_stress)
-from skfem.helpers import dot, ddot, prod, sym_grad, jump
+from skfem.helpers import dot, ddot, prod, sym_grad, jump, mul
 import numpy as np
 from skfem.io import from_meshio
 from skfem.io.json import from_file, to_file
@@ -110,9 +110,8 @@
 def bilin_mortar(u, v, w):
     # jumps
     ju, jv = jump(w, dot(u, w.n), dot(v, w.n))
-    nxn = prod(w.n, w.n)
-    mu = .5 * ddot(nxn, C(sym_grad(u)))
-    mv = .5 * ddot(nxn, C(sym_grad(v)))
+    mu = .5 * dot(w.n, mul(C(sym_grad(u)), w.n))
+    mv = .5 * dot(w.n, mul(C(sym_grad(v)), w.n))
     return ((1. / (alpha * w.h) * ju * jv - mu * jv - mv * ju)
             * (np.abs(w.x[1]) <= limit))
 

skfem/helpers.py

@@ -136,6 +136,15 @@ def prod(u: FieldOrArray,
     return np.einsum('i...,j...,k...->ijk...', u, v, w)
 
 
+def mul(A: FieldOrArray, x: FieldOrArray):
+    """Matrix multiplication."""
+    if isinstance(A, DiscreteField) and A.is_zero():
+        return A
+    if isinstance(x, DiscreteField) and x.is_zero():
+        return x
+    return np.einsum('ij...,j...->i...', A, x)
+
+
 def trace(T):
     """Trace of matrix."""
     return np.einsum('ii...', T)