ebin 1.0__py3-none-any.whl

ebin/gauss_rules.py

@@ -0,0 +1,297 @@
+from functools import partial
+try:
+    from .ftzkiller import FTZKiller
+except ImportError:  # allow running as a loose module
+    from ftzkiller import FTZKiller
+from scipy.linalg import eigh_tridiagonal as scipy_eigh_tridiagonal
+from jax import pure_callback
+from scipy.special import roots_genlaguerre
+from jax.scipy.special import gammaln
+from jax.lax import scan
+from jax import vmap
+import jax.numpy as jnp
+import numpy as np
+import jax
+
+from enum import Enum 
+
+class Rules(Enum):
+    GenLaguerre = 'GenLaguerre'
+    Charlier = 'Charlier'
+    ShiftedJacobi = 'ShiftedJacobi'
+
+
+
+def _calculate_for_single_lambda(lam, diagonal, off_diagonal):
+    """
+    Helper function to compute log(|[v]_1|) for a particualr eigenvalue.
+    """
+
+    # Pre-calculate the first two components (indices 0 and 1)
+    u0 = 1.0
+    u1 = -(diagonal[0] - lam) * u0 / off_diagonal[0]
+
+    initial_carry = (u1, u0)
+    
+    # The scan will compute components u[2] through u[n-1], which is n-2 steps.
+    # All input arrays for the scan must therefore have length n-2.
+    scan_diagonals = diagonal[1:-1]             # Needs a_2, ..., a_{n-1} (0-indexed: diag[1]...diag[n-2])
+    scan_off_diagonals_prev = off_diagonal[:-1] # Needs b_1, ..., b_{n-2} (0-indexed: off[0]...off[n-3])
+    scan_off_diagonals_curr = off_diagonal[1:]  # Needs b_2, ..., b_{n-1} (0-indexed: off[1]...off[n-2])
+
+    def _recurrence_step(carry, scan_inputs):
+        u_curr, u_prev = carry
+        a_k, b_k_minus_1, b_k = scan_inputs
+        
+        u_next = -((a_k - lam) * u_curr + b_k_minus_1 * u_prev) / b_k
+        
+        new_carry = (u_next, u_curr)
+        component_to_collect = u_next
+        return new_carry, component_to_collect
+
+    # Run the scan to get components u_2, ..., u_{n-1}
+    _, u_rest = scan(
+        _recurrence_step,
+        initial_carry,
+        (scan_diagonals, scan_off_diagonals_prev, scan_off_diagonals_curr)
+    )
+
+    # Combine all components into a single vector u
+    u = jnp.concatenate([jnp.array([u0, u1]), u_rest])
+
+    # Stably compute the log of the first component's magnitude 
+    u_max = jnp.max(jnp.abs(u))
+    w = u / u_max
+    
+    sum_w_sq = jnp.dot(w, w)
+    log_s_squared = 2 * jnp.log(u_max) + jnp.log(sum_w_sq)
+    
+    result = jnp.where(u_max > 0, -0.5 * log_s_squared, -jnp.inf)
+    
+    return result
+
+
+
+@jax.jit
+def calc_first_components(diagonal, off_diagonal, eigenvalues):
+    """
+    Calculates the logarithm of the absolute value of the first components of all
+    eigenvectors of a symmetric tridiagonal matrix.
+
+    Args:
+        diagonal (jnp.ndarray): The main diagonal of the matrix (a_1, ..., a_n).
+        off_diagonal (jnp.ndarray): The first off-diagonal (b_1, ..., b_{n-1}).
+        eigenvalues (jnp.ndarray): The pre-computed eigenvalues of the matrix.
+
+    Returns:
+        jnp.ndarray: An array containing log(|[v_j]_1|) for each eigenvector v_j.
+    """
+    n = diagonal.shape[0]
+
+    # Handle edge cases that don't fit the scan logic
+    if n == 0:
+        return jnp.array([])
+    if n == 1:
+        return jnp.array([0.0])
+    if n == 2:
+        def _calc_n2(lam, diag, offdiag):
+            u0 = 1.0
+            u1 = -(diag[0] - lam) * u0 / offdiag[0]
+            u = jnp.array([u0, u1])
+            u_max = jnp.max(jnp.abs(u))
+            w = u / u_max
+            sum_w_sq = jnp.dot(w, w)
+            log_s_squared = 2 * jnp.log(u_max) + jnp.log(sum_w_sq)
+            return -0.5 * log_s_squared
+        
+        return vmap(_calc_n2, in_axes=(0, None, None))(eigenvalues, diagonal, off_diagonal)
+
+    vmapped_calculator = vmap(
+        _calculate_for_single_lambda, in_axes=(0, None, None)
+    )
+    
+    return vmapped_calculator(eigenvalues, diagonal, off_diagonal)
+
+@jax.custom_jvp
+def eigh_tridiagonal(diag, offdiag):
+  def scipy_eigh_vals_only(d, o):
+      return scipy_eigh_tridiagonal(d, o, eigvals_only=True)
+  
+  n = diag.shape[0]
+  with FTZKiller():
+      w = pure_callback(
+          scipy_eigh_vals_only,
+          jax.ShapeDtypeStruct((n,), diag.dtype), # Shape of the output
+          diag, offdiag,
+          vmap_method='sequential',
+      )
+      w = jax.block_until_ready(w)
+  return w
+
+
+@eigh_tridiagonal.defjvp
+def _eigh_tridiagonal_jvp(primals, tangents):
+  diag, offdiag = primals
+  diag_dot, offdiag_dot = tangents
+
+  def scipy_eigh_full(d, o):
+    with FTZKiller():
+        return scipy_eigh_tridiagonal(d, o, eigvals_only=False)
+  
+  n = diag.shape[0]
+  
+  w, v = pure_callback(
+       scipy_eigh_full,
+        (
+            jax.ShapeDtypeStruct((n,), diag.dtype),
+            jax.ShapeDtypeStruct((n, n), diag.dtype)
+        ),
+        diag, offdiag,
+        vmap_method='sequential',
+    )
+  
+  delta_T = jnp.diag(diag_dot) + jnp.diag(offdiag_dot, k=1) + jnp.diag(offdiag_dot, k=-1)
+  
+  w_dot = jnp.einsum('ik,ij,jk->k', v, delta_T, v)
+
+  return w, w_dot
+
+def compute_roots_genlaguerre_scipy(n: int, alpha: float):
+    n = int(n)
+    alpha = float(alpha)
+    
+    with FTZKiller():
+        roots, weights = roots_genlaguerre(n, alpha)
+    return roots.astype(np.float64), weights.astype(np.float64)
+
+def approximate_charliet_smallest_root(n: int, mu: float) -> float:
+    """
+    This function uses the formula:
+        x_1 ≈ μ^n / Σ_{k=1 to n} [ C(n,k) * (k-1)! * μ^(n-k) ]
+    """
+
+    log_mu = jnp.log(mu)
+    log_numerator = n * log_mu
+
+    k_values = jnp.arange(1, n + 1)
+
+    log_binom_coeff = gammaln(n + 1) - gammaln(k_values + 1) - gammaln(n - k_values + 1)
+    
+    log_factorial = gammaln(k_values)
+
+    log_mu_term = (n - k_values) * log_mu
+    log_terms_denominator = log_binom_coeff + log_factorial + log_mu_term
+
+    log_denominator = jax.scipy.special.logsumexp(log_terms_denominator)
+    log_root = log_numerator - log_denominator
+
+    return jnp.exp(log_root)
+
+def construct_jacobi(alpha, beta, n):
+    """
+    Constructs the tridiagonal Jacobi matrix for w(x) = x^(alpha-1)(1-x)^(beta-1).
+    
+    Args:
+        alpha (float): Parameter alpha > 0.
+        beta (float): Parameter beta > 0.
+        n (int): The size of the matrix (n x n).
+
+    Returns:
+        tuple[np.ndarray, np.ndarray]: A tuple containing two arrays:
+                                       - diag (main diagonal, size n)
+                                       - offdiag (sub/super-diagonal, size n-1)
+    """
+    if n <= 0:
+        return jnp.array([]), jnp.array([])
+
+    k = jnp.arange(n, dtype=np.float64)
+    ab = alpha + beta
+
+    diag = jnp.where(
+        k == 0,
+        alpha / ab,
+        0.5 * (1 + ((alpha - 1)**2 - (beta - 1)**2) / ((2*k + ab - 2) * (2*k + ab)))
+    )
+
+    if n == 1:
+        return diag, jnp.array([])
+
+    k_off = k[:-1]
+    c = 2*k_off + ab
+    offdiag = jnp.where(
+        k_off == 0,
+        jnp.sqrt((alpha * beta) / (ab**2 * (ab + 1))),
+        (1/c) * jnp.sqrt(((k_off + 1)*(k_off + alpha)*(k_off + beta)*(k_off + ab - 1)) / (c**2 - 1))
+    )
+
+    return diag, offdiag
+
+@partial(jax.jit, static_argnames=('num_nodes', 'rule', 'norm', 'eigh_solver',
+                                   ))
+def compute_nodes_and_logweights(num_nodes: int, param: float, rule: Rules,
+                                 norm: bool = True, eigh_solver: bool = None):
+    if eigh_solver is None:
+        if rule == Rules.Charlier:
+            eigh_solver = True
+        else:
+            eigh_solver = False
+    k = jnp.arange(num_nodes)
+    lognorm = 0.0
+    if rule == Rules.GenLaguerre:
+        if not norm:
+            lognorm = gammaln(param + 1)
+        diag = 2 * k + param + 1
+        k_off = jnp.arange(1, num_nodes, dtype=float)
+        offdiag = -jnp.sqrt(k_off * (k_off + param))
+    elif rule == Rules.Charlier:
+        diag = k + param
+        k = jnp.arange(1, num_nodes)
+        offdiag = jnp.sqrt(k * param)
+        if not norm:
+            lognorm = param 
+    elif rule == Rules.ShiftedJacobi:
+        diag, offdiag = construct_jacobi(*param, num_nodes)
+        if not norm:
+            lognorm = jax.scipy.special.betaln(*param)
+    if eigh_solver:
+        d = jnp.diag(diag) + jnp.diag(offdiag, 1) + jnp.diag(offdiag, -1)
+        with FTZKiller():
+            nodes, logweights = jnp.linalg.eigh(d)
+            logweights = jnp.log(jnp.abs(logweights.at[0].get()))
+    else:
+        nodes = eigh_tridiagonal(diag, offdiag)
+        logweights = calc_first_components(diag, offdiag, nodes)
+    logweights = lognorm + 2 * logweights
+    return jnp.clip(nodes, 0.0, None), logweights
+
+# import numdifftools as nd 
+
+# jax.config.update('jax_enable_x64', True)
+# jax.config.update('jax_platforms', 'cpu')
+
+
+# @partial(jax.jit, static_argnames=('n', 'mode'))
+# def fun(alpha, mode: bool, n=20):
+#     n, l = compute_nodes_and_logweights(n, alpha, rule=Rules.Charlier,
+#                                         eigh_solver=True)
+#     return ((l  ) ).mean()
+
+# f_old = partial(fun, mode=True)
+# grad_old = jax.jit(jax.jacfwd(f_old, argnums=0))
+# grad_old_nd = nd.Gradient(f_old)
+
+# # f_new= partial(fun, mode=False)
+# # grad_new = jax.jit(jax.jacrev(f_new, argnums=0))
+# # grad_new_nd = nd.Gradient(f_new)
+
+
+# alpha = 1e-3
+
+# print('funs:', f_old(alpha), )
+# print('grad:', grad_old(alpha), )
+# print('grad_nd:', grad_old_nd(alpha), )
+# print('graddiff', np.abs(grad_old(alpha) - grad_old_nd(alpha)))
+
+
+
+

