R/family_loglik.R

In tulpa: Templated Unified Library for Posterior Approximation in Bayesian Hierarchical Models

# family_loglik.R
# ------------------------------------------------------------------------------
# Single source of truth for the R-level family math: the inverse-link mean,
# the log-likelihood, the score (d log-lik / d eta), and the Laplace/IRLS
# working weight (-d^2 log-lik / d eta^2 under the canonical/expected form).
#
# Each family is one entry in `.FAMILY_OPS`, so adding a family is a single
# registry entry -- no per-family branches scattered across fitters. The
# Laplace Newton loop (`glmm_weights()`), the sampler-path log-posterior
# builder, and any diagnostics all read from here, so the mean function and the
# weights cannot drift between code paths.
#
# Parameterization (matches the character-named families the design fitters
# already use): `eta` is the linear predictor, `phi` the dispersion/precision
# (residual variance for gaussian, size for neg_binomial_2, precision for beta),
# `n_trials` the binomial denominator (default 1).
# ------------------------------------------------------------------------------

# --- shared, clamped inverse-link means (used by loglik / score / weight) -----

.mean_binomial <- function(eta, ...) {
  mu <- 1 / (1 + exp(-eta))
  pmin(pmax(mu, 1e-8), 1 - 1e-8)
}
.mean_log <- function(eta, ...) pmax(exp(eta), 1e-8)   # poisson, neg_binomial_2
.mean_identity <- function(eta, ...) eta               # gaussian
.mean_beta <- function(eta, ...) {
  mu <- 1 / (1 + exp(-eta))
  pmin(pmax(mu, 1e-7), 1 - 1e-7)
}

# log(1 - exp(x)) for x <= 0, split at -log(2) so neither branch cancels
# (Maechler 2012): log(-expm1(x)) loses precision as x -> -inf, log1p(-exp(x))
# loses it as x -> 0. Used by the zero-truncated families to subtract
# log P(Y > 0). Returns -Inf at x = 0, the correct log of zero mass.
.log1mexp <- function(x) {
  ifelse(x > -log(2), log(-expm1(x)), log1p(-exp(x)))
}

# Send the log-density to -Inf outside a zero-truncated family's support. The
# truncated formulas evaluate to a finite number at y = 0 rather than vanishing,
# so without this a caller reaching the density directly -- bypassing the front
# door's support check -- would read a plausible value where the family has no
# mass. `y` is recycled against the result so a scalar response broadcasts.
.zero_truncate <- function(ll, y) {
  ll[rep_len(y, length(ll)) < 1] <- -Inf
  ll
}

# One draw per element from a count distribution conditioned on y > 0, by
# inverse CDF on the truncated support: with p0 = P(Y = 0), a uniform draw on
# (p0, 1) pushed through the untruncated quantile function lands on y >= 1.
# `qfun` is stats::qpois / stats::qnbinom; `...` carries its parameters.
.rtrunc_count <- function(qfun, p0, ...) {
  n <- length(p0)
  qfun(p0 + stats::runif(n) * (1 - p0), ...)
}

# Inverse-Gaussian sampler (Michael, Schucany & Haas 1976), parameterized to
# match the `inverse_gaussian` loglik below: mean mu, shape lambda = 1 / phi.
.rinvgauss <- function(n, mu, lambda) {
  y <- stats::rnorm(n)^2
  x <- mu + mu^2 * y / (2 * lambda) -
    mu / (2 * lambda) * sqrt(4 * mu * lambda * y + mu^2 * y^2)
  u <- stats::runif(n)
  ifelse(u <= mu / (mu + x), x, mu^2 / x)
}

# --- family operations registry ----------------------------------------------
# Fields per family:
#   mean(eta, phi)                inverse link, clamped; phi reaches entries
#                                 whose response mean depends on the dispersion
#                                 (lognormal), the rest swallow it via `...`
#   loglik(eta, y, n_trials, phi) elementwise log-density (full normalizer)
#   score(eta, y, n_trials, phi)  elementwise d log-lik / d eta
#   weight(eta, n_trials, phi)    Laplace/IRLS working weight (no y dependence)
#   sample(eta, n_trials, phi)    one y draw per element (posterior predictive)
#   variance(eta, n_trials, phi)  response variance Var(y | eta), elementwise
#   response_mean(eta, n_trials, phi)
#                                 E[y | eta] on the response scale; only where
#                                 it differs from mean() (trial-scaled and
#                                 zero-truncated families)
# Families in .PHI2_FAMILIES additionally take a trailing `phi2` (the second
# dispersion channel: Student-t df); the wrappers pass it only when supplied.

