Custom JutulEquation types blocked by cross-term sparsity merge in MultiModel cases

[gs-arya]

Summary

When defining a custom JutulEquation subtype and registering it on a model inside a MultiModel that also has cross-terms (e.g. a reservoir + facility setup), the custom equation's declared sparsity pattern gets merged into a union pattern that downstream cross-term alignment then walks against caches that were never sized for it. The result is a Jacobian alignment failure at positions the custom equation declares but no cross-term physically writes:

Error: Unable to map cache entry to Jacobian, (<row>, <col>) not allocated in Jacobian matrix.
Error: Jacobian alignment failed.

The root cause appears to be in extra_cross_term_sparsity (in src/multimodel/crossterm.jl), where cross-term sparsity is keyed by equation symbol and applied to the union of all equations sharing that symbol — so a new equation registered under, say, a custom name on the reservoir model still has its declared face-stencil entries merged into the cross-term machinery's view of the cross-term's responsibility.

Motivation

The use case is per-cell variable switching (PVS) — the standard technique used by commercial thermal simulators (CMG STARS, Schlumberger Eclipse Thermal, tNavigator) to handle the narrow-boiling / phase-appearance regime in compositional thermal simulation. In that regime, dH/dT → 0 at the saturation plateau and the Jacobian becomes rank-deficient unless the primary variable is switched on a per-cell basis.

A common workaround in the literature is a "pin" equation: a small per-cell equation that, on cells flagged for switching, freezes one primary variable (e.g. molar enthalpy) at a target value while another (e.g. vapor saturation) becomes the active unknown. Implementing this as a JutulEquation is the most natural fit for Jutul's architecture — but the cross-term sparsity behavior blocks it.

Without a fix here, JutulDarcy users implementing thermal compositional models cannot deploy PVS, which limits the range of thermal recovery scenarios (steam injection / SAGD / wet-steam) that can be simulated reliably.

Reproducer (concept)

using Jutul, JutulDarcy

# Custom equation type
struct PinEquation{D<:Jutul.FlowDiscretization} <: Jutul.JutulEquation
    flow_discretization::D
    target_symbol::Symbol
    pin_scale::Float64
end

# Standard interface methods:
# - associated_entity
# - number_of_equations_per_entity (returns 1)
# - declare_pattern (face-stencil mirroring ConservationLaw via PotentialFlow.half_face_map)
# - update_equation_in_entity! (writes ONLY diagonal: eq_buf[1] = local residual)

# Custom storage that mirrors the face-stencil declaration:
# - diagonal CompactAutoDiffCache
# - extra_pos::Matrix{Int} for face-stencil neighbours
# (Custom align_to_jacobian! using diagonal_alignment! + find_jac_position
#  loop for extras. Custom update_linearized_system_equation! writing diagonal
#  residuals + zero into extras.)

# Register on a reservoir SimulationModel that already has ConservationLaws + cross-terms:
res = model.models[:Reservoir]
res.equations[:custom_pin] = PinEquation(res.domain.discretizations.heat_flow,
                                          :MolarEnthalpy, 1.0)

simulate_reservoir(state0, model, dt_schedule; ...)
# → fails at Jacobian alignment for a cross-term targeting :mass_conservation
#   or :energy_conservation, at a (row, col) that exists in the custom equation's
#   declared pattern but not in the cross-term's cache.

I'm happy to put together a fully runnable minimum reproducer (single-cell mesh, single well, ~50 LOC) if that would help.

Root cause analysis

Three layers of failure observed; the first is user-side-fixable, the second is the architectural blocker, the third is a related smaller issue:

1. (Resolved user-side.) Returning a face-stencil from declare_pattern but only allocating N diagonal entries in CompactAutoDiffCache causes alignment to fail on the declared off-diagonals.

   • Fix: composite storage with diagonal CompactAutoDiffCache plus extra_pos::Matrix{Int} for face-stencil positions. Custom align_to_jacobian! and update_linearized_system_equation!. This works.

2. (Not fixable user-side — this issue.) extra_cross_term_sparsity walks every cross-term targeting :Reservoir and merges its source-impact indices into the entry keyed by the cross-term's target_equation symbol. Downstream alignment for each cross-term then walks against the union of all sparsity registered under that symbol — including entries from the custom equation that no cross-term physically writes to. Cross-term alignment fails at positions outside its own cache.

   • Concrete shape of the failure: a reservoir cell far from any well, but inside the custom equation's face stencil, ends up with (row_in_far_cell, well_primary_column) declared in the union. The cross-term whose source is the well tries to align that position and finds it unallocated.

3. (Related, smaller.) Some cross-term update paths use out[] = <scalar> for the written value, which assumes scalar write to a 1-element view. If a user equation is implemented as N=2 rows per cell instead of N=1 + a separate equation, that path errors out. (Mentioning for completeness; not the main ask here.)

Proposed fix — opt-out hook (smallest change)

The minimal-surface fix is a per-equation hook:

# In Jutul (default behavior unchanged):
equation_participates_in_cross_term_sparsity(::JutulEquation) = true

# In extra_cross_term_sparsity (multimodel/crossterm.jl):
for (eq_sym, eq) in pairs(target_model.equations)
    if !equation_participates_in_cross_term_sparsity(eq)
        continue
    end
    # existing merge logic...
end

A user equation that does its own diagonal-only or face-stencil alignment then opts out:

Jutul.equation_participates_in_cross_term_sparsity(::PinEquation) = false

This is ~3 lines in Jutul + 1 line per user equation. Default behavior is unchanged, so existing models and tests are not affected.

Alternative fix — key sparsity by receiving equation

Structurally cleaner but larger surface area:

# CURRENT (approximate):
push!(eq_d[ct.target_equation], (target_indices, source_indices))

# PROPOSED:
push!(eq_d[(ct.target_equation, objectid(target_model.equations[ct.target_equation]))],
      (target_indices, source_indices))

Then downstream alignment walks consider only sparsity entries actually targeting each receiving equation. This is the right shape architecturally, but would need wider regression testing across JutulDarcy's existing models.

Offer

I have working PVS infrastructure (pin equation, classifier, custom storage) ready to deploy as soon as the cross-term sparsity issue is addressed. Happy to:

• Provide a minimum runnable reproducer (~50 LOC)
• Draft a PR implementing the opt-out hook (Option 1 above) against Jutul.jl, including a test case that exercises the failure mode

If the alternative receiving-equation-keyed approach is preferred, I can prototype that as well, though it'll need more discussion on the regression scope.

Either way — would appreciate maintainer guidance on which fix shape you'd prefer to review before I open a PR.

Cross-references

• This affects any custom JutulEquation registered on a model inside a MultiModel with cross-terms — not specific to thermal.
• PVS background: Coats (1980) "An equation of state compositional model"; Cao & Aziz (2002) "Performance of a fully implicit, parallel compositional simulator"; standard treatment in any commercial thermal simulator manual.