From ffff920d88951610a79e0c1ca75c2f8093a97386 Mon Sep 17 00:00:00 2001 From: Tom Gustafsson Date: Thu, 25 Nov 2021 13:23:25 +0200 Subject: [PATCH] Add matrix multiplication helper (#807) --- docs/examples/ex04.py | 7 +++---- skfem/helpers.py | 9 +++++++++ 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/docs/examples/ex04.py b/docs/examples/ex04.py index d6b988bf2..6b2baed12 100644 --- a/docs/examples/ex04.py +++ b/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)) diff --git a/skfem/helpers.py b/skfem/helpers.py index 0d80ab782..f016f8bd7 100644 --- a/skfem/helpers.py +++ b/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)