.FAMILY_OPS <- list(
  binomial = list(
    mean = .mean_binomial,
    loglik = function(eta, y, n_trials, phi) {
      n <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      mu <- .mean_binomial(eta)
      lchoose(n, y) + y * log(mu) + (n - y) * log1p(-mu)
    },
    score = function(eta, y, n_trials, phi) {
      n <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      y - n * .mean_binomial(eta)
    },
    weight = function(eta, n_trials, phi) {
      n <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      mu <- .mean_binomial(eta)
      n * mu * (1 - mu)
    },
    sample = function(eta, n_trials, phi) {
      n <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      stats::rbinom(length(eta), size = n, prob = .mean_binomial(eta))
    },
    variance = function(eta, n_trials, phi) {
      n <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      mu <- .mean_binomial(eta)
      n * mu * (1 - mu)
    },
    response_mean = function(eta, n_trials, phi) {
      n <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      n * .mean_binomial(eta)
    }
  ),

  poisson = list(
    mean = .mean_log,
    loglik = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      y * log(mu) - mu - lgamma(y + 1)
    },
    score = function(eta, y, n_trials, phi) y - .mean_log(eta),
    weight = function(eta, n_trials, phi) .mean_log(eta),
    sample = function(eta, n_trials, phi) {
      stats::rpois(length(eta), .mean_log(eta))
    },
    variance = function(eta, n_trials, phi) .mean_log(eta)
  ),

  neg_binomial_2 = list(
    mean = .mean_log,
    loglik = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      lgamma(y + phi) - lgamma(phi) - lgamma(y + 1) +
        phi * (log(phi) - log(phi + mu)) +
        y * (log(mu) - log(mu + phi))
    },
    score = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      phi * (y - mu) / (mu + phi)
    },
    weight = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      mu * phi / (mu + phi)
    },
    # Observed curvature -d^2 loglik / d eta^2 at the realized y. Averaging it
    # over y (E[y] = mu) returns the registered `weight`, so the two agree in
    # expectation and differ per observation.
    obs_weight = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      (y + phi) * phi * mu / (mu + phi)^2
    },
    sample = function(eta, n_trials, phi) {
      stats::rnbinom(length(eta), size = phi, mu = .mean_log(eta))
    },
    variance = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      mu + mu^2 / phi
    }
  ),

  # Negative binomial type 1 (log link), phi = dispersion. Mean mu = exp(eta),
  # variance mu (1 + phi) -- linear in the mean, unlike neg_binomial_2's
  # quadratic mu + mu^2/phi. In NB(r, p) form the shape is r = mu / phi and
  # p = 1 / (1 + phi), so the shape grows with the mean and the
  # variance-to-mean ratio stays constant at 1 + phi.
  neg_binomial_1 = list(
    mean = .mean_log,
    loglik = function(eta, y, n_trials, phi) {
      r <- .mean_log(eta) / phi
      lgamma(y + r) - lgamma(r) - lgamma(y + 1) -
        (y + r) * log1p(phi) + y * log(phi)
    },
    score = function(eta, y, n_trials, phi) {
      r <- .mean_log(eta) / phi
      r * (digamma(y + r) - digamma(r) - log1p(phi))
    },
    # Quasi-likelihood IRLS weight (dmu/deta)^2 / V(mu) = mu / (1 + phi), the
    # form glmmTMB uses. Unlike neg_binomial_2 -- where the moment weight
    # phi mu / (mu + phi) coincides with the expected curvature -- this is NOT
    # the Fisher information. Because the shape r = mu / phi moves with the
    # mean, the score carries digammas and the exact expected curvature is
    #   I(eta) = r^2 (trigamma(r) - E[trigamma(y + r)])
    #          = r^2 sum_{k >= 0} P(y > k) / (r + k)^2,
    # a convergent series with no elementary closed form. The moment weight
    # understates it by 0.4% (phi = 0.5, mu = 20) to 58% (phi = 8, mu = 0.37),
    # so the Laplace curvature is correspondingly soft and marginal SEs run
    # wide in that corner of the parameter space.
    weight = function(eta, n_trials, phi) .mean_log(eta) / (1 + phi),
    # Observed curvature at the realized y. Differentiating the score
    # s = r (psi(y + r) - psi(r) - log1p(phi)) with dr/deta = r gives
    #   -ds/deta = -s + r^2 (trigamma(r) - trigamma(y + r)).
    # Its expectation is the exact information r^2 (trigamma(r) -
    # E[trigamma(y + r)]) discussed under `weight`, not the registered
    # moment weight -- so unlike neg_binomial_2 the two do not agree even
    # on average.
    obs_weight = function(eta, y, n_trials, phi) {
      r <- .mean_log(eta) / phi
      s <- r * (digamma(y + r) - digamma(r) - log1p(phi))
      -s + r^2 * (trigamma(r) - trigamma(y + r))
    },
    sample = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      stats::rnbinom(length(eta), size = mu / phi, mu = mu)
    },
    variance = function(eta, n_trials, phi) .mean_log(eta) * (1 + phi)
  ),

  # Zero-truncated Poisson (log link), support y >= 1. The untruncated
  # log-density minus log P(Y > 0) = log(1 - exp(-mu)). The conditional mean
  # E[y | y > 0] = mu / p exceeds mu, so `mean` stays the inverse link and
  # `response_mean` carries the truncated mean.
  truncated_poisson = list(
    mean = .mean_log,
    loglik = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      .zero_truncate(y * log(mu) - mu - lgamma(y + 1) - .log1mexp(-mu), y)
    },
    score = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      y - mu / -expm1(-mu)
    },
    weight = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      mu_c <- mu / -expm1(-mu)
      mu_c - mu_c * mu_c * exp(-mu)
    },
    sample = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      .rtrunc_count(stats::qpois, exp(-mu), lambda = mu)
    },
    # Var(y | y > 0) equals the working weight here: the curvature carries no
    # y dependence, so observed and expected information coincide.
    variance = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      mu_c <- mu / -expm1(-mu)
      mu_c - mu_c * mu_c * exp(-mu)
    },
    response_mean = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      mu / -expm1(-mu)
    }
  ),

  # Zero-truncated negative binomial type 2 (log link), phi = size, support
  # y >= 1. The NB2 log-density minus log P(Y > 0), with
  # P(Y > 0) = 1 - (phi / (phi + mu))^phi = 1 - exp(-phi log1p(mu/phi)).
  truncated_neg_binomial_2 = list(
    mean = .mean_log,
    loglik = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      .zero_truncate(
        lgamma(y + phi) - lgamma(phi) - lgamma(y + 1) +
          phi * (log(phi) - log(phi + mu)) +
          y * (log(mu) - log(mu + phi)) -
          .log1mexp(-phi * log1p(mu / phi)),
        y)
    },
    score = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      p  <- -expm1(-phi * log1p(mu / phi))
      phi * (y - mu / p) / (phi + mu)
    },
    # Expected (Fisher) form w0 (1 - s), matching the registry's y-free
    # contract. This is what Newton uses on both paths: it is positive for
    # every mu, so the Hessian stays PD without a fallback. The observed
    # curvature carries a y term and so cannot be expressed here; it lives in
    # `obs_weight` below, and the zero-inflation mixture uses that one.
    weight = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      a  <- phi * log1p(mu / phi)
      p  <- -expm1(-a)
      w0 <- phi * mu / (p * (phi + mu))
      w0 * (1 - exp(-a) * w0)
    },
    # Observed curvature at the realized y. The truncated density is
    # loglik_nb2 - log P(Y > 0), so the curvature splits as the untruncated
    # observed curvature plus the second derivative of log p, where
    # p = 1 - exp(-a) and a = phi log1p(mu / phi).
    obs_weight = function(eta, y, n_trials, phi) {
      mu  <- .mean_log(eta)
      a   <- phi * log1p(mu / phi)
      p   <- -expm1(-a)
      da  <- phi * mu / (phi + mu)
      d2a <- phi^2 * mu / (phi + mu)^2
      dp  <- exp(-a) * da
      d2p <- exp(-a) * (d2a - da^2)
      (y + phi) * phi * mu / (mu + phi)^2 + (d2p * p - dp^2) / p^2
    },
    sample = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      p0 <- exp(phi * (log(phi) - log(phi + mu)))
      .rtrunc_count(stats::qnbinom, p0, size = phi, mu = mu)
    },
    variance = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      p  <- -expm1(-phi * log1p(mu / phi))
      (mu + mu^2 / phi + mu^2) / p - (mu / p)^2
    },
    response_mean = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      mu / -expm1(-phi * log1p(mu / phi))
    }
  ),

  # Lognormal (identity link on the log scale), phi = log-scale VARIANCE
  # (mirroring the gaussian entry; the SD-parameterized compiled kernels get
  # sqrt(phi) at the front-door boundary). log(y) ~ N(eta, phi); the -log(y)
  # Jacobian is part of the density. E[y] = exp(eta + phi/2).
  lognormal = list(
    mean = function(eta, phi = 1.0, ...) exp(eta + phi / 2),
    loglik = function(eta, y, n_trials, phi) {
      ly <- log(y)
      -ly - 0.5 * log(2 * pi * phi) - (ly - eta)^2 / (2 * phi)
    },
    score = function(eta, y, n_trials, phi) (log(y) - eta) / phi,
    weight = function(eta, n_trials, phi) rep(1 / phi, length(eta)),
    sample = function(eta, n_trials, phi) {
      stats::rlnorm(length(eta), meanlog = eta, sdlog = sqrt(phi))
    },
    variance = function(eta, n_trials, phi) {
      (exp(phi) - 1) * exp(2 * eta + phi)
    }
  ),

  gaussian = list(
    mean = .mean_identity,
    loglik = function(eta, y, n_trials, phi) {
      -0.5 * ((y - eta)^2 / phi + log(2 * pi * phi))
    },
    score = function(eta, y, n_trials, phi) (y - eta) / phi,
    # Laplace/IRLS working weight = -d^2 log-lik / d eta^2 = 1/phi (phi is the
    # residual variance). The dispersion MUST enter here: the marginal-SE Schur
    # (.marginal_H_beta_*) and the dense H_beta path are the only callers, and
    # both need the dispersion-correct curvature. A unit weight here treats phi
    # as 1 and inflates the gaussian fixed-effect SEs by sqrt(1/phi)

    weight = function(eta, n_trials, phi) rep(1 / phi, length(eta)),
    sample = function(eta, n_trials, phi) {
      stats::rnorm(length(eta), mean = eta, sd = sqrt(phi))
    },
    variance = function(eta, n_trials, phi) rep(phi, length(eta))
  ),

  beta = list(
    mean = .mean_beta,
    loglik = function(eta, y, n_trials, phi) {
      mu <- .mean_beta(eta)
      a <- mu * phi
      b <- (1 - mu) * phi
      lgamma(phi) - lgamma(a) - lgamma(b) +
        (a - 1) * log(y) + (b - 1) * log1p(-y)
    },
    score = function(eta, y, n_trials, phi) {
      mu <- .mean_beta(eta)
      a <- mu * phi
      b <- (1 - mu) * phi
      dmu <- mu * (1 - mu)
      phi * (log(y) - log1p(-y) - digamma(a) + digamma(b)) * dmu
    },
    # The expected (Fisher) weight phi^2 (psi'(a) + psi'(b)) dmu^2.
    weight = function(eta, n_trials, phi) {
      mu <- .mean_beta(eta)
      dmu <- mu * (1 - mu)
      tg <- trigamma(mu * phi) + trigamma((1 - mu) * phi)
      phi * phi * tg * dmu * dmu
    },
    # Observed curvature -d^2 loglik / d eta^2 at the realized y: the expected
    # weight minus the score term (1 - 2 mu) dmu g_mu, with g_mu the mu-score.
    # Averaging it over y returns `weight`. This is the weight the engine's
    # Newton builds H from for this family, so it is what a Hessian rebuilt in
    # R has to carry.
    obs_weight = function(eta, y, n_trials, phi) {
      mu <- .mean_beta(eta)
      a <- mu * phi
      b <- (1 - mu) * phi
      dmu <- mu * (1 - mu)
      g_mu <- phi * (log(y) - log1p(-y) - digamma(a) + digamma(b))
      phi * phi * (trigamma(a) + trigamma(b)) * dmu * dmu -
        (1 - 2 * mu) * dmu * g_mu
    },
    sample = function(eta, n_trials, phi) {
      mu <- .mean_beta(eta)
      stats::rbeta(length(eta), mu * phi, (1 - mu) * phi)
    },
    variance = function(eta, n_trials, phi) {
      mu <- .mean_beta(eta)
      mu * (1 - mu) / (phi + 1)
    }
  ),

  # Gamma (log link), phi = shape. Mean mu = exp(eta), variance mu^2 / phi.
  # Mirrors laplace_family_link.h so the R-closure backends agree with the C++
  # Laplace / ModelData-sampler path.
  gamma = list(
    mean = .mean_log,
    loglik = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      phi * log(phi) - lgamma(phi) + (phi - 1) * log(y) -
        phi * log(mu) - phi * y / mu
    },
    score = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      phi * (y - mu) / mu
    },
    weight = function(eta, n_trials, phi) rep(phi, length(eta)),
    sample = function(eta, n_trials, phi) {
      stats::rgamma(length(eta), shape = phi, rate = phi / .mean_log(eta))
    },
    variance = function(eta, n_trials, phi) .mean_log(eta)^2 / phi
  ),

  # Inverse Gaussian (log link), phi = dispersion. Mean mu = exp(eta),
  # variance phi * mu^3.
  inverse_gaussian = list(
    mean = .mean_log,
    loglik = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      -0.5 * log(2 * pi * phi * y^3) - (y - mu)^2 / (2 * phi * mu^2 * y)
    },
    score = function(eta, y, n_trials, phi) {
      mu <- .mean_log(eta)
      (y - mu) / (phi * mu^2)
    },
    weight = function(eta, n_trials, phi) {
      mu <- .mean_log(eta)
      1 / (phi * mu)
    },
    sample = function(eta, n_trials, phi) {
      .rinvgauss(length(eta), mu = .mean_log(eta), lambda = 1 / phi)
    },
    variance = function(eta, n_trials, phi) phi * .mean_log(eta)^3
  ),

  # Beta-binomial (logit link), phi = precision (a + b), n_trials = n.
  # mu = P(success) = a / (a + b); a = mu*phi, b = (1-mu)*phi. Overdispersed
  # binomial: Var(y) = n mu(1-mu) [1 + (n-1)/(phi+1)]. The score is exact; the
  # working weight is the moment-based Fisher weight n mu(1-mu) / D with the
  # overdispersion factor D (matching grad_hess_for_family in the C++ path).
  beta_binomial = list(
    mean = .mean_beta,
    loglik = function(eta, y, n_trials, phi) {
      n  <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      mu <- .mean_beta(eta); a <- mu * phi; b <- (1 - mu) * phi
      lchoose(n, y) + lgamma(y + a) + lgamma(n - y + b) - lgamma(n + a + b) -
        lgamma(a) - lgamma(b) + lgamma(a + b)
    },
    score = function(eta, y, n_trials, phi) {
      n  <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      mu <- .mean_beta(eta); a <- mu * phi; b <- (1 - mu) * phi
      dmu <- mu * (1 - mu)
      phi * (digamma(y + a) - digamma(a) - digamma(n - y + b) + digamma(b)) * dmu
    },
    weight = function(eta, n_trials, phi) {
      n  <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      mu <- .mean_beta(eta); D <- 1 + (n - 1) / (phi + 1)
      n * mu * (1 - mu) / D
    },
    sample = function(eta, n_trials, phi) {
      n  <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      mu <- .mean_beta(eta)
      p  <- stats::rbeta(length(eta), mu * phi, (1 - mu) * phi)
      stats::rbinom(length(eta), size = n, prob = p)
    },
    variance = function(eta, n_trials, phi) {
      n  <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      mu <- .mean_beta(eta)
      n * mu * (1 - mu) * (1 + (n - 1) / (phi + 1))
    },
    response_mean = function(eta, n_trials, phi) {
      n <- if (!is.null(n_trials)) n_trials else rep(1, length(eta))
      n * .mean_beta(eta)
    }
  ),

  # Tweedie compound Poisson-gamma (log link), 1 < p < 2: phi = dispersion,
  # phi2 = power p (REQUIRED -- no silent default). Mean mu = exp(eta),
  # variance phi * mu^p; exact zero mass exp(-lambda) with lambda =
  # mu^(2-p) / (phi (2-p)). The positive-y density is the compound
  # Poisson-gamma series (Dunn & Smyth 2005): N ~ Poisson(lambda) gamma
  # events with shape a = (2-p)/(p-1) and rate b = mu^(1-p) / (phi (p-1)),
  # log-sum-exp'd over N around its dominating term. The score and working
  # weight are the exponential-dispersion closed forms ((y - mu)/(phi V(mu))
  # through the log link), no series needed.
  tweedie = list(
    mean = .mean_log,
    loglik = function(eta, y, n_trials, phi, phi2 = NULL) {
      p <- .tweedie_power(phi2)
      .tweedie_loglik(.mean_log(eta), y, phi, p)
    },
    score = function(eta, y, n_trials, phi, phi2 = NULL) {
      p <- .tweedie_power(phi2)
      mu <- .mean_log(eta)
      (y - mu) / (phi * mu^(p - 1))
    },
    weight = function(eta, n_trials, phi, phi2 = NULL) {
      p <- .tweedie_power(phi2)
      .mean_log(eta)^(2 - p) / phi
    },
    sample = function(eta, n_trials, phi, phi2 = NULL) {
      p  <- .tweedie_power(phi2)
      mu <- .mean_log(eta)
      lam <- mu^(2 - p) / (phi * (2 - p))
      a   <- (2 - p) / (p - 1)
      b   <- mu^(1 - p) / (phi * (p - 1))
      n_ev <- stats::rpois(length(eta), lam)
      out <- numeric(length(eta))
      pos <- n_ev > 0
      out[pos] <- stats::rgamma(sum(pos), shape = n_ev[pos] * a,
                                rate = b[pos])
      out
    },
    variance = function(eta, n_trials, phi, phi2 = NULL) {
      phi * .mean_log(eta)^.tweedie_power(phi2)
    }
  ),

  # Student-t location-scale (identity link), phi = scale, df = phi2 through
  # the second dispersion channel (default .STUDENT_T_DF, the robust drop-in).
  # The score is exact and the working weight is the constant Fisher
  # information (nu+1)/((nu+3) phi^2). Mirrors the C++ `t` branch.
  t = list(
    mean = .mean_identity,
    loglik = function(eta, y, n_trials, phi, phi2 = NULL) {
      nu <- phi2 %||% .STUDENT_T_DF; r <- (y - eta) / phi
      lgamma((nu + 1) / 2) - lgamma(nu / 2) - 0.5 * log(nu * pi * phi^2) -
        0.5 * (nu + 1) * log1p(r^2 / nu)
    },
    score = function(eta, y, n_trials, phi, phi2 = NULL) {
      nu <- phi2 %||% .STUDENT_T_DF; resid <- y - eta
      (nu + 1) * resid / (nu * phi^2 + resid^2)
    },
    weight = function(eta, n_trials, phi, phi2 = NULL) {
      nu <- phi2 %||% .STUDENT_T_DF
      rep((nu + 1) / ((nu + 3) * phi^2), length(eta))
    },
    sample = function(eta, n_trials, phi, phi2 = NULL) {
      eta + phi * stats::rt(length(eta), df = phi2 %||% .STUDENT_T_DF)
    },
    variance = function(eta, n_trials, phi, phi2 = NULL) {
      nu <- phi2 %||% .STUDENT_T_DF
      rep(if (nu > 2) phi^2 * nu / (nu - 2) else Inf, length(eta))
    }
  )
)

