reverted some typhints

Author: srivarra

viscy/representation/embedding_writer.py

@@ -6,7 +6,6 @@
 import numpy as np
 import pandas as pd
 import torch
-import xarray as xr
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import BasePredictionWriter
 from numpy.typing import NDArray
@@ -23,7 +22,7 @@
 _logger = logging.getLogger("lightning.pytorch")
 
 
-def read_embedding_dataset(path: Path) -> xr.Dataset:
+def read_embedding_dataset(path: Path) -> Dataset:
     """Read the embedding dataset written by the EmbeddingWriter callback.
 
     Supports both legacy datasets (without x/y coordinates) and new datasets.
@@ -35,7 +34,7 @@ def read_embedding_dataset(path: Path) -> xr.Dataset:
 
     Returns
     -------
-    xr.Dataset
+    Dataset
         Xarray dataset with features and projections.
     """
     dataset = open_zarr(path)

viscy/representation/evaluation/__init__.py

@@ -18,19 +18,19 @@
 from pathlib import Path
 
 import pandas as pd
-import xarray as xr
 from viscy.data.triplet import TripletDataModule
+from xarray import DataArray
 
 
 def load_annotation(
-    da: xr.DataArray, path: str, name: str, categories: dict | None = None
+    da: DataArray, path: str, name: str, categories: dict | None = None
 ) -> pd.Series:
     """
     Load annotations from a CSV file and map them to the dataset.
 
     Parameters
     ----------
-    da : xr.DataArray
+    da : DataArray
         The dataset array containing 'fov_name' and 'id' coordinates.
     path : str
         Path to the CSV file containing annotations.

viscy/representation/evaluation/clustering.py

@@ -1,3 +1,5 @@
+"""Methods for evaluating clustering performance."""
+
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 from scipy.spatial.distance import cdist
@@ -10,12 +12,16 @@
 from sklearn.neighbors import KNeighborsClassifier
 
 
-def knn_accuracy(embeddings, annotations, k=5):
+def knn_accuracy(embeddings: NDArray, annotations: NDArray, k: int = 5) -> float:
     """
     Evaluate the k-NN classification accuracy.
 
     Parameters
     ----------
+    embeddings : NDArray
+        Embeddings to cluster.
+    annotations : NDArray
+        Ground truth labels.
     k : int, optional
         Number of neighbors to use for k-NN. Default is 5.
 

viscy/representation/evaluation/dimensionality_reduction.py

@@ -1,25 +1,30 @@
+from typing import TYPE_CHECKING
+
 import pandas as pd
 import umap
-import xarray as xr
 from numpy.typing import NDArray
 from sklearn.decomposition import PCA
 from sklearn.preprocessing import StandardScaler
+from xarray import Dataset
+
+if TYPE_CHECKING:
+    from phate import PHATE
 
 
 def compute_phate(
-    embedding_dataset: NDArray | xr.Dataset,
+    embedding_dataset: NDArray | Dataset,
     n_components: int = 2,
     knn: int = 5,
     decay: int = 40,
     update_dataset: bool = False,
     **phate_kwargs,
-) -> tuple[object, NDArray]:
+) -> tuple[PHATE, NDArray]:
     """
     Compute PHATE embeddings for features and optionally update dataset.
 
     Parameters
     ----------
-    embedding_dataset : xr.Dataset | NDArray
+    embedding_dataset : NDArray | Dataset
         The dataset containing embeddings, timepoints, fov_name, and track_id,
         or a numpy array of embeddings.
     n_components : int, optional
@@ -35,7 +40,7 @@ def compute_phate(
 
     Returns
     -------
-    tuple[object, NDArray]
+    tuple[phate.PHATE, NDArray]
         PHATE model and PHATE embeddings
 
     Raises
@@ -53,7 +58,7 @@ def compute_phate(
     # Get embeddings from dataset if needed
     embeddings = (
         embedding_dataset["features"].values
-        if isinstance(embedding_dataset, xr.Dataset)
+        if isinstance(embedding_dataset, Dataset)
         else embedding_dataset
     )
 
@@ -64,7 +69,7 @@ def compute_phate(
     phate_embedding = phate_model.fit_transform(embeddings)
 
     # Update dataset if requested
-    if update_dataset and isinstance(embedding_dataset, xr.Dataset):
+    if update_dataset and isinstance(embedding_dataset, Dataset):
         for i in range(
             min(2, phate_embedding.shape[1])
         ):  # Only update PHATE1 and PHATE2
@@ -73,12 +78,12 @@ def compute_phate(
     return phate_model, phate_embedding
 
 
-def compute_pca(embedding_dataset, n_components=None, normalize_features=True):
+def compute_pca(embedding_dataset: NDArray | Dataset, n_components=None, normalize_features=True):
     """Compute PCA embeddings for features and optionally update dataset.
 
     Parameters
     ----------
-    embedding_dataset : xr.Dataset or NDArray
+    embedding_dataset : Dataset | NDArray
         The dataset containing embeddings, timepoints, fov_name, and track_id,
         or a numpy array of embeddings.
     n_components : int, optional
