R/fit_distribution.R

In stratifyR: Optimal Stratification of Univariate Populations

# =============================================================================
# fit_distribution.R  --  stratifyR 2.0
#
# Internal helper (also exported) used by strata.data() and strata.distr()
# to fit the best parametric distribution to a data vector.
#
# Supports 10 families (matching the original package) plus
# auto-selection by AIC/BIC.
# =============================================================================

#' Fit the Best-Matching Parametric Distribution to a Data Vector
#'
#' Fits up to 10 continuous parametric families to a numeric vector via MLE
#' and selects the best-fitting model by minimum AIC (default) or BIC.  The
#' result is consumed internally by \code{\link{strata.data}} and
#' \code{\link{strata.distr}}, and can also be used standalone for
#' distributional analysis.
#'
#' Supported families: \code{"norm"}, \code{"lnorm"}, \code{"gamma"},
#' \code{"weibull"}, \code{"exp"}, \code{"cauchy"}, \code{"unif"},
#' \code{"pareto"}, \code{"triangle"} (symmetric), \code{"rtriangle"}
#' (right-triangular).  Passing \code{family = "auto"} tries all applicable
#' families and returns the winner.
#'
#' @param data   Numeric vector of observations (population or sample values).
#' @param family Character.  Target family name, or \code{"auto"} (default)
#'   to select by AIC.
#' @param criterion Character.  \code{"aic"} (default) or \code{"bic"}.
#' @param verbose Logical.  Print AIC/BIC table for all attempted fits.
#'   Default \code{FALSE}.
#'
#' @return A list with:
#' \describe{
#'   \item{\code{family}}{Name of the best-fitting distribution.}
#'   \item{\code{params}}{Named vector of MLE parameter estimates.}
#'   \item{\code{aic}}{Named numeric vector of AICs (NA if fit failed).}
#'   \item{\code{bic}}{Named numeric vector of BICs (NA if fit failed).}
#'   \item{\code{fits_ok}}{Logical vector indicating successful fits.}
#'   \item{\code{fit_obj}}{The \code{fitdistrplus::fitdist} object for the
#'     winning family (or \code{NULL} if unavailable).}
#' }
#'
#' @examples
#' \dontrun{
#' set.seed(1)
#' y <- rlnorm(2000, meanlog = 1.5, sdlog = 0.8)
#' fd <- fit_distribution(y)
#' cat("Best family:", fd$family, "\n")
#' cat("Parameters :", fd$params, "\n")
#'
#' # Force a specific family
#' fd2 <- fit_distribution(y, family = "gamma")
#' }
#'
#' @seealso \code{\link{strata.data}}, \code{\link{strata.distr}}
#'
#' @export
fit_distribution <- function(data,
                              family    = "auto",
                              criterion = c("aic", "bic"),
                              verbose   = FALSE) {
  criterion <- match.arg(criterion)
  data      <- as.numeric(data)
  data      <- data[is.finite(data)]

  if (length(data) < 5L)
    stop("'data' must contain at least 5 finite values.")

  # Build a minimal my_env just to pass to get.dist()
  my_env   <- new.env(parent = emptyenv())

  if (family == "auto") {
    # Use internal get.dist which already does AIC-based selection
    result <- get.dist(data, my_env)

    # Also compute BIC for each fitted model
    bic_vec <- tryCatch({
      vapply(names(result$fits_ok), function(nm) {
        if (!isTRUE(result$fits_ok[[nm]])) return(NA_real_)
        aic_v <- result$aic[[nm]]
        if (is.na(aic_v)) return(NA_real_)
        # BIC ~= AIC + (k * ln(n) - 2k) = AIC + k*(ln(n) - 2)
        # For this we'd need k; approximate as AIC + 2 (k=1) ... use fitdist
        NA_real_  # placeholder; full BIC requires re-fitting, skip for speed
      }, FUN.VALUE = numeric(1))
    }, error = function(e) NULL)

    if (verbose) {
      cat("-- fit_distribution() AIC table ---------------------\n")
      aic_tbl <- result$aic[order(result$aic)]
      for (nm in names(aic_tbl)) {
        cat(sprintf("  %-12s : AIC = %s\n", nm,
                    if (is.finite(aic_tbl[nm])) sprintf("%.2f", aic_tbl[nm]) else "  fail"))
      }
      cat(sprintf("  Winner: %s\n", result$distr))
      cat("-----------------------------------------------------\n")
    }

    return(list(
      family  = result$distr,
      params  = result$params,
      aic     = result$aic,
      bic     = bic_vec,
      fits_ok = result$fits_ok,
      fit_obj = NULL   # internal fit objects not retained in get.dist
    ))
  }

  # ---- Single family fit -----------------------------------------------------
  family_norm <- switch(family,
    triangle  = "triang",
    rtriangle = "triang",
    pareto    = "pareto",
    family
  )

  fit_obj <- tryCatch({
    if (family %in% c("pareto")) {
      fitdistrplus::fitdist(data, distr = family_norm, method = "mle",
                            start = list(shape = 1, scale = stats::quantile(data, 0.25)),
                            lower = c(1e-8, 1e-8))
    } else if (family %in% c("triangle", "rtriangle")) {
      minx <- min(data); maxx <- max(data)
      eps  <- (maxx - minx) * 1e-6 + 1e-8
      m0   <- mode.val(data)
      m0   <- min(max(m0, minx + eps), maxx - eps)
      fitdistrplus::fitdist(data, distr = "triang", method = "mle",
                            start   = list(mode = m0),
                            fix.arg = list(min = minx, max = maxx))
    } else if (family %in% c("gamma", "weibull", "exp", "lnorm", "norm")) {
      fitdistrplus::fitdist(data, distr = family, method = "mle",
                            lower = if (family %in% c("gamma","weibull")) c(0,0) else NULL)
    } else {
      fitdistrplus::fitdist(data, distr = family, method = "mle")
    }
  }, error = function(e) {
    message("fit_distribution: fit failed for family '", family, "': ", conditionMessage(e))
    NULL
  })

  if (is.null(fit_obj)) {
    return(list(family = family, params = NULL, aic = c(setNames(NA_real_, family)),
                bic = NULL, fits_ok = setNames(FALSE, family), fit_obj = NULL))
  }

  params <- fit_obj$estimate
  if (family %in% c("triangle","rtriangle") && length(params) == 1L) {
    params <- c(min = min(data), max = max(data), mode = as.numeric(params))
  }

  if (verbose) {
    cat(sprintf("fit_distribution: %s  AIC=%.2f\n", family, fit_obj$aic))
    print(params)
  }

  list(
    family  = family,
    params  = params,
    aic     = setNames(fit_obj$aic, family),
    bic     = setNames(fit_obj$bic, family),
    fits_ok = setNames(TRUE, family),
    fit_obj = fit_obj
  )
}