Remove old JAX restrictions for Winows and test cleanup

noxfile.py

@@ -4,7 +4,6 @@
 from pathlib import Path
 
 
-# Options to modify nox behaviour
 nox.options.default_venv_backend = "virtualenv"
 nox.options.reuse_existing_virtualenvs = True
 if sys.platform != "win32":
@@ -39,7 +38,8 @@ def set_environment_variables(env_dict, session):
 
 @nox.session(name="pybamm-requires")
 def run_pybamm_requires(session):
-    """Download, compile, and install the build-time requirements for Linux and macOS. Supports --install-dir for custom installation paths and --force to force installation."""
+    """Download, compile, and install the build-time requirements for Linux and macOS.
+    Supports --install-dir for custom installation paths and --force to force installation."""
     set_environment_variables(PYBAMM_ENV, session=session)
     if sys.platform != "win32":
         session.install("cmake", silent=False)
@@ -61,27 +61,21 @@ def run_coverage(session):
     """Run the coverage tests and generate an XML report."""
     set_environment_variables(PYBAMM_ENV, session=session)
     session.install("coverage", silent=False)
-    if sys.platform != "win32":
-        session.install("-e", ".[all,dev,jax]", silent=False)
+    if sys.version_info < (3, 9):
+        session.install("-e", ".[all,dev]", silent=False)
     else:
-        if sys.version_info < (3, 9):
-            session.install("-e", ".[all,dev]", silent=False)
-        else:
-            session.install("-e", ".[all,dev,jax]", silent=False)
+        session.install("-e", ".[all,dev,jax]", silent=False)
     session.run("pytest", "--cov=pybamm", "--cov-report=xml", "tests/unit")
 
 
 @nox.session(name="integration")
 def run_integration(session):
     """Run the integration tests."""
     set_environment_variables(PYBAMM_ENV, session=session)
-    if sys.platform != "win32":
-        session.install("-e", ".[all,dev,jax]", silent=False)
+    if sys.version_info < (3, 9):
+        session.install("-e", ".[all,dev]", silent=False)
     else:
-        if sys.version_info < (3, 9):
-            session.install("-e", ".[all,dev]", silent=False)
-        else:
-            session.install("-e", ".[all,dev,jax]", silent=False)
+        session.install("-e", ".[all,dev,jax]", silent=False)
     session.run("python", "run-tests.py", "--integration")
 
 
@@ -96,13 +90,10 @@ def run_doctests(session):
 def run_unit(session):
     """Run the unit tests."""
     set_environment_variables(PYBAMM_ENV, session=session)
-    if sys.platform != "win32":
-        session.install("-e", ".[all,dev,jax]", silent=False)
+    if sys.version_info < (3, 9):
+        session.install("-e", ".[all,dev]", silent=False)
     else:
-        if sys.version_info < (3, 9):
-            session.install("-e", ".[all,dev]", silent=False)
-        else:
-            session.install("-e", ".[all,dev,jax]", silent=False)
+        session.install("-e", ".[all,dev,jax]", silent=False)
     session.run("python", "run-tests.py", "--unit")
 
 
@@ -138,50 +129,36 @@ def set_dev(session):
     # https://bitbucket.org/pybtex-devs/pybtex/issues/169/replace-pkg_resources-with
     # is fixed
     session.run(python, "-m", "pip", "install", "setuptools", external=True)
-    if sys.platform == "linux":
+    if sys.version_info < (3, 9):
         session.run(
             python,
             "-m",
             "pip",
             "install",
             "-e",
-            ".[all,dev,jax]",
+            ".[all,dev]",
             external=True,
         )
     else:
-        if sys.version_info < (3, 9):
-            session.run(
-                python,
-                "-m",
-                "pip",
-                "install",
-                "-e",
-                ".[all,dev]",
-                external=True,
-            )
-        else:
-            session.run(
-                python,
-                "-m",
-                "pip",
-                "install",
-                "-e",
-                ".[all,dev,jax]",
-                external=True,
-            )
+        session.run(
+            python,
+            "-m",
+            "pip",
+            "install",
+            "-e",
+            ".[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,dev,jax]", silent=False)
+    if sys.version_info < (3, 9):
+        session.install("-e", ".[all,dev]", silent=False)
     else:
-        if sys.version_info < (3, 9):
-            session.install("-e", ".[all,dev]", silent=False)
-        else:
-            session.install("-e", ".[all,dev,jax]", silent=False)
+        session.install("-e", ".[all,dev,jax]", silent=False)
     session.run("python", "run-tests.py", "--all")
 
 

pybamm/expression_tree/operations/latexify.py

@@ -79,7 +79,7 @@ def _get_geometry_displays(self, var):
         for _, rng in self.model.default_geometry[var.domain[-1]].items():
             rng_max = get_rng_min_max_name(rng, "max")
 
-        geo_latex = rf"\quad {rng_min} < {name} < {rng_max}"
+        geo_latex = f"\quad {rng_min} < {name} < {rng_max}"
         geo.append(geo_latex)
 
         return geo

pybamm/expression_tree/printing/sympy_overrides.py

@@ -14,7 +14,7 @@ def _print_Derivative(self, expr):
         eqn = super()._print_Derivative(expr)
         if getattr(expr, "force_partial", False) and "partial" not in eqn:
             var1, var2 = re.findall(r"^\\frac{(\w+)}{(\w+) .+", eqn)[0]
-            eqn = eqn.replace(var1, r"\partial").replace(var2, r"\partial")
+            eqn = eqn.replace(var1, "\partial").replace(var2, "\partial")
 
         return eqn
 

pybamm/plotting/quick_plot.py

@@ -160,7 +160,7 @@ def __init__(
             self.spatial_unit = "mm"
         elif spatial_unit == "um":  # micrometers
             self.spatial_factor = 1e6
-            self.spatial_unit = r"$\mu$m"
+            self.spatial_unit = "$\mu$m"
         else:
             raise ValueError(f"spatial unit '{spatial_unit}' not recognized")