IDAKLU-JAX not working with BPX model

[MarkV-ADI]

Hi all, I recently been playing around with the IDAKLU-JAX Solver and very impressed with it's speed! I am trying to target a BPX model cell. However when i do, i get this error: KeyError: 'Current function [A]'

Setup: WSL2 Ubuntu, Python 3.10.12, Pybamm[Jax]==24.5 and BPX==0.4.0

Code

import pybamm
import numpy as np
import time

# Define inputs
inputs = {
    "Current function [A]": 0.222,
    "Separator porosity": 0.3,
}

# Set-up the model
options = {"cell geometry": "arbitrary", "thermal": "lumped"}
model = pybamm.lithium_ion.DFN(options=options)
geometry = model.default_geometry
# param = model.default_parameter_values
param = pybamm.ParameterValues.create_from_bpx(r'lfp_18650_cell_BPX.json')
param.update(inputs)
param.process_geometry(geometry)
param.process_model(model)
var = pybamm.standard_spatial_vars
var_pts = {var.x_n: 20, var.x_s: 20, var.x_p: 20, var.r_n: 10, var.r_p: 10}
mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts)
disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
disc.process_model(model)

# Use a short time-vector for this example
t_eval = np.linspace(0, 360, 10)

# Include the variable of interest in output_variables
output_variables = [
    "Voltage [V]",
    "Current [A]",
    "Time [min]",
    "X-averaged cell temperature [C]" # want to retrieve this and print the output
]

# Create the IDAKLU Solver object
idaklu_solver = pybamm.IDAKLUSolver(rtol=1e-6, atol=1e-6, output_variables=output_variables)

start_time = time.time()

# Perform the simulation
sim = idaklu_solver.solve(
    model,
    t_eval,
    inputs=inputs,
    calculate_sensitivities=True,
)

end_time = time.time()
print(f"Simulation took {end_time - start_time:.2f} seconds.")

# Variable name to extract
var_name = "X-averaged cell temperature [C]"

# Check if the variable is present in the simulation results
if var_name not in sim._variables:
    print(f"Variable '{var_name}' not found in the simulation results.")
else:
    # Extract the variable from the simulation result
    try:
        # Evaluate the variable using the simulation result
        temp_data = sim[var_name](t_eval)
        print("X-averaged cell temperature [C]:")
        print(temp_data)
    except Exception as e:
        print(f"An error occurred while evaluating the variable: {e}")

Response

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[120], [line 44](vscode-notebook-cell:?execution_count=120&line=44)
     [41](vscode-notebook-cell:?execution_count=120&line=41) start_time = time.time()
     [43](vscode-notebook-cell:?execution_count=120&line=43) # Perform the simulation
---> [44](vscode-notebook-cell:?execution_count=120&line=44) sim = idaklu_solver.solve(
     [45](vscode-notebook-cell:?execution_count=120&line=45)     model,
     [46](vscode-notebook-cell:?execution_count=120&line=46)     t_eval,
     [47](vscode-notebook-cell:?execution_count=120&line=47)     inputs=inputs,
     [48](vscode-notebook-cell:?execution_count=120&line=48)     calculate_sensitivities=True,
     [49](vscode-notebook-cell:?execution_count=120&line=49) )
     [51](vscode-notebook-cell:?execution_count=120&line=51) end_time = time.time()
     [52](vscode-notebook-cell:?execution_count=120&line=52) print(f"Simulation took {end_time - start_time:.2f} seconds.")

File [~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:827](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:827), in solve(self, model, t_eval, inputs, nproc, calculate_sensitivities, t_interp)
    [822](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:822) if len(self._model_set_up) > 0:
    [823](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:823)     existing_model = next(iter(self._model_set_up))
    [824](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:824)     raise RuntimeError(
    [825](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:825)         "This solver has already been initialised for model "
    [826](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:826)         f'"{existing_model.name}". Please create a separate '
--> [827](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:827)         "solver for this model"
    [828](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:828)     )
    [829](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:829) # It is assumed that when len(inputs_list) > 1, model set
    [830](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:830) # up (initial condition, time-scale and length-scale) does
    [831](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:831) # not depend on input parameters. Therefore, only `model_inputs[0]`
