Support binary:logistic objective and probability predictions

src/lib.rs

@@ -8,14 +8,14 @@
 //!
 //! - dense in-memory `f64` input
 //! - tree-ensemble prediction
-//! - official upstream `model.json` loading for a narrow regression subset
+//! - official upstream `model.json` loading for a narrow inference subset
 //! - explicit tree construction for tests and adapters
 //!
 //! Current official model support is intentionally narrow:
 //!
 //! - `booster=gbtree`
-//! - `objective=reg:squarederror`
-//! - single-target regression
+//! - `objective=reg:squarederror` or `objective=binary:logistic`
+//! - single-target prediction
 //! - numerical splits only
 //!
 //! # Example

src/model.rs

@@ -8,6 +8,12 @@ use crate::official_model;
 use crate::predict;
 use crate::tree::RegressionTree;
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum PredictionTask {
+    Regression,
+    BinaryLogistic,
+}
+
 /// Inference-only gradient-boosted tree model.
 ///
 /// This type is the stable public model surface. Callers can either load a
@@ -16,8 +22,9 @@ use crate::tree::RegressionTree;
 #[derive(Debug, Clone, PartialEq)]
 pub struct XGBModel {
     trees: Vec<RegressionTree>,
-    base_score: f64,
+    base_margin: f64,
     n_features: usize,
+    task: PredictionTask,
 }
 
 impl XGBModel {
@@ -28,7 +35,16 @@ impl XGBModel {
     /// Returns [`XGBError::InvalidParameter`] if `base_score` is not finite or
     /// if `n_features == 0`.
     pub fn new(base_score: f64, n_features: usize, trees: Vec<RegressionTree>) -> Result<Self> {
-        if !base_score.is_finite() {
+        Self::from_parts(PredictionTask::Regression, base_score, n_features, trees)
+    }
+
+    pub(crate) fn from_parts(
+        task: PredictionTask,
+        base_margin: f64,
+        n_features: usize,
+        trees: Vec<RegressionTree>,
+    ) -> Result<Self> {
+        if !base_margin.is_finite() {
             return Err(XGBError::InvalidParameter {
                 name: "base_score",
                 reason: "must be finite",
@@ -44,15 +60,19 @@ impl XGBModel {
 
         Ok(Self {
             trees,
-            base_score,
+            base_margin,
             n_features,
+            task,
         })
     }
 
-    /// Return the global bias added before tree contributions.
+    /// Return the serialized `XGBoost` `base_score` for supported objectives.
     #[must_use]
     pub fn base_score(&self) -> f64 {
-        self.base_score
+        match self.task {
+            PredictionTask::Regression => self.base_margin,
+            PredictionTask::BinaryLogistic => predict::sigmoid(self.base_margin),
+        }
     }
 
     /// Return the number of feature columns expected by this model.
@@ -67,23 +87,34 @@ impl XGBModel {
         &self.trees
     }
 
-    /// Predict regression values for a dense feature matrix.
+    /// Predict task outputs for a dense feature matrix.
+    ///
+    /// For supported official JSON models this returns:
+    ///
+    /// - regression values for `reg:squarederror`
+    /// - positive-class probabilities for `binary:logistic`
     ///
     /// # Errors
     ///
     /// Returns [`XGBError::FeatureCountMismatch`] if the feature count differs
     /// from the model expectation.
     pub fn predict_dense(&self, features: &DenseMatrix) -> Result<Vec<f64>> {
-        predict::predict_ensemble(self.base_score, &self.trees, features, self.n_features)
+        predict::predict_ensemble(
+            self.task,
+            self.base_margin,
+            &self.trees,
+            features,
+            self.n_features,
+        )
     }
 
     /// Load an official upstream `XGBoost` `model.json` file.
     ///
     /// Currently supported model scope:
     ///
     /// - `booster=gbtree`
-    /// - `objective=reg:squarederror`
-    /// - single-target regression
+    /// - `objective=reg:squarederror` or `objective=binary:logistic`
+    /// - single-target prediction
     /// - numerical splits only
     ///
     /// # Errors

src/official_model/loader.rs

@@ -2,7 +2,7 @@ use std::fs;
 use std::path::Path;
 
 use crate::error::{Result, XGBError};
-use crate::model::XGBModel;
+use crate::model::{PredictionTask, XGBModel};
 use crate::tree::{RegressionTree, TreeNode};
 
 use super::schema::{JsonModel, JsonTree};
@@ -21,12 +21,7 @@ pub(super) fn convert_model(model: JsonModel) -> Result<XGBModel> {
         });
     }
 
-    if model.learner.objective.name != "reg:squarederror" {
-        return Err(XGBError::UnsupportedModel {
-            context: "objective",
-            value: model.learner.objective.name,
-        });
-    }
+    let task = parse_prediction_task(&model.learner.objective.name)?;
 
     let num_target = parse_usize_field(
         &model.learner.model_param.num_target,
@@ -61,7 +56,7 @@ pub(super) fn convert_model(model: JsonModel) -> Result<XGBModel> {
         });
     }
 
-    let base_score = parse_base_score(&model.learner.model_param.base_score)?;
+    let base_margin = parse_base_score(&model.learner.model_param.base_score, task)?;
     let n_features = parse_usize_field(
         &model.learner.model_param.num_feature,
         "learner_model_param.num_feature",
@@ -76,7 +71,18 @@ pub(super) fn convert_model(model: JsonModel) -> Result<XGBModel> {
         .map(|(tree_idx, tree)| convert_tree(tree_idx, tree, n_features))
         .collect::<Result<Vec<_>>>()?;
 
-    XGBModel::new(base_score, n_features, trees)
+    XGBModel::from_parts(task, base_margin, n_features, trees)
+}
+
+fn parse_prediction_task(objective: &str) -> Result<PredictionTask> {
+    match objective {
+        "reg:squarederror" => Ok(PredictionTask::Regression),
+        "binary:logistic" => Ok(PredictionTask::BinaryLogistic),
+        _ => Err(XGBError::UnsupportedModel {
+            context: "objective",
+            value: objective.to_owned(),
+        }),
+    }
 }
 
 fn convert_tree(
@@ -168,14 +174,14 @@ fn convert_node(
     })
 }
 
-fn parse_base_score(raw: &str) -> Result<f64> {
+fn parse_base_score(raw: &str, task: PredictionTask) -> Result<f64> {
     if let Ok(value) = raw.parse::<f64>() {
-        return Ok(xgb_f32(value));
+        return parse_base_score_value(value, task);
     }
 
     if let Ok(values) = serde_json::from_str::<Vec<f64>>(raw) {
         if values.len() == 1 {
-            return Ok(xgb_f32(values[0]));
+            return parse_base_score_value(values[0], task);
         }
     }
 
@@ -184,6 +190,14 @@ fn parse_base_score(raw: &str) -> Result<f64> {
     ))
 }
 
+fn parse_base_score_value(value: f64, task: PredictionTask) -> Result<f64> {
+    let base_score = xgb_f32(value);
+    match task {
+        PredictionTask::Regression => Ok(base_score),
+        PredictionTask::BinaryLogistic => logit(base_score),
+    }
+}
+
 fn parse_usize_field(raw: &str, context: &'static str) -> Result<usize> {
     raw.parse::<usize>()
         .map_err(|_| XGBError::InvalidModelFormat(context))
@@ -235,6 +249,8 @@ fn validate_reachable_subtree(
     visit_state: &mut [VisitState],
     indegree: &mut [usize],
 ) -> Result<()> {
+    // TODO: Replace this recursive walk with an explicit stack so valid but
+    // highly unbalanced trees cannot overflow the loader's call stack.
     match visit_state[node_idx] {
         VisitState::Unvisited => {}
         VisitState::Visiting => {
@@ -312,3 +328,13 @@ fn xgb_f32(value: f64) -> f64 {
     )]
     f64::from(value as f32)
 }
+
+fn logit(probability: f64) -> Result<f64> {
+    if !(0.0..1.0).contains(&probability) {
+        return Err(XGBError::InvalidModelFormat(
+            "binary:logistic base_score must be strictly between 0 and 1",
+        ));
+    }
+
+    Ok((probability / (1.0 - probability)).ln())
+}

src/predict.rs

@@ -2,18 +2,21 @@
 
 use crate::dataset::DenseMatrix;
 use crate::error::{Result, XGBError};
+use crate::model::PredictionTask;
 use crate::tree::RegressionTree;
 
 /// Predict a batch of rows using a fitted tree ensemble.
 ///
-/// Predictions start from `base_score` and then add the contribution of each tree.
+/// Predictions start from `base_margin`, then add the contribution of each tree,
+/// and finally apply the supported task-specific output transform.
 ///
 /// # Errors
 ///
 /// Returns [`XGBError::FeatureCountMismatch`] if `features` does not match the
 /// expected number of feature columns.
 pub fn predict_ensemble(
-    base_score: f64,
+    task: PredictionTask,
+    base_margin: f64,
     trees: &[RegressionTree],
     features: &DenseMatrix,
     expected_feature_count: usize,
@@ -25,11 +28,16 @@ pub fn predict_ensemble(
         });
     }
 
-    let mut predictions = vec![base_score; features.n_rows()];
+    let mut predictions = vec![base_margin; features.n_rows()];
     for (row_idx, prediction) in predictions.iter_mut().enumerate() {
         for tree in trees {
             *prediction += predict_tree(tree, features, row_idx);
         }
+
+        *prediction = match task {
+            PredictionTask::Regression => *prediction,
+            PredictionTask::BinaryLogistic => sigmoid(*prediction),
+        };
     }
 
     Ok(predictions)
@@ -83,3 +91,14 @@ pub fn predict_tree(tree: &RegressionTree, features: &DenseMatrix, row_idx: usiz
 fn xgb_less_than(feature_value: f64, split_value: f64) -> bool {
     (feature_value as f32) < (split_value as f32)
 }
+
+#[must_use]
+pub(crate) fn sigmoid(value: f64) -> f64 {
+    if value >= 0.0 {
+        let exp_neg = (-value).exp();
+        1.0 / (1.0 + exp_neg)
+    } else {
+        let exp_pos = value.exp();
+        exp_pos / (1.0 + exp_pos)
+    }
+}

tests/official_model_json.rs

@@ -40,16 +40,33 @@ fn loads_official_regressor_json_and_predicts() {
 }
 
 #[test]
-fn rejects_unsupported_official_objective() {
-    let error = XGBModel::load_json(fixture_path("classifier.json")).unwrap_err();
+fn loads_official_binary_classifier_json_and_predicts_probabilities() {
+    let model = XGBModel::load_json(fixture_path("classifier.json")).unwrap();
+    let features = DenseMatrix::from_shape_vec(
+        4,
+        3,
+        vec![
+            0.0, 0.0, 0.0, // left-left, left-right
+            0.0, 1.0, 0.0, // left-right, left-right
+            0.0, -2.0, 1.0, // right-left, left-left
+            0.0, 0.0, 1.0, // right-left, right-right
+        ],
+    )
+    .unwrap();
 
-    assert!(matches!(
-        error,
-        XGBError::UnsupportedModel {
-            context: "objective",
-            ..
-        }
-    ));
+    let predictions = model.predict_dense(&features).unwrap();
+
+    assert_vec_close(
+        &predictions,
+        &[0.469_360_77, 0.522_249_44, 0.553_298_23, 0.617_173_93],
+    );
+}
+
+#[test]
+fn binary_classifier_reports_probability_base_score() {
+    let model = XGBModel::load_json(fixture_path("classifier.json")).unwrap();
+
+    assert!((model.base_score() - 0.52).abs() < 1.0e-6);
 }
 
 #[test]