From 1e08d6ef4b21aeaac4be705b167d268e7489f9c4 Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Tue, 20 Aug 2024 10:01:47 +0100 Subject: [PATCH 01/25] working commit, adds [jax] optional dependency to pybamm, example idaklu-jax solver --- examples/scripts/jaxified-idaklu-example.py | 73 +++++++++++++++++++++ pybop/models/base_model.py | 3 +- pyproject.toml | 2 +- 3 files changed, 76 insertions(+), 2 deletions(-) create mode 100644 examples/scripts/jaxified-idaklu-example.py diff --git a/examples/scripts/jaxified-idaklu-example.py b/examples/scripts/jaxified-idaklu-example.py new file mode 100644 index 000000000..fb5cef0d3 --- /dev/null +++ b/examples/scripts/jaxified-idaklu-example.py @@ -0,0 +1,73 @@ +import time + +import numpy as np +import pybamm + +import pybop + +# Parameter set and model definition +parameter_set = pybop.ParameterSet.pybamm("Chen2020") +model = pybop.lithium_ion.SPM(parameter_set=parameter_set) + +output_vars = [ + "Voltage [V]", + "Current [A]", + "Time [s]", +] +solvers = [ + pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6, output_variables=output_vars), +] + +# Fitting parameters +parameters = pybop.Parameters( + pybop.Parameter( + "Negative electrode active material volume fraction", initial_value=0.55 + ), + pybop.Parameter( + "Positive electrode active material volume fraction", initial_value=0.55 + ), +) + +# Define test protocol and generate data +# experiment = pybop.Experiment([("Discharge at 0.5C for 10 minutes (3 second period)")]) +t_eval = np.linspace(0, 10, 100) +values = model.predict( + initial_state={"Initial open-circuit voltage [V]": 4.2}, t_eval=t_eval +) + +# Form dataset +dataset = pybop.Dataset( + { + "Time [s]": values["Time [s]"].data, + "Current function [A]": values["Current [A]"].data, + "Voltage [V]": values["Voltage [V]"].data, + } +) + +# Create the list of input dicts +n = 150 # Number of solves +inputs = list(zip(np.linspace(0.45, 0.6, n), np.linspace(0.45, 0.6, n))) + +# Iterate over the solvers and print benchmarks +for solver in solvers: + model.build( + inputs={ + "Negative electrode active material volume fraction": 0.55, + "Positive electrode active material volume fraction": 0.55, + } + ) + solver = solver.jaxify(model=model.built_model, t_eval=t_eval) + f = solver.get_jaxpr() + print(f"JAX expression: {f}") + model.solver = f + problem = pybop.FittingProblem(model, parameters, dataset) + + start_time = time.time() + for input_values in inputs: + problem.evaluate(inputs=input_values) + print(f"Time Evaluate {solver.name}: {time.time() - start_time:.3f}") + + # start_time = time.time() + # for input_values in inputs: + # problem.evaluateS1(inputs=input_values) + # print(f"Time EvaluateS1 {solver.name}: {time.time() - start_time:.3f}") diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py index 47ae1e076..8104d1f05 100644 --- a/pybop/models/base_model.py +++ b/pybop/models/base_model.py @@ -795,4 +795,5 @@ def solver(self): @solver.setter def solver(self, solver): - self._solver = solver.copy() if solver is not None else None + # self._solver = solver.copy() if solver is not None else None + self._solver = solver if solver is not None else None diff --git a/pyproject.toml b/pyproject.toml index a1b892fd0..d453965ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ ] requires-python = ">=3.9, <3.13" dependencies = [ - "pybamm>=24.5", + "pybamm[jax]>=24.5", "numpy>=1.16, <2.0", "scipy>=1.3", "pints>=0.5", From 79234f9940b49243cc8f063143239fad86c9897a Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Tue, 20 Aug 2024 13:21:54 +0100 Subject: [PATCH 02/25] adds experimental sub-directory, sync commit --- examples/scripts/jaxified-idaklu-example.py | 15 ++++++++------- pybop/__init__.py | 5 +++++ pybop/experimental/__init__.py | 0 pybop/experimental/jax_costs.py | 18 ++++++++++++++++++ 4 files changed, 31 insertions(+), 7 deletions(-) create mode 100644 pybop/experimental/__init__.py create mode 100644 pybop/experimental/jax_costs.py diff --git a/examples/scripts/jaxified-idaklu-example.py b/examples/scripts/jaxified-idaklu-example.py index fb5cef0d3..34edca659 100644 --- a/examples/scripts/jaxified-idaklu-example.py +++ b/examples/scripts/jaxified-idaklu-example.py @@ -1,5 +1,3 @@ -import time - import numpy as np import pybamm @@ -60,12 +58,15 @@ f = solver.get_jaxpr() print(f"JAX expression: {f}") model.solver = f - problem = pybop.FittingProblem(model, parameters, dataset) - start_time = time.time() - for input_values in inputs: - problem.evaluate(inputs=input_values) - print(f"Time Evaluate {solver.name}: {time.time() - start_time:.3f}") + # Setup Fitting Problem + problem = pybop.JaxFittingProblem(model, parameters, dataset) + problem.evaluate([0.55, 0.55]) + + # start_time = time.time() + # for input_values in inputs: + # problem.evaluate(inputs=input_values) + # print(f"Time Evaluate {solver.name}: {time.time() - start_time:.3f}") # start_time = time.time() # for input_values in inputs: diff --git a/pybop/__init__.py b/pybop/__init__.py index 49c7eb83e..547d06a0a 100644 --- a/pybop/__init__.py +++ b/pybop/__init__.py @@ -158,6 +158,11 @@ from .plotting.plot_parameters import plot_parameters from .plotting.plot_problem import quick_plot +# +# Experimental +# +from .experimental.jax_costs import JaxSumSquaredError + # # Remove any imported modules, so we don't expose them as part of pybop # diff --git a/pybop/experimental/__init__.py b/pybop/experimental/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pybop/experimental/jax_costs.py b/pybop/experimental/jax_costs.py new file mode 100644 index 000000000..3251d4efd --- /dev/null +++ b/pybop/experimental/jax_costs.py @@ -0,0 +1,18 @@ +import jax + +from pybop import BaseProblem, Inputs + + +class JaxSumSquaredError: + """ + Jax-based Sum of Squared Error cost function. + """ + + def __init__(self, problem: BaseProblem): + self._target = problem.target + self._problem = problem + + @jax.jit + def __call__(self, inputs: Inputs): + """Evaluate the jax-based cost function""" + # TODO: Jaxified FittingProblem -> Can BaseProblem be @jitted? From f2bf14304945f3fc2f77070dc48f7621f487ed5f Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Sat, 31 Aug 2024 17:04:39 +0100 Subject: [PATCH 03/25] feat: adds Jax functionality via idaklu_jax solver, with example and benchmarks. Adds JaxSumSquaredError and JaxLogNormalLikelihood. --- examples/scripts/jax-solver-example.py | 60 +++++++ .../scripts/jaxified-idaklu-benchmarks.py | 151 ++++++++++++++++++ examples/scripts/jaxified-idaklu-example.py | 74 --------- pybop/__init__.py | 2 +- pybop/experimental/jax_costs.py | 112 ++++++++++++- pybop/models/base_model.py | 30 +++- pybop/models/lithium_ion/base_echem.py | 3 +- pybop/models/lithium_ion/echem.py | 4 +- pybop/problems/base_problem.py | 28 +++- pybop/problems/fitting_problem.py | 33 ++-- 10 files changed, 394 insertions(+), 103 deletions(-) create mode 100644 examples/scripts/jax-solver-example.py create mode 100644 examples/scripts/jaxified-idaklu-benchmarks.py delete mode 100644 examples/scripts/jaxified-idaklu-example.py diff --git a/examples/scripts/jax-solver-example.py b/examples/scripts/jax-solver-example.py new file mode 100644 index 000000000..b60dbd347 --- /dev/null +++ b/examples/scripts/jax-solver-example.py @@ -0,0 +1,60 @@ +import time + +import numpy as np +import pybamm + +import pybop + +# Parameter set and model definition +parameter_set = pybop.ParameterSet.pybamm("Chen2020") + +# Set up the IDAKLU Solver, and enable jax compiliation via the pybop.model arg 'jax' +# The IDAKLU, and it's jaxified version perform very well on the DFN with and without +# gradient calulations +solver = pybamm.IDAKLUSolver() +model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=solver, jax=True) + +# Fitting parameters +parameters = pybop.Parameters( + pybop.Parameter( + "Negative electrode active material volume fraction", initial_value=0.55 + ), + pybop.Parameter( + "Positive electrode active material volume fraction", initial_value=0.55 + ), +) + +# Define test protocol and generate data +t_eval = np.linspace(0, 100, 100) +values = model.predict( + initial_state={"Initial open-circuit voltage [V]": 4.2}, t_eval=t_eval +) + +# Form dataset +dataset = pybop.Dataset( + { + "Time [s]": values["Time [s]"].data, + "Current function [A]": values["Current [A]"].data, + "Voltage [V]": values["Voltage [V]"].data, + } +) + +problem = pybop.FittingProblem(model, parameters, dataset) +cost = pybop.JaxLogNormalLikelihood(problem, sigma0=0.002) + +# Non-gradient optimiser, change to `pybop.AdamW` for gradient-based example +optim = pybop.XNES( + cost, + max_unchanged_iterations=100, + max_iterations=100, +) +optim.pints_optimiser.set_population_size(2) + +start_time = time.time() +x = optim.run() +print(f"Total time: {time.time() - start_time}") +print(f"x:{x}") +print(f"{optim.result.n_iterations}") + +pybop.plot_convergence(optim) +pybop.plot_parameters(optim) diff --git a/examples/scripts/jaxified-idaklu-benchmarks.py b/examples/scripts/jaxified-idaklu-benchmarks.py new file mode 100644 index 000000000..5c9e72145 --- /dev/null +++ b/examples/scripts/jaxified-idaklu-benchmarks.py @@ -0,0 +1,151 @@ +import time + +import jax +import jax.numpy as jnp +import numpy as np +import pybamm + +import pybop + +n = 30 # Number of solves +output_vars = [ + "Voltage [V]", + "Current [A]", + "Time [s]", +] + +solvers = [ + pybamm.CasadiSolver(mode="fast with events", atol=1e-6, rtol=1e-6), + pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6, output_variables=output_vars), +] + +# Parameter set and model definition +parameter_set = pybop.ParameterSet.pybamm("Chen2020") +model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=solvers[0]) + +# Fitting parameters +parameters = pybop.Parameters( + pybop.Parameter( + "Negative electrode active material volume fraction", initial_value=0.55 + ), + pybop.Parameter( + "Positive electrode active material volume fraction", initial_value=0.55 + ), +) + +# Define test protocol and generate data +t_eval = np.linspace(0, 100, 1000) +values = model.predict( + initial_state={"Initial open-circuit voltage [V]": 4.2}, t_eval=t_eval +) + +# Form dataset +dataset = pybop.Dataset( + { + "Time [s]": values["Time [s]"].data, + "Current function [A]": values["Current [A]"].data, + "Voltage [V]": values["Voltage [V]"].data, + } +) + + +# Create inputs function for benchmarking +def inputs(): + return { + "Negative electrode active material volume fraction": 0.55 + + np.random.normal(0, 0.01), + "Positive electrode active material volume fraction": 0.55 + + np.random.normal(0, 0.01), + } + + +# Iterate over the solvers and print benchmarks +for solver in solvers: + # Setup Fitting Problem + model.solver = solver + problem = pybop.FittingProblem(model, parameters, dataset) + cost = pybop.SumSquaredError(problem) + + start_time = time.time() + for _i in range(n): + out = problem.model.simulate(inputs=inputs(), t_eval=t_eval) + print(f"({solver.name}) Time model.simulate: {time.time() - start_time:.4f}") + + start_time = time.time() + for _i in range(n): + out = problem.model.simulateS1(inputs=inputs(), t_eval=t_eval) + print(f"({solver.name}) Time model.SimulateS1: {time.time() - start_time:.4f}") + + start_time = time.time() + for _i in range(n): + out = problem.evaluate(inputs=inputs()) + print(f"({solver.name}) Time problem.evaluate: {time.time() - start_time:.4f}") + + start_time = time.time() + for _i in range(n): + out = problem.evaluateS1(inputs=inputs()) + print(f"({solver.name}) Time Problem.EvaluateS1: {time.time() - start_time:.4f}") + + start_time = time.time() + for _i in range(n): + out = cost(inputs(), calculate_grad=False) + print(f"({solver.name}) Time PyBOP Cost w/o grad: {time.time() - start_time:.4f}") + + start_time = time.time() + for _i in range(n): + out = cost(inputs(), calculate_grad=True) + print(f"({solver.name}) Time PyBOP Cost w/grad: {time.time() - start_time:.4f}") + +# Jaxified benchmarks +ida_solver = pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6, output_variables=output_vars) +solver = ida_solver.jaxify(model=model.built_model, t_eval=t_eval) +solver_no_grad = ida_solver.jaxify( + model=model.built_model, t_eval=t_eval, calculate_sensitivities=False +) +f = solver.get_jaxpr() +k = solver_no_grad.get_jaxpr() + +start_time = time.time() +for _i in range(n): + kout = k(t_eval, inputs()) +print(f"Time jax expression w/o grad: {time.time() - start_time:.4f}") + +start_time = time.time() +for _i in range(n): + fout = f(t_eval, inputs()) +print(f"Time jax expression w/ grad: {time.time() - start_time:.4f}") + +start_time = time.time() +for _i in range(n): + y = solver_no_grad.get_var("Voltage [V]")(t_eval, inputs()) +print(f"Time jax solver_no_grad.get_var: {time.time() - start_time:.4f}") + +start_time = time.time() +for _i in range(n): + y = solver.get_var("Voltage [V]")(t_eval, inputs()) +print(f"Time jax solver.get_var: {time.time() - start_time:.4f}") + + +# Sum-of-squared errors +def sse_no_grad(t_eval, inputs): + y = solver_no_grad.get_var("Voltage [V]")(t_eval, inputs) + r = jnp.asarray([y - problem.target[signal] for signal in problem.signal]) + return jnp.sum(jnp.sum(r**2, axis=0), axis=0) + + +# Sum-of-squared errors +def sse_grad(t_eval, inputs): + y = solver.get_var("Voltage [V]")(t_eval, inputs) + r = jnp.asarray([y - problem.target[signal] for signal in problem.signal]) + return jnp.sum(jnp.sum(r**2, axis=0), axis=0) + + +start_time = time.time() +for _i in range(n): + out = sse_no_grad(t_eval, inputs()) +print(f"Time Jax SumSquaredError w/o grad: {time.time() - start_time:.4f}") + +start_time = time.time() +for _i in range(n): + out = jax.value_and_grad(sse_grad, argnums=1)(t_eval, inputs()) +print(f"Time Jax SumSquaredError w/ grad: {time.time() - start_time:.4f}") diff --git a/examples/scripts/jaxified-idaklu-example.py b/examples/scripts/jaxified-idaklu-example.py deleted file mode 100644 index 34edca659..000000000 --- a/examples/scripts/jaxified-idaklu-example.py +++ /dev/null @@ -1,74 +0,0 @@ -import numpy as np -import pybamm - -import pybop - -# Parameter set and model definition -parameter_set = pybop.ParameterSet.pybamm("Chen2020") -model = pybop.lithium_ion.SPM(parameter_set=parameter_set) - -output_vars = [ - "Voltage [V]", - "Current [A]", - "Time [s]", -] -solvers = [ - pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6, output_variables=output_vars), -] - -# Fitting parameters -parameters = pybop.Parameters( - pybop.Parameter( - "Negative electrode active material volume fraction", initial_value=0.55 - ), - pybop.Parameter( - "Positive electrode active material volume fraction", initial_value=0.55 - ), -) - -# Define test protocol and generate data -# experiment = pybop.Experiment([("Discharge at 0.5C for 10 minutes (3 second period)")]) -t_eval = np.linspace(0, 10, 100) -values = model.predict( - initial_state={"Initial open-circuit voltage [V]": 4.2}, t_eval=t_eval -) - -# Form dataset -dataset = pybop.Dataset( - { - "Time [s]": values["Time [s]"].data, - "Current function [A]": values["Current [A]"].data, - "Voltage [V]": values["Voltage [V]"].data, - } -) - -# Create the list of input dicts -n = 150 # Number of solves -inputs = list(zip(np.linspace(0.45, 0.6, n), np.linspace(0.45, 0.6, n))) - -# Iterate over the solvers and print benchmarks -for solver in solvers: - model.build( - inputs={ - "Negative electrode active material volume fraction": 0.55, - "Positive electrode active material volume fraction": 0.55, - } - ) - solver = solver.jaxify(model=model.built_model, t_eval=t_eval) - f = solver.get_jaxpr() - print(f"JAX expression: {f}") - model.solver = f - - # Setup Fitting Problem - problem = pybop.JaxFittingProblem(model, parameters, dataset) - problem.evaluate([0.55, 0.55]) - - # start_time = time.time() - # for input_values in inputs: - # problem.evaluate(inputs=input_values) - # print(f"Time Evaluate {solver.name}: {time.time() - start_time:.3f}") - - # start_time = time.time() - # for input_values in inputs: - # problem.evaluateS1(inputs=input_values) - # print(f"Time EvaluateS1 {solver.name}: {time.time() - start_time:.3f}") diff --git a/pybop/__init__.py b/pybop/__init__.py index d074ee91e..4f9e7bf28 100644 --- a/pybop/__init__.py +++ b/pybop/__init__.py @@ -162,7 +162,7 @@ # # Experimental # -from .experimental.jax_costs import JaxSumSquaredError +from .experimental.jax_costs import JaxSumSquaredError, JaxLogNormalLikelihood # # Remove any imported modules, so we don't expose them as part of pybop diff --git a/pybop/experimental/jax_costs.py b/pybop/experimental/jax_costs.py index 3251d4efd..eaded91e6 100644 --- a/pybop/experimental/jax_costs.py +++ b/pybop/experimental/jax_costs.py @@ -1,18 +1,114 @@ +from typing import Union + import jax +import jax.numpy as jnp +import numpy as np + +from pybop import BaseCost, BaseLikelihood, BaseProblem, Inputs + + +class BaseJaxCost(BaseCost): + """ + Jax-based Sum of Squared Error cost function. + """ + + def __init__(self, problem: BaseProblem): + super().__init__(problem) + + def __call__( + self, + inputs: Inputs, + calculate_grad: bool = False, + ) -> Union[np.array, tuple[float, np.ndarray]]: + """ + Computes the cost function for the given predictions. + + Parameters + ---------- + y : dict + The dictionary of predictions with keys designating the signals for fitting. + dy : np.ndarray, optional + The corresponding gradient with respect to the parameters for each signal. + calculate_grad : bool, optional + A bool condition designating whether to calculate the gradient. -from pybop import BaseProblem, Inputs + Returns + ------- + float + The Sum of Squared Error. + """ + inputs = self.parameters.verify(inputs) + self._update_solver_sensitivities(calculate_grad) + if calculate_grad: + y, dy = jax.value_and_grad(self.evaluate)(inputs) + return y, np.stack(list(dy.values())) + else: + return self.evaluate(inputs) -class JaxSumSquaredError: + def _update_solver_sensitivities(self, calculate_grad: bool) -> None: + """ + Updates the solver's sensitivity calculation based on the gradient requirement. + + Args: + calculate_grad (bool): Whether gradient calculation is required. + """ + model = self.problem.model + if calculate_grad != model.calculate_sensitivities: + model.jaxify_solver( + t_eval=self.problem.domain_data, calculate_sensitivities=calculate_grad + ) + + @staticmethod + def check_sigma0(sigma0): + if sigma0 is None: + return 0.005 + if not isinstance(sigma0, (int, float)) or sigma0 <= 0: + raise ValueError("sigma0 must be a positive number") + return float(sigma0) + + +class JaxSumSquaredError(BaseJaxCost): """ Jax-based Sum of Squared Error cost function. """ def __init__(self, problem: BaseProblem): - self._target = problem.target - self._problem = problem + super().__init__(problem) + + def evaluate(self, inputs): + # Calculate residuals and error + y = self.problem.jax_evaluate(inputs) + r = jnp.asarray([y - self._target[signal] for signal in self.signal]) + return jnp.sum(jnp.sum(r**2, axis=0), axis=0) + + +class JaxLogNormalLikelihood(BaseJaxCost, BaseLikelihood): + """ + A Log-Normal Likelihood function. This function represents the + underlining observed data sampled from a Log-Normal distribution. + """ + + def __init__(self, problem: BaseProblem, sigma0=None): + super().__init__(problem) + self.sigma = self.check_sigma0(sigma0) + self.sigma2 = jnp.square(self.sigma) + self._offset = 0.5 * self.n_data * jnp.log(2 * jnp.pi) + self._target_as_array = jnp.asarray( + [self._target[signal] for signal in self.signal] + ) + self._log_target_sum = jnp.sum(jnp.log(self._target_as_array)) + self._precompute() + + def _precompute(self): + self._constant_term = ( + -self._offset - self.n_data * jnp.log(self.sigma) - self._log_target_sum + ) - @jax.jit - def __call__(self, inputs: Inputs): - """Evaluate the jax-based cost function""" - # TODO: Jaxified FittingProblem -> Can BaseProblem be @jitted? + def evaluate(self, inputs): + """ + Evaluates the log-normal likelihood. + """ + y = self.problem.jax_evaluate(inputs) + e = jnp.log(self._target_as_array) - jnp.log(y) + return self._constant_term - jnp.sum(jnp.square(e)) / (2 * self.sigma2) diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py index 5789a57a8..53d9488a7 100644 --- a/pybop/models/base_model.py +++ b/pybop/models/base_model.py @@ -71,7 +71,8 @@ def __init__( name: str = "Base Model", parameter_set: Optional[ParameterSet] = None, check_params: Callable = None, - eis=False, + eis: bool = False, + jax: bool = False, ): """ Initialise the BaseModel with an optional name and a parameter set. @@ -107,6 +108,9 @@ def __init__( """ self.name = name self.eis = eis + self.jax = jax + self.calculate_sensitivities = False + if parameter_set is None: self._parameter_set = None elif isinstance(parameter_set, dict): @@ -752,6 +756,30 @@ def predict( "to be specified." ) + def jaxify_solver(self, t_eval, calculate_sensitivities=False): + """ + Jaxify the IDAKLU Solver and store a copy for reconstructing + the jax'ed solver in the future. This is helpful for reconstructing + the solver with `calculate_sensitivities=True` as this solver requires + casting sensitivity information on construction. + """ + self.calculate_sensitivities = calculate_sensitivities + if isinstance(self._solver, pybamm.IDAKLUSolver): + self._IDAKLU_stored = self._solver.copy() + self._solver = self._solver.jaxify( + model=self._built_model, + t_eval=t_eval, + calculate_sensitivities=calculate_sensitivities, + ) + elif isinstance(self._solver, pybamm.solvers.idaklu_jax.IDAKLUJax): + self._solver = self._IDAKLU_stored.jaxify( + model=self._built_model, + t_eval=t_eval, + calculate_sensitivities=calculate_sensitivities, + ) + else: + raise ValueError("Solver is not pybamm.IDAKLUSolver, cannot jaxify it.") + def check_params( self, inputs: Optional[Inputs] = None, diff --git a/pybop/models/lithium_ion/base_echem.py b/pybop/models/lithium_ion/base_echem.py index e919eb317..6fcb96bae 100644 --- a/pybop/models/lithium_ion/base_echem.py +++ b/pybop/models/lithium_ion/base_echem.py @@ -47,9 +47,10 @@ def __init__( spatial_methods=None, solver=None, eis=False, + jax=False, **model_kwargs, ): - super().__init__(name=name, parameter_set=parameter_set, eis=eis) + super().__init__(name=name, parameter_set=parameter_set, eis=eis, jax=jax) model_options = dict(build=False) for key, value in model_kwargs.items(): diff --git a/pybop/models/lithium_ion/echem.py b/pybop/models/lithium_ion/echem.py index aac05b738..fd86e04e7 100644 --- a/pybop/models/lithium_ion/echem.py +++ b/pybop/models/lithium_ion/echem.py @@ -38,13 +38,15 @@ class SPM(EChemBaseModel): def __init__( self, name="Single Particle Model", - eis=False, + eis: bool = False, + jax: bool = False, **model_kwargs, ): super().__init__( pybamm_model=pybamm_lithium_ion.SPM, name=name, eis=eis, + jax=jax, **model_kwargs, ) diff --git a/pybop/problems/base_problem.py b/pybop/problems/base_problem.py index 59c7ce286..501fbd00d 100644 --- a/pybop/problems/base_problem.py +++ b/pybop/problems/base_problem.py @@ -1,5 +1,6 @@ from typing import Optional +import jax.numpy as jnp import numpy as np from pybamm import IDAKLUSolver @@ -69,10 +70,11 @@ def __init__( self.signal = signal or ["Voltage [V]"] self.set_initial_state(initial_state) self._dataset = None - self._domain_data = None self._target = None self.verbose = False self.failure_output = np.asarray([np.inf]) + if not hasattr(self, "_domain_data"): + self._domain_data = None if isinstance(self._model, BaseModel): self.eis = self.model.eis self.domain = "Frequency [Hz]" if self.eis else "Time [s]" @@ -148,6 +150,30 @@ def evaluateS1(self, inputs: Inputs): """ raise NotImplementedError + def jax_evaluate( + self, + inputs: Inputs, + ) -> jnp.ndarray: + """ + Evaluate the model with the given parameters and return the signal + with a Jax model and solver. + + Parameters + ---------- + inputs : Inputs + Parameters for evaluation of the model. + + Returns + ------- + y : jnp.ndarray + The model output y(t) simulated with given inputs. + """ + + y = jnp.squeeze( + self.model.solver.get_vars(self.signal)(self.domain_data, inputs) + ) + return y + def get_target(self): """ Return the target dataset. diff --git a/pybop/problems/fitting_problem.py b/pybop/problems/fitting_problem.py index 25609b6b1..9f2a6432c 100644 --- a/pybop/problems/fitting_problem.py +++ b/pybop/problems/fitting_problem.py @@ -2,6 +2,7 @@ from typing import Optional import numpy as np +from pybamm import IDAKLUSolver from pybop import BaseModel, BaseProblem, Dataset from pybop.parameters.parameter import Inputs, Parameters @@ -77,6 +78,9 @@ def __init__( initial_state=self.initial_state, ) + if isinstance(self._model.solver, IDAKLUSolver) and self.model.jax: + self.model.jaxify_solver(t_eval=self._domain_data) + def set_initial_state(self, initial_state: Optional[dict] = None): """ Set the initial state to be applied to evaluations of the problem. @@ -192,30 +196,27 @@ def evaluateS1(self, inputs: Inputs): t_eval=self._domain_data, initial_state=self.initial_state, ) + except Exception as e: print(f"Error: {e}") return { signal: self.failure_output for signal in self.signal }, self.failure_output - y = {signal: sol[signal].data for signal in self.signal} + y = { + signal: sol[signal].data + for signal in (self.signal + self.additional_variables) + } - # Extract the sensitivities and stack them along a new axis for each signal - dy = np.empty( - ( - sol[self.signal[0]].data.shape[0], - self.n_outputs, - self._n_parameters, - ) - ) + # Pre-allocate the dy array + dy = np.empty((len(self._domain_data), self.n_outputs, self._n_parameters)) + # Extract the sensitivities for all signals at once + param_keys = self.parameters.keys() for i, signal in enumerate(self.signal): - dy[:, i, :] = np.stack( - [ - sol[signal].sensitivities[key].toarray()[:, 0] - for key in self.parameters.keys() - ], - axis=-1, + sensitivities = sol[signal].sensitivities + dy[:, i, :] = np.column_stack( + [sensitivities[key].toarray()[:, 0] for key in param_keys] ) - return y, np.asarray(dy) + return y, dy From d3a18d5a14a2edfb1574236011a31fc69de8c4bc Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Mon, 2 Sep 2024 10:23:23 +0100 Subject: [PATCH 04/25] Post merge fixes, update base_model solver property --- pybop/models/base_model.py | 8 ++++++-- pybop/problems/fitting_problem.py | 10 ++++------ 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py index 53d9488a7..34e60d1a8 100644 --- a/pybop/models/base_model.py +++ b/pybop/models/base_model.py @@ -1002,5 +1002,9 @@ def solver(self): @solver.setter def solver(self, solver): - # self._solver = solver.copy() if solver is not None else None - self._solver = solver if solver is not None else None + if isinstance(solver, pybamm.solvers.idaklu_jax.IDAKLUJax): + self._solver = solver + elif solver is not None: + self._solver = solver.copy() + else: + self._solver = None diff --git a/pybop/problems/fitting_problem.py b/pybop/problems/fitting_problem.py index 9f2a6432c..fa748e473 100644 --- a/pybop/problems/fitting_problem.py +++ b/pybop/problems/fitting_problem.py @@ -203,13 +203,11 @@ def evaluateS1(self, inputs: Inputs): signal: self.failure_output for signal in self.signal }, self.failure_output - y = { - signal: sol[signal].data - for signal in (self.signal + self.additional_variables) - } + y = {sig: sol[sig].data for sig in (self.signal + self.additional_variables)} - # Pre-allocate the dy array - dy = np.empty((len(self._domain_data), self.n_outputs, self._n_parameters)) + # Pre-allocate sensitivities array + y_dims = sol[self.signal[0]].data.shape[0] + dy = np.empty((y_dims, self.n_outputs, self._n_parameters)) # Extract the sensitivities for all signals at once param_keys = self.parameters.keys() From eb08c00a776eddd679f793c06ec613fd8a3d6cb4 Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Mon, 2 Sep 2024 13:20:23 +0100 Subject: [PATCH 05/25] examples: update benchmarking script --- .../scripts/jaxified-idaklu-benchmarks.py | 65 +++---------------- 1 file changed, 10 insertions(+), 55 deletions(-) diff --git a/examples/scripts/jaxified-idaklu-benchmarks.py b/examples/scripts/jaxified-idaklu-benchmarks.py index 5c9e72145..50eddd847 100644 --- a/examples/scripts/jaxified-idaklu-benchmarks.py +++ b/examples/scripts/jaxified-idaklu-benchmarks.py @@ -1,22 +1,14 @@ import time -import jax -import jax.numpy as jnp import numpy as np import pybamm import pybop -n = 30 # Number of solves -output_vars = [ - "Voltage [V]", - "Current [A]", - "Time [s]", -] - +n = 1 # Number of solves solvers = [ pybamm.CasadiSolver(mode="fast with events", atol=1e-6, rtol=1e-6), - pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6, output_variables=output_vars), + pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6), ] # Parameter set and model definition @@ -96,56 +88,19 @@ def inputs(): out = cost(inputs(), calculate_grad=True) print(f"({solver.name}) Time PyBOP Cost w/grad: {time.time() - start_time:.4f}") -# Jaxified benchmarks -ida_solver = pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6, output_variables=output_vars) -solver = ida_solver.jaxify(model=model.built_model, t_eval=t_eval) -solver_no_grad = ida_solver.jaxify( - model=model.built_model, t_eval=t_eval, calculate_sensitivities=False -) -f = solver.get_jaxpr() -k = solver_no_grad.get_jaxpr() - -start_time = time.time() -for _i in range(n): - kout = k(t_eval, inputs()) -print(f"Time jax expression w/o grad: {time.time() - start_time:.4f}") - -start_time = time.time() -for _i in range(n): - fout = f(t_eval, inputs()) -print(f"Time jax expression w/ grad: {time.time() - start_time:.4f}") - -start_time = time.time() -for _i in range(n): - y = solver_no_grad.get_var("Voltage [V]")(t_eval, inputs()) -print(f"Time jax solver_no_grad.get_var: {time.time() - start_time:.4f}") - -start_time = time.time() -for _i in range(n): - y = solver.get_var("Voltage [V]")(t_eval, inputs()) -print(f"Time jax solver.get_var: {time.time() - start_time:.4f}") - - -# Sum-of-squared errors -def sse_no_grad(t_eval, inputs): - y = solver_no_grad.get_var("Voltage [V]")(t_eval, inputs) - r = jnp.asarray([y - problem.target[signal] for signal in problem.signal]) - return jnp.sum(jnp.sum(r**2, axis=0), axis=0) - - -# Sum-of-squared errors -def sse_grad(t_eval, inputs): - y = solver.get_var("Voltage [V]")(t_eval, inputs) - r = jnp.asarray([y - problem.target[signal] for signal in problem.signal]) - return jnp.sum(jnp.sum(r**2, axis=0), axis=0) - +# Recreate for Jax IDAKLU solver +ida_solver = pybamm.IDAKLUSolver(atol=1e-6, rtol=1e-6) +model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=ida_solver, jax=True) +problem = pybop.FittingProblem(model, parameters, dataset) +cost = pybop.JaxSumSquaredError(problem) +# Jaxified benchmarks start_time = time.time() for _i in range(n): - out = sse_no_grad(t_eval, inputs()) + out = cost(inputs(), calculate_grad=False) print(f"Time Jax SumSquaredError w/o grad: {time.time() - start_time:.4f}") start_time = time.time() for _i in range(n): - out = jax.value_and_grad(sse_grad, argnums=1)(t_eval, inputs()) + out = cost(inputs(), calculate_grad=True) print(f"Time Jax SumSquaredError w/ grad: {time.time() - start_time:.4f}") From 3ecd5dd25bde33b6584d35cc0c28ea5e53a9e2b8 Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Tue, 24 Sep 2024 16:32:47 +0100 Subject: [PATCH 06/25] updts docstrings, jax arg to models, adds tests. --- examples/scripts/jax-solver-example.py | 6 +- pybop/experimental/jax_costs.py | 16 +++- pybop/models/lithium_ion/base_echem.py | 4 + pybop/models/lithium_ion/echem.py | 60 +++++++++++-- tests/unit/experimental/test_jax_costs.py | 100 ++++++++++++++++++++++ tests/unit/test_cost.py | 5 -- 6 files changed, 171 insertions(+), 20 deletions(-) create mode 100644 tests/unit/experimental/test_jax_costs.py diff --git a/examples/scripts/jax-solver-example.py b/examples/scripts/jax-solver-example.py index b60dbd347..fc5603548 100644 --- a/examples/scripts/jax-solver-example.py +++ b/examples/scripts/jax-solver-example.py @@ -8,9 +8,9 @@ # Parameter set and model definition parameter_set = pybop.ParameterSet.pybamm("Chen2020") -# Set up the IDAKLU Solver, and enable jax compiliation via the pybop.model arg 'jax' +# Set up the IDAKLU Solver, and enable jax compilation via the pybop.model arg 'jax' # The IDAKLU, and it's jaxified version perform very well on the DFN with and without -# gradient calulations +# gradient calculations solver = pybamm.IDAKLUSolver() model = pybop.lithium_ion.DFN(parameter_set=parameter_set, solver=solver, jax=True) @@ -25,7 +25,7 @@ ) # Define test protocol and generate data -t_eval = np.linspace(0, 100, 100) +t_eval = np.linspace(0, 300, 100) values = model.predict( initial_state={"Initial open-circuit voltage [V]": 4.2}, t_eval=t_eval ) diff --git a/pybop/experimental/jax_costs.py b/pybop/experimental/jax_costs.py index eaded91e6..0f44c6d33 100644 --- a/pybop/experimental/jax_costs.py +++ b/pybop/experimental/jax_costs.py @@ -19,6 +19,7 @@ def __call__( self, inputs: Inputs, calculate_grad: bool = False, + apply_transform: bool = False, ) -> Union[np.array, tuple[float, np.ndarray]]: """ Computes the cost function for the given predictions. @@ -42,7 +43,9 @@ def __call__( if calculate_grad: y, dy = jax.value_and_grad(self.evaluate)(inputs) - return y, np.stack(list(dy.values())) + return y, np.asarray( + list(dy.values()) + ) # Convert grad to numpy for optimisers else: return self.evaluate(inputs) @@ -61,8 +64,6 @@ def _update_solver_sensitivities(self, calculate_grad: bool) -> None: @staticmethod def check_sigma0(sigma0): - if sigma0 is None: - return 0.005 if not isinstance(sigma0, (int, float)) or sigma0 <= 0: raise ValueError("sigma0 must be a positive number") return float(sigma0) @@ -87,9 +88,16 @@ class JaxLogNormalLikelihood(BaseJaxCost, BaseLikelihood): """ A Log-Normal Likelihood function. This function represents the underlining observed data sampled from a Log-Normal distribution. + + Parameters + ----------- + problem: BaseProblem + The problem to fit of type `pybop.BaseProblem` + sigma0: float, optional + The variance in the measured data """ - def __init__(self, problem: BaseProblem, sigma0=None): + def __init__(self, problem: BaseProblem, sigma0=0.02): super().__init__(problem) self.sigma = self.check_sigma0(sigma0) self.sigma2 = jnp.square(self.sigma) diff --git a/pybop/models/lithium_ion/base_echem.py b/pybop/models/lithium_ion/base_echem.py index 751c88789..8eaece07f 100644 --- a/pybop/models/lithium_ion/base_echem.py +++ b/pybop/models/lithium_ion/base_echem.py @@ -29,6 +29,10 @@ class EChemBaseModel(BaseModel): The spatial methods used for discretization. If None, default spatial methods from PyBaMM are used. solver : pybamm.Solver, optional The solver to use for simulating the model. If None, the default solver from PyBaMM is used. + eis : bool, optional + A flag to build the forward model for EIS predictions. Defaults to False. + jax : bool, optional + A flag to build the forward model in preparation for the JAX solvers. Defaults to False. **model_kwargs : optional Valid PyBaMM model option keys and their values. For example, build : bool, optional diff --git a/pybop/models/lithium_ion/echem.py b/pybop/models/lithium_ion/echem.py index fd86e04e7..ab8c71599 100644 --- a/pybop/models/lithium_ion/echem.py +++ b/pybop/models/lithium_ion/echem.py @@ -15,6 +15,10 @@ class SPM(EChemBaseModel): ---------- name : str, optional The name for the model instance, defaulting to "Single Particle Model". + eis : bool, optional + A flag to build the forward model for EIS predictions. Defaults to False. + jax : bool, optional + A flag to build the forward model in preparation for the JAX solvers. Defaults to False. **model_kwargs : optional Valid PyBaMM model option keys and their values, for example: parameter_set : pybamm.ParameterValues or dict, optional @@ -64,6 +68,10 @@ class SPMe(EChemBaseModel): ---------- name: str, optional A name for the model instance, defaults to "Single Particle Model with Electrolyte". + eis : bool, optional + A flag to build the forward model for EIS predictions. Defaults to False. + jax : bool, optional + A flag to build the forward model in preparation for the JAX solvers. Defaults to False. **model_kwargs : optional Valid PyBaMM model option keys and their values, for example: parameter_set : pybamm.ParameterValues or dict, optional @@ -87,11 +95,16 @@ class SPMe(EChemBaseModel): def __init__( self, name="Single Particle Model with Electrolyte", - eis=False, + eis: bool = False, + jax: bool = False, **model_kwargs, ): super().__init__( - pybamm_model=pybamm_lithium_ion.SPMe, name=name, eis=eis, **model_kwargs + pybamm_model=pybamm_lithium_ion.SPMe, + name=name, + eis=eis, + jax=jax, + **model_kwargs, ) @@ -108,6 +121,10 @@ class DFN(EChemBaseModel): ---------- name : str, optional The name for the model instance, defaulting to "Doyle-Fuller-Newman". + eis : bool, optional + A flag to build the forward model for EIS predictions. Defaults to False. + jax : bool, optional + A flag to build the forward model in preparation for the JAX solvers. Defaults to False. **model_kwargs : optional Valid PyBaMM model option keys and their values, for example: parameter_set : pybamm.ParameterValues or dict, optional @@ -131,11 +148,16 @@ class DFN(EChemBaseModel): def __init__( self, name="Doyle-Fuller-Newman", - eis=False, + eis: bool = False, + jax: bool = False, **model_kwargs, ): super().__init__( - pybamm_model=pybamm_lithium_ion.DFN, name=name, eis=eis, **model_kwargs + pybamm_model=pybamm_lithium_ion.DFN, + name=name, + eis=eis, + jax=jax, + **model_kwargs, ) @@ -150,6 +172,10 @@ class MPM(EChemBaseModel): ---------- name : str, optional The name for the model instance, defaulting to "Many Particle Model". + eis : bool, optional + A flag to build the forward model for EIS predictions. Defaults to False. + jax : bool, optional + A flag to build the forward model in preparation for the JAX solvers. Defaults to False. **model_kwargs : optional Valid PyBaMM model option keys and their values, for example: parameter_set : pybamm.ParameterValues or dict, optional @@ -173,12 +199,14 @@ class MPM(EChemBaseModel): def __init__( self, name="Many Particle Model", - eis=False, + eis: bool = False, + jax: bool = False, **model_kwargs, ): super().__init__( pybamm_model=pybamm_lithium_ion.MPM, eis=eis, + jax=jax, name=name, **model_kwargs, ) @@ -195,6 +223,10 @@ class MSMR(EChemBaseModel): ---------- name : str, optional The name for the model instance, defaulting to "Multi Species Multi Reactions Model". + eis : bool, optional + A flag to build the forward model for EIS predictions. Defaults to False. + jax : bool, optional + A flag to build the forward model in preparation for the JAX solvers. Defaults to False. **model_kwargs : optional Valid PyBaMM model option keys and their values, for example: parameter_set : pybamm.ParameterValues or dict, optional @@ -218,13 +250,15 @@ class MSMR(EChemBaseModel): def __init__( self, name="Multi Species Multi Reactions Model", - eis=False, + eis: bool = False, + jax: bool = False, **model_kwargs, ): super().__init__( pybamm_model=pybamm_lithium_ion.MSMR, name=name, eis=eis, + jax=jax, **model_kwargs, ) @@ -237,6 +271,10 @@ class WeppnerHuggins(EChemBaseModel): ---------- name: str, optional A name for the model instance, defaults to "Weppner & Huggins model". + eis : bool, optional + A flag to build the forward model for EIS predictions. Defaults to False. + jax : bool, optional + A flag to build the forward model in preparation for the JAX solvers. Defaults to False. **model_kwargs : optional Valid PyBaMM model option keys and their values, for example: parameter_set : pybamm.ParameterValues or dict, optional @@ -253,7 +291,13 @@ class WeppnerHuggins(EChemBaseModel): The solver to use for simulating the model. If None, the default solver from PyBaMM is used. """ - def __init__(self, name="Weppner & Huggins model", eis=False, **model_kwargs): + def __init__( + self, + name="Weppner & Huggins model", + eis: bool = False, + jax: bool = False, + **model_kwargs, + ): super().__init__( - pybamm_model=BaseWeppnerHuggins, name=name, eis=eis, **model_kwargs + pybamm_model=BaseWeppnerHuggins, name=name, eis=eis, jax=jax, **model_kwargs ) diff --git a/tests/unit/experimental/test_jax_costs.py b/tests/unit/experimental/test_jax_costs.py new file mode 100644 index 000000000..cca314310 --- /dev/null +++ b/tests/unit/experimental/test_jax_costs.py @@ -0,0 +1,100 @@ +import jax.numpy as jnp +import numpy as np +import pybamm +import pytest + +import pybop + + +class TestCosts: + """ + Class for testing Jax-based cost functions + """ + + @pytest.fixture + def model(self, ground_truth): + solver = pybamm.IDAKLUSolver() + model = pybop.lithium_ion.SPM(solver=solver, jax=True) + model.parameter_set["Negative electrode active material volume fraction"] = ( + ground_truth + ) + return model + + @pytest.fixture + def ground_truth(self): + return 0.52 + + @pytest.fixture + def parameters(self): + return pybop.Parameter( + "Negative electrode active material volume fraction", + prior=pybop.Gaussian(0.5, 0.01), + bounds=[0.375, 0.625], + ) + + @pytest.fixture + def experiment(self): + return pybop.Experiment( + [ + ("Discharge at 1C for 10 minutes (20 second period)"), + ] + ) + + @pytest.fixture + def dataset(self, model, experiment): + solution = model.predict(experiment=experiment) + return pybop.Dataset( + { + "Time [s]": solution["Time [s]"].data, + "Current function [A]": solution["Current [A]"].data, + "Voltage [V]": solution["Terminal voltage [V]"].data, + } + ) + + @pytest.fixture + def signal(self): + return "Voltage [V]" + + @pytest.fixture + def problem(self, model, parameters, dataset, signal, request): + problem = pybop.FittingProblem(model, parameters, dataset, signal=signal) + return problem + + @pytest.fixture(params=[pybop.JaxSumSquaredError, pybop.JaxLogNormalLikelihood]) + def cost(self, problem, request): + cls = request.param + return cls(problem) + + @pytest.mark.unit + def test_costs(self, cost): + if isinstance(cost, pybop.BaseLikelihood): + higher_cost = cost([0.52]) + lower_cost = cost([0.55]) + else: + higher_cost = cost([0.55]) + lower_cost = cost([0.52]) + assert higher_cost > lower_cost or ( + higher_cost == lower_cost and not np.isfinite(higher_cost) + ) + + # Test type of returned value + out = cost([0.5]) + assert isinstance(out, jnp.ndarray) + assert out.dtype == jnp.float64 + + # Test UserWarnings + if isinstance(cost, pybop.JaxSumSquaredError): + assert cost([0.5]) >= 0 + + # Test option setting + cost.set_fail_gradient(10) + assert cost._de == 10 + + # Test grad + e, de = cost([0.5], calculate_grad=True) + assert e.shape == () + assert de.shape == (1,) + assert isinstance(e, jnp.ndarray) + assert isinstance(de, np.ndarray) + assert e.dtype == jnp.float64 + assert de.dtype == np.float64 diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py index 2aad7169b..568a91544 100644 --- a/tests/unit/test_cost.py +++ b/tests/unit/test_cost.py @@ -12,11 +12,6 @@ class TestCosts: Class for tests cost functions """ - # Define an invalid likelihood class for MAP tests - class InvalidLikelihood: - def __init__(self, problem, sigma0): - pass - @pytest.fixture def model(self, ground_truth): solver = pybamm.IDAKLUSolver() From 9794910055d74437f0398561fc4145062cfeec94 Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Wed, 25 Sep 2024 12:56:34 +0100 Subject: [PATCH 07/25] add coverage, refactor BaseModel solver setter --- pybop/models/base_model.py | 2 +- tests/unit/experimental/test_jax_costs.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py index d5ff09e5a..84f7415f4 100644 --- a/pybop/models/base_model.py +++ b/pybop/models/base_model.py @@ -1015,7 +1015,7 @@ def solver(self): def solver(self, solver): if isinstance(solver, pybamm.solvers.idaklu_jax.IDAKLUJax): self._solver = solver - elif solver is not None: + elif isinstance(solver, pybamm.BaseSolver): self._solver = solver.copy() else: self._solver = None diff --git a/tests/unit/experimental/test_jax_costs.py b/tests/unit/experimental/test_jax_costs.py index cca314310..72b0ec58c 100644 --- a/tests/unit/experimental/test_jax_costs.py +++ b/tests/unit/experimental/test_jax_costs.py @@ -1,3 +1,5 @@ +import copy + import jax.numpy as jnp import numpy as np import pybamm @@ -98,3 +100,19 @@ def test_costs(self, cost): assert isinstance(de, np.ndarray) assert e.dtype == jnp.float64 assert de.dtype == np.float64 + + @pytest.mark.unit + def test_solver_property(self, problem): + model = problem.model + + # Test getter + assert isinstance(model.solver, pybamm.IDAKLUJax) + + # Test setter + jaxed_solver_copy = copy.copy(model.solver) + model.solver = jaxed_solver_copy + assert isinstance(model.solver, pybamm.IDAKLUJax) + + # Test incorrect setter + model.solver = model.pybamm_model + assert model.solver is None From df1ba0c31980596e86f27b1e43a0308ed3d680e1 Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Wed, 25 Sep 2024 13:26:27 +0100 Subject: [PATCH 08/25] adds coverage --- pybop/models/base_model.py | 2 +- tests/unit/experimental/test_jax_costs.py | 20 ++++++++++++++------ 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py index 84f7415f4..6b3a2aeca 100644 --- a/pybop/models/base_model.py +++ b/pybop/models/base_model.py @@ -784,7 +784,7 @@ def jaxify_solver(self, t_eval, calculate_sensitivities=False): t_eval=t_eval, calculate_sensitivities=calculate_sensitivities, ) - elif isinstance(self._solver, pybamm.solvers.idaklu_jax.IDAKLUJax): + elif isinstance(self._solver, pybamm.IDAKLUJax): self._solver = self._IDAKLU_stored.jaxify( model=self._built_model, t_eval=t_eval, diff --git a/tests/unit/experimental/test_jax_costs.py b/tests/unit/experimental/test_jax_costs.py index 72b0ec58c..bb6fc16ce 100644 --- a/tests/unit/experimental/test_jax_costs.py +++ b/tests/unit/experimental/test_jax_costs.py @@ -84,13 +84,14 @@ def test_costs(self, cost): assert isinstance(out, jnp.ndarray) assert out.dtype == jnp.float64 - # Test UserWarnings - if isinstance(cost, pybop.JaxSumSquaredError): - assert cost([0.5]) >= 0 + # Test option setting + cost.set_fail_gradient(10) + assert cost._de == 10 - # Test option setting - cost.set_fail_gradient(10) - assert cost._de == 10 + # Test incorrect Sigma0 + if isinstance(cost, pybop.JaxLogNormalLikelihood): + with pytest.raises(ValueError, match="sigma0 must be a positive number"): + cost.check_sigma0(-0.01) # Test grad e, de = cost([0.5], calculate_grad=True) @@ -116,3 +117,10 @@ def test_solver_property(self, problem): # Test incorrect setter model.solver = model.pybamm_model assert model.solver is None + + # Test jaxify with incorrect solver + model.solver = pybamm.CasadiSolver() + with pytest.raises( + ValueError, match="Solver is not pybamm.IDAKLUSolver, cannot jaxify it." + ): + model.jaxify_solver(t_eval=np.linspace(0, 1, 100)) From d77031502023f25a663189a8df9d9fd0eeb8481f Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Wed, 25 Sep 2024 13:53:13 +0100 Subject: [PATCH 09/25] convert BaseModel.calculate_sensitivites to property --- pybop/models/base_model.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py index 6b3a2aeca..8ec0677c6 100644 --- a/pybop/models/base_model.py +++ b/pybop/models/base_model.py @@ -109,7 +109,7 @@ def __init__( self.name = name self.eis = eis self.jax = jax - self.calculate_sensitivities = False + self._calculate_sensitivities = False if parameter_set is None: self._parameter_set = None @@ -776,7 +776,7 @@ def jaxify_solver(self, t_eval, calculate_sensitivities=False): the solver with `calculate_sensitivities=True` as this solver requires casting sensitivity information on construction. """ - self.calculate_sensitivities = calculate_sensitivities + self._calculate_sensitivities = calculate_sensitivities if isinstance(self._solver, pybamm.IDAKLUSolver): self._IDAKLU_stored = self._solver.copy() self._solver = self._solver.jaxify( @@ -1011,6 +1011,10 @@ def spatial_methods(self): def solver(self): return self._solver + @property + def calculate_sensitivities(self): + return self._calculate_sensitivities + @solver.setter def solver(self, solver): if isinstance(solver, pybamm.solvers.idaklu_jax.IDAKLUJax): From daae4ba3d3a9b847dd8ec51e03d776db1505b744 Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Wed, 25 Sep 2024 15:33:44 +0100 Subject: [PATCH 10/25] add changelog entry --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1748bcc7b..d1cc09a60 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## Features +- [#481](https://github.com/pybop-team/PyBOP/pull/481) - Adds experimental support for PyBaMM's jaxified IDAKLU solver. Includes Jax-specific cost functions `pybop.JaxSumSquareError` and `pybop.JaxLogNormalLikelihood`. Adds `Jax` optional dependency to PyBaMM dependency. + ## Optimisations - [#512](https://github.com/pybop-team/PyBOP/pull/513) - Refactors `LogPosterior` with attributes pointing to composed likelihood object. From 424699be42ec0cf790a445ba208e0775a5a1317b Mon Sep 17 00:00:00 2001 From: Brady Planden Date: Mon, 30 Sep 2024 14:00:39 +0100 Subject: [PATCH 11/25] feat: removes Jax arg from BaseModel, BaseProblem.model is copied instance --- .../1-single-pulse-circuit-model.ipynb | 84 ++++++++----------- examples/scripts/jax-solver-example.py | 2 +- .../scripts/jaxified-idaklu-benchmarks.py | 2 +- pybop/experimental/jax_costs.py | 4 + pybop/models/base_model.py | 2 - pybop/models/lithium_ion/base_echem.py | 5 +- pybop/models/lithium_ion/echem.py | 25 +----- pybop/problems/base_problem.py | 4 +- pybop/problems/fitting_problem.py | 4 - tests/unit/experimental/test_jax_costs.py | 18 +++- tests/unit/test_problem.py | 4 - 11 files changed, 61 insertions(+), 93 deletions(-) diff --git a/examples/notebooks/1-single-pulse-circuit-model.ipynb b/examples/notebooks/1-single-pulse-circuit-model.ipynb index 2478ce755..c8bc6220f 100644 --- a/examples/notebooks/1-single-pulse-circuit-model.ipynb +++ b/examples/notebooks/1-single-pulse-circuit-model.ipynb @@ -21,40 +21,13 @@ "id": "1", "metadata": {}, "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "/home/engs2510/Documents/Git/PyBOP/.nox/notebooks-overwrite/bin/python3: No module named pip\r\n" - ] - }, { "name": "stdout", "output_type": "stream", "text": [ "Note: you may need to restart the kernel to use updated packages.\n", - "zsh:1: no matches found: pybop[plot]\r\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ + "zsh:1: no matches found: pybop[plot]\r\n", "Note: you may need to restart the kernel to use updated packages.\n", - "/home/engs2510/Documents/Git/PyBOP/.nox/notebooks-overwrite/bin/python3: No module named pip" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\r\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ "Note: you may need to restart the kernel to use updated packages.\n" ] } @@ -213,9 +186,9 @@ { "data": { "text/html": [ - "