Loss_MC3.R

# ----- Utilities -----



# Calculate the probability of each class given intercepts and linear predictor
calc_prob <- function(alpha, mu){
  p_list <- lapply(1:length(mu), function(i){
    torch::torch_diff(torch::torch_sigmoid(alpha - mu[i]))
  })
  return(torch::torch_stack(p_list))
}

# Ramanujan's approximation for log(x!)
rj_log_factorial <- function(x){
  num_1 <- (x + 4 * torch::torch_pow(x, 2) + 8 * torch::torch_pow(x, 3))
  term_1 <- (1 / 6) * torch::torch_log(num_1)
  term_2 <- (log(pi) / 2)
  return(x * torch::torch_log(x) - x + term_1 + term_2)
}

# Loss function
# Multinomial log loss (negative)
# Calculates as though one of each class has been observed 
# to avoid the undefined case for rj_log_factorial(0)
multi_ll <- function(output, target){
  n <- torch::torch_sum(target, dim = 2) + dim(target)[2]
  log_n_fac <- rj_log_factorial(n)
  sum_log_y_fac <- torch::torch_sum(rj_log_factorial(target + 1), dim = 2)
  sum_ylog_p <- torch::torch_sum((target + 1) * torch::torch_log(output), dim = 2)
  ll <- log_n_fac - sum_log_y_fac + sum_ylog_p
  return(torch::torch_sum(ll))
}

# Main MLP
phi_mlp <- torch::nn_module(
  
  # Model name
  "phi_mlp",
  
  # Initialization
  initialize = function(
    features_in
  ) {
    
    # Simple MLP
    self$mlp <- torch::nn_module_list()
    self$mlp$append(torch::nn_flatten())
    self$mlp$append(torch::nn_linear(in_features = features_in, out_features = 16))
    self$mlp$append(torch::nn_dropout(0.6))
    self$mlp$append(torch::nn_batch_norm1d(16))
    self$mlp$append(torch::nn_relu())
    self$mlp$append(torch::nn_linear(in_features = 16, out_features = 32))
    self$mlp$append(torch::nn_dropout(0.6))
    self$mlp$append(torch::nn_batch_norm1d(32))
    self$mlp$append(torch::nn_relu())
    self$mlp$append(torch::nn_linear(in_features = 32, out_features = 16))
    self$mlp$append(torch::nn_dropout(0.6))
    self$mlp$append(torch::nn_batch_norm1d(16))
    self$mlp$append(torch::nn_relu())
    self$mlp$append(torch::nn_linear(in_features = 16, out_features = 1))
    
  },
  
  forward = function(x) {
    for(i in 1:length(self$mlp)){
      #message("MLP module ", i)
      #message(paste0(dim(x), collapse = " "))
      x <- self$mlp[[i]](x)
    }
    return(x)
  }
  
)



# ----- Construct training data -----



# Load the data
climate_data <- data.table::fread("ProcessedData/20250429_climate.csv", data.table = FALSE)
fuel_spread <- data.table::fread("ProcessedData/20250429_fuelSpreadRate.csv", data.table = FALSE)
fuel_id <- sprintf("r%05d_c%05d.png", fuel_spread$row_id, fuel_spread$col_id)
ndvi_pred <- readRDS("ProcessedData/20250430_predClass.rds")

# Match training data
train_imgs <- list.files("C:/Users/CBuglino/Desktop/NDVI_Data")
train_ids <- gsub("^(.*_)(r[0-9]{5}.*)", "\\2", train_imgs)
train_ids <- intersect(train_ids, climate_data$file_id)
train_ids <- intersect(train_ids, fuel_id)
train_ids <- intersect(train_ids, ndvi_pred$img_id)
climate_data <- climate_data[match(train_ids, climate_data$file_id),]
fuel_spread <- fuel_spread[match(train_ids, fuel_id),]
p_hat <- ndvi_pred$p_hat[match(train_ids, ndvi_pred$img_id),]