...
   [1635](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:1635)         )
   [1637](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:1637) if use_jacobian:
   [1638](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/mnt/c/Users/MVeerasi/Documents/Battery-Simulator/~/.local/lib/python3.10/site-packages/pybamm/solvers/base_solver.py:1638)     report(f"Calculating jacobian for {name} using CasADi")

KeyError: 'Current function [A]'

NOTE: The above code works if i use param = model.default_parameter_values instead of param = pybamm.ParameterValues.create_from_bpx(r'lfp_18650_cell_BPX.json')

The BPX file used is found from the about:energy repo for BPX: https://github.com/About-Energy-OpenSource/About-Energy-BPX-Parameterisation/blob/main/LFP/lfp_18650_cell_BPX.json

I thought that i could add in the current function parameter similar to how it was done in chen2020 but after modifying the bpx Parameterisation->cell to have a current function

"Parameterisation": {
            "Cell": {
                  "Ambient temperature [K]": 298.15,
                  "Initial temperature [K]": 298.15,
                  "Reference temperature [K]": 298.15,
                  "Lower voltage cut-off [V]": 2.0,
                  "Upper voltage cut-off [V]": 3.65,
                  "Nominal cell capacity [A.h]": 2,
                  "Specific heat capacity [J.K-1.kg-1]": 999,
                  "Thermal conductivity [W.m-1.K-1]": 1.89,
                  "Density [kg.m-3]": 1940,
                  "Electrode area [m2]": 0.08959998,
                  "Number of electrode pairs connected in parallel to make a cell": 1,
                  "External surface area [m2]": 0.00431,
                  "Volume [m3]": 1.7e-05,
                  "Current function [A]": 5
            },

Response to modified BPX Schema

  extra fields not permitted (type=value_error.extra)
Parameterisation -> Separator
  extra fields not permitted (type=value_error.extra)
Parameterisation -> Cell -> Current function [A]
  extra fields not permitted (type=value_error.extra).....

so on and so on for each parameter...

Is there a way i can use idaklu solver with the bpx model?
should my input dict change to not use current?, or is it that i'm meant to modify the bpx file somehow to target current or is there a general recommend way for me to solve with IDAKULU-JAX for this.

to note: i'm quite new to the domain of battery simulations so i don't believe this is a bug but just a gap in my knowledge on how to perform this type of simulation :)

Thanks all in advance :)

[ejfdickinson]

@MarkV-ADI In relation to the BPX issue, you can't write in new fields in this part of the BPX file or it will be rejected by the schema. The schema errors themselves are not helpful (FaradayInstitution/BPX#52).

BPX permits you to add any additional fields for the BPX file under "Parameterisation" -> "User-defined". PyBaMM imports these directly into the pybamm.ParameterValues object, by default.

Does this error persist if you apply your inputs to the parameter values (as you are doing), assemble a pybamm.Simulation directly and let PyBaMM do the default things, rather than processing the model in place? This might help debug if there is still a PyBaMM issue here. I've put some approximate code below.

# Set-up the model
options = {"cell geometry": "arbitrary", "thermal": "lumped"}
model = pybamm.lithium_ion.DFN(options=options)
geometry = model.default_geometry
# param = model.default_parameter_values
param = pybamm.ParameterValues.create_from_bpx(r'lfp_18650_cell_BPX.json')
param.update(inputs)
var = pybamm.standard_spatial_vars
var_pts = {var.x_n: 20, var.x_s: 20, var.x_p: 20, var.r_n: 10, var.r_p: 10}

# Use a short time-vector for this example
t_eval = np.linspace(0, 360, 10)

# Include the variable of interest in output_variables
output_variables = [
    "Voltage [V]",
    "Current [A]",
    "Time [min]",
    "X-averaged cell temperature [C]" # want to retrieve this and print the output
]

# Create the IDAKLU Solver object
idaklu_solver = pybamm.IDAKLUSolver(rtol=1e-6, atol=1e-6, output_variables=output_variables)

sim = pybamm.Simulation(
    model,
    parameter_values=param,
    geometry=geometry,
    var_pts=var_pts,
    solver=idaklu_solver,
    output_variables=output_variables,
)
sol = sim.solve(
    t_eval=t_eval,
    calculate_sensitivities=True,
)

[pipliggins]

@MarkV-ADI @ejfdickinson's code should work if you don't use the output_variables option in the IDAKLU solver. Using the output_variables option reveals a small bug I've noted in #4439.