feat: Introduce merge_small_segments to collapse empty graph segments and update pseudotime warnings.

Author: LouisFaure

.gitignore

@@ -2,6 +2,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
+.DS_Store
 
 # C extensions
 *.so

scFates/tests/test_merge_segments.py

@@ -0,0 +1,133 @@
+"""Test for merge_small_segments function."""
+import warnings
+import numpy as np
+import scFates as scf
+
+scf.settings.verbosity = 0
+
+
+def test_merge_small_segments():
+    """Test that merge_small_segments handles empty segments correctly."""
+    # Load test data and build a tree
+    adata = scf.datasets.test_adata()
+
+    # Build a tree with many nodes to potentially create short segments
+    scf.tl.tree(
+        adata,
+        Nodes=100,
+        use_rep="pca",
+        method="ppt",
+        device="cpu",
+        ppt_sigma=1,
+        ppt_lambda=200,
+        seed=42,
+    )
+    scf.tl.cleanup(adata)
+    scf.tl.root(adata, adata.uns["graph"]["tips"][0])
+
+    # Run pseudotime and check for warning about empty segments
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        scf.tl.pseudotime(adata)
+        empty_seg_warnings = [
+            warning for warning in w
+            if "Some segs have no cell assigned" in str(warning.message)
+        ]
+
+    # If there are empty segments, test the merge function
+    if len(empty_seg_warnings) > 0:
+        # Get initial state
+        initial_n_nodes = adata.uns["graph"]["B"].shape[0]
+
+        # Run merge_small_segments
+        scf.tl.merge_small_segments(adata)
+
+        # Verify nodes were removed
+        final_n_nodes = adata.uns["graph"]["B"].shape[0]
+        assert final_n_nodes <= initial_n_nodes, "Nodes should be merged"
+
+        # Verify no empty segments remain after re-running pseudotime
+        with warnings.catch_warnings(record=True) as w2:
+            warnings.simplefilter("always")
+            scf.tl.pseudotime(adata)
+            new_warnings = [
+                warning for warning in w2
+                if "Some segs have no cell assigned" in str(warning.message)
+            ]
+        assert len(new_warnings) == 0, "No empty segments should remain"
+    else:
+        # No empty segments - just verify the function runs without error
+        scf.tl.merge_small_segments(adata)
+
+    # Verify basic graph integrity
+    assert adata.uns["graph"]["B"].shape[0] == adata.uns["graph"]["B"].shape[1]
+    assert adata.uns["graph"]["F"].shape[1] == adata.uns["graph"]["B"].shape[0]
+    assert adata.obsm["X_R"].shape[1] == adata.uns["graph"]["B"].shape[0]
+    assert "t" in adata.obs
+    assert "seg" in adata.obs
+
+
+def test_merge_small_segments_synthetic_empty():
+    """Test merge with synthetically created empty segment."""
+    adata = scf.datasets.test_adata()
+
+    scf.tl.tree(
+        adata,
+        Nodes=50,
+        use_rep="pca",
+        method="ppt",
+        device="cpu",
+        ppt_sigma=1,
+        ppt_lambda=100,
+        seed=1,
+    )
+    scf.tl.cleanup(adata)
+    scf.tl.root(adata, adata.uns["graph"]["tips"][0])
+    scf.tl.pseudotime(adata)
+
+    # Force an empty segment by reassigning cells
+    pp_seg = adata.uns["graph"]["pp_seg"]
+    if len(pp_seg) > 1:
+        # Find smallest segment and remove its cells
+        seg_counts = adata.obs.seg.value_counts()
+        smallest_seg = seg_counts.idxmin()
+        
+        # Move cells to another segment
+        other_segs = seg_counts.index[seg_counts.index != smallest_seg]
+        if len(other_segs) > 0:
+            adata.obs.loc[adata.obs.seg == smallest_seg, "seg"] = other_segs[0]
+
+            # Now run merge
+            initial_nodes = adata.uns["graph"]["B"].shape[0]
+            scf.tl.merge_small_segments(adata)
+            
+            # Verify merge happened
+            assert adata.uns["graph"]["B"].shape[0] <= initial_nodes
+
+
+def test_merge_no_empty_segments():
+    """Test that function handles case with no empty segments gracefully."""
+    adata = scf.datasets.test_adata()
+
+    scf.tl.tree(
+        adata,
+        Nodes=30,
+        use_rep="pca",
+        method="ppt",
+        device="cpu",
+        ppt_sigma=1,
+        ppt_lambda=100,
+        seed=1,
+    )
+    scf.tl.cleanup(adata)
+    scf.tl.root(adata, adata.uns["graph"]["tips"][0])
+    scf.tl.pseudotime(adata)
+
+    initial_nodes = adata.uns["graph"]["B"].shape[0]
+
+    # Should complete without error even if no empty segments
+    scf.tl.merge_small_segments(adata)
+
+    # Verify graph is still valid
+    assert adata.uns["graph"]["B"].shape[0] <= initial_nodes
+    assert "t" in adata.obs

