[brosaplanella]

Description

Fixes two bugs in Discretisation(resolve_coupled_variables=True) reported in #6001:

1. A CoupledVariable in a boundary condition raised a DiscretisationError. process_model discretised the boundary conditions before resolving coupled variables. The resolve step now runs before process_boundary_conditions.
2. inplace=False still rewrote the input model's equations. The resolve step wrote into the model passed in rather than into a copy. It now resolves into model.new_copy() when inplace=False. It also assigns new equation dicts instead of writing into the existing ones, because new_copy copies boundary_conditions only one level deep and the per-boundary dicts are shared with the original.

This PR does not change the default of resolve_coupled_variables, and it does not add the shared resolve cache suggested in #6001. Both are left for the discussion in the issue.

Fixes #6001

Type of change

Bug fix. Entry added under ## Bug fixes in CHANGELOG.md.

Tests

Added two regression tests to tests/unit/test_expression_tree/test_coupled_variable.py. Both fail on main and pass with this change:

• test_coupled_variable_in_boundary_condition solves a diffusion problem whose Neumann flux is a CoupledVariable, and checks it against the same problem with the flux written directly.
• test_resolve_does_not_modify_model_when_not_inplace checks that the input model's rhs is untouched after process_model(..., inplace=False).

Important checks:

• No style issues: nox -s pre-commit. Only ruff format and ruff check have been run so far; both are clean.
• All tests pass: nox -s tests.
• The documentation builds: nox -s doctests. Not run yet.
• Code is commented for hard-to-understand areas
• Tests added that prove fix is effective or that feature works

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 98.25%. Comparing base (63476bc) to head (a71183c).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6002      +/-   ##
==========================================
+ Coverage   98.23%   98.25%   +0.01%     
==========================================
  Files         350      350              
  Lines       36173    36174       +1     
==========================================
+ Hits        35533    35541       +8     
+ Misses        640      633       -7     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Bencher Report

Project	PyBaMM
Branch	fix/coupled-variable-resolution
Testbed	bare-metal

⚠️ WARNING: Truncated view!

The full continuous benchmarking report exceeds the maximum length allowed on this platform.

🐰 View full continuous benchmarking report in Bencher

[brosaplanella]

Pushed 25748fc, which makes resolve_coupled_variables=True the default and adds the shared cache suggested in #6001. Below are timings for both settings, to check that the new default causes no regression, including on very large models.

What changed

• Default flipped. Discretisation(resolve_coupled_variables=True) is now the default.
• Shared cache. _resolve_coupled_variables_in_model passes one tree_map cache through every equation, including the recursive resolution of a coupled variable's target. Subtrees shared between equations are now resolved once.
• No copy when there's nothing to resolve. If no equation changes, the model is returned as-is, without new_copy() or rebuilding the equation dicts. Models without CoupledVariables, which includes every built-in model, only pay for one read-only pass over the equations.

Method

• Both settings run on the same commit. With the flag off, the code path is the same as on main.
• Disc off / on: Discretisation.process_model(model, inplace=False, delayed_variable_processing=True), the call Simulation.build makes, on the parameterised model. Each repeat gets a fresh mesh and fresh spatial methods, because FiniteVolumeUnstructured caches assembled operators on the submesh. Off and on runs alternate, with one untimed warm-up of each; median of 5.
• Build + solve off / on: Simulation(..., discretisation_kwargs={"resolve_coupled_variables": flag}).solve([0, 3600]) with the default IDAKLUSolver; median of 5 alternating runs after a warm-up.
• Resolve pass: _resolve_coupled_variables_in_model timed on its own, i.e. the extra work the flag adds; median of 20.
• States: len_rhs_and_alg of the discretised model, identical for both settings.
• 3D pouch: BasicDFNUnstructured({"dimensionality": 2}) with Marquis2019 and 20 points in each of x_n, x_s, x_p, r_n and r_p. The in-plane y × z mesh is scaled up to 2M states. These are discretised only, not solved.

