Register a residual correlation term from another package

Adds a term that correlates the residuals of a gaussian() or student() response, such as brms's sar(), to the terms frm() accepts. Call it from the contributing package's .onLoad(). The term is written in the formula of mu, y ~ x + name(...), as ar() and fcor() are.

Usage

frmtmb_register_autocor(
  name,
  args,
  build,
  whiten,
  unwhiten,
  parse = NULL,
  npar = 0L,
  start = NULL,
  natural = NULL,
  sigma_varies = TRUE,
  bounds = NULL,
  usage = NULL,
  compat_note = NULL,
  jacobian = NULL,
  replace = FALSE,
  package = compat_registrant_name(parent.frame())
)

Arguments

name: The term'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 terms or structures, not a brms name unless core lists it as this package's (sar is frmtmb.spatial's), and not a function formulas call as a term.
args: The term's argument names, in positional order.
build: function(spec, rows, n_all, lookup); see "The term".
whiten, unwhiten: The map and its inverse; see "The map".
parse: NULL, or function(args, env); see "The term". The default keeps the expressions as they are.
npar: The number of entries the term takes in the parameter vector thetaac: a whole number, or function(aux).
start: Starting values: a numeric vector, or function(aux). The default is zeros.
natural: function(theta, aux) giving the term's parameters on their natural scale, named as brms names them (lagsar). It is required when npar can be positive. summary(), confint_varcorr(), variables() and hypothesis() read it.
sigma_varies: TRUE when the term takes a predicted sigma ~ ..., so that whiten and unwhiten receive one sigma per row. FALSE refuses a predicted sigma, and they receive one sigma.
bounds: NULL, or function(aux); see "Bounded parameters and priors".
usage: How the term is written, for messages: "sar(M, type = \"lag\")".
compat_note: The note of the term's row for gaussian() and student() in frm_compat(), where it works: what the term fits. NULL gives a short default. The rows that refuse the other families, the addition terms, rescor, mixtures, non-linear formulas, quadrature and a second residual term are written with the registration.
jacobian: NULL, or function(sigma, theta, aux) giving the Jacobian of whiten in y; see "The Jacobian".
replace, package: As in frmtmb_register_covstruct(). A fit records the package, so that a fit read where it is not installed is refused by its name.

Value

NULL, invisibly. Called for the registration.

Details

A term is ONE affine map. whiten takes the response to innovations u that are independent with unit scale; unwhiten takes innovations back to a response. frmtmb builds everything else from the two:

the density: for gaussian(), \(\sum_i \log \phi(u_i)\) plus the log-Jacobian of the map; for student(), the multivariate Student-t of u with scale matrix I and the response's nu, plus the same log-Jacobian;
the mean of the response, unwhiten(0, ...): for brms's lag SAR this is \((I - \rho W)^{-1} \mu\), which fitted() reports;
the draws of simulate() and predict(), unwhiten() at drawn innovations;
the residual covariance divided by sigma^2 (autocor_matrix()) and the marginal variance that pearson residuals divide by;
the pointwise log-density of frmtmb.sample::log_lik(), brms's leave-one-out conditional density of each row given the others, from the precision matrix K'K of the response, with K the Jacobian of whiten in y (see "The Jacobian").

The term holds for all the rows of a response at once, so the likelihood does not factor into rows. frmtmb therefore refuses the term with weights(), cens(), trunc(), se(), mi(), rescor = TRUE, a mixture, a non-linear formula and quadrature, as it refuses ar(cov = TRUE). The term's fixed data describe the rows of the data, so fitted(), predict() and simulate() take newdata with one row per row of the term's matrix, read in their order (a counterfactual on the same units, as brms takes it). When the fit dropped rows for missing values, newdata covers them too: the term is built again over every row, its parameters held at their natural values (through bounds, for a bounded one), and a row with a missing predictor gives NA. Other newdata is refused, as brms refuses it without a matrix over the new rows in its newdata2.

The term

args names the term's arguments in brms's positional order. parse(args, env) receives the named list of unevaluated argument expressions (NULL for an absent one) and the formula environment; it checks what it can without data and returns the spec. build(spec, rows, n_all, lookup) runs once per fit: rows are the rows of the data that the model keeps (missing values and subset() removed), n_all the number of rows of the data, and lookup(expr, what) resolves an argument expression from data2, the data and the formula environment, as gr(cov = A) resolves its matrix. It returns aux, the fixed data every other function receives. A matrix over the observations is subset with rows, as brms subsets it.

The map

whiten(y, mu, sigma, theta, aux) runs on the TAPE: mu, sigma and theta are RTMB advectors there, y is numeric and mu has one value per row. sigma is ONE value when sigma_varies = FALSE, and one value per row when sigma_varies = TRUE. It returns list(u, logjac): the innovations, one per row, and the log of the absolute determinant of the map's Jacobian in y, one number. Write it without loops over rows; a sparse Matrix times an advector works.

unwhiten(u, mu, sigma, theta, aux) runs on numbers and returns the response with innovations u, with sigma as for whiten. u is a vector, or a matrix with one column per set of innovations, and the result has the same shape: core passes all the columns of the identity at once to read the residual covariance, so a term that factors a matrix factors it once. It must be the inverse of whiten at the same parameters.

When the frame is built, frmtmb checks the map once on numbers: at the start, whiten gives n innovations and one finite log-Jacobian, unwhiten gives n values (and an n x 2 matrix for two columns), and whiten(unwhiten(u)) returns u to a relative error of 1e-6. A term that fails is refused by its name.

The Jacobian

whiten is affine in y, so u = K (y - m) with m the mean unwhiten(0, ...) and K its Jacobian. jacobian(sigma, theta, aux) returns K, n x n, on numbers, with sigma as whiten receives it. A sparse Matrix keeps the pointwise log-density at the cost of its nonzeros, which matters for a large n. Without the slot, frmtmb reads K as the inverse of unwhiten() of the identity: exact, and one dense n x n solve per parameter vector. Measured with sar()'s lag map registered without its slot, one draw takes 0.03 s and 21 Mb at n = 500, 0.5 s and 48 Mb at 1000, 1.7 s and 155 Mb at 2000, and 5 s and 345 Mb at 3000, growing as n^3; with the slot, 1600 rows take 2 ms. Above 1000 rows log_lik() warns once per call (class frmtmb_autocor_dense_fallback), naming the term and the slot. When the frame is built, a given jacobian is checked on numbers at two points: the start with sigma = 1, and a generic point, with sigma 1.37 (or exp(0.3 sin(i)) on row i under sigma_varies = TRUE) and every parameter moved from the start by 0.4 times its position (then -0.4, 0.04 and -0.04, when the map cannot be evaluated there; if none can be, the start, with a warning of class frmtmb_autocor_jacobian_unchecked). At each point K (y - m) must be whiten(y) for two directions of y - m, to a relative error of 1e-6. A slot that forgets sigma, divides rows where it should divide columns, ignores a parameter or transposes an asymmetric matrix is refused by the term's name.

Bounded parameters and priors

Give bounds when each parameter is a bounded quantity mapped as lo + (hi - lo) / (1 + exp(-theta)), as sar()'s rho is: function(aux) returning a matrix with columns lower and upper, one row per parameter, in the order of natural. Then set_prior(class = "<natural name>") puts a density on the natural parameter, carried to theta with the log-Jacobian of the map, and frmtmb_autocor_pars() reports the interval, from which frmtmb.sample takes brms's flat default on it.

Testing a term

Compare the fitted log-likelihood with a dense multivariate normal density, mvtnorm::dmvnorm() at the mean unwhiten(0, ...) and the covariance of the map, and check that whiten(unwhiten(u)) is u. A test that registers calls frmtmb_local_registries() first.

Examples

# a residual correlation with a FIXED matrix V and no parameter,
# myfix(V): y = mu + sigma * L u with L the Cholesky factor of V
myfix_build <- function(spec, rows, n_all, lookup) {
  V <- lookup(spec$V, "myfix()")[rows, rows, drop = FALSE]
  L <- t(chol(V))
  list(L = L, Linv = forwardsolve(L, diag(nrow(L))),
       logdet = sum(log(diag(L))))
}
# sigma_varies = FALSE: sigma is one value
myfix_whiten <- function(y, mu, sigma, theta, aux) {
  u <- as.vector(aux$Linv %*% ((y - mu) / sigma))
  list(u = u, logjac = -aux$logdet - length(y) * log(sigma))
}
myfix_unwhiten <- function(u, mu, sigma, theta, aux) {
  mu + sigma * (aux$L %*% u)
}
# u = Linv (y - mu) / sigma, so the Jacobian is Linv / sigma: right
# for one sigma (sigma_varies = FALSE). A sigma per row divides the
# COLUMNS, Linv %*% diag(1 / sigma), as fcor() does
myfix_jacobian <- function(sigma, theta, aux) aux$Linv / sigma
# a package calls frmtmb_register_autocor() from its .onLoad(); here
# the registration lasts while demo() runs
demo <- function() {
  frmtmb_local_registries()
  frmtmb_register_autocor("myfix", args = "V", build = myfix_build,
                          whiten = myfix_whiten,
                          unwhiten = myfix_unwhiten,
                          jacobian = myfix_jacobian,
                          sigma_varies = FALSE)
  set.seed(1)
  V <- 0.5^abs(outer(1:20, 1:20, "-"))
  d <- data.frame(x = rnorm(20))
  d$y <- 1 + d$x + as.vector(t(chol(V)) %*% rnorm(20))
  fit <- frm(y ~ x + myfix(V), data = d, data2 = list(V = V))
  # the same model as core's own fcor(V)
  c(myfix = logLik(fit),
    fcor = logLik(frm(y ~ x + fcor(V), data = d, data2 = list(V = V))))
}
demo()
#>     myfix      fcor 
#> -21.15892 -21.15892