# Writing Model Codes

This page explains how to describe models using the `rtmb_code` block, which is the core of BayesRTMB, and details the available probability distributions and mathematical functions.

## 1. Overview of the Structure and Role of rtmb_code

`rtmb_code()` is a function that allows you to define Bayesian models using an intuitive syntax similar to Stan. A model is described by dividing it into several “blocks” according to their purposes.

    library(BayesRTMB)

    code <- rtmb_code(
      setup = {
        # Data preprocessing and constant definitions
      },
      parameters = {
        # Declaration of parameters to be estimated
      },
      transform = {
        # Calculation of derived variables using parameters and data
      },
      model = {
        # Definition of priors and likelihoods
      },
      generate = {
        # Calculation of posterior predictions and generated quantities
      }
    )

### 1-1. parameters block

This block is for declaring the parameters you want to estimate in the model. You use the `Dim()` function to specify the dimension and constraints (e.g., lower and upper bounds) of the variables.

Example Code:

    parameters = {
      # Scalar parameter (with lower bound of 0)
      sigma = Dim(lower = 0)

      # Vector parameter (length N)
      mu = Dim(N)

      # Matrix parameter (N rows, M columns)
      X = Dim(c(N, M))

      # Array parameter (I x J x K)
      A = Dim(c(I, J, K))
    }

Here we are defining generic continuous parameters, but by using the `type` argument you can define parameters (types) with more complex structures, such as “correlation matrices” or “sum-to-zero vectors” (see “3. About Parameter Types (Dim)” below for details).

### 1-2. model block

This is the heart of the model where you define priors and the likelihood (the data generating process). Similar to Stan, using the `~` (tilde) sampling syntax is recommended.

    model = {
      # Priors
      mu ~ normal(0, 10)
      sigma ~ exponential(1)

      # Likelihood
      Y ~ normal(mu, sigma)
    }

The same model can also be written by explicitly adding log densities:

    model = {
      lp <- lp + normal_lpdf(mu, 0, 10)
      lp <- lp + exponential_lpdf(sigma, 1)
      lp <- lp + normal_lpdf(Y, mu, sigma)
    }

In other words, `Y ~ normal(mu, sigma)` and `lp <- lp + normal_lpdf(Y, mu, sigma)` are two ways to write the same model. The `~` syntax is usually easier to read, while `lp <- lp + ..._lpdf(...)` is useful when you need to build a custom likelihood. `lp` is reserved for accumulating log probability, so do not declare it in the `parameters` block or use it as an ordinary working variable.

### 1-3. setup block

This block is executed exactly once before the model computations begin. It is primarily used to extract data sizes (number of rows/columns), prepare model variables, and create constants required for computations.

The reserved name `.data` refers to the original object supplied to the `data` argument of `rtmb_model()`. It can therefore be used with a matrix or data frame directly, or with elements of a named list. The `.data` binding is read-only, is available only while `setup` is evaluated, and is not retained in the model data afterward.

Example Code:

    setup = {
      # For data = list(Y = y, X = x)
      Y <- as.numeric(.data$Y)
      X <- as.matrix(.data$X)
      N <- length(Y)
      P <- ncol(X)
    }

When the object passed to `data` is itself a matrix or data frame, use `as.matrix(.data)`:

    setup = {
      Y <- as.matrix(.data)
      N <- nrow(Y)
      P <- ncol(Y)
    }

Columns of a data frame and elements of a named list also remain available directly by name inside `setup`. Using `.data` is useful when the generated code should show explicitly which values came from the original input object.

For common data-only centering operations, BayesRTMB provides `center_grand_mean()` and `center_within_cluster()`. These helpers can be used inside `setup` as well as in ordinary R code.

    setup = {
      x_gmc <- center_grand_mean(.data$x)
      x_cwc <- center_within_cluster(.data$x, .data$group)
    }

### 1-4. transform block

This block is for creating derived variables (such as the mean `mu` or linear predictors) necessary for likelihood calculations from parameters and data. The variables created here can be used directly within the `model` block.

> [!IMPORTANT] The derived variables calculated here are automatically saved inside the model object. Not only can you obtain samples of their posterior distributions after estimation, but when you run MAP estimation (`optimize`), standard errors and 95% confidence intervals for these derived variables are also automatically calculated and output using the delta method.

