Skip to content

Generalising compute barycentric coordinates#245

Draft
achanbour wants to merge 33 commits intomainfrom
achanbour/bary-coords-v2
Draft

Generalising compute barycentric coordinates#245
achanbour wants to merge 33 commits intomainfrom
achanbour/bary-coords-v2

Conversation

@achanbour
Copy link
Copy Markdown

This PR is based on the initial work done in #230. In particular, two sets of new features are introduced:

  1. Computing barycentric coordinates in Hypercubes

Currently, FIAT exposes a method for computing barycentric coordinates on simplicies only. The order in which these barycentric coordinates are listed follows the UFC ordering of facets in FIAT such that the following invariant holds: the i-th barycentric coordinate vanishes on facet i. In order for this invariant to hold not only on simplicies but also on other cells we need: 1) a method to compute barycentric coordinates on these cells 2) a permutation map that reorders the barycentric coordinates in facet order. For 1), I have implemented a method in TensorProductCell that computes barycentric coordinates on each factor and for 2) I have implemented a method compute_facet_permutation that computes the required permutation of coordinates for cells of type Hypercube.

  1. Supporting symbolic inputs (GEM nodes)

In order to ensure the above methods run not only on input points that are numpy arrays but also when these are gem.Node (e.g., runtime-known points), I have extended GEM's indexing interface to support indexing with a slice (needed when extracting coordinates along specific axes individually), an Ellipsis (e..g, points[...,s]) and (optionally) an array of integer indices mimicking numpy's fancy indexing (A[perm] where perm is an integer permutation array). The Ellipsis is expanded into slice(None) which is handled by gem.view which already handles slice indexing. In particular, this introduces a FlexiblyIndexed node in the output ComponentTensor for which I have implemented a handler in gem's evaluator for the tests to work.

To support the integer array indexing, I have implemented a new index type in GEM called ListIndex which can be used alongside Index (free index) and VariableIndex when indexing a GEM tensor. ListIndex acts as lookup table, wrapping a numpy array containing the specific index values to index the tensor at alongside a free index hat drives iteration over those values Indexed(A, (ListIndex(arr),)) means "for each position i in arr, access tensor[arr[I]]", with ListIndex.free_index as the loop variable ranging over len(arr). Indexing a GEM tensor with arr produces a ComponentTensor(Indexed(tensor, (ListIndex(arr),)), (ListIndex.free_index,)) converting the free index back into a shape dimension, such that the result is rank-1 GEM tensor representing the permuted elements. The tests at GEM's evaluator level confirm that ListIndex works as expected - though code generation fails since Loopy does not currently have support for ListIndex.

In the context of this PR's main theme, which is computing barycentric coordinates, whether we really need ListIndex depends on the return type of compute_axis_barycentric_coordinates - the flattened barycentric coordinates array to which the integer facet permutation is applied. If this is a numpy array then all works fine since we can just apply the permutation through numpy. However, if we want to return a GEM node such as gem.ListTensor for instance, and apply the permutation through GEM then we would need ListIndex.

…ion to handle nested tensor-product cells + added tests
@achanbour achanbour requested a review from connorjward April 24, 2026 15:28
Comment thread .gitignore Outdated
Comment thread FIAT/reference_element.py Outdated
Comment thread gem/gem.py Outdated
Comment thread gem/gem.py
if not isinstance(key, tuple):
key = (key,)

# Expand ellipsis -> fill in remaining dimensions with slice(None)
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't quite understand this behaviour. What if I have something like tensor[::2, ..., 3]? I wonder if we should only allow ... if there are no other indices provided.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WDYT?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Currently, there's no support for mixed indexing involving slices and arrays so something like tensor[..., [1,2,3]] isn't supported.

fiat/gem/gem.py

Lines 110 to 114 in a5b8ab0

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.")

The reason for that is that indexing by ... (or by slice) is handled by view() which expresses a reshape while indexing with [1,2,3] is handled via ListIndex which expresses a gather-like operation. I'm not sure if there's a Node in GEM that allows to compose these operations.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be nice to unify the code paths for __getitem__ and gem.view. I would definitely expect something like

tensor_2d[1::3, [2, 3, 4]]

