numba-mwu 0.1.1__tar.gz

numba_mwu-0.1.1/PKG-INFO

@@ -0,0 +1,170 @@
+Metadata-Version: 2.3
+Name: numba-mwu
+Version: 0.1.1
+Summary: Numba-accelerated Mann-Whitney U test with sparse matrix support.
+Author: Noam Teyssier
+Author-email: Noam Teyssier <22600644+noamteyssier@users.noreply.github.com>
+Requires-Dist: numba>=0.64.0
+Requires-Dist: scipy>=1.17.1
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# numba-mwu
+
+Numba-accelerated Mann-Whitney U test.
+Drop-in replacement for `scipy.stats.mannwhitneyu` with parallel batch operations and native sparse matrix support.
+
+All functions use the asymptotic (normal approximation) method and produce results identical to `scipy.stats.mannwhitneyu(..., method="asymptotic")`.
+
+> Note: This is only supported for 1D and 2D inputs.
+
+## Installation
+
+```bash
+uv pip install numba-mwu
+```
+
+## API
+
+Every function returns a `MannWhitneyUResult` named tuple with `statistic` and `pvalue` fields. The batch functions return arrays instead of scalars.
+
+All functions accept `use_continuity` (default `True`) and `alternative` (`"two-sided"`, `"less"`, `"greater"`).
+
+### `mannwhitneyu(x, y)`
+
+Single two-sample test. Equivalent to scipy's `mannwhitneyu`.
+
+```python
+from numba_mwu import mannwhitneyu
+
+result = mannwhitneyu(x, y)
+result.statistic  # U statistic
+result.pvalue     # two-sided p-value
+```
+
+### `mannwhitneyu_rows(X, y)`
+
+Test each row of a 2-D array `X` against a shared reference sample `y`.
+Parallelized across rows.
+
+```python
+from numba_mwu import mannwhitneyu_rows
+
+# X: (n_tests, n1), y: (n2,)
+result = mannwhitneyu_rows(X, y)
+result.statistic  # shape (n_tests,)
+result.pvalue     # shape (n_tests,)
+```
+
+### `mannwhitneyu_columns(X, Y)`
+
+Test each column of `X` against the corresponding column of `Y`.
+Parallelized across columns.
+Designed for the common case of slicing a cells-by-genes matrix into two groups:
+
+```python
+from numba_mwu import mannwhitneyu_columns
+
+# expression: (n_cells, n_genes), labels: (n_cells,)
+X = expression[labels == "A"]  # (n1, n_genes)
+Y = expression[labels == "B"]  # (n2, n_genes)
+
+result = mannwhitneyu_columns(X, Y)
+result.statistic  # shape (n_genes,)
+result.pvalue     # shape (n_genes,)
+```
+
+### `mannwhitneyu_sparse(X, Y)`
+
+Same as `mannwhitneyu_columns` but operates directly on CSR sparse matrices without converting to dense.
+
+Memory overhead per matrix is one `int64` array of length `nnz` (column permutation) plus one `int64` array of length `n_genes + 1` (column pointers).
+No data values are copied.
+
+Requires non-negative data (raw counts, normalized expression, etc.).
+
+> Note: Call `eliminate_zeros()` on each matrix beforehand if it may contain explicitly stored zeros.
+
+```python
+from numba_mwu import mannwhitneyu_sparse
+
+# adata.X is a CSR matrix, adata.obs["group"] has labels
+mask = adata.obs["group"] == "A"
+X = adata.X[mask]    # CSR row-slice is still CSR
+Y = adata.X[~mask]
+
+
+result = mannwhitneyu_sparse(X, Y)
+result.statistic  # shape (n_genes,)
+result.pvalue     # shape (n_genes,)
+```
+
+## Benchmarks
+
+Run benchmarks with:
+
+```bash
+uv run benchmarks/bench_mwu.py
+```
+
+```text
+================================================================================
+SINGLE PAIR BENCHMARKS (overhead comparison)
+================================================================================
+
+--- integer data ---
+scenario                            scipy        numba    speedup
+-----------------------------------------------------------------
+n=20 vs n=20                     223.1 us       3.9 us      56.9x
+n=100 vs n=100                   224.0 us       5.4 us      41.7x
+n=500 vs n=500                   248.3 us      12.6 us      19.7x
+n=1000 vs n=1000                 287.2 us      22.7 us      12.7x
+
+--- float data ---
+scenario                            scipy        numba    speedup
+-----------------------------------------------------------------
+n=20 vs n=20                     212.6 us       3.9 us      53.9x
+n=100 vs n=100                   220.7 us       5.6 us      39.4x
+n=500 vs n=500                   249.4 us      14.7 us      16.9x
+n=1000 vs n=1000                 287.3 us      27.4 us      10.5x
+
+================================================================================
+DENSE MATRIX BENCHMARKS
+================================================================================
+
+--- integer data ---
+scenario                            scipy        numba    speedup
+-----------------------------------------------------------------
+small (100x50)                    11.4 ms      64.1 us     177.8x
+medium (1000x500)                139.5 ms       1.5 ms      94.0x
+large (5000x2000)                 1.01  s      43.7 ms      23.0x
+xlarge (10000x5000)               3.93  s     179.5 ms      21.9x
+
+--- float data ---
+scenario                            scipy        numba    speedup
+-----------------------------------------------------------------
+small (100x50)                    11.1 ms      53.0 us     208.5x
+medium (1000x500)                131.5 ms       1.2 ms     109.1x
+large (5000x2000)                866.6 ms      36.0 ms      24.1x
+xlarge (10000x5000)               3.33  s     151.9 ms      22.0x
+
+================================================================================
+SPARSE MATRIX BENCHMARKS
+================================================================================
+
+--- integer data ---
+scenario                      scipy (dense)   numba sparse    numba dense   sp speedup
+-------------------------------------------------------------------------------------
+small 90% (200x100)                 22.7 ms        51.3 us        84.3 us       442.3x
+medium 90% (2000x1000)             275.5 ms         1.0 ms         3.5 ms       266.9x
+large 95% (5000x2000)              746.8 ms         2.6 ms        20.4 ms       282.1x
+xlarge 95% (10000x5000)             2.80  s        21.1 ms       117.2 ms       132.6x
+
+--- float data ---
+scenario                      scipy (dense)   numba sparse    numba dense   sp speedup
+-------------------------------------------------------------------------------------
+small 90% (200x100)                 22.7 ms        53.2 us        80.7 us       427.0x
+medium 90% (2000x1000)             279.5 ms         1.0 ms         4.3 ms       268.9x
+large 95% (5000x2000)              741.1 ms         3.5 ms        23.7 ms       209.4x
+xlarge 95% (10000x5000)             2.80  s        21.0 ms       111.5 ms       133.0x
+```

