R/asymmetric_laplace.R

In BayesForge: Bayesian Inference using 'numpyro' and 'XLA'

#' @title Asymmetric Laplace distribution
#' @description
#' Samples from an Asymmetric Laplace distribution.

#' The Asymmetric Laplace distribution is a generalization of the Laplace distribution,
#' where the two sides of the distribution are scaled differently. It is defined by
#' a location parameter (loc), a scale parameter (scale), and an asymmetry parameter (asymmetry).
#'
#' \deqn{
#'  f(x, \kappa) = \frac{1}{\kappa+\kappa^{-1}}\exp(-x\kappa),\quad x\ge0
#' }
#' \deqn{
#'  = \frac{1}{\kappa+\kappa^{-1}}\exp(x/\kappa),\quad x<0
#' }
#' \deqn{
#'  \textrm{for } -\infty < x < \infty, \kappa > 0 .
#' }
#'
#' laplace_asymmetric takes 'kappa' as a shape parameter for \eqn{\kappa}. For \eqn{\kappa = 1}, it is identical to a Laplace distribution.
#'
#'
#' @param loc A numeric vector or single numeric value representing the location parameter of the distribution. This corresponds to \eqn{\mu}.
#' @param scale A numeric vector or single numeric value representing the scale parameter of the distribution. This corresponds to \eqn{\sigma}.
#' @param asymmetry A numeric vector or single numeric value representing the asymmetry parameter of the distribution. This corresponds to \eqn{\kappa}.
#' @param shape A numeric vector specifying the shape of the output.  This is used to set the batch shape when \code{sample=FALSE} (model building) or as `sample_shape` to draw a raw JAX array when \code{sample=TRUE} (direct sampling).
#' @param event Integer specifying the number of batch dimensions to reinterpret as event dimensions (used in model building).
#' @param mask A logical vector indicating which observations to mask.
#' @param create_obj Logical; If `TRUE`, returns the raw BI distribution object instead of creating a sample site.
#' @param validate_args Logical: Whether to validate parameter values.  Defaults to `reticulate::py_none()`.
#' @param sample A logical value that controls the function's behavior. If `TRUE`,
#'   the function will directly draw samples from the distribution. If `FALSE`,
#'   it will create a random variable within a model. Defaults to `FALSE`.
#' @param seed An integer used to set the random seed for reproducibility when
#'   `sample = TRUE`. This argument has no effect when `sample = FALSE`, as
#'   randomness is handled by the model's inference engine. Defaults to 0.
#' @param obs A numeric vector or array of observed values. If provided, the
#'   random variable is conditioned on these values. If `NULL`, the variable is
#'   treated as a latent (unobserved) variable. Defaults to `NULL`.
#' @param name A character string representing the name of the random variable
#'   within a model. This is used to uniquely identify the variable. Defaults to 'x'.
#' @param to_jax Boolean. Indicates whether to return a JAX array or not.
#'
#' @return When \code{sample=FALSE}: A BI AsymmetricLaplace distribution object (for model building).
#'         When \code{sample=TRUE}: A JAX array of samples drawn from the AsymmetricLaplace distribution (for direct sampling).
#'         When \code{create_obj=TRUE}: The raw BI distribution object (for advanced use cases).
#' @examples
#' \donttest{
#' library(BayesForge)
#' m <- importBF(platform = "cpu")
#' bf.dist.asymmetric_laplace(sample = TRUE)
#' }
#' @seealso This is a wrapper of  \url{https://num.pyro.ai/en/stable/distributions.html#asymmetriclaplace}
#' @export
bf.dist.asymmetric_laplace <- function(
    loc = 0.0,
    scale = 1.0,
    asymmetry = 1.0,
    validate_args = py_none(),
    name = "x", obs = py_none(),
    mask = py_none(),
    sample = FALSE,
    seed = py_none(),
    shape = c(),
    event = 0,
    create_obj = FALSE, to_jax = TRUE) {
  shape <- do.call(tuple, as.list(as.integer(shape)))
  event <- as.integer(event)
  if (!.BF_env$.py$is_none(seed)) {
    seed <- as.integer(seed)
  }
  .BF_env$.bf_instance$dist$asymmetric_laplace(
    loc = .BF_env$jnp$array(loc),
    scale = .BF_env$jnp$array(scale),
    asymmetry = .BF_env$jnp$array(asymmetry), validate_args = validate_args, name = name, obs = obs, mask = mask, sample = sample, seed = seed, shape = shape, event = event, create_obj = create_obj, to_jax = to_jax
  )
}