Add perturb_fraction / random_state to make_neighbor_perturbation

src/cellina/_spatial_utils.py

@@ -209,7 +209,10 @@ def spatial_neighbors(adata: AnnData,
         return dist
 
 def _node_perturbation(X, var_idx, perturbations, groupby=None, labels=None,
-                       base=np.e, add_shift=False, renormalize=False):
+                       base=np.e, add_shift=False, renormalize=False,
+                       perturb_fraction=1.0, random_state=0):
+    if not 0.0 <= perturb_fraction <= 1.0:
+        raise ValueError(f"perturb_fraction must be in [0, 1], got {perturb_fraction}")
     n_vars = len(var_idx)
     neutral = 0.0 if add_shift else 1.0
 
@@ -236,13 +239,29 @@ def _encode(logfc):
     # Densify once
     X = np.asarray(X.toarray() if issparse(X) else X, dtype=np.float32)
 
+    # Keep the unperturbed expression so that, when only a fraction of cells is
+    # perturbed, the remaining cells can be restored to their original values.
+    X_orig = X.copy() if perturb_fraction < 1.0 else None
+
     if add_shift:  # log1p data + logFC shift
         X = X + transform
     else:          # counts; exp(logFC) applied multiplicatively
         # +1 / -1 avoids zeroing out genes multiplied by zero
         X = (X + 1.0) * transform - 1.0
     X = np.clip(X, 0, None)
 
+    # Fractional perturbation: only a random subset of cells keeps the shift;
+    # the rest are reverted to their original expression. Applied before
+    # renormalisation so reverted cells trivially keep their original row sums.
+    if perturb_fraction < 1.0:
+        n_cells = X.shape[0]
+        n_sel = int(round(perturb_fraction * n_cells))
+        rng = np.random.default_rng(random_state)
+        selected = np.zeros(n_cells, dtype=bool)
+        if n_sel > 0:
+            selected[rng.choice(n_cells, size=n_sel, replace=False)] = True
+        X[~selected] = X_orig[~selected]
+
     if renormalize:
         row_sums_after = np.asarray(X.sum(axis=1)).ravel()
         scale_rows = np.where(row_sums_after == 0, 1.0, row_sums_before / row_sums_after)
@@ -258,6 +277,8 @@ def _make_perturbed_expression(
     base: float = np.e,
     add_shift: bool = True,
     renormalize: bool = True,
+    perturb_fraction: float = 1.0,
+    random_state: int = 0,
 ):
     """Apply Node perturbations to ``adata.X`` and return the modified expression matrix."""
     var_names = list(adata.var_names)
@@ -278,6 +299,7 @@ def _make_perturbed_expression(
     return _node_perturbation(
         X, var_idx=var_idx, perturbations=perturbations,
         groupby=groupby, labels=labels, base=base, add_shift=add_shift, renormalize=renormalize,
+        perturb_fraction=perturb_fraction, random_state=random_state,
     )
 
 
@@ -334,6 +356,8 @@ def make_neighbor_perturbation(
     base: float = np.e,
     add_shift: bool = False,
     renormalize: bool = True,
+    perturb_fraction: float = 1.0,
+    random_state: int = 0,
 ) -> None:
     """
     Apply Node perturbations to neighbour expression and re-aggregate.
@@ -365,6 +389,17 @@ def make_neighbor_perturbation(
         Key in ``adata.obsm`` for the counterfactual spatial features.
     layer_key
         Key in ``adata.layers`` where the perturbed counts are stored.
+    perturb_fraction
+        Fraction of cells (in ``[0, 1]``) that receive the perturbation before
+        re-aggregation. A random subset of this size is perturbed and the rest are
+        left at their original expression, so after the degree-normalised neighbour
+        mean each focal cell effectively sees ``~perturb_fraction`` of its neighbours
+        perturbed. ``1.0`` (default) perturbs every cell (original behaviour); ``0.0``
+        applies no change. The subset is drawn over *all* cells, so when every cell
+        type is covered by ``perturbations`` this equals the fraction of perturbed
+        neighbours.
+    random_state
+        Seed for the random subset selection when ``perturb_fraction < 1``.
 
     Raises
     ------
@@ -385,6 +420,7 @@ def make_neighbor_perturbation(
     adata.layers[layer_key] = _make_perturbed_expression(
         adata, perturbations=perturbations, groupby=groupby,
         base=base, add_shift=add_shift, renormalize=renormalize,
+        perturb_fraction=perturb_fraction, random_state=random_state,
     )
 
     compute_spatial_features(

tests/test_cellina.py

@@ -656,6 +656,30 @@ def test_node_perturbation_row_sum_invariance():
         )
 
 
+def test_node_perturbation_fraction():
+    """perturb_fraction perturbs a seeded random subset of rows; 1.0 is the old behaviour."""
+    from cellina._spatial_utils import _node_perturbation
+
+    rng = np.random.default_rng(0)
+    X = rng.random((100, 5)).astype(np.float32)
+    var_idx = {f"gene{i}": i for i in range(5)}
+    pert = {"gene0": 1.0}
+
+    def run(**kw):
+        return _node_perturbation(X, var_idx, pert, add_shift=True, **kw).toarray()
+
+    np.testing.assert_array_equal(run(perturb_fraction=1.0), run())
+    np.testing.assert_array_equal(run(perturb_fraction=0.0), X)
+
+    half = run(perturb_fraction=0.5, random_state=1)
+    assert (half != X).any(axis=1).sum() == 50
+    np.testing.assert_array_equal(half, run(perturb_fraction=0.5, random_state=1))
+
+    for bad in (-0.1, 1.1):
+        with pytest.raises(ValueError):
+            _node_perturbation(X, var_idx, pert, perturb_fraction=bad)
+
+
 def test_make_neighbor_perturbation_unknown(adata_with_spatial):
     """Unknown cell-type key in perturbations raises ValueError."""
     import pandas as pd