Fix single-cell counting semantics and harden numerical parity validation

CHANGELOG.md

@@ -5,6 +5,21 @@ All notable changes to WASP2 will be documented in this file.
 ## [Unreleased]
 
 ### Fixed
+- Compare validation count columns exactly before floating-point conversion,
+  including UInt64 values above 2**53. Distinguish archived narrow CSV input
+  from explicitly widened reference counts; retain strict legacy discrepancies.
+- Document Rust counting/bulk inference and Python/SciPy single-cell inference,
+  and correct the single-cell AnnData axes. Precision regressions cover balanced
+  likelihoods, overflow and conditional unphased symmetry.
+- Single-cell sparse counts and totals use UInt64 so H5AD output preserves values
+  above 65,535. Metadata conversion no longer requires an undeclared PyArrow install.
+- Automatic single-sample selection now applies the same heterozygote filtering
+  and feature reindexing as explicit selection, correcting inherited Python behavior.
+- Single-cell counting normalizes repeated, comma-separated and file-based sample options.
+- Shared-SNP differential analysis applies minimum depth per region, as archived
+  Python did. Unphased Python likelihoods preserve its sequential arithmetic while
+  retaining log-space fallback for underflow. Expanded comparisons explicitly widen
+  the archive's UInt16 accumulators; this also avoids incidental float32 dispatch.
 - Donor feature deduplication uses the parsed first GT allele, as Python does.
   Original GT tokens remain available for the backend's pre-dedup phasing
   decision; fitting weights, raw counts, QC and provenance reflect retained rows.
@@ -15,8 +30,8 @@ All notable changes to WASP2 will be documented in this file.
   checks, parses numeric first alleles consistently, and rejects malformed mixed
   phased input instead of silently changing a region's likelihood. Interior
   dispersion values are no longer clipped to the endpoint regularization range.
-- Full 32-run public ATAC comparisons execute archived Python in pinned Linux
-  CI, including CLI inference and donor routes. Heterogeneous donor checks cover
+- Full 32-run public ATAC comparisons execute archived Python on wide-integer
+  DataFrames in pinned Linux CI; narrow CSV precision is checked separately, including CLI inference and donor routes. Heterogeneous donor checks cover
   both dispersion models, phasing modes and pseudocounts with independent fits.
 - Container checks exercise reordered ATAC headers, region grouping, phased
   likelihoods and raw counts. Images carry package-version and source-revision

README.md

@@ -56,6 +56,11 @@ WASP2 installs three command-line entry points:
 - `wasp2-count`
 - `wasp2-analyze`
 
+The Python CLI coordinates compiled Rust counting and bulk statistical inference.
+Single-cell counting uses Rust; single-cell imbalance and differential inference
+currently use Python/SciPy. A Python command entry point does not change which
+backend performs the calculations.
+
 ## Quick Start
 
 > **Prerequisite — normalize your VCF first.** WASP2's allelic-imbalance model is biallelic.

docs/source/user_guide/analysis.rst

@@ -5,7 +5,14 @@ Overview
 --------
 
 ``wasp2-analyze`` runs allelic imbalance analysis on bulk or single-cell count
-outputs.
+outputs. The Python CLI calls Rust for bulk inference. Single-cell imbalance
+and differential inference currently run in Python/SciPy; single-cell counting
+uses Rust.
+
+Numerical compatibility retains the archived statistical procedures with wide
+counts and consistent precision. Untouched legacy file readers and explicitly
+widened references are evaluated separately; the latter do not claim to reproduce
+historical UInt16 overflow or mixed scalar/array precision.
 
 Commands:
 
@@ -171,16 +178,14 @@ sets the reference orientation; the first SNP is not symmetrized.
 - Beta-binomial LRT: :math:`\rho` is held at its **null-model MLE** while
   maximizing the alternative likelihood over :math:`\mu` (profile likelihood,
   df = 1). :math:`\rho` is not jointly re-estimated under :math:`H_1`.