ebin/initialize.py

@@ -0,0 +1,186 @@
+"""Starting values for the compound NB-Poisson model.
+
+1. Fit a zero-inflated NB (ZINB) independently to each of the S*B count columns.
+2. Take a W2 barycenter of those NBs (location-scale surrogate: mean of means,
+   mean of sds) as a reference NB.
+3. Map every column onto the reference by the mid-distribution quantile
+   transform, Z = F_ref^{-1}(F_sb(x) - 0.5 f_sb(x)); normalize per column, sum
+   over replicates, and row-normalize to get Pi0.
+4. R0 = r_hat / (Lambda0_s * mean_n Pi0) matches the model means to the
+   marginal fits, and P0 = (1 + R0) / (1/p_hat + R0) matches their dispersion.
+"""
+
+from dataclasses import dataclass, field
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+from jax.scipy.special import gammaln
+from scipy.optimize import minimize
+from scipy.stats import nbinom as sp_nbinom
+
+jax.config.update("jax_enable_x64", True)
+
+
+@jax.jit
+def _zinb_nll_grad(theta, xu, w):
+    def nll(theta):
+        a, c, d = theta
+        r = jnp.exp(a)
+        logp = -jnp.logaddexp(0.0, -c)          # log sigmoid(c)
+        log1mp = -jnp.logaddexp(0.0, c)
+        log_pi0 = -jnp.logaddexp(0.0, -d)
+        log_1mpi0 = -jnp.logaddexp(0.0, d)
+        log_nb = (gammaln(xu + r) - gammaln(r) - gammaln(xu + 1.0)
+                  + r * logp + xu * log1mp)
+        log_zinb = jnp.where(
+            xu == 0,
+            jnp.logaddexp(log_pi0, log_1mpi0 + log_nb),
+            log_1mpi0 + log_nb)
+        return -jnp.sum(w * log_zinb) / jnp.sum(w)
+    return jax.value_and_grad(nll)(theta)
+
+
+def _pad_pow2(a, fill):
+    n = 1
+    while n < a.size:
+        n *= 2
+    out = np.full(n, fill, dtype=np.float64)
+    out[:a.size] = a
+    return out
+
+
+def fit_zinb(x, weights=None):
+    """Fit ZINB(r, p, pi0) to counts ``x``.  Returns (r, p, pi0, nll_per_obs)."""
+    x = np.asarray(x, dtype=np.float64)
+    if weights is None:
+        xu, w = np.unique(x, return_counts=True)
+        w = w.astype(np.float64)
+    else:
+        xu, w = x, np.asarray(weights, dtype=np.float64)
+    # pad to a power of two so jit recompiles only O(log U) shapes
+    xu_p = _pad_pow2(xu, 1.0)
+    w_p = _pad_pow2(w, 0.0)
+
+    wt = w.sum()
+    m = float((xu * w).sum() / wt)
+    v = float((w * (xu - m) ** 2).sum() / wt)
+    frac0 = float(w[xu == 0].sum() / wt) if (xu == 0).any() else 0.0
+    m_pos = m / max(1.0 - frac0, 1e-6)          # crude nonzero-component mean
+    v_pos = max(v, m_pos * 1.5)
+    r0 = np.clip(m_pos ** 2 / max(v_pos - m_pos, 1e-6), 1e-3, 1e6)
+    p0 = np.clip(r0 / (r0 + m_pos), 1e-6, 1 - 1e-6)
+    pi0_0 = np.clip(frac0 * 0.5 + 1e-4, 1e-4, 0.9)
+    theta0 = np.array([np.log(r0), np.log(p0 / (1 - p0)),
+                       np.log(pi0_0 / (1 - pi0_0))])
+
+    def fg(theta):
+        v, g = _zinb_nll_grad(jnp.asarray(theta), jnp.asarray(xu_p),
+                              jnp.asarray(w_p))
+        return float(v), np.asarray(g)
+
+    res = minimize(fg, theta0, jac=True, method="L-BFGS-B",
+                   bounds=[(-12, 18), (-25, 25), (-12, 12)],
+                   options=dict(maxiter=300))
+    a, c, d = res.x
+    return (float(np.exp(a)), float(1 / (1 + np.exp(-c))),
+            float(1 / (1 + np.exp(-d))), float(res.fun))
+
+
+@dataclass
+class InitResult:
+    R: np.ndarray            # (S, B)
+    P: np.ndarray            # (S, B)
+    Pi: np.ndarray           # (N, B)
+    Lambda: np.ndarray       # (S,)
+    phi: float
+    r_hat: np.ndarray = field(repr=False, default=None)   # marginal ZINB fits
+    p_hat: np.ndarray = field(repr=False, default=None)
+    pi0_hat: np.ndarray = field(repr=False, default=None)
+    r_ref: float = None
+    p_ref: float = None
+
+
+def initialize(X, mask=None, lambda_init=50.0, verbose=False):
+    """Starting values (R0, P0, Pi0, Lambda0, phi0) for ``fit_normal``.
+
+    ``lambda_init`` mostly sets the resolution of the latent Poisson; R0
+    absorbs the scale.
+    """
+    X = np.asarray(X, dtype=np.float64)
+    N, S, B = X.shape
+    nan_mask = ~np.isnan(X)
+    mask = nan_mask if mask is None else (np.asarray(mask, bool) & nan_mask)
+    Xz = np.where(mask, np.nan_to_num(X), 0.0)
+
+    # ---- 1. per-column ZINB fits -------------------------------------------
+    r_hat = np.zeros((S, B))
+    p_hat = np.zeros((S, B))
+    pi0_hat = np.zeros((S, B))
+    for s in range(S):
+        for b in range(B):
+            xs = Xz[:, s, b][mask[:, s, b]]
+            r_hat[s, b], p_hat[s, b], pi0_hat[s, b], nll = fit_zinb(xs)
+            if verbose:
+                print(f"  ZINB[{s},{b}]: r={r_hat[s,b]:.3f} "
+                      f"p={p_hat[s,b]:.4f} pi0={pi0_hat[s,b]:.4f} "
+                      f"nll={nll:.4f}")
+
+    # ---- 2. W2 barycenter reference NB -------------------------------------
+    means = r_hat * (1 - p_hat) / p_hat
+    sds = np.sqrt(means / p_hat)
+    mu_ref = float(means.mean())
+    sd_ref = float(sds.mean())
+    var_ref = max(sd_ref ** 2, mu_ref * (1 + 1e-6))
+    p_ref = np.clip(mu_ref / var_ref, 1e-8, 1 - 1e-8)
+    r_ref = mu_ref * p_ref / (1 - p_ref)
+
+    # ---- 3. quantile transform each column onto the reference --------------
+    Z = np.zeros((N, S, B))
+    for s in range(S):
+        for b in range(B):
+            obs = mask[:, s, b]
+            xu, inv = np.unique(Xz[:, s, b], return_inverse=True)
+            u = (sp_nbinom.cdf(xu, r_hat[s, b], p_hat[s, b])
+                 - 0.5 * sp_nbinom.pmf(xu, r_hat[s, b], p_hat[s, b]))
+            u = np.clip(u, 1e-12, 1 - 1e-12)
+            z = sp_nbinom.ppf(u, r_ref, p_ref).astype(np.float64)[inv]
+            fill = z[obs].mean() if obs.any() else 0.0
+            Z[:, s, b] = np.where(obs, z, fill)
+
+    # ---- 4. normalize -> Pi0 ------------------------------------------------
+    Z = Z / np.maximum(Z.sum(axis=0, keepdims=True), 1e-300)
+    Znb = Z.sum(axis=1)                                     # (N, B)
+    row = Znb.sum(axis=-1, keepdims=True)
+    Pi0 = np.where(row > 0, Znb / np.maximum(row, 1e-300), 1.0 / B)
+    Pi0 = np.clip(Pi0, 1e-8, 1.0)
+    Pi0 = Pi0 / Pi0.sum(axis=-1, keepdims=True)
+
+    # ---- 5. Lambda0, R0, P0, phi0 ------------------------------------------
+    Lambda0 = np.broadcast_to(np.asarray(lambda_init, dtype=np.float64),
+                              (S,)).copy()
+    pi_bar = Pi0.mean(axis=0)                               # (B,)
+    R0 = r_hat / (Lambda0[:, None] * pi_bar[None, :])
+    P0 = np.clip((1.0 + R0) / (1.0 / p_hat + R0), 1e-6, 1 - 1e-6)
+    all_zero = (Xz * mask).sum(axis=(1, 2)) == 0
+    phi0 = float(np.clip(all_zero.mean(), 1e-4, 0.5))
+
+    return InitResult(R=R0, P=P0, Pi=Pi0, Lambda=Lambda0, phi=phi0,
+                      r_hat=r_hat, p_hat=p_hat, pi0_hat=pi0_hat,
+                      r_ref=float(r_ref), p_ref=float(p_ref))
+
+
+def init_from_fit(fit):
+    """Build an ``InitResult`` from a previous fit, for warm starts.
+
+    The fitted profile Pi is enough to recover (mu, sigma); R/P/phi carry the
+    global block.  Useful when refitting the same data under a different
+    abundance model.
+    """
+    get = (lambda k: fit[k]) if isinstance(fit, dict) \
+        else (lambda k: getattr(fit, k))
+    Pi = np.asarray(get("Pi"), float).copy()
+    Pi[~np.isfinite(Pi).all(1)] = 1.0 / Pi.shape[1]     # fully-missing rows
+    return InitResult(R=np.asarray(get("R")), P=np.asarray(get("P")), Pi=Pi,
+                      Lambda=np.asarray(get("rates")).reshape(-1),
+                      phi=float(get("phi")))

