R/sampling.R

In rbiogeme: Interface to 'Biogeme'

#' Create a sampling partition for alternatives
#'
#' A sampling partition is a collection of disjoint alternative-ID subsets and
#' the number of alternatives to draw from each subset. The actual random
#' sampling is performed by native Biogeme; this object only records its
#' declarative protocol.
#'
#' @param segments Non-empty list of integer alternative-ID vectors.
#' @param sample_sizes Positive integer sample size for each segment.
#' @param full_set Optional integer vector containing the complete partition
#'   universe. If omitted, it is the union of `segments`.
#' @return A `biogeme_sampling_partition` object.
#' @export
biogeme_sampling_partition <- function(segments, sample_sizes, full_set = NULL) {
  if (!is.list(segments) || length(segments) == 0L) {
    stop("segments must be a non-empty list.", call. = FALSE)
  }
  validate_ids <- function(ids, argument) {
    if (!is.numeric(ids) || length(ids) == 0L || anyNA(ids) ||
        any(!is.finite(ids)) || any(ids != floor(ids)) || anyDuplicated(ids)) {
      stop(argument, " must contain unique finite integer IDs.", call. = FALSE)
    }
    as.integer(ids)
  }
  segments <- lapply(seq_along(segments), function(index) {
    validate_ids(segments[[index]], paste0("segments[[", index, "]"))
  })
  segment_ids <- unlist(segments, use.names = FALSE)
  if (anyDuplicated(segment_ids)) {
    stop("segments must be disjoint.", call. = FALSE)
  }
  if (!is.numeric(sample_sizes) || length(sample_sizes) != length(segments) ||
      anyNA(sample_sizes) || any(!is.finite(sample_sizes)) ||
      any(sample_sizes != floor(sample_sizes)) || any(sample_sizes < 1)) {
    stop("sample_sizes must contain one positive integer per segment.", call. = FALSE)
  }
  sample_sizes <- as.integer(sample_sizes)
  if (is.null(full_set)) {
    full_set <- sort(segment_ids)
  } else {
    full_set <- validate_ids(full_set, "full_set")
  }
  if (!setequal(unlist(segments, use.names = FALSE), full_set)) {
    stop("segments must form a partition of full_set.", call. = FALSE)
  }
  for (index in seq_along(segments)) {
    if (sample_sizes[[index]] > length(segments[[index]])) {
      stop("A sample size cannot exceed its segment size.", call. = FALSE)
    }
  }
  structure(
    list(segments = segments, sample_sizes = sample_sizes, full_set = full_set),
    class = "biogeme_sampling_partition"
  )
}

#' Generate nearly equal sampling segment sizes
#'
#' This matches native `generate_segment_size()`: the remainder is assigned
#' to the first segments. It only creates protocol metadata; native Biogeme
#' still performs the sampling and likelihood calculations.
#'
#' @param sample_size Total number of alternatives to draw.
#' @param number_of_segments Number of partition segments.
#' @return A positive integer vector of segment sizes.
#' @export
sampling_segment_sizes <- function(sample_size, number_of_segments) {
  if (!is.numeric(sample_size) || length(sample_size) != 1L ||
      is.na(sample_size) || !is.finite(sample_size) || sample_size < 1 ||
      sample_size != floor(sample_size)) {
    stop("sample_size must be one positive integer.", call. = FALSE)
  }
  if (!is.numeric(number_of_segments) || length(number_of_segments) != 1L ||
      is.na(number_of_segments) || !is.finite(number_of_segments) ||
      number_of_segments < 1 || number_of_segments != floor(number_of_segments)) {
    stop("number_of_segments must be one positive integer.", call. = FALSE)
  }
  number_of_segments <- as.integer(number_of_segments)
  sample_size <- as.integer(sample_size)
  base <- sample_size %/% number_of_segments
  remainder <- sample_size %% number_of_segments
  sizes <- rep.int(base, number_of_segments)
  if (remainder > 0L) sizes[seq_len(remainder)] <- sizes[seq_len(remainder)] + 1L
  as.integer(sizes)
}

#' Define a cross-variable for sampled-alternative utilities
#'
#' A cross-variable combines individual attributes with alternative
#' attributes. Native Biogeme expands its formula separately for every sampled
#' alternative after the choice set has been generated.
#'
#' @param name Name of the generated cross-variable.
#' @param formula A neutral Biogeme expression involving individual and
#'   alternative variables.
#' @return A `biogeme_cross_variable` specification.
#' @export
cross_variable <- function(name, formula) {
  name <- validate_name(name, "name")
  structure(
    list(name = name, formula = as_biogeme_expression(formula)),
    class = "biogeme_cross_variable"
  )
}

collect_sampling_nest_betas <- function(nests, collected = list()) {
  if (is.null(nests)) return(collected)
  if (inherits(nests, "biogeme_nested_nests")) {
    for (nest in nests$nests) {
      collected <- collect_biogeme_betas(nest$nest_parameter, collected)
    }
    return(collected)
  }
  if (inherits(nests, "biogeme_cross_nested_nests")) {
    for (nest in nests$nests) {
      collected <- collect_biogeme_betas(nest$nest_parameter, collected)
      for (allocation in nest$allocation) {
        collected <- collect_biogeme_betas(allocation, collected)
      }
    }
    return(collected)
  }
  stop("nests must be a native nested or cross-nested nest specification.", call. = FALSE)
}

