shap_rs/tree/

model.rs

use crate::{Predict, Result, ShapError};
use ndarray::{Array1, Array2, ArrayView1, ArrayView2};

#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum MissingBranch {
    Left,
    Right,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum SplitComparison {
    LessThan,
    LessThanOrEqual,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum MissingValuePolicy {
    NaN,
    Zero,
    None,
}

/// A node in a binary regression tree. `cover` is the training weight reaching
/// the node and is used to integrate out features that are not observed.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum Node {
    Leaf {
        values: Vec<f64>,
        cover: f64,
    },
    Split {
        feature: usize,
        threshold: f64,
        left: usize,
        right: usize,
        missing: MissingBranch,
        cover: f64,
    },
    NumericalSplit {
        feature: usize,
        threshold: f64,
        comparison: SplitComparison,
        left: usize,
        right: usize,
        missing: MissingBranch,
        missing_value: MissingValuePolicy,
        cover: f64,
    },
    CategoricalSplit {
        feature: usize,
        categories: Vec<i64>,
        left: usize,
        right: usize,
        missing: MissingBranch,
        missing_value: MissingValuePolicy,
        cover: f64,
    },
}
/// Framework-neutral columnar representation used by model adapters.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TreeArrays {
    pub features: Vec<Option<usize>>,
    pub thresholds: Vec<f64>,
    pub left_children: Vec<Option<usize>>,
    pub right_children: Vec<Option<usize>>,
    pub missing: Vec<MissingBranch>,
    pub leaf_values: Vec<Option<Vec<f64>>>,
    pub covers: Vec<f64>,
    pub root: usize,
    pub n_features: usize,
}
impl Node {
    pub fn cover(&self) -> f64 {
        match self {
            Self::Leaf { cover, .. }
            | Self::Split { cover, .. }
            | Self::NumericalSplit { cover, .. }
            | Self::CategoricalSplit { cover, .. } => *cover,
        }
    }

    pub fn split_feature(&self) -> Option<usize> {
        match self {
            Self::Leaf { .. } => None,
            Self::Split { feature, .. }
            | Self::NumericalSplit { feature, .. }
            | Self::CategoricalSplit { feature, .. } => Some(*feature),
        }
    }

    pub fn children(&self) -> Option<(usize, usize)> {
        match self {
            Self::Leaf { .. } => None,
            Self::Split { left, right, .. }
            | Self::NumericalSplit { left, right, .. }
            | Self::CategoricalSplit { left, right, .. } => Some((*left, *right)),
        }
    }

