Refactoring preprocess_data, calculate_entropy and commitment_score

README.md

@@ -38,6 +38,10 @@ There are several options to intall regvelo:
   pip install git+https://github.com/theislab/regvelo.git@main
 ```
 
+## RegVelo agentic workflow via Biomni lab 
+
+RegVelo pipeline can be run using an agentic workflow via [Biomni lab](https://biomni.phylo.bio/), which allows user to upload custom datasets and run full analysis using a chat-based approach. 
+
 ## Citation
 
 If you find RegVelo useful for your research, please consider citing our work as:

src/regvelo/metrics/_utils.py

@@ -35,9 +35,9 @@ def calculate_entropy(prob_matrix : np.ndarray) -> np.ndarray:
     log_probs = np.zeros_like(prob_matrix)
     mask = prob_matrix != 0
     np.log2(prob_matrix, where=mask, out=log_probs)
-
     entropy = -np.sum(prob_matrix * log_probs, axis=1)
-    return entropy
+    max_entropy = np.log2(prob_matrix.shape[1])
+    return entropy, max_entropy
 
 def get_significance(pvalue: float) -> str:
     """Return significance annotation for a p-value.
@@ -63,4 +63,4 @@ def get_significance(pvalue: float) -> str:
     elif pvalue < 0.1:
         return "*"
     else:
-        return "n.s."
+        return "n.s."

src/regvelo/plotting/_commitment_score.py

@@ -34,7 +34,8 @@ def commitment_score(
         raise KeyError(f"Key '{lineage_key}' not found in `adata.obsm`.")
 
     p = pd.DataFrame(adata.obsm[lineage_key], columns=adata.obsm[lineage_key].names.tolist())
-    score = calculate_entropy(p)
+    entropy, max_entropy = calculate_entropy(p)
+    score = 1 - entropy / max_entropy
     adata.obs["commitment_score"] = np.array(score)
 
     sc.pl.umap(

src/regvelo/preprocessing/_preprocess_data.py

@@ -9,6 +9,7 @@ def preprocess_data(
     adata: AnnData,
     spliced_layer: str = "Ms",
     unspliced_layer: str = "Mu",
+    return_velo_genes = False, 
     min_max_scale: bool = True,
     filter_on_r2: bool = True,
 ) -> AnnData:
@@ -25,6 +26,8 @@ def preprocess_data(
         Name of the spliced layer.
     unspliced_layer
         Name of the unspliced layer.
+    return_velo_genes
+        Return velocity-informative genes (default=False).
     min_max_scale
         Whether to apply min-max scaling to the spliced and unspliced layers.
     filter_on_r2
@@ -34,6 +37,7 @@ def preprocess_data(
     -------
     Preprocessed annotated data object.
     """
+
     if min_max_scale:
         scaler = MinMaxScaler()
         adata.layers[spliced_layer] = scaler.fit_transform(adata.layers[spliced_layer])
@@ -51,4 +55,9 @@ def preprocess_data(
         ].copy()
         adata = adata[:, adata.var.velocity_genes].copy()
 
-    return adata
+    if return_velo_genes:
+        if not filter_on_r2:
+            raise ValueError("return_velo_genes=True requires filter_on_r2=True")
+        return adata.var_names.tolist()
+
+    return adata

src/regvelo/tools/_markov_density_screening.py

@@ -12,7 +12,8 @@
 import cellrank as cr
 import regvelo as rgv
 
-from ..plotting._utils import SIGNIFICANCE_PALETTE, delta_to_probability, smooth_score
+from ..plotting._utils import delta_to_probability, smooth_score
+from ..tools._visits_diff_per_tf import _visits_diff_per_tf
 
 
 @contextmanager
@@ -38,52 +39,6 @@ def _mute_cellrank():
         else:
             os.environ["TQDM_DISABLE"] = old_tqdm_disable
 
-
-def visits_diff_per_tf(
-    adata: AnnData,
-    terminal_states: Sequence[str],
-    dd_sig: np.ndarray,
-    sig_palette: dict,
-) -> tuple[pd.DataFrame, list[str]]:
-    """Collect per-terminal-state visit difference values and significance colours.
-
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData with 'visits_diff' stored in adata.obs.
-    terminal_states : sequence of str
-        Terminal state labels to iterate over.
-    dd_sig : np.ndarray
-        Array of p-values, one per terminal state, from rgv.tl.simulated_visit_diff.
-    sig_palette : dict
-        Mapping from significance label ('n.s.', '*', '**', '***') to hex colour.
-
-    Returns
-    -------
-    df : pd.DataFrame
-        Long-form DataFrame with columns 'Value' (visit diff) and 'Group' (state).
-    palette_rel : list of str
-        Hex colour per terminal state based on significance.
-    """
-    data = []
-    palette_rel = []
-
-    for i, ts in enumerate(terminal_states):
-        p_value = dd_sig[i]
-        terminal_indices_sub = np.where(adata.obs["term_states_fwd"].isin([ts]))[0]
-
-        values = adata.obs["visits_diff"].iloc[terminal_indices_sub]
-        subgroups = [ts] * len(values)
-
-        for val, subgrp in zip(values, subgroups):
-            data.append({"Value": val, "Group": subgrp})
-
-        significance = rgv.mt.get_significance(p_value)
-        palette_rel.append(sig_palette[significance])
-
-    return pd.DataFrame(data), palette_rel
-
-
 def _assign_visits_diff(
     adata: AnnData,
     adata_perturb: AnnData,
@@ -291,7 +246,7 @@ def markov_density_screening(
             res_table.loc[TF, f"dd_sig_{state}"] for state in TERMINAL_STATES
         ])
 
-        df, _ = visits_diff_per_tf(adata, TERMINAL_STATES, dd_sig_tf, SIGNIFICANCE_PALETTE)
+        df, _ = visits_diff_per_tf(adata, TERMINAL_STATES, dd_sig_tf)
         df["Factor"] = TF
 
         for state in TERMINAL_STATES:

src/regvelo/tools/_visits_diff_per_tf.py

@@ -0,0 +1,48 @@
+import numpy as np
+import pandas as pd
+
+from typing import Sequence
+from anndata import AnnData
+from ..metrics._utils import get_significance
+from ..plotting._utils import SIGNIFICANCE_PALETTE
+
+def visits_diff_per_tf(
+    adata: AnnData,
+    terminal_states: Sequence[str],
+    dd_sig: np.ndarray,
+) -> tuple[pd.DataFrame, list[str]]:
+    """Collect per-terminal-state visit difference values and significance colours.
+
+    Parameters
+    ----------
+    adata : AnnData
+        AnnData with 'visits_diff' stored in adata.obs.
+    terminal_states : sequence of str
+        Terminal state labels to iterate over.
+    dd_sig : np.ndarray
+        Array of p-values, one per terminal state, from rgv.tl.simulated_visit_diff.
+
+    Returns
+    -------
+    df : pd.DataFrame
+        Long-form DataFrame with columns 'Value' (visit diff) and 'Group' (state).
+    palette_rel : list of str
+        Hex colour per terminal state based on significance.
+    """
+    data = []
+    palette_rel = []
+
+    for i, ts in enumerate(terminal_states):
+        p_value = dd_sig[i]
+        terminal_indices_sub = np.where(adata.obs["term_states_fwd"].isin([ts]))[0]
+
+        values = adata.obs["visits_diff"].iloc[terminal_indices_sub]
+        subgroups = [ts] * len(values)
+
+        for val, subgrp in zip(values, subgroups):
+            data.append({"Value": val, "Group": subgrp})
+
+        significance = get_significance(p_value)
+        palette_rel.append(SIGNIFICANCE_PALETTE[significance])
+
+    return pd.DataFrame(data), palette_rel

tests/test_markov.py

@@ -0,0 +1,128 @@
+# Testing function for Markov screening and plotting
+import anndata as ad
+import numpy as np
+import pandas as pd
+import torch
+import cellrank as cr
+from scvi.data import synthetic_iid
+import regvelo as rgv
+from regvelo import REGVELOVI
+
+from .src.tools._TFscreening import TFscreening
+from .src.plotting._markov_screen import (
+    _visits_diff_per_tf,
+    _plot_visits_dist,
+    _plot_visits_dist_combined,
+)
+from .src.plotting._driver_TF_ranking import (
+    plot_top_TF,
+    compute_TF_regulon,
+    plot_grn_weight,
+    plot_GRN_per_TF,
+)
+
+# Common variables used in the test
+cluster_key = "cell_type"
+TERMINAL_STATES = ["mNC_head_mesenchymal"]
+STARTING_POINTS = ["start"]
+
+
+def test_markov():
+    # create a small synthetic dataset
+    adata = synthetic_iid()
+    adata.layers["spliced"] = adata.X.copy()
+    adata.layers["unspliced"] = adata.X.copy()
+
+    # give deterministic var names that are guaranteed to exist
+    adata.var_names = pd.Index([f"Gene{i}" for i in range(adata.n_vars)])
+    n_gene = adata.n_vars
+
+    # create a random GRN skeleton (DataFrame) and store in adata.uns
+    grn_matrix = np.random.choice([0, 1], size=(n_gene, n_gene), p=[0.8, 0.2])
+    W_df = pd.DataFrame(grn_matrix, index=adata.var_names, columns=adata.var_names)
+    adata.uns["skeleton"] = W_df
+    TF_list = adata.var_names.tolist()
+
+    # prepare tensor version of W for the model
+    W_tensor = torch.tensor(W_df.values.astype(np.float32))
+
+    # Ensure minimal required obs columns exist so downstream code can run
+    # create a simple clustering annotation and terminal state annotation
+    adata.obs[cluster_key] = np.random.choice(["ctype1", "ctype2"], size=adata.n_obs)
+    # default all to 'other' then mark a few cells as the terminal state used in the test
+    adata.obs["term_states_fwd"] = "other"
+    n_term = max(1, int(0.05 * adata.n_obs))
+    adata.obs.loc[adata.obs.index[:n_term], "term_states_fwd"] = TERMINAL_STATES[0]
+
+    # setup anndata for REGVELOVI and train a small model
+    REGVELOVI.setup_anndata(adata, spliced_layer="spliced", unspliced_layer="unspliced")
+    reg_vae = REGVELOVI(adata, W=W_tensor.T, regulators=TF_list)
+    reg_vae.train()
+
+    reg_vae.get_latent_representation()
+    reg_vae.get_velocity()
+    reg_vae.get_latent_time()
+
+    # CellRank-based macrostate / fate probability estimation
+    vk = cr.kernels.VelocityKernel(adata).compute_transition_matrix()
+    estimator = cr.estimators.GPCCA(vk)
+    estimator.compute_macrostates(n_states=10, cluster_key=cluster_key)
+    # set terminal states by name (this mirrors how the test_regvelo pipeline uses terminal state names)
+    estimator.set_terminal_states(TERMINAL_STATES)
+    estimator.compute_fate_probabilities(tol=1e-5)
+
+    MODEL = reg_vae
+    adata_perturb_dict = {}
+
+    # choose TF candidates from the actual var_names so knockouts exist
+    TF_candidate = [TF_list[0], TF_list[1]]
+
+    for TF in TF_candidate:
+        adata_target_perturb, reg_vae_perturb = rgv.tl.in_silico_block_simulation(model=MODEL, adata=adata, TF=TF, cutoff=0)
+        adata_perturb_dict[TF] = adata_target_perturb
+
+    # Build per-terminal-state cell index map from baseline annotations
+    ct_indices = {
+        ct: adata.obs["term_states_fwd"][adata.obs["term_states_fwd"] == ct].index.tolist()
+        for ct in TERMINAL_STATES
+    }
+
+    # compute macrostates and fate probabilities for each perturbed dataset
+    for TF, adata_target_perturb in adata_perturb_dict.items():
+        vkp = cr.kernels.VelocityKernel(adata_target_perturb).compute_transition_matrix()
+        estimator = cr.estimators.GPCCA(vkp)
+        estimator.compute_macrostates(n_states=10, cluster_key=cluster_key)
+        estimator.set_terminal_states(ct_indices)
+        estimator.compute_fate_probabilities()
+        adata_perturb_dict[TF] = adata_target_perturb
+
+    # run TF screening
+    res_df = TFscreening(
+        adata,
+        adata_perturb_dict,
+        TERMINAL_STATES,
+        STARTING_POINTS,
+        tf_ko_list=TF_candidate,
+        cluster_key=cluster_key,
+        method="stepwise",
+        n_step_to_use=500,
+    )
+
+    # plotting utilities (smoke checks)
+    plot_top_TF(res_df, adata, cluster_key=[cluster_key], threshold=0.1)
+
+    coef_targets, coef_regulators = compute_TF_regulon(
+        adata, MODEL, cluster_key=cluster_key, TF=[TF_candidate[0]], TERMINAL_STATES=TERMINAL_STATES
+    )
+
+    plot_GRN_per_TF(
+        adata,
+        MODEL,
+        cluster_key=[cluster_key],
+        TF=[TF_candidate[0]],
+        TERMINAL_STATES=TERMINAL_STATES,
+        terminal_state_to_plot=TERMINAL_STATES[0],
+        coef_targets=coef_targets,
+        coef_regulators=coef_regulators,
+        n_hits=10,
+    )