Fit an adaptive-shrinkage hidden Markov model

Fits a finite-state hidden Markov model to heteroskedastic normal effect estimates. Each hidden state has an adaptive-shrinkage prior formed from a point mass at its state center and a grid of normal components centered at the same value. Transition probabilities, state-specific mixture weights, and (optionally) supported state centers are learned by generalized EM/Baum–Welch updates.

Usage

fit_ash_hmm(
  y,
  se,
  mu = NULL,
  prior_sd = NULL,
  half_grid = 50L,
  grid_shape = 3,
  grid_expansion = 3,
  grid_max_abs = NULL,
  nonnegative_state_means = FALSE,
  positive_state_means = NULL,
  effect_support = c("auto", "nonnegative", "real"),
  positive_mean_floor = NULL,
  prefilter = NULL,
  min_state_count = 0.5,
  min_state_fraction = 1e-05,
  screening_prior_sd = 0,
  screening_block_size = 5000L,
  sequence_id = rep(1L, length(y)),
  forward_backward_engine = c("auto", "scaled", "log"),
  inference = c("exact", "variational"),
  variational_maxiter = 1L,
  variational_tolerance = 1e-05,
  variational_probability_floor = 1e-12,
  variational_final_exact = TRUE,
  cache_component_emissions = TRUE,
  topology = c("full", "hub"),
  transition_mask = NULL,
  init_transition = NULL,
  init_prob = NULL,
  init_rho = NULL,
  stay_probability = 0.95,
  null_state = c("pointmass", "adaptive"),
  fixed_pointmass_states = NULL,
  shared_mixture = FALSE,
  estimate_init = FALSE,
  transition_prior = 1,
  mixture_prior = 1,
  init_prior = 1,
  learn_state_means = TRUE,
  fixed_mean_states = 1L,
  mean_update_start = 3L,
  mean_min_effective_count = 2,
  mean_min_pointmass_weight = 0,
  mean_min_self_transition = 0.93,
  mean_damping = 1,
  mean_bounds = c("voronoi", "none"),
  prune_states = TRUE,
  prune_start = 5L,
  prune_every = 5L,
  prune_min_state_count = 1,
  prune_min_state_fraction = 1e-04,
  prune_max_fraction = 0.25,
  prune_max_loglik_loss = 0.05,
  merge_distance = NULL,
  null_model_selection = c("none", "bic"),
  null_selection_gamma = 1,
  null_bic_margin = 0,
  parameter_tolerance = 1e-06,
  step_penalty = NULL,
  step_penalty_scale = 1.5,
  maxiter = 20L,
  tolerance = 1e-05,
  verbose = FALSE
)

Arguments

