install options

Author: PierreBoyeau

README.md

@@ -12,6 +12,11 @@ Refer to the preprint and the [project repository](https://github.com/YosefLab/c
 pip install csde
 ```
 
+By default, this installs JAX with CPU support. To enable GPU support (CUDA), install with the appropriate extra (e.g., for CUDA 12):
+```bash
+pip install "csde[cuda12]"
+```
+
 ## Data Requirements
 
 `csde` requires two `AnnData` objects containing gene expression counts. Typically, these are obtained by splitting your full dataset into two groups:

pyproject.toml

@@ -33,6 +33,8 @@ dev = [
     "isort",
     "flake8",
 ]
+cuda12 = ["jax[cuda12]"]
+cuda13 = ["jax[cuda13]"]
 
 [tool.setuptools.packages.find]
 where = ["src"]

tests/test_csde.py

@@ -7,10 +7,8 @@
 
 class TestCSDE(unittest.TestCase):
     def setUp(self):
-        # Create synthetic data
         n_genes = 10
 
-        # adata_pred
         n_pred = 100
         X_pred = np.random.poisson(lam=2.0, size=(n_pred, n_genes)).astype(float)
         obs_pred = pd.DataFrame(
@@ -19,7 +17,6 @@ def setUp(self):
         self.adata_pred = anndata.AnnData(X=X_pred, obs=obs_pred)
         self.adata_pred.var_names = [f"Gene_{i}" for i in range(n_genes)]
 
-        # adata_gt
         n_gt = 50
         X_gt = np.random.poisson(lam=2.0, size=(n_gt, n_genes)).astype(float)
         obs_gt = pd.DataFrame(
@@ -31,8 +28,6 @@ def setUp(self):
         self.adata_gt = anndata.AnnData(X=X_gt, obs=obs_gt)
         self.adata_gt.var_names = [f"Gene_{i}" for i in range(n_genes)]
 
-        # Make sure cell types exist
-        # Force at least one of each to avoid mapping errors in small random samples
         self.adata_pred.obs.iloc[0, 0] = "TypeA"
         self.adata_pred.obs.iloc[1, 0] = "TypeB"
         self.adata_gt.obs.iloc[0, 0] = "TypeA"