How to perform gradient based fitting on data with custom current profile?

[itsjacobhere]

I've found the example here: https://docs.pybamm.org/en/latest/source/examples/notebooks/solvers/idaklu-jax-interface.html showing how to compute the gradient of a model with respect to input parameters, but I am having trouble doing this with a custom current data , custom defined experiment, or current function. I've tried to do this a multitude of ways, mostly by altering the 'Current function [A]' in the parameter values, including using an interpolant, custom function, etc.. but I cannot figure it out. Is there a solution for this? It seems like it should be very simple, but I'm only able to solve this with a static current.

[itsjacobhere]

Update:

I got an example working with a dynamic current:

import pybamm

import time
import numpy as np
import jax
import jax.numpy as jnp

inputs = {
    # "Current function [A]": 0.222,
    # "Separator porosity": 0.48,
    # 'Contact resistance [Ohm]': 0
}

# Set-up the model
model = pybamm.lithium_ion.DFN()
geometry = model.default_geometry

param = pybamm.ParameterValues('Chen2020')

def current_function(t):
    y = 0.01*(t)**1.2+4.1 
    return y

t_eval = np.linspace(0, 100, 100)

param["Current function [A]"] = current_function

param.update({key: "[input]" for key in inputs.keys()})
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)#,inplace=False)

output_variables = [
    "Voltage [V]",
    "Current [A]",
    "Time [min]",
]

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

# Jaxify the IDAKLU solver
jax_solver = idaklu_solver.jaxify(
    model,
    t_eval,
)

f = jax_solver.get_jaxpr()

data = f(t_eval, inputs)
print(data)

#OUTPUT:
[[4.05449876e+000 4.10000000e+000 1.14957963e-311]
 [4.05167363e+000 4.11012133e+000 1.68350168e-002]
 [4.04886734e+000 4.12325272e+000 3.36700337e-002]
 [4.04610103e+000 4.13782538e+000 5.05050505e-002]
 [4.04338093e+000 4.15342072e+000 6.73400673e-002]
 [4.04069778e+000 4.16982353e+000 8.41750842e-002]...]

However when I try to use conditionals in my current function like so:

...
def current_function(t):
    y = 0.5 * (t < 50) + 0.1 * (t >= 50)
    return y
...
f = jax_solver.get_jaxpr()

data = f(t_eval, inputs)
print(data)

I get an error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_2945510/1354928657.py in ?()
      2 f = jax_solver.get_jaxpr()
      3 print(f"JAX expression: {f}")
      4 
      5 # Print all output variables, evaluated over a given time vector