# Families whose ops consume a second dispersion `phi2` (the Student-t degrees
# of freedom; the Tweedie power). Everything else rejects a supplied phi2
# loudly rather than silently ignoring it.
.PHI2_FAMILIES <- c("t", "tweedie")

# The Tweedie power. Required (a silently defaulted variance power would be a
# statistical decision the user never made) and restricted to the compound
# Poisson-gamma range.
.tweedie_power <- function(phi2) {
  if (is.null(phi2)) {
    stop("family = 'tweedie' requires `phi2` (the variance power p, ",
         "1 < p < 2).", call. = FALSE)
  }
  p <- as.numeric(phi2)
  if (!is.finite(p) || p <= 1 || p >= 2) {
    stop("The Tweedie power `phi2` must lie strictly in (1, 2); got ",
         format(p), ".", call. = FALSE)
  }
  p
}

# The compound Poisson-gamma series, enumerated once per observation and read
# two ways: `logsum` is log sum_n W_n, the normalizer of the positive-y density,
# and `nbar` is E_W[n], the weighted mean event count that the dispersion
# derivative needs. Every term carries phi only through lambda and the gamma rate
# b, each exactly 1/phi, so d log sum_n W_n / dphi = -(1 + a) E_W[n] / phi -- the
# same enumeration, no second truncation of the sum.
#
# Terms expand from the dominating index j_max = y^(2-p) / (phi (2-p))
# (Dunn & Smyth 2005), upward then downward, until they fall 37 nats below the
# running maximum.
.tweedie_series <- function(mu, y, phi, p) {
  mu  <- pmin(rep_len(mu, length(y)), 1e10)   # broadcast; clamp eta overflow
  lam <- mu^(2 - p) / (phi * (2 - p))
  a   <- (2 - p) / (p - 1)
  b   <- mu^(1 - p) / (phi * (p - 1))
  logsum <- numeric(length(y))
  nbar   <- numeric(length(y))
  for (i in which(y > 0)) {
    s <- .tweedie_series_at(y[i], lam[i], a, b[i],
                            y[i]^(2 - p) / (phi * (2 - p)))
    logsum[i] <- s$logsum
    nbar[i]   <- s$nbar
  }
  list(lam = lam, a = a, b = b, logsum = logsum, nbar = nbar)
}

