Merge branch 'develop' into issue-1481-casadi-interp

.github/release_checklist.md

@@ -4,4 +4,4 @@
   - `CITATION.cff`
   - `vcpkg.json`
 - Update CHANGELOG.md with a summary of the release
-- Update (and pin) jax and jaxlib to latest version and fix any bugs that arise
+- Update jax and jaxlib to latest version in `pybamm.util.install_jax` and fix any bugs that arise

.github/workflows/build_wheels.yml

@@ -76,7 +76,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Download wheels
-        uses: actions/download-artifact@v1
+        uses: actions/download-artifact@v2
         with:
           name: wheels
 

.github/workflows/test_on_push.yml

@@ -2,21 +2,21 @@ name: PyBaMM
 
 on:
   push:
-    
+
   pull_request:
 
   # everyday at 3 am UTC
   schedule:
     - cron:  '0 3 * * *'
 
 jobs:
-  
+
   style:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
     - name: Setup python
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v2
       with:
         python-version: 3.7
 
@@ -37,19 +37,19 @@ jobs:
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
-  
+
     - name: Install Linux system dependencies
       if: matrix.os == 'ubuntu-latest'
       run: |
         sudo apt-get update
         sudo apt install gfortran gcc libopenblas-dev graphviz
         sudo apt install texlive-full
-    
+
     # Added fixes to homebrew installs:
-    # rm -f /usr/local/bin/2to3 
+    # rm -f /usr/local/bin/2to3
     # (see https://github.com/actions/virtual-environments/issues/2322)
     - name: Install MacOS system dependencies
       if: matrix.os == 'macos-latest'
@@ -62,7 +62,7 @@ jobs:
     - name: Install Windows system dependencies
       if: matrix.os == 'windows-latest'
       run: choco install graphviz --version=2.38.0.20190211
-    
+
     - name: Install standard python dependencies
       run: |
         python -m pip install --upgrade pip wheel setuptools
@@ -72,9 +72,17 @@ jobs:
       if: matrix.os == 'ubuntu-latest'
       run: tox -e pybamm-requires
 
-    - name: Run unit tests for GNU/Linux
+    - name: Run unit tests for GNU/Linux with Python 3.7 and 3.8
+      if: matrix.os == 'ubuntu-latest' && matrix.python-version != 3.9
+      run: python -m tox -e quick
+
+    - name: Run unit tests for GNU/Linux with Python 3.9 and generate coverage report
+      if: matrix.os == 'ubuntu-latest' && matrix.python-version == 3.9
+      run: tox -e coverage
+
+    - name: Run integration tests for GNU/Linux
       if: matrix.os == 'ubuntu-latest'
-      run: python -m tox -e tests
+      run: python -m tox -e integration
 
     - name: Run unit tests for Windows and MacOS
       if: matrix.os != 'ubuntu-latest'
@@ -83,16 +91,11 @@ jobs:
     - name: Install docs dependencies and run doctests
       if: matrix.os == 'ubuntu-latest'
       run: tox -e doctests
-    
+
     - name: Install dev dependencies and run example tests
       if: matrix.os == 'ubuntu-latest'
       run: tox -e examples
-
-    - name: Install and run coverage
-      if: success() && (matrix.os == 'ubuntu-latest' && matrix.python-version == 3.9)
-      run: tox -e coverage
 
     - name: Upload coverage report
       if: matrix.os == 'ubuntu-latest' && matrix.python-version == 3.9
