R/run.R

In IxPopDyMod: Framework for Tick Population and Infection Modeling

#' Run the model
#'
#' @param cfg An `IxPopDyMod::config` object
#' @param progress Boolean indicating whether to log progress every 100 steps
#'
#' @returns Data frame of population of ticks of each life stage each day
#'
#' @examples
#' run(config_ex_1)
#'
#' @export
run <- function(cfg, progress = TRUE) {

  # 00 get valid life stages
  life_stages <- life_stages(cfg$cycle)

  # 01 Precompute predictors and attach to cfg$preds
  if (!is.null(cfg$preds)) {
    attr(cfg$preds, "precomputed") <- precompute_predictors(cfg$preds)
  }

  total_days <- cfg$steps + cfg$max_duration

  # 02 initialize a 2D delay matrix instead of a 3D array
  delay_mat <- matrix(
    0,
    nrow = length(life_stages),
    ncol = total_days,
    dimnames = list(life_stages, NULL)
  )

  # 03 initialize population matrices (padded to total_days to handle future delays safely)
  population <- empty_population_matrix(life_stages = life_stages, steps = total_days)
  population <- set_initial_population(
    population = population, initial_population = cfg$initial_population
  )

  developing_population <- empty_population_matrix(life_stages = life_stages, steps = total_days)

  # at each time step
  for (time in 1:(cfg$steps - 1)) {

    if (progress && time %% 100 == 0) message("Day: ", time)

    # 1. calculate transition probabilities
    trans_matrix <- gen_transition_matrix(
      time = time,
      population = population,
      developing_population = developing_population,
      tick_transitions = cfg$cycle,
      predictors = cfg$preds
    )

    # 2. calculate ticks entering delayed development and update matrices directly
    delays <- update_delays(
      time = time,
      delay_mat = delay_mat,
      population = population,
      developing_population = developing_population,
      tick_transitions = cfg$cycle,
      max_duration = cfg$max_duration,
      predictors = cfg$preds
    )
    delay_mat <- delays$delay_mat
    developing_population <- delays$developing_population

    # 3. calculate the number of ticks at the next time step
    population[, time + 1] <-
      population[, time] %*% trans_matrix + delay_mat[, time + 1]
  }

  # Truncate padded matrices back to exactly cfg$steps before returning
  population <- population[, 1:cfg$steps]
  developing_population <- developing_population[, 1:cfg$steps]

  # Return the total population of ticks each day
  population_matrix_to_output_df(population + developing_population)
}

population_matrix_to_output_df <- function(matrix) {
  df <- as.data.frame(t(matrix))
  df[["day"]] <- seq_len(nrow(df))
  life_stages <- rownames(matrix)
  df <- stats::reshape(
    df,
    direction = "long",
    varying = life_stages,
    v.names = "pop",
    idvar = "day",
    timevar = "stage",
    times = life_stages
  )
  df <- df[order(df[["day"]]), ]
  rownames(df) <- seq_len(nrow(df))
  attr(df, "reshapeLong") <- NULL # nolint: object_name_linter
  df
}

#' Get a predictor from input data
#'
#' @param time Numeric vector of days to get data. Ignored if input is constant
#'   over time (as indicated by NA value in 'j_day' column)
#' @param table input predictors table
#' @param pred string specifying the name of the predictor, e.g. "host_den"
#'
#' @returns a numeric vector of predictor values
get_pred_from_table <- function(time, pred, table) {
  # Check if we have precomputed the fast-lookup list
  precomputed <- attr(table, "precomputed")

  if (is.null(precomputed)) {
    rows <- (is.na(table$j_day) | (table$j_day %in% time)) & (table$pred == pred)
    subset <- table[rows, ]
    pred_values <- subset$value
    pred_names <- subset$pred_subcategory
    if (!all(is.na(pred_names))) {
      names(pred_values) <- pred_names
    }
    return(pred_values)
  }

  p_data <- precomputed[[pred]]

  # Edge case: predictor requested that is not in the table
  if (is.null(p_data)) return(numeric(0))

  if (p_data$is_constant) {
    p_data$data
  } else {

    # Handle single day
    if (length(time) == 1) {
      # Prevent out-of-bounds errors and replicate numeric(0) for missing days
      if (time < 1 || time > length(p_data$data)) return(numeric(0))

      val <- p_data$data[[time]]
      if (is.null(val)) return(numeric(0))
      val

    } else {
      # Handle multiple days
      valid_time <- time[time >= 1 & time <= length(p_data$data)]
      res <- p_data$data[valid_time]

      # Drop missing days to exactly match the old data frame subsetting behavior
      res <- res[!vapply(res, is.null, logical(1))]
      if (length(res) == 0) return(numeric(0))

      vec <- unlist(res, use.names = FALSE)

      # Extract names day-by-day to ensure exact matching, even if days were missing
      if (!is.null(names(res[[1]]))) {
        names(vec) <- unlist(lapply(res, names), use.names = FALSE)
      }
      vec
    }
  }
}