# Remove observations with no target info
y <- as.matrix(fuel_spread[, grepl("spread", names(fuel_spread))])
all_zero_y <- (rowSums(y) == 0)
y <- y[!all_zero_y,]
climate_data <- climate_data[!all_zero_y,]
p_hat <- p_hat[!all_zero_y,]

# Remove observations with no predictor info
elev_na <- is.na(climate_data$elev_avg)
y <- y[!elev_na,]
climate_data <- climate_data[!elev_na,]
p_hat <- p_hat[!elev_na,]

# Remove the "extreme" label since this is never observed
y <- y[, !grepl("extreme", colnames(y))]

# Construct the covariate matrix
cov_inds <- c(2, 4, 6)
x_mu <- apply(climate_data[, cov_inds], 2, mean)
x_sd <- apply(climate_data[, cov_inds], 2, sd)
climate_z <- climate_data
climate_z[, cov_inds] <- lapply(seq_along(cov_inds), function(j){
  (climate_z[, cov_inds[j]] - x_mu[j]) / x_sd[j]
})
X <- model.matrix(~ -1 + airtemp_avg * precip_avg * elev_avg, data = climate_z)
X <- cbind(X, p_hat)
X <- torch::torch_tensor(X)

# Construct the target matrix
y <- torch::torch_tensor(y)

# Consolidate
data_train <- list(
  X = X$clone(),
  y = y$clone()
)



# ----- Construct the test data -----



# Load the data
climate_data <- data.table::fread("ProcessedData/20250429_climate.csv", data.table = FALSE)
fuel_spread <- data.table::fread("ProcessedData/20250429_fuelSpreadRate.csv", data.table = FALSE)
fuel_id <- sprintf("r%05d_c%05d.png", fuel_spread$row_id, fuel_spread$col_id)
ndvi_pred <- readRDS("ProcessedData/20250430_predClass.rds")

# Match test sets
test_ids <- intersect(climate_data$file_id, fuel_id)
test_ids <- intersect(test_ids, ndvi_pred$img_id)
climate_data <- climate_data[match(test_ids, climate_data$file_id),]
fuel_spread <- fuel_spread[match(test_ids, fuel_id),]
p_hat <- ndvi_pred$p_hat[match(test_ids, ndvi_pred$img_id),]

# Remove observations with no target info
y <- as.matrix(fuel_spread[, grepl("spread", names(fuel_spread))])
all_zero_y <- (rowSums(y) == 0)
y <- y[!all_zero_y,]
climate_data <- climate_data[!all_zero_y,]
p_hat <- p_hat[!all_zero_y,]

# Remove observations with no predictor info
elev_na <- is.na(climate_data$elev_avg)
y <- y[!elev_na,]
climate_data <- climate_data[!elev_na,]
p_hat <- p_hat[!elev_na,]

# Remove the "extreme" label since this is never observed
y <- y[, !grepl("extreme", colnames(y))]

# Construct the covariate matrix
cov_inds <- c(2, 4, 6)
climate_z <- climate_data
climate_z[, cov_inds] <- lapply(seq_along(cov_inds), function(j){
  (climate_z[, cov_inds[j]] - x_mu[j]) / x_sd[j]
})
X <- model.matrix(~ -1 + airtemp_avg * precip_avg * elev_avg, data = climate_z)
X <- cbind(X, p_hat)
X <- torch::torch_tensor(X)

# Construct the target matrix
y <- torch::torch_tensor(y)

# Consolidate
data_test <- list(
  X = X$clone(),
  y = y$clone()
)



# ----- Model loss -----



# Compute training and validation loss over training pipeline
design <- data.frame(
  epochs = c(1, seq(5, 125, 5)),
  train_loss = NA,
  val_loss = NA
)