ebin/model.py

@@ -0,0 +1,186 @@
+"""Marginal likelihood of the compound NB-Poisson count model.
+
+    T[n,s,b]            ~ Poisson(a_n * Pi[n,b] * lambda_s)     latent cells
+    X[n,s,b] | T[n,s,b] ~ NB(R[s,b] * T[n,s,b], P[s,b])         reads
+
+Pi[n, :] is the object's bin profile (a simplex row), a_n its abundance and
+lambda_s the per-replicate latent level.  The reads of a cell are a sum over the
+cells that landed in bin b, so the emission compounds with the Poisson: the
+marginal pmf is a sum over the latent count tau, truncated at the point where
+the Poisson tail is below ``tail_eps``.  NB tables are evaluated on the unique
+count values of each channel and gathered back.
+
+lambda is not identified upward -- it rides the R*lambda ridge, and only
+products like c*Pi*lambda with c = R(1-P)/P are identified -- so it is held
+fixed (``lambda_fix``) and R absorbs the scale.
+"""
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+from jax.scipy.special import gammaln, logsumexp
+from scipy.stats import poisson as _sp_poisson
+
+from .gauss_rules import compute_nodes_and_logweights, Rules
+
+jax.config.update("jax_enable_x64", True)
+
+_NEG_INF = -jnp.inf
+
+
+def nb_logpmf(x, r, p):
+    """log NB(x | r, p), scipy convention (mean r(1-p)/p).  r may be
+    non-integer; r == 0 is the point mass at 0."""
+    r_safe = jnp.maximum(r, 1e-300)
+    lp = (gammaln(x + r_safe) - gammaln(r_safe) - gammaln(x + 1.0)
+          + r_safe * jnp.log(p) + x * jnp.log1p(-p))
+    degenerate = jnp.where(x == 0, 0.0, _NEG_INF)
+    return jnp.where(r > 0, lp, degenerate)
+
+
+def poisson_trunc_bound(lam_max, tail_eps=1e-10, slack=5):
+    """Smallest T with P(Poisson(lam_max) > T) < tail_eps, plus slack."""
+    return int(_sp_poisson.isf(tail_eps, max(lam_max, 1e-6))) + 1 + slack
+
+
+def gamma_mixture_rule(alpha, mean, num_nodes):
+    """Gamma(alpha, theta=mean/alpha) discretized on ``num_nodes`` generalized
+    Gauss-Laguerre nodes.  Differentiable in alpha through the custom-JVP
+    tridiagonal eigensolver in gauss_rules.  Returns (nodes, log weights)."""
+    nodes, logw = compute_nodes_and_logweights(
+        num_nodes, alpha - 1.0, rule=Rules.GenLaguerre, norm=True)
+    return nodes * (mean / alpha), logw
+
+
+@jax.checkpoint
+def _logmarg_one_rate(log_em_n, mu_n, tau, lg):
+    """(N,) log sum_tau NB(x | R*tau, P) * Poisson(tau | mu_n) for one rate.
+
+    Checkpointed: the (N, T) intermediates are recomputed in the backward pass
+    instead of being stored."""
+    log_pois = (tau[None, :]
+                * jnp.log(jnp.maximum(mu_n, 1e-300))[:, None]
+                - mu_n[:, None] - lg[None, :])
+    return logsumexp(log_em_n + log_pois, axis=1)
+
+
+class LoglikBuilder:
+    """Precomputes the X-dependent tables and exposes traceable likelihoods.
+
+    X : (N, S, B) counts (NaN = missing), mask : optional (N, S, B) bool.
+    rate_max : upper bound on any latent rate; sizes the tau grid and must
+        dominate the largest rate the optimizer can reach (a_max * lambda when
+        an abundance is modelled).
+    """
+
+    def __init__(self, X, mask=None, rate_max=200.0, tail_eps=1e-10):
+        X = np.asarray(X)
+        if X.ndim != 3:
+            raise ValueError("X must have shape (N, S, B)")
+        nan_mask = ~np.isnan(X) if np.issubdtype(X.dtype, np.floating) else \
+            np.ones(X.shape, bool)
+        mask = nan_mask if mask is None else (np.asarray(mask, bool) & nan_mask)
+        Xi = np.where(mask, np.nan_to_num(X, nan=0.0), 0.0)
+        if not np.allclose(Xi, np.round(Xi)):
+            raise ValueError("X must contain (masked) non-negative integers")
+        Xi = np.round(Xi).astype(np.int64)
+        if (Xi < 0).any():
+            raise ValueError("X must be non-negative")
+
+        self.N, self.S, self.B = Xi.shape
+        self.rate_max = float(rate_max)
+
+        # per-channel unique values: each (s,b) NB table only covers the counts
+        # that actually occur in that channel
+        self.xu = [[None] * self.B for _ in range(self.S)]
+        self.inv = [[None] * self.B for _ in range(self.S)]
+        for s in range(self.S):
+            for b in range(self.B):
+                xu, inv = np.unique(Xi[:, s, b], return_inverse=True)
+                self.xu[s][b] = jnp.asarray(xu, dtype=jnp.float64)
+                self.inv[s][b] = jnp.asarray(inv)
+        self.mask = jnp.asarray(mask)
+        self.X = Xi
+        self.all_zero = jnp.asarray((Xi * mask).sum(axis=(1, 2)) == 0)
+        self.n_observed = int(mask.sum())
+
+        T_full = poisson_trunc_bound(rate_max, tail_eps)
+        self.tau_full = jnp.arange(T_full, dtype=jnp.float64)
+        self.lg_tau_full = gammaln(self.tau_full + 1.0)
+
+    def _log_nb_gathered(self, s, b, r, p, tau):
+        """(N, T) log NB(x_nsb | r * tau, p) via the unique-value table."""
+        log_nb = nb_logpmf(self.xu[s][b][:, None], r * tau[None, :], p)
+        return log_nb[self.inv[s][b]]
+
+    def _acc_over_rates(self, s, tbls, Pi, rates_s):
+        """(N, K) sum over b of the masked log-marginals, one column per rate.
+
+        Scanned with a rematerialized body so only one rate's (N, T)
+        intermediates are live at a time; a plain loop over k lets XLA allocate
+        all K concurrently (tens of GB at N=30000).
+        """
+        tau, lg = self.tau_full, self.lg_tau_full
+        mask_s = [self.mask[:, s, b] for b in range(self.B)]
+
+        def body(carry, rate):
+            acc_k = jnp.zeros(self.N, dtype=jnp.float64)
+            for b in range(self.B):
+                lm = _logmarg_one_rate(tbls[b], Pi[:, b] * rate, tau, lg)
+                acc_k = acc_k + jnp.where(mask_s[b], lm, 0.0)
+            return carry, acc_k
+
+        _, accT = jax.lax.scan(jax.checkpoint(body), None, rates_s)   # (K, N)
+        return accT.T
+
+    def _nb_tables(self, s, R, P):
+        return [self._log_nb_gathered(s, b, R[s, b], P[s, b], self.tau_full)
+                for b in range(self.B)]
+
+    @staticmethod
+    def _zero_inflate(inner, phi, all_zero):
+        phi_c = jnp.clip(phi, 0.0, 1.0 - 1e-12)
+        log_phi = jnp.log(jnp.maximum(phi_c, 1e-300))
+        log_1mphi = jnp.log1p(-phi_c)
+        with_zero = jnp.logaddexp(log_phi, log_1mphi + inner)
+        return jnp.where(all_zero, with_zero, log_1mphi + inner)
+
+    def object_loglik(self, R, P, Pi, rates, log_wmix, phi=0.0):
+        """(N,) per-object log-likelihood.
+
+        Pi : (N, B) rows fed to the latent Poisson (a_n * profile when an
+             abundance is fitted).  rates / log_wmix : (S, K) latent rates and
+             their log-weights (K = 1 for the fixed-rate model).  phi : object
+             level zero-inflation probability.
+        """
+        rates = jnp.atleast_2d(rates)
+        log_wmix = jnp.atleast_2d(log_wmix)
+        inner = jnp.zeros(self.N, dtype=jnp.float64)
+        for s in range(self.S):
+            acc = self._acc_over_rates(s, self._nb_tables(s, R, P), Pi,
+                                       rates[s])                     # (N, K)
+            inner = inner + logsumexp(acc + log_wmix[s][None, :], axis=1)
+        return self._zero_inflate(inner, phi, self.all_zero)
+
+    def object_loglik_abund(self, R, P, Pi, base_lambda, a_nodes, log_wa,
+                            phi=0.0):
+        """(N,) per-object log-likelihood with the abundance integrated/marginalized out:
+
+            a_n ~ sum_k exp(log_wa[k]) delta(a_nodes[k]),
+            T[n,s,b] | a_n ~ Poisson(a_n * Pi[n,b] * base_lambda[s]).
+
+        a_n couples every cell of the object, so the mixture is resolved
+        (logsumexp over nodes) only AFTER summing the per-cell log-marginals
+        over both s and b -- unlike ``object_loglik``, whose rate mixture is
+        resolved within each replicate.  Nodes past the tau grid self-truncate;
+        they carry negligible weight, and truncating only shrinks their
+        contribution, so the bound is conservative.
+        """
+        a_nodes = jnp.asarray(a_nodes)
+        base_lambda = jnp.atleast_1d(jnp.asarray(base_lambda))
+        acc_total = jnp.zeros((self.N, a_nodes.shape[0]), dtype=jnp.float64)
+        for s in range(self.S):
+            acc_total = acc_total + self._acc_over_rates(
+                s, self._nb_tables(s, R, P), Pi, a_nodes * base_lambda[s])
+        inner = logsumexp(acc_total + log_wa[None, :], axis=1)
+        return self._zero_inflate(inner, phi, self.all_zero)