to work

Comment thread gem/gem.py
if has_slice and has_array:
raise NotImplementedError("Mixed slice and array indexing is not supported.")

# Slice indexing -> delegate to view()
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if the implementation of view should live in here

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

then everything is together. We could even deprecate gem.view

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not have everything in gem.view? In this way we can keep the dunder methods in gem short and concise.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't mind that. I just want to unify the code paths somewhat.

Comment thread gem/gem.py Outdated
Comment thread gem/gem.py Outdated
Comment thread gem/gem.py
Comment thread FIAT/reference_element.py
Comment thread FIAT/reference_element.py Outdated
Comment on lines +618 to +619
if len(points) == 0:

if isinstance(points, numpy.ndarray) and len(points) == 0:
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should not start assuming that points may only be a numpy.ndarray. In FIAT most of the time points are treated as tuples, since those are immutable. We try to use numpy for our computations as much as possible, while relying on the casting done by most numpy functions.

Copy link
Copy Markdown

@pbrubeck pbrubeck left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suggest splitting this into two separate PRs. First to get the numerics right for barycentric coordinates on tensor product cells, and the second one getting the GEM side of things right.

Comment thread FIAT/reference_element.py Outdated
Comment thread gem/gem.py Outdated
assert isinstance(index, IndexBase)
if isinstance(index, Index):
index.set_extent(extent)
elif isinstance(index, ListIndex):
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There might be an issue here. Indexed inherits from Scalar. I think we want Indexed(Node, ListIndex) to evaluate to a tensor, so we might need a new class, maybe something similar to ComponentTensor.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If ListIndex is a subclass of Index (my most recent change ) then the extent of this free index would be len(index.index_array) where index_array is the integer index array stored in ListIndex.

fiat/gem/gem.py

Lines 726 to 736 in 6bbddc5

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))

The evaluation into a tensor happens in __getitem__:

fiat/gem/gem.py

Lines 125 to 131 in 6bbddc5

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

Would this be correct?

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think I understand it the logic better now, ListIndex is like a normal Index, but taking values from a list, so I think it is fine to use it with Indexed.

Copy link
Copy Markdown
Author

@achanbour achanbour May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Speaking to @connorjward, he advised to create a Index inside ListIndex to range over the values from the index_array list rather than having it implicit as here so essentially:

  • Option 1 (current): tensor[list_index] is tensor[list_index.index_array[list_index]] where list_index is a free index ranging over 0,1..., len(index_array)

  • Option 2: tensor[list_index] is tensor[list_index.index_array[list_index.free_index]]

In option 2, the inheritance between ListIndex and Index as its parent class may cause issues so a solution could bee to have a separate parent class for both.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both options I think are equivalent but option 2 is clearer.

Comment thread gem/interpreter.py
return Result(val[idx], all_fids)


@_evaluate.register(gem.FlexiblyIndexed)
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've opened #175, which is very similar but didn't merge it because I didn't add tests

Comment thread FIAT/reference_element.py
return unflattening_map


def compute_facet_permutation(unflattening_map, product):
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What are the possible permutations for Prisms and Cubes? Don't they end up in a trivial permutation?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you mean by trivial permutation?

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think he means an identity permutation. I would recommend that you add some sort of check to catch those somewhere.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's identity for simplicies. As written compute_facet_permutation only works for Hypercubes since we only have unflattening_map there. For general TP cells like prisms, my assumptions is that something similar to an unflattening map should be implemented first before generalising compute_facet_permutation.

Copy link
Copy Markdown

@pbrubeck pbrubeck May 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for adding the comment with the quad example. I see the issue now. The problem is that the UFCInterval does not order the facets in the same way UFCTriangle and UFCTetrahedron do. UFCInterval has the vertices as the facets, and vertices are always ordered in vertex order, but facets are ordered by the vertex they exclude, so you get an inconsistency for 1D when the facets are also the vertices.

I think the implementation of compute_barycentric_coordinates on UFCInterval should apply the permutation [1, 0]. Then this gives us the right ordering for UFCQuadrilateral.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The expected ordering for compute_barycentric_coordinates on simplices is not facet-based, but always vertex-based. If you would like a facet-based ordering, then only the interval should be permuted and tensor products would get the facet-based ordering for free.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Pablo. It is indeed very confusing but with your help I am able to see better. I'll try to re-iterate the reasoning here for the sake of clarity.

