feat: add third-party evaluation metrics

MODULE.bazel

@@ -4,7 +4,7 @@ pyrovelocity MODULE
 
 module(
     name = "pyrovelocity",
-    version = "0.4.1",
+    version = "0.4.2",
     compatibility_level = 1,
 )
 

Makefile

@@ -934,8 +934,8 @@ approve-prs: ## Approve github pull requests from bots: PR_ENTRIES="2-5 10 12-18
 		fi; \
 	done
 
-PREVIOUS_VERSION := 0.4.0
-NEXT_VERSION := 0.4.1
+PREVIOUS_VERSION := 0.4.1
+NEXT_VERSION := 0.4.2
 VERSION_FILES := \
 	pyproject.toml \
 	conda/colab/construct.yaml \

conda/colab/construct.yaml

@@ -1,5 +1,5 @@
 name: pyrovelocity-colab
-version: 0.4.1
+version: 0.4.2
 
 channels:
   - pytorch
@@ -89,7 +89,7 @@ specs:
   - brotli-python=1.1.0=py311hfdbb021_2
   - bzip2=1.0.8=h4bc722e_7
   - c-ares=1.34.4=hb9d3cd8_0
-  - ca-certificates=2025.1.31=hbcca0.4.1
+  - ca-certificates=2025.1.31=hbcca054_0
   - cached-property=1.5.2=hd8ed1ab_1
   - cached_property=1.5.2=pyha770c72_1
   - cachetools=5.5.2=pyhd8ed1ab_0
@@ -218,7 +218,7 @@ specs:
   - h2=4.2.0=pyhd8ed1ab_0
   - h5netcdf=1.6.1=pyhd8ed1ab_0
   - h5py=3.13.0=nompi_py311hb639ac4_100
-  - harfbuzz=10.4.1=h76408a6_0
+  - harfbuzz=10.4.0=h76408a6_0
   - hdf5=1.14.3=nompi_h2d575fe_109
   - hicolor-icon-theme=0.17=ha770c72_2
   - hpack=4.1.0=pyhd8ed1ab_0
@@ -294,7 +294,7 @@ specs:
   - libcufile=1.13.1.3=0
   - libcups=2.3.3=h4637d8d_4
   - libcurand=10.3.9.90=0
-  - libcurl=8.12.1=h332b0.4.1
+  - libcurl=8.12.1=h332b0f4_0
   - libcusolver=11.6.1.9=0
   - libcusparse=12.3.1.170=0
   - libdeflate=1.23=h4ddbbb0_0

containers/gpu.Dockerfile

@@ -77,7 +77,7 @@ COPY . /root
 # development
 RUN pip install --no-deps -e .
 # distribution
-# RUN pip install pyrovelocity==0.4.1
+# RUN pip install pyrovelocity==0.4.2
 
 ARG tag
 ENV FLYTE_INTERNAL_IMAGE $tag

containers/pkg.Dockerfile

@@ -34,7 +34,7 @@ COPY . /root
 # development
 RUN pip install --no-deps -e .
 # distribution
-# RUN pip install pyrovelocity==0.4.1
+# RUN pip install pyrovelocity==0.4.2
 
 ARG tag
 ENV FLYTE_INTERNAL_IMAGE $tag

docs/source/notebooks/pyrovelocity_colab_template.ipynb