Results

Model	States	Unique nodes	Resolve pass	Disc off	Disc on	Δ disc	Build + solve off	Build + solve on	Δ build + solve
SPM	42	23	0.25 ms	0.016 s	0.017 s	+7.0%	0.082 s	0.083 s	+2.2%
SPMe	102	100	0.42 ms	0.079 s	0.078 s	−1.9%	0.157 s	0.159 s	+1.0%
DFN	962	335	1.04 ms	0.094 s	0.091 s	−2.9%	0.234 s	0.214 s	−8.5%
DFN + lumped thermal	963	392	1.23 ms	0.107 s	0.114 s	+7.1%	0.450 s	0.491 s	+9.0%
DFN + x-full thermal ¹	1,024	658	1.84 ms	0.162 s	0.168 s	+3.6%	0.446 s	0.459 s	+3.1%
DFN composite anode ²	1,302	439	1.30 ms	0.141 s	0.145 s	+3.1%	0.425 s	0.431 s	+1.4%
SPMe + degradation ³	187	395	1.34 ms	0.169 s	0.159 s	−5.9%	0.627 s	0.632 s	+0.9%
DFN + degradation ³	1,144	419	1.45 ms	0.165 s	0.149 s	−9.6%	0.812 s	0.765 s	−5.7%
DFN + degradation + lumped thermal ³	1,145	493	1.60 ms	0.177 s	0.182 s	+2.7%	2.665 s	2.695 s	+1.1%
DFN 1+1D pouch + x-lumped thermal ¹	9,632	685	2.00 ms	0.383 s	0.382 s	−0.3%	2.817 s	2.853 s	+1.3%
3D pouch DFN, 6 × 4 ¹	23,041	406	1.27 ms	0.262 s	0.270 s	+3.1%	not solved	not solved	—
3D pouch DFN, 24 × 16 ¹	368,641	406	1.22 ms	3.103 s	3.080 s	−0.7%	not solved	not solved	—
3D pouch DFN, 33 × 32 ¹	1,013,761	406	1.25 ms	8.404 s	8.280 s	−1.5%	not solved	not solved	—
3D pouch DFN, 46 × 46 ¹	2,031,361	406	1.28 ms	16.535 s	16.481 s	−0.3%	not solved	not solved	—

¹ Marquis2019. ² Chen2020_composite, "particle phases": ("2", "1"). ³ OKane2022 with SEI (solvent-diffusion limited), SEI porosity change, partially reversible lithium plating, plating porosity change, swelling and cracking, SEI on cracks, and stress-driven loss of active material. All other models use Chen2020.

Takeaways

• No regression. Every Δ is within run-to-run noise: Δ disc ranges from −9.6% to +7.1% with no consistent sign. The extra work the flag adds is the resolve pass, 0.25–2 ms per model.
• The cost does not grow with the number of states. The resolve pass walks the symbolic equations before discretisation, so its cost depends on their size (unique nodes), not on the mesh. For the 3D pouch it stays at about 1.3 ms from 23k to 2M states. At 2M states that is 0.008% of a 16.5 s discretisation.
• The shared cache makes the pass 2–3× faster. Compared with the per-equation cache timings in [Bug]: resolve_coupled_variables=True breaks for boundary conditions and modifies the input model #6001: DFN 2.3 → 1.0 ms, DFN + x-full thermal 5.0 → 1.8 ms, DFN + degradation + thermal 4.7 → 1.6 ms, and 1+1D pouch 5.0 → 2.0 ms. The two sets come from separate runs of the same isolated measurement.
• Solve time is unchanged. Models without coupled variables discretise to identical equations, so the build + solve differences are noise too.

The full unit suite passes on this branch: 5,081 passed, 278 skipped.

Environment: WSL2 on Linux 6.18, 8 cores, 7 GB RAM, Python 3.14.4, commit 25748fc. The 2M-state discretisation peaked at about 1.9 GB of memory.

