Generalising compute barycentric coordinates

.gitignore

@@ -2,7 +2,7 @@
 
 /build/
 /dist/
-/fenics_fiat.egg-info/
+/firedrake_fiat.egg-info/
 
 /.cache/
 /doc/sphinx/source/api-doc

FIAT/reference_element.py

@@ -615,8 +615,10 @@ def get_dimension(self):
 
     def compute_barycentric_coordinates(self, points, entity=None, rescale=False):
         """Returns the barycentric coordinates of a list of points on the complex."""
-        if len(points) == 0:
+
+        if isinstance(points, numpy.ndarray) and len(points) == 0:
             return points
+
         if entity is None:
             entity = (self.get_spatial_dimension(), 0)
         entity_dim, entity_id = entity
@@ -640,8 +642,11 @@ def compute_barycentric_coordinates(self, points, entity=None, rescale=False):
             h = 1 / numpy.linalg.norm(A, axis=1)
             b *= h
             A *= h[:, None]
-        out = numpy.dot(points, A.T)
-        return numpy.add(out, b, out=out)
+        # out = numpy.dot(points, A.T)
+        out = points @ A.T
+
+        # return numpy.add(out, b)
+        return out + b
 
     def compute_bubble(self, points, entity=None):
         """Returns the lowest-order bubble on an entity evaluated at the given
@@ -1406,6 +1411,45 @@ def extrinsic_orientation_permutation_map(self):
     def is_macrocell(self):
         return any(c.is_macrocell() for c in self.cells)
 
+    def compute_factor_barycentric_coordinates(self, points, entity=None, rescale=False):
+        """Compute barycentric coordinates on each axis (factor) of a tensor-product cell.
+
+        Parameters
+        ----------
+        points: numpy.ndarray or GEM.Node
+            The reference coordinates of the points.
+
+        Returns
+        -------
+        numpy.ndarray
+            A flattened array of shape ``(total_bary_coords, )`` and dtype object if points are GEM nodes,
+            otherwise dtype numeric. The i-th entry contains the barycentric coordinates
+            on the i-th factor cell. If factor i is a simplex of dimension d, this will
+            have shape ``(npoints, d+1)``. If factor i is a hypercube of dimension d,
+            this will have shape ``(npoints, 2*d)``.
+        """
+        import gem
+
+        if isinstance(points, numpy.ndarray) and len(points) == 0:
+            return points
+
+        axis_dims = [c.get_spatial_dimension() for c in self.cells]
+        point_slices = TensorProductCell._split_slices(axis_dims)
+
+        result = []
+        for factor, s in zip(self.cells, point_slices):
+            result.append(factor.compute_barycentric_coordinates(points[..., s], entity, rescale))
+
+        # Flatten the array
+        # We cannot construct the flat array directly since we may not know upfront the total number
+        # of barycentric coordinates (e.g., in a simplex it is d+1, in a hypercube it is 2*d)
+        flat_result = numpy.array([bary[j] for bary in result for j in range(bary.shape[0])])
+
+        if isinstance(points, gem.Node):
+            return gem.as_gem(flat_result) # returns a ListTensor wrapping the scalar GEM expr. of bary coords.
+
+        return flat_result
+
 
 class Hypercube(Cell):
     """Abstract class for a reference hypercube"""
@@ -1423,6 +1467,8 @@ def __init__(self, dimension, product):
         self.product = product
         self.unflattening_map = compute_unflattening_map(pt)
 
+        self.facet_perm = compute_facet_permutation(self.unflattening_map, self.product)
+
     def get_dimension(self):
         """Returns the subelement dimension of the cell.  Same as the
         spatial dimension."""
@@ -1521,6 +1567,27 @@ def __ge__(self, other):
     def __le__(self, other):
         return self.product <= other
 
+    def compute_barycentric_coordinates(self, points, entity=None, rescale=False):
+        """Returns the barycentric coordinates of a list of points on the hypercube.
+
+        Parameters
+        ----------
+        points: numpy.ndarray or GEM.Node
+            The reference coordinates of the points.
+
+        Returns
+        -------
+        List of numpy.ndarray or GEM.ComponentTensor
+            Returns a list of barycentric coordinates in local facet order such that for any point
+            lying on local facet `lf` of the cell, the barycentric coordinate at index `lf` vanishes.
+        """
+        if isinstance(points, numpy.ndarray) and len(points) == 0:
+            return points
+
+        tp_bary_coords = self.product.compute_factor_barycentric_coordinates(points, entity, rescale)
+
+        return tp_bary_coords[self.facet_perm] # A[[1, 3, 5, 4]]
+
 
 class UFCHypercube(Hypercube):
     """Reference UFC Hypercube
@@ -1839,6 +1906,60 @@ def compute_unflattening_map(topology_dict):
     return unflattening_map
 
 
