Serialisation of models

docs/source/api/expression_tree/operations/index.rst

@@ -8,4 +8,5 @@ Classes and functions that operate on the expression tree
   evaluate
   jacobian
   convert_to_casadi
+  serialise
   unpack_symbol

docs/source/api/expression_tree/operations/serialise.rst

@@ -0,0 +1,5 @@
+Serialise
+=========
+
+.. autoclass:: pybamm.expression_tree.operations.serialise.Serialise
+  :members:

docs/source/examples/index.rst

@@ -62,6 +62,7 @@ The notebooks are organised into subfolders, and can be viewed in the galleries
     notebooks/models/MSMR.ipynb
     notebooks/models/pouch-cell-model.ipynb
     notebooks/models/rate-capability.ipynb
+    notebooks/models/saving_models.ipynb
     notebooks/models/SEI-on-cracks.ipynb
     notebooks/models/simulating-ORegan-2022-parameter-set.ipynb
     notebooks/models/SPM.ipynb

docs/source/examples/notebooks/models/saving_models.ipynb

@@ -0,0 +1,270 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Saving PyBaMM models to file"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Models which are discretised (i.e. ready to solve/ previously solved, see [this notebook](https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/examples/notebooks/spatial_methods/finite-volumes.ipynb) for more information on the pybamm.Discretisation class) can be serialised and saved to a JSON file, ready to be read in again either in PyBaMM, or a different modelling library. \n",
+    "\n",
+    "In the example below, we build and solve a basic DFN model, and then save the model out to `sim_model_example.json`, which should have appear in the 'models' directory."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install pybamm -q    # install PyBaMM if it is not installed\n",
+    "import pybamm\n",
+    "\n",
+    "# do the example\n",
+    "dfn_model = pybamm.lithium_ion.DFN()\n",
+    "dfn_sim = pybamm.Simulation(dfn_model)\n",
+    "dfn_sim.solve([0, 3600])\n",
+    "\n",
+    "dfn_sim.save_model(\"sim_model_example\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This model file can then be read in and solved by choosing a solver, and running as below."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Recreate the pybamm model from the JSON file\n",
+    "new_dfn_model = pybamm.load_model(\"sim_model_example.json\")\n",
+    "\n",
+    "sim_reloaded = pybamm.Simulation(new_dfn_model)\n",
+    "sim_reloaded.solve([0, 3600])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "It would be nice to plot the results of the two models, to confirm that they are producing the same result.\n",
+    "\n",
+    "However, notice that running the code below generates an error stating that the model variables were not provided during the reading in of the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dfn_models = [dfn_model, new_dfn_model]\n",
+    "sims = []\n",
+    "for model in dfn_models:\n",
+    "    plot_sim = pybamm.Simulation(model)\n",
+    "    plot_sim.solve([0, 3600])\n",
+    "    sims.append(plot_sim)\n",
+    "\n",
+    "pybamm.dynamic_plot(sims, time_unit=\"seconds\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To be able to plot the results from a serialised model, the mesh and model variables need to be saved alongside the model itself.\n",
+    "\n",
+    "To do this, set the `variables` option to `True` when saving the model as in the example below; notice how the models will now plot nicely."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# using the first simulation, save a new file which includes a list of the model variables\n",
+    "dfn_sim.save_model(\"sim_model_variables\", variables=True)\n",
+    "\n",
+    "# read the model back in\n",
+    "model_with_vars = pybamm.load_model(\"sim_model_variables.json\")\n",
+    "\n",
+    "# plot the pre- and post-serialisation models together to prove they behave the same\n",
+    "models = [dfn_model, model_with_vars]\n",
+    "sims = []\n",
+    "for model in models:\n",
+    "    sim = pybamm.Simulation(model)\n",
+    "    sim.solve([0, 3600])\n",
+    "    sims.append(sim)\n",
+    "\n",
+    "pybamm.dynamic_plot(sims, time_unit=\"seconds\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Saving from Model\n",
+    "\n",
+    "Alternatively, the model can be saved directly from the Model class.\n",
+    "\n",
+    "Note that at the moment, only models derived from the BaseBatteryModel class can be serialised; those built from scratch using pybamm.BaseModel() are currently unsupported.\n",
+    "\n",
+    "First set up the model, as explained in detail for the [SPM](https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/examples/notebooks/models/SPM.ipynb)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# create the model\n",
+    "spm_model = pybamm.lithium_ion.SPM()\n",
+    "\n",
+    "# set up and discretise ready to solve\n",
+    "geometry = spm_model.default_geometry\n",
+    "param = spm_model.default_parameter_values\n",
+    "param.process_model(spm_model)\n",
+    "param.process_geometry(geometry)\n",
+    "mesh = pybamm.Mesh(geometry, spm_model.default_submesh_types, spm_model.default_var_pts)\n",
+    "disc = pybamm.Discretisation(mesh, spm_model.default_spatial_methods)\n",
+    "disc.process_model(spm_model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Then save the model. Note that in this case the model variables and the mesh must be provided directly."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Serialise the spm_model, providing the varaibles and the mesh\n",
+    "spm_model.save_model(\"example_model\", variables=spm_model.variables, mesh=mesh)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now you can read the model back in, solve and plot."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# read back in\n",
+    "new_spm_model = pybamm.load_model(\"example_model.json\")\n",
+    "\n",
+    "# select a solver and solve\n",
+    "new_spm_solver = new_spm_model.default_solver\n",
+    "new_spm_solution = new_spm_solver.solve(new_spm_model, [0, 3600])\n",
+    "\n",
+    "# plot the solution\n",
+    "new_spm_solution.plot()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Making edits to a serialised model\n",
+    "\n",
+    "As mentioned at the begining of this notebook, only models which have already been discretised can be serialised and readh back in. This means that after serialisation, the model *cannot be edited*, as the non-discretised elements of the model such as the original rhs are not saved.\n",
+    "\n",
+    "If you are likely to want to save a model and then edit it in the future, you may wish to use the `Simulation.save()` functionality to pickle your simulation, as described in [tutorial 6](https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/examples/notebooks/getting_started/tutorial-6-managing-simulation-outputs.ipynb)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Before finishing we will remove the data files we saved so that we leave the directory as we found it"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.remove(\"example_model.json\")\n",
+    "os.remove(\"sim_model_example.json\")\n",
+    "os.remove(\"sim_model_variables.json\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## References\n",
+    "\n",
+    "The relevant papers for this notebook are:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[1] Joel A. E. Andersson, Joris Gillis, Greg Horn, James B. Rawlings, and Moritz Diehl. CasADi – A software framework for nonlinear optimization and optimal control. Mathematical Programming Computation, 11(1):1–36, 2019. doi:10.1007/s12532-018-0139-4.\n",
+      "[2] Marc Doyle, Thomas F. Fuller, and John Newman. Modeling of galvanostatic charge and discharge of the lithium/polymer/insertion cell. Journal of the Electrochemical society, 140(6):1526–1533, 1993. doi:10.1149/1.2221597.\n",
+      "[3] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, and others. Array programming with NumPy. Nature, 585(7825):357–362, 2020. doi:10.1038/s41586-020-2649-2.\n",
+      "[4] Scott G. Marquis, Valentin Sulzer, Robert Timms, Colin P. Please, and S. Jon Chapman. An asymptotic derivation of a single particle model with electrolyte. Journal of The Electrochemical Society, 166(15):A3693–A3706, 2019. doi:10.1149/2.0341915jes.\n",
+      "[5] Valentin Sulzer, Scott G. Marquis, Robert Timms, Martin Robinson, and S. Jon Chapman. Python Battery Mathematical Modelling (PyBaMM). Journal of Open Research Software, 9(1):14, 2021. doi:10.5334/jors.309.\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "pybamm.print_citations()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "dev",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.16"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

pybamm/__init__.py

@@ -188,6 +188,11 @@
     UserSupplied2DSubMesh,
 )
 
