added atac_layer argument to train_model tasks and made tests for it

src/pyrovelocity/tasks/train.py

@@ -279,6 +279,7 @@ def check_shared_time(posterior_samples, adata):
 @beartype
 def train_model(
     adata: str | AnnData,
+    atac_layer: Optional[str] = None,
     guide_type: str = "auto",
     model_type: str = "auto",
     batch_size: int = -1,
@@ -305,6 +306,7 @@ def train_model(
 
     Args:
         adata (str | AnnData): Path to a file that can be read to an AnnData object or an AnnData object.
+        atac_layer (Optional[str], optional): Name of AnnData layer that contains atac data, if present.
         guide_type (str, optional): The type of guide function for the Pyro model. Default is "auto".
         model_type (str, optional): The type of Pyro model. Default is "auto".
         batch_size (int, optional): Batch size for training. Default is -1, which indicates using the full dataset.
@@ -347,6 +349,11 @@ def train_model(
         >>> copy_raw_counts(adata)
         >>> _, model, posterior_samples = train_model(adata, use_gpu="auto", seed=99, max_epochs=200, loss_plot_path=loss_plot_path)
     """
+
+    if atac_layer:
+        logger.info(
+            "Multiome model not yet implemented. Proceeding without atac data."
+        )
     if isinstance(adata, str):
         adata = load_anndata_from_path(adata)
 

src/pyrovelocity/tests/tasks/__init__.py

@@ -0,0 +1 @@
+"""Unit test package for pyrovelocity."""

src/pyrovelocity/tests/tasks/test_train_model.py

@@ -0,0 +1,20 @@
+"""Tests for `pyrovelocity._train_model` task."""
+
+from pyrovelocity.tasks.train import train_model
+from pyrovelocity.utils import generate_sample_data
+from pyrovelocity.tasks.preprocess import copy_raw_counts
+
+
+def test_train_model(tmp_path):
+    loss_plot_path = str(tmp_path) + "/loss_plot_docs.png"
+    print(loss_plot_path)
+    adata = generate_sample_data(random_seed=99)
+    copy_raw_counts(adata)
+    _, model, posterior_samples = train_model(
+        adata,
+        atac_layer="atac",
+        use_gpu="auto",
+        seed=99,
+        max_epochs=200,
+        loss_plot_path=loss_plot_path,
+    )