Refactor for arbitrary components in tensor prod

.github/workflows/test.yml

@@ -180,6 +180,17 @@ jobs:
         run: |
           . venv/bin/activate
           python3 -m pip install --break-system-packages -e './fiat-repo'
+
+      - uses: actions/checkout@v5
+        with:
+          repository: firedrakeproject/ufl
+          path: ufl-repo
+          ref: refs/heads/indiamai/hash_directional_sobolev
+
+      - name: Install checked out UFL
+        run: |
+          . venv/bin/activate
+          python3 -m pip install --break-system-packages -e './ufl-repo'
 
       - name: Run tests
         run: |

fuse/dof.py

@@ -3,6 +3,7 @@
 from fuse.traces import TrH1
 import numpy as np
 import sympy as sp
+import numbers
 
 
 class Pairing():
@@ -35,7 +36,7 @@ def __call__(self, kernel, v, cell):
         return v(*kernel.pt)
 
     def tabulate(self):
-        return 1
+        return np.eye(self.entity.dim())
 
     def add_entity(self, entity):
         res = DeltaPairing()
@@ -80,8 +81,8 @@ def tabulate(self):
         if self.orientation:
             new_bvs = np.array(self.entity.orient(self.orientation).basis_vectors())
             basis_change = np.matmul(np.linalg.inv(new_bvs), bvs)
-            return basis_change
-        return np.eye(bvs.shape[0])
+            return (1/self.entity.volume())*basis_change
+        return (1/self.entity.volume())*np.eye(bvs.shape[0])
 
     def add_entity(self, entity):
         res = L2Pairing()
@@ -194,7 +195,7 @@ def evaluate(self, Qpts, Qwts, basis_change, immersed, dim, value_shape):
             comps = [[tuple()] for pt in Qpts]
         else:
             comps = [[(i,) for v in value_shape for i in range(v)] for pt in Qpts]
-        if isinstance(self.pt, tuple) or isinstance(self.pt, int):
+        if isinstance(self.pt, tuple) or isinstance(self.pt, numbers.Number):
             return Qpts, np.array([wt*self.pt for wt in Qwts]).astype(np.float64), comps
         if not immersed:
             return Qpts, np.array([wt*np.matmul(self.pt, basis_change) for wt in Qwts]).astype(np.float64), comps
@@ -432,13 +433,7 @@ def convert_to_fiat(self, ref_el, interpolant_degree, value_shape=tuple()):
     def to_quadrature(self, arg_degree, value_shape):
         Qpts, Qwts = self.cell_defined_on.quadrature(self.kernel.degree(arg_degree))
         Qwts = Qwts.reshape(Qwts.shape + (1,))
-        dim = self.cell_defined_on.get_spatial_dimension()
-        if dim > 0:
-            bvs = np.array(self.cell_defined_on.basis_vectors())
-            new_bvs = np.array(self.cell_defined_on.orient(self.pairing.orientation).basis_vectors())
-            basis_change = np.matmul(np.linalg.inv(new_bvs), bvs)
-        else:
-            basis_change = np.eye(dim)
+        basis_change = self.pairing.tabulate()
 
         if self.immersed and (isinstance(self.kernel, VectorKernel) or isinstance(self.kernel, BarycentricPolynomialKernel) or isinstance(self.kernel, PolynomialKernel)):
             def immersed(pt):

fuse/enriched.py

@@ -10,7 +10,7 @@ class EnrichedElement(ElementTriple):
     In general, FUSE element triples should be represented nodally,
     however this may not be possible for all constructions.
 