numba_mwu-0.1.1/README.md

@@ -0,0 +1,159 @@
+# numba-mwu
+
+Numba-accelerated Mann-Whitney U test.
+Drop-in replacement for `scipy.stats.mannwhitneyu` with parallel batch operations and native sparse matrix support.
+
+All functions use the asymptotic (normal approximation) method and produce results identical to `scipy.stats.mannwhitneyu(..., method="asymptotic")`.
+
+> Note: This is only supported for 1D and 2D inputs.
+
+## Installation
+
+```bash
+uv pip install numba-mwu
+```
+
+## API
+
+Every function returns a `MannWhitneyUResult` named tuple with `statistic` and `pvalue` fields. The batch functions return arrays instead of scalars.
+
+All functions accept `use_continuity` (default `True`) and `alternative` (`"two-sided"`, `"less"`, `"greater"`).
+
+### `mannwhitneyu(x, y)`
+
+Single two-sample test. Equivalent to scipy's `mannwhitneyu`.
+
+```python
+from numba_mwu import mannwhitneyu
+
+result = mannwhitneyu(x, y)
+result.statistic  # U statistic
+result.pvalue     # two-sided p-value
+```
+
+### `mannwhitneyu_rows(X, y)`
+
+Test each row of a 2-D array `X` against a shared reference sample `y`.
+Parallelized across rows.
+
+```python
+from numba_mwu import mannwhitneyu_rows
+
+# X: (n_tests, n1), y: (n2,)
+result = mannwhitneyu_rows(X, y)
+result.statistic  # shape (n_tests,)
+result.pvalue     # shape (n_tests,)
+```
+
+### `mannwhitneyu_columns(X, Y)`
+
+Test each column of `X` against the corresponding column of `Y`.
+Parallelized across columns.
+Designed for the common case of slicing a cells-by-genes matrix into two groups:
+
+```python
+from numba_mwu import mannwhitneyu_columns
+
+# expression: (n_cells, n_genes), labels: (n_cells,)
+X = expression[labels == "A"]  # (n1, n_genes)
+Y = expression[labels == "B"]  # (n2, n_genes)
+
+result = mannwhitneyu_columns(X, Y)
+result.statistic  # shape (n_genes,)
+result.pvalue     # shape (n_genes,)
+```
+
+### `mannwhitneyu_sparse(X, Y)`
+
+Same as `mannwhitneyu_columns` but operates directly on CSR sparse matrices without converting to dense.
+
+Memory overhead per matrix is one `int64` array of length `nnz` (column permutation) plus one `int64` array of length `n_genes + 1` (column pointers).
+No data values are copied.
+
+Requires non-negative data (raw counts, normalized expression, etc.).
+
+> Note: Call `eliminate_zeros()` on each matrix beforehand if it may contain explicitly stored zeros.
+
+```python
+from numba_mwu import mannwhitneyu_sparse
+
+# adata.X is a CSR matrix, adata.obs["group"] has labels
+mask = adata.obs["group"] == "A"
+X = adata.X[mask]    # CSR row-slice is still CSR
+Y = adata.X[~mask]
+
+
+result = mannwhitneyu_sparse(X, Y)
+result.statistic  # shape (n_genes,)
+result.pvalue     # shape (n_genes,)
+```
+
+## Benchmarks
+
+Run benchmarks with:
+
+```bash
+uv run benchmarks/bench_mwu.py
+```
+
+```text
+================================================================================
+SINGLE PAIR BENCHMARKS (overhead comparison)
+================================================================================
+
+--- integer data ---
+scenario                            scipy        numba    speedup
+-----------------------------------------------------------------
+n=20 vs n=20                     223.1 us       3.9 us      56.9x
+n=100 vs n=100                   224.0 us       5.4 us      41.7x
+n=500 vs n=500                   248.3 us      12.6 us      19.7x
+n=1000 vs n=1000                 287.2 us      22.7 us      12.7x
+
+--- float data ---
+scenario                            scipy        numba    speedup
+-----------------------------------------------------------------
+n=20 vs n=20                     212.6 us       3.9 us      53.9x
+n=100 vs n=100                   220.7 us       5.6 us      39.4x
+n=500 vs n=500                   249.4 us      14.7 us      16.9x
+n=1000 vs n=1000                 287.3 us      27.4 us      10.5x
+
+================================================================================
+DENSE MATRIX BENCHMARKS
+================================================================================
+
+--- integer data ---
+scenario                            scipy        numba    speedup
+-----------------------------------------------------------------
+small (100x50)                    11.4 ms      64.1 us     177.8x
+medium (1000x500)                139.5 ms       1.5 ms      94.0x
+large (5000x2000)                 1.01  s      43.7 ms      23.0x
+xlarge (10000x5000)               3.93  s     179.5 ms      21.9x
+
+--- float data ---
+scenario                            scipy        numba    speedup
+-----------------------------------------------------------------
+small (100x50)                    11.1 ms      53.0 us     208.5x
+medium (1000x500)                131.5 ms       1.2 ms     109.1x
+large (5000x2000)                866.6 ms      36.0 ms      24.1x
+xlarge (10000x5000)               3.33  s     151.9 ms      22.0x
+
+================================================================================
+SPARSE MATRIX BENCHMARKS
+================================================================================
+
+--- integer data ---
+scenario                      scipy (dense)   numba sparse    numba dense   sp speedup
+-------------------------------------------------------------------------------------
+small 90% (200x100)                 22.7 ms        51.3 us        84.3 us       442.3x
+medium 90% (2000x1000)             275.5 ms         1.0 ms         3.5 ms       266.9x
+large 95% (5000x2000)              746.8 ms         2.6 ms        20.4 ms       282.1x
+xlarge 95% (10000x5000)             2.80  s        21.1 ms       117.2 ms       132.6x
+
+--- float data ---
+scenario                      scipy (dense)   numba sparse    numba dense   sp speedup
+-------------------------------------------------------------------------------------
+small 90% (200x100)                 22.7 ms        53.2 us        80.7 us       427.0x
+medium 90% (2000x1000)             279.5 ms         1.0 ms         4.3 ms       268.9x
+large 95% (5000x2000)              741.1 ms         3.5 ms        23.7 ms       209.4x
+xlarge 95% (10000x5000)             2.80  s        21.0 ms       111.5 ms       133.0x
+```