+def compute_facet_permutation(unflattening_map, product):
+    """
+    Returns a permutation mapping each facet of a `~.Hypercube` to the index of the
+    barycentric coordinate that vanishes on it.
+
+    The order of barycentric coordinates returned by `compute_factor_barycentric_coordinates`
+    is determined by axis structure, not by facet numbering. Reordering them by this permutation
+    yields the invariant: the i-th barycentric coordinate vanishes on the i-th facet.
+    """
+    # First compute axis offsets into the flattened barycentric coordinate array.
+    axis_offsets = []
+    offset = 0
+    for axis_cell in product.cells:
+        axis_offsets.append(offset)
+        offset += axis_cell.get_dimension() + 1
+
+    # Initialise the integer permutation array
+    sd = len(product.cells)
+    num_facets = 2 * sd
+    perm = numpy.zeros(num_facets, dtype=int)
+
+    for f in range(num_facets):
+        # Recover the tensor-product representation of the facet as given by the unflattening map
+        dim_tuple, tp_entity = unflattening_map[(sd - 1, f)]
+
+        # Determine the axis that's orthogonal to the facet
+        # E.g., in a quad:
+        # if dim_tuple = (0,1) -> facet has dimension 0 on the first component -> fixed at x = 0 or x = 1
+        # if dim_tuple = (1,0) -> facet has dimension 0 on the second component -> fixed at y = 0 or y = 1
+        axis = next(
+            i for i, d in enumerate(dim_tuple)
+            if d == product.cells[i].get_dimension() - 1
+        )
+
+        # Determine the index of the endpoint that produces the facet
+        # which gives the local facet number in the axis space
+        entity_shape = tuple(
+            len(c.get_topology()[d])
+            for c, d in zip(product.cells, dim_tuple)
+        )
+        tuple_ei = numpy.unravel_index(tp_entity, entity_shape)
+        local_facet = tuple_ei[axis]
+
+        # For a simplex (UFCInterval, UFCTriangle), the barycentric coordinate that vanishes on local facet i
+        # corresponds to the ID of the vertex that doesn't belong to that facet
+        all_vertices = set(product.cells[axis].get_topology()[0].keys())
+        facet_vertices = set(product.cells[axis].get_topology()[0][local_facet])
+        bary_index = next(iter(all_vertices - facet_vertices))
+
+        perm[f] = axis_offsets[axis] + bary_index
+
+    return perm
+
+
 def max_complex(complexes):
     max_cell = max(complexes)
     if all(max_cell >= b for b in complexes):