.tweedie_series_at <- function(y, lam, a, b, jmax) {
  logterm <- function(n) {
    n * log(lam) - lgamma(n + 1) + n * a * log(b) +
      (n * a - 1) * log(y) - lgamma(n * a)
  }
  n0 <- max(1L, as.integer(round(jmax)))
  lt <- logterm(n0)
  lmax <- lt; terms <- lt; ns <- n0
  n <- n0
  repeat {
    n <- n + 1L
    lt <- logterm(n)
    terms <- c(terms, lt); ns <- c(ns, n)
    lmax <- max(lmax, lt)
    if (lt < lmax - 37) break
  }
  n <- n0
  while (n > 1L) {
    n <- n - 1L
    lt <- logterm(n)
    terms <- c(terms, lt); ns <- c(ns, n)
    lmax <- max(lmax, lt)
    if (lt < lmax - 37) break
  }
  w <- exp(terms - lmax)
  list(logsum = lmax + log(sum(w)), nbar = sum(ns * w) / sum(w))
}

# Tweedie log-density at mean mu (vectorized over observations). y = 0 has the
# exact Poisson-zero mass; y > 0 is the series above minus lambda and b y.
.tweedie_loglik <- function(mu, y, phi, p) {
  s   <- .tweedie_series(mu, y, phi, p)
  out <- s$logsum - s$lam - s$b * y
  out[y == 0] <- -s$lam[y == 0]
  out[y < 0]  <- -Inf                      # outside the support
  out
}

# d log f / d phi, elementwise. With 1 + a = 1 / (p - 1),
#
#   dl/dphi = [lambda + b y - E_W[n] / (p - 1)] / phi
#
# which at y = 0 reduces to lambda / phi: no series, no event, and the exact
# derivative of the Poisson-zero mass -lambda.
.tweedie_dloglik_dphi <- function(mu, y, phi, p) {
  s <- .tweedie_series(mu, y, phi, p)
  (s$lam + s$b * y - s$nbar / (p - 1)) / phi
}

# Default degrees of freedom for the Student-t family (matches the C++
# kStudentTDf), used when no `phi2` is supplied.
.STUDENT_T_DF <- 4


#' Supported R-level family names
#'
#' The response families `family = ` accepts on the engine's front doors, in
#' their canonical (default-link) spelling. A family's `*_<link>` variants
#' (e.g. `"gamma_inverse"`) are accepted as well and not listed here.
#'
#' @return A character vector of family names.
#' @examples
#' family_names()
#' @export
family_names <- function() names(.FAMILY_OPS)


# Families carrying a dispersion / precision parameter `phi` that enters the
# likelihood, score, and working weight. A non-positive `phi` makes the variance
# / log-density singular (gaussian: r^2 / (2 phi^2); gamma: lgamma(phi)), so it
# must be validated at the front door rather than flowing into a NaN log-lik
#. The `*_<link>` suffix forms (e.g. "gamma_inverse") share
# the base family's dispersion, so membership is tested on the prefix.
.PHI_FAMILIES <- c("gaussian", "lognormal", "gamma", "neg_binomial_2",
                   "neg_binomial_1", "inverse_gaussian",
                   "beta", "beta_binomial", "t", "tweedie",
                   "truncated_neg_binomial_2",
                   "interval_gaussian", "truncated_gaussian")

# Families whose `phi` is the residual VARIANCE on the R side while the compiled
# kernels parameterize by the residual SD. Every R -> kernel boundary converts
# through the two helpers below rather than testing the family name in place, so
# a `<family>_<link>` spelling converts like its base family. Testing the name
# directly fitted `gaussian_log` at variance `phi^2` when the caller asked for
# `phi`, and the error was invisible under the canonical link, where the
# observed-minus-working curvature is identically zero and a wrong dispersion
# cancels out of the mode motion.
.VARIANCE_PHI_FAMILIES <- c("gaussian", "lognormal")

# Does this family's R-side `phi` denote a variance whose kernel takes the SD?
#' @keywords internal
.phi_is_variance <- function(family) {
  .family_base(family) %in% .VARIANCE_PHI_FAMILIES
}