#' Get tick density for specified time and life stages
#'
#' @param time Numeric vector of length one indicating day to get tick density
#' @param pred Character vector indicating life stages to get tick density of
#' @param population Matrix of number of ticks not currently undergoing
#'   development per life stage per day
#' @param developing_population Matrix of number of currently developing ticks
#'   per life stage per day
#' @returns Numeric vector of length one indicating current number of ticks in
#'   given life stages
#' @noRd
get_tick_den <- function(time, pred, population, developing_population) {
  stopifnot(length(time) == 1)
  life_stages <- rownames(population)
  life_stages_to_sum <- grep(pred, life_stages)
  total_population <- population + developing_population
  sum(total_population[life_stages_to_sum, time])
}

#' Get the value of a predictor
#'
#' @param time First day to get predictor value(s)
#' @param pred A `predictor_spec`
#' @param is_delay Boolean indicating whether the predictor is for a transition
#'   involving a delay
#' @param population Tick population matrix
#' @param developing_population Matrix of currently developing ticks
#' @param max_duration Numeric vector of length one. Determines the maximum
#' number of days that a duration-based transition can last.
#' @param predictors Table of predictor values
#' @noRd
#'
#' @returns a vector of a predictor at time time. The vector's length is based
#' on whether the transition is_delay.
get_pred <- function(
    time, pred, is_delay, population, developing_population, max_duration, predictors
) {

  life_stages <- rownames(population)

  if (pred$pred %in% valid_predictors_from_table(predictors)) {
    if (is_delay && !pred$first_day_only) {
      time <- time:(time + max_duration)
    }
    get_pred_from_table(time, pred$pred, predictors)
  } else if (any(grepl(pred$pred, life_stages))) {
    get_tick_den(time, pred$pred, population, developing_population)
  } else {
    # Validation should prevent hitting this case
    stop("Failed to match predictor: \"", pred$pred, "\"", call. = FALSE)
  }
}



#' Get the value determining probability or duration of a transition
#'
#' This generic function pulls out the functional form and parameters needed to
#' evaluate the transition function, then evaluates it using supplied predictors
#' like temperature or host density.
#'
#' @param time First day to get predictor value(s)
#' @param transition A `transition`
#' @param population Tick population matrix
#' @param developing_population Matrix of currently developing ticks
#' @param max_duration Numeric vector of length one. Determines the maximum
#' number of days that a duration-based transition can last.
#' @param predictors Table of predictor values
#' @noRd
#'
#' @returns Numeric vector indicating probability or duration of a transition.
get_transition_value <- function(
    time, transition, predictors, max_duration, population, developing_population
) {

  inputs <- get_transition_inputs_unevaluated(
    time = time,
    transition = transition,
    predictors = predictors,
    max_duration = max_duration,
    population = population,
    developing_population = developing_population
  )

  value <- do.call(inputs[["function"]], c(inputs[["parameters"]], inputs[["predictors"]]))

  validate_transition_value(
    transition = transition, value = value, max_duration = max_duration
  )

  value
}

get_transition_inputs_unevaluated <- function(
    time, transition, predictors, max_duration, population, developing_population
) {
  f <- transition$fun

  # a list of parameter values, each of which could be a scalar or named numeric vector
  params <- transition$parameters

  # a list of predictor values, each of which could be a scalar or named numeric vector
  predictor_values <- lapply(
    transition$predictors,
    function(pred) {
      get_pred(
        time = time,
        pred = pred,
        is_delay = transition$transition_type == "duration",
        population = population,
        developing_population = developing_population,
        max_duration = max_duration,
        predictors = predictors
      )
    }
  )

  list("function" = f, "parameters" = params, "predictors" = predictor_values)
}