numba_mwu-0.1.1/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "numba-mwu"
+version = "0.1.1"
+description = "Numba-accelerated Mann-Whitney U test with sparse matrix support."
+readme = "README.md"
+authors = [
+    { name = "Noam Teyssier", email = "22600644+noamteyssier@users.noreply.github.com" }
+]
+requires-python = ">=3.11"
+dependencies = [
+    "numba>=0.64.0",
+    "scipy>=1.17.1",
+]
+
+[build-system]
+requires = ["uv_build>=0.9.22,<0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "marimo>=0.20.2",
+    "pytest>=9.0.2",
+    "ruff>=0.15.2",
+    "ty>=0.0.18",
+]

numba_mwu-0.1.1/src/numba_mwu/__init__.py

@@ -0,0 +1,281 @@
+"""Numba-accelerated Mann-Whitney U test."""
+
+from collections import namedtuple
+
+import numpy as np
+
+from ._batch import _mannwhitneyu_columns, _mannwhitneyu_rows
+from ._core import GREATER, LESS, TWO_SIDED, _mannwhitneyu_single
+from ._sparse import _build_col_index, _sparse_mwu_batch
+
+__all__ = [
+    "MannWhitneyUResult",
+    "SparseColumnIndex",
+    "mannwhitneyu",
+    "mannwhitneyu_rows",
+    "mannwhitneyu_columns",
+    "mannwhitneyu_sparse",
+    "sparse_column_index",
+]
+
+MannWhitneyUResult = namedtuple("MannWhitneyUResult", ("statistic", "pvalue"))
+SparseColumnIndex = namedtuple(
+    "SparseColumnIndex", ("data", "col_indptr", "col_order", "n_rows", "n_cols")
+)
+
+_ALTERNATIVE_MAP = {
+    "two-sided": TWO_SIDED,
+    "less": LESS,
+    "greater": GREATER,
+}
+
+
+def _validate_alternative(alternative):
+    alt = alternative.lower()
+    if alt not in _ALTERNATIVE_MAP:
+        raise ValueError(
+            f"`alternative` must be one of {set(_ALTERNATIVE_MAP)}, got {alternative!r}"
+        )
+    return _ALTERNATIVE_MAP[alt]
+
+
+def _validate_1d(arr, name):
+    arr = np.asarray(arr, dtype=np.float64)
+    if arr.ndim != 1:
+        raise ValueError(f"`{name}` must be 1-dimensional, got ndim={arr.ndim}")
+    if arr.shape[0] == 0:
+        raise ValueError(f"`{name}` must be of nonzero size.")
+    if np.any(np.isnan(arr)):
+        raise ValueError(f"`{name}` must not contain NaNs.")
+    return arr
+
+
+def mannwhitneyu(x, y, use_continuity=True, alternative="two-sided"):
+    """Perform the Mann-Whitney U rank test on two independent samples.
+
+    Numba-accelerated asymptotic implementation equivalent to
+    ``scipy.stats.mannwhitneyu(..., method='asymptotic')``.
+
+    Parameters
+    ----------
+    x, y : array_like
+        1-D arrays of samples.
+    use_continuity : bool, optional
+        Whether a continuity correction (1/2) should be applied. Default True.
+    alternative : {'two-sided', 'less', 'greater'}, optional
+        Defines the alternative hypothesis. Default is 'two-sided'.
+
+    Returns
+    -------
+    result : MannWhitneyUResult
+        Named tuple with ``statistic`` (U for sample x) and ``pvalue``.
+    """
+    x = _validate_1d(x, "x")
+    y = _validate_1d(y, "y")
+    alt = _validate_alternative(alternative)
+    stat, pval = _mannwhitneyu_single(x, y, use_continuity, alt)
+    return MannWhitneyUResult(stat, pval)
+
+
+def mannwhitneyu_rows(X, y, use_continuity=True, alternative="two-sided"):
+    """Run Mann-Whitney U test for each row of X against y (parallelized).
+
+    Parameters
+    ----------
+    X : array_like, shape (n_tests, n1)
+        2-D array where each row is a sample to test against y.
+    y : array_like, shape (n2,)
+        1-D reference sample.
+    use_continuity : bool, optional
+        Whether a continuity correction (1/2) should be applied. Default True.
+    alternative : {'two-sided', 'less', 'greater'}, optional
+        Defines the alternative hypothesis. Default is 'two-sided'.
+
+    Returns
+    -------
+    result : MannWhitneyUResult
+        Named tuple with ``statistic`` and ``pvalue`` arrays of shape (n_tests,).
+    """
+    X = np.asarray(X, dtype=np.float64)
+    if X.ndim != 2:
+        raise ValueError(f"`X` must be 2-dimensional, got ndim={X.ndim}")
+    if X.shape[0] == 0:
+        raise ValueError("`X` must have at least one row.")
+    if np.any(np.isnan(X)):
+        raise ValueError("`X` must not contain NaNs.")
+    y = _validate_1d(y, "y")
+    alt = _validate_alternative(alternative)
+    stats, pvals = _mannwhitneyu_rows(X, y, use_continuity, alt)
+    return MannWhitneyUResult(stats, pvals)
+
+
+def _validate_2d(arr, name):
+    arr = np.asarray(arr, dtype=np.float64)
+    if arr.ndim != 2:
+        raise ValueError(f"`{name}` must be 2-dimensional, got ndim={arr.ndim}")
+    if arr.shape[0] == 0:
+        raise ValueError(f"`{name}` must have at least one row.")
+    if np.any(np.isnan(arr)):
+        raise ValueError(f"`{name}` must not contain NaNs.")
+    return arr
+
+
+def mannwhitneyu_columns(X, Y, use_continuity=True, alternative="two-sided"):
+    """Run Mann-Whitney U test on each column of X vs corresponding column of Y.
+
+    Designed for the common workflow where a matrix has been sliced by
+    group membership into two separate matrices (or views).
+
+    Parameters
+    ----------
+    X : array_like, shape (n1, n_genes)
+        2-D array for group A (one column per gene/feature).
+    Y : array_like, shape (n2, n_genes)
+        2-D array for group B. Must have the same number of columns as X.
+    use_continuity : bool, optional
+        Whether a continuity correction (1/2) should be applied. Default True.
+    alternative : {'two-sided', 'less', 'greater'}, optional
+        Defines the alternative hypothesis. Default is 'two-sided'.
+
+    Returns
+    -------
+    result : MannWhitneyUResult
+        Named tuple with ``statistic`` and ``pvalue`` arrays of shape (n_genes,).
+    """
+    X = _validate_2d(X, "X")
+    Y = _validate_2d(Y, "Y")
+    if X.shape[1] != Y.shape[1]:
+        raise ValueError(
+            f"`X` and `Y` must have the same number of columns, "
+            f"got {X.shape[1]} and {Y.shape[1]}."
+        )
+    alt = _validate_alternative(alternative)
+    stats, pvals = _mannwhitneyu_columns(X, Y, use_continuity, alt)
+    return MannWhitneyUResult(stats, pvals)
+
+
+def _validate_csr(X, name):
+    from scipy.sparse import issparse, isspmatrix_csr
+
+    if not issparse(X):
+        raise TypeError(f"`{name}` must be a scipy sparse matrix.")
+    if not (isspmatrix_csr(X) or X.format == "csr"):
+        raise TypeError(
+            f"`{name}` must be in CSR format. Convert with `{name}.tocsr()` if needed."
+        )
+    if X.data.size > 0 and X.data.min() < 0:
+        raise ValueError(
+            f"Sparse MWU requires non-negative data in `{name}`. "
+            "For data with negative values, convert to dense and use "
+            "mannwhitneyu_columns."
+        )
+    return X
+
+
+def sparse_column_index(X):
+    """Precompute a column index for a CSR sparse matrix.
+
+    The returned ``SparseColumnIndex`` can be passed to
+    ``mannwhitneyu_sparse`` in place of a raw CSR matrix, avoiding
+    redundant index construction when the same matrix is reused
+    across many comparisons::
+
+        ref_idx = sparse_column_index(ref_matrix)
+        for label in labels:
+            result = mannwhitneyu_sparse(group_matrix, ref_idx)
+
+    Parameters
+    ----------
+    X : scipy.sparse.csr_matrix or csr_array
+        Sparse matrix in CSR format with non-negative values.
+        Call ``X.eliminate_zeros()`` beforehand if it may contain
+        explicitly stored zeros.
+
+    Returns
+    -------
+    SparseColumnIndex
+        Precomputed index that can be passed to ``mannwhitneyu_sparse``.
+    """
+    X = _validate_csr(X, "X")
+    if X.shape[0] == 0:
+        raise ValueError("`X` must have at least one row.")
+    data = np.ascontiguousarray(X.data, dtype=np.float64)
+    indptr = np.ascontiguousarray(X.indptr)
+    indices = np.ascontiguousarray(X.indices)
+    col_indptr, col_order = _build_col_index(indptr, indices, X.shape[1])
+    return SparseColumnIndex(data, col_indptr, col_order, X.shape[0], X.shape[1])
+
+
+def _resolve_sparse(arg, name):
+    """Convert a CSR matrix or SparseColumnIndex to a SparseColumnIndex."""
+    if isinstance(arg, SparseColumnIndex):
+        return arg
+    return sparse_column_index(arg)
+
+
+def mannwhitneyu_sparse(X, Y, use_continuity=True, alternative="two-sided"):
+    """Run Mann-Whitney U test for each gene (column) of two sparse matrices.
+
+    Designed for single-cell expression matrices where rows are cells and
+    columns are genes. Each matrix represents one group — typically sliced
+    from a full expression matrix by cell labels::
+
+        X = full_matrix[labels == "A"]  # CSR row-slice is still CSR
+        Y = full_matrix[labels == "B"]
+
+    Works directly with CSR format without converting to CSC or dense.
+    The only allocation per matrix is a column-index permutation array
+    (one int per nonzero entry) and column pointers (one int per gene).
+
+    Requires non-negative data — zeros must be the smallest values so they
+    form a contiguous block at the start of the sorted order. This holds
+    for raw counts, normalized expression, and any non-negative transformation.
+
+    Both ``X`` and ``Y`` can be either a CSR matrix or a precomputed
+    ``SparseColumnIndex`` (from ``sparse_column_index``). Precomputing
+    is useful when the same matrix is compared against many others.
+
+    Parameters
+    ----------
+    X : csr_matrix, csr_array, or SparseColumnIndex, shape (n1, n_genes)
+        Sparse expression matrix for group A. If a CSR matrix, must have
+        non-negative values. Call ``X.eliminate_zeros()`` beforehand if
+        the matrix may contain explicitly stored zeros.
+    Y : csr_matrix, csr_array, or SparseColumnIndex, shape (n2, n_genes)
+        Sparse expression matrix for group B. Must have the same number
+        of columns as X.
+    use_continuity : bool, optional
+        Whether to apply continuity correction. Default True.
+    alternative : {'two-sided', 'less', 'greater'}, optional
+        Alternative hypothesis. Default 'two-sided'.
+
+    Returns
+    -------
+    result : MannWhitneyUResult
+        Named tuple with ``statistic`` and ``pvalue`` arrays of shape
+        (n_genes,).
+    """
+    idx_a = _resolve_sparse(X, "X")
+    idx_b = _resolve_sparse(Y, "Y")
+
+    if idx_a.n_cols != idx_b.n_cols:
+        raise ValueError(
+            f"`X` and `Y` must have the same number of columns, "
+            f"got {idx_a.n_cols} and {idx_b.n_cols}."
+        )
+
+    alt = _validate_alternative(alternative)
+
+    stats, pvals = _sparse_mwu_batch(
+        idx_a.data,
+        idx_a.col_indptr,
+        idx_a.col_order,
+        idx_a.n_rows,
+        idx_b.data,
+        idx_b.col_indptr,
+        idx_b.col_order,
+        idx_b.n_rows,
+        use_continuity,
+        alt,
+    )
+    return MannWhitneyUResult(stats, pvals)

