Pdex 0.3.0

.github/workflows/ci.yml

@@ -1,6 +1,9 @@
 name: "CI"
 
-on: [push, pull_request]
+on:
+  push:
+    branches: [main]
+  pull_request:
 
 jobs:
   all_jobs:
@@ -18,7 +21,7 @@ jobs:
 
   semver-check:
     runs-on: ubuntu-latest
-    if: github.event_name == 'pull_request'
+    if: github.event_name == 'pull_request' && github.base_ref == 'main'
 
     steps:
       - uses: actions/checkout@v4

CHANGELOG.md

@@ -0,0 +1,48 @@
+# Changelog
+
+All notable changes to this project are documented in this file.
+
+## [0.3.0] 
+
+### Removed
+
+- **Breaking:** the deprecated `fold_change` output column has been removed. It was an
+  alias for `log2_fold_change` (identical values); use `log2_fold_change` directly. The
+  `FutureWarning` previously emitted on every `pdex(...)` call is also gone.
+
+### Changed
+
+- `"all"` mode (1-vs-rest) is now a genuine one-shot computation instead of a per-group
+  loop: the expression matrix is materialized once, each gene is ranked once via
+  `mwu_one_vs_rest()`, and each group's "rest" pseudobulk/CPM is derived algebraically
+  from the global sum rather than by re-slicing and re-ranking a fresh "rest" matrix.
+  This turns `_pdex_all` from `O(n_groups × n_obs)` into ~`O(n_obs)`, which matters most
+  for screens with many groups (e.g. guides) and/or large cell counts.
+- Bumped the `numba-mwu` dependency floor to `>=0.2.0` (required by the `"all"`-mode
+  optimization above).
+
+### Fixed
+
+- Guarded against floating-point cancellation noise in the algebraic "rest" pseudobulk
+  and CPM derivation (`_pdex_all`): the rest mean of non-negative data can never be
+  legitimately negative, so it is now clipped to `>= 0` before feeding
+  `log2_fold_change`/`percent_change`, preventing spurious negative values from
+  floating-point noise.
+- `_available_cpus()` now checks `hasattr(os, "sched_getaffinity")` before calling it,
+  instead of relying on `AttributeError` from the call itself, fixing a case where the
+  attribute exists but raises for an unrelated reason on some platforms.
+
+### Tests
+
+- Added regression coverage for `"all"` mode with multiple groups and edge cases
+  (`tests/test_math.py`, `tests/test_pdex.py`, `tests/conftest.py`).
+
+### CI
+
+- `semver-check` now only runs on pull requests targeting `main` (previously it could
+  run — and fail spuriously — on unrelated PR bases).
+- CI workflows now avoid redundant duplicate runs on a single PR push.
+
+## [0.2.5] - previous release
+
+See git history prior to this file's introduction for earlier changes.

CLAUDE.md

@@ -73,15 +73,19 @@ checked empirically (inspect the per-gene CPM distribution).
 | File                   | Role                                                                                                    |
 | ---------------------- | ------------------------------------------------------------------------------------------------------- |
 | `src/pdex/__init__.py` | `pdex()` entry point and full pipeline logic                                                            |
-| `src/pdex/_math.py`    | Numba JIT-compiled `fold_change()`, `percent_change()`, and `mwu()` wrappers; `pseudobulk()` dispatcher; `cpm_bulk()` pooled-CPM view for the filter |
+| `src/pdex/_math.py`    | Numba JIT-compiled `log2_fold_change()`, `percent_change()`, and `mwu()`/`mwu_one_vs_rest()` wrappers over `numba-mwu`; `pseudobulk()` dispatcher; `cpm_bulk()` pooled-CPM view for the filter; the `bulk_matrix_pre_transform_mean()`/`pseudobulk_from_pre_mean()`/`cpm_from_gene_means()` trio powering `"all"` mode's global-sum optimization |
 | `src/pdex/_utils.py`   | `set_numba_threadpool()` — sets Numba thread count before JIT warmup; `_available_cpus()` — affinity-aware CPU count (respects cgroup/SLURM limits); `_detect_is_log1p()` heuristic |
 
 ### Performance Design
 
