Fix hierarchical ancestor missed in the transport to model

.Rbuildignore

@@ -14,3 +14,5 @@
 ^CRAN-SUBMISSION$
 ^revdep$
 ^vignettes/*_files$
+^\.claude$
+^\.positai$

.gitignore

@@ -15,3 +15,5 @@ tabnet_*.tar.gz
 tabnet.Rproj
 po/glossary.csv
 inst/IMPORTLIST
+.positai
+tools

DESCRIPTION

@@ -46,6 +46,7 @@ Imports:
     zeallot
 Suggests: 
     cli,
+    fBasics,
     knitr,
     modeldata,
     patchwork,
@@ -68,5 +69,5 @@ Config/testthat/parallel: false
 Config/testthat/start-first: interface, explain, params
 Encoding: UTF-8
 Roxygen: list(markdown = TRUE)
-RoxygenNote: 7.3.3
 Language: en-US
+Config/roxygen2/version: 8.0.0

NAMESPACE

@@ -24,6 +24,7 @@ S3method(tabnet_pretrain,recipe)
 S3method(update,tabnet)
 export("%>%")
 export(attention_width)
+export(build_ancestor_matrix_from_outcomes)
 export(cat_emb_dim)
 export(check_compliant_node)
 export(checkpoint_epochs)
@@ -39,6 +40,9 @@ export(mlp_activation)
 export(mlp_hidden_multiplier)
 export(momentum)
 export(nn_aum_loss)
+export(nn_mc_loss)
+export(nnf_mc_loss)
+export(nnf_multilabel_one_hot)
 export(node_to_df)
 export(num_independent)
 export(num_independent_decoder)
@@ -71,6 +75,18 @@ importFrom(rlang,.data)
 importFrom(stats,predict)
 importFrom(stats,update)
 importFrom(tidyr,replace_na)
+importFrom(torch,as_array)
+importFrom(torch,nn_module)
 importFrom(torch,nn_prune_head)
+importFrom(torch,torch_argsort)
+importFrom(torch,torch_cat)
+importFrom(torch,torch_long)
+importFrom(torch,torch_matmul)
+importFrom(torch,torch_mean)
+importFrom(torch,torch_minimum)
+importFrom(torch,torch_mul)
+importFrom(torch,torch_std)
+importFrom(torch,torch_sum)
+importFrom(torch,torch_tensor)
 importFrom(tune,min_grid)
 importFrom(zeallot,"%<-%")

NEWS.md

@@ -1,5 +1,9 @@
 # tabnet (development version)
 
+## Bugfixes
+
+* Ancestor matrix is now taken into account for hierarchical classification (#188).
+
 # tabnet 0.8.0
 
 ## New features

R/hardhat.R

@@ -162,16 +162,14 @@ tabnet_fit.Node <- function(x, tabnet_model = NULL, config = tabnet_config(), ..
   # get tree leaves and extract attributes into data.frames
   xy_df <- node_to_df(x)
   processed <- hardhat::mold(xy_df$x, xy_df$y)
-  # Given n classes, M is an (n x n) matrix where M_ij = 1 if class i is descendant of class j
-  ancestor <- data.tree::ToDataFrameNetwork(x) %>%
-   mutate_if(is.character, ~.x %>% as.factor %>% as.numeric)
-  # TODO check correctness
-  # embed the M matrix in the config$ancestor variable
-  dims <- c(max(ancestor), max(ancestor))
-  ancestor_m <- Matrix::sparseMatrix(ancestor$from, ancestor$to, dims = dims, x = 1)
   check_type(processed$outcomes)
-
+  
   config <- merge_config_and_dots(config, ...)
+  # add ancestor boolean sparse matrix to config
+  # check_dag_compliance(xy_df$y)
+  config$ancestor <- build_ancestor_matrix_from_outcomes(x, processed$outcomes)
+  # make outcomes levels available so that batched y could be one-hot encoded.
+  config$outcomes <- processed$outcomes
   tabnet_bridge(processed, config = config, tabnet_model, from_epoch, task = "supervised")
 }
 
@@ -272,7 +270,7 @@ tabnet_pretrain.default <- function(x, ...) {
 
 #' @export
 #' @rdname tabnet_pretrain
-tabnet_pretrain.data.frame <- function(x, y, tabnet_model = NULL, config = tabnet_config(), ..., from_epoch = NULL) {
+tabnet_pretrain.data.frame <- function(x, y = NULL, tabnet_model = NULL, config = tabnet_config(), ..., from_epoch = NULL) {
   processed <- hardhat::mold(x, y)
 
   config <- merge_config_and_dots(config, ...)
@@ -309,8 +307,7 @@ tabnet_pretrain.Node <- function(x, tabnet_model = NULL, config = tabnet_config(
   check_compliant_node(x)
   # get tree leaves and extract attributes into data.frames
   xy_df <- node_to_df(x)
-  tabnet_pretrain(xy_df$x, xy_df$y, tabnet_model = tabnet_model, config = config, ..., from_epoch = from_epoch)
-
+  tabnet_pretrain(xy_df$x, tabnet_model = tabnet_model, config = config, ..., from_epoch = from_epoch)
 }
 
 new_tabnet_pretrain <- function(pretrain, blueprint) {
@@ -418,13 +415,15 @@ tabnet_bridge <- function(processed, config = tabnet_config(), tabnet_model, fro
 #' @importFrom stats predict
 #' @export
 predict.tabnet_fit <- function(object, new_data, type = NULL, ..., epoch = NULL) {
-  if (inherits(new_data, "Node")) {
+  if (inherits(new_data, "Node") && !is.null(object$fit$config$ancestor)) {
     new_data_df <- node_to_df(new_data)$x
+    # Enforces column order, type, column names, etc
+    processed <- hardhat::forge(new_data_df, object$blueprint)
+
   } else {
     new_data_df <- new_data
+    processed <- hardhat::forge(new_data, object$blueprint)
   }
-  # Enforces column order, type, column names, etc
-  processed <- hardhat::forge(new_data_df, object$blueprint)
   batch_size <- object$fit$config$batch_size
   out <- predict_tabnet_bridge(type, object, processed$predictors, epoch, batch_size)
   hardhat::validate_prediction_size(out, new_data_df)
@@ -436,8 +435,8 @@ predict_tabnet_bridge <- function(type, object, predictors, epoch, batch_size) {
   type <- check_type(object$blueprint$ptypes$outcomes, type)
   is_multi_outcome <- ncol(object$blueprint$ptypes$outcomes) > 1
   outcome_nlevels <- NULL
-  if (is_multi_outcome & type != "numeric") {
-    outcome_nlevels <- purrr::map_dbl(object$blueprint$ptypes$outcomes, ~length(levels(.x)))
+  if (is_multi_outcome && type != "numeric") {
+    outcome_nlevels <- purrr::map_dbl(object$blueprint$ptypes$outcomes, ~nlevels(.x))
   }
 
   if (!is.null(epoch)) {
@@ -458,6 +457,7 @@ predict_tabnet_bridge <- function(type, object, predictors, epoch, batch_size) {
     object$fit$network$load_state_dict(m$state_dict())
   }
 
+
   type_multioutcome <- paste0(type, "_", is_multi_outcome)
   switch(
     type_multioutcome,
@@ -604,4 +604,81 @@ nn_prune_head.tabnet_pretrain <- function(x, head_size) {
     nn_prune_head(x$fit$network, head_size=head_size)
   }
 
-}
+}
+
+
+#' Build ancestor matrix aligned with observed outcome classes
+#'
+#' Extracts class names from the outcome tibble (factor levels) and builds
+#' the ancestor matrix only for classes that actually appear in the data.
+#'
+#' @param x A `data.tree::Node` object.
+#' @param outcomes A tibble with factor columns (one per hierarchy level),
+#'   as returned by `hardhat::mold()$outcomes`.
+#' @param device Torch device ("cpu" or "cuda").
+#' @return A `torch_tensor` of shape `(1, n_classes, n_classes)`.
+#' @export
+build_ancestor_matrix_from_outcomes <- function(x, outcomes, device = "cpu") {
+  # 1. Extract all class names from factor levels (preserving order)
+  #    outcomes is a tibble with one factor column per hierarchy level
+  level_cols <- names(outcomes)
+  all_class_names <- unlist(lapply(outcomes, levels), use.names = FALSE)
+  n_classes <- length(all_class_names)
+
+  if (n_classes == 0L) {
+    runtime_error("No factor levels found in outcomes : {str(outcomes)}")
+  }
+
+  # 2. Build a lookup: class_name -> data.tree Node
+  all_nodes <- data.tree::Traverse(x, traversal = "pre-order")
+  all_nodes <- unname(all_nodes)
+  level_lengths <- lengths(lapply(outcomes, levels))
+  lvl_vector <- rep(seq_along(level_cols) + 1L, level_lengths)
+
+  # 3. Resolve each class name to its Node
+  class_nodes <- lapply(seq_along(all_class_names), function(k) {
+    nm <- all_class_names[k]
+    lvl <- lvl_vector[k]
+
+    candidates <- Filter(function(n) n$level == lvl && n$name == nm, all_nodes)
+    if (length(candidates) == 0) {
+      runtime_error("Factor level {.var {nm}} not found at tree level {lvl} (outcomes column {.var {level_cols[lvl - 1L]}})")
+    }
+    candidates[[1]]
+  })
+
+  # 4. Create 1-based index mapping
+  class_map <- setNames(seq_len(n_classes), all_class_names)
+
+  # 5. Collect (descendant, ancestor) pairs by climbing up
+  row_list <- vector("list", n_classes)
+  col_list <- vector("list", n_classes)
+
+  for (i in seq_len(n_classes)) {
+    current <- class_nodes[[i]]
+    anc_indices <- integer()
+
+    repeat {
+      idx <- class_map[current$name]
+      if (!is.null(idx)) {
+        anc_indices <- c(anc_indices, idx)
+      }
+      if (current$isRoot || is.null(current$parent)) break
+      current <- current$parent
+    }
+
+    row_list[[i]] <- rep(i, length(anc_indices))
+    col_list[[i]] <- anc_indices
+  }
+
+  # 6. Fill matrix
+  R <- matrix(0L, nrow = n_classes, ncol = n_classes)
+  rows <- unlist(row_list, use.names = FALSE)
+  cols <- unlist(col_list, use.names = FALSE)
+  if (length(rows) > 0) R[cbind(rows, cols)] <- 1L
+
+  # 7. Convert to torch
+  R_torch <- torch::torch_tensor(R, dtype = torch::torch_double(), device = device)
+  R_torch$unsqueeze(1)
+}
+