scFates/tools/__init__.py

@@ -5,6 +5,7 @@
     attach_tree,
     simplify,
     convert_to_soft,
+    merge_small_segments,
 )
 from .graph_fitting import tree, curve, circle, explore_sigma
 from .root import root, roots

scFates/tools/graph_operations.py

@@ -589,6 +589,157 @@ def simplify(adata: AnnData, n_nodes: int = 10, copy: bool = False):
     return adata if copy else None
 
 
+def merge_small_segments(adata: AnnData, copy: bool = False):
+    """\
+    Merge segments that have no cells assigned into their neighboring nodes.
+
+    Small segments (typically between two high-degree nodes like forks) can end up
+    with no cells assigned during pseudotime calculation. This function merges
+    such segments by collapsing intermediate nodes into the earlier milestone,
+    resolving downstream analysis failures.
+
+    This function should be called after pseudotime calculation if warnings about
+    segments having no cells assigned are encountered.
+
+    Parameters
+    ----------
+    adata
+        Annotated data matrix with computed pseudotime.
+    copy
+        Return a copy instead of writing to adata.
+
+    Returns
+    -------
+    adata : anndata.AnnData
+        if `copy=True` it returns or else updates fields in `adata`:
+
+        `.uns['graph']['B']`
+            updated adjacency matrix with merged nodes.
+        `.uns['graph']['F']`
+            updated coordinates with merged nodes.
+        `.obsm['X_R']`
+            updated soft assignment matrix with merged nodes.
+        `.uns['graph']['pp_seg']`
+            updated segment definitions.
+        `.uns['graph']['pp_info']`
+            updated node information.
+    """
+
+    logg.info("merging small segments without cells", reset=True)
+
+    adata = adata.copy() if copy else adata
+
+    if "t" not in adata.obs:
+        raise ValueError(
+            "You need to run `tl.pseudotime` before merging small segments."
+        )
+
+    graph = adata.uns["graph"]
+    pp_seg = graph["pp_seg"].copy()
+    pp_info = graph["pp_info"].copy()
+    B = graph["B"].copy()
+    F = graph["F"].copy()
+    R = adata.obsm["X_R"].copy()
+
+    # Identify segments with no cells assigned
+    cell_segs = adata.obs.seg.value_counts()
+    all_segs = pp_seg.n.values
+    empty_segs = [s for s in all_segs if s not in cell_segs.index]
+
+    if len(empty_segs) == 0:
+        logg.info("    no empty segments found", time=False)
+        return adata if copy else None
+
+    logg.info(f"    found {len(empty_segs)} empty segment(s): {empty_segs}", time=False)
+
+    # Build graph for path finding
+    g = igraph.Graph.Adjacency((B > 0).tolist(), mode="undirected")
+
+    # Track nodes to remove
+    nodes_to_remove = set()
+
+    for seg_id in empty_segs:
+        seg_row = pp_seg.loc[pp_seg.n == seg_id].iloc[0]
+        from_node = int(seg_row["from"])
+        to_node = int(seg_row["to"])
+
+        # Get all nodes in this segment
+        seg_nodes = pp_info.index[pp_info.seg == seg_id].tolist()
+
+        # Identify intermediate nodes (not the from/to milestones)
+        intermediate_nodes = [n for n in seg_nodes if n != from_node and n != to_node]
+
+        # The from_node is earlier in pseudotime, we keep it
+        # Merge intermediate nodes and to_node into from_node
+
+        # Combine soft assignments: add R columns of removed nodes to from_node
+        for node in intermediate_nodes + [to_node]:
+            if node < R.shape[1]:
+                R[:, from_node] = R[:, from_node] + R[:, node]
+                nodes_to_remove.add(node)
+
+        # Update adjacency: connect from_node to to_node's neighbors
+        if to_node < B.shape[0]:
+            to_neighbors = np.where(B[to_node, :] > 0)[0]
+            for neighbor in to_neighbors:
+                if neighbor != from_node and neighbor not in nodes_to_remove:
+                    B[from_node, neighbor] = 1
+                    B[neighbor, from_node] = 1
+
+    if len(nodes_to_remove) == 0:
+        logg.info("    no nodes to merge", time=False)
+        return adata if copy else None
+
+    # Remove nodes from matrices
+    nodes_to_keep = [i for i in range(B.shape[0]) if i not in nodes_to_remove]
+    nodes_to_keep = np.array(nodes_to_keep)
+
+    B = B[np.ix_(nodes_to_keep, nodes_to_keep)]
+    F = F[:, nodes_to_keep]
+    R = R[:, nodes_to_keep]
+
+    # Normalize R
+    Rsum = R.sum(axis=1)
+    Rsum[Rsum == 0] = 1  # Avoid division by zero
+    R = R / Rsum.reshape(-1, 1)
+
+    # Update graph structures
+    g = igraph.Graph.Adjacency((B > 0).tolist(), mode="undirected")
+    tips = np.argwhere(np.array(g.degree()) == 1).flatten()
+    forks = np.argwhere(np.array(g.degree()) > 2).flatten()
+
+    # Update adata
+    adata.uns["graph"]["B"] = B
+    adata.uns["graph"]["F"] = F
+    adata.uns["graph"]["tips"] = tips
+    adata.uns["graph"]["forks"] = forks
+    adata.obsm["X_R"] = R
+
+    # Create mapping from old to new indices
+    old_to_new = {old: new for new, old in enumerate(nodes_to_keep)}
+
+    # Update milestones
+    if "milestones" in graph:
+        new_milestones = {}
+        for name, old_idx in graph["milestones"].items():
+            if old_idx in old_to_new:
+                new_milestones[name] = old_to_new[old_idx]
+        adata.uns["graph"]["milestones"] = new_milestones
+
+    # Update root
+    if "root" in graph and graph["root"] in old_to_new:
+        adata.uns["graph"]["root"] = old_to_new[graph["root"]]
+
+    # Recalculate pseudotime with new graph structure
+    root(adata, adata.uns["graph"]["root"])
+    pseudotime(adata)
+
+    n_removed = len(nodes_to_remove)
+    logg.info("    finished", time=True, end=" " if settings.verbosity > 2 else "\n")
+    logg.hint(f"merged {n_removed} nodes from {len(empty_segs)} empty segment(s)")
+
+    return adata if copy else None
+
 def getpath(adata, root_milestone, milestones, include_root=False):
     """\
     Obtain dataframe of cell of a given path.

scFates/tools/pseudotime.py

@@ -157,7 +157,8 @@ def pseudotime(
 
         if allsegs.shape[1]!=allsegs_complete.shape[1]:
             missing = allsegs_complete.columns[~allsegs_complete.columns.isin(allsegs.columns)]
-            message=f"Some segs have no cell assigned: {missing.tolist()}"
+            message=f"Some segs have no cell assigned: {missing.tolist()}. " \
+                    "This may cause downstream failures. Consider running `scf.tl.merge_small_segments(adata)`."
             warnings.warn(message)
 
         for c in allsegs.columns:
@@ -248,7 +249,8 @@ def pseudotime(
     adata.obs["milestones"] = milestones_str.astype("category")
 
     if adata.obs["milestones"].isna().sum()>0:
-        message = "Some cells have no milestones assigned. This is likely due to the fact that these uniquely compose a segment."
+        message = "Some cells have no milestones assigned. This is likely due to the fact that these uniquely compose a segment. " \
+                  "Consider running `scf.tl.merge_small_segments(adata)`."
         warnings.warn(message)
 
     adata.uns["graph"]["milestones"] = dict(