Revamp CI

.github/workflows/code-style.yaml

@@ -0,0 +1,35 @@
+name: CodeStyle
+
+on:
+  push:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  check-and-test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+
+    name: Code style
+    steps:
+      #----------------------------------------------
+      #       check-out repo and set-up python
+      #----------------------------------------------
+    - uses: actions/checkout@v4
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v6
+      with:
+        version: "0.9.6"
+      #----------------------------------------------
+      #       install
+      #----------------------------------------------
+    - name: Install dependencies
+      run: UV_PYTHON=python3.11 make sync
+      #----------------------------------------------
+      #       lint
+      #----------------------------------------------
+    - name: Run lint
+      run: UV_PYTHON=python3.11 make lint

.github/workflows/docs.yaml

@@ -10,19 +10,21 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       #----------------------------------------------
-      #       check-out repo and set-up python
+      #       check out repo and set up uv
       #----------------------------------------------
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: astral-sh/setup-uv@v6
       with:
-        python-version: "3.10"
+        version: "0.9.6"
       #----------------------------------------------
       #       install
       #----------------------------------------------
+    - name: Set up venv
+      run: UV_PYTHON=3.11 make venv
     - name: Install dependencies
       run: |
         make install
-        pip install -r docs/requirements.txt
+        uv pip install -r docs/requirements.txt
     - name: Build documentation
       run: |
         make docs

.github/workflows/quality-checks.yaml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11"]
       fail-fast: false
 
     name: Python-${{ matrix.python-version }}
@@ -20,18 +20,18 @@ jobs:
       #       check-out repo and set-up python
       #----------------------------------------------
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v6
       with:
-        python-version: ${{ matrix.python-version }}
+        version: "0.9.6"
       #----------------------------------------------
       #       install
       #----------------------------------------------
     - name: Install dependencies
-      run: GK_REQUIREMENTS=test_requirements-${{ matrix.python-version }}.txt make install
+      run: UV_PYTHON=python${{ matrix.python-version }} make sync
       #----------------------------------------------
-      #       Lint and test
+      #       test
       #----------------------------------------------
-    - name: Run lint
-      run: make lint
     - name: Run tests
-      run: make test
+      run: UV_PYTHON=python${{ matrix.python-version }} make test

Makefile

@@ -1,38 +1,47 @@
 .PHONY: help docs install format lint test
 
