Merge pull request #3417 from arjxn-py/install_odes

Extend `pybamm_install_odes` to include support for macOS systems
pybamm-team · Jan 6, 2024 · 0b87513 · 0b87513
2 parents 0369f0c + 738cd57
commit 0b87513
Show file tree

Hide file tree

Showing 12 changed files with 130 additions and 42 deletions.
diff --git a/.github/workflows/run_periodic_tests.yml b/.github/workflows/run_periodic_tests.yml
@@ -153,3 +153,50 @@ jobs:
           eval "$(pyenv init -)"
           pyenv activate pybamm-${{ matrix.python-version }}
           pyenv uninstall -f $( python --version )
+
+  test_install_odes:
+    runs-on: ${{ matrix.os }}
+    strategy:
+        matrix:
+            os: [ubuntu-latest, macos-latest]
+            python-version: ["3.8", "3.9", "3.10", "3.11"]
+        fail-fast: false
+    name: Test pybamm_install_odes on ${{ matrix.os }}
+
+    steps:
+      - name: Check out PyBaMM repository
+        uses: actions/checkout@v4
+
+      - name: Install Linux system dependencies
+        if: matrix.os == 'ubuntu-latest'
+        run: |
+          sudo apt-get update
+          sudo apt-get install gfortran gcc libopenblas-dev
+      - name: Install macOS system dependencies
+        if: matrix.os == 'macos-latest'
+        env:
+          # Homebrew environment variables
+          HOMEBREW_NO_INSTALL_CLEANUP: 1
+          HOMEBREW_NO_AUTO_UPDATE: 1
+          HOMEBREW_NO_COLOR: 1
+          # Speed up CI
+          NONINTERACTIVE: 1
+        run: |
+          brew analytics off
+          brew install openblas
+          brew reinstall gcc gfortran
+
+      - name: Set up Python ${{ matrix.python-version }}
+        id: setup-python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install PyBaMM
+        run: python -m pip install -e .
+
+      - name: Test pybamm_install_odes on ${{ matrix.os }}
+        run: |
+          python -m pip cache purge
+          python -m pip install wget cmake
+          pybamm_install_odes
diff --git a/.gitignore b/.gitignore
@@ -107,6 +107,9 @@ KLU_module_deps
 # setup
 setup.log
 
+# odes setup
+scikits_odes_setup.log
+
 # test
 test.c
 test.json

diff --git a/.lycheeignore b/.lycheeignore
@@ -1,4 +1,5 @@
 # a list of links/files to be ignored by lychee link checker (see workflow file)
+https://github.com/LLNL/sundials/releases/download/v%7BSUNDIALS_VERSION%7D/sundials-%7BSUNDIALS_VERSION%7D.tar.gz
 
 # Errors in docs/source/user_guide/getting_started.md
 file:///home/runner/work/PyBaMM/PyBaMM/docs/source/user_guide/api_docs

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Features
 
