DICErClust Illustrated Example: Heart Failure Risk Stratification

## Install from local tarball (run once):
# install.packages(
#   "/path/to/DICErClust_0.1.1.tar.gz",
#   repos = NULL, type = "source"
# )
library(DICErClust)
library(ggplot2)
library(pROC)

hf_url  <- paste0(
  "https://archive.ics.uci.edu/ml/",
  "machine-learning-databases/00519/",
  "heart_failure_clinical_records_dataset.csv"
)
hf_dest <- tempfile(fileext = ".csv")
download.file(hf_url, hf_dest, quiet = TRUE)
hf <- read.csv(hf_dest)

cat(sprintf("Rows: %d   Columns: %d\n", nrow(hf), ncol(hf)))
print(table(DEATH_EVENT = hf$DEATH_EVENT))

## Continuous lab features → LSTM encoder (data_x)
x_cols <- c("age", "creatinine_phosphokinase", "ejection_fraction",
            "platelets", "serum_creatinine", "serum_sodium", "time")

## Binary demographic indicators → outcome head (data_v)
v_cols <- c("anaemia", "diabetes", "high_blood_pressure", "sex", "smoking")

## Min-max scale continuous features to [0, 1].
## Scaling prevents any single lab value from dominating the MSE
## reconstruction loss relative to others.
scale_01 <- function(x) {
  r <- range(x, na.rm = TRUE)
  if (diff(r) == 0) return(x * 0)
  (x - r[1]) / diff(r)
}

X_x <- apply(as.matrix(hf[, x_cols]), 2, scale_01)  # 299 × 7, numeric
X_v <- apply(as.matrix(hf[, v_cols]), 2, as.numeric) # 299 × 5, binary as float

## Note: data_v *must* be stored as numeric (float), not integer.
## torch_tensor() infers dtype from R storage mode; integer columns produce
## int64 tensors that are incompatible with the float32 model weights.

cat(sprintf("data_x: %d × %d\ndata_v: %d × %d\n",
            nrow(X_x), ncol(X_x), nrow(X_v), ncol(X_v)))

n_x <- ncol(X_x)  # 7  continuous predictors
n_v <- ncol(X_v)  # 5  binary demographics
outcome <- hf$DEATH_EVENT

set.seed(1111)
idx_death <- which(outcome == 1)
idx_alive <- which(outcome == 0)

train_idx <- sort(c(
  sample(idx_death, floor(0.70 * length(idx_death))),
  sample(idx_alive, floor(0.70 * length(idx_alive)))
))
test_idx <- setdiff(seq_len(nrow(hf)), train_idx)

cat(sprintf("Train: %d patients  (deaths: %d, %.0f%%)\n",
            length(train_idx), sum(outcome[train_idx]),
            100 * mean(outcome[train_idx])))
cat(sprintf("Test : %d patients  (deaths: %d, %.0f%%)\n",
            length(test_idx),  sum(outcome[test_idx]),
            100 * mean(outcome[test_idx])))

data_dir <- file.path(tempdir(), "dice_hf")
dir.create(data_dir, showWarnings = FALSE)

saveRDS(
  list(X_x[train_idx, ], X_v[train_idx, ], as.integer(outcome[train_idx])),
  file.path(data_dir, "hf_train.rds")
)
saveRDS(
  list(X_x[test_idx, ], X_v[test_idx, ], as.integer(outcome[test_idx])),
  file.path(data_dir, "hf_test.rds")
)

args <- list(
  seed              = 1111,          # reproducibility seed
  input_path        = data_dir,      # directory containing RDS files
  filename_train    = "hf_train.rds",
  filename_test     = "hf_test.rds",

  ## ── Architecture ──────────────────────────────────────────
  n_input_fea       = n_x,          # 7 continuous LSTM input features
  n_hidden_fea      = 4,            # LSTM latent dimension (7 → 4)
  lstm_layer        = 1,            # single LSTM layer
  lstm_dropout      = 0.0,          # no dropout (small dataset)
  K_clusters        = 2,            # binary risk partition: high vs. low

  ## ── Auxiliary features ────────────────────────────────────
  n_dummy_demov_fea = n_v,          # 5 binary demographic covariates

  ## ── Hardware ──────────────────────────────────────────────
  cuda              = FALSE,        # set TRUE for GPU acceleration

  ## ── Optimiser ─────────────────────────────────────────────
  lr                = 1e-4,         # Adam learning rate

  ## ── Training schedule ─────────────────────────────────────
  init_AE_epoch     = 5,            # Stage 1: autoencoder warm-up epochs
  iter              = 30,           # Stage 2: number of clustering iterations
  epoch_in_iter     = 2,            # gradient-update epochs per iteration

  ## ── Loss weights ──────────────────────────────────────────
  ## Combined loss: L = λ_AE·L_AE + λ_clf·L_classifier
  ##                  + λ_out·L_outcome + λ_p·L_p_value
  ## L_p_value = 3.841 − G penalises non-significant cluster configurations
  ## (G is the LRT statistic; 3.841 is the χ²(1) critical value at α = 0.05)
  lambda_AE         = 1.0,
  lambda_classifier = 1.0,
  lambda_outcome    = 1.0,
  lambda_p_value    = 1.0
)

## DICEr writes output files relative to the working directory.
## We temporarily switch to tempdir() to keep them self-contained.
old_wd <- setwd(tempdir())
suppressWarnings(DICEr(args))
setwd(old_wd)

part2_dir <- file.path(tempdir(), "hn_4_K_2", "part2_AE_nhidden_4")

if (!file.exists(file.path(part2_dir, "data_train_iter.rds"))) {
  stop(
    "No checkpoint found — the p < 0.05 criterion was not met in ",
    args$iter, " iterations.  Increase args$iter and rerun."
  )
}

res_train <- readRDS(file.path(part2_dir, "data_train_iter.rds"))
res_test  <- readRDS(file.path(part2_dir, "data_test_iter.rds"))