Refactor for arbitrary components in tensor prod

.github/workflows/test.yml

@@ -180,9 +180,20 @@ 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: |
           . venv/bin/activate
           cd ./fuse-repo
-          firedrake-run-split-tests 1 1 -n 8 "$EXTRA_PYTEST_ARGS" ./test
+          firedrake-run-split-tests 1 1 -n 8 "$EXTRA_PYTEST_ARGS" ./test

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
@@ -264,14 +265,14 @@ def evaluate(self, Qpts, bary_pts, Qwts, basis_change, immersed, dim, value_shap
         return Qpts, np.array(wts).astype(np.float64), comps
 
     def _to_dict(self):
-        o_dict = {"fn": self.fn}
+        o_dict = {"fn": self.fn, "syms": self.syms}
         return o_dict
 
     def dict_id(self):
         return "BarycentricPolynomialKernel"
 
     def _from_dict(obj_dict):
-        return BarycentricPolynomialKernel(obj_dict["fn"])
+        return BarycentricPolynomialKernel(obj_dict["fn"], symbols=obj_dict["syms"])
 
 
 class PolynomialKernel(BaseKernel):
@@ -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):
@@ -471,6 +466,7 @@ def immersed(pt):
             J_det = self.cell.attachment_J_det(self.cell.id, self.cell_defined_on.id)
             if not np.allclose(J_det, 1):
                 raise ValueError("Jacobian Determinant is not 1 did you do something wrong")
+            J_det = 1
             # if self.pairing.orientation:
             #     immersion = self.target_space.tabulate(wts, self.pairing.entity.orient(self.pairing.orientation))[0]
             # else:

fuse/element_construction.py

@@ -254,7 +254,6 @@ def vector_basis_fns(cell, deg, rot=False, interior_only=False):
 
     All transformation groups are S1 as these are interior to the cell.
     """
-    print(cell, deg)
     edge = cell.edges()[0]
     face = cell.d_entities(2)[0]
     facet_cell = edge
@@ -280,7 +279,7 @@ def vector_basis_fns(cell, deg, rot=False, interior_only=False):
                         xs1 = [DOF(L2Pairing(), BarycentricPolynomialKernel(pf_bf*new_bf(o), symbols=symbols))]
                         dofs += [DOFGenerator(xs1, pf_grp, S1)]
                         counter += (pf_grp).size()
-    print("facet dofs: ", counter)
+    # print("facet dofs: ", counter)
     counter = 0
     interior_deg = deg - 2
 
@@ -293,7 +292,7 @@ def vector_basis_fns(cell, deg, rot=False, interior_only=False):
                 counter += len(new_dofs)
                 dofs += [DOFGenerator(new_dofs, S1, S1)]
         interior_deg = deg - 3
-        print("interior facet dofs:", counter)
+        # print("interior facet dofs:", counter)
         counter = 0
 
     basis_funcs, groups, symbols = lagrange_barycentric_basis(cell.dimension, cell.ordered_vertex_coords(), interior_deg)
@@ -328,8 +327,8 @@ def vector_basis_fns(cell, deg, rot=False, interior_only=False):
                 xs1 = [DOF(L2Pairing(), BarycentricPolynomialKernel(np.prod(symbols)*vec_bf*bf, symbols=symbols))]
                 dofs += [DOFGenerator(xs1, grp, g2)]
                 counter += (grp).size()
-    print("interior dofs:", counter)
-    print("end cell", cell)
+    # print("interior dofs:", counter)
+    # print("end cell", cell)
     return dofs
 
 

fuse/enriched.py

@@ -10,12 +10,13 @@ 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):
         from fuse.tensor_products import TensorProductTriple
-        if not isinstance(A, TensorProductTriple) or not isinstance(B, TensorProductTriple):
+        valid_types = (TensorProductTriple, EnrichedElement)
+        if not isinstance(A, valid_types) or not isinstance(B, valid_types):
             raise ValueError("EnrichedElement should only be used for Tensor product elements. Use + between triples for enrichment.")
         self.A = A
         self.B = B
@@ -25,6 +26,9 @@ def __init__(self, A, B, flat=False, symmetric=True, matrices=True):
         if A.cell.flat != B.cell.flat:
             raise ValueError("Tensor products must both be flat or both not flat for enrichment.")
         self.cell = A.cell
+        self.flat = flat
+        if hasattr(A, "unflat_cell"):
+            self.unflat_cell = A.unflat_cell
         self.symmetric = symmetric
         self.apply_matrices = matrices
         if self.apply_matrices:
@@ -36,9 +40,20 @@ def __init__(self, A, B, flat=False, symmetric=True, matrices=True):
     def sub_elements(self):
         return [self.A, self.B]
 
+    def get_value_shape(self):
+        if str(self.spaces[1]) in ("HDiv", "HCurl"):
+            return (self.cell.get_spatial_dimension(),)
+        return super().get_value_shape()
+
     def __repr__(self):
         return "Enriched(%s, %s)" % (repr(self.A), repr(self.B))
 
+    def __add__(self, other):
+        assert self.spaces[0].set_shape == other.spaces[0].set_shape
+        assert str(self.spaces[1]) == str(other.spaces[1])
+        return EnrichedElement(self, other, symmetric=self.symmetric and other.symmetric,
+                               matrices=self.apply_matrices or other.apply_matrices)
+
     def setup_matrices(self):
         if self.cell.flat and not self.symmetric:
             raise NotImplementedError("Matrices for flattened cells that are not symmetric not supported")
@@ -50,8 +65,12 @@ def setup_matrices(self):
         else:
             cell = self.cell
         top = cell.to_fiat().get_topology()
+        seen_total_dims = set()
         for dim in top.keys():
             total_dim = sum(dim) if self.cell.flat else dim
+            if total_dim in seen_total_dims:
+                continue
+            seen_total_dims.add(total_dim)
             ents = self.entity_dofs[total_dim].keys()
             # comp_os = cell.component_orientations()
             for e_idx, e in enumerate(ents):
@@ -87,3 +106,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/groups.py

@@ -308,6 +308,16 @@ def __repr__(self):
             return self.name + str(self.size())
         return "GS" + str(self.size())
 
+    def _to_dict(self):
+        return {"members": [m.perm.array_form for m in self._members]}
+
+    def dict_id(self):
+        return "PermutationSet"
+
+    def _from_dict(o_dict):
+        perm_set = [Permutation(m) for m in o_dict["members"]]
+        return PermutationSetRepresentation(perm_set)
+
 
 class GroupRepresentation(PermutationSetRepresentation):
     """