----> 6 data = f(t_eval, inputs)
      7 print(data)

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/pybamm/solvers/idaklu_jax.py in ?(t, inputs)
    628             """
    629             logger.info("f")
    630             flatargs, treedef = tree_flatten((t, inputs))
    631             self.jax_inputs = inputs
--> 632             out = f_p.bind(*flatargs)
    633             return out

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/jax/_src/core.py in ?(self, *args, **params)
    384   def bind(self, *args, **params):
    385     assert (not config.enable_checks.value or
    386             all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args
--> 387     return self.bind_with_trace(find_top_trace(args), args, params)

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/jax/_src/core.py in ?(self, trace, args, params)
    389   def bind_with_trace(self, trace, args, params):
    390     with pop_level(trace.level):
--> 391       out = trace.process_primitive(self, map(trace.full_raise, args), params)
    392     return map(full_lower, out) if self.multiple_results else full_lower(out)

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/jax/_src/core.py in ?(self, primitive, tracers, params)
    875       # Import here to avoid circular imports
    876       from jax.experimental.key_reuse._core import call_impl_with_key_reuse_checks  # pytype: disable=import-error
    877       return call_impl_with_key_reuse_checks(primitive, primitive.impl, *tracers, **params)
    878     else:
--> 879       return primitive.impl(*tracers, **params)

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/pybamm/solvers/idaklu_jax.py in ?(t, *inputs)
    637         @f_p.def_impl
    638         def f_impl(t, *inputs):
    639             """Concrete implementation of Primitive (used for non-jitted evaluation)"""
    640             logger.info("f_impl")
--> 641             term_v = self._jaxify_solve(t, None, *inputs)
    642             logger.debug(f"f_impl [exit]: {type(term_v)}, {term_v}")
    643             return term_v

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/pybamm/solvers/idaklu_jax.py in ?(self, t, invar, *inputs_values)
    368         logger.debug(f"  t: {t}")
    369         logger.debug(f"  invar: {invar}")
    370         logger.debug(f"  inputs: {dict(d)}")
    371         logger.debug(f"  calculate_sensitivities: {invar is not None}")
--> 372         sim = IDAKLUJax._cached_solve(
    373             self.solver,
    374             self.jax_model,
    375             tuple(self.jax_t_eval),

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/pybamm/solvers/idaklu_jax.py in ?(self, model, t_hashable, *args, **kwargs)
    346     @lru_cache(maxsize=1)  # noqa: B019
    347     def _cached_solve(self, model, t_hashable, *args, **kwargs):
    348         """Cache the last solve for reuse"""
--> 349         return self.solve(model, t_hashable, *args, **kwargs)

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/pybamm/solvers/base_solver.py in ?(self, model, t_eval, inputs, nproc, calculate_sensitivities)
    946                 # update y0 (for DAE solvers, this updates the initial guess for the
    947                 # rootfinder)
    948                 model.y0 = last_state
    949                 if len(model.algebraic) > 0:
--> 950                     model.y0 = self.calculate_consistent_state(
    951                         model, t_eval[end_index], model_inputs_list[0]
    952                     )
    953         solve_time = timer.time()

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/pybamm/solvers/base_solver.py in ?(self, model, time, inputs)
    684         if self.root_method is None:
    685             return model.y0
    686         try:
    687             root_sol = self.root_method._integrate(model, np.array([time]), inputs)
--> 688         except pybamm.SolverError as e:
    689             raise pybamm.SolverError(
    690                 f"Could not find consistent states: {e.args[0]}"
    691             ) from e

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/pybamm/solvers/casadi_algebraic_solver.py in ?(self, model, t_eval, inputs_dict)
     93         # If the upper bound is negative then the variable must always be negative
     94         constraints[model.bounds[1] <= 0] = -1
     95 
     96         # Set up rootfinder
---> 97         roots = casadi.rootfinder(
     98             "roots",
     99             "newton",
    100             dict(x=y_alg_sym, p=t_sym, g=alg),

~/anaconda3/envs/pybamm/lib/python3.11/site-packages/casadi/casadi.py in ?(*args)
  30641 
  30642 
  30643 
  30644     """
> 30645     return _casadi.rootfinder(*args)

RuntimeError: .../casadi/core/function_internal.cpp:146: Error calling Newton::init for 'roots':
.../casadi/core/rootfinder.cpp:218: Assertion "n_ == oracle_.nnz_in(iin_)" failed:
Dimension mismatch. Input size is 0, while output size is 100

Is there a way to pass a pybamm experiment as the current function? Is this error because the function is not differentiable for all t?

[martinjrobins]

As you suggest, something to do with the discontinuity in your current function is causing the solver to fail, but the error message is not very good here!

You can smooth off your current function using pybamm.sigmoid, like so:

def current_function(t):
    y = 0.5 * (1.0 - pybamm.sigmoid(t, 50, 10)) + 0.1 * pybamm.sigmoid(t, 50, 10)
    return y

This works for me.

I'm not sure what you mean by "pass a pybamm experiment at the current function". If you want to calculate the derivatve of a pybamm solve using pybamm.Simulation and pybamm.Experiment, then this is currently not implemented. However, this PR should enable this functionality (#4415).

[itsjacobhere]

Yes! I would like to compute gradients while also running a simulation based on a pybamm.Experiment that is passed into either the solver or simulation class to describe the current profile, instead of direcly passing the current profile.

Thanks for the response! Has that PR been implemented yet? Is it only available for the pybamm.ScipySolver?