pybamm-team · valentinsulzer · Jan 20, 2022 · Jan 1, 2022 · Jan 1, 2022 · Jan 3, 2022
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,9 @@
+# [Unreleased](https://github.com/pybamm-team/PyBaMM/)
+
+## Features
+
+-   Added an option to force install compatible versions of jax and jaxlib if already installed using CLI ([#1881](https://github.com/pybamm-team/PyBaMM/pull/1881))
+
 # [v21.12](https://github.com/pybamm-team/PyBaMM/tree/v21.11) - 2021-12-29
 
 ## Features

diff --git a/pybamm/input/parameters/lithium_ion/cells/LGM50_ORegan2021/__init__.py b/pybamm/input/parameters/lithium_ion/cells/LGM50_ORegan2021/__init__.py
diff --git a/pybamm/input/parameters/lithium_ion/electrolytes/lipf6_Valoen2005/__init__.py b/pybamm/input/parameters/lithium_ion/electrolytes/lipf6_Valoen2005/__init__.py
diff --git a/pybamm/input/parameters/lithium_ion/negative_electrodes/graphite_ORegan2021/__init__.py b/pybamm/input/parameters/lithium_ion/negative_electrodes/graphite_ORegan2021/__init__.py
diff --git a/pybamm/input/parameters/lithium_ion/negative_electrodes/li_metal_Xu2019/__init__.py b/pybamm/input/parameters/lithium_ion/negative_electrodes/li_metal_Xu2019/__init__.py
diff --git a/pybamm/input/parameters/lithium_ion/positive_electrodes/NMC532_Xu2019/__init__.py b/pybamm/input/parameters/lithium_ion/positive_electrodes/NMC532_Xu2019/__init__.py
diff --git a/pybamm/input/parameters/lithium_ion/positive_electrodes/nmc_ORegan2021/__init__.py b/pybamm/input/parameters/lithium_ion/positive_electrodes/nmc_ORegan2021/__init__.py
diff --git a/pybamm/input/parameters/lithium_ion/separators/separator_ORegan2021/__init__.py b/pybamm/input/parameters/lithium_ion/separators/separator_ORegan2021/__init__.py
diff --git a/pybamm/models/full_battery_models/lithium_metal/__init__.py b/pybamm/models/full_battery_models/lithium_metal/__init__.py
diff --git a/pybamm/parameters_cli.py b/pybamm/parameters_cli.py
@@ -143,7 +143,7 @@ def edit_parameter(arguments=None):
 
     .. code::
 
-      edit_param(["lithium_ion"])
+      edit_parameter(["lithium_ion"])
 
     will create the directory structure::
 

diff --git a/pybamm/util.py b/pybamm/util.py
@@ -4,6 +4,7 @@
 # The code in this file is adapted from Pints
 # (see https://github.com/pints-team/pints)
 #
+import argparse
 import importlib.util
 import numbers
 import os
@@ -15,9 +16,9 @@
 import warnings
 from collections import defaultdict
 from platform import system
-import pkg_resources
 
 import numpy as np
+import pkg_resources
 
 import pybamm
 
@@ -364,23 +365,50 @@ def is_jax_compatible():
     )
 
 
-def install_jax():  # pragma: no cover
-    """Install jax, jaxlib"""
+def install_jax(arguments=None):  # pragma: no cover
+    """
+    Install compatible versions of jax, jaxlib.
+
+    Command Line Interface:
+    -----------------------
+    >>> pybamm_install_jax
+
+    optional arguments:
+    -h, --help   show help message
+    -f, --force  force install compatible versions of jax and jaxlib
+    """
+    parser = argparse.ArgumentParser(description="Install jax and jaxlib")
+    parser.add_argument(
+        "-f",
+        "--force",
+        action="store_true",
+        help=f"force install compatible versions of jax ({JAX_VERSION}) and jaxlib ({JAXLIB_VERSION})",  # noqa: E501
+    )
+
+    args = parser.parse_args(arguments)
+
     if system() == "Windows":
         raise NotImplementedError("Jax is not available on Windows")
+
+    # Raise an error if jax and jaxlib are already installed, but incompatible
+    # and --force is not set
     elif importlib.util.find_spec("jax") is not None:
-        if not is_jax_compatible():
+        if not args.force and not is_jax_compatible():
             raise ValueError(
-                "Jax is already installed but the installed version of jax or jaxlib is not supported by PyBaMM",  # noqa: E501
+                f"""
+Jax is already installed but the installed version of jax or jaxlib is not supported by PyBaMM.
+
+You can force install compatible versions of jax ({JAX_VERSION}) and jaxlib ({JAXLIB_VERSION}) using the following command:
+pybamm_install_jax --force""",  # noqa: E501
             )
-    else:
-        subprocess.check_call(
-            [
-                sys.executable,
-                "-m",
-                "pip",
-                "install",
-                f"jax=={JAX_VERSION}",
-                f"jaxlib=={JAXLIB_VERSION}",
-            ]
-        )
+
+    subprocess.check_call(
+        [
+            sys.executable,
+            "-m",
+            "pip",
+            "install",
+            f"jax=={JAX_VERSION}",
+            f"jaxlib=={JAXLIB_VERSION}",
+        ]
+    )