+- The `pybamm_install_odes` command now includes support for macOS systems and can be used to set up SUNDIALS and install the `scikits.odes` solver on macOS ([#3417](https://github.com/pybamm-team/PyBaMM/pull/3417))
 - Added support for Python 3.12 ([#3531](https://github.com/pybamm-team/PyBaMM/pull/3531))
 - Added method to get QuickPlot axes by variable ([#3596](https://github.com/pybamm-team/PyBaMM/pull/3596))
 - Added custom experiment terminations ([#3596](https://github.com/pybamm-team/PyBaMM/pull/3596))

diff --git a/README.md b/README.md
@@ -104,7 +104,7 @@ PyBaMM makes releases every four months and we use [CalVer](https://calver.org/)
 
 PyBaMM is available on GNU/Linux, MacOS and Windows.
 We strongly recommend to install PyBaMM within a python virtual environment, in order not to alter any distribution python files.
-For instructions on how to create a virtual environment for PyBaMM, see [the documentation](https://docs.pybamm.org/en/latest/source/user_guide/installation/GNU-linux.html#user-install).
+For instructions on how to create a virtual environment for PyBaMM, see [the documentation](https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#user-install).
 
 ### Using pip
 
@@ -130,8 +130,8 @@ conda install -c conda-forge pybamm
 
 Following GNU/Linux and macOS solvers are optionally available:
 
-- [scikits.odes](https://scikits-odes.readthedocs.io/en/latest/)-based solver, see [the documentation](https://docs.pybamm.org/en/latest/source/user_guide/installation/GNU-linux.html#optional-scikits-odes-solver).
-- [jax](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html)-based solver, see [the documentation](https://docs.pybamm.org/en/latest/source/user_guide/installation/GNU-linux.html#optional-jaxsolver).
+- [scikits.odes](https://scikits-odes.readthedocs.io/en/latest/)-based solver, see [the documentation](https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-scikits-odes-solver).
+- [jax](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html)-based solver, see [the documentation](https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver).
 
 ## 📖 Citing PyBaMM
 

diff --git a/...rce/user_guide/installation/GNU-linux.rst → ...user_guide/installation/gnu-linux-mac.rst b/...rce/user_guide/installation/GNU-linux.rst → ...user_guide/installation/gnu-linux-mac.rst
@@ -1,4 +1,4 @@
-GNU-Linux & MacOS
+GNU/Linux & macOS
 =================
 
 .. contents::
@@ -102,17 +102,15 @@ For an introduction to virtual environments, see
 Optional - scikits.odes solver
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Users can install `scikits.odes <https://github.com/bmcage/odes>`__ in
-order to use the wrapped SUNDIALS ODE and DAE
-`solvers <https://docs.pybamm.org/en/latest/source/api/solvers/scikits_solvers.html>`__.
+Users can install `scikits.odes <https://github.com/bmcage/odes>`__ to utilize its interfaced SUNDIALS ODE and DAE `solvers <https://docs.pybamm.org/en/latest/source/api/solvers/scikits_solvers.html>`__ wrapped in PyBaMM.
 
 .. note::
 
    Currently, only GNU/Linux and macOS are supported.
 
 .. note::
 
-   The ``scikits.odes`` solver is not supported on Python 3.12 yet, please refer to https://github.com/bmcage/odes/issues/162.
+   The ``scikits.odes`` solver is not supported on Python 3.12 yet. Please refer to https://github.com/bmcage/odes/issues/162.
    There is support for Python 3.8, 3.9, 3.10, and 3.11.
 
 .. tab:: GNU/Linux
@@ -121,10 +119,9 @@ order to use the wrapped SUNDIALS ODE and DAE
 
    .. code:: bash
 
-	   apt install libopenblas-dev
-	   pybamm_install_odes
+      apt install libopenblas-dev
+      pybamm_install_odes
 
-   The ``pybamm_install_odes`` command is installed with PyBaMM. It automatically downloads and installs the SUNDIALS library on your
    system (under ``~/.local``), before installing ``scikits.odes``. (Alternatively, one can install SUNDIALS without this script and run ``pip install pybamm[odes]`` to install ``pybamm`` with ``scikits.odes``.)
 
 .. tab:: macOS
@@ -133,9 +130,29 @@ order to use the wrapped SUNDIALS ODE and DAE
 
    .. code:: bash
 
-	  pip install scikits.odes
+      brew install openblas gcc gfortran
+      pybamm_install_odes
 
-   Assuming that SUNDIALS was installed as described :ref:`above<user-install-label>`.
+The ``pybamm_install_odes`` command, installed with PyBaMM, automatically downloads and installs the SUNDIALS library on your
+system (under ``~/.local``), before installing `scikits.odes <https://scikits-odes.readthedocs.io/en/stable/installation.html>`__ . (Alternatively, one can install SUNDIALS without this script and run ``pip install pybamm[odes]`` to install ``pybamm`` with `scikits.odes <https://scikits-odes.readthedocs.io/en/stable/installation.html>`__)
+
+To avoid installation failures when using ``pip install pybamm[odes]``, make sure to set the ``SUNDIALS_INST`` environment variable. If you have installed SUNDIALS using Homebrew, set the variable to the appropriate location. For example:
+
+.. code:: bash
+
+   export SUNDIALS_INST=$(brew --prefix sundials)
+
+Ensure that the path matches the installation location on your system. You can verify the installation location by running:
+
+.. code:: bash
+
+   brew info sundials
+
+Look for the installation path, and use that path to set the ``SUNDIALS_INST`` variable.
+
+Note: The location where Homebrew installs SUNDIALS might vary based on the system architecture (ARM or Intel). Adjust the path in the ``export SUNDIALS_INST`` command accordingly.
+
+To avoid manual setup of path the ``pybamm_install_odes`` is recommended for a smoother installation process, as it takes care of automatically downloading and installing the SUNDIALS library on your system.
 
 Optional - JaxSolver
 ~~~~~~~~~~~~~~~~~~~~

diff --git a/docs/source/user_guide/installation/index.rst b/docs/source/user_guide/installation/index.rst
@@ -47,8 +47,8 @@ Optional solvers
 
 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://docs.pybamm.org/en/latest/source/user_guide/installation/GNU-linux.html#optional-scikits-odes-solver>`_.
-*  `jax <https://jax.readthedocs.io/en/latest/notebooks/quickstart.html>`_ -based solver, see `Optional - JaxSolver <https://docs.pybamm.org/en/latest/source/user_guide/installation/GNU-linux.html#optional-jaxsolver>`_.
+*  `scikits.odes <https://scikits-odes.readthedocs.io/en/latest/>`_ -based solver, see `Optional - scikits.odes solver <https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-scikits-odes-solver>`_.
+*  `jax <https://jax.readthedocs.io/en/latest/notebooks/quickstart.html>`_ -based solver, see `Optional - JaxSolver <https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver>`_.
 
 Dependencies
 ------------
@@ -233,15 +233,15 @@ odes dependencies
 
 Installable with ``pip install "pybamm[odes]"``
 
-================================================================================================================================ ================== ================== =============================
-Dependency                                                                                                                       Minimum Version    pip extra          Notes
-================================================================================================================================ ================== ================== =============================
-`scikits.odes <https://docs.pybamm.org/en/latest/source/user_guide/installation/GNU-linux.html#optional-scikits-odes-solver>`__  \-                 odes               For scikits ODE & DAE solvers
-================================================================================================================================ ================== ================== =============================
+======================================================================================================================================= ================== ================== =============================
+Dependency                                                                                                                              Minimum Version    pip extra          Notes
+======================================================================================================================================= ================== ================== =============================
+`scikits.odes <https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-scikits-odes-solver>`__     \-                 odes               For scikits ODE & DAE solvers
+======================================================================================================================================= ================== ================== =============================
 
 .. note::
 
-   Before running ``pip install "pybamm[odes]"``, make sure to install ``scikits.odes`` build-time requirements as described `here <https://docs.pybamm.org/en/latest/source/user_guide/installation/GNU-linux.html#optional-scikits-odes-solver>`_ .
+   Before running ``pip install "pybamm[odes]"``, make sure to install ``scikits.odes`` build-time requirements as described `here <https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-scikits-odes-solver>`_ .
 
 Full installation guide
 -----------------------
@@ -251,7 +251,7 @@ Installing a specific version? Installing from source? Check the advanced instal
 .. toctree::
    :maxdepth: 1
 
-   GNU-linux
+   gnu-linux-mac
    windows
    windows-wsl
    install-from-source

diff --git a/docs/source/user_guide/installation/windows-wsl.rst b/docs/source/user_guide/installation/windows-wsl.rst
@@ -37,7 +37,7 @@ Get PyBaMM's Source Code
 5. Follow the Installation Steps
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Follow the `installation instructions for PyBaMM on Linux <GNU-linux.html>`__.
+Follow the `installation instructions for PyBaMM on Linux <gnu-linux-mac.html>`__.
 
 Using Visual Studio Code with the WSL
 ---------------------------------------

diff --git a/pybamm/expression_tree/operations/evaluate_python.py b/pybamm/expression_tree/operations/evaluate_python.py
@@ -41,7 +41,7 @@ class JaxCooMatrix:
     def __init__(self, row, col, data, shape):
         if not pybamm.have_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.html#optional-jaxsolver"
+                "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
             )
 
         self.row = jax.numpy.array(row)
@@ -522,7 +522,7 @@ class EvaluatorJax:
     def __init__(self, symbol):
         if not pybamm.have_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.html#optional-jaxsolver"
+                "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
             )
 
         constants, python_str = pybamm.to_python(symbol, debug=False, output_jax=True)

diff --git a/pybamm/install_odes.py b/pybamm/install_odes.py
@@ -9,6 +9,14 @@
 
 from pybamm.util import root_dir
 
+if sys.platform == "win32":
+    raise Exception("pybamm_install_odes is not supported on Windows.")
+
+SUNDIALS_VERSION = "6.5.0"
+
+# Build in parallel wherever possible
+os.environ["CMAKE_BUILD_PARALLEL_LEVEL"] = str(cpu_count())
+
 try:
     # wget module is required to download SUNDIALS or SuiteSparse.
     import wget
@@ -37,19 +45,13 @@ def download_extract_library(url, directory):
 def install_sundials(download_dir, install_dir):
     # Download the SUNDIALS library and compile it.
     logger = logging.getLogger("scikits.odes setup")
-    sundials_version = "6.5.0"
 
     try:
         subprocess.run(["cmake", "--version"])
     except OSError:
         raise RuntimeError("CMake must be installed to build SUNDIALS.")
 
-    url = (
-        "https://github.com/LLNL/"
-        + "sundials/releases/download/v{}/sundials-{}.tar.gz".format(
-            sundials_version, sundials_version
-        )
-    )
+    url = f"https://github.com/LLNL/sundials/releases/download/v{SUNDIALS_VERSION}/sundials-{SUNDIALS_VERSION}.tar.gz"
     logger.info("Downloading sundials")
     download_extract_library(url, download_dir)
 
@@ -58,7 +60,7 @@ def install_sundials(download_dir, install_dir):
         "-DSUNDIALS_INDEX_SIZE=32",
         "-DBUILD_ARKODE:BOOL=OFF",
         "-DEXAMPLES_ENABLE:BOOL=OFF",
-        "-DCMAKE_INSTALL_PREFIX=" + install_dir,
+        f"-DCMAKE_INSTALL_PREFIX={install_dir}",
     ]
 
     # SUNDIALS are built within directory 'build_sundials' in the PyBaMM root
@@ -70,7 +72,7 @@ def install_sundials(download_dir, install_dir):
 
     print("-" * 10, "Running CMake prepare", "-" * 40)
     subprocess.run(
-        ["cmake", f"../sundials-{sundials_version}", *cmake_args],
+        ["cmake", f"../sundials-{SUNDIALS_VERSION}", *cmake_args],
         cwd=build_directory,
         check=True,
     )
@@ -81,23 +83,40 @@ def install_sundials(download_dir, install_dir):
 
 
 def update_LD_LIBRARY_PATH(install_dir):
-    # Look for current python virtual env and add export statement
-    # for LD_LIBRARY_PATH in activate script.  If no virtual env found,
-    # then the current user's .bashrc file is modified instead.
+    # Look for the current python virtual env and add an export statement
+    # for LD_LIBRARY_PATH in the activate script. If no virtual env is found,
+    # the current user's .bashrc file is modified instead.
 
     export_statement = f"export LD_LIBRARY_PATH={install_dir}/lib:$LD_LIBRARY_PATH"
 
+    home_dir = os.environ.get("HOME")
+    bashrc_path = os.path.join(home_dir, ".bashrc")
+    zshrc_path = os.path.join(home_dir, ".zshrc")
     venv_path = os.environ.get("VIRTUAL_ENV")
+
     if venv_path:
         script_path = os.path.join(venv_path, "bin/activate")
     else:
-        script_path = os.path.join(os.environ.get("HOME"), ".bashrc")
+        if os.path.exists(bashrc_path):
+            script_path = os.path.join(os.environ.get("HOME"), ".bashrc")
+        elif os.path.exists(zshrc_path):
+            script_path = os.path.join(os.environ.get("HOME"), ".zshrc")
+        elif os.path.exists(bashrc_path) and os.path.exists(zshrc_path):
+            print(
+                "Both .bashrc and .zshrc found in the home directory. Setting .bashrc as path"
+            )
+            script_path = os.path.join(os.environ.get("HOME"), ".bashrc")
+        else:
+            print("Neither .bashrc nor .zshrc found in the home directory.")
 
     if os.getenv("LD_LIBRARY_PATH") and f"{install_dir}/lib" in os.getenv(
         "LD_LIBRARY_PATH"
     ):
         print(f"{install_dir}/lib was found in LD_LIBRARY_PATH.")
-        print("--> Not updating venv activate or .bashrc scripts")
+        if os.path.exists(bashrc_path):
+            print("--> Not updating venv activate or .bashrc scripts")
+        if os.path.exists(zshrc_path):
+            print("--> Not updating venv activate or .zshrc scripts")
     else:
         with open(script_path, "a+") as fh:
             # Just check that export statement is not already there.
@@ -140,7 +159,7 @@ def main(arguments=None):
         else os.path.join(pybamm_dir, args.install_dir)
     )
 
-    # Check is sundials is already installed
+    # Check if sundials is already installed
     SUNDIALS_LIB_DIRS = [join(os.getenv("HOME"), ".local"), "/usr/local", "/usr"]
 
     if args.sundials_libs:

diff --git a/pybamm/solvers/jax_bdf_solver.py b/pybamm/solvers/jax_bdf_solver.py
@@ -1008,7 +1008,7 @@ def jax_bdf_integrate(func, y0, t_eval, *args, rtol=1e-6, atol=1e-6, mass=None):
     """
     if not pybamm.have_jax():
         raise ModuleNotFoundError(
-            "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/GNU-linux.html#optional-jaxsolver"
+            "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
         )
 
     def _check_arg(arg):

diff --git a/pybamm/solvers/jax_solver.py b/pybamm/solvers/jax_solver.py
@@ -61,7 +61,7 @@ def __init__(
     ):
         if not pybamm.have_jax():
             raise ModuleNotFoundError(
-                "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/GNU-linux.html#optional-jaxsolver"
+                "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
             )
 
         # note: bdf solver itself calculates consistent initial conditions so can set