Benchmark script

"""
Time discretisation, and build + solve, with ``resolve_coupled_variables`` off and on.

Run with ``uv run --no-sync python local/benchmark_resolve_default.py``.
"""

from __future__ import annotations

import argparse
import copy
import gc
import platform
import statistics
import subprocess
import time

import pybamm

pybamm.set_logging_level("ERROR")

DEGRADATION = {
    "SEI": "solvent-diffusion limited",
    "SEI porosity change": "true",
    "lithium plating": "partially reversible",
    "lithium plating porosity change": "true",
    "particle mechanics": ("swelling and cracking", "swelling only"),
    "SEI on cracks": "true",
    "loss of active material": "stress-driven",
}


def pouch_3d_var_pts(y, z):
    return {"x_n": 20, "x_s": 20, "x_p": 20, "r_n": 20, "r_p": 20, "y": y, "z": z}


# (name, model factory, parameter set, var_pts, solve)
CONFIGS = [
    ("SPM", lambda: pybamm.lithium_ion.SPM(), "Chen2020", None, True),
    ("SPMe", lambda: pybamm.lithium_ion.SPMe(), "Chen2020", None, True),
    ("DFN", lambda: pybamm.lithium_ion.DFN(), "Chen2020", None, True),
    (
        "DFN + lumped thermal",
        lambda: pybamm.lithium_ion.DFN({"thermal": "lumped"}),
        "Chen2020",
        None,
        True,
    ),
    (
        "DFN + x-full thermal",
        lambda: pybamm.lithium_ion.DFN({"thermal": "x-full"}),
        "Marquis2019",
        None,
        True,
    ),
    (
        "DFN composite anode",
        lambda: pybamm.lithium_ion.DFN(
            {
                "particle phases": ("2", "1"),
                "open-circuit potential": (("single", "current sigmoid"), "single"),
            }
        ),
        "Chen2020_composite",
        None,
        True,
    ),
    (
        "SPMe + degradation",
        lambda: pybamm.lithium_ion.SPMe(DEGRADATION),
        "OKane2022",
        None,
        True,
    ),
    (
        "DFN + degradation",
        lambda: pybamm.lithium_ion.DFN(DEGRADATION),
        "OKane2022",
        None,
        True,
    ),
    (
        "DFN + degradation + lumped thermal",
        lambda: pybamm.lithium_ion.DFN({**DEGRADATION, "thermal": "lumped"}),
        "OKane2022",
        None,
        True,
    ),
    (
        "DFN 1+1D pouch + x-lumped thermal",
        lambda: pybamm.lithium_ion.DFN(
            {
                "current collector": "potential pair",
                "dimensionality": 1,
                "thermal": "x-lumped",
            }
        ),
        "Marquis2019",
        None,
        True,
    ),
    *(
        (
            f"3D pouch DFN (unstructured, {y}×{z})",
            lambda: pybamm.lithium_ion.BasicDFNUnstructured({"dimensionality": 2}),
            "Marquis2019",
            pouch_3d_var_pts(y, z),
            False,
        )
        for y, z in [(6, 4), (24, 16), (33, 32), (46, 46)]
    ),
]


def make_simulation(model_factory, parameter_set, var_pts, flag=None):
    kwargs = {} if flag is None else {"resolve_coupled_variables": flag}
    return pybamm.Simulation(
        model_factory(),
        parameter_values=pybamm.ParameterValues(parameter_set),
        var_pts=var_pts,
        discretisation_kwargs=kwargs,
    )


def unique_nodes(model):
    exprs = [
        *model.rhs.values(),
        *model.algebraic.values(),
        *model.initial_conditions.values(),
        *(e for bcs in model.boundary_conditions.values() for e, _ in bcs.values()),
    ]
    return len({node for expr in exprs for node in expr.pre_order()})


def time_it(fn):
    gc.collect()
    start = time.perf_counter()
    result = fn()
    return time.perf_counter() - start, result


