[R] Add class names to coefficients (dmlc#10745)

Author: david-cortes

R-package/R/xgb.Booster.R

@@ -1109,17 +1109,25 @@ coef.xgb.Booster <- function(object, ...) {
   if (n_cols == 1L) {
     out <- c(intercepts, coefs)
     if (add_names) {
-      names(out) <- feature_names
+      .Call(XGSetVectorNamesInplace_R, out, feature_names)
     }
   } else {
     coefs <- matrix(coefs, nrow = num_feature, byrow = TRUE)
     dim(intercepts) <- c(1L, n_cols)
     out <- rbind(intercepts, coefs)
+    out_names <- vector(mode = "list", length = 2)
     if (add_names) {
-      row.names(out) <- feature_names
+      out_names[[1L]] <- feature_names
     }
-    # TODO: if a class names attributes is added,
-    # should use those names here.
+    if (inherits(object, "xgboost")) {
+      metadata <- attributes(object)$metadata
+      if (NROW(metadata$y_levels)) {
+        out_names[[2L]] <- metadata$y_levels
+      } else if (NROW(metadata$y_names)) {
+        out_names[[2L]] <- metadata$y_names
+      }
+    }
+    .Call(XGSetArrayDimNamesInplace_R, out, out_names)
   }
   return(out)
 }

R-package/tests/testthat/test_basic.R

@@ -750,6 +750,19 @@ test_that("Coefficients from gblinear have the expected shape and names", {
   pred_auto <- predict(model, x, outputmargin = TRUE)
   pred_manual <- unname(mm %*% coefs)
   expect_equal(pred_manual, pred_auto, tolerance = 1e-7)
+
+  # xgboost() with additional metadata
+  model <- xgboost(
+    iris[, -5],
+    iris$Species,
+    booster = "gblinear",
+    objective = "multi:softprob",
+    nrounds = 3,
+    nthread = 1
+  )
+  coefs <- coef(model)
+  expect_equal(row.names(coefs), c("(Intercept)", colnames(x)))
+  expect_equal(colnames(coefs), levels(iris$Species))
 })
 
 test_that("Deep copies work as expected", {