numba_mwu-0.1.1/src/numba_mwu/_batch.py

@@ -0,0 +1,60 @@
+"""Parallel batch Mann-Whitney U tests using numba.prange."""
+
+import numba as nb
+import numpy as np
+
+from ._core import _mannwhitneyu_single
+
+
+@nb.njit(parallel=True)
+def _mannwhitneyu_rows(X, y, use_continuity, alternative):
+    """Run Mann-Whitney U test for each row of X against y.
+
+    Parameters
+    ----------
+    X : 2D float64 array (n_tests, n1)
+    y : 1D float64 array (n2,)
+    use_continuity : bool
+    alternative : int (0=two-sided, 1=less, 2=greater)
+
+    Returns
+    -------
+    U_out : float64 array (n_tests,)
+    p_out : float64 array (n_tests,)
+    """
+    n_tests = X.shape[0]
+    U_out = np.empty(n_tests, dtype=np.float64)
+    p_out = np.empty(n_tests, dtype=np.float64)
+    for i in nb.prange(n_tests):  # type: ignore
+        U_out[i], p_out[i] = _mannwhitneyu_single(
+            X[i].copy(), y, use_continuity, alternative
+        )
+    return U_out, p_out
+
+
+@nb.njit(parallel=True)
+def _mannwhitneyu_columns(X, Y, use_continuity, alternative):
+    """Run Mann-Whitney U test on each column of X vs corresponding column of Y.
+
+    Parameters
+    ----------
+    X : 2D float64 array (n1, n_genes)
+        First group (one column per gene).
+    Y : 2D float64 array (n2, n_genes)
+        Second group (one column per gene). Must have same number of columns as X.
+    use_continuity : bool
+    alternative : int (0=two-sided, 1=less, 2=greater)
+
+    Returns
+    -------
+    U_out : float64 array (n_genes,)
+    p_out : float64 array (n_genes,)
+    """
+    n_genes = X.shape[1]
+    U_out = np.empty(n_genes, dtype=np.float64)
+    p_out = np.empty(n_genes, dtype=np.float64)
+    for i in nb.prange(n_genes):  # type: ignore
+        x = X[:, i].copy()  # ensure contiguous
+        y = Y[:, i].copy()
+        U_out[i], p_out[i] = _mannwhitneyu_single(x, y, use_continuity, alternative)
+    return U_out, p_out