#' Define a sampled-alternative estimation model
#'
#' The model stores ordinary R data and neutral expressions. During
#' [estimate_sampled_alternatives()], the bridge constructs native Biogeme
#' `SamplingContext`, `ChoiceSetsGeneration`, and `GenerateModel` objects.
#' Alternative sampling, sampled likelihood construction, differentiation, and
#' optimization therefore remain entirely in native Biogeme.
#'
#' @param alternatives Numeric data frame with one row per alternative.
#' @param individuals Numeric data frame with one row per decision maker.
#' @param choice_column Choice column in `individuals` containing alternative IDs.
#' @param id_column Unique alternative-ID column in `alternatives`.
#' @param utility Generic utility expression before sampled-alternative suffixes
#'   are added.
#' @param partition A [biogeme_sampling_partition()] for the main sample.
#' @param biogeme_file_name Native sampled-data filename. It is always
#'   regenerated by default; it is not silently recycled.
#' @param model_type One of `"logit"`, `"nested"`, or `"cnl"`.
#' @param combined_variables Optional list of [cross_variable()] objects.
#' @param mev_partition Optional MEV-term sampling partition.
#' @param mev_sample_sizes Optional positive integer MEV sample sizes.
#' @param nests Nested or cross-nested specification for the corresponding
#'   model type.
#' @param control Optional [biogeme_control()] object.
#' @return A `biogeme_sampled_alternatives_model`.
#' @export
sampled_alternatives_model <- function(
    alternatives,
    individuals,
    choice_column,
    id_column,
    utility,
    partition,
    biogeme_file_name,
    model_type = c("logit", "nested", "cnl"),
    combined_variables = list(),
    mev_partition = NULL,
    mev_sample_sizes = NULL,
    nests = NULL,
    control = NULL
) {
  alternatives <- validate_biogeme_data(alternatives)
  individuals <- validate_biogeme_data(individuals)
  if (!is.character(choice_column) || length(choice_column) != 1L ||
      is.na(choice_column) || !nzchar(choice_column) ||
      !choice_column %in% names(individuals)) {
    stop("choice_column must name a column in individuals.", call. = FALSE)
  }
  if (!is.character(id_column) || length(id_column) != 1L || is.na(id_column) ||
      !nzchar(id_column) || !id_column %in% names(alternatives)) {
    stop("id_column must name a column in alternatives.", call. = FALSE)
  }
  if (anyNA(alternatives[[id_column]]) ||
      any(alternatives[[id_column]] != floor(alternatives[[id_column]])) ||
      anyDuplicated(alternatives[[id_column]])) {
    stop("The alternative ID column must contain unique integer IDs.", call. = FALSE)
  }
  if (anyNA(individuals[[choice_column]]) ||
      any(individuals[[choice_column]] != floor(individuals[[choice_column]]))) {
    stop("The choice column must contain finite integer alternative IDs.", call. = FALSE)
  }
  if (!inherits(partition, "biogeme_sampling_partition")) {
    stop("partition must be a biogeme_sampling_partition object.", call. = FALSE)
  }
  available_ids <- as.integer(alternatives[[id_column]])
  if (!all(partition$full_set %in% available_ids)) {
    stop("partition contains alternative IDs absent from alternatives.", call. = FALSE)
  }
  if (any(!individuals[[choice_column]] %in% available_ids)) {
    stop("individual choices must be present in alternatives.", call. = FALSE)
  }
  utility <- as_biogeme_expression(utility)
  if (!is.list(combined_variables) || any(!vapply(
      combined_variables,
      inherits,
      logical(1),
      what = "biogeme_cross_variable"
  ))) {
    stop("combined_variables must contain cross_variable() objects.", call. = FALSE)
  }
  model_type <- match.arg(model_type)
  if (!is.character(biogeme_file_name) || length(biogeme_file_name) != 1L ||
      is.na(biogeme_file_name) || !nzchar(biogeme_file_name)) {
    stop("biogeme_file_name must be one non-empty path.", call. = FALSE)
  }
  biogeme_file_name <- normalizePath(
    path.expand(biogeme_file_name),
    mustWork = FALSE
  )
  if (!is.null(mev_partition) &&
      !inherits(mev_partition, "biogeme_sampling_partition")) {
    stop("mev_partition must be a biogeme_sampling_partition object.", call. = FALSE)
  }
  if (!is.null(mev_partition)) {
    if (is.null(mev_sample_sizes)) mev_sample_sizes <- mev_partition$sample_sizes
    if (!is.numeric(mev_sample_sizes) ||
        length(mev_sample_sizes) != length(mev_partition$segments) ||
        anyNA(mev_sample_sizes) || any(!is.finite(mev_sample_sizes)) ||
        any(mev_sample_sizes != floor(mev_sample_sizes)) || any(mev_sample_sizes < 1)) {
      stop("mev_sample_sizes must contain one positive integer per MEV segment.", call. = FALSE)
    }
    mev_sample_sizes <- as.integer(mev_sample_sizes)
  } else if (!is.null(mev_sample_sizes)) {
    stop("mev_sample_sizes requires mev_partition.", call. = FALSE)
  }
  if (model_type == "nested" && !inherits(nests, "biogeme_nested_nests")) {
    stop("nested models require nested_nests().", call. = FALSE)
  }
  if (model_type == "cnl" && !inherits(nests, "biogeme_cross_nested_nests")) {
    stop("CNL models require cross_nested_nests().", call. = FALSE)
  }
  if (model_type == "logit" && !is.null(nests)) {
    stop("nests must be NULL for logit models.", call. = FALSE)
  }
  if (!is.null(control) && !is.list(control)) {
    stop("control must be a list returned by biogeme_control().", call. = FALSE)
  }

  parameters <- collect_biogeme_betas(utility, list())
  for (item in combined_variables) {
    parameters <- collect_biogeme_betas(item$formula, parameters)
  }
  parameters <- collect_sampling_nest_betas(nests, parameters)
  structure(
    list(
      kind = "sampled_alternatives",
      alternatives = alternatives,
      individuals = individuals,
      choice_column = choice_column,
      id_column = id_column,
      utility = utility,
      partition = partition,
      biogeme_file_name = biogeme_file_name,
      model_type = model_type,
      combined_variables = combined_variables,
      mev_partition = mev_partition,
      mev_sample_sizes = mev_sample_sizes,
      nests = nests,
      control = control,
      parameters = parameters
    ),
    class = c("biogeme_sampled_alternatives_model", "biogeme_model")
  )
}