The expected ordering for compute_barycentric_coordinates on simplices is not facet-based, but always vertex-based.

  • This is consistent with the fact that $\lambda_i(P_j) = \delta_{ij}$ i.e., the i-th barycentric coordinate is 1 on vertex i and 0 on all other vertices. Hence it is 0 on the facet excluding vertex i.
  • Since in UFCTriangle and UFCTetrahedron order facets by the vertex they exclude we get that facet i excludes vertex i, i.e., $\lambda_i = 0$ on facet i.
  • Since in UFCInterval the vertices are the facets, we have that facet 1 contains vertex 1 and facet 2 contains vertex 2. This translates to $\lambda_1$ vanishing on facet 2 and $\lambda_2$ vanishing on facet 1 (reversed).

only the interval should be permuted and tensor products would get the facet-based ordering for free.

  • Since compute_factor_barycentric_coordinates recurses into calling compute_barycentric_coordinates on simplices forming the factors of the tensor-product cell, we just need to ensure that the facet-ordering is preserved on 1D intervals as the facet ordering already holds in all other simplices like UFCTriangle and UFCTetrahedron.

@achanbour
Copy link
Copy Markdown
Author

@connorjward I changed the definition of ListIndex in gem/gem.py from being itself a free Index to wrapping a free_index (Index) to iterate over the index array it stores. This eliminated the need to specialise the optimisation / folding rules over the free indices that already exist in GEM such as

fiat/gem/gem.py

Lines 795 to 797 in 82d51a9

# 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.

and similarly those is gem/optimise.py (e.g., replace_indices_indexed).

There's one behaviour that I noted though when testing it through the compilation pipeline. Suppose gem_expr = ComponentTensor(Indexed(ListTensor, ListIndex), Index(13)) where Index(13) is ListIndex.free_index which is what

fiat/gem/gem.py

Lines 125 to 131 in 82d51a9

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[perm[i]] is a scalar
retval = ComponentTensor(indexed, (list_index.free_index,)) # CT(A[perm[i]], i) -> A[perm]
return retval

currently produces. Now if I then do evaluation_expr = Indexed(gem_expr, Index('k')) so that I can then evaluate the CT into an output tensor A indexed by the same free index Index('k') (effectively doing A[k] = gem_expr[k]), I notice that the substitution rule of Indexed (that I reference above) does the following: it tries substituting Index(13) with Index(k) but since ListIndex does not expose its internal free index Index(k) ends up being discarded leaving Index(13) as the free index (which is then used as loop index in the accumulation A[k] = gem_expr[k]). This doesn't raise an error when gem_expr is a ComponentTensor but when it is a ListTensor which is what compute_factor_barycentric_coordinates in fiat/reference_element.py for general TP cells (non-Hypercubes) returns then what happens is that the evaluation expression evaluation_expr ends up having two free indices: Index('k') (the loop index) and Index(13) (the internal free index attached to ListIndex) which then triggers an error at impero level

fiat/gem/impero.py

Lines 95 to 103 in 82d51a9

class Return(Terminal):
"""Save value of GEM expression into an lvalue. Used to "return"
values from a kernel."""
__slots__ = ('variable', 'expression')
__front__ = ('variable', 'expression')
def __init__(self, variable, expression):
assert set(variable.free_indices) >= set(expression.free_indices)

I think returning a ListTensor is a temporary solution (until I find the permutation to apply for general TP cells) since ultimately I think we want all compute_barycentric_coordinates methods to return a ComponentTensor in which case its substitution rules would handle free indices correctly and the compilation into A[k] = gem_expr[k] would work.

@connorjward
Copy link
Copy Markdown

but since ListIndex does not expose its internal free index Index(k) ends up being discarded leaving Index(13) as the free index (which is then used as loop index in the accumulation A[k] = gem_expr[k]).

I think you need to set free_indices on ListIndex:

self.free_indices = (self.free_index,)

