Expressing barycentric coordinates symbolically in GEM

[achanbour]

This PR introduces new features that enable the barycentric coordinates of points expressed in cell-local reference coordinates to be expressed symbolically in GEM.

Previously, FIAT’s compute_barycentric_coordinates routines assumed that input points are NumPy arrays and performed eager NumPy operations such as dot and array slicing. Passing a gem.Node object resulted in NumPy doing object manipulations rather than constructing meaningful GEM expressions.

Broadly two sets of changes were made:

1. In FIAT.reference_element.py: new methods were introduced to enable the computation of barycentric coordinates on tensor product cells which previously did not exist. Additionally, all operations within those methods support both the case where the input array of points is a NumPy array or a GEM tensor expression. Some operations such as slicing in the case of tensor product cells required introducing gem operations which makes FIAT depended on GEM. These could potentially be removed if the equivalent NumPy operations can be re-implemented in GEM directly.

2. In gem.gem.py: minor changes were made to __matmul__. First, the method now supports multiplication between a GEM tensor and a NumPy array. Second, due to some compilation issues raised by loopy, a special case was introduced to handle tensor contractions over singleton dimensions. This essentially replaces the index summation with direct indexing which eliminates the need to build trivial contraction loops (iterating over indices of extent 1).

[pbrubeck]

Should we move the Ellipsis suport to gem.view?

[achanbour]

To me, gem.view looks like a lower-level operation which accepts one slice per each dimension to return a (new) tensor with the restricted view. The ellipsis, on the other hand, is more a syntactic convenience for indexing, mimicking that of NumPy. It basically just creates under hood the list of slices that view then accepts to produce the resulting tensor view. I would personally keep it in __getitem__.

[achanbour]

Update: gem.view returns a ComponentTensor with FlexiblyIndexed node which apparently gem's evaluator cannot evaluate. Instead of calling gem.view when doing slice indexing, an alternative is to build a ListTensor from the components directly:

fiat/gem/gem.py

Lines 131 to 146 in 76e63f7

	if any(isinstance(k, slice) for k in key):
	if len(key) != len(self.shape):
	raise IndexError("Expects one key per dimension.")
	# Unpack slices into ranges
	ranges = [
	range(
	k.start if k.start is not None else 0,
	k.stop if k.stop is not None else s,
	k.step if k.step is not None else 1
	) for k, s in zip(key, self.shape)
	]
	# Extract tensor indices from the cartesian product of ranges
	components = numpy.array(
	[Indexed(self, idx) for idx in itertools.product(*ranges)],
	dtype=object).reshape([len(r) for r in ranges])
	return as_gem(components)

This seems to work with evaluate.

[pbrubeck]

I think we should discuss a protocol to support symbolic computations with FIAT.

FIAT is written in numpy, so one approach is to assume that all inputs and outputs are (or can be casted as) numpy arrays. This is the approach taken in #172.

In particular, let's take as an example the main part of that PR, which is point_evaluation in finat/fiat_elements.py .

fiat/finat/fiat_elements.py

Lines 148 to 151 in 5c231cc

	Xi = tuple(gem.Indexed(refcoords, i) for i in np.ndindex(refcoords.shape))
	ps = PointSingleton(Xi)
	result = self.basis_evaluation(order, ps, entity=entity,
	coordinate_mapping=coordinate_mapping)

There, finat converts a tensor-valued GEM expression (a list of coordinates of shape (num_points, d)) into a d-tuple with each component (each element of the tuple is a the partially indexed view of the list of coordinates of shape (num_points,)). This tuple is passed onto FIAT, which outputs the tabulations as a (dict of) numpy array. Finally finat is the responsible to cast the result as a GEM ListTensor.

However, this means that we will have expanded expressions that involve the individual coordinates, that prevent GEM contractions along the coordinate Index. I don't think this is a big issue in dimension<=3, as the compiler unrolls IndexSums for small index extents anyway.

[achanbour]

To support GEM tensor expressions in barycentric computations, the main idea was to substitute NumPy operations with equivalent tensor operations that have implementations in both NumPy and GEM. This resulted in barycentric coordinate computations supporting inputs that are both NumPy arrays and GEM Nodes. Where array-style manipulations were required, equivalent functionality was introduced directly in GEM. An example of that is array slicing which appears in TensorProductCell.compute_axis_barycentric_coordinates.

The only explicit GEM dependency added to FIAT appears in Hypercube.compute_barycentric_coordinates, where facet permutations require indexing and rebuilding tensors.

fiat/FIAT/reference_element.py

Lines 1637 to 1646 in ddf6eb4

	if isinstance(tp_bary_coords[0], gem.Node):
	# breakpoint()
	components = [
	gem.Indexed(tp_bary_coords[axis], (bary_index,))
	for axis, bary_index in self.facet_perm
	]
	bary_coords = gem.ListTensor(components)
	else:
	components = [tp_bary_coords[axis][..., bary_index] for axis, bary_index in self.facet_perm]
	bary_coords = numpy.stack(components, axis=-1)

This raises a design question: whether it is acceptable to keep this GEM dependency in FIAT (since barycentric computations are inherently tied to reference cell geometry, which lives in FIAT), or whether it should instead be routed through FInAT to preserve the traditional separation between FIAT and GEM. Introducing an indirection layer in FInAT to dispatch back into FIAT for geometry-related symbolic operations would be a far more complex approach, at least in the context of the barycentric coordinate computation.

[pbrubeck]

Ideally we would like to support any input that can be casted into numpy.ndarray, we often work with points as list of tuples, mainly because numpy types are not immutable.

[pbrubeck]

Suggested change

	if isinstance(points, numpy.ndarray) and len(points) == 0:
	if numpy.size(points) == 0:

[pbrubeck]

Suggested change

	if isinstance(points, numpy.ndarray) and len(points) == 0:
	if numpy.size(points) == 0:

[pbrubeck]

Suggested change

point = np.asarray(point)

[pbrubeck]

Suggested change

point = np.asarray(point)

[pbrubeck]

Suggested change

point = np.asarray(point)