#' @export
print.biogeme_sampled_alternatives_model <- function(x, ...) {
  cat(
    "Biogeme sampled-alternative ", x$model_type, " model with ",
    length(x$parameters), " parameter definition(s)\n", sep = ""
  )
  invisible(x)
}

#' Estimate a sampled-alternative model with native Biogeme
#'
#' The native sampling context always regenerates the sampled data file with
#' recycling disabled. Native Biogeme then constructs the sampled logit,
#' nested-logit, or CNL likelihood and performs estimation. No R expression or
#' callback is evaluated inside the sampler, likelihood, derivative, or
#' optimizer.
#'
#' @param model A [sampled_alternatives_model()] object.
#' @param model_name Native Biogeme model name.
#' @param controls Named native Biogeme controls.
#' @param starting_values Optional named finite numeric vector.
#' @param run_bootstrap Whether to run native bootstrap re-estimation.
#' @param control Optional [biogeme_control()] object; an alias for `controls`.
#' @return An object of class `biogeme_fit` with sampling metadata attached.
#' @export
estimate_sampled_alternatives <- function(
    model,
    model_name = "rbiogeme_sampled",
    controls = list(),
    starting_values = NULL,
    run_bootstrap = FALSE,
    control = NULL
) {
  if (!inherits(model, "biogeme_sampled_alternatives_model")) {
    stop("model must be a sampled_alternatives_model object.", call. = FALSE)
  }
  if (!is.null(control)) {
    if (!is.null(controls) && length(controls) > 0L) {
      stop("Provide either controls or control, not both.", call. = FALSE)
    }
    controls <- control
  } else if (length(controls) == 0L && !is.null(model$control)) {
    controls <- model$control
  }
  controls <- if (is.null(controls)) list() else controls
  controls <- biogeme_prepare_output_controls(controls)
  if (identical(model_name, "rbiogeme_sampled") && !is.null(controls$model_name)) {
    model_name <- controls$model_name
  }
  model_name <- validate_name(model_name, "model_name")
  starting_values <- validate_starting_values(starting_values)
  if (!is.logical(run_bootstrap) || length(run_bootstrap) != 1L || is.na(run_bootstrap)) {
    stop("run_bootstrap must be one non-missing logical value.", call. = FALSE)
  }
  raw <- biogeme_estimate_sampled_alternatives_model(
    model = model,
    model_name = model_name,
    controls = controls,
    starting_values = starting_values,
    run_bootstrap = run_bootstrap
  )
  if (!is.list(raw) || is.null(raw$result)) {
    biogeme_abort(
      "Biogeme returned an invalid sampled-alternative result.",
      class = "biogeme_result_error",
      operation = "sampled-alternative result conversion",
      suggestion = "enable debug mode and report the complete result structure"
    )
  }
  fit <- as_biogeme_fit(raw$result, model)
  fit$sampling_context <- raw$sampling_context
  fit$sampled_data_file <- as.character(raw$sampled_data_file)
  fit$main_sample_size <- as.integer(raw$main_sample_size)
  fit$mev_sample_size <- as.integer(raw$mev_sample_size)
  fit$sampled_database_rows <- as.integer(raw$sampled_database_rows)
  fit
}