Skip to content

fix: correct pvalue logic and add discretization step #54

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
drbh wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: correct pvalue logic and add discretization step #54

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 89 additions & 16 deletions src/pdex/_single_cell.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from collections.abc import Iterator
from functools import partial
from multiprocessing.shared_memory import SharedMemory
from typing import Callable

import anndata as ad
import os
Expand Down Expand Up @@ -443,6 +444,9 @@ def prepare_ranksum_buffers(X_target, X_ref):
return K_cols, pool_cnt, pool_cnt_t


SQRT2 = math.sqrt(2.0)


@njit(parallel=True, fastmath=True)
def ranksum_kernel_with_pool(X_target, X_ref, K_cols, pool_cnt, pool_cnt_t):
n_t = X_target.shape[0]
Expand All @@ -452,7 +456,7 @@ def ranksum_kernel_with_pool(X_target, X_ref, K_cols, pool_cnt, pool_cnt_t):
p_values = np.empty(n_genes, dtype=np.float64)
u_stats = np.empty(n_genes, dtype=np.float64)

for j in prange(n_genes):
for j in prange(n_genes): # type: ignore[not-iterable]
tid = get_thread_id() # Numba ≥ 0.56
cnt = pool_cnt[tid]
cnt_t = pool_cnt_t[tid]
Expand Down Expand Up @@ -482,19 +486,17 @@ def ranksum_kernel_with_pool(X_target, X_ref, K_cols, pool_cnt, pool_cnt_t):

# U and p
u = rank_sum_target - 0.5 * n_t * (n_t + 1)
u_stats[j] = u

N = n_t + n_r
if N > 1:
tie_adj = tie_sum / (N * (N - 1))
sigma2 = (n_t * n_r) * ((N + 1) - tie_adj) / 12.0
if sigma2 > 0.0:
z = (u - 0.5 * n_t * n_r) / math.sqrt(sigma2)
p_values[j] = math.erfc(abs(z) / math.sqrt(2.0))
else:
p_values[j] = 1.0
else:
p_values[j] = 1.0
u2 = n_t * n_r - u
u_min = u if u > u2 else u2

n = n_t + n_r
s = np.sqrt(n_t * n_r / 12 * ((n + 1) - tie_sum / (n * (n - 1))))
mean = n_t * n_r / 2.0
cc = 0.5 if u_min < mean else -0.5 if u_min > mean else 0.0
z_corr = (u_min - mean + cc) / s
p_corr = math.erfc(abs(z_corr) / SQRT2)
u_stats[j] = u_min
p_values[j] = p_corr if s > 0.0 else 1.0

