Merge pull request #45 from camlab-bioml/cklamann/bugfixes

Issues 42, 43, 44

Author: kieranrcampbell

Dockerfile

@@ -35,6 +35,10 @@ RUN python3 -m venv $VIRTUAL_ENV && \
 
 USER $USERNAME
 
-COPY --chown=${USER_UID}:${USER_GID} . .
+# prevent full rebuilds every time code changes
+COPY --chown=${USER_UID}:${USER_GID} pyproject.toml poetry.lock README.md /code/
+COPY --chown=${USER_UID}:${USER_GID} starling/__init__.py /code/starling/__init__.py
 
 RUN poetry install --with docs,dev
+
+COPY . .

starling/starling.py

@@ -160,8 +160,8 @@ def prepare_data(self) -> None:
             self.adata.uns["init_cell_size_variances"] = np.array(init_sv)
         else:
             # init_cell_size_centroids = None; init_cell_size_variances = None
-            self.adata.varm["init_cell_size_centroids"] = None
-            self.adata.varm["init_cell_size_variances"] = None
+            self.adata.uns["init_cell_size_centroids"] = None
+            self.adata.uns["init_cell_size_variances"] = None
             self.train_df = utility.ConcatDataset([self.X, tr_fy, tr_fl])
 
         # model_params = utility.model_paramters(self.init_e, self.init_v, self.init_s, self.init_sv)

starling/utility.py