gem/gem.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 """GEM is the intermediate language of TSFC for describing
 tensor-valued mathematical expressions and tensor operations.
 It is similar to Einstein's notation.
@@ -20,6 +22,8 @@
 from operator import attrgetter
 from numbers import Integral, Number
 
+from types import EllipsisType
+
 import numpy
 from numpy import asarray
 
@@ -32,7 +36,7 @@
            'Variable', 'Sum', 'Product', 'Division', 'FloorDiv', 'Remainder', 'Power',
            'MathFunction', 'MinValue', 'MaxValue', 'Comparison',
            'LogicalNot', 'LogicalAnd', 'LogicalOr', 'Conditional',
-           'Index', 'VariableIndex', 'Indexed', 'ComponentTensor',
+           'Index', 'VariableIndex', 'ListIndex', 'Indexed', 'ComponentTensor',
            'IndexSum', 'ListTensor', 'Concatenate', 'Delta', 'OrientationVariableIndex',
            'index_sum', 'partial_indexed', 'reshape', 'view',
            'indices', 'as_gem', 'FlexiblyIndexed',
@@ -81,12 +85,53 @@ def is_equal(self, other):
             self.children = other.children
         return result
 
-    def __getitem__(self, indices):
-        try:
-            indices = tuple(indices)
-        except TypeError:
-            indices = (indices, )
-        return Indexed(self, indices)
+    def __getitem__(
+        self,
+        key: IndexT | tuple[IndexT, ...],
+    ) -> ComponentTensor | Indexed:
+        """A generalised interface for indexing GEM tensors"""
+        if not isinstance(key, tuple):
+            key = (key,)
+
+        # Expand ellipsis -> fill in remaining dimensions with slice(None)
+        if any(k is Ellipsis for k in key):
+            if key.count(Ellipsis) > 1:
+                raise NotImplementedError("Multiple ellipses are not supported.")
+            ellipsis_pos = key.index(Ellipsis)
+            remaining_dims = len(self.shape) - (len(key) - 1)
+            if remaining_dims < 0:
+                raise IndexError("Too many indices provided.")
+            key = (
+                key[:ellipsis_pos]
+                + (slice(None), ) * remaining_dims
+                + key[ellipsis_pos + 1:]
+            )
+
+        has_slice = any(isinstance(k, slice) for k in key)
+        has_array = any(isinstance(k, (numpy.ndarray, list)) for k in key)
+
+        if has_slice and has_array:
+            raise NotImplementedError("Mixed slice and array indexing is not supported.")
+
+        # Slice indexing -> delegate to view()
+        if has_slice:
+            # view expects one slice for each axis/dim of the tensor
+            if len(key) != len(self.shape):
+                raise IndexError("Expects the number of slices to match the gem.Node tensor rank")
+            return view(self, *key)
+
+        # Support a list or array of integer indices
+        # Previously, built a ListTensor out of Indexed nodes, one for each element of the permutation
+        if has_array:
+            arr_pos = next(i for i, k in enumerate(key) if isinstance(k, (numpy.ndarray, list)))
+            list_index = ListIndex(key[arr_pos])
+            new_key = key[:arr_pos] + (list_index,) + key[arr_pos+1:]
+            indexed = Indexed(self, new_key) # A[i] scalar
+            retval = ComponentTensor(indexed, (list_index,)) # CT(A[i], i) -> A[index_array[i]]
+            return retval
+
+        # Point indexing
+        return Indexed(self, key)
 
     def __neg__(self):
         return componentwise(Product, minus, self)
@@ -117,6 +162,7 @@ def __matmul__(self, other):
             raise ValueError("Both objects must have shape for matmul")
         elif self.shape[-1] != other.shape[0]:
             raise ValueError(f"Mismatching shapes {self.shape} and {other.shape} in matmul")
+
         *i, k = indices(len(self.shape))
         _, *j = indices(len(other.shape))
         expr = Product(Indexed(self, (*i, k)), Indexed(other, (k, *j)))
@@ -677,6 +723,45 @@ def __repr__(self):
     def __reduce__(self):
         return type(self), (self.expression,)
 
+class ListIndex(Index):
+    """ListIndex is a free index representing array indexing such that tensor[list_index] is tensor[index_array[list_index]]
+    where index_array is the integer index array and list_index is a free index ranging over 0,1..., len(index_array)"""
+
+    __slots__ = Index.__slots__ + ('index_array',)
+
+    def __init__(self, index_array, name=None):
+        index_array = numpy.asarray(index_array)
+        assert numpy.issubdtype(index_array.dtype, numpy.integer)
+        self.index_array = index_array
+        super().__init__(name=name, extent=len(index_array))
+
+    def __str__(self):
+        if self.name is None:
+            return f"{self.index_array.tolist()}[i_{self.count}]"
+        return f"{self.index_array.tolist()}[{self.name}]"
+
+    def __repr__(self):
+        if self.name is None:
+            return f"ListIndex({self.index_array.tolist()}, {self.count})"
+        return f"ListIndex({self.index_array.tolist()}, {self.name})"
+
+    # def __eq__(self, other):
+    #     if type(self) is not type(other):
+    #         return False
+    #     return numpy.array_equal(self.index_array, other.index_array)
+
+    # def __ne__(self, other):
+    #     return not self.__eq__(other)
+
+    # def __hash__(self):
+    #     return hash((type(self), self.index_array.tobytes()))
+
+    # def __reduce__(self):
+    #     return type(self), (self.index_array, )
+
+
+IndexT = int | Index | VariableIndex | ListIndex | slice | EllipsisType | list | numpy.ndarray
+
 
 class Indexed(Scalar):
     __slots__ = ('children', 'multiindex', 'indirect_children')
@@ -693,6 +778,9 @@ def __new__(cls, aggregate, multiindex):
             assert isinstance(index, IndexBase)
             if isinstance(index, Index):
                 index.set_extent(extent)
+            # elif isinstance(index, ListIndex):
+            #     if numpy.any(index.index_array < 0) or numpy.any(index.index_array >= extent):
+            #         raise IndexError("Invalid index in ListIndex")
             elif isinstance(index, int) and not (0 <= index < extent):
                 raise IndexError("Invalid literal index")
 
@@ -707,7 +795,7 @@ def __new__(cls, aggregate, multiindex):
         # Simplify Indexed(ComponentTensor(Indexed(C, kk), jj), ii) -> Indexed(C, ll)
         # This pattern corresponds to an index replacement rule jj -> ii applied to
         # the innermost multiindex kk to produce ll.
-        if isinstance(aggregate, ComponentTensor):
+        if isinstance(aggregate, ComponentTensor) and not any(isinstance(i, ListIndex) for i in aggregate.multiindex):
             B, = aggregate.children
             jj = aggregate.multiindex
             ii = multiindex
@@ -740,6 +828,8 @@ def __new__(cls, aggregate, multiindex):
                 new_indices.append(i)
             elif isinstance(i, VariableIndex):
                 new_indices.extend(i.expression.free_indices)
+            elif isinstance(i, ListIndex):
+                new_indices.append(i)
         self.free_indices = unique(aggregate.free_indices + tuple(new_indices))
 
         return self
@@ -752,6 +842,8 @@ def index_ordering(self):
                 free_indices.append(i)
             elif isinstance(i, VariableIndex):
                 free_indices.extend(i.expression.free_indices)
+            elif isinstance(i, ListIndex):
+                free_indices.append(i.free_index)
         return tuple(free_indices)
 
 
@@ -875,7 +967,9 @@ def __new__(cls, expression, multiindex):
 
         # Index folding
         if isinstance(expression, Indexed):
-            if multiindex == expression.multiindex:
+            # NOTE: The folding rule CT(A[i], i) -> A is correct for when i is an Index as it ranges over every position in A
+            # but is not right when i is a ListIndex since A[i] means element A[index_array[i]]
+            if multiindex == expression.multiindex and not any(isinstance(i, ListIndex) for i in multiindex):
                 return expression.children[0]
 
         self = super(ComponentTensor, cls).__new__(cls)