y: Numeric vector of noisy effect estimates.
se: Numeric vector of known standard errors. It must have the same length as `y`, contain only finite values, and be strictly positive.
mu: Optional numeric vector of initial state centers. State 1 is the zero/hub state. When `NULL`, centers are constructed and screened automatically. Supply `mu` when using state-sized custom initial values.
prior_sd: Optional nondecreasing numeric vector of ash prior standard deviations. Its first value must be zero, representing a point mass. When `NULL`, a Stephens-style geometric scale grid is constructed automatically.
half_grid: Number of nonnegative candidates in the automatic mean grid, including zero. Signed mode appends their nonzero negatives.
grid_shape: Positive power-grid shape parameter. Values above one put relatively more candidates near the outer part of the range.
grid_expansion: Positive multiplier applied to the empirical maximum absolute observation when constructing the automatic mean grid.
grid_max_abs: Optional finite positive endpoint used instead of `max(abs(y))` when constructing the automatic grid.
nonnegative_state_means: Logical; if `TRUE` (default), use one state centered at zero, constrain every non-null state center to be strictly positive, and truncate all continuous ash components below at zero. Thus both prior and posterior effects have support on `[0, Inf)`. If `FALSE`, allow signed state centers and effects on the real line.
positive_state_means: Deprecated compatibility alias for `nonnegative_state_means`. Leave as `NULL` in new code. If both arguments are supplied, they must agree.
effect_support: Either `"auto"` (default), `"nonnegative"`, or `"real"`. `"auto"` follows `nonnegative_state_means`. An explicit value must agree with that argument. The nonnegative model uses exact truncated-normal marginal emissions and posterior moments; it does not clip an unconstrained posterior.
positive_mean_floor: Strict lower bound for learned non-null centers in nonnegative mode. The default is a scale-aware numerical value.
prefilter: Logical controlling independent state-grid screening before HMM fitting. The default is `TRUE` for an automatic grid and `FALSE` for a supplied grid.
min_state_count, min_state_fraction: A prefilter candidate is retained when its independent soft count is at least the larger of these absolute and fractional thresholds. The zero state is always retained.
screening_prior_sd: Nonnegative extra prior standard deviation used only in the independent prefilter score.
screening_block_size: Positive integer block size for memory-efficient prefilter calculations.
sequence_id: Vector identifying independent sequences. Adjacent equal values belong to one sequence; transitions are not counted across changes in `sequence_id`.
forward_backward_engine: Forward-backward implementation: `"auto"` (default) uses the fast scaled matrix recursion for dense transition masks and the sparse log-domain recursion for sparse masks; `"scaled"` requests the fast recursion with automatic log-domain fallback; `"log"` always uses the reference log-domain recursion.
inference: State-inference method. `"exact"` uses ordinary forward–backward in every EM iteration. `"variational"` uses a checkerboard coordinate-ascent approximation `q(Q[1:T]) = prod_t q_t(Q[t])` during training and, by default, performs one exact forward–backward pass for the returned posterior and marginal likelihood. Variational inference requires a full transition topology.
variational_maxiter: Maximum checkerboard coordinate sweeps inside each variational E-step. The small default exploits warm starts between outer iterations; increase it when a tighter training ELBO matters more than speed.
variational_tolerance: Positive convergence tolerance for the variational state ELBO.
variational_probability_floor: Tiny lower bound on variational state probabilities. It prevents exactly zero transition counts.
variational_final_exact: Logical; after variational training, run one exact forward–backward pass under the fitted parameters. This is strongly recommended and is always done when BIC null selection is requested.
cache_component_emissions: Logical; cache the state-by-scale component log densities and refresh only states whose centers move. This can reduce density evaluations substantially at a modest memory cost.
topology: Transition topology used when `transition_mask` is `NULL`: `"full"` allows all transitions; `"hub"` disallows direct transitions between distinct non-null states.
transition_mask: Optional logical square matrix specifying allowed transitions between the supplied/retained states.
init_transition: Optional initial row-stochastic transition matrix.
init_prob: Optional initial state-probability vector.
init_rho: Optional initial state-by-scale matrix of ash mixture weights.
stay_probability: Initial self-transition probability used when constructing a transition matrix automatically. A scalar is recycled by state.
null_state: Either `"pointmass"` (default), which fixes the null prior to the Dirac mass at zero, or `"adaptive"`, which learns an ash mixture centered at zero.
fixed_pointmass_states: Optional integer indices of states whose ash mixture is fixed entirely on its point component. `NULL` derives the value from `null_state`; an explicit value is an advanced override.
shared_mixture: Logical; if `TRUE`, all free states share one learned vector of ash mixture weights. The default learns each state separately.
estimate_init: Logical; estimate initial state probabilities rather than keeping their initialized values fixed.
transition_prior: Dirichlet parameters for allowed transition probabilities. Supply a scalar or a state-by-state matrix; values must be at least one.
mixture_prior: Dirichlet parameters for ash mixture weights. Supply a scalar, a vector with one entry per scale, or a state-by-scale matrix; values must be at least one.
init_prior: Dirichlet parameters for the initial distribution, used only when `estimate_init = TRUE`.
learn_state_means: Logical; learn eligible non-null state centers after the fixed-grid warm-up.
fixed_mean_states: Integer state indices whose centers are never moved. State 1 should normally remain fixed.
mean_update_start: First EM iteration at which state-center updates and center-eligibility decisions are allowed.
mean_min_effective_count: Minimum smoothed state occupancy required for a center update.
mean_min_pointmass_weight: Minimum fitted weight on the state's point component required for center learning.
mean_min_self_transition: Minimum self-transition probability required for center learning. This favors persistent plateau-like states.
mean_damping: Number in `(0, 1]` multiplying each eligible center move.
mean_bounds: Either `"voronoi"` (default), which keeps centers in nonoverlapping anchor cells, or `"none"`.
prune_states: Logical; enable likelihood-checked dynamic state pruning.
prune_start: First EM iteration eligible for dynamic pruning.
prune_every: Positive number of EM iterations between pruning checks.
prune_min_state_count, prune_min_state_fraction: A state is a low-occupancy pruning candidate below the larger absolute/fractional cutoff.
prune_max_fraction: Maximum fraction of current states considered in one deletion batch. It must lie strictly between zero and one.
prune_max_loglik_loss: Maximum allowed decrease in the deletion-check criterion. This is the full HMM marginal log likelihood under exact training and the mean-field state ELBO under variational training. The historical argument name is retained for compatibility.
merge_distance: Distance below which adjacent fitted centers are considered near duplicates. `NULL` uses `0.05 * median(se)`.
null_model_selection: Either `"bic"` for the strict all-zero safety comparison or `"none"` to retain the fitted HMM unconditionally.
null_selection_gamma: Nonnegative multiplier for the additional candidate-grid term in the strict-null information criterion.
null_bic_margin: Nonnegative margin favoring retention of the fitted HMM over the strict null.
parameter_tolerance: Positive threshold used when counting numerically active parameters for the strict-null information criterion.
step_penalty: Nonnegative penalty for each state change in the separate penalized decoder. `NULL` uses `step_penalty_scale * log(length(y))`.
step_penalty_scale: Nonnegative multiplier for the default step penalty. The default `1.5` accounts for both an added segment level and an unknown change location, reducing transient one-observation steps.
maxiter: Maximum number of generalized EM iterations.
tolerance: Positive relative convergence tolerance for the penalized fixed-dimensional objective.
verbose: Logical; print grid selection and per-iteration diagnostics.