### 1-5. generate block

This block is used for creating predictive distributions for Posterior Predictive Checks (PPC) and for generating new quantities of interest calculated from estimated parameters.

> [!IMPORTANT] Computations written in the `generate` block are not included in the likelihood evaluation (Automatic Differentiation calculations) and can be executed post-hoc after the estimation is fully complete. Therefore, no matter how computationally heavy the operations written here are, they will not negatively affect the estimation speed (MAP or MCMC). Also, if you forgot to write something during model definition or later decide you want to compute another quantity, you can calculate these retroactively by calling the `fit$generated_quantities(new_generate_code)` method on the fitted object.

## 2. Available Probability Distributions

In BayesRTMB, you can use many probability distributions within the `model` block. Most univariate distributions are vectorized, meaning if `Y` and `mu` are vectors, a single statement efficiently calculates the sum of log-densities for all elements.

### 2-1. Common Probability Density and Mass Functions

Continuous Probability Distributions (LPDF)

Function | Description
--- | ---
`normal(mean, sd)` | Normal distribution.
`lognormal(meanlog, sdlog)` | Log-normal distribution.
`exp_mod_normal(mu, sigma, lambda)` | Exponentially modified normal distribution.
`exponential(rate)` | Exponential distribution.
`cauchy(location, scale)` | Cauchy distribution.
`student_t(df, mu, sigma)` | Student’s t-distribution.
`gamma(shape, rate)` | Gamma distribution.
`inverse_gamma(shape, scale)` | Inverse-gamma distribution.
`beta(a, b)` | Beta distribution.

Discrete Probability Distributions (LPMF)

Function | Description
--- | ---
`bernoulli(prob)` | Bernoulli distribution (binary data).
`binomial(size, prob)` | Binomial distribution.
`poisson(mean)` | Poisson distribution (count data).
`neg_binomial_2(mu, size)` | Negative binomial distribution (mean-variance parameterization).

### 2-2. Advanced Probability Distributions

Distributions are also provided to improve computational stability or for specific modeling purposes.

Function | Description
--- | ---
`bernoulli_logit(eta)` / `binomial_logit(size, eta)` | Functions that take the linear predictor (`eta`) on the logit scale directly instead of probabilities. Numerically stable.
`ordered_logistic(eta, cutpoints)` | Used for ordinal categorical data, such as in ordered logistic regression.
`lkj_corr(eta)` | LKJ prior distribution for correlation matrices.

### 2-3. Probability Distributions for Special Types

Distributions tailored for matrices, multivariate data, and specific model structures.

Function | Description
--- | ---
`multi_normal(mean, Sigma)` | Standard multivariate normal distribution.
`multi_normal_CF(mean, sd, CF_Omega)` | Multivariate normal distribution using a Cholesky factor parameterization.
`dirichlet(alpha)` | Dirichlet distribution for a simplex (a vector summing to 1).
`lower_tri_normal(mean, sd)` | Normal distribution over the elements of a lower triangular matrix.
`centered_tri_multi_normal(sigma)` | Multivariate normal distribution for a centered triangular matrix (used for identifiability constraints).
`sufficient_multi_normal_fa(S_mat, N, y_bar, mu, psi, Lambda)` | Likelihood for factor analysis using sufficient statistics. Extremely computationally efficient for large samples.
`normal_mixture(pi_w, mean, sd)` | Normal mixture distribution combining multiple normal distributions.
`diffusion(alpha, tau, beta, delta)` | Drift diffusion model for response times and binary choices. Use as `obs(RT, Choice) ~ diffusion(...)`.
`wishart(n, V)` | Wishart distribution.

### 2-4. User-Defined Distributions

You can also define a custom log-density function in the `setup` block and use it with the sampling syntax. The function name should end in `_lpdf` for continuous distributions or `_lpmf` for discrete distributions. For example, `exgaussian_lpdf()` can be used in the model block as `Y ~ exgaussian(mu, sigma, lambda)`.