@@ -70,7 +70,7 @@
         }
       ],
       "source": [
-        "pyrovelocity_version = \"0.4.1\"\n",
+        "pyrovelocity_version = \"0.4.2\"\n",
         "pyrovelocity_colab_script_url = (\n",
         "    \"https://storage.googleapis.com/pyrovelocity/data/scripts/\"\n",
         "    + f\"pyrovelocity-colab-{pyrovelocity_version}-Linux-x86_64.sh\"\n",

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyrovelocity"
-version = "0.4.1"
+version = "0.4.2"
 packages = [{ include = "pyrovelocity", from = "src" }]
 description = "A multivariate RNA Velocity model to estimate future cell states with uncertainty using probabilistic modeling with pyro."
 authors = ["pyrovelocity team"]

scripts/conda

@@ -3,7 +3,7 @@
 set -euo pipefail
 
 PACKAGE_NAME="pyrovelocity"
-PACKAGE_VERSION="0.4.1"
+PACKAGE_VERSION="0.4.2"
 CONDA_BUILD_STRING="pyhff70e4c"
 CONDA_BUILD_NUMBER="0"
 # CONDA_CHANNEL_LABEL="pyrovelocity_dev"
@@ -32,7 +32,7 @@ Example:
 
     ./conda \\
     --name pyrovelocity \\
-    --version 0.4.1.dev1 \\
+    --version 0.4.2.dev1 \\
     --build-string pyhff70e4c \\
     --build-number 0 \\
     --label pyrovelocity_dev
@@ -67,9 +67,9 @@ PACKAGE_SPEC="conda-forge/label/\
 $CONDA_CHANNEL_LABEL::\
 $PACKAGE_NAME=$PACKAGE_VERSION=$CONDA_BUILD_STRING"_"$CONDA_BUILD_NUMBER"
 
-BLUE="\0.4.1;34;1m"
+BLUE="\0.4.2;34;1m"
 BOLD="\033[1m"
-NO_COLOR="\0.4.1m"
+NO_COLOR="\0.4.2m"
 if [ "$USE_COLOR" = false ]; then
     BLUE=""
     BOLD=""

src/pyrovelocity/metrics/trajectory.py

@@ -0,0 +1,200 @@
+"""
+Trajectory evaluation metrics for velocity models.
+
+This module contains metrics for evaluating velocity model trajectories,
+including directional correctness and coherence measures taken from:
+
+> Qiao C, Huang Y. Representation learning of RNA velocity reveals robust cell
+> transitions. Proc Natl Acad Sci U S A. 2021;118. doi:10.1073/pnas.2105859118.
+"""
+
+import numpy as np
+from anndata import AnnData
+from beartype import beartype
+from beartype.typing import Dict, List, Tuple, Union
+from scipy import sparse
+from sklearn.metrics.pairwise import cosine_similarity
+
+
+@beartype
+def keep_type(
+    adata, nodes: np.ndarray, target: str, k_cluster: str
+) -> np.ndarray:
+    """
+    Select cells of given target type.
+
+    This implementation is included as a dependency of cross_boundary_correctness.
+    It is based on the original implementation:
+    https://github.com/qiaochen/VeloAE/blob/v0.2.0/veloproj/eval_util.py#L28-L41
+
+    Args:
+        adata: AnnData object
+        nodes: Indexes for cells
+        target: Cluster name
+        k_cluster: Cluster key in adata.obs dataframe
+
+    Returns:
+        Selected cells matching the target cluster
+    """
+    return nodes[adata.obs[k_cluster][nodes].values == target]
+
+
+@beartype
+def cross_boundary_correctness(
+    adata: AnnData,
+    k_cluster: str,
+    cluster_edges: List[Tuple[str, str]],
+    k_velocity: str = "velocity",
+    return_raw: bool = False,
+    x_emb: str = "X_umap",
+) -> Union[
+    Dict[Tuple[str, str], List[float]],
+    Tuple[Dict[Tuple[str, str], float], float],
+]:
+    """
+    Cross-Boundary Direction Correctness Score.
+
+    Calculates how well velocity vectors point toward neighboring cells in
+    adjacent clusters, measuring the model's ability to predict correct
+    developmental trajectories.
+
+    Qiao C, Huang Y. Representation learning of RNA velocity reveals robust cell
+    transitions. Proc Natl Acad Sci U S A. 2021;118. doi:10.1073/pnas.2105859118
+
+    This implementation is based on the original implementation:
+    https://github.com/qiaochen/VeloAE/blob/v0.2.0/veloproj/eval_util.py#L146-L200
+
+    Args:
+        adata: AnnData object
+        k_cluster: Key to the cluster column in adata.obs DataFrame
+        cluster_edges: Pairs of clusters with transition direction A->B
+        k_velocity: Key to the velocity matrix in adata.obsm
+        return_raw: Whether to return raw cell scores or aggregated scores
+        x_emb: Key to embedding for visualization
+
+    Returns:
+        Raw cell scores by cluster edge or mean scores by cluster edge and overall mean
+    """
+    if "neighbors" not in adata.uns:
+        raise ValueError("AnnData object must have neighbors computed")
+
+    if "indices" not in adata.uns["neighbors"]:
+        k = adata.uns["neighbors"]["params"]["n_neighbors"]
+        connectivities = adata.obsp["connectivities"]
+
+        if sparse.issparse(connectivities):
+            connectivities_array = connectivities.toarray()
+        else:
+            connectivities_array = connectivities
+
+        neighbor_indices = np.argsort(-connectivities_array, axis=1)[:, :k]
+        adata.uns["neighbors"]["indices"] = neighbor_indices
+
+    scores = {}
+    all_scores = {}
+
+    x_emb_data = adata.obsm[x_emb]
+
+    if x_emb == "X_umap":
+        v_emb = adata.obsm[f"{k_velocity}_umap"]
+    else:
+        v_emb = adata.obsm[
+            [key for key in adata.obsm if key.startswith(k_velocity)][0]
+        ]
+
+    for u, v in cluster_edges:
+        sel = adata.obs[k_cluster] == u
+        nbs = adata.uns["neighbors"]["indices"][sel]  # [n * 30]
+
+        boundary_nodes = map(
+            lambda nodes: keep_type(adata, nodes, v, k_cluster), nbs
+        )
+        x_points = x_emb_data[sel]
+        x_velocities = v_emb[sel]
+
+        type_score = []
+        for x_pos, x_vel, nodes in zip(x_points, x_velocities, boundary_nodes):
+            if len(nodes) == 0:
+                continue
+
+            position_dif = x_emb_data[nodes] - x_pos
+            dir_scores = cosine_similarity(
+                position_dif, x_vel.reshape(1, -1)
+            ).flatten()
+            type_score.append(np.mean(dir_scores))
+
+        scores[(u, v)] = np.mean(type_score) if type_score else 0.0
+        all_scores[(u, v)] = type_score
+
+    if return_raw:
+        return all_scores
+
+    return scores, np.mean([sc for sc in scores.values()]) if scores else 0.0
+
+
+@beartype
+def inner_cluster_coherence(
+    adata, k_cluster: str, k_velocity: str, return_raw: bool = False
+) -> Union[Dict[str, List[float]], Tuple[Dict[str, float], float]]:
+    """
+    In-cluster Coherence Score.
+
+    Measures how aligned velocity vectors are within the same cluster,
+    indicating consistency in predicted cellular trajectories.
+
+    Qiao C, Huang Y. Representation learning of RNA velocity reveals robust cell
+    transitions. Proc Natl Acad Sci U S A. 2021;118. doi:10.1073/pnas.2105859118
+
+    This implementation is based on the original implementation:
+    https://github.com/qiaochen/VeloAE/blob/v0.2.0/veloproj/eval_util.py#L203-L237
+
+    Args:
+        adata: AnnData object
+        k_cluster: Key to the cluster column in adata.obs DataFrame
+        k_velocity: Key to the velocity matrix in adata.layers
+        return_raw: Whether to return raw scores or aggregated scores
+
+    Returns:
+        Raw scores by cluster or mean scores by cluster and overall mean
+    """
+    if "neighbors" not in adata.uns:
+        raise ValueError("AnnData object must have neighbors computed")
+
+    if "indices" not in adata.uns["neighbors"]:
+        k = adata.uns["neighbors"]["params"]["n_neighbors"]
+        connectivities = adata.obsp["connectivities"]
+
+        if sparse.issparse(connectivities):
+            connectivities_array = connectivities.toarray()
+        else:
+            connectivities_array = connectivities
+
+        neighbor_indices = np.argsort(-connectivities_array, axis=1)[:, :k]
+        adata.uns["neighbors"]["indices"] = neighbor_indices
+
+    clusters = np.unique(adata.obs[k_cluster])
+    scores = {}
+    all_scores = {}
+
+    for cat in clusters:
+        sel = adata.obs[k_cluster] == cat
+        nbs = adata.uns["neighbors"]["indices"][sel]
+        same_cat_nodes = map(
+            lambda nodes: keep_type(adata, nodes, cat, k_cluster), nbs
+        )
+        velocities = adata.layers[k_velocity]
+        cat_vels = velocities[sel]
+
+        cat_score = [
+            cosine_similarity(cat_vels[[ith]], velocities[nodes]).mean()
+            for ith, nodes in enumerate(same_cat_nodes)
+            if len(nodes) > 0
+        ]
+
+        all_scores[cat] = cat_score
+        scores[cat] = np.mean(cat_score) if cat_score else 0.0
+
+    if return_raw:
+        return all_scores
+
+    return scores, np.mean([sc for sc in scores.values()]) if scores else 0.0

src/pyrovelocity/tasks/__init__.py

@@ -1,12 +1,13 @@
 import pyrovelocity.tasks.data
+import pyrovelocity.tasks.evaluate
 import pyrovelocity.tasks.postprocess
 import pyrovelocity.tasks.preprocess
 import pyrovelocity.tasks.summarize
 import pyrovelocity.tasks.train
 
-
 __all__ = [
     "data",
+    "evaluate",
     "postprocess",
     "preprocess",
     "summarize",