Change: formatted and reconfigured zarr imports for zarr 3 (#42)

* Change: formatted and reconfigured imports for zarr 3

* Change: reverse superfluous formatting

---------

Co-authored-by: Max Hargreaves <hargreaves.max@gene.com>

Author: DiDeoxy

setup.cfg

@@ -58,7 +58,7 @@ install_requires =
     tiledb-vector-search>=0.11.0
     torch>=1.10.1
     tqdm
-    zarr>=2.6.1,<3.0.0
+    zarr>=2.6.1
     importlib-metadata; python_version<"3.8"
 
 [options.packages.find]

setup.py

@@ -1,11 +1,12 @@
 """
-    Setup file for scimilarity.
-    Use setup.cfg to configure your project.
+Setup file for scimilarity.
+Use setup.cfg to configure your project.
 
-    This file was generated with PyScaffold 4.5.
-    PyScaffold helps you to put up the scaffold of your new Python project.
-    Learn more under: https://pyscaffold.org/
+This file was generated with PyScaffold 4.5.
+PyScaffold helps you to put up the scaffold of your new Python project.
+Learn more under: https://pyscaffold.org/
 """
+
 from setuptools import setup
 
 if __name__ == "__main__":

src/scimilarity/anndata_data_models.py

@@ -4,7 +4,7 @@
 import pytorch_lightning as pl
 import torch
 from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 from .utils import align_dataset
 from .ontologies import (
@@ -13,8 +13,21 @@
     find_most_viable_parent,
 )
 
+if TYPE_CHECKING:
+    import numpy
+    from numpy.typing import NDArray
+    from typing import Any
 
-class scDataset(Dataset):
+    Index = (
+        NDArray[numpy.integer[Any]]
+        | NDArray[numpy.bool_]
+        | tuple[NDArray[numpy.integer[Any]] | NDArray[numpy.bool_], ...]
+    )
+
+
+class scDataset(
+    Dataset[tuple["numpy.ndarray", "numpy.ndarray", Optional["numpy.ndarray"]]]
+):
     """A class that represents a single cell dataset.
 
     Parameters
@@ -27,17 +40,28 @@ class scDataset(Dataset):
         The study identifier for every cell.
     """
 
-    def __init__(self, X, Y, study=None):
+    def __init__(
+        self,
+        X: "numpy.ndarray",
+        Y: "numpy.ndarray",
+        study: Optional["numpy.ndarray"] = None,
+    ):
         self.X = X
         self.Y = Y
         self.study = study
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.Y)
 
-    def __getitem__(self, idx):
+    def __getitem__(
+        self, idx: "Index"
+    ) -> tuple["numpy.ndarray", "numpy.ndarray", Optional["numpy.ndarray"]]:
         # data, label, study
-        return self.X[idx].toarray().flatten(), self.Y[idx], self.study[idx]
+        return (
+            self.X[idx].toarray().flatten(),
+            self.Y[idx],
+            self.study[idx] if self.study is not None else None,
+        )
 
 
 class scCollator:
@@ -51,11 +75,7 @@ class scCollator:
         Use sparse matrices.
     """
 
-    def __init__(
-        self,
-        label2int: dict,
-        sparse: bool = False,
-    ):
+    def __init__(self, label2int: dict, sparse: bool = False):
         self.label2int = label2int
         self.sparse = sparse
 
@@ -65,11 +85,8 @@ def __call__(self, batch):
         X = torch.squeeze(torch.Tensor(np.vstack(profiles)))
         if self.sparse:
             X = X.to_sparse()
-        return (
-            X,
-            torch.Tensor([self.label2int[l] for l in labels]),
-            np.array(studies),
-        )
+
+        return (X, torch.Tensor([self.label2int[l] for l in labels]), np.array(studies))
 
 
 class MetricLearningDataModule(pl.LightningDataModule):

src/scimilarity/cell_annotation.py

@@ -1,7 +1,13 @@
-from typing import Optional, Union, List, Set, Tuple
+from typing import Optional, Union, List, Set, Tuple, TYPE_CHECKING
+
 
 from .cell_search_knn import CellSearchKNN
 
+if TYPE_CHECKING:
+    import anndata
+    import numpy
+    import pandas
+
 
 class CellAnnotation(CellSearchKNN):
     """A class that annotates cells using a cell embedding and then knn search.
