RandomEffectStateManagement: Random Effects State Management Routines

In stochtree: Stochastic Tree Ensembles (XBART and BART) for Supervised Learning and Causal Inference

Random Effects State Management Routines

Description

A forest sampler features two types of state: ephemeral and persistent. Persistent state includes objects like ForestSamples and RandomEffectSamples which constitute part of the final sampled model. Ephemeral state supports the sampling computations, but is not retained after the sampler finishes.

The two primary random-effects-based bits of ephemeral state are the RandomEffectsModel and RandomEffectsTracker classes, which represent the current state of a random effects model and its corresponding tracking data structures.

In a linear sampling loop, this ephemeral state is updated with each iteration of the sampler and any retained forests are copied to a RandomEffectSamples object. However, in multi-chain settings, the state of a random effects model must typically be "reset" at the beginning of a new chain. These function enable this process by synchronizing the state of a RandomEffectsModel and RandomEffectsTracker with a corresponding element of a RandomEffectSamples object, or by resetting both to their default (root) state.

resetRandomEffectsModel resets a RandomEffectsModel object based on the parameters indexed by sample_num in a RandomEffectsSamples object. resetRandomEffectsTracker resets a RandomEffectsTracker object based on the parameters indexed by sample_num in a RandomEffectsSamples object.

rootResetRandomEffectsModel resets a RandomEffectsModel object to its "default" state. rootResetRandomEffectsTracker resets a RandomEffectsTracker object to its "default" state.

These functions are intended for advanced use cases in which users require detailed control of sampling algorithms and data structures. Minimal input validation and error checks are performed – users are responsible for providing the correct inputs. For tutorials on the "proper" usage of the stochtree's advanced workflow, we provide several vignettes at https://stochtree.ai/

Usage

resetRandomEffectsModel(rfx_model, rfx_samples, sample_num, sigma_alpha_init)

resetRandomEffectsTracker(
  rfx_tracker,
  rfx_model,
  rfx_dataset,
  residual,
  rfx_samples
)

rootResetRandomEffectsModel(
  rfx_model,
  alpha_init,
  xi_init,
  sigma_alpha_init,
  sigma_xi_init,
  sigma_xi_shape,
  sigma_xi_scale
)

rootResetRandomEffectsTracker(rfx_tracker, rfx_model, rfx_dataset, residual)

Arguments

rfx_model	Object of type RandomEffectsModel.
rfx_samples	Object of type RandomEffectSamples.
sample_num	Index of sample stored in rfx_samples from which to reset the state of a random effects model. Zero-indexed, so resetting based on the first sample would require setting sample_num = 0.
sigma_alpha_init	Initial value of the "working parameter" scale parameter.
rfx_tracker	Object of type RandomEffectsTracker.
rfx_dataset	Object of type RandomEffectsDataset.
residual	Object of type Outcome.
alpha_init	Initial value of the "working parameter".
xi_init	Initial value of the "group parameters".
sigma_xi_init	Initial value of the "group parameters" scale parameter.
sigma_xi_shape	Shape parameter for the inverse gamma variance model on the group parameters.
sigma_xi_scale	Scale parameter for the inverse gamma variance model on the group parameters.

Value

All four functions have no return type and operate in-place on the relevant RandomEffectsModel or RandomEffectsTracker objects

Examples

n <- 100
p <- 10
rfx_group_ids <- sample(1:2, size = n, replace = TRUE)
rfx_basis <- matrix(rep(1.0, n), ncol=1)
rfx_dataset <- createRandomEffectsDataset(rfx_group_ids, rfx_basis)
y <- (-2*(rfx_group_ids==1)+2*(rfx_group_ids==2)) + rnorm(n)
y_std <- (y-mean(y))/sd(y)
outcome <- createOutcome(y_std)
rng <- createCppRNG(1234)
num_groups <- length(unique(rfx_group_ids))
num_components <- ncol(rfx_basis)
rfx_model <- createRandomEffectsModel(num_components, num_groups)
rfx_tracker <- createRandomEffectsTracker(rfx_group_ids)
rfx_samples <- createRandomEffectSamples(num_components, num_groups, rfx_tracker)
alpha_init <- rep(1,num_components)
xi_init <- matrix(rep(alpha_init, num_groups),num_components,num_groups)
sigma_alpha_init <- diag(1,num_components,num_components)
sigma_xi_init <- diag(1,num_components,num_components)
sigma_xi_shape <- 1
sigma_xi_scale <- 1
rfx_model$set_working_parameter(alpha_init)
rfx_model$set_group_parameters(xi_init)
rfx_model$set_working_parameter_cov(sigma_alpha_init)
rfx_model$set_group_parameter_cov(sigma_xi_init)
rfx_model$set_variance_prior_shape(sigma_xi_shape)
rfx_model$set_variance_prior_scale(sigma_xi_scale)
for (i in 1:3) {
    rfx_model$sample_random_effect(rfx_dataset=rfx_dataset, residual=outcome,
                                   rfx_tracker=rfx_tracker, rfx_samples=rfx_samples,
                                   keep_sample=TRUE, global_variance=1.0, rng=rng)
}
resetRandomEffectsModel(rfx_model, rfx_samples, 0, 1.0)
resetRandomEffectsTracker(rfx_tracker, rfx_model, rfx_dataset, outcome, rfx_samples)
rootResetRandomEffectsModel(rfx_model, alpha_init, xi_init, sigma_alpha_init,
                            sigma_xi_init, sigma_xi_shape, sigma_xi_scale)
rootResetRandomEffectsTracker(rfx_tracker, rfx_model, rfx_dataset, outcome)

stochtree documentation built on April 29, 2026, 1:08 a.m.