onnx-export-rs 0.1.1

use crate::{Error, Result};

/// A recursive tree accepted by [`flatten_tree`].
#[derive(Clone, Debug, PartialEq)]
pub enum RecursiveNode {
    /// A split taking the left child when `feature <= threshold`.
    Branch {
        /// Zero-based input feature.
        feature: usize,
        /// Split threshold.
        threshold: f64,
        /// True/left child.
        left: Box<Self>,
        /// False/right child.
        right: Box<Self>,
    },
    /// A split taking the left child when `feature < threshold`.
    BranchLessThan {
        /// Zero-based input feature.
        feature: usize,
        /// Split threshold.
        threshold: f64,
        /// True/left child.
        left: Box<Self>,
        /// False/right child.
        right: Box<Self>,
    },
    /// A leaf with one value per target/class score.
    Leaf(Vec<f64>),
}

/// A flat tree node with sequential ONNX node IDs.
#[derive(Clone, Debug, PartialEq)]
pub struct TreeNode {
    /// Sequential node ID.
    pub id: i64,
    /// Split feature; zero for a leaf.
    pub feature_id: i64,
    /// Split threshold; zero for a leaf.
    pub threshold: f32,
    /// True-child ID; zero for a leaf.
    pub true_child_id: i64,
    /// False-child ID; zero for a leaf.
    pub false_child_id: i64,
    /// Comparison used by a branch; ignored for leaves.
    pub branch_mode: BranchMode,
    /// Leaf values. Empty for a branch.
    pub leaf_values: Vec<f32>,
}

/// Comparison performed by a tree branch.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum BranchMode {
    /// Take the true child when the feature is at most the threshold.
    #[default]
    LessOrEqual,
    /// Take the true child when the feature is strictly below the threshold.
    LessThan,
}

impl BranchMode {
    pub(crate) const fn as_onnx(self) -> &'static [u8] {
        match self {
            Self::LessOrEqual => b"BRANCH_LEQ",
            Self::LessThan => b"BRANCH_LT",
        }
    }
}

impl TreeNode {
    /// Returns whether this node is a leaf.
    #[must_use]
    pub fn is_leaf(&self) -> bool {
        !self.leaf_values.is_empty()
    }
}

/// A single flat decision tree.
#[derive(Clone, Debug, PartialEq)]
pub struct TreeStructure {
    /// Nodes ordered by their sequential IDs.
    pub nodes: Vec<TreeNode>,
}

/// Forest score aggregation.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AggregationMode {
    /// Sum tree outputs.
    Sum,
    /// Average tree outputs.
    Average,
    /// Pick the minimum output.
    Min,
    /// Pick the maximum output.
    Max,
}

impl AggregationMode {
    pub(crate) const fn as_onnx(self) -> &'static str {
        match self {
            Self::Sum => "SUM",
            Self::Average => "AVERAGE",
            Self::Min => "MIN",
            Self::Max => "MAX",
        }
    }
}

/// One or more trees and their aggregation behavior.
#[derive(Clone, Debug, PartialEq)]
pub struct ForestStructure {
    /// Constituent trees.
    pub trees: Vec<TreeStructure>,
    /// Aggregation behavior.
    pub aggregation: AggregationMode,
    /// Number of regression targets or class scores per leaf.
    pub n_targets: usize,
}

/// Type of prediction produced by a tree ensemble.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TreeTask {
    /// Continuous score output.
    Regression,
    /// Class-score output. Consumers may apply `ArgMax` for labels.
    Classification,
}

/// Score transformation applied after tree aggregation.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum PostTransform {
    /// Return raw aggregate scores.
    #[default]
    None,
    /// Apply a logistic transform.
    Logistic,
    /// Apply softmax across targets.
    Softmax,
}

impl PostTransform {
    pub(crate) const fn as_onnx(self) -> &'static [u8] {
        match self {
            Self::None => b"NONE",
            Self::Logistic => b"LOGISTIC",
            Self::Softmax => b"SOFTMAX",
        }
    }
}

/// A gradient-boosted tree ensemble.
#[derive(Clone, Debug, PartialEq)]
pub struct GradientBoostedEnsemble {
    /// Trees whose leaf contributions are summed.
    pub trees: Vec<TreeStructure>,
    /// Initial score for every target.
    pub base_values: Vec<f64>,
    /// Multiplier applied to every tree leaf.
    pub learning_rate: f64,
    /// Number of regression targets or class scores.
    pub n_targets: usize,
    /// Output behavior.
    pub task: TreeTask,
    /// Transformation applied after summing base and tree scores.
    pub post_transform: PostTransform,
}

/// Flattens a recursive tree in preorder with sequential node IDs.
pub fn flatten_tree(root: &RecursiveNode) -> Result<TreeStructure> {
    fn visit(node: &RecursiveNode, output: &mut Vec<TreeNode>) -> Result<i64> {
        let id = i64::try_from(output.len())
            .map_err(|_| Error::InvalidModel("tree has too many nodes".into()))?;
        output.push(TreeNode {
            id,
            feature_id: 0,
            threshold: 0.0,
            true_child_id: 0,
            false_child_id: 0,
            branch_mode: BranchMode::LessOrEqual,
            leaf_values: Vec::new(),
        });
        match node {
            RecursiveNode::Leaf(values) => {
                if values.is_empty() {
                    return Err(Error::InvalidModel("tree leaf has no values".into()));
                }
                output[id as usize].leaf_values = values.iter().map(|&v| v as f32).collect();
            }
            RecursiveNode::Branch {
                feature,
                threshold,
                left,
                right,
            }
            | RecursiveNode::BranchLessThan {
                feature,
                threshold,
                left,
                right,
            } => {
                let true_id = visit(left, output)?;
                let false_id = visit(right, output)?;
                let flat = &mut output[id as usize];
                flat.feature_id = i64::try_from(*feature)
                    .map_err(|_| Error::InvalidModel("feature index is too large".into()))?;
                flat.threshold = *threshold as f32;
                flat.true_child_id = true_id;
                flat.false_child_id = false_id;
                flat.branch_mode = if matches!(node, RecursiveNode::BranchLessThan { .. }) {
                    BranchMode::LessThan
                } else {
                    BranchMode::LessOrEqual
                };
            }
        }
        Ok(id)
    }

    let mut nodes = Vec::new();
    visit(root, &mut nodes)?;
    Ok(TreeStructure { nodes })
}