    pub(crate) fn hot_child(&self, value: f64) -> Option<usize> {
        fn is_missing(value: f64, policy: MissingValuePolicy) -> bool {
            match policy {
                MissingValuePolicy::NaN => value.is_nan(),
                MissingValuePolicy::Zero => value.is_nan() || value == 0.0,
                MissingValuePolicy::None => false,
            }
        }
        fn missing_child(branch: MissingBranch, left: usize, right: usize) -> usize {
            match branch {
                MissingBranch::Left => left,
                MissingBranch::Right => right,
            }
        }
        match self {
            Self::Leaf { .. } => None,
            Self::Split {
                threshold,
                left,
                right,
                missing,
                ..
            } => Some(if value.is_nan() {
                missing_child(*missing, *left, *right)
            } else if value <= *threshold {
                *left
            } else {
                *right
            }),
            Self::NumericalSplit {
                threshold,
                comparison,
                left,
                right,
                missing,
                missing_value,
                ..
            } => Some(if is_missing(value, *missing_value) {
                missing_child(*missing, *left, *right)
            } else if match comparison {
                SplitComparison::LessThan => value < *threshold,
                SplitComparison::LessThanOrEqual => value <= *threshold,
            } {
                *left
            } else {
                *right
            }),
            Self::CategoricalSplit {
                categories,
                left,
                right,
                missing,
                missing_value,
                ..
            } => Some(if is_missing(value, *missing_value) {
                missing_child(*missing, *left, *right)
            } else if value.is_finite()
                && value.fract() == 0.0
                && value >= i64::MIN as f64
                && value <= i64::MAX as f64
                && categories.contains(&(value as i64))
            {
                *left
            } else {
                *right
            }),
        }
    }
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct Tree {
    nodes: Vec<Node>,
    root: usize,
    n_features: usize,
    n_outputs: usize,
}
#[derive(serde::Deserialize)]
struct TreePayload {
    nodes: Vec<Node>,
    root: usize,
    n_features: usize,
    n_outputs: usize,
}
impl<'de> serde::Deserialize<'de> for Tree {
    fn deserialize<D: serde::Deserializer<'de>>(
        deserializer: D,
    ) -> std::result::Result<Self, D::Error> {
        let payload = <TreePayload as serde::Deserialize>::deserialize(deserializer)?;
        let tree = Self::new(payload.nodes, payload.root, payload.n_features)
            .map_err(serde::de::Error::custom)?;
        if tree.n_outputs() != payload.n_outputs {
            return Err(serde::de::Error::custom(
                "serialized tree output count does not match its leaves",
            ));
        }
        Ok(tree)
    }
}
impl Tree {
    /// Revalidates a tree after deserialization or adapter conversion.
    pub fn validate(&self) -> Result<()> {
        Self::new(self.nodes.clone(), self.root, self.n_features).map(|_| ())
    }
    pub fn from_arrays(a: TreeArrays) -> Result<Self> {
        let n = a.features.len();
        if [
            a.thresholds.len(),
            a.left_children.len(),
            a.right_children.len(),
            a.missing.len(),
            a.leaf_values.len(),
            a.covers.len(),
        ]
        .iter()
        .any(|&x| x != n)
        {
            return Err(ShapError::DimensionMismatch {
                expected: format!("{n} entries in every tree array"),
                found: "inconsistent tree array lengths".into(),
            });
        }
        let mut nodes = Vec::with_capacity(n);
        for i in 0..n {
            match (&a.features[i], &a.leaf_values[i]) {
                (None, Some(values)) => nodes.push(Node::Leaf {
                    values: values.clone(),
                    cover: a.covers[i],
                }),
                (Some(feature), None) => nodes.push(Node::Split {
                    feature: *feature,
                    threshold: a.thresholds[i],
                    left: a.left_children[i].ok_or_else(|| {
                        ShapError::InvalidConfiguration(format!("split {i} has no left child"))
                    })?,
                    right: a.right_children[i].ok_or_else(|| {
                        ShapError::InvalidConfiguration(format!("split {i} has no right child"))
                    })?,
                    missing: a.missing[i],
                    cover: a.covers[i],
                }),
                _ => {
                    return Err(ShapError::InvalidConfiguration(format!(
                        "node {i} must be exactly one of leaf or split"
                    )))
                }
            }
        }
        Self::new(nodes, a.root, a.n_features)
    }
    pub fn new(nodes: Vec<Node>, root: usize, n_features: usize) -> Result<Self> {
        if nodes.is_empty() || root >= nodes.len() {
            return Err(ShapError::InvalidConfiguration(
                "tree must have a valid root".into(),
            ));
        }
        if n_features == 0 {
            return Err(ShapError::InvalidConfiguration(
                "tree must have at least one feature".into(),
            ));
        }
        let mut outputs = None;
        for (i, node) in nodes.iter().enumerate() {
            if !node.cover().is_finite() || node.cover() < 0.0 {
                return Err(ShapError::InvalidConfiguration(format!(
                    "node {i} has invalid cover"
                )));
            }
            match node {
                Node::Leaf { values, .. } => {
                    if values.is_empty() || values.iter().any(|v| !v.is_finite()) {
                        return Err(ShapError::InvalidConfiguration(format!(
                            "leaf {i} has invalid values"
                        )));
                    }
                    if outputs
                        .replace(values.len())
                        .is_some_and(|n| n != values.len())
                    {
                        return Err(ShapError::InvalidConfiguration(
                            "all leaves must have the same output count".into(),
                        ));
                    }
                }
                Node::Split {
                    feature,
                    left,
                    right,
                    threshold,
                    ..
                } => {
                    if *feature >= n_features
                        || *left >= nodes.len()
                        || *right >= nodes.len()
                        || !threshold.is_finite()
                    {
                        return Err(ShapError::InvalidConfiguration(format!(
                            "split node {i} is invalid"
                        )));
                    }
                }
                Node::NumericalSplit {
                    feature,
                    left,
                    right,
                    threshold,
                    ..
                } => {
                    if *feature >= n_features
                        || *left >= nodes.len()
                        || *right >= nodes.len()
                        || !threshold.is_finite()
                    {
                        return Err(ShapError::InvalidConfiguration(format!(
                            "numerical split node {i} is invalid"
                        )));
                    }
                }
                Node::CategoricalSplit {
                    feature,
                    categories,
                    left,
                    right,
                    ..
                } => {
                    if *feature >= n_features
                        || *left >= nodes.len()
                        || *right >= nodes.len()
                        || categories.is_empty()
                    {
                        return Err(ShapError::InvalidConfiguration(format!(
                            "categorical split node {i} is invalid"
                        )));
                    }
                }
            }
        }
        let n_outputs =
            outputs.ok_or_else(|| ShapError::InvalidConfiguration("tree has no leaves".into()))?;
        let tree = Self {
            nodes,
            root,
            n_features,
            n_outputs,
        };
        tree.validate_graph()?;
        Ok(tree)
    }
    fn validate_graph(&self) -> Result<()> {
        fn visit(t: &Tree, i: usize, state: &mut [u8]) -> Result<()> {
            if state[i] == 1 {
                return Err(ShapError::InvalidConfiguration(
                    "tree contains a cycle".into(),
                ));
            }
            if state[i] == 2 {
                return Err(ShapError::InvalidConfiguration(
                    "tree node has multiple parents".into(),
                ));
            }
            state[i] = 1;
            if let Some((left, right)) = t.nodes[i].children() {
                visit(t, left, state)?;
                visit(t, right, state)?;
            }
            state[i] = 2;
            Ok(())
        }
        let mut state = vec![0; self.nodes.len()];
        visit(self, self.root, &mut state)?;
        if state.contains(&0) {
            return Err(ShapError::InvalidConfiguration(
                "tree contains nodes unreachable from the root".into(),
            ));
        }
        Ok(())
    }
    pub fn nodes(&self) -> &[Node] {
        &self.nodes
    }
    pub fn root(&self) -> usize {
        self.root
    }
    pub fn n_features(&self) -> usize {
        self.n_features
    }
    pub fn n_outputs(&self) -> usize {
        self.n_outputs
    }
    pub fn predict_row(&self, x: ArrayView1<'_, f64>) -> Result<&[f64]> {
        if x.len() != self.n_features {
            return Err(ShapError::DimensionMismatch {
                expected: format!("{} features", self.n_features),
                found: format!("{}", x.len()),
            });
        }
        let mut i = self.root;
        loop {
            match &self.nodes[i] {
                Node::Leaf { values, .. } => return Ok(values),
                node => i = node.hot_child(x[node.split_feature().unwrap()]).unwrap(),
            }
        }
    }
    pub fn expected_value(&self) -> Vec<f64> {
        fn rec(t: &Tree, i: usize) -> Vec<f64> {
            match &t.nodes[i] {
                Node::Leaf { values, .. } => values.clone(),
                node => {
                    let (left, right) = node.children().unwrap();
                    let a = rec(t, left);
                    let b = rec(t, right);
                    let total = t.nodes[left].cover() + t.nodes[right].cover();
                    let p = if total > 0.0 {
                        t.nodes[left].cover() / total
                    } else {
                        0.5
                    };
                    a.into_iter()
                        .zip(b)
                        .map(|(x, y)| p * x + (1.0 - p) * y)
                        .collect()
                }
            }
        }
        rec(self, self.root)
    }
    #[cfg(test)]
    pub(crate) fn conditional_value(&self, x: ArrayView1<'_, f64>, present: &[bool]) -> Vec<f64> {
        fn rec(t: &Tree, i: usize, x: ArrayView1<'_, f64>, p: &[bool]) -> Vec<f64> {
            match &t.nodes[i] {
                Node::Leaf { values, .. } => values.clone(),
                node => {
                    let feature = node.split_feature().unwrap();
                    let (left, right) = node.children().unwrap();
                    if p[feature] {
                        let c = node.hot_child(x[feature]).unwrap();
                        rec(t, c, x, p)
                    } else {
                        let a = rec(t, left, x, p);
                        let b = rec(t, right, x, p);
                        let total = t.nodes[left].cover() + t.nodes[right].cover();
                        let q = if total > 0.0 {
                            t.nodes[left].cover() / total
                        } else {
                            0.5
                        };
                        a.into_iter()
                            .zip(b)
                            .map(|(u, v)| q * u + (1.0 - q) * v)
                            .collect()
                    }
                }
            }
        }
        rec(self, self.root, x, present)
    }
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct TreeEnsemble {
    trees: Vec<(Tree, f64)>,
    output_groups: Vec<Option<usize>>,
    base_values: Array1<f64>,
    n_features: usize,
}
#[derive(serde::Deserialize)]
struct TreeEnsemblePayload {
    trees: Vec<(Tree, f64)>,
    #[serde(default)]
    output_groups: Vec<Option<usize>>,
    base_values: Array1<f64>,
    n_features: usize,
}
impl<'de> serde::Deserialize<'de> for TreeEnsemble {
    fn deserialize<D: serde::Deserializer<'de>>(
        deserializer: D,
    ) -> std::result::Result<Self, D::Error> {
        let payload = <TreeEnsemblePayload as serde::Deserialize>::deserialize(deserializer)?;
        let ensemble = if payload.output_groups.is_empty() {
            Self::new(payload.trees, payload.base_values.to_vec())
        } else {
            Self::new_with_output_groups(
                payload.trees,
                payload.base_values.to_vec(),
                payload.output_groups,
            )
        }
        .map_err(serde::de::Error::custom)?;
        if ensemble.n_features() != payload.n_features {
            return Err(serde::de::Error::custom(
                "serialized ensemble feature count does not match its trees",
            ));
        }
        Ok(ensemble)
    }
}
impl TreeEnsemble {
    pub fn new(trees: Vec<(Tree, f64)>, base_values: Vec<f64>) -> Result<Self> {
        let groups = vec![None; trees.len()];
        Self::new_with_output_groups(trees, base_values, groups)
    }

