From 02dcf7a1d21b23e5802bbbdc3f2c69566706eca4 Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Mon, 30 Oct 2023 10:54:37 +0000
Subject: [PATCH] move derivative tabulation from PolynomialSet to ExpansionSet

---
 FIAT/expansions.py     | 19 ++++++++++++++++++-
 FIAT/polynomial_set.py | 30 +++---------------------------
 2 files changed, 21 insertions(+), 28 deletions(-)

diff --git a/FIAT/expansions.py b/FIAT/expansions.py
index f40d5d322..42c2a33af 100644
--- a/FIAT/expansions.py
+++ b/FIAT/expansions.py
@@ -268,8 +268,8 @@ def _tabulate_duffy(self, n, pts):
         eta = (lambda x: x, lambda x: x, eta_square, eta_cube)[dim](xi)
         basis = [dubiner_1d(n, k, eta[k]) for k in range(dim)]
         derivs = [dubiner_deriv_1d(n, k, eta[k]) for k in range(dim)]
-        alphas = mis(dim, 0) + mis(dim, 1)
         tabulations = {}
+        alphas = mis(dim, 0) + mis(dim, 1)
         for alpha in alphas:
             V = [v if a == 0 else dv for a, v, dv in zip(alpha, basis, derivs)]
             phi = V[0]
@@ -311,6 +311,23 @@ def make_dmats(self, degree):
         dmats = numpy.linalg.solve(v, dv)
         return cache.setdefault(key, dmats)
 
+    def _tabulate_jet(self, degree, pts, order=0):
+        from FIAT.polynomial_set import mis
+        result = {}
+        base_vals = self.tabulate(degree, pts)
+        dmats = self.make_dmats(degree) if order > 0 else []
+        for i in range(order + 1):
+            alphas = mis(self.ref_el.get_spatial_dimension(), i)
+            for alpha in alphas:
+                beta = next((beta for beta in sorted(result, reverse=True)
+                             if all(bj <= aj for bj, aj in zip(beta, alpha))), (0,) * len(alpha))
+                vals = base_vals if sum(beta) == 0 else result[beta]
+                for dmat, start, end in zip(dmats, beta, alpha):
+                    for j in range(start, end):
+                        vals = numpy.dot(dmat.T, vals)
+                result[alpha] = vals
+        return result
+
 
 class PointExpansionSet(ExpansionSet):
     """Evaluates the point basis on a point reference element."""
diff --git a/FIAT/polynomial_set.py b/FIAT/polynomial_set.py
index d1ecf00c7..502675cda 100644
--- a/FIAT/polynomial_set.py
+++ b/FIAT/polynomial_set.py
@@ -69,23 +69,9 @@ def tabulate_new(self, pts):
 
     def tabulate(self, pts, jet_order=0):
         """Returns the values of the polynomial set."""
-        result = {}
-        base_vals = self.expansion_set.tabulate(self.embedded_degree, pts)
-        dmats = self.get_dmats() if jet_order > 0 else self.dmats
-        for i in range(jet_order + 1):
-            alphas = mis(self.ref_el.get_spatial_dimension(), i)
-            for alpha in alphas:
-                if sum(alpha) == 0:
-                    D = numpy.eye(len(base_vals))
-                elif len(dmats) > 0:
-                    D = form_matrix_product(dmats, alpha)
-                else:
-                    # special for vertex without defined point location
-                    assert pts == [()]
-                    D = numpy.eye(1)
-                result[alpha] = numpy.dot(self.coeffs,
-                                          numpy.dot(numpy.transpose(D),
-                                                    base_vals))
+        result = self.expansion_set._tabulate_jet(self.embedded_degree, pts, order=jet_order)
+        for alpha in result:
+            result[alpha] = numpy.dot(self.coeffs, result[alpha])
         return result
 
     def get_expansion_set(self):
@@ -175,16 +161,6 @@ def project(f, U, Q):
     return coeffs
 
 
-def form_matrix_product(mats, alpha):
-    """Forms product over mats[i]**alpha[i]"""
-    m = mats[0].shape[0]
-    result = numpy.eye(m)
-    for i in range(len(alpha)):
-        for j in range(alpha[i]):
-            result = numpy.dot(mats[i], result)
-    return result
-
-
 def polynomial_set_union_normalized(A, B):
     """Given polynomial sets A and B, constructs a new polynomial set
     whose span is the same as that of span(A) union span(B).  It may