As with any contribution to PyBaMM, please follow the workflow in CONTRIBUTING.md. In particular, start by creating an issue to discuss what you want to do - this is a good way to avoid wasted coding hours!
All models in PyBaMM are implemented as expression trees. After the model has been created, parameters have been set, and the model has been discretised, the model is now a linear algebra object with the following attributes:
A pybamm.Symbol node that can be evaluated at a state (t, y) and returns the value of all the differential equations at that state, concatenated into a single vector
A pybamm.Symbol node that can be evaluated at a state (t, y) and returns the value of all the algebraic equations at that state, concatenated into a single vector
A numpy array of initial conditions for all the differential and algebraic equations, concatenated into a single vector
A dictionary of pybamm.Symbol nodes representing events at which the solver should terminate. Specifically, the solver should terminate when any of the events in model.events.values() evaluate to zero
The role of solvers is to solve a model at a given set of time points, returning a vector of times t and a matrix of states y.
There is one general base solver class, pybamm.BaseSolver, which sets up some useful solver properties such as tolerances and implement a method self.solve() that solves a model at a given set of time points.
The solve method unpacks the model, simplifies it by removing extraneous operations, (optionally) creates or calls the mass matrix and/or jacobian, and passes the appropriate attributes to another method, called integrate, which does the time-stepping. The role of specific solver classes is simply to implement this integrate method for an arbitrary set of derivative function, initial conditions etc.
The base solver class also computes a consistent set of initial conditions for the algebraic equations, using model.concatenated_initial_conditions as an initial guess.
To add a new solver (e.g. My Fast DAE Solver), first create a new file (my_fast_dae_solver.py) in pybamm/solvers/,
with a single class that inherits from pybamm.BaseSolver. For example:
def MyFastDaeSolver(pybamm.BaseSolver):
Also add the class to pybamm/__init__.py:
from .solvers.my_fast_dae_solver import MyFastDaeSolver
You can then start implementing the solver by adding the integrate function to the class.
For an example of an existing solver implementation, see the Scikits DAE solver API docs and notebook.
For the new solver to be added to PyBaMM, you must add unit tests to demonstrate that it behaves as expected
(see, for example, the Scikits solver tests).
The best way to get started would be to create a file test_my_fast_solver.py in tests/unit/test_solvers/ that performs at least the
following checks:
The integrate method works on a simple ODE/DAE model with/without jacobian, mass matrix and/or events as appropriate
The solve method works on a simple model (in theory, if the integrate method works then the solve method should always work)
If the solver is expected to converge in a certain way as the time step is changed, you could also add a convergence test in tests/convergence/solvers/.
In theory, any existing model can now be solved using MyFastDaeSolver instead of their default solvers, with no extra work from here. To test this, add something like the following test to one of the model test files (e.g. DFN):
def test_my_fast_solver(self):
model = pybamm.lithium_ion.DFN()
solver = pybamm.MyFastDaeSolver()
modeltest = tests.StandardModelTest(model, solver=solver)
modeltest.test_all()
This will check that the model can run with the new solver (but not that it gives a sensible answer!).
Once you have performed the above checks, you are almost ready to merge your code into the core PyBaMM - see CONTRIBUTING.md workflow for how to do this.