    pub fn new_with_output_groups(
        trees: Vec<(Tree, f64)>,
        base_values: Vec<f64>,
        output_groups: Vec<Option<usize>>,
    ) -> Result<Self> {
        if trees.is_empty() {
            return Err(ShapError::InvalidConfiguration(
                "ensemble must contain a tree".into(),
            ));
        }
        let nf = trees[0].0.n_features();
        let no = trees[0].0.n_outputs();
        if base_values.len() != no
            || output_groups.len() != trees.len()
            || output_groups.iter().flatten().any(|&group| group >= no)
            || base_values.iter().any(|v| !v.is_finite())
            || trees
                .iter()
                .any(|(t, w)| t.n_features() != nf || t.n_outputs() != no || !w.is_finite())
        {
            return Err(ShapError::DimensionMismatch {
                expected: format!("trees with {nf} features and {no} outputs"),
                found: "inconsistent ensemble".into(),
            });
        }
        Ok(Self {
            trees,
            output_groups,
            base_values: Array1::from(base_values),
            n_features: nf,
        })
    }
    pub fn trees(&self) -> &[(Tree, f64)] {
        &self.trees
    }
    pub fn output_groups(&self) -> &[Option<usize>] {
        &self.output_groups
    }
    pub fn base_offset(&self) -> ndarray::ArrayView1<'_, f64> {
        self.base_values.view()
    }
    pub fn n_features(&self) -> usize {
        self.n_features
    }
    pub fn n_outputs(&self) -> usize {
        self.base_values.len()
    }
    pub fn expected_value(&self) -> Array1<f64> {
        let mut v = self.base_values.clone();
        for (t, w) in &self.trees {
            for (o, x) in t.expected_value().into_iter().enumerate() {
                v[o] += w * x
            }
        }
        v
    }
    /// Predicts raw tree outputs with a per-sample base margin replacing the
    /// ensemble's fixed base offset (matching XGBoost `base_margin` semantics).
    pub fn predict_with_base_margin(
        &self,
        x: ArrayView2<'_, f64>,
        base_margin: ArrayView2<'_, f64>,
    ) -> Result<Array2<f64>> {
        if base_margin.dim() != (x.nrows(), self.n_outputs()) {
            return Err(ShapError::DimensionMismatch {
                expected: format!("({}, {}) base margins", x.nrows(), self.n_outputs()),
                found: format!("{:?}", base_margin.dim()),
            });
        }
        if base_margin.iter().any(|value| !value.is_finite()) {
            return Err(ShapError::InvalidConfiguration(
                "base margins must be finite".into(),
            ));
        }
        let mut prediction = self.predict(x)?;
        for row in 0..prediction.nrows() {
            for output in 0..prediction.ncols() {
                prediction[[row, output]] += base_margin[[row, output]] - self.base_values[output];
            }
        }
        Ok(prediction)
    }
    /// Revalidates all trees and ensemble dimensions after deserialization.
    pub fn validate(&self) -> Result<()> {
        for (t, w) in &self.trees {
            t.validate()?;
            if !w.is_finite() {
                return Err(ShapError::InvalidConfiguration(
                    "tree weight must be finite".into(),
                ));
            }
        }
        Self::new_with_output_groups(
            self.trees.clone(),
            self.base_values.to_vec(),
            self.output_groups.clone(),
        )
        .map(|_| ())
    }
}