The following example simulates response times in seconds from an ex-Gaussian distribution and fits the same model. Here, `lambda` is the rate of the exponential component.

    library(BayesRTMB)

    simulate_exgaussian <- function(N, mu, sigma, lambda) {
      rnorm(N, mean = mu, sd = sigma) + rexp(N, rate = lambda)
    }

    set.seed(1234)

    true <- c(mu = 0.45, sigma = 0.06, lambda = 8)
    rt <- simulate_exgaussian(
      N = 500,
      mu = true["mu"],
      sigma = true["sigma"],
      lambda = true["lambda"]
    )

    hist(rt)

The user-defined distribution is placed in `setup`. The model block can then use the same sampling syntax as built-in distributions.

    code_exgaussian <- rtmb_code(
      setup = {
        Y <- rt

        exgaussian_lpdf <- function(x, mu, sigma, lambda, sum = TRUE) {
          z <- (x - mu) / sigma - lambda * sigma

          res <-
            log(lambda) +
            lambda * (mu - x) +
            0.5 * (lambda * sigma)^2 +
            RTMB::pnorm(z, log.p = TRUE)

          if (sum) sum(res) else res
        }
      },

      parameters = {
        mu     <- Dim(lower = 0)
        sigma  <- Dim(lower = 0)
        lambda <- Dim(lower = 0)
      },

      model = {
        Y ~ exgaussian(mu, sigma, lambda)

        mu     ~ exponential(1)
        sigma  ~ exponential(1)
        lambda ~ exponential(1 / 10)
      },

      generate = {
        log_lik <- exgaussian_lpdf(Y, mu, sigma, lambda, sum = FALSE)
        report(log_lik)
      }
    )

For this model, it is helpful to provide initial values on the scale of the observed response times.

    init_exgaussian <- list(
      mu = median(rt) - sd(rt) / 2,
      sigma = sd(rt) / 2,
      lambda = 1 / (sd(rt) / 2)
    )

    mdl_exgaussian <- rtmb_model(
      data = list(rt = rt),
      code = code_exgaussian,
      init = init_exgaussian
    )

    fit_map <- mdl_exgaussian$optimize(num_estimate = 5)

    est_map <- unlist(fit_map$estimate(pars = "parameters"))[names(true)]
    data.frame(
      parameter = names(true),
      true = as.numeric(true),
      estimate = as.numeric(est_map),
      error = as.numeric(est_map - true)
    )

The same model can be sampled with MCMC or fit with variational inference.

    fit_mcmc <- mdl_exgaussian$sample(parallel = TRUE)
    fit_mcmc
    fit_mcmc$WAIC()

    fit_vb <- mdl_exgaussian$variational()
    fit_vb

For a custom model, posterior predictions can be supplied after fitting by passing a `generate` block. The block must return one replicated outcome for each posterior draw. Because `exgaussian_lpdf()` ends in `_lpdf`, `type = "auto"` selects the continuous density check.

    predict_exgaussian <- rtmb_code(
      generate = {
        y_rep <- rnorm(length(Y), mu, sigma) + rexp(length(Y), rate = lambda)
        report(y_rep)
      }
    )

    yrep <- fit_mcmc$posterior_predict(
      code = predict_exgaussian,
      draws = 100,
      seed = 123
    )
    fit_mcmc$pp_check(
      code = predict_exgaussian,
      observed = "Y",
      type = "auto",
      draws = 100,
      seed = 123
    )

When WAIC is needed, the `generate` block should save an observation-wise `log_lik` by calling the custom `_lpdf` with `sum = FALSE`.

## 3. About Parameter Types (Dim)

In the `parameters` block, you define the types and dimensions of parameters using `Dim()`.

- Scalar: `Dim()` or `Dim(1)`
- Vector: `Dim(N)`
- Matrix: `Dim(c(N, M))`
- Array: `Dim(c(I, J, K))`

You can also enforce various constraints using arguments: * `lower = 0, upper = 1`: Specifies the range of values (e.g., for variances or probabilities). * `random = TRUE`: Designates the parameter as a random effect subject to Laplace approximation in hierarchical models.

Furthermore, by specifying a specific keyword in the `type` argument, you can define complex types with specific structures.

`type` Specification | Description
--- | ---
`type = "simplex"` | A vector where all elements are positive and sum to 1 (used for probability assignments, etc.).
`type = "centered"` | A vector where the sum of all elements is zero (used for main effects in ANOVA-type models, etc.).
`type = "corr_matrix"` | A correlation matrix with 1s on the diagonal satisfying positive definiteness.
`type = "CF_corr"` | The Cholesky factor of a correlation matrix (sometimes more computationally efficient than the correlation matrix itself).
`type = "lower_tri"` | A lower triangular matrix.
`type = "centered_tri"` | A lower triangular matrix constrained so the column sums are zero (useful for identification constraints in factor analysis).

