Fix ZTNB EM numerical stability: initialize new_inner_logL and cap initw

[chrishaak]

Hi Francis - you might remember the missing 'innerlogL' error that I receive sporadically - this is an attempt at addressing it. I want to be certain that you think this is actually a reasonable thing to do - in particular the capping of infinite probabilty values. And similarly that you think that the selected cap value is reasonable (if indeed you think that the overall approach is reasonable at all). Details below:

Two related numerical stability fixes in the ZTNB EM fitting loop:

1. Initialize new_inner_logL <- -Inf before each ZTNB EM while-loop (initfit_fn and update_Xcoefsspp_fn). The variable is only assigned inside the loop body after a successful gam() fit. If gam() fails and breaks on the first iteration, new_inner_logL is never created, causing 'object new_inner_logL not found' when referenced after the loop.

2. Cap EM augmentation weights (initw) at 100 after each E-step in three locations. initw = P(Y=0)/P(Y>0) diverges to Inf when the NB model assigns near-certainty to Y=0, causing the augmented gam() to produce degenerate estimates. A cap of 100 (~99% posterior probability) is already an extremely strong signal and prevents instability with no effect on well-behaved fits.

Background: How CBFM Fits the ZTNB

The zero-truncated negative binomial (ZTNB) distribution has no native gam() family in
mgcv, so CBFM fits it via an EM algorithm that exploits the relationship between the
zero-truncated NB and the full NB. The key insight is that ZTNB data can be viewed as full NB data where zeros were generated
but never observed. The EM algorithm "fills in" those missing zeros iteratively:

• E-step: For each positive observation, compute initw — the posterior odds ratio
  P(Y=0) / P(Y>0) under the current NB parameters. This represents how many "missing"
  zeros should be imputed per observed positive. For example, if the model estimates
  P(Y=0) = 0.8 at some covariate combination, then initw = 0.8/0.2 = 4: four missing
  zeros are imputed for every observed positive there.

• M-step: Fit a standard gam(..., family = nb()) to an augmented dataset of 2n
  rows — the original observations stacked on top of n synthetic zeros — using initw
  as weights on the imputed zero rows. This is equivalent to maximising the ZTNB likelihood
  on the positive observations alone.

• Iterate until convergence.

This EM loop appears in two functions:

• initfit_fn — computes per-species starting values before the outer PQL loop
• update_Xcoefsspp_fn — performs per-species coefficient updates within each PQL iteration

---

Bug 1: Uninitialized Variable new_inner_logL

Symptom

Error in initfit_fn(j = j, formula = formula, control = control) :
  task N failed - "object 'new_inner_logL' not found"

Root Cause

Inside the ZTNB EM loop, the gam() call at each M-step is wrapped in try(). If this
call fails (e.g. due to a degenerate species — sparse positive counts, no intercept in the
formula, or extreme observation weights), the code executes break to exit the loop early.

The variable new_inner_logL is only assigned inside the loop body, after a successful
gam() fit. If the loop breaks on the very first iteration, new_inner_logL is never
created. The code immediately after the loop then attempts:

fit0$logLik <- new_inner_logL   # crashes: object not found

Possible Fix?

Initialize new_inner_logL <- -Inf immediately before each ZTNB EM while loop. This
ensures the variable always exists when the post-loop code runs. A value of -Inf is the
correct sentinel: it signals that no valid log-likelihood was obtained, allowing downstream
code to handle the failure gracefully (e.g. falling back to random starting values in
initfit_fn).

Locations — search strings to find the insertion points:

Location	Function	Search for line above insertion
1	initfit_fn	cw_inner_logL <- -Inf (first ZTNB EM block)
2	update_Xcoefsspp_fn	cw_inner_logL <- -Inf (ZTNB EM block)

Change at each location:

# Before (original):
inner_err <- Inf
inner_inner_counter <- 0
cw_inner_logL <- -Inf
while(inner_err > 1e-4) {

# After (fixed):
inner_err <- Inf
inner_inner_counter <- 0
cw_inner_logL <- -Inf
new_inner_logL <- -Inf      # initialize before loop in case gam() fails and breaks early
while(inner_err > 1e-4) {

Behavioral impact: None on the normal code path. Only affects the failure case, where
previously the code crashed;

---

Bug 2: Infinite EM Augmentation Weights (initw)

Symptom

The augmented gam() inside the ZTNB EM loop hangs, produces NaN coefficients, or
causes downstream numerical failures. No explicit error message — the fit silently produces
garbage or the outer PQL loop diverges.

Root Cause

The E-step weight is:

initw <- dnbinom(0, mu = exp(cw_eta), size = theta) /
         (1 - dnbinom(0, mu = exp(cw_eta), size = theta))

This is P(Y=0) / P(Y>0). When the NB model assigns near-certainty to Y=0 at a
particular covariate combination (i.e. P(Y=0) → 1), this ratio diverges to Inf.

An Inf weight in the augmented gam() means one synthetic zero row has infinite
influence, completely dominating the M-step fit. The IRLS solver inside gam() then
produces degenerate or NaN coefficient estimates.

This becomes more likely when:

• Theta is small (heavy overdispersion): the NB assigns high probability to zero across
  a wide range of means, so initw is large for many observations
• Observation weights are low: downweighting positive observations in a sparse species
  pulls the fitted mean lower, further increasing P(Y=0) and initw
• The two effects compound: observation weights × EM weights, so a species that was
  borderline stable without observation weights tips over to Inf with them

Fix

Cap initw at a finite maximum immediately after each E-step computation:

initw <- pmin(initw, 100)

A cap of 100 means the augmented gam() sees at most 100 synthetic zeros per observed
positive — corresponding to a posterior probability of ~99% that the observation "should
not exist." This is already an extremely strong signal; allowing initw to exceed this
provides no meaningful information and introduces severe numerical instability.

The cap value of 100 was chosen conservatively (higher than strictly necessary) to avoid
any effect on well-behaved fits. In practice, initw > 20 is already rare; initw > 100
essentially never occurs except in the degenerate cases this fix is designed to handle.

Locations — search strings to find the insertion points:

Location	Function	Search for line above insertion
1	initfit_fn	initw <- dnbinom(0, mu = exp(cw_eta), size = fit0$family$getTheta(TRUE)) / (1 - dnbinom(... (before EM loop)
2	initfit_fn	initw <- dnbinom(0, mu = exp(cw_eta), size = 1/fit0$dispparam) / (1 - dnbinom(... (inside EM loop)
3	update_Xcoefsspp_fn	initw <- initw / (1 - initw) (before initw <- initw[find_nonzeros])

Change at each location:

# Before (original):
initw <- dnbinom(0, ...) / (1 - dnbinom(0, ...))
w1 <- ...   # or initw <- initw[find_nonzeros]

# After (fixed):
initw <- dnbinom(0, ...) / (1 - dnbinom(0, ...))
initw <- pmin(initw, 100)  # cap posterior odds to prevent Inf when P(Y=0) ≈ 1
w1 <- ...   # or initw <- initw[find_nonzeros]

Behavioral impact: None on well-behaved fits where initw is finite and reasonable.
Only affects cases where initw would have been Inf or very large, converting a silent
numerical failure into a stable (if imprecise) fit.

---

Relationship Between the Two Bugs

Both bugs share a common upstream cause: the gam() call inside the ZTNB EM loop
encountering degenerate inputs. Bug 2 (initw = Inf) is one specific mechanism that can
cause the gam() to fail, which then triggers Bug 1 (new_inner_logL uninitialized) when
the loop breaks early. In practice:

• Bug 1 alone can be triggered without Bug 2 (e.g. by a formula with no intercept)
• Bug 2 alone can cause silent numerical corruption without triggering Bug 1 (if the
  augmented gam() completes but with degenerate weights)
• Both can occur together when low observation weights drive initw to Inf

Fixing both together provides defense in depth: Bug 2's fix prevents the degenerate input;
Bug 1's fix ensures graceful failure even if some other pathway produces a degenerate
gam() in the future.

---

Summary of Changes

#	Function	Change	Lines affected
1	initfit_fn	Add new_inner_logL <- -Inf before ZTNB EM loop	1 line inserted
2	update_Xcoefsspp_fn	Add new_inner_logL <- -Inf before ZTNB EM loop	1 line inserted
3	initfit_fn	Add initw <- pmin(initw, 100) after first E-step	1 line inserted
4	initfit_fn	Add initw <- pmin(initw, 100) after second E-step (inside loop)	1 line inserted
5	update_Xcoefsspp_fn	Add initw <- pmin(initw, 100) after E-step	1 line inserted

Total: 5 single-line insertions, no deletions, no changes to function signatures or
existing logic.

[fhui28]

Bug 1 has been resolved. I will think about bug 2 a bit more