-    In particulr, we need to preserve tensor product structure.
+    In particular, we need to preserve tensor product structure.
     """
 
     def __init__(self, A, B, flat=False, symmetric=True, matrices=True):
@@ -87,3 +87,9 @@ def generate(self):
     def to_ufl(self):
         ufl_sub_elements = [e.to_ufl() for e in self.sub_elements]
         return finat.ufl.EnrichedElement(*ufl_sub_elements, triple=self)
+
+    def flatten(self):
+        return EnrichedElement(self.A.flatten(), self.B.flatten(), flat=True, symmetric=self.symmetric, matrices=self.matrices)
+
+    def unflatten(self):
+        return EnrichedElement(self.A.unflatten(), self.B.unflatten(), flat=False, symmetric=self.symmetric, matrices=self.matrices)

fuse/spaces/polynomial_spaces.py

@@ -133,6 +133,9 @@ def __hash__(self):
     def restrict(self, mindegree, maxdegree):
         return PolynomialSpace(maxdegree, contains=-1, mindegree=mindegree, set_shape=self.set_shape)
 
+    def to_vector(self):
+        return PolynomialSpace(self.maxdegree, self.contains, self.mindegree, set_shape=True)
+
     def _to_dict(self):
         return {"set_shape": self.set_shape, "min": self.mindegree, "contains": self.contains, "max": self.maxdegree}
 
@@ -231,6 +234,9 @@ def __add__(self, x):
         s.extend([x])
         return ConstructedPolynomialSpace(w, s)
 
+    def to_vector(self):
+        return ConstructedPolynomialSpace(self.weights, [space.to_vector() for space in self.spaces])
+
     def _to_dict(self):
         super_dict = super(ConstructedPolynomialSpace, self)._to_dict()
         super_dict["spaces"] = self.spaces

fuse/tensor_products.py

@@ -4,18 +4,21 @@
 from fuse.enriched import EnrichedElement
 import numpy as np
 from finat.ufl import TensorProductElement, FuseElement, HDivElement, HCurlElement
+from itertools import product
+from functools import reduce
+from collections import defaultdict
 
 
-def tensor_product(A, B, matrices=True):
-    if not (isinstance(A, ElementTriple) and isinstance(B, ElementTriple)):
-        raise ValueError("Both components of Tensor Product need to be a Fuse Triple.")
-    return TensorProductTriple(A, B, matrices=matrices)
+def tensor_product(*factors, matrices=True):
+    if not all(isinstance(f, ElementTriple) for f in factors):
+        raise ValueError("All components of Tensor Product need to be a Fuse Triple.")
+    return TensorProductTriple(*factors, matrices=matrices)
 
 
-def symmetric_tensor_product(A, B, matrices=True):
-    if not (isinstance(A, ElementTriple) and isinstance(B, ElementTriple)):
-        raise ValueError("Both components of Tensor Product need to be a Fuse Triple.")
-    return TensorProductTriple(A, B, matrices=matrices, symmetric=True)
+def symmetric_tensor_product(*factors, matrices=True):
+    if not all(isinstance(f, ElementTriple) for f in factors):
+        raise ValueError("All components of Tensor Product need to be a Fuse Triple.")
+    return TensorProductTriple(*factors, matrices=matrices, symmetric=True)
 
 
 def flatten_dictionary(tensor_dict):
@@ -34,15 +37,14 @@ def flatten_dictionary(tensor_dict):
 
 class TensorProductTriple(ElementTriple):
 
-    def __init__(self, A, B, flat=False, symmetric=True, matrices=True):
-        self.A = A
-        self.B = B
+    def __init__(self, *factors, flat=False, symmetric=True, matrices=True):
+        self.factors = factors
         self.spaces = []
-        for (a, b) in zip(self.A.spaces, self.B.spaces):
-            self.spaces.append(a if a >= b else b)
+        for i in range(len(self.factors[0].spaces)):
+            self.spaces.append(max(f.spaces[i] for f in self.factors))
 
-        self.DOFGenerator = [A.DOFGenerator, B.DOFGenerator]
-        self.cell = TensorProductPoint(A.cell, B.cell)
+        self.DOFGenerator = [f.DOFGenerator for f in self.factors]
+        self.cell = TensorProductPoint(*[f.cell for f in factors])
         self.symmetric = symmetric
         self.flat = flat
         if self.flat:
@@ -59,51 +61,49 @@ def __init__(self, A, B, flat=False, symmetric=True, matrices=True):
 
     @property
     def sub_elements(self):
-        return [self.A, self.B]
+        return self.factors
 
     def __repr__(self):
-        return "TensorProd(%s, %s)" % (repr(self.A), repr(self.B))
+        return f"TensorProd({','.join(['{}' for f in self.factors])})".format(*(repr(f) for f in self.factors))
+
+    def _entity_associations(self, dofs, overall=True):
+        return self.entity_assocs, None, None
 
     def setup_matrices(self):
-        if self.A.cell.dimension > 1 or self.B.cell.dimension > 1:
-            raise NotImplementedError("Combining of matrices not implemented in 3D")
         if self.cell.flat and not self.symmetric:
             raise NotImplementedError("Matrices for flattened cells that are not symmetric not supported")
-        self.A.to_ufl()
-        self.B.to_ufl()
+        for f in self.factors:
+            f.to_ufl()
         oriented_mats_by_entity, flat_by_entity = self._initialise_entity_dicts(self.generate(), tensor=True)
         if self.flat:
             cell = self.unflat_cell
         else:
             cell = self.cell
         top = cell.to_fiat().get_topology()
-        for dim in top.keys():
-            total_dim = sum(dim) if self.flat else dim
-            a_ents = self.A.cell.get_topology()[dim[0]].keys()
-            b_ents = self.B.cell.get_topology()[dim[1]].keys()
-            ents = [(a, b) for a in a_ents for b in b_ents]
-            comp_os = cell.component_orientations()
-            for e, (a, b) in enumerate(ents):
-                ent_dofs = self.entity_dofs[total_dim][self.ent_mapping[dim][(a, b)]]
-                if len(ent_dofs) >= 1:
-                    sub_mat = oriented_mats_by_entity[dim][e]
-                    a_mat = self.A.matrices[dim[0]][a]
-                    a_ent_ids = self.A.entity_ids[dim[0]][a]
-                    b_mat = self.B.matrices[dim[1]][b]
-                    b_ent_ids = self.B.entity_ids[dim[1]][b]
-
-                    os = [(o_a, o_b) for o_a in a_mat.keys() for o_b in b_mat.keys()]
-                    for o in os:
-                        a_sub_mat = a_mat[o[0]][np.ix_(a_ent_ids, a_ent_ids)]
-                        b_sub_mat = b_mat[o[1]][np.ix_(b_ent_ids, b_ent_ids)]
-                        if self.mat_transformer is not None:
-                            o_classes = (self.A.cell.group.get_member_by_val(o[0]), self.B.cell.group.get_member_by_val(o[1]))
-                            combined_sub_mat = self.mat_transformer(a_sub_mat, b_sub_mat, o_classes)
-                        else:
-                            combined_sub_mat = np.kron(a_sub_mat, b_sub_mat)
-                        new_o = comp_os[dim][o]
-                        sub_mat[new_o][np.ix_(ent_dofs, ent_dofs)] = np.matmul(sub_mat[new_o][np.ix_(ent_dofs, ent_dofs)], combined_sub_mat)
-                        # sub_mat[new_o][np.ix_(ent_dofs, ent_dofs)] = np.eye(np.matmul(sub_mat[new_o][np.ix_(ent_dofs, ent_dofs)], combined_sub_mat).shape[0])
+        if len(self.factors) == 2:
+            for dim in top.keys():
+                total_dim = sum(dim) if self.flat else dim
+                f_ents = [f.cell.get_topology()[d].keys() for f, d in zip(self.factors, dim)]
+                ents = list(product(*(f_ents)))
+                comp_os = cell.component_orientations()
+                for e, sub_ents in enumerate(ents):
+                    ent_dofs = self.entity_dofs[total_dim][self.ent_mapping[dim][sub_ents]]
+                    if len(ent_dofs) >= 1:
+                        sub_mat = oriented_mats_by_entity[dim][e]
+                        mats = [f.matrices[d][ent] for f, d, ent in zip(self.factors, dim, sub_ents)]
+                        ent_ids = [f.entity_dofs[d][ent] for f, d, ent in zip(self.factors, dim, sub_ents)]
+                        os = list(product(*([mat.keys() for mat in mats])))
+                        for o in os:
+                            sub_mats = [mat[o_f][np.ix_(ent_id, ent_id)] for mat, o_f, ent_id in zip(mats, o, ent_ids)]
+                            if self.mat_transformer is not None:
+                                o_classes = (f.cell.group.get_member_by_val(o_f) for f, o_f in zip(self.factors, o))
+                                combined_sub_mat = self.mat_transformer(*sub_mats, o_classes)
+                            else:
+                                combined_sub_mat = reduce(lambda acc, x: np.kron(acc, x), sub_mats)
+                            new_o = comp_os[dim][o]
+                            if new_o in sub_mat.keys():
+                                sub_mat[new_o][np.ix_(ent_dofs, ent_dofs)] = np.matmul(sub_mat[new_o][np.ix_(ent_dofs, ent_dofs)], combined_sub_mat)
+                            # sub_mat[new_o][np.ix_(ent_dofs, ent_dofs)] = np.eye(np.matmul(sub_mat[new_o][np.ix_(ent_dofs, ent_dofs)], combined_sub_mat).shape[0])
 
         if self.cell.flat:
             oriented_mats_by_entity = flatten_dictionary(oriented_mats_by_entity)
@@ -112,39 +112,41 @@ def setup_matrices(self):
         self.reversed_matrices = self.reverse_dof_perms(self.matrices)
 
     def generate(self):
-        a_dofs = self.A.generate()
-        b_dofs = self.B.generate()
-        a_ent_assocs, _, _ = self.A._entity_associations(a_dofs, overall=False)
-        b_ent_assocs, _, _ = self.B._entity_associations(b_dofs, overall=False)
+        dofs = [f.generate() for f in self.factors]
+        ent_assocs = [f._entity_associations(dofs_f, overall=False)[0] for f, dofs_f in zip(self.factors, dofs)]
         if self.flat:
             top = self.unflat_cell.to_fiat().get_topology()
         else:
             top = self.cell.to_fiat().get_topology()
-        self.entity_dofs = {}
+        self.entity_dofs = defaultdict(dict)
         self.ent_mapping = {}
+        self.entity_assocs = defaultdict(dict)
+        self.dof_ids = {}
         dofs = []
-        ent_counter = {}
+        ent_counter = defaultdict(lambda: 0)
         dof_counter = 0
         for dim in top.keys():
             total_dim = sum(dim) if self.flat else dim
-            ents_A = a_ent_assocs[dim[0]].keys()
-            ents_B = b_ent_assocs[dim[1]].keys()
-            if total_dim not in self.entity_dofs.keys():
-                self.entity_dofs[total_dim] = {}
-                ent_counter[total_dim] = 0
+            ents = [ent_assoc[d].keys() for ent_assoc, d in zip(ent_assocs, dim)]
+            # if total_dim not in self.entity_dofs.keys():
+            #     self.entity_dofs[total_dim] = {}
+            #     self.entity_assocs[total_dim] = {}
             self.ent_mapping[dim] = {}
             ent_list = []
-            for i, ent in enumerate([(a_e, b_e) for a_e in ents_A for b_e in ents_B]):
+            for i, ent in enumerate(list(product(*ents))):
                 self.ent_mapping[dim][ent] = i + ent_counter[total_dim] if self.flat else ent
-                self.entity_dofs[total_dim][self.ent_mapping[dim][ent]] = tuple()
+                self.entity_dofs[total_dim][self.ent_mapping[dim][ent]] = []
                 ent_list += [ent]
-            for a_e, b_e in ent_list:
-                a_dofs = [d for dofs in a_ent_assocs[dim[0]][a_e].values() for d in dofs]
-                b_dofs = [d for dofs in b_ent_assocs[dim[1]][b_e].values() for d in dofs]
-                new_dofs = [(a, b) for a in a_dofs for b in b_dofs]
+            for es in ent_list:
+                e_dofs = [[d for dofs in ent_assoc[d][e].values() for d in dofs] for ent_assoc, d, e in zip(ent_assocs, dim, es)]
+                new_dofs = list(product(*e_dofs))
                 dofs += new_dofs
-                self.entity_dofs[total_dim][self.ent_mapping[dim][(a_e, b_e)]] = [i + dof_counter for i in range(len(new_dofs))]
-                dof_counter += len(new_dofs)
+                dof_gens = "(" + "*".join([",".join(list(ent_assoc[d][e].keys())) for ent_assoc, d, e in zip(ent_assocs, dim, es)]) + ")"
+                self.entity_assocs[total_dim][self.ent_mapping[dim][es]] = {dof_gens: new_dofs}
+                self.entity_dofs[total_dim][self.ent_mapping[dim][es]] += [i + dof_counter for i in range(len(new_dofs))]
+                for d in new_dofs:
+                    self.dof_ids[d] = dof_counter
+                    dof_counter += 1
                 ent_counter[total_dim] += 1
 
         return dofs
@@ -163,10 +165,10 @@ def __add__(self, other):
         return EnrichedElement(self, other, symmetric=self.symmetric and other.symmetric, matrices=self.apply_matrices or other.apply_matrices)
 
     def flatten(self):
-        return TensorProductTriple(self.A, self.B, flat=True, symmetric=self.symmetric, matrices=self.apply_matrices)
+        return TensorProductTriple(*self.factors, flat=True, symmetric=self.symmetric, matrices=self.apply_matrices)
 
     def unflatten(self):
-        return TensorProductTriple(self.A, self.B, flat=False, symmetric=self.symmetric, matrices=self.apply_matrices)
+        return TensorProductTriple(*self.factors, flat=False, symmetric=self.symmetric, matrices=self.apply_matrices)
 
 
 def compute_matrix_transform(trace, cell, o):
@@ -188,13 +190,17 @@ def compute_matrix_transform(trace, cell, o):
 class HDiv(TensorProductTriple):
 
     def __init__(self, tensor_element):
+        self.base_element = tensor_element
         self.gem_transformer, self.mat_transformer = self.select_fuse_hdiv_transformer(tensor_element)
         self.trace = TrHDiv
-        super(HDiv, self).__init__(tensor_element.A, tensor_element.B, tensor_element.flat, tensor_element.symmetric, tensor_element.matrices)
-        self.spaces[1] = TrHDiv
+        super(HDiv, self).__init__(*tensor_element.factors, flat=tensor_element.flat, symmetric=tensor_element.symmetric, matrices=tensor_element.matrices)
+        # self.spaces = (self.spaces[0], TrHDiv, self.spaces[2])
 
     def to_ufl(self):
-        return HDivElement(super(HDiv, self).to_ufl(), self.gem_transformer)
+        return HDivElement(super(HDiv, self).to_ufl(), transform=self.gem_transformer)
+
+    def repr(self):
+        return "HDiv(" + super(HDiv, self).repr() + ")"
 
     def select_fuse_hdiv_transformer(self, element):
         # Assume: something x interval
@@ -211,14 +217,15 @@ def select_fuse_hdiv_transformer(self, element):
             # Make the scalar value the right hand rule normal on the
             # y-aligned edges.
             cell = element.sub_elements[1].cell
-            bv = cell.basis_vectors()[0]
-            return lambda v: [gem.Product(gem.Literal(bv[0]), v), gem.Zero()], lambda m_a, m_b, o: np.kron(transform(cell, o[1]) @ m_a, m_b)
+            bv = cell.basis_vectors()[0][0]
+            return lambda v: [gem.Product(gem.Literal(bv), v), gem.Zero()], lambda m_a, m_b, o: np.kron(transform(cell, o[1]) @ m_a, m_b)
         elif ks == (1, 0):
             # Make the scalar value the upward-pointing normal on the
             # x-aligned edges.
             cell = element.sub_elements[0].cell
             bv = cell.basis_vectors()[0][0]
-            return lambda v: [gem.Zero(), gem.Product(gem.Literal(bv), v)], lambda m_a, m_b, o: np.kron(m_a, transform(cell, o[0]) @ m_b)
+            return lambda v: [gem.Zero(), v], lambda m_a, m_b, o: np.kron(m_a, transform(cell, o[0]) @ m_b)
+            # return lambda v: [gem.Zero(), gem.Product(gem.Literal(bv), v)], lambda m_a, m_b, o: np.kron(m_a, transform(cell, o[0]) @ m_b)
         # elif ks == (2, 0):
         #     # Same for 3D, so z-plane.
         #     return lambda v: [gem.Zero(), gem.Zero(), v]
@@ -241,17 +248,26 @@ def select_fuse_hdiv_transformer(self, element):
             raise NotImplementedError("Unexpected original mapping!")
             assert False, "Unexpected form degree combination!"
 
+    def flatten(self):
+        return HDiv(self.base_element.flatten())
+
+    def unflatten(self):
+        return HDiv(self.base_element.unflatten())
+
 
 class HCurl(TensorProductTriple):
 
     def __init__(self, tensor_element):
         self.gem_transformer, self.mat_transformer = self.select_fuse_hcurl_transformer(tensor_element)
         self.trace = TrHCurl
-        super(HDiv, self).__init__(tensor_element.A, tensor_element.B, tensor_element.flat, tensor_element.symmetric, tensor_element.matrices)
-        self.spaces[1] = TrHCurl
+        super(HDiv, self).__init__(*tensor_element.factors, tensor_element.flat, tensor_element.symmetric, tensor_element.matrices)
+        self.spaces = (self.spaces[0], TrHCurl, self.spaces[2])
 
     def to_ufl(self):
-        return HCurlElement(super(HDiv, self).to_ufl(), self.gem_transformer)
+        return HCurlElement(super(HCurl, self).to_ufl(), self.gem_transformer)
+
+    def repr(self):
+        return "HCurl(" + super(HCurl, self).repr() + ")"
 
     def select_fuse_hcurl_transformer(element):
         import gem
@@ -292,3 +308,9 @@ def select_fuse_hcurl_transformer(element):
         else:
             raise NotImplementedError("Unexpected original mapping!")
             assert False, "Unexpected original mapping!"
+
+    def flatten(self):
+        return HCurl(self.base_element.flatten())
+
+    def unflatten(self):
+        return HCurl(self.base_element.unflatten())