---
title: "Writing Model Codes"
description: "A guide on how to write model codes in BayesRTMB."
output: 
  rmarkdown::html_vignette:
    toc: true
    toc_depth: 3
vignette: >
  %\VignetteIndexEntry{Writing Model Codes}
  %\VignetteEncoding{UTF-8}
  %\VignetteEngine{knitr::rmarkdown}
editor_options: 
  chunk_output_type: console
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE, eval = FALSE)
```

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.

```{r eval=FALSE}
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:**
```{r eval=FALSE}
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.

```{r eval=FALSE}
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:

```{r eval=FALSE}
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) as variables from the `data` list passed from R, and to create constants required for computations.

**Example Code:**
```{r eval=FALSE}
setup = {
  N <- length(Y)  # Get the sample size of the observed data
  P <- ncol(X)    # Get the number of columns in the design matrix (number of predictors)
}
```

### 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.

```{r}
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.

```{r}
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.

```{r}
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.

```{r}
fit_mcmc <- mdl_exgaussian$sample(parallel = TRUE)
fit_mcmc
fit_mcmc$WAIC()

fit_vb <- mdl_exgaussian$variational()
fit_vb
```

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()`.

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

```{r eval=FALSE}
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.

```{r eval=FALSE}
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](wrapper_functions.html). For mixed models and GLMMs in particular, see [Hierarchical Models and GLMMs](rtmb_glmer.html). For the internal inference pipeline, see [RTMB Internals and Inference Algorithms](rtmb_internals.html). For function-level reference material, see [Analysis Reference](analysis_reference.html).