numba_mwu-0.1.1/src/numba_mwu/_core.py

@@ -0,0 +1,137 @@
+"""Numba-accelerated core functions for Mann-Whitney U test."""
+
+import math
+
+import numba as nb
+import numpy as np
+
+# Alternative hypothesis encoding
+TWO_SIDED = 0
+LESS = 1
+GREATER = 2
+
+
+@nb.njit
+def _ndtr(x: float) -> float:
+    """Normal CDF using math.erfc."""
+    return 0.5 * math.erfc(-x / math.sqrt(2.0))
+
+
+@nb.njit
+def _rankdata_avg(x):
+    """Rank data using average method, returning ranks and tie group sizes.
+
+    Parameters
+    ----------
+    x : 1D float64 array
+
+    Returns
+    -------
+    ranks : float64 array of same length as x
+    tie_sizes : float64 array of length n_groups (number of distinct values)
+    """
+    n = x.shape[0]
+    idx = np.argsort(x)
+    ranks = np.empty(n, dtype=np.float64)
+    tie_sizes = np.empty(n, dtype=np.float64)
+    n_ties = 0
+
+    i = 0
+    while i < n:
+        j = i
+        while j < n - 1 and x[idx[j]] == x[idx[j + 1]]:
+            j += 1
+        avg_rank = (i + j) / 2.0 + 1.0
+        tie_count = j - i + 1
+        for k in range(i, j + 1):
+            ranks[idx[k]] = avg_rank
+        tie_sizes[n_ties] = float(tie_count)
+        n_ties += 1
+        i = j + 1
+
+    return ranks, tie_sizes[:n_ties]
+
+
+@nb.njit
+def _mannwhitneyu_single(x, y, use_continuity, alternative):
+    """Compute Mann-Whitney U test for two 1D samples (asymptotic method).
+
+    Parameters
+    ----------
+    x : 1D float64 array
+    y : 1D float64 array
+    use_continuity : bool
+        Whether to apply continuity correction.
+    alternative : int
+        0 = two-sided, 1 = less, 2 = greater
+
+    Returns
+    -------
+    U1 : float
+        U statistic for sample x.
+    pvalue : float
+    """
+    n1 = x.shape[0]
+    n2 = y.shape[0]
+    n = n1 + n2
+
+    # Concatenate and rank
+    xy = np.empty(n, dtype=np.float64)
+    for i in range(n1):
+        xy[i] = x[i]
+    for i in range(n2):
+        xy[n1 + i] = y[i]
+
+    ranks, t = _rankdata_avg(xy)
+
+    # Sum of ranks for first sample
+    R1 = 0.0
+    for i in range(n1):
+        R1 += ranks[i]
+
+    U1 = R1 - n1 * (n1 + 1.0) / 2.0
+    U2 = n1 * n2 - U1
+
+    # Select statistic and multiplier based on alternative
+    if alternative == GREATER:
+        U = U1
+        f = 1.0
+    elif alternative == LESS:
+        U = U2
+        f = 1.0
+    else:  # TWO_SIDED
+        U = max(U1, U2)
+        f = 2.0
+
+    # Tie correction term: sum(t_i^3 - t_i)
+    tie_term = 0.0
+    for i in range(t.shape[0]):
+        ti = t[i]
+        tie_term += ti * ti * ti - ti
+
+    # Standard deviation with tie correction
+    mu = n1 * n2 / 2.0
+    denom = n * (n - 1.0)
+    var_inner = (n + 1.0) - tie_term / denom
+    s_sq = n1 * n2 / 12.0 * var_inner
+
+    if s_sq <= 0.0:
+        # All values identical — no evidence of difference
+        return U1, 1.0
+
+    s = math.sqrt(s_sq)
+
+    numerator = U - mu
+    if use_continuity:
+        numerator -= 0.5
+
+    z = numerator / s
+    p = _ndtr(-z) * f
+
+    # Clip to [0, 1]
+    if p > 1.0:
+        p = 1.0
+    if p < 0.0:
+        p = 0.0
+
+    return U1, p

