Allow splitting Basis created using ElementVector (#811)

kinnala · Nov 29, 2021 · be9642f · be9642f
1 parent 264c9e9
commit be9642f
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/skfem/assembly/basis/abstract_basis.py b/skfem/assembly/basis/abstract_basis.py
@@ -5,7 +5,8 @@
 import numpy as np
 from numpy import ndarray
 from skfem.assembly.dofs import Dofs, DofsView
-from skfem.element import DiscreteField, Element, ElementComposite
+from skfem.element import (DiscreteField, Element, ElementComposite,
+                           ElementVector)
 from skfem.mapping import Mapping
 from skfem.mesh import Mesh
 from skfem.quadrature import get_quadrature
@@ -360,11 +361,17 @@ def linear_combination(n, refn):
 
     def split_indices(self) -> List[ndarray]:
         """Return indices for the solution components."""
-        if isinstance(self.elem, ElementComposite):
+        if ((isinstance(self.elem, ElementComposite)
+             or isinstance(self.elem, ElementVector))):
+            nelems = (len(self.elem.elems)
+                      if isinstance(self.elem, ElementComposite)
+                      else self.mesh.dim())
             o = np.zeros(4, dtype=np.int64)
             output: List[ndarray] = []
-            for k in range(len(self.elem.elems)):
-                e = self.elem.elems[k]
+            for k in range(nelems):
+                e = (self.elem.elems[k]
+                     if isinstance(self.elem, ElementComposite)
+                     else self.elem.elem)
                 output.append(np.concatenate((
                     self.nodal_dofs[o[0]:(o[0] + e.nodal_dofs)].flatten('F'),
                     self.edge_dofs[o[1]:(o[1] + e.edge_dofs)].flatten('F'),
@@ -385,6 +392,10 @@ def split_bases(self) -> List['AbstractBasis']:
             return [type(self)(self.mesh, e, self.mapping,
                                quadrature=self.quadrature)
                     for e in self.elem.elems]
+        elif isinstance(self.elem, ElementVector):
+            return [type(self)(self.mesh, self.elem.elem, self.mapping,
+                               quadrature=self.quadrature)
+                    for _ in range(self.mesh.dim())]
         raise ValueError("AbstractBasis.elem has only a single component!")
 
     @property

diff --git a/tests/test_basis.py b/tests/test_basis.py
@@ -68,6 +68,13 @@ def bilinf(u, p, v, q, w):
         self.assertTrue(abs((C1 - C2).min()) < 1e-10)
         self.assertTrue(abs((C1 - C2).max()) < 1e-10)
 
+        # test splitting ElementVector
+        (ux, uxbasis), (uy, uybasis) = u_basis.split(u)
+        assert_allclose(ux[uxbasis.nodal_dofs[0]], u[u_basis.nodal_dofs[0]])
+        assert_allclose(ux[uxbasis.facet_dofs[0]], u[u_basis.facet_dofs[0]])
+        assert_allclose(uy[uybasis.nodal_dofs[0]], u[u_basis.nodal_dofs[1]])
+        assert_allclose(uy[uybasis.facet_dofs[0]], u[u_basis.facet_dofs[1]])
+
 
 class TestCompositeFacetAssembly(TestCase):