#' Raise error if returned value from a transition is invalid
#'
#' @param transition A transition
#' @param value The evaluated value
#' @param max_duration max_duration parameter of a config. Used for validating
#'   duration-based transitions.
#'
#' @returns nothing, called for side effects
#' @noRd
validate_transition_value <- function(transition, value, max_duration) {
  if (!is.numeric(value)) {
    stop(
      "Transitions must evaluate to a numeric. The return type was `", class(value),
      "` for the transition: ", format.transition(transition),
      call. = FALSE
    )
  }

  if (transition$transition_type == "probability" && length(value) != 1) {
    stop(
      "Probability type transitions must evaluate to a vector of length `1`. ",
      "The returned length was `", length(value), "` for the transition: ",
      format.transition(transition),
      call. = FALSE
    )
  }

  if (
    transition$transition_type == "duration" &&
    transition_is_mortality(transition) &&
    length(value) != 1
  ) {
    # TODO we could consider support this if it'd be more biologically realistic
    stop(
      "Found non-constant mortality for a duration-based transition.",
      "Only scalar mortality values are supported. This occured in the transition: ",
      format.transition(transition),
      call. = FALSE
    )
  }

  if (transition$transition_type == "duration" && !(length(value) %in% c(1, max_duration + 1))) {
    stop(
      "Duration type transitions must evaluate to a vector of length `1`, or of length ",
      "`max_duration + 1`. The returned length was `", length(value),
      "` for the transition: ", format.transition(transition),
      call. = FALSE
    )
  }
}

#' Generate a matrix of transition probabilities between tick life stages
#'
#' @param time Numeric vector indicating day to get transition probabilities
#' @param population Tick population matrix. See get_tick_den for details.
#' @param developing_population Matrix of currently developing ticks. See
#'   get_tick_den for details.
#' @param tick_transitions A \code{\link{life_cycle}} object
#' @param predictors A \code{\link{predictors}} object
#'
#' @returns Matrix of transition probabilities, indicating the probabilities of
#'   transitioning from each stage (axis 1) to each stage (axis 2).
#'
#' @noRd
gen_transition_matrix <- function(
  time, population, developing_population, tick_transitions, predictors
) {
  life_stages <- rownames(population)
  trans_matrix <- empty_transition_matrix(life_stages)

  transitions <- tick_transitions |>
    query_transitions_by_mortality(mortality = FALSE) |>
    query_transitions("transition_type", "probability")

  mort <- tick_transitions |>
    query_transitions_by_mortality(mortality = TRUE) |>
    query_transitions("transition_type", "probability")

  for (i in transitions) {
    trans_matrix[i$from, i$to] <- get_transition_value(
      time = time,
      transition = i,
      predictors = predictors,
      max_duration = NULL,
      population = population,
      developing_population = developing_population
    )
  }

  for (i in mort) {
    mortality <- get_transition_value(
      time = time,
      transition = i,
      predictors = predictors,
      max_duration = NULL,
      population = population,
      developing_population = developing_population
    )

    # The max(0, 1- ...) structure should ensure that survival is between 0
    # and 1. This line also ensures that when a reproductive tick lays
    # (multiple) eggs, the reproductive tick will not survive to the next
    # stage. This is the desired behavior because all ticks are semelparous.
    mortality <- max(0, 1 - sum(trans_matrix[i$from, ]) - mortality)

    trans_matrix[i$from, i$from] <- mortality
  }

  trans_matrix
}


#' Update the delay matrices directly for the current time
#' @noRd
update_delays <- function(
    time, delay_mat, population, developing_population, tick_transitions, max_duration, predictors
) {
  life_stages <- rownames(population)
  transitions <- query_transitions(tick_transitions, "transition_type", "duration")

  for (from_stage in life_stages(transitions)) {
    trans <- transitions |>
      query_transitions("from", from_stage) |>
      query_transitions_by_mortality(mortality = FALSE) |>
      # there can only be one duration-based transition from each life stage, so
      # unlisting should just give the first element
      unlist(recursive = FALSE)

    val <- get_transition_value(
      time = time, transition = trans, predictors = predictors,
      max_duration = max_duration, population = population,
      developing_population = developing_population
    )

    days_to_next <- get_transition_duration(val = val, max_duration = max_duration)

    mort_transition <- transitions |>
      query_transitions("from", from_stage) |>
      query_transitions_by_mortality(mortality = TRUE) |>
      unlist(recursive = FALSE)

    surv_to_next <- get_transition_survival(
      mort_transition = mort_transition, time = time, predictors = predictors,
      max_duration = max_duration, population = population,
      developing_population = developing_population, days_to_next = days_to_next
    )

    # Calculate exactly how many ticks are entering the delay
    new_delayed <- population[from_stage, time] * surv_to_next
    emerge_day <- time + days_to_next

    # Add them to the specific day they will emerge
    delay_mat[trans[["to"]], emerge_day] <- delay_mat[trans[["to"]], emerge_day] + new_delayed

    # Add them to the developing pool for the exact days they are "in the oven"
    # (Starting tomorrow, ending the day before they emerge)
    if (days_to_next > 1) {
      dev_days <- (time + 1):(emerge_day - 1)
      developing_population[from_stage, dev_days] <-
        developing_population[from_stage, dev_days] + new_delayed
    }
  }

  list(delay_mat = delay_mat, developing_population = developing_population)
}