fuse/serialisation.py

@@ -31,6 +31,7 @@ def __init__(self):
                           "Edge": Edge,
                           "Triple": ElementTriple,
                           "Group": GroupRepresentation,
+                          "PermutationSet": PermutationSetRepresentation,
                           "SobolevSpace": ElementSobolevSpace,
                           "InterpolationSpace": InterpolationSpace,
                           "PolynomialSpace": PolynomialSpace,
@@ -40,8 +41,11 @@ def __init__(self):
                           "DOFGen": DOFGenerator,
                           "Delta": DeltaPairing,
                           "L2Inner": L2Pairing,
+                          "BarycentricPolynomialKernel": BarycentricPolynomialKernel,
+                          "VectorKernel": VectorKernel,
                           "PolynomialKernel": PolynomialKernel,
                           "PointKernel": PointKernel,
+                          "ComponentKernel": ComponentKernel,
                           "Trace": Trace
                           }
 
@@ -65,6 +69,10 @@ def encode_traverse(self, obj, path=[]):
                 res_array[i] = dfs_res
             return res_array
 
+        # Some sympy objects are not hashable, so we must accept we may serialise them twice.
+        if isinstance(obj, sp.core.containers.Tuple) or isinstance(obj, sp.Expr) or isinstance(obj, sp.Matrix) or isinstance(obj, sp.Poly):
+            return "Sympy/" + sp.srepr(obj)
+
         if obj in self.seen_objs.keys():
             return self.seen_objs[obj]["id"]
 
@@ -75,9 +83,6 @@ def encode_traverse(self, obj, path=[]):
             self.store_obj(obj, obj.dict_id(), obj_id, obj_dict, path)
             return obj.dict_id() + "/" + str(obj_id)
 
-        if isinstance(obj, sp.core.containers.Tuple) or isinstance(obj, sp.Expr):
-            return "Sympy/" + sp.srepr(obj)
-
         return obj
 
     def get_id(self, obj):

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