## 4. About Mathematical Functions

Numerous numerically stable utility functions are available to prevent overflow and underflow during Automatic Differentiation (AD) calculations.

Link and Inverse Link Functions

Function | Description
--- | ---
`logit(x)` | Computes the logit transformation `log(x/(1-x))`.
`inv_logit(x)` | Computes the inverse logit (logistic) transformation `1/(1+exp(-x))`.

Functions for Numerical Stabilization

Function | Description
--- | ---
`log_sum_exp(x)` | Safely computes `log(sum(exp(x)))` using the log-sum-exp trick.
`log1p_exp(x)` | Stably computes `log(1 + exp(x))`.
`log1m_exp(x)` | Stably computes `log(1 - exp(x))` for `x < 0`.
`log_softmax(x)` | Computes the logarithm of the softmax function.
`fabs(x)` | A smoothed version of `abs(x)` that guarantees differentiability near zero.

Matrix and Vector Transformations

These are useful when transforming an unconstrained vector into a matrix with a specific structure.

Function | Description
--- | ---
`centered(x)` | Transforms a vector of length K-1 into a vector of length K that sums to zero.
`to_lower_tri(x, M, D)` | Creates an M x D lower triangular matrix from the vector `x`.
`to_centered_matrix(x, R, C)` | Creates an R x C matrix where the sum of each column is zero.
`to_centered_tri(x, R, C)` | Creates a matrix with lower triangular elements whose column sums are zero, used for identification constraints in factor analysis.

Linear Algebra for AD

Function | Description
--- | ---
`log_det_chol(L)` | Computes the log-determinant of a covariance matrix from its Cholesky factor `L`.
`quad_form_chol(x, L)` | Computes a quadratic form using the Cholesky factor `L`.
`distance(x, y)` | Computes the Euclidean distance between two vectors, adding a small epsilon for numerical stability.

## 5. AD-Compatible Working Containers

In most models, the fastest and simplest approach is to write vectorized code. For example, prefer a single vectorized likelihood such as `Y ~ normal(mu, sigma)` over manually filling a temporary vector inside a loop.

When you do need a mutable working vector or array whose elements will later be assigned AD values, use `rtmb_vector()` or `rtmb_array()` instead of `numeric()`, `matrix()`, or `array()`.

    model = {
      eta <- rtmb_vector(0, N)
      for (i in seq_len(N)) {
        eta[i] <- alpha + X[i, ] %*% beta
      }
      Y ~ normal(eta, sigma)
    }

    model = {
      logit_x <- rtmb_array(0, dim = c(N_time, C, D))
      for (t in seq_len(N_time)) {
        for (c in seq_len(C)) {
          for (d in seq_len(D)) {
            logit_x[t, c, d] <- alpha[d] + beta[d] * X[t, c]
          }
        }
      }
    }

Inside `rtmb_code()`, BayesRTMB usually provides the AD seed automatically from the model parameters, so you normally do not need to set the `seed` argument. These helpers are mainly for assignment-heavy code. If the same calculation can be expressed directly with vectorized matrix or array operations, the vectorized form is usually preferable for both tape construction and repeated AD evaluation.

## 6. Upgrading Saved Fit Objects

Saved fit objects created by older BayesRTMB versions may not automatically have methods added in newer versions. Use `upgrade_fit()` to rebuild an existing MCMC, VB, MAP, or classic fit object with the currently loaded class definitions.

    fit2 <- upgrade_fit(fit)

This does not refit the model. It copies the stored estimation results into a new fit object, and by default also rebuilds the embedded model object when possible.

## Conclusion

Rather than writing a complex model from the beginning, it’s better to start estimating with simple models first. The error messages are still under development, and we plan to make them more user-friendly in the future.

For standard wrapper models, see Wrapper Functions. For mixed models and GLMMs in particular, see Hierarchical Models and GLMMs. For the internal inference pipeline, see RTMB Internals and Inference Algorithms. For function-level reference material, see Analysis Reference.