Fix Integral/BdIntegral JIT cache collisions

[gthyagi]

Summary

This PR fixes an order-dependent JIT callback bug in uw.maths.Integral(...) and uw.maths.BdIntegral(...).

The failure showed up most clearly on spherical-shell meshes: evaluating a volume integral and a boundary integral on the same live mesh could return 0.0 depending on which one was called first.

Changes in this PR:

• fix JIT cache-key collisions between volume and boundary callback builds in getext()
• preserve the primary-field list before hashing/codegen so mesh-variable boundary integrals continue to bind to petsc_u[...]
• add spherical-shell regression tests that exercise both call orders

Bug

The bug was not in the benchmark physics or in PETSc's spherical geometry handling. It was in Underworld's JIT extension cache.

Observed behaviour

On a spherical shell mesh, the following orders were inconsistent:

1. BdIntegral(...) only: correct
2. BdIntegral(...) then Integral(...): the later volume integral could return 0.0
3. Integral(...) then BdIntegral(...): the later boundary integral could return 0.0

This was reproducible with a minimal generic spherical-shell script using fn=1.0 and no Stokes solve.

Root cause

underworld3.utilities._jitextension.getext() was building its cache key from a flattened tuple of expanded expressions.

That loses the role of each expression:

• residual
• BC
• jacobian
• boundary residual
• boundary jacobian

So the same symbolic expression, for example 1.0, could be compiled in two incompatible callback contexts and still reuse the same cached module.

In practice that meant a boundary callback build could reuse a residual-only extension, or vice versa, which then wired the wrong callback slot into PETSc.

While fixing that, a second issue also appeared: primary_field_list can be a generator in these call paths. If it is consumed while constructing the cache signature, _createext() no longer sees the primary fields. That causes mesh-variable boundary integrals to compile against petsc_a[...] instead of petsc_u[...].

Fix

The fix is in src/underworld3/utilities/_jitextension.py.

1. Preserve callback role in the cache key

getext() now:

• materializes each callback list separately
• expands each callback list structurally by role
• builds a structured signature that keeps the five callback categories separate

That means identical expressions no longer collide across incompatible callback types.

2. Preserve the primary-field list for code generation

primary_field_list is converted to a tuple immediately.

This prevents generator consumption during cache-key construction and ensures _createext() still receives the full ordered set of primary fields for correct petsc_u[...] mapping.

Tests

Added two regression tests in tests/test_0502_boundary_integrals.py:

• test_spherical_bd_then_integral_does_not_poison_volume_path()
• test_spherical_integral_then_bd_does_not_poison_boundary_path()

These use a spherical shell mesh and assert that:

• BdIntegral(mesh, 1.0, "Lower") remains nonzero before and after a volume integral
• Integral(mesh, 1.0) remains nonzero after a boundary integral
• the boundary measure is order-independent

I also reran the full boundary-integral regression module after the fix:

cd /Users/tgol0006/uw_folder/uw3_git_gthyagi_latest/underworld3
/Users/tgol0006/.pixi/bin/pixi run -e amr-dev pytest tests/test_0502_boundary_integrals.py -q

Result:

• 21 passed

Why this matters

This is a core correctness issue in JIT callback dispatch. It is not limited to the Thieulot benchmark.

Any workflow that mixes Integral(...) and BdIntegral(...) on the same live mesh can be affected when the same symbolic form appears in both paths. The spherical-shell case made it easy to reproduce, but the underlying bug lives in shared JIT callback caching.

[lmoresi]

Excellent - this works well.

I'll merge this and I have some additional changes to add in a new PR