Rename have_x to has_x to improve how logic reads (#4398)

* Rename functions

* style: pre-commit fixes

* Update CHANGELOG.md

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

CHANGELOG.md

@@ -19,6 +19,8 @@
 
 ## Breaking changes
 
+- Replaced `have_jax` with `has_jax`, `have_idaklu` with `has_idaklu`, and
+  `have_iree` with `has_iree` ([#4398](https://github.com/pybamm-team/PyBaMM/pull/4398))
 - Remove deprecated function `pybamm_install_jax` ([#4362](https://github.com/pybamm-team/PyBaMM/pull/4362))
 - Removed legacy python-IDAKLU solver. ([#4326](https://github.com/pybamm-team/PyBaMM/pull/4326))
 

docs/source/api/util.rst

@@ -16,6 +16,6 @@ Utility functions
 
 .. autofunction:: pybamm.load
 
-.. autofunction:: pybamm.have_jax
+.. autofunction:: pybamm.has_jax
 
 .. autofunction:: pybamm.is_jax_compatible

examples/scripts/compare_dae_solver.py

@@ -28,7 +28,7 @@
 casadi_sol = pybamm.CasadiSolver(atol=1e-8, rtol=1e-8).solve(model, t_eval)
 solutions = [casadi_sol]
 
-if pybamm.have_idaklu():
+if pybamm.has_idaklu():
     klu_sol = pybamm.IDAKLUSolver(atol=1e-8, rtol=1e-8).solve(model, t_eval)
     solutions.append(klu_sol)
 else:

src/pybamm/__init__.py

@@ -15,7 +15,7 @@
 )
 from .util import (
     get_parameters_filepath,
-    have_jax,
+    has_jax,
     import_optional_dependency,
     is_jax_compatible,
     get_git_commit_info,
@@ -170,7 +170,7 @@
 from .solvers.jax_bdf_solver import jax_bdf_integrate
 
 from .solvers.idaklu_jax import IDAKLUJax
-from .solvers.idaklu_solver import IDAKLUSolver, have_idaklu, have_iree
+from .solvers.idaklu_solver import IDAKLUSolver, has_idaklu, has_iree
 
 # Experiments
 from .experiment.experiment import Experiment

src/pybamm/expression_tree/operations/evaluate_python.py

@@ -11,7 +11,7 @@
 
 import pybamm
 
-if pybamm.have_jax():
+if pybamm.has_jax():
     import jax
 
     platform = jax.lib.xla_bridge.get_backend().platform.casefold()
@@ -43,7 +43,7 @@ class JaxCooMatrix:
     def __init__(
         self, row: ArrayLike, col: ArrayLike, data: ArrayLike, shape: tuple[int, int]
     ):
-        if not pybamm.have_jax():  # pragma: no cover
+        if not pybamm.has_jax():  # pragma: no cover
             raise ModuleNotFoundError(
                 "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
             )
@@ -527,7 +527,7 @@ class EvaluatorJax:
     """
 
     def __init__(self, symbol: pybamm.Symbol):
-        if not pybamm.have_jax():  # pragma: no cover
+        if not pybamm.has_jax():  # pragma: no cover
             raise ModuleNotFoundError(
                 "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
             )

src/pybamm/solvers/idaklu_jax.py

@@ -22,7 +22,7 @@
     except ImportError:  # pragma: no cover
         idaklu_spec = None
 
-if pybamm.have_jax():
+if pybamm.has_jax():
     import jax
     from jax import lax
     from jax import numpy as jnp
@@ -57,11 +57,11 @@ def __init__(
         calculate_sensitivities=True,
         t_interp=None,
     ):
-        if not pybamm.have_jax():
+        if not pybamm.has_jax():
             raise ModuleNotFoundError(
                 "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
             )  # pragma: no cover
-        if not pybamm.have_idaklu():
+        if not pybamm.has_idaklu():
             raise ModuleNotFoundError(
                 "IDAKLU is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/index.html"
             )  # pragma: no cover

src/pybamm/solvers/idaklu_solver.py

@@ -14,7 +14,7 @@
 import warnings
 
 
-if pybamm.have_jax():
+if pybamm.has_jax():
     import jax
     from jax import numpy as jnp
 
@@ -33,11 +33,11 @@
         idaklu_spec = None
 
 
-def have_idaklu():
+def has_idaklu():
     return idaklu_spec is not None
 
 
-def have_iree():
+def has_iree():
     try:
         import iree.compiler  # noqa: F401
 

src/pybamm/solvers/jax_bdf_solver.py

@@ -7,7 +7,7 @@
 
 import pybamm
 
-if pybamm.have_jax():
+if pybamm.has_jax():
     import jax
     import jax.numpy as jnp
     from jax import core, dtypes
@@ -1007,7 +1007,7 @@ def jax_bdf_integrate(func, y0, t_eval, *args, rtol=1e-6, atol=1e-6, mass=None):
         calculated state vector at each of the m time points
 
     """
-    if not pybamm.have_jax():
+    if not pybamm.has_jax():
         raise ModuleNotFoundError(
             "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
         )

src/pybamm/solvers/jax_solver.py

@@ -6,7 +6,7 @@
 
 import pybamm
 
-if pybamm.have_jax():
+if pybamm.has_jax():
     import jax
     import jax.numpy as jnp
     from jax.experimental.ode import odeint
@@ -59,7 +59,7 @@ def __init__(
         extrap_tol=None,
         extra_options=None,
     ):
-        if not pybamm.have_jax():
+        if not pybamm.has_jax():
             raise ModuleNotFoundError(
                 "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
             )

src/pybamm/util.py

@@ -264,7 +264,7 @@ def get_parameters_filepath(path):
         return os.path.join(pybamm.__path__[0], path)
 
 
-def have_jax():
+def has_jax():
     """
     Check if jax and jaxlib are installed with the correct versions
 

tests/integration/test_models/test_full_battery_models/test_lithium_ion/base_lithium_ion_tests.py

@@ -22,7 +22,7 @@ def test_sensitivities(self):
         param = pybamm.ParameterValues("Ecker2015")
         rtol = 1e-6
         atol = 1e-6
-        if pybamm.have_idaklu():
+        if pybamm.has_idaklu():
             solver = pybamm.IDAKLUSolver(rtol=rtol, atol=atol)
         else:
             solver = pybamm.CasadiSolver(rtol=rtol, atol=atol)
@@ -53,7 +53,7 @@ def test_optimisations(self):
         to_python = optimtest.evaluate_model(to_python=True)
         np.testing.assert_array_almost_equal(original, to_python)
 
-        if pybamm.have_jax():
+        if pybamm.has_jax():
             to_jax = optimtest.evaluate_model(to_jax=True)
             np.testing.assert_array_almost_equal(original, to_jax)
 

tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_mpm.py

@@ -25,7 +25,7 @@ def test_optimisations(self):
         to_python = optimtest.evaluate_model(to_python=True)
         np.testing.assert_array_almost_equal(original, to_python)
 
-        if pybamm.have_jax():
+        if pybamm.has_jax():
             to_jax = optimtest.evaluate_model(to_jax=True)
             np.testing.assert_array_almost_equal(original, to_jax)
 

tests/integration/test_solvers/test_idaklu.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 
-@pytest.mark.skipif(not pybamm.have_idaklu(), reason="idaklu solver is not installed")
+@pytest.mark.skipif(not pybamm.has_idaklu(), reason="idaklu solver is not installed")
 class TestIDAKLUSolver:
     def test_on_spme(self):
         model = pybamm.lithium_ion.SPMe()

tests/unit/test_citations.py

@@ -423,14 +423,14 @@ def test_solver_citations(self):
         assert "Virtanen2020" in citations._papers_to_cite
         assert "Virtanen2020" in citations._citation_tags.keys()
 
-        if pybamm.have_idaklu():
+        if pybamm.has_idaklu():
             citations._reset()
             assert "Hindmarsh2005" not in citations._papers_to_cite
             pybamm.IDAKLUSolver()
             assert "Hindmarsh2005" in citations._papers_to_cite
             assert "Hindmarsh2005" in citations._citation_tags.keys()
 
-    @pytest.mark.skipif(not pybamm.have_jax(), reason="jax or jaxlib is not installed")
+    @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
     def test_jax_citations(self):
         citations = pybamm.citations
         citations._reset()

tests/unit/test_experiments/test_simulation_with_experiment.py

@@ -168,7 +168,7 @@ def test_run_experiment_multiple_times(self):
             sol1["Voltage [V]"].data, sol2["Voltage [V]"].data
         )
 
-    @unittest.skipIf(not pybamm.have_idaklu(), "idaklu solver is not installed")
+    @unittest.skipIf(not pybamm.has_idaklu(), "idaklu solver is not installed")
     def test_run_experiment_cccv_solvers(self):
         experiment_2step = pybamm.Experiment(
             [

tests/unit/test_expression_tree/test_operations/test_evaluate_python.py

@@ -12,7 +12,7 @@
 from collections import OrderedDict
 import re
 
-if pybamm.have_jax():
+if pybamm.has_jax():
     import jax
 from tests import (
     function_test,
@@ -446,7 +446,7 @@ def test_evaluator_python(self):
             result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
-    @pytest.mark.skipif(not pybamm.have_jax(), reason="jax or jaxlib is not installed")
+    @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
     def test_find_symbols_jax(self):
         # test sparse conversion
         constant_symbols = OrderedDict()
@@ -459,7 +459,7 @@ def test_find_symbols_jax(self):
             next(iter(constant_symbols.values())).toarray(), A.entries.toarray()
         )
 
-    @pytest.mark.skipif(not pybamm.have_jax(), reason="jax or jaxlib is not installed")
+    @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
     def test_evaluator_jax(self):
         a = pybamm.StateVector(slice(0, 1))
         b = pybamm.StateVector(slice(1, 2))
@@ -621,7 +621,7 @@ def test_evaluator_jax(self):
                 result = evaluator(t=t, y=y)
                 np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
-    @pytest.mark.skipif(not pybamm.have_jax(), reason="jax or jaxlib is not installed")
+    @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
     def test_evaluator_jax_jacobian(self):
         a = pybamm.StateVector(slice(0, 1))
         y_tests = [np.array([[2.0]]), np.array([[1.0]]), np.array([1.0])]
@@ -636,7 +636,7 @@ def test_evaluator_jax_jacobian(self):
             result_true = evaluator_jac(t=None, y=y)
             np.testing.assert_allclose(result_test, result_true)
 
-    @pytest.mark.skipif(not pybamm.have_jax(), reason="jax or jaxlib is not installed")
+    @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
     def test_evaluator_jax_jvp(self):
         a = pybamm.StateVector(slice(0, 1))
         y_tests = [np.array([[2.0]]), np.array([[1.0]]), np.array([1.0])]
@@ -656,23 +656,23 @@ def test_evaluator_jax_jvp(self):
             np.testing.assert_allclose(result_test, result_true)
             np.testing.assert_allclose(result_test_times_v, result_true_times_v)
 
-    @pytest.mark.skipif(not pybamm.have_jax(), reason="jax or jaxlib is not installed")
+    @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
     def test_evaluator_jax_debug(self):
         a = pybamm.StateVector(slice(0, 1))
         expr = a**2
         y_test = np.array([2.0, 3.0])
         evaluator = pybamm.EvaluatorJax(expr)
         evaluator.debug(y=y_test)
 
-    @pytest.mark.skipif(not pybamm.have_jax(), reason="jax or jaxlib is not installed")
+    @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
     def test_evaluator_jax_inputs(self):
         a = pybamm.InputParameter("a")
         expr = a**2
         evaluator = pybamm.EvaluatorJax(expr)
         result = evaluator(inputs={"a": 2})
         assert result == 4
 
-    @pytest.mark.skipif(not pybamm.have_jax(), reason="jax or jaxlib is not installed")
+    @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
     def test_evaluator_jax_demotion(self):
         for demote in [True, False]:
             pybamm.demote_expressions_to_32bit = demote  # global flag
@@ -734,7 +734,7 @@ def test_evaluator_jax_demotion(self):
                 assert all(str(c_i.dtype)[-2:] == target_dtype for c_i in c_demoted.col)
             pybamm.demote_expressions_to_32bit = False
 
-    @pytest.mark.skipif(not pybamm.have_jax(), reason="jax or jaxlib is not installed")
+    @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
     def test_jax_coo_matrix(self):
         A = pybamm.JaxCooMatrix([0, 1], [0, 1], [1.0, 2.0], (2, 2))
         Adense = jax.numpy.array([[1.0, 0], [0, 2.0]])

tests/unit/test_solvers/test_base_solver.py

@@ -355,12 +355,12 @@ def test_multiprocess_context(self):
         assert solver.get_platform_context("Linux") == "fork"
         assert solver.get_platform_context("Darwin") == "fork"
 
-    @unittest.skipIf(not pybamm.have_idaklu(), "idaklu solver is not installed")
+    @unittest.skipIf(not pybamm.has_idaklu(), "idaklu solver is not installed")
     def test_sensitivities(self):
         def exact_diff_a(y, a, b):
             return np.array([[y[0] ** 2 + 2 * a], [y[0]]])
 
-        @unittest.skipIf(not pybamm.have_jax(), "jax or jaxlib is not installed")
+        @unittest.skipIf(not pybamm.has_jax(), "jax or jaxlib is not installed")
         def exact_diff_b(y, a, b):
             return np.array([[y[0]], [0]])
 

tests/unit/test_solvers/test_idaklu_jax.py

@@ -9,7 +9,7 @@
 import unittest
 
 testcase = []
-if pybamm.have_idaklu() and pybamm.have_jax():
+if pybamm.has_idaklu() and pybamm.has_jax():
     from jax.tree_util import tree_flatten
     import jax
     import jax.numpy as jnp
@@ -87,7 +87,7 @@ def no_jit(f):
 
 # Check the interface throws an appropriate error if either IDAKLU or JAX not available
 @unittest.skipIf(
-    pybamm.have_idaklu() and pybamm.have_jax(),
+    pybamm.has_idaklu() and pybamm.has_jax(),
     "Both IDAKLU and JAX are available",
 )
 class TestIDAKLUJax_NoJax(unittest.TestCase):
@@ -97,7 +97,7 @@ def test_instantiate_fails(self):
 
 
 @unittest.skipIf(
-    not pybamm.have_idaklu() or not pybamm.have_jax(),
+    not pybamm.has_idaklu() or not pybamm.has_jax(),
     "IDAKLU Solver and/or JAX are not available",
 )
 class TestIDAKLUJax(unittest.TestCase):