feat: update python (#84)

* feat: make jaxmd a project dependency

* feat: upgrade python version to 3.11

* feat: add optional gpu dependencies

* feat: test tox without dependencies

* feat: attempt fixing ci tests

* docs: update README.md

* ci: try uv venv 3.13

* ci: switch back to tox

* docs: fix badge

* docs: update README.md

* chore: remove tox pr trigger

* docs: update installation

* docs: remove editable install

Author: lwalew

.github/workflows/tests_and_linters.yaml

@@ -53,7 +53,7 @@ jobs:
 
       - name: Run tests 🧪
         run: |
-          uv sync --group dev --group jax_md
+          uv sync --only-dev
           uv run pytest --verbose --cov-report xml:coverage.xml \
                             --cov-report term-missing \
                             --junitxml=pytest.xml \

.github/workflows/tox_tests.yaml

@@ -5,26 +5,6 @@ on:
     branches: [ main, develop ]
 
 jobs:
-  tox_py310:
-    runs-on: instadeep-ci-4
-    container:
-      image: python:3.10-slim-bullseye
-
-    steps:
-      - name: Git setup
-        run: |
-          apt-get update && apt-get install -y coreutils git
-
-      - name: Checkout code  📦
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: '0'
-
-      - name: Run tests with tox 🧪
-        run: |
-          pip install tox==4.30.3
-          tox -e py310
-
   tox_py311:
     runs-on: instadeep-ci-4
     container:

README.md

@@ -1,7 +1,7 @@
 #  🔬 MLIPAudit:  A library to validate and benchmark MLIP models
 
 [![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
-[![Python 3.11](https://img.shields.io/badge/python-3.10%20%7C%203.11-blue)](https://www.python.org/downloads/release/python-3110/)
+[![Python 3.11](https://img.shields.io/badge/python-3.11%20%7C%203.12%20%7C%203.13-blue))](https://www.python.org/downloads/release/python-3110/)
 [![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit)](https://github.com/pre-commit/pre-commit)
 ![badge](https://img.shields.io/endpoint?url=https://gist.githubusercontent.com/mlipbot/e7c79b17c0a9d47bc826100ef880a16f/raw/pytest-coverage-comment.json)
 [![Tests and Linters 🧪](https://github.com/instadeepai/mlipaudit-open/actions/workflows/tests_and_linters.yaml/badge.svg?branch=main)](https://github.com/instadeepai/mlipaudit-open/actions/workflows/tests_and_linters.yaml)
@@ -29,21 +29,13 @@ MLIPAudit can be installed via pip:
 pip install mlipaudit
 ```
 
-However, this command **only installs the regular CPU version** of JAX.
-We recommend that MLIPAudit is run on GPU. Also, some benchmarks will require
-[JAX-MD](https://github.com/jax-md/jax-md) as a dependency. As the newest
-version of JAX-MD is not available on PyPI yet, this dependency will not
-be shipped with MLIPAudit automatically and instead must be installed
-directly from the GitHub repository.
-
-Therefore, we recommend running
-
+However, this command **only installs the regular CPU version** of JAX. If benchmarking
+native JAX models, we recommend installing the core library along with the GPU
+dependencies (`jax[cuda12]` and `jaxlib`) with the following command:
 ```bash
-pip install -U "jax[cuda12]" git+https://github.com/jax-md/jax-md.git
+pip install mlipaudit[gpu]
 ```
 
-to install both of these additional packages.
-
 ## 📖 Documentation
 
 The detailed code documentation that also contains descriptions for each benchmark and

docs/source/installation/index.rst

@@ -13,23 +13,9 @@ After installation and activating the respective Python environment, the command
 tools `mlipaudit` and `mlipauditapp` should be available.
 
 However, the command above **only installs the regular CPU version** of JAX.
-We recommend that the library is run on GPU.
-This requires also installing the necessary versions
-of `jaxlib <https://pypi.org/project/jaxlib/>`_ which can also be installed via pip. See
-the `installation guide of JAX <https://docs.jax.dev/en/latest/installation.html>`_ for
-more information.
-At time of release, the following install command is supported:
+If benchmarking native JAX models, we recommend installing the core library
+along with the GPU dependencies (`jax[cuda12]` and `jaxlib`) with the following command:
 
 .. code-block:: bash
 
-    pip install -U "jax[cuda12]"
-
-Also, some benchmarks require `JAX-MD <https://github.com/jax-md/jax-md>`_ as a
-dependency. As the newest
-version of JAX-MD is not available on PyPI yet, this dependency will not
-be shipped with *mlipaudit* automatically and instead must be installed
-directly from the GitHub repository, like this:
-
-.. code-block:: bash
-
-    pip install git+https://github.com/jax-md/jax-md.git
+    pip install mlipaudit[gpu]

pyproject.toml

@@ -6,7 +6,7 @@ readme = "README.md"
 authors = [
     { name = "InstaDeep", email = "[email protected]" }
 ]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 dependencies = [
     "huggingface-hub>=0.33.4",
     "mdtraj>=1.10.3",
@@ -16,6 +16,7 @@ dependencies = [
     "vl-convert-python>=1.8.0",
     "mdtraj>=1.10.3",
     "tmtools>=0.2.0",
+    "jax-md>=0.2.26"
 ]
 
 [project.scripts]
@@ -25,6 +26,12 @@ mlipaudit = "mlipaudit.main:main"
 requires = ["hatchling"]
 build-backend = "hatchling.build"
 
+[project.optional-dependencies]
+gpu = [
+    "jax[cuda12]>=0.4.33",
+    "jaxlib>=0.4.33"
+]
+
 [dependency-groups]
 dev = [
     "notebook>=7.4.4",
@@ -39,19 +46,9 @@ dev = [
 scripts = [
     "requests>=2.32.4",
 ]
-gpu = [
-    "jax[cuda12]>=0.4.33",
-    "jaxlib>=0.4.33"
-]
-jax_md = [
-    "jax-md",
-]
 
 [tool.coverage.run]
 omit = [
     "main.py",
     "app.py"
 ]
-
-[tool.uv.sources]
-jax-md = { git = "https://github.com/jax-md/jax-md.git", tag = "jax-md-v0.2.25" }

tox.ini

@@ -1,10 +1,10 @@
 [tox]
 isolated_build = True
-env_list = py{310,311,312,313}
+env_list = py{311,312,313}
 
 [testenv]
 deps =
     pytest>=8.4.1
     pytest-mock>=3.14.1
-    git+https://github.com/jax-md/jax-md.git
+description = Test various python versions
 commands = pytest {posargs}