Value

An object of class `ash_hmm_fit`, which is a list containing:

call: The matched function call.
state_probability: An observation-by-state matrix of smoothed probabilities `Pr(Q[t] = m | y)`.
posterior: A list containing the marginal posterior `mean`, `sd`, `probability_ge_zero`, `probability_le_zero`, `probability_zero`, and local false sign rate `lfsr` for every observation.
viterbi_state: The ordinary joint MAP state path under the fitted transition matrix.
penalized_state: The state path from the explicit change-penalized decoder.
step_selection: A list with a segment table, per-sequence counts, total `step_count`, total `change_count`, `occupied_state_count`, decoded `state`, and the applied `penalty`.
boundary_probability: Posterior probabilities of a state change between adjacent observations; entries spanning independent sequences are `NA`.
fitted: Fitted centers and scale grid, mixture weights, transition matrix and mask, initial probabilities, state identifiers and occupancies, null-state type, effect support, constraints, and mean-learning settings.
grid: Automatic-grid settings, original and retained candidates, screening statistics, selected scale grid, and pruning history.
log_likelihood,log_null,log_evidence_ratio: The fitted HMM marginal log likelihood, exact all-zero log likelihood, and their difference.
variational_bound,training_objective,inference: The final variational training bound (or `NA` for exact training), the criterion used during training, and diagnostics recording the inference method, inner convergence, and whether a final exact smoothing pass was used.
model_selection: Details of the optional strict-null comparison, including criteria, effective dimension, and whether the result was collapsed to the exact null.
history: Per-iteration likelihood, objective, number of states, moved centers, and pruned-state count.
mean_history: Long-form history of state centers and occupancies.
pruning_history: One row per removed state, recording its persistent identifier, original grid index, centers, occupancy, and reason.
converged: Logical convergence indicator.
iterations: Number of completed generalized EM/model-reduction iterations recorded after initialization.

Details

If `mu` is `NULL`, a dense grid of candidate state centers is constructed from `y` and screened before the first forward-backward pass. With `nonnegative_state_means = TRUE`, state 1 is fixed at zero, all non-null state centers are constrained to be strictly positive, and every continuous ash component is a normal distribution truncated to `[0, Inf)`. Consequently the latent effects and posterior means are nonnegative as well. Set `nonnegative_state_means = FALSE` to use ordinary Gaussian ash components on the real line and construct a signed symmetric center grid.

Low-occupancy or nearly duplicated states can be proposed for pruning. A proposed deletion is accepted only after recomputing the complete HMM marginal likelihood. After fitting, an optional information-criterion gate can replace the fitted model by the exact all-zero HMM, and a separate change-penalized decoder reports contiguous steps.

Examples

if (FALSE) { # \dontrun{
set.seed(1)
truth <- c(rep(0, 50), rep(2, 50), rep(0, 50))
se <- rep(0.5, length(truth))
y <- rnorm(length(truth), truth, se)

fit <- fit_ash_hmm(y, se, nonnegative_state_means = TRUE)
fit$posterior$mean
fit$step_selection$segments

# Use a signed grid when negative state centers are scientifically possible.
signed_fit <- fit_ash_hmm(y - 1, se, nonnegative_state_means = FALSE)
} # }