Resolve conflicts

pybamm-team · Nov 10, 2021 · 0e08601 · 0e08601
2 parents 0a57e76 + 337e602
commit 0e08601
Show file tree

Hide file tree

Showing 31 changed files with 1,395 additions and 1,186 deletions.
diff --git a/.github/release_checklist.md b/.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
diff --git a/.github/workflows/test_on_push.yml b/.github/workflows/test_on_push.yml
@@ -16,7 +16,7 @@ jobs:
     steps:
     - uses: actions/checkout@v2
     - name: Setup python
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v2
       with:
         python-version: 3.7
 
@@ -37,7 +37,7 @@ 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 }}
 
@@ -49,7 +49,7 @@ jobs:
         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'
@@ -98,4 +98,4 @@ jobs:
 
     - name: Upload coverage report
       if: matrix.os == 'ubuntu-latest' && matrix.python-version == 3.9
-      uses: codecov/codecov-action@v1
+      uses: codecov/codecov-action@v2.1.0
diff --git a/CHANGELOG.md b/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))
 
@@ -13,6 +14,7 @@
 ## Breaking changes
 
 -   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
 
@@ -25,7 +27,7 @@
 -   SEI models can now be included in the half-cell model ([#1705](https://github.com/pybamm-team/PyBaMM/pull/1705))
 
 ## Bug fixes
- 
+
 -   Half-cell model and lead-acid models can now be simulated with `Experiment`s ([#1759](https://github.com/pybamm-team/PyBaMM/pull/1759))
 -   Removed in-place modification of the solution objects by `QuickPlot` ([#1747](https://github.com/pybamm-team/PyBaMM/pull/1747))
 -   Fixed vector-vector multiplication bug that was causing errors in the SPM with constant voltage or power ([#1735](https://github.com/pybamm-team/PyBaMM/pull/1735))

diff --git a/README.md b/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
 

diff --git a/docs/index.rst b/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
 ============

diff --git a/docs/install/GNU-linux.rst b/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
 -----------------
 

diff --git a/docs/requirements.txt b/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

diff --git a/examples/scripts/experiment_drive_cycle.py b/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()
diff --git a/pybamm/__init__.py b/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
 

diff --git a/pybamm/experiments/experiment.py b/pybamm/experiments/experiment.py
@@ -195,6 +195,7 @@ def read_string(self, cond, drive_cycles):
                 "electric": op_CC["electric"] + op_CV["electric"],
                 "time": op_CV["time"],
                 "period": op_CV["period"],
+                "dc_data": None,
             }, event_CV
         # Read period
         if " period)" in cond:
@@ -223,26 +224,29 @@ def read_string(self, cond, drive_cycles):
                     drive_cycles[cond_list[1]], end_time
                 )
                 # Drive cycle as numpy array
+                dc_name = cond_list[1] + "_ext_{}".format(end_time)
                 dc_data = ext_drive_cycle
                 # Find the type of drive cycle ("A", "V", or "W")
                 typ = cond_list[2][1]
-                electric = (dc_data, typ)
+                electric = (dc_name, typ)
                 time = ext_drive_cycle[:, 0][-1]
                 period = np.min(np.diff(ext_drive_cycle[:, 0]))
                 events = None
             else:
                 # e.g. Run US06
                 # Drive cycle as numpy array
+                dc_name = cond_list[1]
                 dc_data = drive_cycles[cond_list[1]]
                 # Find the type of drive cycle ("A", "V", or "W")
                 typ = cond_list[2][1]
-                electric = (dc_data, typ)
+                electric = (dc_name, typ)
                 # Set time and period to 1 second for first step and
                 # then calculate the difference in consecutive time steps
                 time = drive_cycles[cond_list[1]][:, 0][-1]
                 period = np.min(np.diff(drive_cycles[cond_list[1]][:, 0]))
                 events = None
         else:
+            dc_data = None
             if "for" in cond and "or until" in cond:
                 # e.g. for 3 hours or until 4.2 V
                 cond_list = cond.split()
@@ -273,7 +277,12 @@ def read_string(self, cond, drive_cycles):
                     )
                 )
 
-        return {"electric": electric, "time": time, "period": period}, events
+        return {
+            "electric": electric,
+            "time": time,
+            "period": period,
+            "dc_data": dc_data,
+        }, events
 
     def extend_drive_cycle(self, drive_cycle, end_time):
         "Extends the drive cycle to enable for event"