-SUCCESS='\033[0;32m'
+SUCCESS = \033[0;32m
+RESET   = \033[0m
 
 SHELL=/bin/bash
-PYVERSION:=$(shell python -c "import sys;t='{v[0]}.{v[1]}'.format(v=list(sys.version_info[:2]));sys.stdout.write(t)")
-GK_REQUIREMENTS?=test_requirements-$(PYVERSION).txt
+UV        ?= uv
+UV_RUN    ?= uv run
+VENV_DIR  ?= .venv
+UV_PYTHON ?= 3.11
 
 help: ## Shows this help message
 	# $(MAKEFILE_LIST) is set by make itself; the following parses the `target:  ## help line` format and adds color highlighting
 	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-24s\033[0m %s\n", $$1, $$2}'
 
 docs:
 	(cd docs ; make clean; make doctest; make html)
-	@echo "${SUCCESS}============== Docs are available at docs/_build/html/index.html ============== ${SUCCESS}"
+	@echo -e "$(SUCCESS)============== Docs are available at docs/_build/html/index.html ==============$(RESET)"
 
+venv: ## Create a virtualenv (UV_PYTHON=3.10 to force)
+	@$(UV) venv --seed $(if $(strip $(UV_PYTHON)),--python $(UV_PYTHON),) $(VENV_DIR)
+	@echo -e "$(SUCCESS)Virtualenv ready in $(VENV_DIR)$(RESET)"
 
-install:  ## Install repo for developement (Only for Linux)
-	@echo "=== pip install package with dev requirements (using $(GK_REQUIREMENTS)) =============="
-	pip install --upgrade pip
-	pip install --upgrade --upgrade-strategy eager --no-cache-dir -r $(GK_REQUIREMENTS) | cat
-	pip install -e .
+sync: ## Resolve + install project and dev deps for development
+	@$(UV) sync --dev
+	@echo -e "$(SUCCESS)Environment synced from pyproject.toml$(RESET)"
 
-format:  ## Formats code with `autoflake`, `black` and `isort`
-	autoflake --remove-all-unused-imports --recursive --remove-unused-variables --in-place geometric_kernels tests --exclude=__init__.py
-	black geometric_kernels tests
-	isort geometric_kernels tests
+install: sync  ## Backward-compat
 
-lint:
-	flake8 geometric_kernels tests
-	black geometric_kernels tests --check --diff
-	isort geometric_kernels tests --check-only --diff
-	mypy --namespace-packages geometric_kernels
 
+format: sync ## Formats code with `autoflake`, `black` and `isort`
+	@$(UV_RUN) autoflake --remove-all-unused-imports --recursive --remove-unused-variables --in-place geometric_kernels tests --exclude=__init__.py
+	@$(UV_RUN) black geometric_kernels tests
+	@$(UV_RUN) isort geometric_kernels tests
+	@echo -e "$(SUCCESS)Format done$(RESET)"
 
-test:  ## Run the tests, start with the failing ones and break on first fail.
-	pytest -v -x --ff -rN -Wignore -s --tb=short --durations=0 --cov --cov-report=xml tests
-	pytest --nbmake --nbmake-kernel=python3 --durations=0 --nbmake-timeout=1000 --ignore=notebooks/frontends/GPJax.ipynb notebooks/
+lint: sync
+	@$(UV_RUN) flake8 geometric_kernels tests
+	@$(UV_RUN) black geometric_kernels tests --check --diff
+	@$(UV_RUN) isort geometric_kernels tests --check-only --diff
+	@$(UV_RUN) mypy --namespace-packages geometric_kernels
+	@echo -e "$(SUCCESS)Lint done$(RESET)"
+
+test: sync ## Run the tests, start with the failing ones and break on first fail.
+	@$(UV_RUN) pytest -v -x --ff -rN -Wignore -s --tb=short --durations=0 --cov --cov-report=xml tests
+	@$(UV_RUN) pytest --nbmake --nbmake-kernel=python3 --durations=0 --nbmake-timeout=1000 --ignore=notebooks/frontends/GPJax.ipynb notebooks/
+	@echo -e "$(SUCCESS)Tests done$(RESET)"

README.md

@@ -13,7 +13,12 @@ This enables kernel methods — in particular Gaussian process models &mdash
 
 0. [Optionally] create and activate a new virtual environment.
 
-    You can use Conda
+    You can use [uv](https://github.com/astral-sh/uv)
+    ```bash
+    uv venv
+    ```
+
+    or conda
 
     ```bash
     conda create -n [env_name] python=3.[version]
@@ -33,6 +38,8 @@ This enables kernel methods — in particular Gaussian process models &mdash
     pip install geometric_kernels
     ```
 
+    **NOTE**: If you use `uv`, swap `pip` with `uv pip` everywhere. If you initialized a project with `uv init`, use `uv add` instead.
+
     If you want to install specific GitHub branch called `[branch]`, run
 
     ```bash
@@ -142,15 +149,33 @@ The documentation for GeometricKernels is available on a [separate website](http
 
 ## For development and running the tests
 
-Run these commands from the root directory of the repository.
+If you want to contribute to the library, thank you! Follow these instructions to set up the environment and run the tests and the code formatting.
+
+Initialize the virtual environment.
+
+```bash
+make venv
+
+```
 
-Install all backends and the dev requirements (Pytest, black, etc.)
+You can change the python version and the venv directory like this:
+
+```bash
+make venv UV_PYTHON=3.11 VENV_DIR=.venv
+```
+
+**NOTE**: We use [uv](https://github.com/astral-sh/uv) for development. It is not strictly necessary, and if you don't want to use it, you can still run `make lint` and `make test` if you set `UV_RUN=` to be empty. You will need to set up the environment yourself.
+
+Install all backends and the dev requirements (Pytest, black, etc.). This will install all the backends.
 
 ```bash
 make install
 ```
 
-Run style checks
+**NOTE**: If not using `uv`, you can still install the dev requirements via `pip install -e .[dev]`.
+
+Run the style checks
+
 ```bash
 make lint
 ```
@@ -161,6 +186,8 @@ Run the tests
 make test
 ```
 
+If you want to run Jupyter with your `uv` development environment, check out [this page](https://docs.astral.sh/uv/guides/integration/jupyter/).
+
 **Example:** If you want to learn how to implement your own space or kernel component, checkout the [CustomSpacesAndKernels.ipynb](https://github.com/geometric-kernels/GeometricKernels/blob/main/notebooks/CustomSpacesAndKernels.ipynb) notebook.
 
 ## If you have a question

docs/Makefile

@@ -4,7 +4,7 @@
 # You can set these variables from the command line, and also
 # from the environment for the first two.
 SPHINXOPTS    ?=
-SPHINXBUILD   ?= sphinx-build
+SPHINXBUILD   ?= uv run sphinx-build
 SOURCEDIR     = .
 BUILDDIR      = _build
 

docs/README.md

@@ -1,12 +1,27 @@
-# Compiling documentation
+# Compiling Documentation
 
-First, execute doctests by running the following command in this dirrectory
+## Installing Depdenencies
+
+If you are using `uv`, first install docs building depedenencies **while in the `docs/` dirrectory**.
+```bash
+uv add -r requirements.txt
 ```
-make doctest
+If you are using `pip`, you can run
+```bash
+pip install -r requirements.txt
 ```
 
-If all tests are passed, run
+## Running Doctests
+
+Execute doctests by running the following command **while in the `docs/` dirrectory**.
+```bash
+make doctest
 ```
+
+## Compiling Documentation
+
+If all tests are passed, run the following command **while in the `docs/` dirrectory**.
+```bash
 make html
 ```
 

docs/index.rst

@@ -16,6 +16,31 @@ This is a **Python 3** library.
 
 Before doing anything, you might want to create and activate a new virtual environment:
 
+.. raw:: html
+
+   <div class="bootstrap">
+   <div class="accordion" id="virtualenvs">
+     <div class="accordion-item" style="background-color: var(--color-background-primary);">
+       <h2 class="accordion-header mb-0" id="virtualenvsHeadingOne">
+         <button class="accordion-button" type="button" data-bs-toggle="collapse" data-bs-target="#virtualenvsCollapseOne" aria-expanded="true" aria-controls="virtualenvsCollapseOne" style="background-color: var(--color-background-secondary); color: var(--color-foreground-primary);">
+           uv
+         </button>
+       </h2>
+       <div id="virtualenvsCollapseOne" class="accordion-collapse collapse show" aria-labelledby="virtualenvsHeadingOne" data-bs-parent="#virtualenvs">
+         <div class="accordion-body pb-0">
+
+..  code-block:: bash
+
+   uv venv --python python[version] [venv_dir]
+
+where [env_dir] is the directory of the environment and [version] is the version of Python you want to use, we currently support 3.9, 3.10, 3.11.
+
+.. raw:: html
+
+         </div>
+       </div>
+     </div>
+
 .. raw:: html
 
    <div class="bootstrap">
@@ -74,6 +99,9 @@ To install GeometricKernels, run
 
     pip install geometric_kernels
 
+.. note::
+   If you use `uv`, swap `pip` with `uv pip` everywhere. Additionally, if you ran `uv init`, use `uv add` instal, to add the requirement in your `pyproject.toml` and install it.
+
 .. note::
   If you want to install specific GitHub branch called `[branch]`, run
 

geometric_kernels/lab_extras/jax/extras.py

@@ -30,7 +30,7 @@ def trapz(y: B.JAXNumeric, x: _Numeric, dx: _Numeric = 1.0, axis: int = -1):  #
     """
     Integrate along the given axis using the trapezoidal rule.
     """
-    return jnp.trapz(y, x, dx, axis)
+    return jnp.trapezoid(y, x, dx, axis)
 
 
 @dispatch

pyproject.toml

@@ -21,9 +21,9 @@ classifiers = [
 keywords=[
     "geometric-kernels",
 ]
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 dependencies = [
-    "backends>=1.7",
+    "backends", # >=1.7",
     "einops",
     "geomstats",
     "numpy>=1.16",
@@ -59,3 +59,53 @@ allow_redefinition = true
 [tool.black]
 line-length = 88
 target-version = ['py38', 'py39', 'py310', 'py311']
+
+
+[tool.uv]
+default-groups = []
+
+[dependency-groups]
+dev = [
+  # --- shared from test_requirements.txt ---
+  "ipykernel",
+  "backends>=1.5.4",
+  "plotly",
+  "kaleido",
+  "black==24.3.0",
+  "flake8==7.0.0",
+  "isort==5.13.2",
+  "autoflake",
+  "pytest",
+  "pytest-cov",
+  "nbmake",
+  "mypy",
+  "scikit-learn",
+  "geoopt",
+  'scipy',
+
+  # typing-extensions only required on newer pins in 3.10/3.11 files
+  'typing_extensions>=4.6; python_version >= "3.10" and python_version < "3.12"',
+
+  # Torch
+  "torch",
+
+  # gpytorch
+  'gpytorch',
+  "botorch>=0.9",
+
+  # TensorFlow / GPflow / TFP split
+  'tensorflow==2.13.0; python_version=="3.8"',
+  'tensorflow==2.13.1; python_version=="3.9"',
+  'tensorflow==2.15;   python_version>="3.10" and python_version<"3.12"',
+  'tensorflow-probability==0.20.1; python_version < "3.10"',
+  'tensorflow-probability==0.23;   python_version >= "3.10" and python_version < "3.12"',
+  'gpflow==2.9.0; python_version < "3.10"',
+  'gpflow==2.9;    python_version >= "3.10" and python_version < "3.12"',
+
+  # JAX family
+  'jax',
+  'jaxlib',
+  'jaxtyping==0.2.25; python_version == "3.9"',
+  'jaxtyping;         python_version != "3.9"',
+  'optax',
+]