Move grouping of integral from build_integral_data to group_form_integrals

[jorgensd]

Move group to resolve #304 by grouping prior to pull-back.
This will resolve FEniCS/ffcx#716 as well.

I think this is a better fix than what I came up with in https://github.com/FEniCS/ufl/pull/92/files
as we now group the integrals in the function that is supposed to do grouping, and it will reduce the load of the other functions (pull-back, geometry_lowering etc.) since it happens at the start of compute_form_data.

[jorgensd]

Seems like one should use the IndexRenumberingTransform:

ufl/ufl/algorithms/renumbering.py

Lines 35 to 39 in 004a678

	class IndexRenumberingTransformer(VariableRenumberingTransformer):
	"""Index renumbering transformer.
	This is a poorly designed algorithm. It is used in some tests,
	please do not use for anything else.

@wence- do you have any ideas on how to improve it?

[wence-]

Seems like one should use the IndexRenumberingTransform:

ufl/ufl/algorithms/renumbering.py

Lines 35 to 39 in 004a678

	class IndexRenumberingTransformer(VariableRenumberingTransformer):
	"""Index renumbering transformer.
	This is a poorly designed algorithm. It is used in some tests,
	please do not use for anything else.

@wence- do you have any ideas on how to improve it?

It doesn't look that bad. I would forward-port the model to the MultiFunction case and use map_expr_dags instead.

What you want to know is "are two dags equal to each other, up to index relabelling?". This is achievable by visiting the dag, child before parent, and replacing free indices with a unique new index.

class IndexRelabeller(MultiFunction):
    def __init__(self):
        super().__init__()
        count = itertools.count()
        self.index_cache = defaultdict(lambda: Index(next(count)))

    expr = MultiFunction.reuse_if_untouched
   
    def multi_index(self, o):
        return type(o)(tuple(self.index_cache[i] if isinstance(i, Index) else i for i in o.indices()))

    def zero(self, o):
        fi = o.ufl_free_indices
        fid = o.ufl_free_index_dimensions
        new_indices = [self.index_cache[Index(i)].count() for i in fi]
        new_fi, new_fid = zip(*sorted(zip(new_indices, fid), key=lambda x: x[0]))
        return type(o)(o.ufl_shape, tuple(new_fi), tuple(new_fid))

map_expr_dags(IndexRelabeller(), whatever)

Would probably work, untested.

[jorgensd]

Seems like one should use the IndexRenumberingTransform:

ufl/ufl/algorithms/renumbering.py

Lines 35 to 39 in 004a678

	class IndexRenumberingTransformer(VariableRenumberingTransformer):
	"""Index renumbering transformer.
	This is a poorly designed algorithm. It is used in some tests,
	please do not use for anything else.

@wence- do you have any ideas on how to improve it?

It doesn't look that bad. I would forward-port the model to the MultiFunction case and use map_expr_dags instead.

What you want to know is "are two dags equal to each other, up to index relabelling?". This is achievable by visiting the dag, child before parent, and replacing free indices with a unique new index.

class IndexRelabeller(MultiFunction):
    def __init__(self):
        super().__init__()
        count = itertools.count()
        self.index_cache = defaultdict(lambda: Index(next(count)))

    expr = MultiFunction.reuse_if_untouched
   
    def multi_index(self, o):
        return type(o)(tuple(self.index_cache[i] if isinstance(i, Index) else i for i in o.indices()))

    def zero(self, o):
        fi = o.ufl_free_indices
        fid = o.ufl_free_index_dimensions
        new_indices = [self.index_cache[Index(i)].count() for i in fi]
        new_fi, new_fid = zip(*sorted(zip(new_indices, fid), key=lambda x: x[0]))
        return type(o)(o.ufl_shape, tuple(new_fi), tuple(new_fid))

map_expr_dags(IndexRelabeller(), whatever)

Would probably work, untested.

Thanks Lawrence! I will test this tomorrow!

[jorgensd]

Thanks for the suggestions @wence- . I have now implemented this in the code, replacing the old transformer.
The test reduces the number of integrals computed in integral data from 3 on main to 1.
This should give a significant performance boost for both Firedrake and DOLFINx for problems with multiple subdomain ids, as there will be less forms to apply attach_estimated_degrees, apply_function_pullbacks, apply_integral_scaling etc. on, as well as fewer kernels generated by TSFC/FFCx

[garth-wells]

@jorgensd is this ready, or does it need more feedback?

[jorgensd]

@jorgensd is this ready, or does it need more feedback?

I believe this is ready for merging. @jpdean have you tested it for the use-cases that spawned revisiting this issue?

[jpdean]

Yes, this has resolved FEniCS/ffcx#716!