impl Predict for TreeEnsemble {
    fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array2<f64>> {
        self.validate()?;
        if x.ncols() != self.n_features {
            return Err(ShapError::DimensionMismatch {
                expected: format!("{} features", self.n_features),
                found: format!("{}", x.ncols()),
            });
        }
        crate::error::checked_f64_shape(&[x.nrows(), self.base_values.len()], "tree prediction")?;
        let mut out = Array2::from_shape_fn((x.nrows(), self.base_values.len()), |(_, o)| {
            self.base_values[o]
        });
        for i in 0..x.nrows() {
            for (t, w) in &self.trees {
                for (o, v) in t.predict_row(x.row(i))?.iter().enumerate() {
                    out[[i, o]] += w * v
                }
            }
        }
        Ok(out)
    }
    fn n_features(&self) -> Option<usize> {
        Some(self.n_features)
    }
    fn n_outputs(&self) -> Option<usize> {
        Some(self.base_values.len())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn rejects_unreachable_and_shared_nodes() {
        let leaf = || Node::Leaf {
            values: vec![0.],
            cover: 1.,
        };
        let unreachable = Tree::new(
            vec![
                Node::Split {
                    feature: 0,
                    threshold: 0.,
                    left: 1,
                    right: 2,
                    missing: MissingBranch::Left,
                    cover: 2.,
                },
                leaf(),
                leaf(),
                leaf(),
            ],
            0,
            1,
        );
        assert!(unreachable.is_err());

        let shared = Tree::new(
            vec![
                Node::Split {
                    feature: 0,
                    threshold: 0.,
                    left: 1,
                    right: 1,
                    missing: MissingBranch::Left,
                    cover: 2.,
                },
                leaf(),
            ],
            0,
            1,
        );
        assert!(shared.is_err());
    }

    #[test]
    fn rejects_non_finite_ensemble_base_values() {
        let tree = Tree::new(
            vec![Node::Leaf {
                values: vec![1.],
                cover: 1.,
            }],
            0,
            1,
        )
        .unwrap();
        assert!(TreeEnsemble::new(vec![(tree, 1.)], vec![f64::NAN]).is_err());
    }
}