# R registry convention (variance) -> compiled-kernel convention (SD).
#' @keywords internal
.phi_to_kernel <- function(family, phi) {
  if (.phi_is_variance(family)) sqrt(phi) else phi
}

# Compiled-kernel convention (SD) -> R registry convention (variance).
#' @keywords internal
.phi_to_registry <- function(family, phi) {
  if (.phi_is_variance(family)) phi^2 else phi
}

# Count families whose likelihood is defined only at non-negative integer `y`.
# The C++ kernels cast the response to `int` (laplace_family_link.h /
# laplace_likelihoods.h), silently flooring a continuous response into a biased
# log-likelihood, so a non-integer `y` is rejected at the front door.
.COUNT_FAMILIES <- c("poisson", "binomial", "neg_binomial_2", "neg_binomial_1",
                     "beta_binomial",
                     "truncated_poisson", "truncated_neg_binomial_2")

# Zero-truncated count families: the density conditions on y > 0, so a zero
# response is outside the support rather than merely unlikely.
.TRUNCATED_FAMILIES <- c("truncated_poisson", "truncated_neg_binomial_2")

# Families supported on the strictly positive reals. The log-density is
# undefined at zero -- log(y) for lognormal, the y^(shape - 1) and 1/y factors
# for gamma and inverse gaussian -- so a zero is outside the support rather
# than merely unlikely. tweedie is deliberately absent: for 1 < p < 2 it is
# compound Poisson-gamma and carries a point mass at zero.
.POSITIVE_FAMILIES <- c("gamma", "lognormal", "inverse_gaussian")

# Families carried by the R registry but not by the compiled kernels. Every
# shipped family is now wired into src/laplace_family_link.h, so this is empty;
# it stays as the gate a newly registered family passes through until its
# kernels exist. The compiled dispatch no longer falls through to Poisson --
# an unrecognized name raises there (unknown_family_stop) -- so a family added
# to the registry alone fails loudly rather than fitting the wrong likelihood.
.R_ONLY_FAMILIES <- character(0)

#' Reject families the compiled kernels do not implement
#'
#' @noRd
.validate_family_compiled <- function(family) {
  if (!(.family_base(family) %in% .R_ONLY_FAMILIES)) return(invisible(TRUE))
  stop(sprintf(paste0(
    "family = '%s' is defined in the R family registry but is not yet wired ",
    "into the compiled kernels, so tulpa() cannot fit it. The C++ family ",
    "dispatch would silently substitute a Poisson likelihood. Use the ",
    "family_*() R interface for density, score, weight, and variance ",
    "evaluation in the meantime."), family), call. = FALSE)
}

# Base family of a `family` / `family_<link>` code (the part before the link
# suffix), so a custom link form validates against the same rules.
#
# Candidates are matched longest-first because family names are prefixes of one
# another: "beta_binomial" starts with "beta_", so a registration-order scan
# returns "beta" and every `.family_base`-driven validator then applies beta's
# rules to a beta-binomial. Ordering by decreasing name length makes the
# longest -- i.e. the most specific -- family win, and an exact hit always beats
# a prefix hit.
.family_base <- function(family) {
  if (!is.character(family) || length(family) != 1L) return(NA_character_)
  # Resolve a short spelling first, so a caller that skipped canonicalization
  # still gets the family's own rules rather than silently falling through
  # every membership test below.
  family <- .canonical_family(family)
  cand <- unique(c(.PHI_FAMILIES, .COUNT_FAMILIES, family_names()))
  if (family %in% cand) return(family)
  cand <- cand[order(nchar(cand), decreasing = TRUE)]
  for (fam in cand) {
    if (startsWith(family, paste0(fam, "_"))) return(fam)
  }
  family
}

# Does `family` read a per-observation trial count? The binomial families do;
# every other family ignores `n_trials` entirely. The one predicate for that
# question: a fit's `$n_trials` field is not the answer, because the design
# bundle carries a `rep(1, n)` default for every family and `.finalize_fit()`
# stamps it on every fit (gcol33/tulpa#781), so a non-NULL `$n_trials` says
# nothing about whether the model uses one (gcol33/tulpa#837).
.family_reads_trials <- function(family) {
  .family_base(family) %in% c("binomial", "beta_binomial")
}

# One R-side reading of `n_trials`, so every door answers the same
# (gcol33/tulpa#677): NULL stays NULL; otherwise the family must read a
# denominator, a scalar is recycled to `n_obs`, and the result is a length-
# `n_obs` integer vector of positive counts (NA passes to the finite guard).
#' @keywords internal
.normalize_n_trials <- function(family, n_trials, n_obs) {
  if (is.null(n_trials)) return(NULL)
  if (!.family_reads_trials(family)) {
    stop(sprintf(paste0(
      "`n_trials` is the binomial denominator and is not read by ",
      "family = '%s'. Drop it, or use family = 'binomial'."), family),
      call. = FALSE)
  }
  n_trials <- as.integer(n_trials)
  if (length(n_trials) == 1L) n_trials <- rep(n_trials, n_obs)
  if (length(n_trials) != n_obs) {
    stop(sprintf(paste0(
      "`n_trials` must have length 1 (recycled) or nrow(data) (%d); got %d."),
      n_obs, length(n_trials)), call. = FALSE)
  }
  if (any(!is.na(n_trials) & n_trials < 1L)) {
    stop("`n_trials` must be a positive integer count.", call. = FALSE)
  }
  n_trials
}

#' Validate the dispersion parameter `phi` for a family
#'
#' Errors when `family` carries a dispersion / precision parameter and `phi` is
#' not a positive finite scalar. A no-op for families without dispersion
#' (binomial, poisson) and for unknown / model-package families. Shared by the
#' [tulpa()] front door and [tulpa_laplace()] so the rule lives in one place
#'.
#' @noRd
.validate_family_phi <- function(family, phi) {
  if (!(.family_base(family) %in% .PHI_FAMILIES)) return(invisible(TRUE))
  if (!is.numeric(phi) || length(phi) != 1L || !is.finite(phi) || phi <= 0) {
    stop(sprintf("`phi` must be a positive finite scalar for family = '%s'.",
                 family), call. = FALSE)
  }
  invisible(TRUE)
}

#' Validate the response `y` for a count family
#'
#' Errors when `family` is a count family and `y` carries negative or
#' non-integer values, which the integer-casting kernels would silently floor
#' into a biased likelihood. A no-op for continuous families.
#'
#' @param zi Whether a zero-inflation component is being fitted alongside the
#'   family. With a zero-truncated family this is the hurdle model, where the
#'   zeros belong to the zero component rather than to the truncated count
#'   density, so the `y >= 1` requirement does not apply to them.
#' @noRd
.validate_family_counts <- function(family, y, zi = FALSE) {
  base <- .family_base(family)
  if (!(base %in% .COUNT_FAMILIES)) return(invisible(TRUE))
  yf <- y[is.finite(y)]
  if (length(yf) && (any(yf < 0) || any(yf != round(yf)))) {
    stop(sprintf(paste0(
      "family = '%s' requires non-negative integer counts in `y`; got negative ",
      "or non-integer values. The likelihood casts `y` to an integer, which ",
      "would silently floor a continuous response into a biased log-likelihood."),
      family), call. = FALSE)
  }
  if (!zi && base %in% .TRUNCATED_FAMILIES && length(yf) && any(yf < 1)) {
    stop(sprintf(paste0(
      "family = '%s' is zero-truncated and requires `y >= 1`; got %d zero(s). ",
      "The density conditions on y > 0, so a zero carries no likelihood. Use ",
      "the untruncated family, or a hurdle model if the zeros are structural."),
      family, sum(yf < 1)), call. = FALSE)
  }
  invisible(TRUE)
}


