[Bug]: IDAKLUSolver with output_variables provides accesss to alternative variables

[BradyPlanden]

PyBaMM Version

[Develop]

Python Version

3.11

Describe the bug

Using the output_variable functionality with the IDAKLUSolver appears to provide access to the full state matrix. It is currently possible to access variables not defined within the output_variables dictionary as shown below.

Steps to Reproduce

import pybamm
import numpy as np

# construct model
model = pybamm.lithium_ion.DFN()
geometry = model.default_geometry
param = model.default_parameter_values
input_parameters = {}  # Sensitivities dictionary
param.update({key: "[input]" for key in input_parameters})
param.process_model(model)
param.process_geometry(geometry)
var_pts = {"x_n": 50, "x_s": 50, "x_p": 50, "r_n": 5, "r_p": 5}
mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts)
disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
disc.process_model(model)
t_eval = np.linspace(0, 3600, 100)

options = {
    'linear_solver': 'SUNLinSol_KLU',
    'jacobian': 'sparse',
    'num_threads': 4,
}

# Set output variables
output_variables = [
            "Voltage [V]",
            "Time [min]",
            "x [m]",
            "Negative particle flux [mol.m-2.s-1]",
            "Throughput capacity [A.h]",  # ExplicitTimeIntegral
        ]

# Solve for a subset of variables and compare results
solver = pybamm.IDAKLUSolver(
    atol=1e-8, rtol=1e-8,
    options=options,
    output_variables=output_variables,
)
sol = solver.solve(
    model,
    t_eval,
    inputs=input_parameters,
)

# Test output variables 
print(sol["Voltage [V]"].data) # OK
print(sol["Discharge capacity [A.h]"].data) # Should not work

Relevant log output

No response

[jsbrittain]

@BradyPlanden From a quick glance I don't think that the solver is returning the full state matrix - this is not returned to python from IDAKLU when output_variables is specified; for example, if you run the solver with output_variables=['Voltage [V]'] only, the solver will report values for Voltage [V], Time [min] and x [m], but will fail with a casadi error when attempting to return Negative particle flux [mol.m-2.s-1]. In-fact, if you request a variable that is in the output_variables list then the class returned is of type ProcessedVariableComputed (representing pre-computed variables), whereas if the variable was not in output_variables then the return is of type ProcessedVariable. The latter class will attempt to compute the variable and add it to the Solution dictionary (this is normal behaviour when output_variables is not specified). This makes me think that the non-output_variables items are being calculated in python via casadi functions - however, because the full state matrix is not available they will sometimes/often(?) fail. Either way this needs a more informative accompanying error.

This does leave open a design decision though: 1) strictly only allow access to those variables that were in output_variables, or 2) attempt to resolve all variables as is currently being performed, but fall-over more gracefully when this is not possible. I will take a closer look at these options in the coming days.

[jsbrittain]

@BradyPlanden I've taken a closer look at this and it is the model variables, or those that do not require access to the state vector, that are being returned. Below in a code snippet that demonstrates this by checking each variable in a full model against a reduced (minimal) model. None of the variables that return in the reduced model reference the state vector (just by identifying any instances of "y[" in the equations as a spot check for now), while all those that fail to return do reference it.

import pybamm
import numpy as np

input_parameters = {}  # Sensitivities dictionary
t_eval = np.linspace(0, 3600, 100)


# construct model
def construct_model():
    model = pybamm.lithium_ion.DFN()
    geometry = model.default_geometry
    param = model.default_parameter_values
    param.update({key: "[input]" for key in input_parameters})
    param.process_model(model)
    param.process_geometry(geometry)
    var_pts = {"x_n": 50, "x_s": 50, "x_p": 50, "r_n": 5, "r_p": 5}
    mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts)
    disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
    disc.process_model(model)
    return model


options = {
    'linear_solver': 'SUNLinSol_KLU',
    'jacobian': 'sparse',
    'num_threads': 4,
}

# Solve for all variables
model_all = construct_model()
solver_all = pybamm.IDAKLUSolver(
    atol=1e-8, rtol=1e-8,
    options=options,
)
sol_all = solver_all.solve(
    model_all,
    t_eval,
    inputs=input_parameters,
)

# Solve for a subset of variables
output_variables = [
    "Current [A]",
]
model = construct_model()
solver = pybamm.IDAKLUSolver(
    atol=1e-8, rtol=1e-8,
    options=options,
    output_variables=output_variables,
)
sol = solver.solve(
    model,
    t_eval,
    inputs=input_parameters,
)

found = 0
assert_failed = 0
for var in model.variables_and_events:
    try:
        sol[var].data
        if not np.allclose(sol[var].data, sol_all[var].data):
            print(f"ASSERT FAILED '{var}'")
            assert_failed += 1
            continue
        if "y[" in model.variables[var].__str__():
            print("===> eqn does reference state vector")
            assert_failed += 1
            continue
        print(f"              '{var}'")
        found += 1
    except RuntimeError:
        if "Event:" not in var and "y[" not in model.variables[var].__str__():
            print("===> eqn does not reference state vector")
            assert_failed += 1
        print(f"NOT FOUND     '{var}'")

print(f"Found {found} and failed to find {len(model.variables_and_events) - found} of "
      f"{len(model.variables_and_events)} variables.")
print(f"Assert failed {assert_failed} times.")

which returns (excerpt)...

Found 325 and failed to find 222 of 547 variables.
Assert failed 0 times.

I will add some checks to return a more sensible error message when an uncomputed (and unavailable) variable is requested (currently casadi throws a runtime error) and adjust the unit tests to cover these scenarios. I also suspect that there is some cross-contamination of the model instances in the unit tests (between solver_all and solver).