Add single-threaded njit Xphi kernel for single-process path

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

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(

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)