-      uses: codecov/codecov-action@v1
-
+      uses: codecov/[email protected]

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+-   `Experiment`s with drive cycles can be solved ([#1793](https://github.com/pybamm-team/PyBaMM/pull/1793))
 -   Added surface area to volume ratio as a factor to the SEI equations ([#1790](https://github.com/pybamm-team/PyBaMM/pull/1790))
 -   Half-cell SPM and SPMe have been implemented ([#1731](https://github.com/pybamm-team/PyBaMM/pull/1731))
 
@@ -14,6 +15,7 @@
 
 -   Raise error when trying to convert an `Interpolant` with the "pchip" interpolator to CasADI ([#1791](https://github.com/pybamm-team/PyBaMM/pull/1791))
 -   Raise error if `Concatenation` is used directly with `Variable` objects (`concatenation` should be used instead) ([#1789](https://github.com/pybamm-team/PyBaMM/pull/1789))
+-   Made jax and the PyBaMM JaxSolver optional ([#1767](https://github.com/pybamm-team/PyBaMM/pull/1767))
 
 # [v21.10](https://github.com/pybamm-team/PyBaMM/tree/v21.9) - 2021-10-31
 

CONTRIBUTING.md

@@ -280,7 +280,7 @@ Major PyBaMM features are showcased in [Jupyter notebooks](https://jupyter.org/)
 
 All example notebooks should be listed in [examples/README.md](https://github.com/pybamm-team/PyBaMM/blob/develop/examples/notebooks/README.md). Please follow the (naming and writing) style of existing notebooks where possible.
 
-Where possible, notebooks are tested daily. A list of slow notebooks (which time-out and fail tests) is maintained in `.slow-books`, these notebooks will be excluded from daily testing.
+All the notebooks are tested daily.
 
 ## Citations
 

README.md

@@ -87,7 +87,9 @@ conda install -c conda-forge pybamm
 ```
 
 ### Optional solvers
-On GNU/Linux and MacOS, an optional [scikits.odes](https://scikits-odes.readthedocs.io/en/latest/)-based solver is available, see [the documentation](https://pybamm.readthedocs.io/en/latest/install/GNU-linux.html#scikits-odes-label).
+Following GNU/Linux and macOS solvers are optionally available:
+- [scikits.odes](https://scikits-odes.readthedocs.io/en/latest/)-based solver, see [the documentation](https://pybamm.readthedocs.io/en/latest/install/GNU-linux.html#optional-scikits-odes-solver).
+- [jax](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html)-based solver, see [the documentation](https://pybamm.readthedocs.io/en/latest/install/GNU-linux.html#optional-jaxsolver).
 
 ## 📖 Citing PyBaMM
 

docs/index.rst

@@ -46,7 +46,10 @@ PyBaMM is available as a conda package through the conda-forge channel.
 
 Optional solvers
 -----------------
-On GNU/Linux and MacOS, an optional `scikits.odes <https://scikits-odes.readthedocs.io/en/latest/>`_ -based solver is available, see :ref:`scikits.odes-label`.
+Following GNU/Linux and macOS solvers are optionally available:
+
+*  `scikits.odes <https://scikits-odes.readthedocs.io/en/latest/>`_ -based solver, see `Optional - scikits.odes solver <https://pybamm.readthedocs.io/en/latest/install/GNU-linux.html#optional-scikits-odes-solver>`_.
+*  `jax <https://jax.readthedocs.io/en/latest/notebooks/quickstart.html>`_ -based solver, see `Optional - JaxSolver <https://pybamm.readthedocs.io/en/latest/install/GNU-linux.html#optional-jaxsolver>`_.
 
 Installation
 ============

docs/install/GNU-linux.rst

@@ -45,8 +45,8 @@ User install
 ------------
 
 We recommend to install PyBaMM within a virtual environment, in order
-not to alter any distribution python files. 
-First, make sure you are using python 3.7, 3.8, or 3.9. 
+not to alter any distribution python files.
+First, make sure you are using python 3.7, 3.8, or 3.9.
 To create a virtual environment ``env`` within your current directory type:
 
 .. code:: bash
@@ -116,9 +116,24 @@ macOS
 .. code:: bash
 
 	  pip install scikits.odes
-   
+
 Assuming that the SUNDIALS were installed as described :ref:`above<user-install-label>`.
 
+Optional - JaxSolver
+--------------------
+
+Users can install ``jax`` and ``jaxlib`` to use the Jax solver.
+Currently, only GNU/Linux and macOS are supported.
+
+GNU/Linux and macOS
+~~~~~~~~~~~~~~~~~~~
+
+.. code:: bash
+
+	  pybamm_install_jax
+
+The ``pybamm_install_jax`` command is installed with PyBaMM. It automatically downloads and installs jax and jaxlib on your system.
+
 Developer install
 -----------------
 

docs/requirements.txt

@@ -7,8 +7,6 @@ autograd >= 1.2
 scikit-fem >= 0.2.0
 casadi >= 3.5.0
 imageio>=2.9.0
-jax==0.2.12
-jaxlib==0.1.70
 jupyter  # For example notebooks
 pybtex
 # Note: Matplotlib is loaded for debug plots but to ensure pybamm runs
@@ -19,4 +17,4 @@ matplotlib >= 2.0
 #
 guzzle-sphinx-theme
 sphinx>4.0
-sympy==1.8
+sympy==1.9

examples/scripts/experiment_drive_cycle.py

@@ -0,0 +1,52 @@
+#
+# Constant-current constant-voltage charge with US06 Drive Cycle using Experiment Class.
+#
+import pybamm
+import pandas as pd
+import os
+
+os.chdir(pybamm.__path__[0] + "/..")
+
+pybamm.set_logging_level("INFO")
+
+# import drive cycle from file
+drive_cycle_current = pd.read_csv(
+    "pybamm/input/drive_cycles/US06.csv", comment="#", header=None
+).to_numpy()
+
+
+# Map Drive Cycle
+def map_drive_cycle(x, min_op_value, max_op_value):
+    min_ip_value = x[:, 1].min()
+    max_ip_value = x[:, 1].max()
+    x[:, 1] = (x[:, 1] - min_ip_value) / (max_ip_value - min_ip_value) * (
+        max_op_value - min_op_value
+    ) + min_op_value
+    return x
+
+
+# Map current drive cycle to voltage and power
+drive_cycle_power = map_drive_cycle(drive_cycle_current, 1.5, 3.5)
+
+experiment = pybamm.Experiment(
+    [
+        "Charge at 1 A until 4.0 V",
+        "Hold at 4.0 V until 50 mA",
+        "Rest for 30 minutes",
+        "Run US06_A (A)",
+        "Rest for 30 minutes",
+        "Run US06_W (W)",
+        "Rest for 30 minutes",
+    ],
+    drive_cycles={
+        "US06_A": drive_cycle_current,
+        "US06_W": drive_cycle_power,
+    },
+)
+
+model = pybamm.lithium_ion.DFN()
+sim = pybamm.Simulation(model, experiment=experiment, solver=pybamm.CasadiSolver())
+sim.solve()
+
+# Show all plots
+sim.plot()

pybamm/CITATIONS.txt

@@ -192,17 +192,17 @@
 }
 
 @article{Kirk2021,
-	author = {Toby L. Kirk and Colin P. Please and S. Jon Chapman},
-	title = {Physical Modelling of the Slow Voltage Relaxation Phenomenon in Lithium-Ion Batteries},
-	journal = {Journal of The Electrochemical Society},
-	year = 2021,
-	month = {jun},
-	publisher = {The Electrochemical Society},
-	volume = {168},
-	number = {6},
-	pages = {060554},
-	doi = {10.1149/1945-7111/ac0bf7},
-	url = {https://doi.org/10.1149/1945-7111/ac0bf7},
+  author = {Toby L. Kirk and Colin P. Please and S. Jon Chapman},
+  title = {Physical Modelling of the Slow Voltage Relaxation Phenomenon in Lithium-Ion Batteries},
+  journal = {Journal of The Electrochemical Society},
+  year = 2021,
+  month = {jun},
+  publisher = {The Electrochemical Society},
+  volume = {168},
+  number = {6},
+  pages = {060554},
+  doi = {10.1149/1945-7111/ac0bf7},
+  url = {https://doi.org/10.1149/1945-7111/ac0bf7},
 }
 
 @article{Lain2019,
@@ -283,17 +283,17 @@
 }
 
 @article{OKane2020,
-	doi = {10.1149/1945-7111/ab90ac},
-	url = {https://doi.org/10.1149/1945-7111/ab90ac},
-	year = {2020},
-	month = {may},
-	publisher = {The Electrochemical Society},
-	volume = {167},
-	number = {9},
-	pages = {090540},
-	author = {Simon E. J. O'Kane and Ian D. Campbell and Mohamed W. J. Marzook and Gregory J. Offer and Monica Marinescu},
-	title = {Physical Origin of the Differential Voltage Minimum Associated with Lithium Plating in Li-Ion Batteries},
-	journal = {Journal of The Electrochemical Society}
+  doi = {10.1149/1945-7111/ab90ac},
+  url = {https://doi.org/10.1149/1945-7111/ab90ac},
+  year = {2020},
+  month = {may},
+  publisher = {The Electrochemical Society},
+  volume = {167},
+  number = {9},
+  pages = {090540},
+  author = {Simon E. J. O'Kane and Ian D. Campbell and Mohamed W. J. Marzook and Gregory J. Offer and Monica Marinescu},
+  title = {Physical Origin of the Differential Voltage Minimum Associated with Lithium Plating in Li-Ion Batteries},
+  journal = {Journal of The Electrochemical Society}
 }
 
 @article{ORegan2021,

pybamm/__init__.py

@@ -7,7 +7,7 @@
 #
 import sys
 import os
-import platform
+
 
 #
 # Version info
@@ -66,7 +66,7 @@ def version(formatted=False):
 #
 from .util import Timer, TimerTime, FuzzyDict
 from .util import root_dir, load_function, rmse, get_infinite_nested_dict, load
-from .util import get_parameters_filepath
+from .util import get_parameters_filepath, have_jax, install_jax
 from .logger import logger, set_logging_level
 from .settings import settings
 from .citations import Citations, citations, print_citations
@@ -102,12 +102,8 @@ def version(formatted=False):
     EvaluatorPython,
 )
 
-if not (
-    platform.system() == "Windows"
-    or (platform.system() == "Darwin" and "ARM64" in platform.version())
-):
-    from .expression_tree.operations.evaluate_python import EvaluatorJax
-    from .expression_tree.operations.evaluate_python import JaxCooMatrix
+from .expression_tree.operations.evaluate_python import EvaluatorJax
+from .expression_tree.operations.evaluate_python import JaxCooMatrix
 
 from .expression_tree.operations.jacobian import Jacobian
 from .expression_tree.operations.convert_to_casadi import CasadiConverter
@@ -226,13 +222,8 @@ def version(formatted=False):
 from .solvers.scikits_ode_solver import ScikitsOdeSolver, have_scikits_odes
 from .solvers.scipy_solver import ScipySolver
 
-# Jax not supported under windows
-if not (
-    platform.system() == "Windows"
-    or (platform.system() == "Darwin" and "ARM64" in platform.version())
-):
-    from .solvers.jax_solver import JaxSolver
-    from .solvers.jax_bdf_solver import jax_bdf_integrate
+from .solvers.jax_solver import JaxSolver
+from .solvers.jax_bdf_solver import jax_bdf_integrate
 
 from .solvers.idaklu_solver import IDAKLUSolver, have_idaklu