-- Numba JIT compilation accelerates per-cell/per-gene math (`fold_change`, `percent_change`, `_log1p_col_mean`, `_expm1_vec`)
-- `numba-mwu` (external dep) provides a Numba-accelerated Mann-Whitney U implementation
+- Numba JIT compilation accelerates per-cell/per-gene math (`log2_fold_change`, `percent_change`, `_log1p_col_mean`, `_expm1_vec`)
+- `numba-mwu` (external dep, `>=0.2.0`) provides Numba-accelerated Mann-Whitney U kernels for **both** the pairwise case (`mannwhitneyu_columns`/`mannwhitneyu_sparse`, used by `"ref"` and `"on_target"` modes) and the one-vs-rest case (`mannwhitneyu_one_vs_rest`/`_sparse`, used by `"all"` mode) — the one-vs-rest kernels originated in pdex and were upstreamed into `numba-mwu` since the optimization is domain-agnostic (see that package's CLAUDE.md for the algorithm).
 - Sparse CSR matrices are handled by reusing pre-computed non-targeting column indices to avoid redundant dense conversion
 - Parallelism is controlled via `threads` passed to `set_numba_threadpool()`
+- **`"all"` mode (1-vs-rest) is a one-shot computation, not a per-group loop over `numba-mwu`.** Because `group ∪ rest` is always the full (non-filtered) dataset regardless of which group is being tested, `_pdex_all` materializes the expression matrix exactly once (instead of once per group) and:
+  - Ranks each gene once across all cells via `mwu_one_vs_rest()` (`_math.py`, a thin wrapper over `numba_mwu.mannwhitneyu_one_vs_rest`/`_sparse`) and reduces to every group's rank-sum in the same pass, rather than re-ranking group+rest from scratch per group.
+  - Derives each group's "rest" pseudobulk and CPM algebraically as `(global_sum - group_sum) / n_rest` (`bulk_matrix_pre_transform_mean()`/`pseudobulk_from_pre_mean()`/`cpm_from_gene_means()`) instead of recomputing over a freshly sliced "rest" matrix.
+  - This turns `_pdex_all` from `O(n_groups × n_obs)` matrix I/O and ranking into ~`O(n_obs)` total, which matters most for screens with many groups (e.g. guides) and/or large cell counts.
 
 ### Output Schema
 
@@ -95,7 +99,6 @@ The returned Polars DataFrame (or pandas DataFrame when `as_pandas=True`) has co
 | `ref_mean`          | float | Pseudobulk mean for the reference, always in natural (count) space    |
 | `target_membership` | int   | Number of cells in the target group                                   |
 | `ref_membership`    | int   | Number of cells in the reference                                      |
-| `fold_change`       | float | **Deprecated** alias for `log2_fold_change` (identical values). Retained for one release; emits a `FutureWarning` on every `pdex(...)` call and will be removed in pdex 0.3.0. |
 | `log2_fold_change`  | float | log2((target_mean + epsilon) / (ref_mean + epsilon)) — computed from pseudobulk means. `epsilon` defaults to `1e-9` (finite-guard), so by default there are no `±inf`/`NaN`: one-sided zeros become large-but-finite and `0/0` is `0.0`. With `epsilon == 0`, `0/0` is still defined as `0.0` (not `NaN`) but one-sided zeros yield `±inf`. |
 | `percent_change`    | float | (target_mean - ref_mean) / (ref_mean + epsilon) — computed from pseudobulk means. With the default `epsilon=1e-9` there are no non-finite values; with `epsilon == 0`, `0/0` is `0.0` (not `NaN`) and a zero reference with nonzero target yields `+inf`. |
 | `p_value`           | float | Mann-Whitney U p-value (per-cell vectors)                             |

README.md

@@ -103,7 +103,6 @@ Returns a Polars DataFrame (or pandas if `as_pandas=True`) with one row per (gro
 | `ref_mean`          | Pseudobulk mean for the reference (count space)    |
 | `target_membership` | Number of cells in the target group                |
 | `ref_membership`    | Number of cells in the reference                   |
-| `fold_change`       | **Deprecated alias** for `log2_fold_change` (identical values). Will be removed in pdex 0.3.0. |
 | `log2_fold_change`  | log2((target_mean + epsilon) / (ref_mean + epsilon)). With default `epsilon=1e-9`, one-sided zeros are large finite values; with `epsilon=0.0`, one-sided zeros yield `±inf`. Genes unexpressed in both groups (0/0) report `0.0`, not `NaN`. |
 | `percent_change`    | (target_mean - ref_mean) / (ref_mean + epsilon). With default `epsilon=1e-9`, zero-reference cases are finite; with `epsilon=0.0`, a zero reference with nonzero target yields `+inf`. Genes unexpressed in both groups (0/0) report `0.0`, not `NaN`. |
 | `p_value`           | Mann-Whitney U p-value                             |

pyproject.toml

@@ -1,14 +1,14 @@
 [project]
 name = "pdex"
-version = "0.2.5"
+version = "0.3.0"
 description = "Parallel differential expression for single-cell perturbation sequencing"
 readme = "README.md"
 authors = [{ name = "noam teyssier", email = "noam.teyssier@arcinstitute.org" }]
 requires-python = ">=3.11"
 dependencies = [
     "anndata>=0.12.10",
     "numba>=0.64.0",
-    "numba-mwu>=0.1.1",
+    "numba-mwu>=0.2.0",
     "numpy>=2.4.2",
     "pandas>=2.3.3",
     "polars>=1.38.1",

src/pdex/__init__.py

@@ -14,11 +14,15 @@
 
 from pdex._math import (
     bulk_matrix_arithmetic,
+    bulk_matrix_pre_transform_mean,
     cpm_bulk,
+    cpm_from_gene_means,
     log2_fold_change,
     mwu,
+    mwu_one_vs_rest,
     percent_change,
     pseudobulk,
+    pseudobulk_from_pre_mean,
 )
 
 from ._utils import _detect_is_log1p, set_numba_threadpool
@@ -252,7 +256,7 @@ def pdex(
     epsilon:
         Pseudocount added to the **native (count-space) means** — the denominator
         (and, for ``log2_fold_change``, the numerator) — before computing
-        ``fold_change`` and ``percent_change``. It is never applied to the CPM view
+        ``log2_fold_change`` and ``percent_change``. It is never applied to the CPM view
         used by ``cpm_filter``. When ``epsilon > 0``, extreme values from near-zero
         reference means (scRNA-seq sparsity artifact) are dampened toward zero, and
         one-sided zeros become large-but-finite instead of ``±inf``. Has no effect on
@@ -303,7 +307,7 @@ def pdex(
     pl.DataFrame | pd.DataFrame
         One row per (group, feature) pair with columns: ``target``, ``feature``,
         ``target_mean``, ``ref_mean``, ``target_membership``, ``ref_membership``,
-        ``fold_change``, ``log2_fold_change``, ``percent_change``, ``p_value``,
+        ``log2_fold_change``, ``percent_change``, ``p_value``,
         ``statistic``, ``fdr``.
 
         ``target_mean`` and ``ref_mean`` are always in **natural (count) space**.
@@ -318,12 +322,8 @@ def pdex(
         and ``percent_change`` define this as ``0.0`` (no change) rather than
         ``NaN``. One-sided zeros still produce ``±inf``.
 
-        ``fold_change`` is a **deprecated** alias for ``log2_fold_change``
-        (identical values). It is retained for one release to ease migration
-        and will be removed in pdex 0.3.0. New code should read
-        ``log2_fold_change`` directly. A :class:`FutureWarning` is emitted
-        on every ``pdex(...)`` call.  The MWU ``p_value`` and
-        ``statistic`` are computed directly on the per-cell expression vectors.
+        The MWU ``p_value`` and ``statistic`` are computed directly on the
+        per-cell expression vectors.
 
         For ``mode="ref"``, the reference group itself is excluded from the output.
 
@@ -341,14 +341,6 @@ def pdex(
     if epsilon < 0:
         raise ValueError(f"epsilon must be non-negative, got {epsilon}")
 
-    warnings.warn(
-        "The `fold_change` column in pdex output is deprecated and will be "
-        "removed in pdex 0.3.0. Use `log2_fold_change` instead — it contains "
-        "the same values (`log2(target_mean / ref_mean)`).",
-        FutureWarning,
-        stacklevel=2,
-    )
-
     # Set the global threadpool for numba
     set_numba_threadpool(threads)
 
@@ -498,7 +490,6 @@ def _assemble_group_frame(
             "ref_mean": ref_mean,
             "target_membership": np.full(feature.shape[0], target_membership),
             "ref_membership": np.full(feature.shape[0], ref_membership),
-            "fold_change": lfc,
             "log2_fold_change": lfc,
             "percent_change": pc,
             "p_value": pvalue,
@@ -594,43 +585,79 @@ def _pdex_all(
     cpm_filter: float | None = None,
 ) -> pl.DataFrame:
     unique_groups, unique_group_indices = _unique_groups(adata.obs, groupby)
-    log.info("Found %d groups for 1-vs-rest comparison", len(unique_groups))
+    n_groups = len(unique_groups)
+    log.info("Found %d groups for 1-vs-rest comparison", n_groups)
+
+    if n_groups < 2:
+        raise ValueError(f"mode='all' requires at least 2 groups, found {n_groups}")
 
     feature_names = adata.var_names
 
+    # group ∪ rest is always the full (non-filtered) dataset in "all" mode, so the
+    # matrix is materialized once here rather than once per group (see CLAUDE.md).
+    valid_mask = np.flatnonzero(unique_group_indices >= 0)
+    codes_valid = unique_group_indices[valid_mask]
+    n_valid = valid_mask.size
+
+    global_matrix = _isolate_matrix(adata, valid_mask)
+
+    # One-shot 1-vs-rest MWU: each gene is ranked once, not once per group.
+    mwu_result = mwu_one_vs_rest(global_matrix, codes_valid, n_groups)
+    all_statistic = mwu_result.statistic
+    all_pvalue = np.asarray(mwu_result.pvalue).clip(0, 1)
+
+    # Pseudobulk pre-transform mean, computed once; "rest" is derived per group via
+    # (global_sum - group_sum) / n_rest rather than pseudobulk() on a fresh rest slice.
+    global_pre_mean = bulk_matrix_pre_transform_mean(
+        global_matrix, geometric_mean=geometric_mean, is_log1p=is_log1p
+    )
+    global_pre_sum = global_pre_mean * n_valid
+
+    if cpm_filter is not None:
+        global_arith_mean = bulk_matrix_arithmetic(global_matrix, is_log1p)
+        global_arith_sum = global_arith_mean * n_valid
+
     results = []
     for group_idx in tqdm(
-        range(len(unique_groups)),
+        range(n_groups),
         desc="Running parallel differential expression (1 vs Rest)",
     ):
         group_name = unique_groups[group_idx]
 
-        group_mask = np.flatnonzero(unique_group_indices == group_idx)
-        rest_mask = np.flatnonzero(
-            (unique_group_indices != group_idx) & (unique_group_indices >= 0)
-        )
+        local_group_mask = np.flatnonzero(codes_valid == group_idx)
+        n_group = local_group_mask.size
+        n_rest = n_valid - n_group
 
-        group_matrix = _isolate_matrix(adata, group_mask)
-        rest_matrix = _isolate_matrix(adata, rest_mask)
+        group_matrix = global_matrix[local_group_mask]
 
-        group_bulk = pseudobulk(
+        group_pre_mean = bulk_matrix_pre_transform_mean(
             group_matrix, geometric_mean=geometric_mean, is_log1p=is_log1p
         )
-        rest_bulk = pseudobulk(
-            rest_matrix, geometric_mean=geometric_mean, is_log1p=is_log1p
+        # Clip: rest_pre_mean is a mean of non-negative data, so it can never be
+        # legitimately negative. This guards against floating-point noise which can
+        # otherwise corrupt log2_fold_change/percent_change.
+        rest_pre_mean = np.clip(
+            (global_pre_sum - group_pre_mean * n_group) / n_rest, 0, None
         )
 
+        group_bulk = pseudobulk_from_pre_mean(group_pre_mean, geometric_mean)
+        rest_bulk = pseudobulk_from_pre_mean(rest_pre_mean, geometric_mean)
+
         lfc = log2_fold_change(group_bulk, rest_bulk, epsilon)
         pc = percent_change(group_bulk, rest_bulk, epsilon)
-        mwu_result = mwu(group_matrix, rest_matrix)
 
-        mwu_statistic = mwu_result.statistic
-        mwu_pvalue = np.asarray(mwu_result.pvalue).clip(0, 1)
+        mwu_statistic = all_statistic[group_idx]
+        mwu_pvalue = all_pvalue[group_idx]
 
         if cpm_filter is None:
             keep = None
         else:
-            rest_cpm = cpm_bulk(rest_matrix, is_log1p)
+            group_arith_mean = bulk_matrix_arithmetic(group_matrix, is_log1p)
+            # Same floating-point-cancellation guard as rest_pre_mean above.
+            rest_arith_mean = np.clip(
+                (global_arith_sum - group_arith_mean * n_group) / n_rest, 0, None
+            )
+            rest_cpm = cpm_from_gene_means(rest_arith_mean)
             keep = _cpm_keep_mask(group_matrix, rest_cpm, is_log1p, cpm_filter)
 
         results.append(
@@ -639,8 +666,8 @@ def _pdex_all(
                 feature_names=feature_names,
                 target_bulk=group_bulk,
                 ref_bulk=rest_bulk,
-                target_membership=group_mask.size,
-                ref_membership=rest_mask.size,
+                target_membership=n_group,
+                ref_membership=n_rest,
                 lfc=lfc,
                 pc=pc,
                 pvalue=mwu_pvalue,
@@ -750,7 +777,6 @@ def _pdex_on_target(
                 "ref_mean": ref_mean,
                 "target_membership": group_mask.size,
                 "ref_membership": ref_membership,
-                "fold_change": lfc,
                 "log2_fold_change": lfc,
                 "percent_change": pc,
                 "p_value": p_value,
@@ -769,7 +795,6 @@ def _pdex_on_target(
                 "ref_mean": pl.Float64,
                 "target_membership": pl.Int64,
                 "ref_membership": pl.Int64,
-                "fold_change": pl.Float64,
                 "log2_fold_change": pl.Float64,
                 "percent_change": pl.Float64,
                 "p_value": pl.Float64,