# Compute
pb <- txtProgressBar(0, nrow(design), style = 3)
for(n in 1:nrow(design)){
  
  # Load the model onto the GPU
  nom <- sprintf("model_climate_e%03d.pt", design$epochs[n])
  model <- torch::torch_load(nom, device = "cpu")
  nom <- sprintf("alpha_climate_e%03d.pt", design$epochs[n])
  alpha <- torch::torch_load(nom, device = "cpu")
  model$eval()
  
  # Manual batch learning to avoid overflow when computing the loss
  batch_size <- 2^7
  block_inds <- sample.int(nrow(data_train$X), nrow(data_train$X), replace = FALSE)
  batch_list <- split(block_inds, ceiling(seq_along(block_inds) / batch_size))
  
  # Mini batch learning
  ll_list <- list()
  for(j in seq_along(batch_list)){
    
    # Select the data
    batch_inds <- batch_list[[j]]
    batch_x <- torch::torch_tensor(data_train$X[batch_inds,])
    batch_y <- torch::torch_tensor(data_train$y[batch_inds,])
    
    # Learn on entire dataset
    mu <- model(batch_x)
    p_hat <- calc_prob(alpha, mu)
    
    # Compute loss
    loss <- multi_ll(p_hat, batch_y)
    
    # Update step
    ll_list[[length(ll_list) + 1]] <- as.numeric(loss$detach())
    
  }
  
  # Record total log likelihood
  ll <- sum(unlist(ll_list))
  design$train_loss[n] <- ll / nrow(data_train$X)
  
  # Manual batch learning to avoid overflow when computing the loss
  batch_size <- 2^7
  block_inds <- sample.int(nrow(data_test$X), nrow(data_test$X), replace = FALSE)
  batch_list <- split(block_inds, ceiling(seq_along(block_inds) / batch_size))
  
  # Mini batch learning
  ll_list <- list()
  for(j in seq_along(batch_list)){
    
    # Select the data
    batch_inds <- batch_list[[j]]
    batch_x <- torch::torch_tensor(data_test$X[batch_inds,])
    batch_y <- torch::torch_tensor(data_test$y[batch_inds,])
    
    # Learn on entire dataset
    mu <- model(batch_x)
    p_hat <- calc_prob(alpha, mu)
    
    # Compute loss
    loss <- multi_ll(p_hat, batch_y)
    
    # Update step
    ll_list[[length(ll_list) + 1]] <- as.numeric(loss$detach())
    
  }
  
  # Record total log likelihood
  ll <- sum(unlist(ll_list))
  design$val_loss[n] <- ll / nrow(data_test$X)
  
  # Save design
  dnom <- sprintf("design_climate_%03d.csv", n)
  write.csv(design, dnom, row.names = FALSE, quote = FALSE)
  setTxtProgressBar(pb, n)
  
}



# ----- Naive loss -----



# Calculate observed proportions from training data
y_train <- as.matrix(data_train$y)
p_hat_naive <- colSums(y_train) / sum(y_train)

# Naive loss on the validation data
# Calculate validation loss per observation
# Manual batch learning to avoid overflow when computing the loss
batch_size <- 2^7
block_inds <- sample.int(nrow(data_test$X), nrow(data_test$X), replace = FALSE)
batch_list <- split(block_inds, ceiling(seq_along(block_inds) / batch_size))

# Mini batch learning
ll_list <- list()
for(j in seq_along(batch_list)){
  
  # Select the data
  batch_inds <- batch_list[[j]]
  batch_x <- torch::torch_tensor(data_test$X[batch_inds,])
  batch_y <- torch::torch_tensor(data_test$y[batch_inds,])
  
  # Learn on entire dataset
  p_hat <- matrix(p_hat_naive, nrow(batch_x), length(p_hat_naive), byrow = TRUE)
  p_hat <- torch::torch_tensor(p_hat)
  
  # Compute loss
  loss <- multi_ll(p_hat, batch_y)
  
  # Update step
  ll_list[[length(ll_list) + 1]] <- as.numeric(loss$detach())
  
}

# Record total log likelihood
ll <- sum(unlist(ll_list))

# Save results
dat <- data.frame(
  epochs = 0,
  train_loss = NA,
  val_loss = ll / nrow(data_test$X)
)
loss_data <- rbind(dat, design)
write.csv(loss_data, file = "loss_data_climate3.csv", row.names = FALSE, quote = FALSE)