--- title: "Guide to adding a distribution" output: rmarkdown::html_vignette vignette: > %\VignetteIndexEntry{adding-a-distribution} %\VignetteEngine{knitr::rmarkdown} %\VignetteEncoding{UTF-8} --- ```{r, include = FALSE} knitr::opts_chunk$set( collapse = TRUE, comment = "#>" ) ``` A new distribution is one R file, one test file, and one line in the distribution list. The complete example below is meant to be copied; the sections after it cover the cases it does not. ## Calling other distributions `RTMB` records a computation tape by operator overloading, so a density has to be ordinary R arithmetic on whatever it is handed. That is less restrictive than it sounds. `RTMB` promotes around 37 `stats` functions (`dnorm`, `dgamma`, `pbeta`, `lgamma`, `besselK` and so on) to S4 generics with a method for AD variables, so you can just write the bare name. With `@import RTMB` on your function, `dgamma(x, shape)` dispatches on its arguments: ordinary numbers get the usual `stats` behaviour, AD variables get the tape-aware method. ```r logdens <- dgamma(x, shape = shape, scale = scale, log = TRUE) # yes logdens <- stats::dgamma(x, shape, scale, log = TRUE) # no - breaks taping logdens <- RTMB::dgamma(x, shape, scale, log = TRUE) # unnecessary ``` `stats::` pins the non-AD version and the tape breaks. `RTMB::` works, but it is noise and it hides the dispatch. For a function `RTMB` does not cover, see [When `RTMB` has no AD version](#when-rtmb-has-no-ad-version). ## A complete example This is the whole of `R/foo.R`, for a one-parameter distribution with density \eqn{f(x;\theta) = \theta e^{-\theta x}} on \eqn{x \ge 0}. ```r #' Foo distribution #' #' Density, distribution function, quantile function, and random generation #' for the Foo distribution. #' #' @details #' The Foo distribution with rate \eqn{\theta > 0} has density #' \deqn{f(x;\,\theta) = \theta e^{-\theta x}, \quad x \ge 0.} #' #' @param x,q vector of quantiles #' @param p vector of probabilities #' @param n number of random values to return #' @param theta positive rate parameter #' @param log,log.p logical; if \code{TRUE}, probabilities/ densities are returned as \eqn{\log(p)}. #' @param lower.tail logical; if \code{TRUE} (default), probabilities are \eqn{P[X \le x]}, otherwise \eqn{P[X > x]}. #' #' @return #' \code{dfoo} gives the density, \code{pfoo} gives the distribution function, #' \code{qfoo} gives the quantile function, and \code{rfoo} generates random deviates. #' #' @examples #' set.seed(123) #' x <- rfoo(5, theta = 2) #' d <- dfoo(x, theta = 2) #' p <- pfoo(x, theta = 2) #' q <- qfoo(p, theta = 2) #' @name foo NULL #' @rdname foo #' @export #' @import RTMB dfoo <- function(x, theta, log = FALSE) { if (!ad_context()) { args <- as.list(environment()) simulation_check(args) # informative error if the likelihood is written in the wrong order if (any(theta <= 0)) stop("theta must be > 0") } # potentially escape to RNG or CDF if (inherits(x, "simref")) return(dGenericSim("dfoo", x = x, theta = theta, log = log)) if (inherits(x, "osa")) return(dGenericOSA("dfoo", x = x, theta = theta, log = log)) logdens <- log(theta) - theta * x + log(1 - smaller(x, 0)) # 0 outside x >= 0 if (log) return(logdens) return(exp(logdens)) } #' @rdname foo #' @export pfoo <- function(q, theta, lower.tail = TRUE, log.p = FALSE) { if (!ad_context()) { if (any(theta <= 0)) stop("theta must be > 0") } p <- (1 - smaller(q, 0)) * (-expm1(-theta * q)) if (!lower.tail) p <- 1 - p if (log.p) p <- log(p) return(p) } #' @rdname foo #' @export qfoo <- function(p, theta, lower.tail = TRUE, log.p = FALSE) { if (!ad_context()) { if (any(theta <= 0)) stop("theta must be > 0") } if (log.p) p <- exp(p) # undo the transformations first ... if (!lower.tail) p <- 1 - p # ... in the reverse order to pfoo if (!ad_context()) { if (any(p < 0 | p > 1)) stop("p must be in [0, 1]") } -log1p(-p) / theta } #' @rdname foo #' @export #' @importFrom stats runif rfoo <- function(n, theta) { if (any(theta <= 0)) stop("theta must be > 0") n <- ceiling(n) qfoo(runif(n), theta) } ``` Five details in there are easy to get wrong: - Validation and `simulation_check` belong inside `if (!ad_context())`. On AD variables they either fail or give the wrong answer. - The string passed to `dGenericSim`/`dGenericOSA` has to match the function name exactly. `RTMB` uses it to find your `p` and `r` functions. - Work in log space and call `exp()` once, at the end. - Check `p` after undoing `log.p` and `lower.tail`, not before. Checking first makes `log.p = TRUE` impossible, since `log(p)` is negative. - Match the support indicator to the boundary. Foo has support \eqn{x \ge 0}, so `1 - smaller(x, 0)` is right. `greater(x, 0)` is 0 at `x = 0` and would give a density of 0 there instead of \eqn{\theta}. Then the test file, and one bullet in `vignettes/distlist.Rmd`: ``` * [`foo(theta)`](../reference/foo.html): Foo distribution parameterised by rate `theta` ``` ## Tests `tests/testthat/test-foo.R`. Two parameter settings plus a gradient check is the minimum. ```r test_that("foo passes standard distribution checks (theta = 1)", { check_continuous_dist(dfoo, pfoo, qfoo, xs = c(0.2, 0.5, 1, 2), lower = 0, upper = Inf, theta = 1) }) test_that("foo AD gradient has no NaN", { check_ad_gradient(dfoo, rfoo, theta = 1) }) ``` | Your distribution is | Use | |---|---| | continuous | `check_continuous_dist(dfun, pfun, qfun, xs, lower, upper, ...)` | | continuous, no quantile function | same, with `qfun = NULL` | | discrete | `check_discrete_dist(dfun, pfun, xs_int, sum_support, ...)` | | discrete, no cdf | same, with `pfun = NULL` | | zero-inflated continuous | `check_zeroinfl_dist` | | inflated at 0 and/or 1 | `check_inflated_dist` | | AD-compatible cdf | also `check_ad_cdf(pfun, dfun, q, ...)`, with `.dq = FALSE` if discrete | | OSA residuals | also `check_osa_cdf(dfun, pfun, y, ...)`, with `.discrete = TRUE` if discrete | Choose `xs` so `pfoo(xs)` stays away from 0 and 1 — the `q(p(x))` round-trip cannot recover `x` once the cdf saturates in double precision. ## Numerical stability | Instead of | Write | |---|---| | `exp(x) - 1` | `expm1(x)` | | `log(1 + x)` | `log1p(x)` | | `1 - exp(-x)` | `-expm1(-x)` | | `x^2`, `x^3` | `x * x`, `x * x * x` | | `x^a`, non-integer `a` | `exp(a * log(x))` | | `log(pnorm(x))` | `pnorm(x, log.p = TRUE)`, accurate where `pnorm(x)` underflows | The one to watch for is `0 * (-Inf)`. That is `NaN`, and it takes out the whole gradient rather than the single element you were multiplying. It shows up when an indicator multiplies a log that is `-Inf` outside the support. Either put the indicator inside the log, ```r logdens + log(greater(x, 0)) # rather than greater(x, 0) * logdens ``` or clamp with `as.finite()`, or offset the argument with `log(x + .Machine$double.xmin)`. ## Branching without breaking the tape `if`/`else` and `ifelse()` on a taped value do not work: the branch is never recorded. Use the arithmetic indicators, which carry gradients. | Helper | Returns 1 when | |---|---| | `iszero(x)` | `x == 0` | | `isnonzero(x)` | `x != 0` | | `ispos(x)`, `ispos_strict(x)` | `x > 0` (note: **0 at `x = 0`**) | | `isneg(x)` | `x < 0` | | `greater(x, val)` | `x > val` | | `smaller(x, val)` | `x < val` | `pmax.ad()` and `pmin.ad()` are the tape-safe `pmax`/`pmin`. Outside a tape, `if (any(...)) stop(...)` is fine. That is what the `!ad_context()` guard is for. One more thing about `ifelse()`, even where it is legitimate: its result takes the length of the *test*. If the test is a parameter shorter than `x`, the answer is silently truncated. Recycle everything to a common length first. ```r n <- max(lengths(list(x, a, b))) x <- rep_len(x, n); a <- rep_len(a, n); b <- rep_len(b, n) ``` ## Zero-modified distributions `RTMBdist` ships four families built on a base distribution. The prefixes are `zi` (zero-inflated), `zt` (zero-truncated) and `h` (hurdle, also called zero-altered). Write \eqn{\pi_0} for the probability of a zero under the **base** distribution and \eqn{p_0} for the `zeroprob` parameter. | Family | At `x = 0` | At `x > 0` | Helper | |---|---|---|---| | zero-inflated, discrete | \eqn{p_0 + \pi_0} | \eqn{(1-p_0)\,f(x)} | `log_zi_discrete(x, logdens, zeroprob)` | | zero-inflated, continuous | \eqn{p_0} | \eqn{(1-p_0)\,f(x)} | `log_zi(x, logdens, zeroprob)` | | zero-truncated | 0 | \eqn{f(x)/(1-\pi_0)} | write it out, see below | | hurdle | \eqn{p_0} | \eqn{(1-p_0)\,f(x)/(1-\pi_0)} | `log_hurdle(x, logdens, p0, zeroprob)` | Zero-inflation and a hurdle only differ for discrete distributions. Inflation adds zeros on top of the ones the base already produces, so `zeroprob` is not the probability of a zero; a hurdle replaces them, so it is. For a continuous base the two are the same thing, because the continuous part puts no mass at zero. That is why `gamlss.dist` calls those zero-adjusted rather than zero-inflated, and why there is no `h` version of a continuous distribution here. Zero-truncated, using the Poisson as the example: ```r dztfoo <- function(x, theta, log = FALSE) { # ... guards and escapes as usual ... logdens <- dfoo(x, theta, log = TRUE) - log1p(-p0) + # p0 = P(X = 0) under the base distribution log(ispos_strict(x)) # -Inf below the support, so the density is 0 there if (log) return(logdens) return(exp(logdens)) } ``` Hurdle, which reuses the same `p0`: ```r dhfoo <- function(x, theta, zeroprob = 0.5, log = FALSE) { # ... guards and escapes as usual ... logdens <- log_hurdle(x, dfoo(x, theta, log = TRUE), p0, zeroprob) if (log) return(logdens) return(exp(logdens)) } ``` `log_hurdle()` combines the two branches with `logspace_add()`, so either can be exactly zero, `zeroprob = 0` and `zeroprob = 1` included, without producing `NaN`. The cdf is `greater(q, -1) * (zeroprob + (1 - zeroprob) * pztfoo(q, theta))`. The RNG draws a uniform, keeps the zeros, and fills the rest from `rztfoo`. ## Reparameterised variants Transform the parameters at the top and delegate. The `2` suffix means a reparameterisation, usually in terms of the mean. ```r dfoo2 <- function(x, mean, sd, log = FALSE) { if (!ad_context()) { ... } if (inherits(x, "simref")) return(dGenericSim("dfoo2", ...)) if (inherits(x, "osa")) return(dGenericOSA("dfoo2", ...)) shape <- mean * mean / (sd * sd) scale <- sd * sd / mean dfoo(x, shape = shape, scale = scale, log = log) } ``` Keep the `simulation_check` and the simref/osa escapes even though the base function has them. The escapes need to report `dfoo2`, not `dfoo`. ## Distributions to read The example above is deliberately simple. When you hit something it does not cover, the shortest route is usually to open a distribution that already solves the same problem. | If you are writing | Read | |---|---| | a plain continuous distribution, closed forms throughout | `R/gompertz.R`, `R/llogis.R` | | a discrete distribution | `R/ztpois.R`, `R/zipois.R` | | a reparameterisation of an existing one | `R/nbinom2.R`, `R/beta2.R` | | zero-inflated, zero-truncated or hurdle | `R/zipois.R`, `R/ztpois.R`, `R/hpois.R` | | something with an awkward support boundary | `R/pareto.R` (\eqn{x > 1}), `R/zibeta.R` (excluding \eqn{x = 1}) | | a term that has to be clamped to stay finite | `R/kumar.R` | | a quantile function that needs root finding | `R/exgauss.R` | | a density that needs branching on a parameter | `R/bccg.R`, `R/bcpe.R` | | an AD version of a `stats` function | `R/ad-dispatch.R`, `R/geom-ad.R` | ## Before you are done - `devtools::document()`, so `NAMESPACE` and the `man/` pages are regenerated. - `devtools::test()` for the new file, then the whole suite. A new distribution that breaks an old test usually means a shared helper changed behaviour. - `devtools::check()`. Watch in particular for `checking Rd cross-references`, which catches a `@seealso` pointing at a topic that does not exist. - Check the `NAMESPACE` diff. If it contains anything beyond your own exports, something imported more than you meant it to. ## When the cdf has no closed form A distribution function is needed for OSA residuals, so avoid shipping without one. Two internal helpers build it from the density. `numerical_cdf(dfun, q, pars, centre, scale, lower, upper, ...)` integrates a continuous density with `RTMB`'s AD-compatible `integrate()`. Call it only in AD context and keep a faster non-AD cdf, e.g. from another package, for ordinary use, as `pskewnorm` does. If that cdf is inaccurate in the tails, pass its values through `refine_tails()`, which integrates only where the requested tail is small. Pass `centre` close to the mode and `scale` of the order of the spread; the comments in `R/numerical_cdf.R` explain why. `discrete_cdf(dfun, q, pars, upper, ...)` sums a probability mass function over \eqn{0, \ldots, q}. It is AD-compatible in the parameters, while `q` must be numeric, which is how `RTMB`'s CDF-based OSA residuals use it. Pass `upper` (e.g. `size`) for a bounded support, so that the cdf is exactly 1 at the upper end. Otherwise rounding can push it above 1, and the OSA residuals become `NaN`. `pbetabinom` is a short example. Both take the parameters as a named list rather than `...`, so that a parameter called `df` cannot partially match the argument `dfun`. ## When `RTMB` has no AD version {#when-rtmb-has-no-ad-version} `RTMB` covers most of `stats` but not all of it. `dgeom`, `pt` and `plnorm` have no AD method. Exporting your own version of one would mask the `stats` function for everyone who loads `RTMBdist` and break their ordinary calls. `R/ad-dispatch.R` avoids that. Write the AD implementation as `.ad`, export it, and register an unexported S4 generic: ```r setGeneric("pt", signature = c("q", "df")) setMethod("pt", signature(q = "num", df = "num."), stats::pt) # plain numbers setMethod("pt", signature(q = "ad", df = "ad." ), # any AD variable function(q, df, ncp, lower.tail = TRUE, log.p = FALSE) { ... pt.ad(q, df) ... }) ``` `num` and `ad` are `RTMB`'s dispatch classes: `ad` covers AD variables *and* numbers, `num` covers only numbers, so the more specific `num` method wins whenever nothing is being taped. A trailing dot (`num.`, `ad.`) also allows the argument to be missing. Since the generic is not exported, a user's `stats::pt()` is untouched, while inside the package a bare `pt()` dispatches correctly. It costs one S4 dispatch, under a microsecond, paid once when the tape is built rather than on every gradient evaluation. One import to avoid: never write `@importFrom stats pnorm` to get hold of a `pnorm`. That replaces `RTMB`'s AD-aware `pnorm` across the entire package and silently breaks taping in every file. Call `stats::pnorm()` at the point of use instead.