Merge pull request #13 from mscroggs/restructure-and-matrix-shape

Restructure and matrix shape

Author: mscroggs

symfem/core/finite_element.py

@@ -7,17 +7,32 @@
 class FiniteElement:
     """Abstract finite element."""
 
-    def __init__(self, reference, basis, dofs, domain_dim, range_dim):
+    def __init__(self, reference, basis, dofs, domain_dim, range_dim,
+                 range_shape=None):
         assert len(basis) == len(dofs)
         self.reference = reference
         self.basis = basis
         self.dofs = dofs
         self.domain_dim = domain_dim
         self.range_dim = range_dim
+        self.range_shape = range_shape
         self.space_dim = len(dofs)
         self._basis_functions = None
+        self._reshaped_basis_functions = None
 
-    def get_basis_functions(self):
+    def get_polynomial_basis(self, reshape=True):
+        """Get the polynomial basis for the element."""
+        if reshape and self.range_shape is not None:
+            if len(self.range_shape) != 2:
+                raise NotImplementedError
+            assert self.range_shape[0] * self.range_shape[1] == self.range_dim
+            return [sympy.Matrix(
+                [b[i * self.range_shape[1]: (i + 1) * self.range_shape[1]]
+                 for i in range(self.range_shape[0])]) for b in self.basis]
+
+        return self.basis
+
+    def get_basis_functions(self, reshape=True):
         """Get the basis functions of the element."""
         if self._basis_functions is None:
             mat = []
@@ -44,6 +59,14 @@ def get_basis_functions(self):
                             b[j] += c * d_j
                     self._basis_functions.append(b)
 
+        if reshape and self.range_shape is not None:
+            if len(self.range_shape) != 2:
+                raise NotImplementedError
+            assert self.range_shape[0] * self.range_shape[1] == self.range_dim
+            return [sympy.Matrix(
+                [b[i * self.range_shape[1]: (i + 1) * self.range_shape[1]]
+                 for i in range(self.range_shape[0])]) for b in self._basis_functions]
+
         return self._basis_functions
 
     def tabulate_basis(self, points, order="xyzxyz"):
@@ -52,7 +75,7 @@ def tabulate_basis(self, points, order="xyzxyz"):
             output = []
             for p in points:
                 row = []
-                for b in self.get_basis_functions():
+                for b in self.get_basis_functions(False):
                     row.append(subs(b, x, p))
                 output.append(row)
             return output
@@ -62,15 +85,15 @@ def tabulate_basis(self, points, order="xyzxyz"):
             for p in points:
                 row = []
                 for d in range(self.range_dim):
-                    for b in self.get_basis_functions():
+                    for b in self.get_basis_functions(False):
                         row.append(subs(b[d], x, p))
                 output.append(row)
             return output
         if order == "xyzxyz":
             output = []
             for p in points:
                 row = []
-                for b in self.get_basis_functions():
+                for b in self.get_basis_functions(False):
                     for i in subs(b, x, p):
                         row.append(i)
                 output.append(row)
@@ -79,7 +102,7 @@ def tabulate_basis(self, points, order="xyzxyz"):
             output = []
             for p in points:
                 row = []
-                for b in self.get_basis_functions():
+                for b in self.get_basis_functions(False):
                     row.append(subs(b, x, p))
                 output.append(row)
             return output

symfem/elements/regge.py

@@ -36,7 +36,8 @@ def __init__(self, reference, order):
                                       for j, o in enumerate(entity.origin)),
                                 tangent, entity_dim=edim))
 
-        super().__init__(reference, poly, dofs, reference.tdim, reference.tdim ** 2)
+        super().__init__(reference, poly, dofs, reference.tdim, reference.tdim ** 2,
+                         (reference.tdim, reference.tdim))
 
     names = ["Regge"]
     min_order = 0

test/test_against_basix.py

@@ -60,9 +60,9 @@ def test_against_basix(cell, symfem_type, basix_type, order):
     result = space.tabulate(0, points)[0]
 
     if element.range_dim == 1:
-        basis = element.get_basis_functions()
+        basis = element.get_basis_functions(False)
         sym_result = [[float(subs(b, x, p)) for b in basis] for p in points]
     else:
-        basis = element.get_basis_functions()
+        basis = element.get_basis_functions(False)
         sym_result = [[float(subs(b, x, p)[j]) for j in range(element.range_dim) for b in basis] for p in points]
     assert np.allclose(result, sym_result)