def benchmark(name, model_factory, parameter_set, var_pts, solve, repeats):
    # parameterise once as Simulation.build does; inplace=False leaves it reusable
    sim = make_simulation(model_factory, parameter_set, var_pts)
    sim._set_parameters()
    model = sim._model_with_set_params

    resolve_times = [
        time_it(
            lambda: pybamm.Discretisation()._resolve_coupled_variables_in_model(
                model, inplace=False
            )
        )[0]
        for _ in range(20)
    ]

    # fresh mesh and spatial methods per repeat, as the unstructured finite volume
    # method caches assembled operators on the submesh; repeat 0 is a warm-up
    disc_times = {False: [], True: []}
    states = {}
    for repeat in range(repeats + 1):
        for flag in (False, True):
            # Mesh pops "tabs" from the geometry it is given
            mesh = pybamm.Mesh(
                copy.deepcopy(sim._geometry), sim._submesh_types, sim._var_pts
            )
            disc = pybamm.Discretisation(
                mesh,
                copy.deepcopy(sim._spatial_methods),
                resolve_coupled_variables=flag,
            )
            elapsed, built = time_it(
                lambda: disc.process_model(
                    model, inplace=False, delayed_variable_processing=True
                )
            )
            if repeat > 0:
                disc_times[flag].append(elapsed)
            states[flag] = built.len_rhs_and_alg
    if states[False] != states[True]:
        raise RuntimeError(f"{name}: state counts differ between flag values")

    e2e_times = {False: [], True: []}
    if solve:
        for repeat in range(repeats + 1):
            for flag in (False, True):
                sim_e2e = make_simulation(model_factory, parameter_set, var_pts, flag)
                elapsed = time_it(lambda: sim_e2e.solve([0, 3600]))[0]
                if repeat > 0:
                    e2e_times[flag].append(elapsed)

    def median(times):
        return statistics.median(times) if times else float("nan")

    return {
        "name": name,
        "states": states[True],
        "nodes": unique_nodes(model),
        "resolve_ms": 1e3 * median(resolve_times),
        "disc_off": median(disc_times[False]),
        "disc_on": median(disc_times[True]),
        "e2e_off": median(e2e_times[False]),
        "e2e_on": median(e2e_times[True]),
    }


def format_row(r):
    def pct(on, off):
        return "—" if off != off else f"{100 * (on - off) / off:+.1f}%"

    def secs(t):
        return "not solved" if t != t else f"{t:.3f} s"

    return (
        f"| {r['name']} | {r['states']:,} | {r['nodes']:,} | {r['resolve_ms']:.2f} ms "
        f"| {r['disc_off']:.3f} s | {r['disc_on']:.3f} s "
        f"| {pct(r['disc_on'], r['disc_off'])} "
        f"| {secs(r['e2e_off'])} | {secs(r['e2e_on'])} "
        f"| {pct(r['e2e_on'], r['e2e_off'])} |"
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--repeats", type=int, default=5)
    parser.add_argument("--only", nargs="*", help="substrings of config names to run")
    args = parser.parse_args()

    commit = subprocess.run(
        ["git", "rev-parse", "--short", "HEAD"], capture_output=True, text=True
    ).stdout.strip()
    print(
        f"<!-- commit {commit}, Python {platform.python_version()}, "
        f"repeats {args.repeats} -->"
    )
    print(
        "| Model | States | Unique nodes | Resolve pass | Disc off | Disc on | Δ disc "
        "| Build + solve off | Build + solve on | Δ build + solve |"
    )
    print("|---|---|---|---|---|---|---|---|---|---|")
    for name, factory, parameter_set, var_pts, solve in CONFIGS:
        if args.only and not any(s in name for s in args.only):
            continue
        result = benchmark(name, factory, parameter_set, var_pts, solve, args.repeats)
        print(format_row(result), flush=True)


if __name__ == "__main__":
    main()

🤖 Generated with Claude Code