# clear just the touched slice
for v in range(Kp1_use):
Expand Down Expand Up @@ -525,6 +527,8 @@ def _process_single_target_vectorized(
obs_values: np.ndarray,
X: np.ndarray,
X_ref: np.ndarray,
X_disc: np.ndarray,
X_ref_disc: np.ndarray,
means_ref: np.ndarray,
gene_names: np.ndarray,
is_log1p: bool,
Expand All @@ -538,6 +542,7 @@ def _process_single_target_vectorized(
# Get target data
target_mask = obs_values == target
X_target = X[target_mask, :]
X_target_disc = X_disc[target_mask, :]

# Vectorized means calculation
if is_log1p:
Expand All @@ -563,7 +568,8 @@ def _process_single_target_vectorized(
pcc = np.where(means_ref == 0, np.nan, pcc)

# Statistical tests across all genes simultaneously
p_values, statistics = _vectorized_ranksum_test_numba(X_target, X_ref)
p_values, statistics = _vectorized_ranksum_test_numba(X_target_disc, X_ref_disc)
pairwise_fdr = false_discovery_control(p_values, method="bh")

# Build results for all genes at once using vectorized operations
target_results = [
Expand All @@ -577,6 +583,7 @@ def _process_single_target_vectorized(
"fold_change": fc[i],
"p_value": p_values[i],
"statistic": statistics[i],
"pairwise_fdr": pairwise_fdr[i],
}
for i in range(len(gene_names))
]
Expand Down Expand Up @@ -637,6 +644,68 @@ def parallel_differential_expression_vec(
targets_to_process = [target for target in unique_targets if target != reference]
gene_names = adata.var.index.values

is_discrete = np.issubdtype(X.dtype, np.integer) and np.issubdtype(
X_ref.dtype, np.integer
)
if is_discrete:
logger.info("Data appears discrete - using direct integer representation")
X_disc = X.astype(np.int32, copy=False)
X_ref_disc = X_ref.astype(np.int32, copy=False)
else:
logger.info("Data appears continuous - discretizing to integer representation")
import time

start_discretization = time.time()

# Discretize data to the smallest integer range while preserving order
X_f = np.asarray(X, dtype=np.float32, order="C")

# Handle NaNs/Infs
if not np.isfinite(X_f).all():
X_f = np.nan_to_num(
X_f,
nan=0.0,
posinf=np.finfo(np.float32).max,
neginf=np.finfo(np.float32).min,
)

X_min = float(np.min(X_f))
X_max = float(np.max(X_f))

# Shift if needed
if X_min < 0.0:
X_f = X_f - X_min
X_ref = X_ref - X_min
X_max = X_max - X_min

# Degenerate case: all values identical
if X_max <= 0.0:
X_disc = np.zeros_like(X_f, dtype=np.int32)
X_ref_disc = np.zeros_like(X_ref, dtype=np.int32)
else:
# Scale into the *smallest* integer range: 0..N where N = number of distinct values - 1
# If continuous, fallback to full dynamic range of int32
unique_vals = np.unique(X_f)
if unique_vals.shape[0] < 1_000_000: # heuristic: treat as discrete
mapping = {v: i for i, v in enumerate(np.unique(unique_vals))}
# Vectorized mapping with np.searchsorted
idxs = np.searchsorted(unique_vals, X_f.ravel()).reshape(X_f.shape)
X_disc = idxs.astype(np.int32)
idxs_ref = np.searchsorted(unique_vals, X_ref.ravel()).reshape(
X_ref.shape
)
X_ref_disc = idxs_ref.astype(np.int32)
else:
# Continuous → scale to full int32 range
scale = np.iinfo(np.int32).max / X_max
X_disc = np.rint(X_f * scale).astype(np.int32, copy=False)
X_ref_disc = np.rint(X_ref * scale).astype(np.int32, copy=False)

end_discretization = time.time()
logger.info(
f"Discretization completed in {end_discretization - start_discretization:.2f} seconds"
)

# Process targets sequentially with numba functions
logger.info(f"Processing {len(targets_to_process)} targets")
all_results = []
Expand All @@ -647,6 +716,8 @@ def parallel_differential_expression_vec(
obs_values=obs_values, # type: ignore
X=X,
X_ref=X_ref,
X_disc=X_disc,
X_ref_disc=X_ref_disc,
means_ref=means_ref,
gene_names=gene_names, # type: ignore
is_log1p=is_log1p,
Expand Down Expand Up @@ -695,4 +766,6 @@ def parallel_differential_expression_vec_wrapper(

if use_experimental:
logger.warning("Using experimental features")
parallel_differential_expression = parallel_differential_expression_vec_wrapper
parallel_differential_expression: Callable = (
parallel_differential_expression_vec_wrapper
)
38 changes: 27 additions & 11 deletions tests/bench_expr.py → tests/test_experimental_correctness.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,25 +62,30 @@ def test_correctness_comparison():
adata = _create_test_data()
ref_result, exp_result = _run_both_modes(adata)

# Check structure matches
assert ref_result.shape == exp_result.shape
assert list(ref_result.columns) == list(exp_result.columns)
ref_columns = ref_result.columns
exp_columns = exp_result.columns

# check that the columns are the same except that exp_result has a "pairwise_fdr" column
if "pairwise_fdr" in exp_columns:
exp_columns = exp_columns.drop("pairwise_fdr")

assert set(ref_columns) == set(exp_columns)

# Check means match exactly
for col in ["reference_mean", "target_mean"]:
np.testing.assert_allclose(ref_result[col], exp_result[col], rtol=1e-5)
np.testing.assert_allclose(ref_result[col], exp_result[col], rtol=1e-6)

# Check fold changes match (excluding NaNs)
mask = ~(ref_result["fold_change"].isna() | exp_result["fold_change"].isna())
np.testing.assert_allclose(
ref_result.loc[mask, "fold_change"],
exp_result.loc[mask, "fold_change"],
rtol=1e-5,
rtol=1e-6,
)

# P-values should match within 1e-6 absolute tolerance
np.testing.assert_allclose(
ref_result["p_value"], exp_result["p_value"], rtol=1e-2, atol=1e-6
ref_result["p_value"], exp_result["p_value"], rtol=1e-6, atol=1e-12
)


Expand Down Expand Up @@ -112,12 +117,12 @@ def test_detailed_correctness_metrics():

# Check fold changes match
np.testing.assert_allclose(
ref_result["fold_change"], exp_result["fold_change"], rtol=1e-5
ref_result["fold_change"], exp_result["fold_change"], rtol=1e-6, atol=1e-12
)

# P-values should match within 1e-6 absolute tolerance
np.testing.assert_allclose(
ref_result["p_value"], exp_result["p_value"], rtol=1e-2, atol=1e-6
ref_result["p_value"], exp_result["p_value"], rtol=1e-6, atol=1e-12
)


Expand Down Expand Up @@ -146,15 +151,22 @@ def test_edge_cases_numerical_stability():

# Fold changes should match
np.testing.assert_allclose(
ref_small["fold_change"], exp_small["fold_change"], rtol=1e-5
ref_small["fold_change"], exp_small["fold_change"], rtol=1e-6
)

# P-values should match within 1e-6 absolute tolerance
np.testing.assert_allclose(
ref_small["p_value"], exp_small["p_value"], rtol=1e-2, atol=1e-6
ref_small["p_value"], exp_small["p_value"], rtol=1e-6, atol=1e-12
)

assert ref_large.shape == exp_large.shape
assert ref_large.shape[0] == exp_large.shape[0]

# the shape [1] is different because the experimental mode adds a pairwise_fdr column
ref_columns = set(ref_large.columns)
exp_columns = set(exp_large.columns)
if "pairwise_fdr" in exp_columns:
exp_columns.remove("pairwise_fdr")
assert ref_columns == exp_columns


@pytest.mark.parametrize(
Expand Down Expand Up @@ -210,3 +222,7 @@ def test_benchmark_parameterized_datasets(
assert len(result) > 0
assert "p_value" in result.columns
assert "fold_change" in result.columns


if __name__ == "__main__":
pytest.main([__file__])
Loading