+#
+# Serialisation
+#
+from .models.base_model import load_model
+
 #
 # Spatial Methods
 #

pybamm/expression_tree/array.py

@@ -57,6 +57,30 @@ def __init__(
             name, domain=domain, auxiliary_domains=auxiliary_domains, domains=domains
         )
 
+    @classmethod
+    def _from_json(cls, snippet: dict):
+        instance = cls.__new__(cls)
+
+        if isinstance(snippet["entries"], dict):
+            matrix = csr_matrix(
+                (
+                    snippet["entries"]["data"],
+                    snippet["entries"]["row_indices"],
+                    snippet["entries"]["column_pointers"],
+                ),
+                shape=snippet["entries"]["shape"],
+            )
+        else:
+            matrix = snippet["entries"]
+
+        instance.__init__(
+            matrix,
+            name=snippet["name"],
+            domains=snippet["domains"],
+        )
+
+        return instance
+
     @property
     def entries(self):
         return self._entries
@@ -128,6 +152,30 @@ def to_equation(self):
         entries_list = self.entries.tolist()
         return sympy.Array(entries_list)
 
+    def to_json(self):
+        """
+        Method to serialise an Array object into JSON.
+        """
+
+        if isinstance(self.entries, np.ndarray):
+            matrix = self.entries.tolist()
+        elif isinstance(self.entries, csr_matrix):
+            matrix = {
+                "shape": self.entries.shape,
+                "data": self.entries.data.tolist(),
+                "row_indices": self.entries.indices.tolist(),
+                "column_pointers": self.entries.indptr.tolist(),
+            }
+
+        json_dict = {
+            "name": self.name,
+            "id": self.id,
+            "domains": self.domains,
+            "entries": matrix,
+        }
+
+        return json_dict
+
 
 def linspace(start, stop, num=50, **kwargs):
     """