Merge branch 'develop' into issue-3361-improve-print_parameter_info-f…

…unctionality
agriyakhetarpal · Dec 10, 2023 · 53bedfd · 53bedfd
2 parents d365023 + 4ad2853
commit 53bedfd
Show file tree

Hide file tree

Showing 44 changed files with 660 additions and 137 deletions.
diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
@@ -48,8 +48,7 @@ jobs:
             echo "tag=all" >> "$GITHUB_OUTPUT"
           fi
 
-      - name: Build and push Docker image to Docker Hub (no solvers)
-        if: matrix.build-args == 'No solvers'
+      - name: Build and push Docker image to Docker Hub (${{ matrix.build-args }})
         uses: docker/build-push-action@v5
         with:
           context: .
@@ -58,29 +57,5 @@ jobs:
           push: true
           platforms: linux/amd64, linux/arm64
 
-      - name: Build and push Docker image to Docker Hub (with ODES and IDAKLU solvers)
-        if: matrix.build-args == 'ODES' || matrix.build-args == 'IDAKLU'
-        uses: docker/build-push-action@v5
-        with:
-          context: .
-          file: scripts/Dockerfile
-          tags: pybamm/pybamm:${{ steps.tags.outputs.tag }}
-          push: true
-          build-args: ${{ matrix.build-args }}=true
-          platforms: linux/amd64, linux/arm64
-
-      - name: Build and push Docker image to Docker Hub (with ALL and JAX solvers)
-        if: matrix.build-args == 'ALL' || matrix.build-args == 'JAX'
-        uses: docker/build-push-action@v5
-        with:
-          context: .
-          file: scripts/Dockerfile
-          tags: pybamm/pybamm:${{ steps.tags.outputs.tag }}
-          push: true
-          build-args: ${{ matrix.build-args }}=true
-          # exclude arm64 for JAX and ALL builds for now, see
-          # https://github.com/google/jax/issues/13608
-          platforms: linux/amd64
-
       - name: List built image(s)
         run: docker images
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,19 +2,31 @@
 
 ## Features
 