@@ -21,10 +27,7 @@ class CellAnnotation(CellSearchKNN):
     """
 
     def __init__(
-        self,
-        model_path: str,
-        use_gpu: bool = False,
-        filenames: Optional[dict] = None,
+        self, model_path: str, use_gpu: bool = False, filenames: Optional[dict] = None
     ):
         import os
 

src/scimilarity/cell_embedding.py

@@ -1,4 +1,8 @@
-from typing import Optional, Tuple, Union
+from typing import Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import scipy.sparse
+    import numpy
 
 
 class CellEmbedding:
@@ -112,9 +116,9 @@ def get_embeddings(
         if (
             (isinstance(X, csr_matrix) or isinstance(X, csc_matrix))
             and (
-                isinstance(X.data, zarr.core.Array)
-                or isinstance(X.indices, zarr.core.Array)
-                or isinstance(X.indptr, zarr.core.Array)
+                isinstance(X.data, zarr.Array)
+                or isinstance(X.indices, zarr.Array)
+                or isinstance(X.indptr, zarr.Array)
             )
             and num_cells <= buffer_size
         ):
@@ -141,7 +145,7 @@ def get_embeddings(
                 embedding_parts.append(self.model(profiles))
 
         if not embedding_parts:
-            raise RuntimeError(f"No valid cells detected.")
+            raise RuntimeError("No valid cells detected.")
 
         if self.use_gpu:
             # detach, move from gpu into cpu, return as numpy array
@@ -151,6 +155,6 @@ def get_embeddings(
             embedding = torch.vstack(embedding_parts).detach().numpy()
 
         if np.isnan(embedding).any():
-            raise RuntimeError(f"NaN detected in embeddings.")
+            raise RuntimeError("NaN detected in embeddings.")
 
         return embedding

src/scimilarity/cell_query.py

@@ -1,7 +1,12 @@
-from typing import Dict, List, Optional, Tuple, Union, Set
+from typing import Dict, List, Optional, Tuple, Union, TYPE_CHECKING
 
 from .cell_search_knn import CellSearchKNN
 
+if TYPE_CHECKING:
+    import anndata
+    import numpy
+    import pandas
+
 
 class CellQuery(CellSearchKNN):
     """A class that searches for similar cells using a cell embedding.
@@ -39,8 +44,6 @@ def __init__(
         load_knn: bool = True,
     ):
         import os
-        import numpy as np
-        import pandas as pd
         import tiledb
 
         super().__init__(
@@ -340,7 +343,7 @@ def search_centroid_nearest(
 
         import numpy as np
         from scipy.cluster.vq import kmeans
-        from .utils import get_centroid, get_dist2centroid
+        from .utils import get_centroid
 
         cells = adata[adata.obs[centroid_key] == 1].copy()
         centroid = get_centroid(cells.layers["counts"])
@@ -623,7 +626,7 @@ def search_centroid_exhaustive(
 
         import numpy as np
         from scipy.cluster.vq import kmeans
-        from .utils import get_centroid, get_dist2centroid
+        from .utils import get_centroid
 
         cells = adata[adata.obs[centroid_key] == 1].copy()
         centroid = get_centroid(cells.layers["counts"])

src/scimilarity/cell_search_knn.py

@@ -1,7 +1,10 @@
-from typing import Optional, Tuple, Union
+from typing import Tuple, TYPE_CHECKING
 
 from .cell_embedding import CellEmbedding
 
+if TYPE_CHECKING:
+    import numpy
+
 
 class CellSearchKNN(CellEmbedding):
     """A class for searching similar cells using cell embeddings kNN.

src/scimilarity/interpreter.py

@@ -1,5 +1,11 @@
 from torch import nn
-from typing import Optional, Union
+from typing import Optional, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch
+    import numpy
+    import pandas
+    import scipy.sparse
 
 
 class SimpleDist(nn.Module):
@@ -15,11 +21,7 @@ def __init__(self, encoder: "torch.nn.Module"):
         super().__init__()
         self.encoder = encoder
 
-    def forward(
-        self,
-        anchors: "torch.Tensor",
-        negatives: "torch.Tensor",
-    ):
+    def forward(self, anchors: "torch.Tensor", negatives: "torch.Tensor"):
         """Forward.
 
         Parameters
@@ -55,11 +57,7 @@ class Interpreter:
     >>> interpreter = Interpreter(CellEmbedding("/opt/data/model").model)
     """
 
-    def __init__(
-        self,
-        encoder: "torch.nn.Module",
-        gene_order: list,
-    ):
+    def __init__(self, encoder: "torch.nn.Module", gene_order: list):
         from captum.attr import IntegratedGradients
 
         self.encoder = encoder

src/scimilarity/ontologies.py

@@ -1,7 +1,11 @@
 import networkx as nx
 import obonet
 import pandas as pd
-from typing import Union, Tuple, List
+from typing import Union, Tuple, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import pandas
+    import numpy
 
 
 def subset_nodes_to_set(nodes, restricted_set: Union[list, set]) -> nx.DiGraph:

src/scimilarity/tiledb_data_models.py

@@ -1,12 +1,12 @@
 import numpy as np
-import os, re
+import os
 import pandas as pd
 import pytorch_lightning as pl
 from scipy.sparse import coo_matrix, diags
 import tiledb
 import torch
 from torch.utils.data import Dataset, DataLoader, Sampler
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, TYPE_CHECKING
 
 from .utils import query_tiledb_df
 from .ontologies import (
@@ -19,6 +19,9 @@
 
 log = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    import pandas
+
 
 class scDataset(Dataset):
     """A class that represents cells in TileDB.
@@ -29,10 +32,7 @@ class scDataset(Dataset):
         Pandas dataframe of valid cells.
     """
 
-    def __init__(
-        self,
-        data_df: "pandas.DataFrame",
-    ):
+    def __init__(self, data_df: "pandas.DataFrame"):
         self.data_df = data_df
 
     def __len__(self):