Installation with various extras_requires for different CUDA support (

#97)

* Add pip install user

* Use extras_require for setuptools

* Check and print CUDA usage after importing

* Use BaseException

* Update README

.gitignore

@@ -4,3 +4,4 @@ __pycache__
 *.eggs
 *.svg
 docs/build
+build

README.md

@@ -1,4 +1,5 @@
 # Appletree
+
 A high-Performance Program simuLatEs and fiTs REsponse of xEnon.
 
 [![DOI](https://zenodo.org/badge/534803881.svg)](https://zenodo.org/badge/latestdoi/534803881)
@@ -11,37 +12,52 @@ A high-Performance Program simuLatEs and fiTs REsponse of xEnon.
 ## Installation and Set-Up
 
 ### Regular installation:
+
+With cpu support:
+
+```
+pip install appletree[cpu]
+```
+
+With CUDA Toolkit 11.2 support:
+
+```
+pip install appletree[cuda112] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
-pip install appletree -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+With CUDA Toolkit 12.1 support:
+
+```
+pip install appletree[cuda121] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 ### Developer setup:
+
 Clone the repository:
 
 ```
 git clone https://github.com/XENONnT/appletree
 cd appletree
 ```
 
-Install the package and requirements in your environment:
+To install the package and requirements in your environment, replace `pip install appletree[*]` to `python3 -m pip install .[*] --user` in the above `pip` commands.
 
-```
-pip install -r requirements.txt --user
-python3 -m pip install ./ --user
-```
+To install appletree in editable mode, insert `--editable` argument after `install` in the above `pip install` or `python3 -m pip install` commands.
 
-If you wanna install appletree in editable mode, replace the last line with
+For example, to install in your environment and in editable mode with CUDA Toolkit 12.1 support:
 
 ```
-python3 -m pip install --editable ./ --user
+python3 -m pip install --editable .[cuda121] --user -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
-You are now good to go!
+Then you are now good to go!
 
 ## Usage
+
 The best way to start with the `appletree` package is to have a look at the tutorial `notebooks`. 
 
 ## Contributing
+
 Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.
 
 Please make sure to update tests as appropriate.

appletree/__init__.py

@@ -41,6 +41,21 @@
 from . import context
 from .context import *
 
+# check CUDA support setup
+from warnings import warn
+platform = utils.get_platform()
+if platform == 'cpu':
+    warning = 'You are running appletree on CPU, which usually results in low performance.'
+    warn(warning)
+try:
+    import jax
+    # try allocate something
+    jax.numpy.ones(1)
+except BaseException:
+    if platform == 'gpu':
+        print('Can not allocate memory on GPU, please check your CUDA version.')
+    raise ImportError(f'Appletree is not correctly setup to be used on {platform.upper()}.')
+
 try:
     import aptext
     HAVE_APTEXT = True

docs/source/installation.rst

@@ -11,7 +11,24 @@ Appletree can be found on `pypi <https://pypi.org/project/appletree/>`_ now. To
 .. code-block:: console
 
     pip install --upgrade pip
-    pip install appletree -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+With cpu support:
+
+.. code-block:: console
+
+    pip install appletree[cpu]
+
+With CUDA Toolkit 11.2 support:
+
+.. code-block:: console
+
+    pip install appletree[cuda112] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+With CUDA Toolkit 12.1 support:
+
+.. code-block:: console
+
+    pip install appletree[cuda121] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 Option 2: install from source code
 -----------------------------------------------------
@@ -20,11 +37,13 @@ Option 2: install from source code
 
     git clone https://github.com/XENONnT/appletree
     cd appletree
-    pip install -r requirements.txt --user
-    python3 -m pip install ./ --user
 
-If you wanna install appletree in editable mode, replace the last line with
+To install the package and requirements in your environment, replace `pip install appletree[*]` to `python3 -m pip install .[*] --user` in the above `pip` commands.
+
+To install appletree in editable mode, insert `--editable` argument after `install` in the above `pip install` or `python3 -m pip install` commands.
+
+For example, to install in your environment and in editable mode with CUDA Toolkit 12.1 support:
 
 .. code-block:: console
 
-    python3 -m pip install --editable ./ --user
+    python3 -m pip install --editable .[cuda121] --user -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

requirements.txt

@@ -4,8 +4,7 @@ GOFevaluation
 graphviz
 h5py
 immutabledict
---find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-jax[cuda12_pip]==0.4.10
+jax
 matplotlib
 multihist
 numpy

setup.py

@@ -28,10 +28,18 @@ def open_requirements(path):
     install_requires=requires,
     python_requires='>=3.8',
     extras_require={
-        'doc': [],
-        'test': [
-            'pytest',
-            'flake8',
+        'cpu': [
+            # pip install appletree[cpu]
+            'jax[cpu]',
+        ],
+        'cuda112': [
+            # pip install appletree[cuda112] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+            'jax[cuda]==0.3.15',
+        ],
+        'cuda121': [
+            # pip install appletree[cuda121] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+            # according to this commit, jax[cuda12_pip]==0.4.10 is valid for CUDA Toolkit 12.1
+            'jax[cuda12_pip]',
         ],
     },
     packages=setuptools.find_packages(),