Fix #4930: discretise Variable.reference / scale so r/x in initial concentration don't leak

[gaoflow]

Description

Fixes #4930.

pybamm.lithium_ion.DFN (and adjacent reduced-order models) crashes with

NotImplementedError: method self.evaluate() not implemented for symbol r_n
of type <class 'pybamm.expression_tree.independent_variable.SpatialVariable'>

as soon as Initial concentration in <Domain> electrode [mol.m-3] is set to a Python callable of (r, x). The exact reproducer from the issue:

import pybamm

def initial_concentration(r, x):
    return 17038.0 + r + x

parameter_values = pybamm.ParameterValues("Chen2020")
parameter_values.update(
    {"Initial concentration in negative electrode [mol.m-3]": initial_concentration}
)

model = pybamm.lithium_ion.DFN()
experiment = pybamm.Experiment(["Rest for 1 sec"])
sim = pybamm.Simulation(
    model, experiment=experiment, parameter_values=parameter_values
)
solution = sim.solve()      # ← NotImplementedError on r_n

pybamm.lithium_ion.BasicDFN works fine with the same callable.

Root cause

LithiumIonParameters builds c_init as a FunctionParameter evaluated on pybamm.SpatialVariable arguments (r, PrimaryBroadcast(x, ...)), so the user's f(r, x) becomes a symbolic expression carrying r_n / x_n leaves. Its derived average

c_init_av = pybamm.xyz_average(pybamm.r_average(self.c_init))

does not fold to a scalar at parameter-substitution time — PyBaMM keeps xyz_average / r_average as deferred Integral operators. U_init is computed from c_init_av, and every electrolyte- / electrode-potential pybamm.Variable is constructed with reference=-self.param.n.prim.U_init (see electrolyte_conductivity/full_conductivity.py:35), so its .reference attribute carries r_n / x_n leaves whenever c_init does.

Discretisation._process_symbol for Variable then did:

return symbol.reference + symbol.scale * pybamm.StateVector(*y_slices, ...)

without first running reference / scale through process_symbol. The un-discretised SpatialVariable leaves leaked into the discretised rhs of Porosity times concentration [mol.m-3] (its source term references the electrolyte potential, which carries the leaky reference) and into the discretised initial conditions of Positive electrode potential / Electrolyte potential (process_equation subtracts name.reference when ics=True). The model then crashed at test_shape (line 880) → check_initial_conditions (line 1292) at runtime.

Fix

Discretise reference and scale through process_symbol before combining them with the state-vector slice / initial-condition equation, in three branches of _process_symbol:

• Variable
• VariableDot
• ConcatenationVariable

and the matching processed_eqn - name.reference / processed_eqn / name.scale step in process_equation. xyz_average / r_average get evaluated against the actual mesh, folding r_n / x_n leaves into numeric vectors.

The change is a no-op for the common case where reference / scale are Scalar (process_symbol caches the result in _discretised_symbols so it adds no measurable overhead) and unblocks callable initial concentrations on both Simulation and Simulation+Experiment build paths.

Tests

New regression file tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_initial_concentration_function_4930.py (+157 lines, 5 tests):

• test_dfn_solves_with_function_of_r_and_x_for_initial_concentration — OP's reproducer minus the experiment.
• test_dfn_with_experiment_solves_with_function_of_r_and_x — OP's exact reproducer.
• test_spm_and_spme_solve_with_function_of_r_and_x — guards reduced-order models too (the discretisation fix is on the generic Variable branch).
• test_discretised_rhs_has_no_spatial_variable_leaves — direct discretisation invariant: no rhs / algebraic / initial_conditions expression may carry an un-folded SpatialVariable leaf after Simulation.build().
• test_positive_electrode_callable_initial_concentration — symmetric positive-electrode side, pins both self.param.n and self.param.p routes.

All 5 new tests pass. Pre-existing test suites stay green:

• tests/unit/test_discretisations/ — 40 passed (+1 pre-existing skfem optional-dep skip).
• tests/unit/test_models/test_full_battery_models/test_lithium_ion/ — 634 passed, 12 pre-existing skips.
• tests/unit/test_experiments/ + tests/unit/test_simulation.py — 210 passed (+4 pre-existing tqdm / skfem optional-dep failures, present on main HEAD too).

ruff check clean.

Type of change

• New feature
• Optimization
• Bug fix

Key checklist

• ruff check passes on touched files
• All tests pass for the touched suites
• Regression tests added (5 tests, including the OP's reproducer)
• CHANGELOG entry under # [Unreleased] → ## Bug fixes