#' Validate the response against a family's support
#'
#' The one support rule set, shared by the [tulpa()] front door and by the
#' engine fitters ([tulpa_laplace()], [tulpa_laplace_beta()],
#' [tulpa_nuts_beta()]). Holding it in a single function is what keeps a
#' boundary response rejected with the same message wherever it enters: a
#' per-fitter copy is only as good as the set of fitters that carry it, and
#' `tulpa()` reaches beta-with-RE through a sampler that carried none.
#'
#' Counts delegate to [.validate_family_counts()]; the continuous families and
#' the binomial denominator add their rules here.
#'
#' @param family Family identifier, canonical or aliased.
#' @param y Response vector, or `NULL` (a model with no response passes).
#' @param n_trials Binomial denominators, when known. Supplied, the binomial
#'   rule is `0 <= y <= n_trials`; absent, the response is the 0/1 form.
#' @param zi Whether a zero-inflation component is fitted alongside the family,
#'   passed through to the count rule for the hurdle case.
#' @noRd
.validate_family_support <- function(family, y, n_trials = NULL, zi = FALSE) {
  if (is.null(y)) return(invisible(TRUE))
  base <- .family_base(family)
  yv <- suppressWarnings(as.numeric(y))
  yf <- yv[is.finite(yv)]

  .validate_family_counts(family, yv, zi = zi)

  if (identical(base, "beta") && length(yf) &&
      (min(yf) <= 0 || max(yf) >= 1)) {
    stop(sprintf(paste0(
      "family = 'beta' requires `y` strictly in (0, 1); got range [%s, %s]. ",
      "Use cover(positive = 'beta') for hurdle handling of 0/1, or transform ",
      "the boundary observations inward, e.g. (y * (n - 1) + 0.5) / n ",
      "(Smithson & Verkuilen 2006)."),
      format(min(yf)), format(max(yf))), call. = FALSE)
  }

  if (base %in% .POSITIVE_FAMILIES && length(yf) && min(yf) <= 0) {
    stop(sprintf(paste0(
      "family = '%s' requires `y` strictly positive; got %d value(s) <= 0 ",
      "(minimum %s). The log-density is undefined at zero, so these ",
      "observations are outside the support."),
      family, sum(yf <= 0), format(min(yf))), call. = FALSE)
  }

  if (identical(base, "binomial") && length(yf)) {
    if (is.null(n_trials)) {
      if (any(yf != 0 & yf != 1)) {
        stop(paste0(
          "family = 'binomial' with no `n_trials` requires a 0/1 response. ",
          "Supply `n_trials` for aggregated counts, or write the response as ",
          "cbind(successes, failures)."), call. = FALSE)
      }
    } else {
      nt <- suppressWarnings(as.numeric(n_trials))
      if (length(nt) == 1L) nt <- rep(nt, length(yv))
      if (length(nt) != length(yv)) {
        stop(sprintf(paste0(
          "`n_trials` has length %d but the response has %d observation(s)."),
          length(nt), length(yv)), call. = FALSE)
      }
      ok <- is.finite(yv) & is.finite(nt)
      if (any(nt[ok] < 0) || any(nt[ok] != round(nt[ok]))) {
        stop("`n_trials` must be non-negative integers.", call. = FALSE)
      }
      if (any(yv[ok] > nt[ok])) {
        bad <- which(ok & yv > nt)
        stop(sprintf(paste0(
          "family = 'binomial' requires `y <= n_trials`; %d observation(s) ",
          "exceed their denominator (first at row %d: y = %s, n_trials = %s)."),
          length(bad), bad[1L], format(yv[bad[1L]]), format(nt[bad[1L]])),
          call. = FALSE)
      }
    }
  }

  invisible(TRUE)
}


#' Reject non-finite fitting inputs
#'
#' The design is built with `na.action = na.pass`, so a missing predictor or
#' response survives into `X` / `y`. No fitter drops incomplete cases, so an
#' NA/NaN/Inf would propagate into the C++ kernels as a NaN estimate. Fail
#' loudly with the offending row instead.
#'
#' Reached from [tulpa()], from the direct beta-only doors, and from
#' [.validate_glm_design()], which is what carries it to every fitter taking a
#' `(y, X, n_trials)` bundle.
#'
#' Every per-observation numeric the kernels index is checked here, not only
#' the design and the response: `n_trials` and the offset enter eta and the
#' likelihood the same way, and an NA in either came back as a silent all-zero
#' coefficient vector on the Laplace door (gcol33/tulpa#665).
#'
#' @param X,y Design matrix and response; either may be `NULL` to skip that arm.
#' @param n_trials,offset Per-observation binomial denominators and offset, or
#'   `NULL` to skip that arm.
#' @param where Caller name for the message prefix, or `NULL` for none.
#' @noRd
.assert_finite_model_inputs <- function(X, y, n_trials = NULL, offset = NULL,
                                        where = NULL) {
  pre <- if (is.null(where)) "" else paste0(where, ": ")
  if (!is.null(X)) {
    ok_row <- if (is.matrix(X)) .all_finite_rows(X) else is.finite(X)
    if (!all(ok_row)) {
      bad <- which(!ok_row)
      stop(sprintf(paste0(
        "%sNon-finite value(s) in the model matrix (%d row(s), first at row ",
        "%d). tulpa does not drop incomplete cases; remove or impute ",
        "NA/NaN/Inf in the predictors before fitting."),
        pre, length(bad), bad[1L]),
        call. = FALSE)
    }
  }
  if (!is.null(y)) {
    bad <- which(!is.finite(as.numeric(y)))
    if (length(bad)) {
      stop(sprintf(paste0(
        "%sNon-finite value(s) in the response (%d, first at row %d). Remove ",
        "or impute NA/NaN/Inf in the response before fitting."),
        pre, length(bad), bad[1L]), call. = FALSE)
    }
  }
  for (arm in list(list(v = n_trials, nm = "n_trials"),
                   list(v = offset, nm = "offset"))) {
    if (is.null(arm$v)) next
    bad <- which(!is.finite(suppressWarnings(as.numeric(arm$v))))
    if (length(bad)) {
      stop(sprintf(paste0(
        "%sNon-finite value(s) in `%s` (%d, first at row %d). Remove or ",
        "impute NA/NaN/Inf before fitting."),
        pre, arm$nm, length(bad), bad[1L]), call. = FALSE)
    }
  }
  invisible(TRUE)
}

#' The response as the numeric vector the kernels read
#'
#' A logical response is the binary idiom `glm()` accepts for a binomial
#' family, `TRUE` a success; it is read as 0/1 there. Anywhere else, and for
#' any other non-numeric response (a factor, a character vector), the family
#' has no reading of it, so it is refused by name rather than left to fail
#' inside a backend in that backend's words -- `is.numeric(y) is not TRUE` on
#' the Laplace door, a failed inner solve on EB, non-finite integration nodes
#' on the structured tier (gcol33/tulpa#880).
#'
#' @param family Family identifier, canonical or aliased.
#' @param y The response, or `NULL`.
#' @noRd
.numeric_response <- function(family, y) {
  if (is.null(y) || is.numeric(y)) return(y)
  if (is.logical(y) && .family_reads_trials(family)) {
    return(as.integer(y))
  }
  hint <- if (is.logical(y)) {
    "Use family = 'binomial' for a binary response, or convert it with as.numeric()."
  } else if (is.factor(y) || is.character(y)) {
    paste0("Code a binary response as 0/1 (e.g. as.integer(y == \"<level>\")), ",
           "or use family = 'multinomial' / 'ordinal' for a categorical one.")
  } else {
    "Convert it to a numeric vector."
  }
  stop(sprintf("The response is %s, which family = '%s' does not read. %s",
               if (is.logical(y)) "logical" else paste0("a ", class(y)[1L]),
               family, hint), call. = FALSE)
}