@@ -38,7 +38,7 @@ def __len__(self):
 def init_clustering(
     initial_clustering_method: Literal["User", "KM", "GMM", "FS", "PG"],
     adata: AnnData,
-    k: Union[int, None],
+    k: Union[int, None] = None,
     labels: Optional[np.ndarray] = None,
 ) -> AnnData:
     """Compute initial cluster centroids, variances & labels
@@ -49,7 +49,8 @@ def init_clustering(
         ``FS`` (FlowSOM), ``User`` (user-provided), or ``PG`` (PhenoGraph).
     :param k: The number of clusters, must be ``n_components`` when ``initial_clustering_method`` is ``GMM`` (required),
         ``k`` when ``initial_clustering_method`` is ``KM`` (required), ``k`` when ``initial_clustering_method``
-        is ``FS`` (required), ``?`` when  ``initial_clustering_method`` is ``PG`` (optional)
+        is ``FS`` (required), ``?`` when  ``initial_clustering_method`` is ``PG`` (optional), and can be ommited when
+        ``initial_clustering_method`` is "User", because user will be passing in their own labels.
     :param labels: optional, user-provided labels
 
     :raises: ValueError
@@ -67,6 +68,11 @@ def init_clustering(
             "k cannot be ommitted for KMeans, FlowSOM, or Gaussian Mixture"
         )
 
+    if initial_clustering_method == "User" and labels is None:
+        raise ValueError(
+            "labels must be provided when initial_clustering_method is set to 'User'"
+        )
+
     if initial_clustering_method == "KM":
         kms = KMeans(k).fit(adata.X)
         init_l = kms.labels_
@@ -90,12 +96,13 @@ def init_clustering(
         else:
             init_l = labels
 
-        k = len(np.unique(init_l))
+        classes = np.unique(init_l)
+        k = len(classes)
         init_e = np.zeros((k, adata.X.shape[1]))
         init_ev = np.zeros((k, adata.X.shape[1]))
-        for c in range(k):
-            init_e[c, :] = adata.X[init_l == c].mean(0)
-            init_ev[c, :] = adata.X[init_l == c].var(0)
+        for i, c in enumerate(classes):
+            init_e[i, :] = adata.X[init_l == c].mean(0)
+            init_ev[i, :] = adata.X[init_l == c].var(0)
 
     elif initial_clustering_method == "FS":
         ## needs to output to csv first

tests/test_sanity.py

@@ -1,16 +1,16 @@
 from os.path import dirname, join
 
 import anndata as ad
+import numpy as np
 import pandas as pd
-import pytorch_lightning as pl
 from lightning_lite import seed_everything
 from pytorch_lightning.callbacks import EarlyStopping
 
 from starling import starling, utility
 
 
 def test_can_run_km(tmpdir):
-    """Temporary sanity check"""
+    """Test that we can run with the KM setting in init_clustering"""
     seed_everything(10, workers=True)
 
     raw_adata = ad.read_h5ad(join(dirname(__file__), "fixtures", "sample_input.h5ad"))
@@ -51,7 +51,7 @@ def test_can_run_km(tmpdir):
 
 
 def test_can_run_gmm(tmpdir):
-    """Temporary sanity check"""
+    """Test that we can run with the GMM setting in init_clustering"""
     seed_everything(10, workers=True)
     adata = utility.init_clustering(
         "GMM",
@@ -89,7 +89,7 @@ def test_can_run_gmm(tmpdir):
 
 
 def test_can_run_pg(tmpdir):
-    """Temporary sanity check"""
+    """Test that we can run with the PG setting in init_clustering"""
     seed_everything(10, workers=True)
     adata = utility.init_clustering(
         "PG",
@@ -111,7 +111,6 @@ def test_can_run_pg(tmpdir):
     ## initial expression centriods (p x c) matrix
     init_cent = pd.DataFrame(result.varm["init_exp_centroids"], index=result.var_names)
 
-    # j seems to vary here
     assert init_cent.shape[0] == 24
 
     ## starling expression centriods (p x c) matrix
@@ -125,3 +124,28 @@ def test_can_run_pg(tmpdir):
     )
 
     assert prom_mat.shape[0] == 13685
+
+
+def test_can_run_pg_without_cell_size(tmpdir):
+    """Test that we can run the model with model_cell_size=False in ST"""
+    seed_everything(10, workers=True)
+    adata = utility.init_clustering(
+        "PG",
+        ad.read_h5ad(join(dirname(__file__), "fixtures", "sample_input.h5ad")),
+        k=10,
+    )
+    st = starling.ST(adata, model_cell_size=False)
+    cb_early_stopping = EarlyStopping(monitor="train_loss", mode="min", verbose=False)
+
+    ## train ST
+    st.train_and_fit(
+        max_epochs=2,
+        callbacks=[cb_early_stopping],
+        default_root_dir=tmpdir,
+    )
+
+    result = st.result()
+
+    exp_cent = pd.DataFrame(result.varm["st_exp_centroids"], index=result.var_names)
+
+    assert exp_cent.shape[0] == 24

tests/test_utility.py

@@ -1,3 +1,4 @@
+import numpy as np
 from anndata import AnnData
 
 from starling.utility import init_clustering, validate_starling_arguments
@@ -6,9 +7,11 @@
 def assert_annotated(adata: AnnData, k):
     assert "init_exp_centroids" in adata.varm
     assert adata.varm["init_exp_centroids"].shape == (adata.X.shape[1], k)
+    assert not np.any(np.isnan(adata.varm["init_exp_centroids"]))
 
     assert "init_exp_centroids" in adata.varm
     assert adata.varm["init_exp_variances"].shape == (adata.X.shape[1], k)
+    assert not np.any(np.isnan(adata.varm["init_exp_variances"]))
 
     assert "init_label" in adata.obs
     assert adata.obs["init_label"].shape == (adata.X.shape[0],)
@@ -32,6 +35,25 @@ def test_init_clustering_pg(simple_adata):
     assert_annotated(initialized, k)
 
 
+def test_init_clustering_user(simple_adata):
+    k = 3
+    initialized = init_clustering(
+        "User", simple_adata, labels=np.random.randint(k, size=32)
+    )
+    assert_annotated(initialized, k)
+
+
+def test_init_clustering_user_string(simple_adata):
+    k = 3
+    initialized = init_clustering(
+        "User",
+        simple_adata,
+        labels=np.random.choice(np.array(["a", "b", "c"]), size=32),
+    )
+
+    assert_annotated(initialized, k)
+
+
 def test_validation_passes_with_no_size(simple_adata):
     validate_starling_arguments(
         adata=simple_adata,