Add pre-commit configuration, fix dependencies, and add CI workflow

Author: PierreBoyeau

.github/workflows/test.yml

@@ -0,0 +1,40 @@
+name: Tests
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11", "3.12"]
+
+    steps:
+    - uses: actions/checkout@v4
+    
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+        cache: 'pip'
+        
+    - name: Install dependencies
+      run: |
+        python -m pip install -U pip
+        pip install -e ".[dev]"
+        
+    - name: Lint with flake8
+      run: |
+        # stop the build if there are Python syntax errors or undefined names
+        flake8 src tests --count --select=E9,F63,F7,F82 --show-source --statistics
+        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+        flake8 src tests --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+        
+    - name: Test with pytest
+      run: |
+        pytest tests/

.pre-commit-config.yaml

@@ -0,0 +1,15 @@
+repos:
+  - repo: https://github.com/psf/black
+    rev: 25.1.0
+    hooks:
+      - id: black
+  - repo: https://github.com/pycqa/isort
+    rev: 6.0.1
+    hooks:
+      - id: isort
+        args: ["--profile", "black"]
+  - repo: https://github.com/pycqa/flake8
+    rev: 7.1.1
+    hooks:
+      - id: flake8
+        args: ["--max-line-length=127", "--extend-ignore=E203"]

README.md

@@ -1,5 +1,7 @@
 # CSDE: Corrected Spatial Differential Expression
 
+[![Tests](https://github.com/YosefLab/csde/actions/workflows/test.yml/badge.svg)](https://github.com/YosefLab/csde/actions/workflows/test.yml)
+
 `csde` (Corrected Spatial Differential Expression) is a Python package designed to **identify differentially expressed (DE) genes between spatially-resolved cell populations** (e.g., T-cells inside vs. outside a tumor).
 
 Standard analysis relies on cell population assignments (e.g., "infiltrating" vs. "non-infiltrating") obtained automatically from clustering/ML that are often prone to errors. `csde` corrects for these inaccuracies by leveraging a small subset of validated "ground-truth" data, providing rigorous statistical guarantees for spatially-resolved DE analyses.

pyproject.toml

@@ -8,14 +8,14 @@ version = "0.1.0"
 description = "Corrected Spatial Differential Expression analysis package"
 readme = "README.md"
 requires-python = ">=3.10"
-license = {text = "MIT"}
+license = {text = "BSD-3-Clause"}
 authors = [
   {name = "Pierre Boyeau"},
 ]
 dependencies = [
     "anndata>=0.8.0",
     "scanpy>=1.9.0",
-    "jax>=0.4.0",
+    "jax>=0.4.0,<0.10.0",
     "flax>=0.7.0",
     "optax>=0.1.0",
     "numpyro>=0.12.0",

src/csde/api.py

@@ -1,4 +1,4 @@
-from typing import Optional, Tuple, List
+from typing import Optional
 
 import anndata
 import numpy as np

src/csde/model.py

@@ -1,10 +1,7 @@
-import os
-from typing import Union, Optional, Tuple, List, Dict, Any
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import flax.linen as nn
 import jax
-
-jax.config.update("jax_enable_x64", True)
 import jax.numpy as jnp
 import numpy as np
 import pandas as pd
@@ -14,6 +11,8 @@
 
 from csde.optimization import _zstat_generic2, optimize_ppi, optimize_ppi_gd
 
+jax.config.update("jax_enable_x64", True)
+
 
 class PPIAbstractClass:
     """
@@ -413,14 +412,14 @@ def likelihood(model_params, x, y):
 
         all_grads = np.zeros((n_obs, self.n_params))
         for i in tqdm(range(0, n_obs, batch_size), desc="Gradient computation"):
-            x_batch = x[i : i + batch_size]
-            y_batch = y[i : i + batch_size]
+            x_batch = x[i:i+batch_size]
+            y_batch = y[i:i+batch_size]
             n_obs_batch = x_batch.shape[0]
             score = self.jit(jax.jacfwd(likelihood))
             grads = score(self.model_params, x_batch, y_batch)
             grad_mu = np.array(grads["params"]["mu"].reshape(n_obs_batch, -1))
             grad_mu0 = np.array(grads["params"]["mu0"].reshape(n_obs_batch, -1))
-            all_grads[i : i + batch_size] = np.hstack([grad_mu, grad_mu0])
+            all_grads[i:i+batch_size] = np.hstack([grad_mu, grad_mu0])
         return np.array(all_grads)
 
     def _construct_contrast(self, feature_id: int, idx_a: int) -> np.ndarray:

src/csde/optimization.py

@@ -1,5 +1,5 @@
 import time
-from typing import Optional, Callable, Dict, Any, Tuple
+from typing import Any, Optional, Tuple
 
 import jax
 import jax.numpy as jnp

tox.ini

@@ -0,0 +1,16 @@
+[tox]
+envlist = py311, py312
+isolated_build = True
+toxworkdir = /tmp/tox/{env:USER}/{toxinidir:name}
+
+[testenv]
+description = Run flake8 and pytest
+extras = 
+    dev
+commands =
+    # stop the build if there are Python syntax errors or undefined names
+    flake8 src tests --count --select=E9,F63,F7,F82 --show-source --statistics
+    # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+    flake8 src tests --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+    # Test with pytest
+    pytest {posargs:tests}