-- Scalar dispersion is fit on :math:`(0,1)` with boundary values clamped for
-  numerical stability. The linear predictor retains the :math:`\pm 10` guard. The
-  linear model uses raw-depth standardization,
-  :math:`\rho(N)=\operatorname{expit}(d_1 + d_2(N-c)/s)`, and persists the fitted
-  center :math:`c` and positive scale :math:`s` with :math:`d_1` and :math:`d_2`.
-- FDR correction uses :func:`scipy.stats.false_discovery_control` with
-  ``method='bh'`` (Benjamini–Hochberg). ``scipy`` raises on NaN p-values,
-  but a hand-written BH loop using ``np.minimum.accumulate`` silently
-  propagates NaN through the cumulative minimum — always drop or impute
-  NaN p-values before BH correction.
+- Scalar dispersion is fit on :math:`(0,1)`; exact endpoints receive numerical
+  regularization. The linear fit uses pseudocounted depth and a :math:`\pm 10`
+  predictor guard during fitting. Stored coefficients preserve that predictor
+  using :math:`\rho(N)=\operatorname{expit}(d_1+d_2(N-c)/s)` with raw depth
+  :math:`N`, center :math:`c=-2\,\mathrm{pseudocount}` and scale :math:`s=1`.
+- Both backends use Benjamini–Hochberg FDR correction. Bulk uses the Rust
+  implementation; single-cell analysis uses
+  :func:`scipy.stats.false_discovery_control` with ``method='bh'``.
 
 See the upstream mapping-filter docs (:doc:`/methods/mapping_filter`) for the
 WASP2 canonical filter contract that applies to both remapping and counting

docs/source/user_guide/single_cell.rst

@@ -4,7 +4,15 @@ Single-Cell Analysis
 Overview
 --------
 
-WASP2 provides specialized tools for allele-specific analysis in single-cell RNA-seq (scRNA-seq) data. This guide covers the barcode file format requirements and single-cell-specific workflows.
+WASP2 supports barcode-based allele counting and group-level imbalance analysis.
+The validated workflow is scATAC-seq. Rust performs allele counting; Python/SciPy
+performs single-cell imbalance and differential inference through the Python CLI.
+The current counter does not implement UMI-aware molecule counting for scRNA-seq.
+
+Counting uses wide integer storage and inference uses wide count sums. Statistical
+compatibility follows the archived likelihood and testing procedures with consistent
+precision; it does not reproduce historical UInt16 overflow or mixed-precision
+rounding. The original and explicitly widened archive references are tested separately.
 
 Barcode File Format
 -------------------
@@ -300,12 +308,14 @@ The single-cell counting module outputs an AnnData (``.h5ad``) file containing:
 
 **Observations (obs):**
 
-* SNP information (chrom, pos, ref, alt)
-* Aggregate counts per SNP
+* Cell barcodes and cell-level annotations
 
 **Variables (var):**
 
-* Cell barcodes
+* SNP information (chrom, pos, ref, alt)
+* Aggregate counts per SNP
+
+Count layers have shape ``cells × SNPs``.
 
 **Unstructured (uns):**
 

scripts/compare_archive_analysis.py

@@ -9,6 +9,7 @@
 
 from __future__ import annotations
 
+import ast
 import contextlib
 import io
 import subprocess
@@ -27,11 +28,30 @@
 ROOT = Path(__file__).resolve().parents[1]
 ARCHIVE_REF = "41354fc"
 STAT_COLUMNS = ["null_ll", "alt_ll", "mu", "lrt", "pval", "fdr_pval"]
-COUNT_COLUMNS = ["ref_count", "alt_count", "N", "snp_count", "num_snps"]
-
-
-def load_archive(ref: str = ARCHIVE_REF) -> dict[str, types.ModuleType]:
-    """Load unmodified historical modules, restoring their legacy import names."""
+COUNT_COLUMNS = [
+    "ref_count",
+    "alt_count",
+    "N",
+    "snp_count",
+    "num_snps",
+    "num_snps_group1",
+    "num_snps_group2",
+]
+
+
+def load_archive(
+    ref: str = ARCHIVE_REF,
+    *,
+    widen_single_cell_sums: bool = False,
+    widen_bulk_csv_counts: bool = False,
+) -> dict[str, types.ModuleType]:
+    """Load historical modules with separately selectable count-width adapters.
+
+    The explicit dtype repair prevents count overflow and SciPy float32 dispatch.
+    Single-cell changes only sum accumulator keywords; bulk CSV changes only
+    the three count-column reader dtypes. Neither changes statistical algorithms.
+    The default executes the unmodified archive.
+    """
     names = ["as_analysis", "as_analysis_sc", "compare_ai"]
     saved = {name: sys.modules.get(name) for name in names}
     modules = {}
@@ -43,7 +63,60 @@ def load_archive(ref: str = ARCHIVE_REF) -> dict[str, types.ModuleType]:
             module = types.ModuleType(f"archive_{name}")
             modules[name] = module
             sys.modules[name] = module
