diff --git a/schpf/hpf_numba.py b/schpf/hpf_numba.py index 4b7fef3..b0bbd49 100644 --- a/schpf/hpf_numba.py +++ b/schpf/hpf_numba.py @@ -113,6 +113,89 @@ def compute_Xphi_data(X_data, X_row, X_col, return Xphi +@numba.njit(nogil=True) +def _compute_Xphi_data_serial(X_data, X_row, X_col, + theta_vi_shape, theta_vi_rate, + beta_vi_shape, beta_vi_rate): + """Single-threaded twin of compute_Xphi_data. + + Byte-for-byte the body of compute_Xphi_data with numba.prange + replaced by range. The Xphi computation is an independent + per-nonzero map-each iteration writes Xphi[i,:] from a + per-iteration logrho, with no cross-iteration reduction-and the + two e_logx precomputes are independent per-cell/gene writes, so the + serial loop produces results bit-identical to the parallel kernel. + Same decorator as compute_Xphi_data minus parallel=True. + """ + # convenience + ncells, ngenes = (theta_vi_shape.shape[0], beta_vi_shape.shape[0]) + nfactors, nnz = (theta_vi_shape.shape[1], X_data.shape[0]) + dtype = theta_vi_shape.dtype + + # precompute theta.e_logx + theta_e_logx = np.zeros_like(theta_vi_shape, dtype=dtype) + for i in range(ncells): + for k in range(nfactors): + theta_e_logx[i,k] = psi(theta_vi_shape[i,k]) \ + - np.log(theta_vi_rate[i,k]) + + # precompute beta.e_logx + beta_e_logx = np.zeros_like(beta_vi_shape, dtype=dtype) + for i in range(ngenes): + for k in range(nfactors): + beta_e_logx[i,k] = psi(beta_vi_shape[i,k]) \ + - np.log(beta_vi_rate[i,k]) + + # compute Xphi + Xphi = np.zeros((X_row.shape[0], theta_e_logx.shape[1]), dtype=dtype) + for i in range(nnz): + logrho = np.zeros((Xphi.shape[1]), dtype=dtype) + for k in range(nfactors): + logrho[k] = theta_e_logx[X_row[i],k] + beta_e_logx[X_col[i], k] + + #log normalizer trick + rho_shift = np.zeros((Xphi.shape[1]), dtype=dtype) + normalizer = np.zeros(1, dtype=dtype)[0] + largest_in = np.max(logrho) + for k in range(nfactors): + rho_shift[k] = np.exp(logrho[k] - largest_in) + normalizer += rho_shift[k] + + for k in range(nfactors): + Xphi[i,k] = X_data[i] * rho_shift[k] / normalizer + + return Xphi + + +def compute_Xphi_data_serial(X, theta, beta, theta_ix=None): + """Single-threaded numba version of the Xphi computation. + + Object-signature wrapper (same interface as compute_Xphi_data_numpy) + around the single-threaded njit kernel _compute_Xphi_data_serial. + Used by the single_process=True path so it runs JIT'd numba instead + of the pure-numpy compute_Xphi_data_numpy. + + Feeds the kernel the same arrays the parallel compute_Xphi_data + receives (theta.vi_shape[theta_ix], etc.). Result is + bit-identical to the parallel kernel and, as a consequence, + single_process=True now matches single_process=False exactly (where + the numpy fallback diverged at the digamma source and the softmax + formulation). + + Parameters mirror compute_Xphi_data_numpy: X is the sparse (coo) + count matrix; theta, beta are HPF_Gamma; theta_ix optionally subsets + theta's rows (used in projection / minibatch dispatch). + """ + if theta_ix is None: + theta_vi_shape = theta.vi_shape + theta_vi_rate = theta.vi_rate + else: + theta_vi_shape = theta.vi_shape[theta_ix] + theta_vi_rate = theta.vi_rate[theta_ix] + return _compute_Xphi_data_serial( + X.data, X.row, X.col, + theta_vi_shape, theta_vi_rate, + beta.vi_shape, beta.vi_rate) def compute_Xphi_data_numpy(X, theta, beta, theta_ix=None): """Single-threaded version of compute_Xphi_data diff --git a/schpf/scHPF_.py b/schpf/scHPF_.py index dcb327c..665333c 100644 --- a/schpf/scHPF_.py +++ b/schpf/scHPF_.py @@ -655,7 +655,7 @@ def _fit(self, X, freeze_genes=False, reinit=True, loss_function=None, Xphi_data = X_batch.data[:,None] * random_phi else: if single_process: - Xphi_data = compute_Xphi_data_numpy(X_batch, theta, beta, + Xphi_data = compute_Xphi_data_serial(X_batch, theta, beta, theta_ix=batch_ix) else: Xphi_data = compute_Xphi_data( diff --git a/tests/test_xphi_serial.py b/tests/test_xphi_serial.py new file mode 100644 index 0000000..e778811 --- /dev/null +++ b/tests/test_xphi_serial.py @@ -0,0 +1,86 @@ +"""single-threaded njit Xphi equivalence. + +compute_Xphi_data_serial is compute_Xphi_data's body with prange->range, so: + + - it is bit-identical to the parallel kernel (the Xphi step is an + independent per-nonzero map; prange changes execution order, not the + per-element arithmetic) -> asserted by equality, with theta_ix both + None and a non-trivial reorder; + - it is close to the numpy fallback, because the two existing + kernels already diverge: the parallel/serial kernels derive e_logx from + vi_shape/vi_rate via numba psi, while the numpy version reads the + object's precomputed e_logx and uses logsumexp. single_process + now agrees with single_process=False. + +The numpy function is kept (additive change), still callable here as +the comparison reference. +""" + +import numpy as np +import pytest + +schpf = pytest.importorskip("schpf") +from scipy.sparse import coo_matrix +from schpf.scHPF_ import HPF_Gamma +from schpf.hpf_numba import ( + compute_Xphi_data, + compute_Xphi_data_numpy, + compute_Xphi_data_serial, +) + +RANDOM_STATE = 11 + + +def _problem(N=80, M=50, K=6, density=0.3, seed=RANDOM_STATE): + rng = np.random.default_rng(seed) + dense = (rng.random((N, M)) < density) * rng.integers(1, 6, (N, M)) + X = coo_matrix(dense.astype(np.float64)) + u = lambda *s: rng.uniform(0.5, 1.5, s) + theta = HPF_Gamma(u(N, K), u(N, K)) + beta = HPF_Gamma(u(M, K), u(M, K)) + return X, theta, beta + + +def test_serial_equals_parallel_no_theta_ix(): + """0 ULP vs the parallel kernel (theta_ix=None).""" + X, theta, beta = _problem() + serial = compute_Xphi_data_serial(X, theta, beta) + parallel = compute_Xphi_data( + X.data, X.row, X.col, + theta.vi_shape, theta.vi_rate, + beta.vi_shape, beta.vi_rate, + ) + assert serial.shape == parallel.shape + assert np.array_equal(serial, parallel) + + +def test_serial_equals_parallel_with_theta_ix(): + """0 ULP vs the parallel kernel under a non-trivial theta_ix.""" + X, theta, beta = _problem() + N = theta.vi_shape.shape[0] + theta_ix = np.random.default_rng(RANDOM_STATE + 1).permutation(N) + + serial = compute_Xphi_data_serial(X, theta, beta, theta_ix=theta_ix) + parallel = compute_Xphi_data( + X.data, X.row, X.col, + theta.vi_shape[theta_ix], theta.vi_rate[theta_ix], + beta.vi_shape, beta.vi_rate, + ) + assert np.array_equal(serial, parallel) + + +def test_serial_approx_numpy(): + """approx the numpy fallback.""" + X, theta, beta = _problem() + serial = compute_Xphi_data_serial(X, theta, beta) + numpy_ref = compute_Xphi_data_numpy(X, theta, beta) + assert serial.shape == numpy_ref.shape + assert np.allclose(serial, numpy_ref, rtol=1e-9, atol=1e-12) + + +def test_numpy_function_kept_unchanged(): + """The additive change keeps compute_Xphi_data_numpy callable.""" + X, theta, beta = _problem() + out = compute_Xphi_data_numpy(X, theta, beta) + assert out.shape == (X.nnz, theta.vi_shape.shape[1]) + assert np.allclose(out.sum(axis=1), X.data, rtol=1e-9, atol=1e-12)