#' Reject an incomplete random-effect design
#'
#' The formula's grouping factors and random-slope covariates are built with
#' `na.action = na.pass` as the fixed design is, so an NA group or slope value
#' survives into `re_terms`. There it reached each backend in its own words --
#' an `re_idx_list` index message on the Laplace door, "NAs are not allowed in
#' subscripted assignments" on the samplers (gcol33/tulpa#880). Refused here
#' with the message `.assert_finite_model_inputs()` gives the predictors.
#'
#' @param re_terms The `re_terms` of [tulpa_build_model_data()].
#' @param where Caller name for the message prefix, or `NULL` for none.
#' @noRd
.assert_complete_groups <- function(re_terms, where = NULL) {
  pre <- if (is.null(where)) "" else paste0(where, ": ")
  for (rt in re_terms) {
    bad <- which(is.na(rt$group_idx))
    if (length(bad)) {
      stop(sprintf(paste0(
        "%sMissing value(s) in the grouping variable `%s` (%d row(s), first ",
        "at row %d). tulpa does not drop incomplete cases; remove those rows ",
        "or give them a group before fitting."),
        pre, rt$group_var %||% "?", length(bad), bad[1L]), call. = FALSE)
    }
    if (!is.null(rt$slope_matrix)) {
      ok <- .all_finite_rows(as.matrix(rt$slope_matrix))
      if (!all(ok)) {
        bad <- which(!ok)
        stop(sprintf(paste0(
          "%sNon-finite value(s) in the random-slope covariate(s) of `%s` ",
          "(%d row(s), first at row %d). tulpa does not drop incomplete ",
          "cases; remove or impute NA/NaN/Inf before fitting."),
          pre, rt$group_var %||% "?", length(bad), bad[1L]), call. = FALSE)
      }
    }
  }
  invisible(TRUE)
}

# Row-wise all-finite test for a numeric matrix.
.all_finite_rows <- function(X) {
  fin <- is.finite(X)
  dim(fin) <- dim(X)
  .rowSums(fin, nrow(X), ncol(X)) == ncol(X)
}


#' Look up a family's operation set, with a clear error for unknown families
#' @noRd
.family_ops <- function(family) {
  ops <- .FAMILY_OPS[[family]]
  # A `<base>_<link>` code resolves through the link layer (family_link.R),
  # which composes the base family's mu-space form with the link. A code naming
  # a family's own canonical link comes back as the hand-written entry above, so
  # the canonical path is untouched.
  if (is.null(ops)) ops <- .linked_family_ops(family)
  if (is.null(ops)) {
    .family_or_stop(family)
  }
  ops
}

# Short and historical family spellings, mapped onto the registry's canonical
# names. These entered from the lme4/glmmTMB idiom and from earlier front
# doors, and several inner fitters grew their own `c("neg_binomial_2",
# "negbin")` branch to accept one of them, which left the same string accepted
# by a fitter and rejected by tulpa(). Resolution happens once, where a family
# string enters, so every branch downstream tests one spelling.
.FAMILY_ALIASES <- c(
  negbin            = "neg_binomial_2",
  negative_binomial = "neg_binomial_2",
  nbinom2           = "neg_binomial_2",
  nb2               = "neg_binomial_2",
  nb1               = "neg_binomial_1",
  nbinom1           = "neg_binomial_1",
  student_t         = "t",
  tpoisson          = "truncated_poisson",
  tnbinom2          = "truncated_neg_binomial_2"
)

# A stats::family() object as the engine's family name.
#
# `family = binomial()` is the natural thing to reach for and used to reach an
# `if (family == ...)` on a five-element list, erroring "the condition has
# length > 1" (gcol33/tulpa#679). The link rides the name in this engine's own
# `<base>_<link>` convention, so a non-canonical link is carried rather than
# silently fitted at the canonical one.
#
# The suffix is decided against THIS engine's default link (`.LINK_DEFAULTS`),
# not against stats' canonical one: the two differ for gamma and inverse
# gaussian, where stats defaults to the inverse / 1/mu^2 link and tulpa's bare
# name means log. Comparing against stats' table mapped `Gamma()` to the bare
# `"gamma"` and fitted a log link with no warning (gcol33/tulpa#873). stats'
# spellings of the family (`inverse.gaussian`) and of a link (`1/mu^2`) are
# translated to the registry's before either is compared, and a link the base
# does not admit is refused here, naming the object the user wrote.
.STATS_FAMILY_NAMES <- c(inverse.gaussian = "inverse_gaussian")
.STATS_LINK_NAMES   <- c(`1/mu^2` = "1mu2")

#' @keywords internal
.family_object_to_name <- function(fam) {
  nm <- tolower(fam$family %||% "")
  if (!nzchar(nm)) {
    stop("`family` object carries no $family name.", call. = FALSE)
  }
  base <- unname(.STATS_FAMILY_NAMES[nm])
  if (is.na(base)) base <- nm
  lk <- fam$link
  if (is.null(lk) || !(base %in% names(.LINK_DEFAULTS))) return(base)
  link <- unname(.STATS_LINK_NAMES[lk])
  if (is.na(link)) link <- lk
  if (!(link %in% .OK_LINKS[[base]])) {
    stop(sprintf(
      "`family = %s(link = \"%s\")`: the %s link is not available for %s. ",
      fam$family, lk, lk, base),
      "Supported links: ", paste(.OK_LINKS[[base]], collapse = ", "), ".",
      call. = FALSE)
  }
  if (identical(link, unname(.LINK_DEFAULTS[[base]]))) base
  else paste0(base, "_", link)
}

#' Resolve a family spelling to its canonical registry name
#'
#' A no-op for a canonical name, for a `<family>_<link>` code, and for anything
#' unrecognized, which reaches `.family_or_stop()` and errors there against the
#' canonical list.
#'
#' Case is normalized as a last resort, after the exact spelling and the alias
#' table both miss: `stats::Gamma()$family` is `"Gamma"`, capitalized, and
#' `.family_object_to_name()` lowercases it for the `family = Gamma(...)`
#' object form, so the equivalent string form (`family = "Gamma"`) has to reach
#' the same registry entry rather than being refused as unknown
#' (gcol33/tulpa#806). Applied to a `<family>_<link>` code too
#' (`"Gamma_log"` -> `"gamma_log"`), never to a name that is already
#' recognized as it stands -- so a registry name that happens to collide
#' case-insensitively with another spelling is never silently rewritten.
#' @param family Family identifier as supplied by the caller.
#' @noRd
.canonical_family <- function(family) {
  if (!is.character(family) || length(family) != 1L || is.na(family)) {
    return(family)
  }
  hit <- unname(.FAMILY_ALIASES[family])
  if (!is.na(hit)) return(hit)
  if (!is.null(.FAMILY_OPS[[family]]) || !is.null(.linked_family_ops(family))) {
    return(family)
  }
  lower <- tolower(family)
  if (!identical(lower, family) &&
      (!is.null(.FAMILY_OPS[[lower]]) || !is.null(.linked_family_ops(lower)))) {
    return(lower)
  }
  family
}

# A hurdle is a composition here, not a family: the zero-inflation mixture over
# a zero-truncated base degenerates exactly to the two-part likelihood
# (family_zi.R:13-16), so there is no `hurdle_*` entry to alias to. Naming one is
# still the natural thing to reach for, so it gets the composition rather than
# the generic list of everything the registry does have.
.HURDLE_SPELLINGS <- c(
  hurdle_poisson  = "truncated_poisson",
  hurdle_nbinom2  = "truncated_neg_binomial_2",
  hurdle_negbin   = "truncated_neg_binomial_2",
  hurdle_nb2      = "truncated_neg_binomial_2"
)