-            exec(compile(source, f"{ref}:src/analysis/{name}.py", "exec"), module.__dict__)
+            tree = ast.parse(source)
+            if widen_bulk_csv_counts and name == "as_analysis":
+                changed = set()
+                function = next(
+                    node
+                    for node in tree.body
+                    if isinstance(node, ast.FunctionDef) and node.name == "get_imbalance"
+                )
+                for node in ast.walk(function):
+                    if not (
+                        isinstance(node, ast.Call)
+                        and isinstance(node.func, ast.Attribute)
+                        and isinstance(node.func.value, ast.Name)
+                        and node.func.value.id == "pd"
+                        and node.func.attr == "read_csv"
+                    ):
+                        continue
+                    for keyword in node.keywords:
+                        if keyword.arg != "dtype" or not isinstance(keyword.value, ast.Dict):
+                            continue
+                        for key, value in zip(
+                            keyword.value.keys, keyword.value.values, strict=True
+                        ):
+                            if (
+                                isinstance(key, ast.Constant)
+                                and key.value in {"ref_count", "alt_count", "other_count"}
+                                and isinstance(value, ast.Attribute)
+                                and isinstance(value.value, ast.Name)
+                                and value.value.id == "np"
+                                and value.attr == "uint16"
+                            ):
+                                value.attr = "uint64"
+                                changed.add(key.value)
+                if changed != {"ref_count", "alt_count", "other_count"}:
+                    raise ValueError("Cannot locate the three archived bulk CSV count dtypes")
+            if widen_single_cell_sums and name in {"as_analysis_sc", "compare_ai"}:
+                for node in ast.walk(tree):
+                    if not (
+                        isinstance(node, ast.Call)
+                        and isinstance(node.func, ast.Attribute)
+                        and node.func.attr == "sum"
+                    ):
+                        continue
+                    for keyword in node.keywords:
+                        value = keyword.value
+                        if (
+                            keyword.arg == "dtype"
+                            and isinstance(value, ast.Attribute)
+                            and isinstance(value.value, ast.Name)
+                            and value.value.id == "np"
+                            and value.attr == "uint16"
+                        ):
+                            value.attr = "uint64"
+            exec(compile(tree, f"{ref}:src/analysis/{name}.py", "exec"), module.__dict__)
     finally:
         for name, module in saved.items():
             if module is None:
@@ -219,12 +292,24 @@ def compare_frames(expected: pd.DataFrame, observed: pd.DataFrame, key: str = "r
         if column not in observed:
             failures.append(f"missing {column}")
             continue
-        left = expected[column].to_numpy(dtype=float)
-        right = observed[column].to_numpy(dtype=float)
-        deltas[column] = float(np.max(np.abs(left - right))) if len(left) else 0.0
         if column in COUNT_COLUMNS:
-            equal = np.array_equal(left, right)
+            # Python integers retain UInt64 values and signed/unsigned differences.
+            # Casting first to float64 hides one-count changes above 2**53.
+            left_counts = expected[column].tolist()
+            right_counts = observed[column].tolist()
+            valid = all(
+                pd.notna(a) and pd.notna(b) for a, b in zip(left_counts, right_counts, strict=True)
+            )
+            equal = valid and left_counts == right_counts
+            deltas[column] = (
+                max((abs(a - b) for a, b in zip(left_counts, right_counts, strict=True)), default=0)
+                if valid
+                else float("nan")
+            )
         else:
+            left = expected[column].to_numpy(dtype=float)
+            right = observed[column].to_numpy(dtype=float)
+            deltas[column] = float(np.max(np.abs(left - right))) if len(left) else 0.0
             equal = (
                 np.all(np.isfinite(left))
                 and np.all(np.isfinite(right))
@@ -242,6 +327,81 @@ def compare_frames(expected: pd.DataFrame, observed: pd.DataFrame, key: str = "r
     }
 
 
+def classify_reference_differences(
+    legacy: pd.DataFrame, widened: pd.DataFrame, observed: pd.DataFrame, key: str = "region"
+) -> dict:
+    """Retain strict failures and localize changes caused by a count-width adapter.
+
+    A complementary estimate is only a candidate: equivalence requires a separate
+    check of the actual anchored objective. This helper never alters an output.
+    """
+    strict = compare_frames(legacy, observed, key)
+    corrected = compare_frames(widened, observed, key)
+    frames = [
+        frame.set_index(key, drop=False).rename_axis(None) for frame in (legacy, widened, observed)
+    ]
+    if any(not frame.index.is_unique for frame in frames):
+        raise ValueError("Reference classification requires unique result identities")
+    rows = []
+    for identity in sorted(set().union(*(frame.index for frame in frames))):
+        present = [identity in frame.index for frame in frames]
+        if not all(present):
+            rows.append(
+                {
+                    "identity": identity,
+                    "stage": "tested_identity",
+                    "present": present,
+                    "reference_width_explains": present[1:] == [True, True]
+                    and corrected["status"] == "PASS",
+                }
+            )
+            continue
+        parts = [frame.loc[[identity]] for frame in frames]
+        comparison = compare_frames(parts[0], parts[2], key)
+        if comparison["status"] == "PASS":
+            continue
+        fields = comparison.get("failures", [])
+        stages = []
+        if any(field in COUNT_COLUMNS for field in fields):
+            stages.append("count_precision")
+        if any(field in fields for field in ("null_ll", "alt_ll", "lrt")):
+            stages.append("likelihood_precision")
+        effects = [field for field in ("mu", "mu1", "mu2", "combined_mu") if field in fields]
+        if effects:
+            stages.append("effect_estimate_precision")
+        if "pval" in fields:
+            stages.append("p_value_precision")
+        if "fdr_pval" in fields or any(field.startswith("significance at") for field in fields):
+            stages.append("multiple_testing")
+        candidates = [
+            field
+            for field in effects
+            if np.isclose(
+                float(parts[0][field].iloc[0]) + float(parts[2][field].iloc[0]),
+                1.0,
+                rtol=5e-5,
+                atol=5e-5,
+            )
+        ]
+        rows.append(
+            {
+                "identity": identity,
+                "stages": stages,
+                "fields": fields,
+                "complementary_estimate_candidates": candidates,
+                "reference_width_explains": compare_frames(parts[1], parts[2], key)["status"]
+                == "PASS",
+                "max_absolute_differences": comparison["max_absolute_differences"],
+            }
+        )
+    return {
+        "legacy_strict_comparison": strict,
+        "widened_comparison": corrected,
+        "mismatched_rows": rows,
+        "unexplained_rows": sum(not row["reference_width_explains"] for row in rows),
+    }
+
+
 def prepare_bulk_reference_counts(
     counts: pd.DataFrame,
     *,

src/analysis/__main__.py

@@ -306,9 +306,7 @@ def find_imbalance_sc(
     try:
         from .run_analysis_sc import run_ai_analysis_sc
     except (ImportError, ModuleNotFoundError) as error:
-        raise typer.BadParameter(
-            "Single-cell analysis requires: pip install 'wasp2[single-cell]'"
-        ) from error
+        raise typer.BadParameter("Single-cell analysis requires: pip install wasp2") from error
 
     run_ai_analysis_sc(
         count_file=counts,
@@ -404,9 +402,7 @@ def compare_imbalance(
     try:
         from .run_compare_ai import run_ai_comparison
     except (ImportError, ModuleNotFoundError) as error:
-        raise typer.BadParameter(
-            "Single-cell comparison requires: pip install 'wasp2[single-cell]'"
-        ) from error
+        raise typer.BadParameter("Single-cell comparison requires: pip install wasp2") from error
 
     run_ai_comparison(
         count_file=counts,

src/analysis/as_analysis.py

@@ -173,6 +173,19 @@ def opt_unphased_dp(
     # Get likelihood of first pos
     first_ll = opt_prob(prob, first_disp, first_ref[0], first_n[0])
 
+    # Preserve the archived sequential arithmetic while products remain normal.
+    # Equivalent log-space operations can select another symmetric optimum.
+    forward = np.atleast_1d(opt_prob(prob, phase_disp, phase_ref, phase_n, log=False))
+    reverse = np.atleast_1d(opt_prob(1 - prob, phase_disp, phase_ref, phase_n, log=False))
+    previous = 1.0
+    for p1, p2 in zip(forward, reverse, strict=True):
+        phase1, phase2 = previous * p1, previous * p2
+        previous = 0.5 * phase1 + 0.5 * phase2
+        if previous < np.finfo(float).tiny:
+            break
+    else:
+        return float(first_ll - np.log(previous))
+
     # Get LOG likelihood with regard to phasing of first pos.
     # opt_prob(..., log=True) returns the NEGATIVE log-pmf, so negate to get the
     # log-pmf itself for stable log-space accumulation.

src/analysis/as_analysis_sc.py

@@ -3,7 +3,7 @@
 Provides functions for analyzing allelic imbalance in single-cell data
 stored in AnnData format with SNP counts in layers.
 
-Orientation contract (scverse convention, see ISSUE2_ORIENTATION_DECISION §2b):
+Orientation contract (scverse convention):
 
     obs = cells (barcodes); var = SNPs (variants).
 
@@ -126,7 +126,7 @@ def get_imbalance_sc(
     n_counts = ref_counts + alt_counts
 
     # Calculate dispersion across dataset
-    def opt_disp(rho: float, ref_data: NDArray[np.uint16], n_data: NDArray[np.uint16]) -> float:
+    def opt_disp(rho: float, ref_data: NDArray[np.uint64], n_data: NDArray[np.uint64]) -> float:
         rho_safe = float(clamp_rho(rho))  # Prevent division by zero
         return float(
             -np.sum(