diff --git a/src/cellina/_spatial_utils.py b/src/cellina/_spatial_utils.py index 395cbfe..dd419b8 100644 --- a/src/cellina/_spatial_utils.py +++ b/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,6 +239,10 @@ 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 @@ -243,6 +250,18 @@ def _encode(logfc): 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( diff --git a/tests/test_cellina.py b/tests/test_cellina.py index 1c26d77..bf1fe11 100644 --- a/tests/test_cellina.py +++ b/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