+- Added method to get QuickPlot axes by variable ([#3596](https://github.com/pybamm-team/PyBaMM/pull/3596))
+- Added custom experiment terminations ([#3596](https://github.com/pybamm-team/PyBaMM/pull/3596))
+- Mechanical parameters are now a function of stoichiometry and temperature ([#3576](https://github.com/pybamm-team/PyBaMM/pull/3576))
 - Added a new unary operator, `EvaluateAt`, that evaluates a spatial variable at a given position ([#3573](https://github.com/pybamm-team/PyBaMM/pull/3573))
 - Added a method, `insert_reference_electrode`, to `pybamm.lithium_ion.BaseModel` that insert a reference electrode to measure the electrolyte potential at a given position in space and adds new variables that mimic a 3E cell setup. ([#3573](https://github.com/pybamm-team/PyBaMM/pull/3573))
 - Serialisation added so models can be written to/read from JSON ([#3397](https://github.com/pybamm-team/PyBaMM/pull/3397))
 - Added a `get_parameter_info` method for models and modified "print_parameter_info" functionality to extract all parameters and their type in a tabular and readable format ([#3361](https://github.com/pybamm-team/PyBaMM/pull/3584))
 - Mechanical parameters are now a function of stoichiometry and temperature ([#3576](https://github.com/pybamm-team/PyBaMM/pull/3576))
 
+
 ## Bug fixes
 
 - Fixed a bug where simulations using the CasADi-based solvers would fail randomly with the half-cell model ([#3494](https://github.com/pybamm-team/PyBaMM/pull/3494))
 - Fixed bug that made identical Experiment steps with different end times crash ([#3516](https://github.com/pybamm-team/PyBaMM/pull/3516))
 - Fixed bug in calculation of theoretical energy that made it very slow ([#3506](https://github.com/pybamm-team/PyBaMM/pull/3506))
 - The irreversible plating model now increments `f"{Domain} dead lithium concentration [mol.m-3]"`, not `f"{Domain} lithium plating concentration [mol.m-3]"` as it did previously. ([#3485](https://github.com/pybamm-team/PyBaMM/pull/3485))
 
+## Optimizations
+
+- Updated `jax` and `jaxlib` to the latest available versions and added Windows (Python 3.9+) support for the Jax solver ([#3550](https://github.com/pybamm-team/PyBaMM/pull/3550))
+
+## Breaking changes
+
+- Dropped support for the `[jax]` extra, i.e., the Jax solver when running on Python 3.8. The Jax solver is now available on Python 3.9 and above ([#3550](https://github.com/pybamm-team/PyBaMM/pull/3550))
+
 # [v23.9](https://github.com/pybamm-team/PyBaMM/tree/v23.9) - 2023-10-31
 
 ## Features

diff --git a/docs/conf.py b/docs/conf.py
@@ -76,6 +76,7 @@
 
 doctest_global_setup = """
 from docs import *
+import pybamm
 """
 
 # Add any paths that contain templates here, relative to this directory.

diff --git a/docs/source/api/experiment/experiment_steps.rst b/docs/source/api/experiment/experiment_steps.rst
@@ -18,3 +18,28 @@ directly:
 
 .. autoclass:: pybamm.step._Step
     :members:
+
+Step terminations
+-----------------
+
+Standard step termination events are implemented by the following classes, which are
+called when the termination is specified by a specific string. These classes can be
+either be called directly or via the string format specified in the class docstring
+
+.. autoclass:: pybamm.step.CrateTermination
+    :members:
+
+.. autoclass:: pybamm.step.CurrentTermination
+    :members:
+
+.. autoclass:: pybamm.step.VoltageTermination
+    :members:
+
+The following classes can be used to define custom terminations for an experiment
+step:
+
+.. autoclass:: pybamm.step.BaseTermination
+    :members:
+
+.. autoclass:: pybamm.step.CustomTermination
+    :members:
diff --git a/docs/source/api/plotting/quick_plot.rst b/docs/source/api/plotting/quick_plot.rst
@@ -7,3 +7,6 @@ Quick Plot
   :members:
 
 .. autofunction:: pybamm.dynamic_plot
+
+.. autoclass:: pybamm.QuickPlotAxes
+  :members:
diff --git a/docs/source/examples/index.rst b/docs/source/examples/index.rst
@@ -85,6 +85,17 @@ The notebooks are organised into subfolders, and can be viewed in the galleries
     notebooks/parameterization/parameter-values.ipynb
     notebooks/parameterization/parameterization.ipynb
 
+.. nbgallery::
+    :caption: Simulations and Experiments
+    :glob:
+
+    notebooks/simulations_and_experiments/callbacks.ipynb
+    notebooks/simulations_and_experiments/custom-experiments.ipynb
+    notebooks/simulations_and_experiments/experiments-start-time.ipynb
+    notebooks/simulations_and_experiments/rpt-experiment.ipynb
+    notebooks/simulations_and_experiments/simulating-long-experiments.ipynb
+    notebooks/simulations_and_experiments/simulation-class.ipynb
+
 .. nbgallery::
     :caption: Plotting
     :glob:
@@ -111,11 +122,6 @@ The notebooks are organised into subfolders, and can be viewed in the galleries
     :glob:
 
     notebooks/batch_study.ipynb
-    notebooks/callbacks.ipynb
     notebooks/change-settings.ipynb
     notebooks/initialize-model-with-solution.ipynb
-    notebooks/rpt-experiment.ipynb
-    notebooks/simulating-long-experiments.ipynb
-    notebooks/simulation-class.ipynb
     notebooks/solution-data-and-processed-variables.ipynb
-    notebooks/experiments-start-time.ipynb
diff --git a/docs/source/examples/notebooks/batch_study.ipynb b/docs/source/examples/notebooks/batch_study.ipynb
@@ -523,7 +523,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The difference in the individual plots is not very well visible in the above slider plot, but we can access all the simulations created by `BatchStudy` (`batch_study.sims`) and pass it to `pybamm.plot_summary_variables` to plot the summary variables (more details on \"summary variables\" are available in the [`simulating-long-experiments`](https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/examples/notebooks/simulating-long-experiments.ipynb) notebook).\n",
+    "The difference in the individual plots is not very well visible in the above slider plot, but we can access all the simulations created by `BatchStudy` (`batch_study.sims`) and pass it to `pybamm.plot_summary_variables` to plot the summary variables (more details on \"summary variables\" are available in the [`simulating-long-experiments`](./simulations_and_experiments/simulating-long-experiments.ipynb) notebook).\n",
     "\n",
     "## Comparing summary variables"
    ]

diff --git a/docs/source/examples/notebooks/change-settings.ipynb b/docs/source/examples/notebooks/change-settings.ipynb
@@ -7,7 +7,7 @@
    "source": [
     "# Changing settings when solving PyBaMM models\n",
     "\n",
-    "[This](https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/examples/notebooks/models/SPM.ipynb) example notebook showed how to run an SPM battery model, using the default parameters, discretisation and solvers that were defined for that particular model. Naturally we would like the ability to alter these options on a case by case basis, and this notebook gives an example of how to do this, again using the SPM model. In this notebook we explicitly handle all the stages of setting up, processing and solving the model in order to explain them in detail. However, it is often simpler in practice to use the `Simulation` class, which handles many of the stages automatically, as shown [here](https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/examples/notebooks/simulation-class.ipynb).\n",
+    "[This](./models/SPM.ipynb) example notebook showed how to run an SPM battery model, using the default parameters, discretisation and solvers that were defined for that particular model. Naturally we would like the ability to alter these options on a case by case basis, and this notebook gives an example of how to do this, again using the SPM model. In this notebook we explicitly handle all the stages of setting up, processing and solving the model in order to explain them in detail. However, it is often simpler in practice to use the `Simulation` class, which handles many of the stages automatically, as shown [here](./simulations_and_experiments/simulation-class.ipynb).\n",
     "\n",
     "\n",
     "### Table of Contents\n",

diff --git a/docs/source/examples/notebooks/getting_started/tutorial-1-how-to-run-a-model.ipynb b/docs/source/examples/notebooks/getting_started/tutorial-1-how-to-run-a-model.ipynb
@@ -205,7 +205,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.0"
+   "version": "3.11.3"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/examples/notebooks/models/DFN.ipynb b/docs/source/examples/notebooks/models/DFN.ipynb
@@ -107,7 +107,7 @@
    "source": [
     "Below we show how to solve the DFN model, using the default geometry, mesh, parameters, discretisation and solver provided with PyBaMM. For a more detailed example, see the notebook on the [SPM](https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/examples/notebooks/models/SPM.ipynb).\n",
     "\n",
-    "In order to show off all the different points at which the process of setting up and solving a model in PyBaMM can be customised we explicitly handle the stages of choosing a geometry, setting parameters, discretising the model and solving the model. However, it is often simpler in practice to use the `Simulation` class, which handles many of the stages automatically, as shown [here](../simulation-class.ipynb).\n",
+    "In order to show off all the different points at which the process of setting up and solving a model in PyBaMM can be customised we explicitly handle the stages of choosing a geometry, setting parameters, discretising the model and solving the model. However, it is often simpler in practice to use the `Simulation` class, which handles many of the stages automatically, as shown [here](../simulations_and_experiments/simulation-class.ipynb).\n",
     "\n",
     "First we need to import pybamm, along with numpy which we will use in this notebook."
    ]

diff --git a/docs/source/examples/notebooks/models/SPM.ipynb b/docs/source/examples/notebooks/models/SPM.ipynb
@@ -54,7 +54,7 @@
    "source": [
     "## Example solving SPM using PyBaMM\n",
     "\n",
-    "Below we show how to solve the Single Particle Model, using the default geometry, mesh, parameters, discretisation and solver provided with PyBaMM. In this notebook we explicitly handle all the stages of setting up, processing and solving the model in order to explain them in detail. However, it is often simpler in practice to use the `Simulation` class, which handles many of the stages automatically, as shown [here](../simulation-class.ipynb).\n",
+    "Below we show how to solve the Single Particle Model, using the default geometry, mesh, parameters, discretisation and solver provided with PyBaMM. In this notebook we explicitly handle all the stages of setting up, processing and solving the model in order to explain them in detail. However, it is often simpler in practice to use the `Simulation` class, which handles many of the stages automatically, as shown [here](../simulations_and_experiments/simulation-class.ipynb).\n",
     "\n",
     "First we need to import `pybamm`, and then change our working directory to the root of the pybamm folder. "
    ]

diff --git a/...source/examples/notebooks/callbacks.ipynb → ...mulations_and_experiments/callbacks.ipynb b/...source/examples/notebooks/callbacks.ipynb → ...mulations_and_experiments/callbacks.ipynb
diff --git a/docs/source/examples/notebooks/simulations_and_experiments/custom-experiments.ipynb b/docs/source/examples/notebooks/simulations_and_experiments/custom-experiments.ipynb
diff --git a/...es/notebooks/experiments-start-time.ipynb → ..._experiments/experiments-start-time.ipynb b/...es/notebooks/experiments-start-time.ipynb → ..._experiments/experiments-start-time.ipynb
diff --git a/...e/examples/notebooks/rpt-experiment.ipynb → ...ions_and_experiments/rpt-experiment.ipynb b/...e/examples/notebooks/rpt-experiment.ipynb → ...ions_and_experiments/rpt-experiment.ipynb
diff --git a/...tebooks/simulating-long-experiments.ipynb → ...riments/simulating-long-experiments.ipynb b/...tebooks/simulating-long-experiments.ipynb → ...riments/simulating-long-experiments.ipynb
diff --git a/...examples/notebooks/simulation-class.ipynb → ...ns_and_experiments/simulation-class.ipynb b/...examples/notebooks/simulation-class.ipynb → ...ns_and_experiments/simulation-class.ipynb
@@ -6,7 +6,7 @@
    "metadata": {},
    "source": [
     "# A step-by-step look at the Simulation class\n",
-    "The simplest way to solve a model is to use the `Simulation` class. This automatically processes the model (setting of parameters, setting up the mesh and discretisation, etc.) for you, and provides built-in functionality for solving and plotting. Changing things such as parameters in handled by passing options to the `Simulation`, as shown in the [Getting Started](getting_started/tutorial-1-how-to-run-a-model.ipynb) guides, [example notebooks](../index.rst) and [documentation](https://docs.pybamm.org/en/latest/source/api/simulation.html).\n",
+    "The simplest way to solve a model is to use the `Simulation` class. This automatically processes the model (setting of parameters, setting up the mesh and discretisation, etc.) for you, and provides built-in functionality for solving and plotting. Changing things such as parameters in handled by passing options to the `Simulation`, as shown in the [Getting Started](../getting_started/tutorial-1-how-to-run-a-model.ipynb) guides, [example notebooks](../../index.rst) and [documentation](https://docs.pybamm.org/en/latest/source/api/simulation.html).\n",
     "\n",
     "In this notebook we show how to solve a model using a `Simulation` and compare this to manually handling the different stages of the process, such as setting parameters, ourselves step-by-step."
    ]
@@ -152,7 +152,7 @@
    "metadata": {},
    "source": [
     "## Processing the model step-by-step\n",
-    "One way of gaining more control over the simulation processing is by passing options, as outlined in the [documentation](https://docs.pybamm.org/en/latest/source/api/simulation.html). However, you can also process the model step-by-step yourself. A detailed example of this can be found in the [SPM notebook](./models/SPM.ipynb), but here we outline the basic steps.\n",
+    "One way of gaining more control over the simulation processing is by passing options, as outlined in the [documentation](https://docs.pybamm.org/en/latest/source/api/simulation.html). However, you can also process the model step-by-step yourself. A detailed example of this can be found in the [SPM notebook](../models/SPM.ipynb), but here we outline the basic steps.\n",
     "\n",
     "First we pick a model"
    ]
@@ -171,7 +171,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Next we must set up the geometry. We'll use the default geometry for the SPM. In all of the following steps we will also use the default settings provided by the model. For a look at changing these options, see the [change settings](./change-settings.ipynb) notebook."
+    "Next we must set up the geometry. We'll use the default geometry for the SPM. In all of the following steps we will also use the default settings provided by the model. For a look at changing these options, see the [change settings](../change-settings.ipynb) notebook."
    ]
   },
   {

diff --git a/docs/source/user_guide/installation/GNU-linux.rst b/docs/source/user_guide/installation/GNU-linux.rst
@@ -133,7 +133,10 @@ Optional - JaxSolver
 ~~~~~~~~~~~~~~~~~~~~
 
 Users can install ``jax`` and ``jaxlib`` to use the Jax solver.
-Currently, only GNU/Linux and macOS are supported.
+
+.. note::
+
+   The Jax solver is not supported on Python 3.8. It is supported on Python 3.9, 3.10, and 3.11.
 
 .. code:: bash
 

diff --git a/docs/source/user_guide/installation/index.rst b/docs/source/user_guide/installation/index.rst
@@ -217,13 +217,13 @@ Dependency                                                  Minimum Version    p
 Jax dependencies
 ^^^^^^^^^^^^^^^^^
 
-Installable with ``pip install "pybamm[jax]"``
+Installable with ``pip install "pybamm[jax]"``, currently supported on Python 3.9-3.11.
 
 ========================================================================= ================== ================== =======================
 Dependency                                                                Minimum Version    pip extra          Notes
 ========================================================================= ================== ================== =======================
-`JAX <https://jax.readthedocs.io/en/latest/notebooks/quickstart.html>`__  0.4.8              jax                For JAX solvers
-`jaxlib <https://pypi.org/project/jaxlib/>`__                             0.4.7              jax                Support library for JAX
+`JAX <https://jax.readthedocs.io/en/latest/notebooks/quickstart.html>`__  0.4.20             jax                For the JAX solver
+`jaxlib <https://pypi.org/project/jaxlib/>`__                             0.4.20             jax                Support library for JAX
 ========================================================================= ================== ================== =======================
 
 .. _install.odes_dependencies:

diff --git a/docs/source/user_guide/installation/windows.rst b/docs/source/user_guide/installation/windows.rst
@@ -66,6 +66,21 @@ installed automatically when you install PyBaMM using ``pip``.
 For an introduction to virtual environments, see
 (https://realpython.com/python-virtual-environments-a-primer/).
 
+Optional - JaxSolver
+~~~~~~~~~~~~~~~~~~~~
+
+Users can install ``jax`` and ``jaxlib`` to use the Jax solver.
+
+.. note::
+
+   The Jax solver is not supported on Python 3.8. It is supported on Python 3.9, 3.10, and 3.11.
+
+.. code:: bash
+
+	  pip install "pybamm[jax]"
+
+The ``pip install "pybamm[jax]"`` command automatically downloads and installs ``pybamm`` and the compatible versions of ``jax`` and ``jaxlib`` on your system. (``pybamm_install_jax`` is deprecated.)
+
 Uninstall PyBaMM
 ----------------
 

diff --git a/noxfile.py b/noxfile.py
@@ -61,9 +61,12 @@ def run_coverage(session):
     set_environment_variables(PYBAMM_ENV, session=session)
     session.install("coverage", silent=False)
     if sys.platform != "win32":
-        session.install("-e", ".[all,odes,jax]", silent=False)
+        session.install("-e", ".[all,jax,odes]", silent=False)
     else:
-        session.install("-e", ".[all]", silent=False)
+        if sys.version_info < (3, 9):
+            session.install("-e", ".[all]", silent=False)
+        else:
+            session.install("-e", ".[all,jax]", silent=False)
     session.run("coverage", "run", "run-tests.py", "--nosub")
     session.run("coverage", "combine")
     session.run("coverage", "xml")
@@ -74,9 +77,12 @@ def run_integration(session):
     """Run the integration tests."""
     set_environment_variables(PYBAMM_ENV, session=session)
     if sys.platform != "win32":
-        session.install("-e", ".[all,odes,jax]", silent=False)
+        session.install("-e", ".[all,jax,odes]", silent=False)
     else:
-        session.install("-e", ".[all]", silent=False)
+        if sys.version_info < (3, 9):
+            session.install("-e", ".[all]", silent=False)
+        else:
+            session.install("-e", ".[all,jax]", silent=False)
     session.run("python", "run-tests.py", "--integration")
 
 
@@ -92,9 +98,12 @@ def run_unit(session):
     """Run the unit tests."""
     set_environment_variables(PYBAMM_ENV, session=session)
     if sys.platform != "win32":
-        session.install("-e", ".[all,odes,jax]", silent=False)
+        session.install("-e", ".[all,jax,odes]", silent=False)
     else:
-        session.install("-e", ".[all]", silent=False)
+        if sys.version_info < (3, 9):
+            session.install("-e", ".[all]", silent=False)
+        else:
+            session.install("-e", ".[all,jax]", silent=False)
     session.run("python", "run-tests.py", "--unit")
 
 
@@ -144,17 +153,37 @@ def set_dev(session):
             external=True,
         )
     else:
-        session.run(python, "-m", "pip", "install", "-e", ".[all,dev]", external=True)
+        if sys.version_info < (3, 9):
+            session.run(
+                python,
+                "-m",
+                "pip",
+                "install",
+                ".[all,dev]",
+                external=True,
+            )
+        else:
+            session.run(
+                python,
+                "-m",
+                "pip",
+                "install",
+                ".[all,dev,jax]",
+                external=True,
+            )
 
 
 @nox.session(name="tests")
 def run_tests(session):
     """Run the unit tests and integration tests sequentially."""
     set_environment_variables(PYBAMM_ENV, session=session)
     if sys.platform != "win32":
-            session.install("-e", ".[all,odes,jax]", silent=False)
+            session.install("-e", ".[all,jax,odes]", silent=False)
     else:
-        session.install("-e", ".[all]", silent=False)
+        if sys.version_info < (3, 9):
+            session.install("-e", ".[all]", silent=False)
+        else:
+            session.install("-e", ".[all,jax]", silent=False)
     session.run("python", "run-tests.py", "--all")