feat(models): add k-NN model

.gitignore

@@ -1,5 +1,6 @@
 #
 /archive/
+example_notebooks/*
 
 #
 .DS_Store

.vscode/settings.json

@@ -56,6 +56,7 @@
   "search.followSymlinks": false,
   "terminal.integrated.fontSize": 14,
   "terminal.integrated.scrollback": 100000,
+  "python.terminal.activateEnvironment": false,
   "workbench.colorTheme": "Catppuccin Mocha",
   "workbench.iconTheme": "vscode-icons",
   // Passing --no-cov to pytestArgs is required to respect breakpoints

src/pyrovelocity/models/__init__.py

@@ -7,13 +7,14 @@
 from pyrovelocity.models._deterministic_simulation import (
     solve_transcription_splicing_model_analytical,
 )
-from pyrovelocity.models._transcription_dynamics import mrna_dynamics
+from pyrovelocity.models._transcription_dynamics import mrna_dynamics, atac_mrna_dynamics
 from pyrovelocity.models._velocity import PyroVelocity
 
 
 __all__ = [
     deterministic_transcription_splicing_probabilistic_model,
     mrna_dynamics,
+    atac_mrna_dynamics,
     PyroVelocity,
     solve_transcription_splicing_model,
     solve_transcription_splicing_model_analytical,

src/pyrovelocity/models/_trainer.py

@@ -271,8 +271,8 @@ def train_faster(
             if scipy.sparse.issparse(self.adata.layers["raw_spliced"])
             else self.adata.layers["raw_spliced"],
             dtype=torch.float32,
-        ).to(device)
-
+        ).to(device)            
+        
         epsilon = 1e-6
 
         log_u_library_size = np.log(
@@ -335,60 +335,127 @@ def train_faster(
 
         losses = []
         patience = patient_init
-        for step in range(max_epochs):
-            if cell_state is None:
-                elbos = (
-                    svi.step(
-                        u,
-                        s,
-                        u_library.reshape(-1, 1),
-                        s_library.reshape(-1, 1),
-                        u_library_mean.reshape(-1, 1),
-                        s_library_mean.reshape(-1, 1),
-                        u_library_scale.reshape(-1, 1),
-                        s_library_scale.reshape(-1, 1),
-                        None,
-                        None,
+
+        if not self.adata.uns['atac']:
+
+            for step in range(max_epochs):
+                if cell_state is None:
+                    elbos = (
+                        svi.step(
+                            u,
+                            s,
+                            u_library.reshape(-1, 1),
+                            s_library.reshape(-1, 1),
+                            u_library_mean.reshape(-1, 1),
+                            s_library_mean.reshape(-1, 1),
+                            u_library_scale.reshape(-1, 1),
+                            s_library_scale.reshape(-1, 1),
+                            None,
+                            None,
+                        )
+                        / normalizer
                     )
-                    / normalizer
-                )
-            else:
-                elbos = (
-                    svi.step(
-                        u,
-                        s,
-                        u_library.reshape(-1, 1),
-                        s_library.reshape(-1, 1),
-                        u_library_mean.reshape(-1, 1),
-                        s_library_mean.reshape(-1, 1),
-                        u_library_scale.reshape(-1, 1),
-                        s_library_scale.reshape(-1, 1),
-                        None,
-                        cell_state.reshape(-1, 1),
+                else:
+                    elbos = (
+                        svi.step(
+                            u,
+                            s,
+                            u_library.reshape(-1, 1),
+                            s_library.reshape(-1, 1),
+                            u_library_mean.reshape(-1, 1),
+                            s_library_mean.reshape(-1, 1),
+                            u_library_scale.reshape(-1, 1),
+                            s_library_scale.reshape(-1, 1),
+                            None,
+                            cell_state.reshape(-1, 1),
+                        )
+                        / normalizer
+                    )
+                if (step == 0) or (
+                    ((step + 1) % log_every == 0) and ((step + 1) < max_epochs)
+                ):
+                    mlflow.log_metric("-ELBO", -elbos, step=step + 1)
+                    logger.info(
+                        f"step {step + 1: >4d} loss = {elbos:0.6g} patience = {patience}"
+                    )
+                if step > log_every:
+                    if (losses[-1] - elbos) < losses[-1] * patient_improve:
+                        patience -= 1
+                    else:
+                        patience = patient_init
+                if patience <= 0:
+                    break
+                losses.append(elbos)
+
+        else:
+
+            c = torch.tensor(
+            np.array(
+                self.adata.layers["atac"].toarray(), dtype="float32"
+            )
+            if scipy.sparse.issparse(self.adata.layers["atac"])
+            else self.adata.layers["atac"],
+            dtype=torch.float32,
+        ).to(device) 
+
+
+            for step in range(max_epochs):
+                if cell_state is None:
+                    elbos = (
+                        svi.step(
+                            c,
+                            u,
+                            s,
+                            u_library.reshape(-1, 1),
+                            s_library.reshape(-1, 1),
+                            u_library_mean.reshape(-1, 1),
+                            s_library_mean.reshape(-1, 1),
+                            u_library_scale.reshape(-1, 1),
+                            s_library_scale.reshape(-1, 1),
+                            None,
+                            None,
+                        )
+                        / normalizer
                     )
-                    / normalizer
-                )
-            if (step == 0) or (
-                ((step + 1) % log_every == 0) and ((step + 1) < max_epochs)
-            ):
-                mlflow.log_metric("-ELBO", -elbos, step=step + 1)
-                logger.info(
-                    f"step {step + 1: >4d} loss = {elbos:0.6g} patience = {patience}"
-                )
-            if step > log_every:
-                if (losses[-1] - elbos) < losses[-1] * patient_improve:
-                    patience -= 1
                 else:
-                    patience = patient_init
-            if patience <= 0:
-                break
-            losses.append(elbos)
+                    elbos = (
+                        svi.step(
+                            c,
+                            u,
+                            s,
+                            u_library.reshape(-1, 1),
+                            s_library.reshape(-1, 1),
+                            u_library_mean.reshape(-1, 1),
+                            s_library_mean.reshape(-1, 1),
+                            u_library_scale.reshape(-1, 1),
+                            s_library_scale.reshape(-1, 1),
+                            None,
+                            cell_state.reshape(-1, 1),
+                        )
+                        / normalizer
+                    )
+                if (step == 0) or (
+                    ((step + 1) % log_every == 0) and ((step + 1) < max_epochs)
+                ):
+                    mlflow.log_metric("-ELBO", -elbos, step=step + 1)
+                    logger.info(
+                        f"step {step + 1: >4d} loss = {elbos:0.6g} patience = {patience}"
+                    )
+                if step > log_every:
+                    if (losses[-1] - elbos) < losses[-1] * patient_improve:
+                        patience -= 1
+                    else:
+                        patience = patient_init
+                if patience <= 0:
+                    break
+                losses.append(elbos)
+
         mlflow.log_metric("-ELBO", -elbos, step=step + 1)
         mlflow.log_metric("real_epochs", step + 1)
         logger.info(
             f"step {step + 1: >4d} loss = {elbos:0.6g} patience = {patience}"
         )
-        return losses
+        return losses        
 
     def train_faster_with_batch(
         self,