@@ -93,7 +98,7 @@ def compute_pca(embedding_dataset, n_components=None, normalize_features=True):
     """
     embeddings = (
         embedding_dataset["features"].values
-        if isinstance(embedding_dataset, xr.Dataset)
+        if isinstance(embedding_dataset, Dataset)
         else embedding_dataset
     )
 
@@ -107,7 +112,7 @@ def compute_pca(embedding_dataset, n_components=None, normalize_features=True):
     pc_features = PCA_features.fit_transform(scaled_features)
 
     # Create base dictionary with id and fov_name
-    if isinstance(embedding_dataset, xr.Dataset):
+    if isinstance(embedding_dataset, Dataset):
         pca_dict = {
             "id": embedding_dataset["id"].values,
             "fov_name": embedding_dataset["fov_name"].values,
@@ -139,13 +144,13 @@ def _fit_transform_umap(
 
 
 def compute_umap(
-    embedding_dataset: xr.Dataset, normalize_features: bool = True
+    embedding_dataset: Dataset, normalize_features: bool = True
 ) -> tuple[umap.UMAP, umap.UMAP, pd.DataFrame]:
     """Compute UMAP embeddings for features and projections.
 
     Parameters
     ----------
-    embedding_dataset : xr.Dataset
+    embedding_dataset : Dataset
         Xarray dataset with features and projections.
     normalize_features : bool, optional
         Scale the input to zero mean and unit variance before fitting UMAP,

viscy/representation/evaluation/distance.py

@@ -2,14 +2,32 @@
 from typing import Literal
 
 import numpy as np
-import xarray as xr
+from numpy.typing import NDArray
 from sklearn.metrics.pairwise import cosine_similarity
+from xarray import Dataset
 
 
 def calculate_cosine_similarity_cell(
-    embedding_dataset: xr.Dataset, fov_name: str, track_id: int
-):
-    """Extract embeddings and calculate cosine similarities for a specific cell"""
+    embedding_dataset: Dataset, fov_name: str, track_id: int
+) -> tuple[NDArray, NDArray]:
+    """
+    
+    Extract embeddings and calculate cosine similarities for a specific cell
+    
+    Parameters
+    ----------
+    embedding_dataset : Dataset
+        Dataset containing embeddings and metadata
+    fov_name : str
+        Field of view identifier
+    track_id : int
+        Track identifier for the specific cell
+
+    Returns
+    -------
+    tuple[NDArray, NDArray]
+        Time points and cosine similarities for the specific cell
+    """
     filtered_data = embedding_dataset.where(
         (embedding_dataset["fov_name"] == fov_name)
         & (embedding_dataset["track_id"] == track_id),
@@ -25,7 +43,7 @@ def calculate_cosine_similarity_cell(
 
 
 def compute_displacement(
-    embedding_dataset: xr.Dataset,
+    embedding_dataset: Dataset,
     distance_metric: Literal["euclidean_squared", "cosine"] = "euclidean_squared",
 ) -> dict[int, list[float]]:
     """Compute the displacement or mean square displacement (MSD) of embeddings.
@@ -37,15 +55,13 @@ def compute_displacement(
 
     Parameters
     ----------
-    embedding_dataset : xarray.Dataset
+    embedding_dataset : Dataset
         Dataset containing embeddings and metadata
-    distance_metric : str
+    distance_metric : Literal["euclidean_squared", "cosine"]
         The metric to use for computing distances between embeddings.
         Valid options are:
-        - "euclidean": Euclidean distance (L2 norm)
         - "euclidean_squared": Squared Euclidean distance (for MSD, default)
         - "cosine": Cosine similarity
-        - "cosine_dissimilarity": 1 - cosine similarity
 
     Returns
     -------
@@ -152,13 +168,13 @@ def compute_dynamic_range(mean_displacement_per_tau: dict[int, float]):
     return max(displacements) - min(displacements)
 
 
-def compute_rms_per_track(embedding_dataset: xr.Dataset):
+def compute_rms_per_track(embedding_dataset: Dataset):
     """
     Compute RMS of the time derivative of embeddings per track.
 
     Parameters
     ----------
-    embedding_dataset : xarray.Dataset
+    embedding_dataset : Dataset
         The dataset containing embeddings, timepoints, fov_name, and track_id.
 
     Returns
@@ -204,13 +220,13 @@ def compute_rms_per_track(embedding_dataset: xr.Dataset):
 
 
 def calculate_normalized_euclidean_distance_cell(
-    embedding_dataset: xr.Dataset, fov_name: str, track_id: int
+    embedding_dataset: Dataset, fov_name: str, track_id: int
 ):
     """Calculate normalized euclidean distance for a specific cell track.
 
     Parameters
     ----------
-    embedding_dataset : xr.Dataset
+    embedding_dataset : Dataset
         Dataset containing embedding data with fov_name and track_id coordinates
     fov_name : str
         Field of view identifier

viscy/representation/evaluation/lca.py

@@ -5,18 +5,18 @@
 import pandas as pd
 import torch
 import torch.nn as nn
-import xarray as xr
 from captum.attr import IntegratedGradients, Occlusion
 from numpy.typing import NDArray
 from sklearn.linear_model import LogisticRegression
 from sklearn.metrics import classification_report
 from sklearn.preprocessing import StandardScaler
 from torch import Tensor
 from viscy.representation.contrastive import ContrastiveEncoder
+from xarray import DataArray
 
 
 def fit_logistic_regression(
-    features: xr.DataArray,
+    features: DataArray,
     annotations: pd.Series,
     train_fovs: list[str],
     remove_background_class: bool = True,
@@ -32,7 +32,7 @@ def fit_logistic_regression(
 
     Parameters
     ----------
-    features : xr.DataArray
+    features : DataArray
         Xarray of features.
     annotations : pd.Series
         Categorical class annotations with label values starting from 0.