# Validate a family identifier against the registry (the one unknown-family
# error, shared by every front door); returns the identifier invisibly.
#' @keywords internal
.family_or_stop <- function(family) {
  ok <- is.character(family) && length(family) == 1L &&
    (!is.null(.FAMILY_OPS[[family]]) || !is.null(.linked_family_ops(family)))
  if (!ok) {
    nm <- as.character(family)[1L]
    base <- unname(.HURDLE_SPELLINGS[nm])
    if (length(base) == 1L && !is.na(base)) {
      stop(sprintf(paste0(
        "Family '%s' is written as a composition here: family = '%s' with a ",
        "`ziformula`. The zero-inflation mixture over a zero-truncated base ",
        "degenerates exactly to the hurdle likelihood, so there is no separate ",
        "hurdle family to select."), nm, base), call. = FALSE)
    }
    stop(sprintf(paste0(
      "Unknown family '%s'. Supported: %s. A non-canonical link is written as ",
      "'<family>_<link>' (links: %s); %s take one."),
      nm,
      paste(family_names(), collapse = ", "),
      paste(link_names(), collapse = ", "),
      paste(names(.LINK_DEFAULTS), collapse = ", ")
    ), call. = FALSE)
  }
  invisible(family)
}


# Resolve a supplied second dispersion against the family: NULL passes, a
# phi2-aware family gets it, anything else errors rather than silently
# ignoring a parameter the user thought they set.
.phi2_or_stop <- function(family, phi2) {
  if (is.null(phi2)) return(NULL)
  if (!.family_base(family) %in% .PHI2_FAMILIES) {
    stop(sprintf(
      "Family '%s' takes no second dispersion `phi2` (supported: %s).",
      family, paste(.PHI2_FAMILIES, collapse = ", ")), call. = FALSE)
  }
  if (!is.numeric(phi2) || length(phi2) != 1L || !is.finite(phi2) || phi2 <= 0) {
    stop("`phi2` must be a positive finite scalar.", call. = FALSE)
  }
  phi2
}


# Backends that thread the second dispersion into their likelihood. The
# log-posterior builder and the ModelData samplers carry it by construction;
# "laplace" reaches it through the non-spatial compiled kernel, and "eb" /
# "re_cov_nested" through that same kernel as their inner solve. Derived from the
# registry rather than restated, so a new sampler backend inherits phi2 with its
# input contract. A function, not a constant: this file sources before
# inference_modes.R.
#' @keywords internal
.phi2_backends <- function() {
  c("laplace", "eb", "re_cov_nested",
    names(BACKEND_REGISTRY)[vapply(
      BACKEND_REGISTRY,
      function(b) isTRUE(b$input %in% c("logpost", "modeldata")) || isTRUE(b$phi2),
      logical(1))])
}


#' Inverse-link mean for a family
#'
#' `phi` reaches entries whose response mean depends on the dispersion
#' (lognormal); the others ignore it via `...`.
#' @noRd
family_mean <- function(eta, family, phi = 1.0) {
  .family_ops(family)$mean(eta, phi)
}


#' Elementwise log-likelihood for a family
#' @noRd
family_loglik <- function(eta, y, family, n_trials = NULL, phi = 1.0,
                          phi2 = NULL) {
  ops <- .family_ops(family)
  if (is.null(.phi2_or_stop(family, phi2))) {
    return(ops$loglik(eta, y, n_trials, phi))
  }
  ops$loglik(eta, y, n_trials, phi, phi2)
}


#' Score (d log-likelihood / d eta), elementwise
#' @noRd
family_score_eta <- function(eta, y, family, n_trials = NULL, phi = 1.0,
                             phi2 = NULL) {
  ops <- .family_ops(family)
  if (is.null(.phi2_or_stop(family, phi2))) {
    return(ops$score(eta, y, n_trials, phi))
  }
  ops$score(eta, y, n_trials, phi, phi2)
}


#' Laplace/IRLS working weight (-d^2 log-lik / d eta^2), elementwise
#' @noRd
family_weight <- function(eta, family, n_trials = NULL, phi = 1.0,
                          phi2 = NULL) {
  ops <- .family_ops(family)
  if (is.null(.phi2_or_stop(family, phi2))) {
    return(ops$weight(eta, n_trials, phi))
  }
  ops$weight(eta, n_trials, phi, phi2)
}


#' One response draw per element of `eta` (posterior predictive), elementwise
#' @noRd
family_sample <- function(eta, family, n_trials = NULL, phi = 1.0,
                          phi2 = NULL) {
  ops <- .family_ops(family)
  if (is.null(ops$sample)) {
    stop(sprintf("Family '%s' has no sampling function registered.", family),
         call. = FALSE)
  }
  if (is.null(.phi2_or_stop(family, phi2))) {
    return(ops$sample(eta, n_trials, phi))
  }
  ops$sample(eta, n_trials, phi, phi2)
}


#' Response variance Var(y | eta) for a family, elementwise
#' @noRd
family_variance <- function(eta, family, n_trials = NULL, phi = 1.0,
                            phi2 = NULL) {
  ops <- .family_ops(family)
  if (is.null(ops$variance)) {
    stop(sprintf("Family '%s' has no variance function registered.", family),
         call. = FALSE)
  }
  if (is.null(.phi2_or_stop(family, phi2))) {
    return(ops$variance(eta, n_trials, phi))
  }
  ops$variance(eta, n_trials, phi, phi2)
}


#' Observed curvature (-d^2 log-lik / d eta^2 at the realized `y`), elementwise
#'
#' One path, the compiled dispatch (`obs_grad_hess_for_family`), for every
#' family. It used to answer from the registry's `obs_weight` closure where a
#' family registered one and substitute the y-free EXPECTED weight where none
#' was registered. That substitution is exact only where the response enters
#' the log-likelihood linearly in `eta` (Poisson, binomial, the zero-truncated
#' Poisson); for beta, gamma, inverse_gaussian, beta_binomial, tweedie and t it
#' is a different function, and can carry the opposite sign -- at `y = 12`,
#' `eta = -0.9` the beta_binomial observed curvature is `-0.125` against an
#' expected weight of `+0.658`. beta_binomial is in `.ZI_FAMILIES`, so the
#' zero-inflation mixture in `R/family_zi.R` differentiated through the wrong
#' one (gcol33/tulpa#824).
#'
#' The registry's `obs_weight` closures (neg_binomial_2, neg_binomial_1,
#' truncated_neg_binomial_2, beta) stay as the R oracle the compiled dispatch
#' is pinned against in `test-family-registry-compiled.R`.
#' @noRd
.family_obs_weight <- function(eta, y, family, n_trials = NULL, phi = 1.0,
                               phi2 = NULL) {
  .family_or_stop(family)
  .phi2_or_stop(family, phi2)
  n   <- max(length(eta), length(y), length(n_trials))
  eta <- rep_len(as.numeric(eta), n)
  y   <- rep_len(as.numeric(y), n)
  nt  <- rep_len(as.integer(if (is.null(n_trials)) 1L else n_trials), n)
  cpp_family_obs_weight(y, nt, eta, family,
                        .phi_to_kernel(family, phi),
                        if (is.null(phi2)) NA_real_ else as.numeric(phi2))
}


#' Response-scale mean of y given eta for a family (trial-scaled where relevant)
#' @noRd
family_response_mean <- function(eta, family, n_trials = NULL, phi = 1.0) {
  ops <- .family_ops(family)
  if (!is.null(ops$response_mean)) return(ops$response_mean(eta, n_trials, phi))
  ops$mean(eta, phi)
}