numba_mwu-0.1.1/src/numba_mwu/_sparse.py

@@ -0,0 +1,236 @@
+"""Sparse-aware Mann-Whitney U test kernels for CSR matrices.
+
+Works directly with CSR format (the standard for single-cell data) by
+building a lightweight column index — a permutation array and column
+pointers — without copying the data values. Memory overhead is one int
+array of length nnz plus one int array of length (n_cols + 1) per matrix.
+"""
+
+import math
+
+import numba as nb
+import numpy as np
+
+from ._core import GREATER, LESS, _ndtr
+
+
+@nb.njit
+def _build_col_index(csr_indptr, csr_indices, n_cols):
+    """Build column pointers and a permutation from CSR arrays.
+
+    Returns
+    -------
+    col_indptr : int64 array (n_cols + 1)
+        col_indptr[j] .. col_indptr[j+1] gives the range of entries for
+        column j in the col_order array.
+    col_order : int64 array (nnz)
+        Indices into the CSR data arrays, grouped by column.
+        data[col_order[col_indptr[j]:col_indptr[j+1]]] gives the nonzero
+        values for column j.
+    """
+    nnz = csr_indices.shape[0]
+
+    # Count entries per column
+    col_counts = np.zeros(n_cols, dtype=np.int64)
+    for k in range(nnz):
+        col_counts[csr_indices[k]] += 1
+
+    # Build col_indptr via cumulative sum
+    col_indptr = np.empty(n_cols + 1, dtype=np.int64)
+    col_indptr[0] = 0
+    for j in range(n_cols):
+        col_indptr[j + 1] = col_indptr[j] + col_counts[j]
+
+    # Fill col_order (scatter each entry to its column bucket)
+    pos = np.empty(n_cols, dtype=np.int64)
+    for j in range(n_cols):
+        pos[j] = col_indptr[j]
+
+    col_order = np.empty(nnz, dtype=np.int64)
+    n_rows = csr_indptr.shape[0] - 1
+    for i in range(n_rows):
+        for k in range(csr_indptr[i], csr_indptr[i + 1]):
+            c = csr_indices[k]
+            col_order[pos[c]] = k
+            pos[c] += 1
+
+    return col_indptr, col_order
+
+
+@nb.njit
+def _sparse_mwu_column(vals_a, vals_b, n_a, n_b, use_continuity, alternative):
+    """Compute Mann-Whitney U for a single gene from two groups' nonzero values.
+
+    Zeros are treated analytically — they form a contiguous block at the
+    start of the sorted order (requires non-negative data).
+
+    Parameters
+    ----------
+    vals_a : float64 1-D array
+        Nonzero values for this gene in group A.
+    vals_b : float64 1-D array
+        Nonzero values for this gene in group B.
+    n_a, n_b : int
+        Total number of cells in group A and group B (including zeros).
+    use_continuity : bool
+    alternative : int (0=two-sided, 1=less, 2=greater)
+
+    Returns
+    -------
+    U1 : float64
+    pvalue : float64
+    """
+    n = n_a + n_b
+    nnz_a = vals_a.shape[0]
+    nnz_b = vals_b.shape[0]
+    nnz = nnz_a + nnz_b
+    nz = n - nnz  # total implicit zeros
+    nz_a = n_a - nnz_a  # zeros in group A
+
+    # --- All zeros: no evidence of difference ---
+    if nnz == 0:
+        return n_a * n_b / 2.0, 1.0
+
+    # --- Merge nonzero values into a single array with group tags ---
+    all_vals = np.empty(nnz, dtype=np.float64)
+    is_a = np.empty(nnz, dtype=np.int8)
+    for k in range(nnz_a):
+        all_vals[k] = vals_a[k]
+        is_a[k] = 1
+    for k in range(nnz_b):
+        all_vals[nnz_a + k] = vals_b[k]
+        is_a[nnz_a + k] = 0
+
+    # --- Sort nonzero values ---
+    sort_idx = np.argsort(all_vals)
+
+    # --- Walk sorted nonzeros: compute local ranks, tie correction,
+    #     and sum of global ranks for group A nonzeros ---
+    tie_term_nz = 0.0
+    sum_global_ranks_a = 0.0
+
+    i = 0
+    while i < nnz:
+        j = i
+        while j < nnz - 1 and all_vals[sort_idx[j]] == all_vals[sort_idx[j + 1]]:
+            j += 1
+
+        tie_count = float(j - i + 1)
+        tie_term_nz += tie_count * tie_count * tie_count - tie_count
+
+        # Local average rank (1-based among nonzeros), shifted to global
+        local_avg_rank = (i + j) / 2.0 + 1.0
+        global_avg_rank = nz + local_avg_rank
+
+        for k in range(i, j + 1):
+            if is_a[sort_idx[k]] == 1:
+                sum_global_ranks_a += global_avg_rank
+
+        i = j + 1
+
+    # --- R1: total rank sum for group A ---
+    zero_avg_rank = (nz + 1.0) / 2.0
+    R1 = nz_a * zero_avg_rank + sum_global_ranks_a
+
+    # --- U statistic ---
+    U1 = R1 - n_a * (n_a + 1.0) / 2.0
+    U2 = n_a * n_b - U1
+
+    # --- Tie correction (zero block + nonzero tie groups) ---
+    tie_term = (float(nz) * float(nz) * float(nz) - float(nz)) + tie_term_nz
+
+    # --- Variance and z-score ---
+    mu = n_a * n_b / 2.0
+    denom = float(n) * float(n - 1)
+    var_inner = (n + 1.0) - tie_term / denom
+    s_sq = n_a * n_b / 12.0 * var_inner
+
+    if s_sq <= 0.0:
+        return U1, 1.0
+
+    s = math.sqrt(s_sq)
+
+    if alternative == GREATER:
+        U = U1
+        f = 1.0
+    elif alternative == LESS:
+        U = U2
+        f = 1.0
+    else:  # TWO_SIDED
+        U = max(U1, U2)
+        f = 2.0
+
+    numerator = U - mu
+    if use_continuity:
+        numerator -= 0.5
+
+    z = numerator / s
+    p = _ndtr(-z) * f
+
+    if p > 1.0:
+        p = 1.0
+    if p < 0.0:
+        p = 0.0
+
+    return U1, p
+
+
+@nb.njit
+def _gather_col_vals(csr_data, col_order, start, end):
+    """Gather nonzero values for a single column from CSR data via col_order."""
+    n = end - start
+    vals = np.empty(n, dtype=np.float64)
+    for k in range(n):
+        vals[k] = csr_data[col_order[start + k]]
+    return vals
+
+
+@nb.njit(parallel=True)
+def _sparse_mwu_batch(
+    data_a,
+    col_indptr_a,
+    col_order_a,
+    n_a,
+    data_b,
+    col_indptr_b,
+    col_order_b,
+    n_b,
+    use_continuity,
+    alternative,
+):
+    """Run sparse MWU test on each gene using two CSR matrices' column indices.
+
+    Parameters
+    ----------
+    data_a : float64 1-D array — CSR data for group A
+    col_indptr_a : int64 1-D array (n_genes + 1) — column pointers for A
+    col_order_a : int64 1-D array (nnz_a) — column permutation for A
+    n_a : int — total rows in A
+    data_b : float64 1-D array — CSR data for group B
+    col_indptr_b : int64 1-D array (n_genes + 1) — column pointers for B
+    col_order_b : int64 1-D array (nnz_b) — column permutation for B
+    n_b : int — total rows in B
+    use_continuity : bool
+    alternative : int
+
+    Returns
+    -------
+    U_out : float64 array (n_genes,)
+    p_out : float64 array (n_genes,)
+    """
+    n_genes = col_indptr_a.shape[0] - 1
+    U_out = np.empty(n_genes, dtype=np.float64)
+    p_out = np.empty(n_genes, dtype=np.float64)
+
+    for j in nb.prange(n_genes):  # type: ignore
+        vals_a = _gather_col_vals(
+            data_a, col_order_a, col_indptr_a[j], col_indptr_a[j + 1]
+        )
+        vals_b = _gather_col_vals(
+            data_b, col_order_b, col_indptr_b[j], col_indptr_b[j + 1]
+        )
+        U_out[j], p_out[j] = _sparse_mwu_column(
+            vals_a, vals_b, n_a, n_b, use_continuity, alternative
+        )
+
+    return U_out, p_out