If you grep gem.py for free_indices you'll see that it's this recursively defined object. So I think we just need to plug into that infrastructure.

@achanbour
Copy link
Copy Markdown
Author

achanbour commented May 6, 2026

but since ListIndex does not expose its internal free index Index(k) ends up being discarded leaving Index(13) as the free index (which is then used as loop index in the accumulation A[k] = gem_expr[k]).

I think you need to set free_indices on ListIndex:

self.free_indices = (self.free_index,)

If you grep gem.py for free_indices you'll see that it's this recursively defined object. So I think we just need to plug into that infrastructure.

Tried this and it doesn't change anything to the behaviour I described above.

Just to clarify: ListIndex is still treated as an index i.e., you can use it when indexing a GEM Node Indexed(Node, ListIndex). Also its free index gets appended to the free_indices of Indexed:

fiat/gem/gem.py

Lines 826 to 833 in 82d51a9

for i in multiindex:
if isinstance(i, Index):
new_indices.append(i)
elif isinstance(i, VariableIndex):
new_indices.extend(i.expression.free_indices)
elif isinstance(i, ListIndex):
new_indices.append(i.free_index)
self.free_indices = unique(aggregate.free_indices + tuple(new_indices))

which explains why Indexed(ListTensor((Indexed(CT, ListIndex.free_index), Index(k)) has ListIndex.free_index and Index(k) in Indexed.free_indices. This produces the impero error I mention above. It does not happen when indexing a ComponentTensor since, in that case, the substitution rule ensures that Index(k) is discarded and only the free index of ListIndex remains so we have a single free index in the expression we're evaluating.

@connorjward
Copy link
Copy Markdown

Are you treating ListIndex as a free index in places? I think that that might be wrong. A ListIndex carries a free index but is not itself one.

Comment thread FIAT/reference_element.py Outdated
Comment thread FIAT/reference_element.py Outdated
Comment thread FIAT/reference_element.py
@achanbour
Copy link
Copy Markdown
Author

achanbour commented May 6, 2026

ListIndex

I don't think I do other than I have it sub-classed from IndexBase so that I can use it inside an Indexed node. Other than that, I refer to ListIndex.free_index every time.

Regarding my first sentence: I wonder if instead of Indexed(ComponentTensor, ListIndex) you meant having something like Indexed(ComponentTensor, ListIndex.free_index). But in that case, the code wouldn't be able to tell that ListIndex.free_index is not just any other free index but one attached to ListIndex. So I think indexing with ListIndex is fine and the logic about how that gets handled involves ListIndex.free_index.

Comment thread FIAT/reference_element.py Outdated
@achanbour
Copy link
Copy Markdown
Author

@pbrubeck Thanks again for helping me work this out. I've now made the required changes such that the barycentric coordinates are now listed in facet order, consistently across all cells. I added/modified the tests in test_reference_element.py to ensure it works on simplicies and hypercubes. I believe this should now generalise to other TP cells easily? Will add further tests shortly.

@achanbour
Copy link
Copy Markdown
Author

achanbour commented May 6, 2026

@pbrubeck Thanks again for helping me work this out. I've now made the required changes such that the barycentric coordinates are now listed in facet order, consistently across all cells. I added/modified the tests in test_reference_element.py to ensure it works on simplicies and hypercubes. I believe this should now generalise to other TP cells easily? Will add further tests shortly.

I have just added a test in test_reference_element.py called test_tp_bary_coords_are_in_facet_order which essentially verifies the facet ordering of barycentric coordinates in general (non-hypercube) tensor-product cells. It calls compute_factor_barycentric_coordinates which returns the barycentric coordinates grouped by factor. It then checks that along each factor the barycentric coordinates are listed in facet order.

I cannot think right now of a better way to test this since entities in general TP cells are not indexed by integers but by dimension tuples. For example in prisms which is a triangle x interval we have:

(1, 1) is edge x edge = the 3 quadrilateral side faces (triangle edge x interval)
(2, 0) is triangle x vertex = the 2 triangular faces

The concept of "bary coord i vanishes on facet i" only makes sense for cells that have a flat integer facet numbering like SimplicialComplex and Hypercube. Therefore, testing on a per-factor basis seems to be most reasonable in this case.

@pbrubeck
Copy link
Copy Markdown

pbrubeck commented May 6, 2026

I cannot think right now of a better way to test this since entities in general TP cells are not indexed by integers but by dimension tuples. For example in prisms which is a triangle x interval we have:

(1, 1) is edge x edge = the 3 quadrilateral side faces (triangle edge x interval)
(2, 0) is triangle x vertex = the 2 triangular faces

The concept of "bary coord i vanishes on facet i" only makes sense for cells that have a flat integer facet numbering like SimplicialComplex and Hypercube. Therefore, testing on a per-factor basis seems to be most reasonable in this case.

The facet numbering that we use is lexicographical, first ordered by the tuple dimension, then by entity number:

ordered_facets = [(dim, entity) for dim in sorted(top) for entity in sorted(top[dim]) if sum(dim) == sd - 1]

Copy link
Copy Markdown

@pbrubeck pbrubeck left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cannot make code suggestions at the moment. To be more general, the permutation we want might not always be indices[::-1]. Instead we need something along the lines of

        if facet_ordering:
            facet_ids = self.connectivity[(entity_dim, entity_dim-1)][entity_id]
            facets = [top[entity_dim-1][f] for f in facet_ids]
            perm = [facets.index(tuple(set(indices) - {v})) for v in indices]
            indices = [indices[p] for p in perm]

And with this we could set facet_ordering=True for any dimension

@achanbour
Copy link
Copy Markdown
Author

I cannot make code suggestions at the moment. To be more general, the permutation we want might not always be indices[::-1]. Instead we need something along the lines of

        if facet_ordering:
            facet_ids = self.connectivity[(entity_dim, entity_dim-1)][entity_id]
            facets = [top[entity_dim-1][f] for f in facet_ids]
            perm = [facets.index(tuple(set(indices) - {v})) for v in indices]
            indices = [indices[p] for p in perm]

And with this we could set facet_ordering=True for any dimension

Yeah this makes sense to me. I am not sure how this would extend to multi-cell simplicies as I haven't worked with these yet but I noticed one thing: indices was set earlier to be the local vertex indices on the sub-entity: indices = [i for i, v in enumerate(top[sd][cell_id]) if v in subcomplex] but I guess that facets here expresses each facet using the vertex ids of the whole simplex (since top = self.get_topology()). Would we want instead to built perm out of global_indices = top[entity_dim][entity_id] and since indices is built out of the same order the permutation would still apply to it?

… sub-entities of simplicies + added notes + modified test checking facet ordering of bary coords on tp cells
@achanbour
Copy link
Copy Markdown
Author

achanbour commented May 7, 2026

In latest commit, I have generalised the facet ordering permutation of barycentric coordinates in SimplicialComplex.compute_barycentric_coordinates following Pablo's comment to work for any entity_dim.

Since we allow computing barycentric coordinates on sub-entities, I think a missing piece here is to implement a sub-entity decomposition in the case of tensor product cells (i.e., express the sub-entity of a tensor product cell in terms of its factors)

fiat/FIAT/reference_element.py

Lines 1453 to 1456 in 8084e8a

# NOTE: entity is a tuple key of the tensor product cell's topology. It should be decomposed into per-factor sub-entities
# before being passed to factor.compute_barycentric_coordinates
if entity is not None:
raise NotImplementedError("Tensor-product sub-entity decomposition is not implemented yet.")

Following this, we may want to pass facet_ordering=True when computing barycentric coordinates on that sub-entity extending from the 1D case:

fiat/FIAT/reference_element.py

Lines 1463 to 1467 in 8084e8a

else:
# For higher dimensional simplicies: vertex ordering is already guaranteed.
# NOTE: But what if we want to compute barycentric coordinates on a sub-entity and still want facet ordering?
# Then we need to pass facet_ordering=True even when factor_sd > 1
factor_bary_coords = factor.compute_barycentric_coordinates(points[..., s], entity, rescale)

@pbrubeck
Copy link
Copy Markdown

pbrubeck commented May 7, 2026

Thanks for the hard work you've put into this. I still think we need to split this into FIAT-only and GEM-only PRs. Would that be okay?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants