Register a covariance structure from another package

Adds a structured covariance, or a precision matrix that depends on parameters, to the structures frm() accepts. Call it from the contributing package's .onLoad(). A structure is written in a formula in one of two ways:

Usage

frmtmb_register_covstruct(
  name,
  lpdf,
  vcov,
  npar,
  start = NULL,
  sd_idx = NULL,
  natural = NULL,
  draw = NULL,
  levels_coupled = FALSE,
  expand = NULL,
  term = NULL,
  from_natural = NULL,
  chol = NULL,
  full_rank = NULL,
  replace = FALSE,
  package = compat_registrant_name(parent.frame())
)

Arguments

name: The structure's name, as a formula writes it, without parentheses. It must be a syntactic name that nothing else in a formula means: not one of frmtmb's own structures, specials, families or addition terms, not a brms name, not a name core lists as another extension's (only that package registers it), and not a function that formulas call as a term, such as poly, ns, log or factor.
lpdf: function(b, theta, blk): the log density of the block's b, summed over levels. AD-safe. The frame checks only that it is finite at b = 0 and the start: a correct density can be flat along a direction, centered away from zero or bimodal, so its shape tells nothing, and a density with the wrong sign is for the structure's own test to catch ("Testing a structure").
vcov: function(theta, blk): the numeric dim x dim covariance within one level, with the coefficient names as dimnames. VarCorr() reports it, and an unseen level is drawn from it unless levels_coupled is TRUE.
npar: The number of entries the structure takes in theta: a whole number, or function(blk).
start: Starting values for theta: a numeric vector, or function(blk). The default is zeros.
sd_idx: The positions within the block's theta that are LOG standard deviations: an integer vector, or function(blk). They take the SE checks' boundary reading and class "sd" priors. The default is none.
natural: NULL, or function(theta, blk) for a field; see "Rows a field reports".
draw: NULL, or function(theta, blk) returning one draw of the block's whole b. Give it when the levels are correlated; the default draws each level from vcov.
levels_coupled: TRUE when a fixed matrix correlates the LEVELS of the block. The block then has no law for an unseen level (prediction gives one no block contribution), it is one unit for vcov_cluster() and the sandwich, and importance = refuses it.
expand: NULL, or a list for a block whose coefficients are a fixed linear map P b of its parameters: apply(b, blk) computes P b (AD-safe), jacobian(blk) returns P as triplets list(i, j, x) in within-block positions, and condsd(v, blk) gives the conditional standard deviations of P b from the variances v of b (optional; without it ranef() reports NA). It acts on blocks whose build returned expands = TRUE.
term: NULL for a bar structure, or the field term; see "A field term".
from_natural: NULL, or function(sds, C, blk) returning the whole theta segment from natural-scale standard deviations and a correlation matrix, for frm_simulate()'s newparams in brms's names.
chol: NULL, or function(theta, blk) returning the lower Cholesky factor of vcov, on the tape. With it, and with every theta entry a standard deviation, frmtmb.sample::frm_sample() can sample the block non-centered.
full_rank: TRUE when an effect a row owns can reach any value with a finite density: the covariance is positive definite at some parameter value and the density is positive everywhere, as for a gaussian or Student-t latent. A weighted model whose every row owns such an effect has no maximum when the weights sum to more than the rows, and frm() refuses it, as it refuses (1 | obs). FALSE when the effects cannot (a singular covariance, a latent with bounded support). NULL, the default, states nothing, and the check stays open for the structure. A package whose structure is full rank should state full_rank = TRUE, so that a weighted model with no maximum is refused rather than fitted to a runaway.
replace: A registered name is replaced when the package that registered it registers it again, so a reload is not an error. A name another package, or the session, registered is replaced only with replace = TRUE; without it the call refuses.
package: The package that owns the structure. A fit records it for each block, so that a fit read where that package is not installed is refused by the package's name before any method uses the structure. The default is the package whose namespace calls, and NULL outside a package.

Value

NULL, invisibly. Called for the registration.

Details

as a BAR STRUCTURE, name(terms | g), like ar1(0 + t | g): the block has one coefficient per term at each level of g, and the structure says how the coefficients of one level covary. Give no term.
as a FIELD, a predictor special such as car(M, gr = g): the block is one intercept per level of the grouping variable gr, and the structure says how the levels covary, usually through a fixed matrix the special names. Give term.

Everything else frmtmb does with a block reads the entry this call builds: the objective, the standard errors and their checks, VarCorr(), confint_varcorr(), ranef(), predict() with and without new levels, simulate(), conditional_effects(), vcov_cluster() and frmtmb.sample::frm_sample(). The built-in structures are entries of the same shape.

The block

Each function receives blk, the block: a list with dim (the number of coefficients per level), n_levels, levels, cnms (the coefficient names), theta_idx and b_idx, and for a field aux, the fixed data term$build returned. b is level-major: the dim coefficients of one level are contiguous, so matrix(b, nrow = dim) has one level per column. Read the block with [[, never $; see frm_hazard_reads().

lpdf runs on the TAPE: b and theta are RTMB advectors there. Use RTMB::matrix() and RTMB::dmvnorm() or RTMB::dgmrf(), and write the density without loops over levels. The other functions run on numbers.

A field term

term is a list:

args: The special's argument names in positional order. It must include "gr", the grouping variable. frmtmb matches a call against them by position and by exact name, and refuses an unknown or misspelled argument.
parse: function(args, env): args is the named list of unevaluated argument expressions (NULL for an absent one), and env the formula environment. Check the arguments and return a list, the spec, which build receives. Do not evaluate data here; the data are not known yet.
build: function(spec, gr, lookup): gr holds the grouping values of the rows, and lookup(expr, what) resolves an argument expression from data2, the data and the formula environment, as gr(cov = A) resolves its matrix. Return list(levels, index, aux): levels the block's level labels, index the level of each row (whole numbers into levels), and aux the fixed data the other functions read as blk[["aux"]]. Add expands = TRUE when the block's coefficients are a fixed linear map of its parameters (see expand).
usage: How the special is written, for messages: "car(M, gr = g)".
levels_of: Optional. function(gr) giving the level labels of grouping values, when they are not as.character(gr). It is used on new data too.
levels_complete: Optional. TRUE when levels is the whole index set of the structure (every mesh node), so a value outside it is not an unseen level.

Rows a field reports

A field's theta is usually not a set of brms sd_<group>__<coef> and cor_ parameters. Give natural for such a structure: function(theta, blk) returning list(term, type, est, grad), one entry per reported row. type is "sd" or "range" (estimated on the log scale), "prop" (logit), "cor" (Fisher z) or "raw"; est is the value on that scale and grad its jacobian in theta, a matrix with one row per term. confint_varcorr() reports the rows with delta-method intervals. A structure with natural is a field for brms's names too: ngrps(), VarCorr() and variables() do not count it as a grouping factor.

Testing a structure

Compare the fitted log-likelihood with one you compute independently: for a gaussian response, integrate the block out by hand and evaluate the marginal density with mvtnorm::dmvnorm() or a dense solve(). A test that registers calls frmtmb_local_registries() first, which restores every registry when the test ends.

Examples

# an exchangeable block with a FIXED correlation of 0.5 and one
# estimated standard deviation: xcs(0 + t | g)
xcs_lpdf <- function(b, theta, blk) {
  d <- blk[["dim"]]
  C <- RTMB::matrix(0.5, d, d) + diag(0.5, d)
  Sigma <- exp(2 * theta[1]) * C
  dim(b) <- c(d, length(b) %/% d)
  sum(RTMB::dmvnorm(t(b), 0, Sigma, log = TRUE))
}
xcs_vcov <- function(theta, blk) {
  d <- blk[["dim"]]
  V <- exp(2 * theta[1]) * (matrix(0.5, d, d) + diag(0.5, d))
  dimnames(V) <- list(blk[["cnms"]], blk[["cnms"]])
  V
}
if (FALSE) { # \dontrun{
# in the contributing package's .onLoad()
frmtmb_register_covstruct("xcs", lpdf = xcs_lpdf, vcov = xcs_vcov,
                          npar = 1L, sd_idx = 1L)
} # }