Add MSE + covariance penalty term custom objective function (#21)

* Add draft of cov penalty

* Add roxygen rd file

* Update pkgdown yml

* Lint

* Add nolint

* Edit nolint

* Attempt to fix data leakage

* Edit comments

* Add comments

* Add spacing

* Fix spacing

* Lint

* Decrease line length

* Remove zero grad lines

* Remove zero_grad code

* Remove type=raw spec

* Switch obj to objective

* Swap yc to y_centered

* Swap in for ==

* Point to actions branch

* Point back to main

* Switch default value check to null check

* Remove nolint tag

* Add y mean check

* Version bump

* Add new tests

* Remove hessian test

* Remove hessian tset

* Update formatting

Author: wagnerlmichael

DESCRIPTION

@@ -1,7 +1,7 @@
 Package: lightsnip
 Type: Package
 Title: LightGBM Model Engine for Parsnip
-Version: 0.0.6
+Version: 0.0.7
 Authors@R: c(
     person(given = "Dan", family = "Snow", email="daniel.snow@cookcountyil.gov", role=c("aut", "cre")),
     person(given = "Nicole", family = "Jardine", email="nicole.jardine@cookcountyil.gov", role=c("aut", "ctb")),
@@ -48,5 +48,6 @@ Suggests:
     tune,
     workflows,
     yardstick
-RoxygenNote: 7.2.3
+Roxygen: list(markdown = TRUE)
 Config/testthat/edition: 3
+Config/roxygen2/version: 8.0.0

NAMESPACE

@@ -16,6 +16,7 @@ export(lambda_l2)
 export(learning_rate)
 export(lgbm_load)
 export(lgbm_save)
+export(make_objective_mse_cov)
 export(max_bin)
 export(max_cat_threshold)
 export(max_depth)

R/lightgbm.R

@@ -111,8 +111,8 @@ add_boost_tree_lightgbm <- function() {
 #'   but it means that the Dataset object cannot be changed after it has been
 #'   constructed. If you'd prefer to be able to change the Dataset object after
 #'   construction, set \code{free_raw_data = FALSE}. Useful for debugging.
-#' @param verbose Integer. < 0: Fatal, = 0: Error (Warning), = 1: Info,
-#'   > 1: Debug.
+#' @param verbose Integer. Verbosity level: < 0 = Fatal, 0 = Error (Warning),
+#'   1 = Info, > 1 = Debug.
 #' @param save_tree_error Boolean. Whether or not to use the training set
 #'   to compute errors for each tree that will be stored on the record_evals
 #'   attribute. Note that this parameter is mutually exclusive with
@@ -146,8 +146,36 @@ train_lightgbm <- function(x,
   force(y)
   others <- list(...)
 
-  # Set training objective (always regression)
-  if (!any(names(others) %in% c("objective"))) {
+  # Custom objective handling. `mse_cov_rho` is a lightsnip-specific engine
+  # arg used only when `objective == "mse_cov"`; pop it off so it is not
+  # forwarded to lgb.train (which would error on an unknown parameter).
+  mse_cov_rho <- others$mse_cov_rho
+  others$mse_cov_rho <- NULL
+
+  custom_objective <- NULL
+
+  # Sentinel that gates two downstream branches:
+  #   - whether to fall back to the default "regression" objective
+  #   - whether to construct the mse_cov objective callback
+  mse_cov_rho_val <- NULL
+
+  if (!is.null(others$objective) && identical(others$objective, "mse_cov")) {
+    if (is.null(mse_cov_rho)) {
+      rlang::abort(
+        "`objective = \"mse_cov\"` requires `mse_cov_rho` to be set."
+      )
+    }
+    mse_cov_rho_val <- as.numeric(mse_cov_rho)
+    # Clear `objective`/`num_class` so lgb.train doesn't reject the unknown
+    # name when we hand it the callback via `obj`.
+    others$objective <- NULL
+    others$num_class <- NULL
+  }
+
+  # Set training objective default (always regression) when not specified.
+  # Skipped when a custom `obj` callback is in use, since lgb.train will then
+  # supply the gradient/hessian itself and `objective` must be unset.
+  if (is.null(mse_cov_rho_val) && !any(names(others) == "objective")) {
     others$num_class <- 1
     others$objective <- "regression"
   }
@@ -235,6 +263,17 @@ train_lightgbm <- function(x,
     trn_index <- 1:n
   }
 
+  # Build the mse_cov callback against training rows only — `y[val_index]`
+  # is held out for lgb.train's early stopping, so including those labels
+  # in `y_mean` would leak the holdout's label mean into the centering
+  # term used by the covariance penalty on every boosting iteration.
+  if (!is.null(mse_cov_rho_val)) {
+    custom_objective <- make_objective_mse_cov(
+      rho    = mse_cov_rho_val,
+      y_mean = mean(y[trn_index])
+    )
+  }
+
   d <- lightgbm::lgb.Dataset(
     data = as.matrix(x[trn_index, , drop = FALSE]),
     label = y[trn_index],
@@ -269,6 +308,10 @@ train_lightgbm <- function(x,
   if (!is.null(early_stop) && validation > 0) {
     main_args$early_stopping_rounds <- early_stop
   }
+  # Wire in the custom objective callback (if any) under lgb.train's `obj` arg
+  if (!is.null(custom_objective)) {
+    main_args$obj <- quote(custom_objective)
+  }
 
   call <- parsnip::make_call(fun = "lgb.train", ns = "lightgbm", main_args)
   rlang::eval_tidy(call, env = rlang::current_env())

R/objectives.R

@@ -0,0 +1,76 @@
+#' Custom LightGBM objective: MSE with a squared covariance penalty
+#'
+#' @description Build a custom LightGBM objective callback that minimizes a
+#' standard squared-error loss plus a soft penalty on the covariance between
+#' the per-sample residual \code{r = y_pred - y_true} and the (centered) labels
+#' `y_true`. The penalty pushes the model toward "vertical equity" by
+#' discouraging residuals that systematically scale with `y`.
+#'
+#' This is an R port of the `LGBCovPenalty` objective from an active
+#' collaboration (see the
+# nolint start: line_length_linter.
+#' [soft-vertical-equity-constrained-mass-appraissal](https://github.com/nicacevedo/soft-vertical-equity-constrained-mass-appraissal)
+# nolint end
+#' repository). It is intended to be used when the model is trained in
+#' log-space (so the "diff" residual is equivalent to a log-ratio).
+#'
+#' Penalty (using mean-centered labels `y_centered = y_true - mean(y_true)`):
+#' \deqn{\mathrm{cov} = \frac{1}{n} \sum_{i=1}^{n} r_i \tilde{y}_i}{
+#'   cov = (1/n) * sum_i r_i * y_centered_i}
+#' \deqn{\mathrm{penalty} = \frac{\rho}{2} \, n \, \mathrm{cov}^2}{
+#'   penalty = 0.5 * rho * n * cov^2}
+#' where \eqn{\tilde{y}}{y_centered} is the vector of mean-centered labels.
+#' A diagonal Hessian approximation is used (matches the reference Python
+#' implementation).
+#'
+#' @param rho Numeric. Non-negative penalty weight. `rho = 0` recovers plain
+#'   MSE.
+#' @param y_mean Numeric. Mean of the training labels. Should be computed once
+#'   from the training set and captured here so the centering is stable across
+#'   iterations.
+#'
+#' @return A function with signature `function(preds, dtrain)` suitable for
+#'   passing as the `obj` argument of [lightgbm::lgb.train].
+#'
+#' @export
+make_objective_mse_cov <- function(rho, y_mean) {
+  rho <- as.numeric(rho)
+  y_mean <- as.numeric(y_mean)
+  if (length(rho) != 1L || is.na(rho) || rho < 0) {
+    rlang::abort("`rho` must be a single non-negative numeric value.")
+  }
+  if (length(y_mean) != 1L || is.na(y_mean)) {
+    rlang::abort("`y_mean` must be a single non-missing numeric value.")
+  }
+
+  function(preds, dtrain) {
+    y_true <- lightgbm::get_field(dtrain, "label")
+    y_pred <- as.numeric(preds)
+    n <- length(y_pred)
+
+    # Centered labels (training-set mean is captured at construction time so
+    # the penalty geometry stays stable across boosting iterations)
+    y_centered <- y_true - y_mean
+
+    # Residual ("diff" mode); in log-space training this is the log-ratio
+    r <- y_pred - y_true
+
+    # Covariance estimate (E[y_centered] is ~0 by construction)
+    cov_val <- mean(r * y_centered)
+
+    # Base squared-error grad/hess
+    grad_base <- 2.0 * (y_pred - y_true)
+    hess_base <- rep(2.0, n)
+
+    # Penalty grad/hess (diagonal approximation)
+    # dc/dy_pred_i = (1/n) * y_centered_i  (since dr_i/dy_pred_i = 1)
+    a <- y_centered / n
+    grad_pen <- rho * n * cov_val * a
+    hess_pen <- rho * n * (a^2)
+
+    grad <- grad_base + grad_pen
+    hess <- hess_base + hess_pen
+
+    list(grad = grad, hess = hess)
+  }
+}

_pkgdown.yml

@@ -4,6 +4,7 @@ llm-docs: true
 
 template:
   bootstrap: 5
+  math-rendering: katex
 
 authors:
   sidebar:
@@ -26,9 +27,15 @@ reference:
   - train_lightgbm
   - pred_lgb_reg_num
   - multi_predict._lgb.Booster
+- subtitle: Custom objectives
+  desc: >
+    Factory functions that return custom LightGBM objective callbacks for
+    use as the `obj` argument of `lgb.train`.
+- contents:
+  - make_objective_mse_cov
 - subtitle: LightGBM hyperparameters
-  desc:: >
-    LightGBM dials:: paramter functions that can be used with tune_ functions.
+  desc: >
+    LightGBM dials:: parameter functions that can be used with tune_ functions.
     Pass them to set_engine to use.
 - contents:
   - param_lgbm

man/axe_recipe.Rd

@@ -80,6 +80,44 @@ test_that("lightgbm alternate objective", {
   expect_equal(info$objective, "huber")
 })
 
+test_that("lightgbm mse_cov custom objective", {
+  skip_if_not_installed("lightgbm")
+
+  fit_with_rho <- function(rho) {
+    parsnip::boost_tree(trees = 50) %>%
+      parsnip::set_engine(
+        engine = "lightgbm",
+        objective = "mse_cov", mse_cov_rho = rho, verbose = -1,
+        max_depth = 15, feature_fraction = 1, min_data_in_leaf = 1
+      ) %>%
+      parsnip::set_mode("regression") %>%
+      parsnip::fit(mpg ~ ., data = mtcars)
+  }
+
+  # predict.lgb.Booster warns that class/response prediction types are
+  # unsupported for custom objectives; numeric predictions still work
+  lgb_fit <- fit_with_rho(0.5)
+  pred <- suppressWarnings(predict(lgb_fit, mtcars[, -1]))
+
+  expect_equal(nrow(pred), nrow(mtcars))
+  expect_true(all(is.finite(pred$.pred)))
+  expect_not_constant_predictions(pred$.pred)
+
+  # The penalty weight should reach the objective: different rho values
+  # must produce different fits
+  pred_plain <- suppressWarnings(predict(fit_with_rho(0), mtcars[, -1]))
+  expect_all_preds_differ(list(pred$.pred, pred_plain$.pred))
+})
+
+test_that("lightgbm mse_cov without mse_cov_rho throws error", {
+  expect_error(
+    parsnip::boost_tree(trees = 10, mode = "regression") %>%
+      parsnip::set_engine("lightgbm", objective = "mse_cov", verbose = -1) %>%
+      parsnip::fit(mpg ~ ., data = mtcars),
+    regexp = "requires `mse_cov_rho`"
+  )
+})
+
 test_that("lightgbm with save_tree_error saves record_evals", {
   model_fit <- parsnip::boost_tree(trees = 50) %>%
     parsnip::set_engine(