Conflicting __eq__/__hash__ in basix.ufl._BlockedElement

[conpierce8]

Describe the issue

Using the gdim argument in basix.ufl.element leads to conflicting results from _BlockedElement.__eq__ and _BlockedElement.__hash__, as illustrated in the following MWE.

import basix, basix.ufl
elem1 = basix.ufl.element(basix.ElementFamily.P, basix.CellType.triangle, 1, shape=(2,))
elem2 = basix.ufl.element(basix.ElementFamily.P, basix.CellType.triangle, 1, shape=(2,), gdim=2)
assert (elem1 == elem2 and hash(elem1) == hash(elem2)) or (elem1 != elem2 and hash(elem1) != hash(elem2))

results in

Traceback (most recent call last):
  File "/root/mwe.py", line 4, in <module>
    assert (elem1 == elem2 and hash(elem1) == hash(elem2)) or (elem1 != elem2 and hash(elem1) != hash(elem2))
AssertionError

Expected behavior

This MWE should run successfully, since the hash values should be equal whenever __eq__ returns True.

The only required property is that objects which compare equal have the same hash value

Though not strictly required, I would expect them to be different when __eq__ returns False.

Proposed solution(s)

The solution could take one of the following two forms, depending on the intended functionality of the library:

1. Add an extra condition to _BlockedElement.__eq__ incorporating the gdim. This is the correct solution if two elements identical except for gdim should be considered different. (I.e. elem1 == elem2 evaluates to False in the above MWE.) I expect this is the desired solution, but wanted to confirm that this is the intended functionality.

       def __eq__(self, other) -> bool:
        """Check if two elements are equal."""
        return (
            isinstance(other, _BlockedElement) and self._block_size == other._block_size
            and self.block_shape == other.block_shape and self.sub_element == other.sub_element
            and self._gdim == other._gdim)

2. Remove any reference to gdim in the repr of _BlockedElement, since the hash is based on repr. (I.e. elem1 == elem2 evaluates to True in the above MWE.)

We may also need to take a look at _BasixElement, _ComponentElement, and _MixedElement to make sure they are handling gdim correctly. Furthermore, we might need to check _QuadratureElement for a similar issue, since the cell name and the pullback are included in the repr but not in __eq__.

Also related, the repr of _RealElement is wrong. Hopefully this MWE illustrates it sufficiently 😂 :

import basix, basix.ufl
real_element = basix.ufl.real_element(basix.CellType.triangle, (2,))
repr(real_element)

producing:

'RealElement(<functools._lru_cache_wrapper object at 0x7f51cfb3e770>)'

[conpierce8]

A further note: the inconsistency between __eq__ and __hash__ is indirectly responsible for some weird errors, as illustrated in the MWE below. I tracked this down to ufl.utils.sorting.topological_sorting(nodes, edges): when __eq__ returns True for a pair of finite elements but __hash__ returns a different value for each element, some of the elements in nodes can be erroneously omitted from the topological sort.

import basix, basix.ufl
import dolfinx as dfx
import ufl
from mpi4py import MPI

msh = dfx.mesh.create_unit_square(MPI.COMM_WORLD, 4, 4)
elem_scalar = basix.ufl.element(basix.ElementFamily.P, msh.basix_cell(), 1, gdim=2)
elem_vector = basix.ufl.blocked_element(elem_scalar, shape=(2,), gdim=2)
elem_tensor = basix.ufl.blocked_element(elem_scalar, shape=(2, 2), symmetry=True, gdim=2)
elem_mixed = basix.ufl.mixed_element([elem_vector, elem_tensor])
V = dfx.fem.functionspace(msh, elem_mixed)
u1, u2 = ufl.TrialFunctions(V)
v1, v2 = ufl.TestFunctions(V)
L = ufl.inner(ufl.grad(u1), v2) * ufl.dx
dfx.fem.form(L)

producing

Traceback (most recent call last):
  File "/root/mwe2.py", line 15, in <module>
    dfx.fem.form(L)
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/forms.py", line 188, in form
    return _create_form(form)
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/forms.py", line 183, in _create_form
    return _form(form)
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/forms.py", line 141, in _form
    ufcx_form, module, code = jit.ffcx_jit(mesh.comm, form,
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/jit.py", line 56, in mpi_jit
    return local_jit(*args, **kwargs)
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/jit.py", line 204, in ffcx_jit
    r = ffcx.codegeneration.jit.compile_forms([ufl_object], options=p_ffcx, **p_jit)
  File "/usr/local/lib/python3.10/dist-packages/ffcx/codegeneration/jit.py", line 199, in compile_forms
    raise e
  File "/usr/local/lib/python3.10/dist-packages/ffcx/codegeneration/jit.py", line 190, in compile_forms
    impl = _compile_objects(decl, forms, form_names, module_name, p, cache_dir,
  File "/usr/local/lib/python3.10/dist-packages/ffcx/codegeneration/jit.py", line 260, in _compile_objects
    _, code_body = ffcx.compiler.compile_ufl_objects(ufl_objects, prefix=module_name, options=options)
  File "/usr/local/lib/python3.10/dist-packages/ffcx/compiler.py", line 102, in compile_ufl_objects
    ir = compute_ir(analysis, object_names, prefix, options, visualise)
  File "/usr/local/lib/python3.10/dist-packages/ffcx/ir/representation.py", line 189, in compute_ir
    ir_elements = [_compute_element_ir(e, analysis.element_numbers, finite_element_names)
  File "/usr/local/lib/python3.10/dist-packages/ffcx/ir/representation.py", line 189, in <listcomp>
    ir_elements = [_compute_element_ir(e, analysis.element_numbers, finite_element_names)
  File "/usr/local/lib/python3.10/dist-packages/ffcx/ir/representation.py", line 241, in _compute_element_ir
    ir["sub_elements"] = [finite_element_names[e] for e in element.sub_elements]
  File "/usr/local/lib/python3.10/dist-packages/ffcx/ir/representation.py", line 241, in <listcomp>
    ir["sub_elements"] = [finite_element_names[e] for e in element.sub_elements]
KeyError: Basix element (P, triangle, 1, gll_warped, unset, False)

[mscroggs]

Yes, gdim should be added to the __eq__ checks. gdim was added later than some of the other inputs, and looks like we missed adding it there.

The other updates you suggest sound correct too.

Do you want to open a PR to fix these?

[conpierce8]

Do you want to open a PR to fix these?

Sure, I'll have a go at it.

[conpierce8]

I've opened a draft pull request with an initial implementation. Would it be good to add a test (maybe in test_ufl_wrapper.py?) to check that the hashes are always equal when __eq__ returns True?

[mscroggs]

I've opened a draft pull request with an initial implementation. Would it be good to add a test (maybe in test_ufl_wrapper.py?) to check that the hashes are always equal when __eq__ returns True?

Yes, it would be sensible to add a test there