From cc99ad97b335c1698816c19b83fa46dc1683b6f0 Mon Sep 17 00:00:00 2001
From: Tom Gustafsson <jkinnala@gmail.com>
Date: Mon, 8 Oct 2018 22:09:46 +0300
Subject: [PATCH] Added multielement support to MappingAffine.F and tests for
 InteriorBasis.interpolator.

---
 skfem/mapping.py       |  7 ++++++-
 tests/test_assembly.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/skfem/mapping.py b/skfem/mapping.py
index ae9d07886..c6e52923a 100644
--- a/skfem/mapping.py
+++ b/skfem/mapping.py
@@ -368,7 +368,12 @@ def F(self, X, tind=None):
         else:
             A, b = self.A[:, :, tind], self.b[:, tind]
 
-        return (np.einsum('ijk,jl', A, X).T + b.T).T
+        if len(X.shape) == 2:
+            return (np.einsum('ijk,jl', A, X).T + b.T).T
+        elif len(X.shape) == 3:
+            return (np.einsum('ijk,jkl->ikl', A, X).T + b.T).T
+        else:
+            raise Exception("Wrong dimension of input.")
 
     def invF(self, x, tind=None):
         if tind is None:
diff --git a/tests/test_assembly.py b/tests/test_assembly.py
index 6e9b8ad99..c3b2bb1f3 100644
--- a/tests/test_assembly.py
+++ b/tests/test_assembly.py
@@ -89,3 +89,40 @@ class IntegrateFuncOverBoundaryPartTetP1(IntegrateFuncOverBoundaryPart):
 
 class IntegrateFuncOverBoundaryPartTetP0(IntegrateFuncOverBoundaryPart):
     case = (MeshTet, ElementTetP0)
+
+
+class BasisInterpolator(unittest.TestCase):
+    case = (MeshTri, ElementTriP1)
+
+    def initOnes(self, basis):
+        return np.ones(basis.N)
+
+    def runTest(self):
+        mtype, etype = self.case
+        m = mtype()
+        m.refine(3)
+        e = etype()
+        ib = InteriorBasis(m, e)
+
+        x = self.initOnes(ib)
+        f = ib.interpolator(x)
+
+        self.assertTrue(np.sum(f(np.array([np.sin(m.p[0, :]), np.sin(3.0*m.p[1, :])]))-1.0) < 1e-10)
+
+class BasisInterpolatorMorley(BasisInterpolator):
+    case = (MeshTri, ElementTriMorley)
+
+    def initOnes(self, basis):
+        @bilinear_form
+        def mass(u, du, ddu, v, dv, ddv, w):
+            return u*v
+
+        @linear_form
+        def ones(v, dv, ddv, w):
+            return 1.0*v
+
+        M = asm(mass, basis)
+        f = asm(ones, basis)
+
+        return solve(M, f)
+