get_transition_survival <- function(
  mort_transition, time, predictors, max_duration, population, developing_population, days_to_next
) {
  if (is.null(mort_transition)) {
    # Case where there's no explicit mortality
    return(1)
  }

  # Case where there's a transition representing mortality
  mort <- get_transition_value(
    time = time,
    transition = mort_transition,
    predictors = predictors,
    max_duration = max_duration,
    population = population,
    developing_population = developing_population
  )

  if (mort_transition[["mortality_type"]] == "throughout_transition") {
    # Apply scalar mortality once during the transition
    return(1 - mort)
  }

  # Apply scalar mortality every day
  (1 - mort) ^ days_to_next
}

get_transition_duration <- function(val, max_duration) {
  # Duration-type transitions return either a vector of length 1, or of length
  # max_duration + 1. If the return value is of length 1, we interpret this
  # as the daily rate that the transition takes place. In this case, we
  # increase the length of the vector, so that we can take a cumulative sum
  # and determine the first day that the cumulative sum >= 1.
  # We add 1 to the length for consistency with output vector length from
  # transitions that use predictor data in a table, for which the length of
  # the output vector is determined in `get_pred()`.
  if (length(val) == 1) {
    val <- rep(val, max_duration + 1)
  }

  # With tiny tolerance for floating-point accumulation errors
  days <- cumsum(val) >= (1 - 1e-9)
  #days <- cumsum(val) >= 1

  if (!any(days)) {
    # Note that this has to be a run-time check, because it's dependent on
    # model state that varies throughout the model run
    stop(
      "Cumulative sum of daily transition probabilities never reached 1, ",
      "max_duration may be too small",
      call. = FALSE
    )
  }

  # Transition duration is the number of days until the first day when the sum
  # of the daily probabilities >= 1
  min(which(days))
}


# Create an empty (zero population) population matrix
empty_population_matrix <- function(life_stages, steps) {
  matrix(
    data = 0,
    nrow = length(life_stages),
    ncol = steps,
    dimnames = list(life_stages)
  )
}

# Create an empty matrix of transition probabilities between life stages
empty_transition_matrix <- function(life_stages) {
  n_life_stages <- length(life_stages)
  matrix(
    0,
    ncol = n_life_stages,
    nrow = n_life_stages,
    dimnames = list(life_stages, life_stages)
  )
}

# Set initial population for each life stage - zero if not specified in cfg
set_initial_population <- function(population, initial_population) {
  population[, 1] <- vapply(
    rownames(population),
    function(stage) {
      ifelse(stage %in% names(initial_population), initial_population[[stage]], 0)
    },
    FUN.VALUE = numeric(1L)
  )
  population
}


#' Precompute predictors for O(1) lookups
#' @param table input predictors table
#' @returns a nested list structure for fast extraction
#' @noRd
precompute_predictors <- function(table) {
  if (is.null(table)) return(NULL)

  preds <- unique(table$pred)
  out <- list()

  for (p in preds) {
    sub <- table[table$pred == p, ]

    if (all(is.na(sub$j_day))) {
      # Handle constant predictors
      vals <- sub$value
      if (!all(is.na(sub$pred_subcategory))) {
        names(vals) <- sub$pred_subcategory
      }
      out[[p]] <- list(is_constant = TRUE, data = vals)

    } else {
      # Handle variable predictors using a list indexed by day
      max_day <- max(sub$j_day, na.rm = TRUE)
      day_list <- vector("list", max_day)

      for (d in 1:max_day) {
        day_sub <- sub[sub$j_day == d, ]
        if (nrow(day_sub) > 0) {
          vals <- day_sub$value
          if (!all(is.na(day_sub$pred_subcategory))) {
            names(vals) <- day_sub$pred_subcategory
          }
          day_list[[d]] <- vals
        }
      }
      out[[p]] <- list(is_constant = FALSE, data = day_list)
    }
  }
  out
}