shrinkage

Empirical-Bayes shrinkage strength estimation (spec section 4).

The support weight w = tau^2 / (tau^2 + v) is the posterior weight an empirical-Bayes observer puts on a fitted bin value versus the population prior 0 (terms are centered). tau^2 is the between-bin signal variance, estimated per term; v is the bin's sampling variance.

Shrinkage dataclass

Shrinkage(tau2: float, sigma2: float, k: float)

Per-term shrinkage summary (audit payload, spec section 4.4).

estimate_tau2

estimate_tau2(
    f: ndarray,
    v: ndarray,
    n: ndarray,
    *,
    method: str = "mom"
) -> float

Estimate the between-bin signal variance tau^2 for one term.

Parameters:

Name	Type	Description	Default
f	ndarray	Bin values, bin sampling variances, bin training masses (any shape, flattened internally).	required
v	ndarray	Bin values, bin sampling variances, bin training masses (any shape, flattened internally).	required
n	ndarray	Bin values, bin sampling variances, bin training masses (any shape, flattened internally).	required
method	str	"mom" (default): mass-weighted method of moments, max(0, sum(pi f^2) - sum(pi v)). "mle": 1-D marginal maximum likelihood, more stable for high-cardinality terms.	'mom'

Source code in src/triadxai/shrinkage.py

def estimate_tau2(f: np.ndarray, v: np.ndarray, n: np.ndarray, *, method: str = "mom") -> float:
    """Estimate the between-bin signal variance tau^2 for one term.

    Parameters
    ----------
    f, v, n
        Bin values, bin sampling variances, bin training masses (any shape,
        flattened internally).
    method
        ``"mom"`` (default): mass-weighted method of moments,
        ``max(0, sum(pi f^2) - sum(pi v))``. ``"mle"``: 1-D marginal
        maximum likelihood, more stable for high-cardinality terms.
    """
    pi = _mass_weights(n)
    f_ = np.ravel(f).astype(float)
    v_ = np.ravel(v).astype(float)
    if method == "mom":
        return max(0.0, float((pi * f_**2).sum() - (pi * v_).sum()))
    if method == "mle":
        m = pi > 0.0

        def nll(log_tau2: float) -> float:
            s = math.exp(log_tau2) + v_[m]
            return float((pi[m] * (np.log(s) + f_[m] ** 2 / s)).sum())

        res = optimize.minimize_scalar(
            nll, bounds=(math.log(1e-12), math.log(1e6)), method="bounded"
        )
        tau2 = float(math.exp(res.x))
        return 0.0 if tau2 <= _TAU2_FLOOR else tau2
    raise ValueError(f"unknown tau^2 method {method!r}; expected 'mom' or 'mle'")

estimate_tau2_shap

estimate_tau2_shap(
    phi_ref: ndarray, v_ref: ndarray
) -> np.ndarray

Per-feature tau^2 for approximate mode, from a reference population.

Method of moments about the zero prior: tau2_j = max(0, mean(phi_j^2) - mean(v_j)). SHAP attributions are centered on the training population, so a strongly nonzero mean signals a reference-population mismatch and is logged as a warning.

Source code in src/triadxai/shrinkage.py

def estimate_tau2_shap(phi_ref: np.ndarray, v_ref: np.ndarray) -> np.ndarray:
    """Per-feature tau^2 for approximate mode, from a reference population.

    Method of moments about the zero prior:
    ``tau2_j = max(0, mean(phi_j^2) - mean(v_j))``. SHAP attributions are
    centered on the training population, so a strongly nonzero mean signals
    a reference-population mismatch and is logged as a warning.
    """
    mean = phi_ref.mean(axis=0)
    std = phi_ref.std(axis=0)
    off = np.abs(mean) > 0.05 * np.where(std > 0, std, 1.0)
    if off.any():
        logger.warning(
            "SHAP attributions off-center for feature indices %s; "
            "reference population may not match training data",
            np.flatnonzero(off).tolist(),
        )
    return np.maximum(0.0, (phi_ref**2).mean(axis=0) - v_ref.mean(axis=0))

fit_shrinkage

fit_shrinkage(
    term: TermData, *, method: str = "mom"
) -> Shrinkage

Estimate tau^2 and the count-form shrinkage strength k for one term.

Source code in src/triadxai/shrinkage.py

def fit_shrinkage(term: TermData, *, method: str = "mom") -> Shrinkage:
    """Estimate tau^2 and the count-form shrinkage strength k for one term."""
    tau2 = estimate_tau2(term.f, term.v, term.n, method=method)
    pi = _mass_weights(term.n)
    sigma2 = float((pi * np.ravel(term.n) * np.ravel(term.v)).sum())
    if tau2 == 0.0:
        logger.warning(
            "term %d: tau^2 = 0 (pure-noise feature); entire contribution -> D",
            term.term_idx,
        )
        return Shrinkage(tau2=0.0, sigma2=sigma2, k=math.inf)
    return Shrinkage(tau2=tau2, sigma2=sigma2, k=sigma2 / tau2)