use std::cmp::Ordering;
use runmat_builtins::{
BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
CharArray, LogicalArray, ObjectInstance, ResolveContext, StringArray, StructValue, Tensor,
Type, Value,
};
use runmat_macros::runtime_builtin;
use crate::builtins::common::tensor;
use crate::builtins::table::{
is_tabular_object, table_variable_names_from_object, table_variables,
};
use crate::{build_runtime_error, gather_if_needed_async, BuiltinResult, RuntimeError};
pub(crate) const CLASSIFICATION_TREE_CLASS: &str = "ClassificationTree";
const FITCTREE_NAME: &str = "fitctree";
const PREDICT_NAME: &str = "predict";
const EPS: f64 = 1.0e-12;
const MAX_TREE_DEPTH: usize = 4096;
const OUTPUT_TREE: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
name: "Mdl",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Required,
default: None,
description: "ClassificationTree object containing class names, split rules, and posterior probabilities.",
}];
const PARAM_TBL_OR_X: BuiltinParamDescriptor = BuiltinParamDescriptor {
name: "tblOrX",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Required,
default: None,
description: "Input table or numeric predictor matrix.",
};
const PARAM_Y_OR_RESPONSE: BuiltinParamDescriptor = BuiltinParamDescriptor {
name: "yOrResponse",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Optional,
default: None,
description: "Response vector, response variable name, or table formula.",
};
const PARAM_REST: BuiltinParamDescriptor = BuiltinParamDescriptor {
name: "options",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Variadic,
default: None,
description: "Name-value options such as ClassNames, PredictorNames, ResponseName, MaxNumSplits, MinLeafSize, MinParentSize, SplitCriterion, and Weights.",
};
const FITCTREE_INPUTS_X_Y: [BuiltinParamDescriptor; 2] = [PARAM_TBL_OR_X, PARAM_Y_OR_RESPONSE];
const FITCTREE_INPUTS_FULL: [BuiltinParamDescriptor; 3] =
[PARAM_TBL_OR_X, PARAM_Y_OR_RESPONSE, PARAM_REST];
const FITCTREE_SIGNATURES: [BuiltinSignatureDescriptor; 4] = [
BuiltinSignatureDescriptor {
label: "Mdl = fitctree(Tbl, ResponseVarName)",
inputs: &FITCTREE_INPUTS_X_Y,
outputs: &OUTPUT_TREE,
},
BuiltinSignatureDescriptor {
label: "Mdl = fitctree(Tbl, formula)",
inputs: &FITCTREE_INPUTS_X_Y,
outputs: &OUTPUT_TREE,
},
BuiltinSignatureDescriptor {
label: "Mdl = fitctree(X, Y)",
inputs: &FITCTREE_INPUTS_X_Y,
outputs: &OUTPUT_TREE,
},
BuiltinSignatureDescriptor {
label: "Mdl = fitctree(___, Name, Value)",
inputs: &FITCTREE_INPUTS_FULL,
outputs: &OUTPUT_TREE,
},
];
const ERROR_FITCTREE_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.FITCTREE.INVALID_ARGUMENT",
identifier: Some("RunMat:fitctree:InvalidArgument"),
when:
"Inputs, response labels, dimensions, or name-value options are malformed or unsupported.",
message: "fitctree: invalid argument",
};
const ERROR_FITCTREE_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.FITCTREE.INTERNAL",
identifier: Some("RunMat:fitctree:Internal"),
when: "RunMat cannot construct the ClassificationTree result.",
message: "fitctree: internal error",
};
const ERROR_PREDICT_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.PREDICT.INVALID_ARGUMENT",
identifier: Some("RunMat:predict:InvalidArgument"),
when: "Classification-tree prediction inputs are malformed or incompatible with the model.",
message: "predict: invalid argument",
};
const ERROR_PREDICT_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.PREDICT.INTERNAL",
identifier: Some("RunMat:predict:Internal"),
when: "Classification-tree prediction cannot construct its output.",
message: "predict: internal error",
};
const FITCTREE_ERRORS: [BuiltinErrorDescriptor; 2] =
[ERROR_FITCTREE_INVALID_ARGUMENT, ERROR_FITCTREE_INTERNAL];
pub const FITCTREE_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
signatures: &FITCTREE_SIGNATURES,
output_mode: BuiltinOutputMode::Fixed,
completion_policy: BuiltinCompletionPolicy::Public,
errors: &FITCTREE_ERRORS,
};
fn fitctree_type(_args: &[Type], _ctx: &ResolveContext) -> Type {
Type::Unknown
}
fn fitctree_error(
message: impl Into<String>,
descriptor: &'static BuiltinErrorDescriptor,
) -> RuntimeError {
let mut builder = build_runtime_error(message).with_builtin(FITCTREE_NAME);
if let Some(identifier) = descriptor.identifier {
builder = builder.with_identifier(identifier);
}
builder.build()
}
fn fitctree_invalid(message: impl Into<String>) -> RuntimeError {
fitctree_error(message, &ERROR_FITCTREE_INVALID_ARGUMENT)
}
fn fitctree_internal(message: impl Into<String>) -> RuntimeError {
fitctree_error(message, &ERROR_FITCTREE_INTERNAL)
}
fn predict_invalid(message: impl Into<String>) -> RuntimeError {
predict_error(message, &ERROR_PREDICT_INVALID_ARGUMENT)
}
fn predict_internal(message: impl Into<String>) -> RuntimeError {
predict_error(message, &ERROR_PREDICT_INTERNAL)
}
fn predict_error(
message: impl Into<String>,
descriptor: &'static BuiltinErrorDescriptor,
) -> RuntimeError {
let mut builder = build_runtime_error(message).with_builtin(PREDICT_NAME);
if let Some(identifier) = descriptor.identifier {
builder = builder.with_identifier(identifier);
}
builder.build()
}
#[derive(Clone, Debug)]
struct FitOptions {
class_names: Option<Vec<ClassLabel>>,
predictor_names: Option<Vec<String>>,
response_name: Option<String>,
max_num_splits: Option<usize>,
min_leaf_size: usize,
min_parent_size: usize,
split_criterion: SplitCriterion,
weights: Option<WeightSpec>,
}
impl Default for FitOptions {
fn default() -> Self {
Self {
class_names: None,
predictor_names: None,
response_name: None,
max_num_splits: None,
min_leaf_size: 1,
min_parent_size: 10,
split_criterion: SplitCriterion::Gdi,
weights: None,
}
}
}
#[derive(Clone, Debug)]
enum WeightSpec {
Values(Vec<f64>),
Variable(String),
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum SplitCriterion {
Gdi,
Deviance,
}
#[derive(Clone, Debug, PartialEq)]
enum ClassLabel {
Numeric(f64),
Text(String),
Logical(bool),
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum LabelKind {
Numeric,
Text,
Logical,
}
#[derive(Clone, Debug)]
struct LabelVector {
labels: Vec<ClassLabel>,
kind: LabelKind,
}
#[derive(Clone, Debug)]
struct FitSpec {
x: Tensor,
class_indices: Vec<usize>,
class_names: Vec<ClassLabel>,
class_kind: LabelKind,
predictor_names: Vec<String>,
response_name: String,
weights: Vec<f64>,
options: FitOptions,
}
#[derive(Clone, Debug)]
struct Node {
left: Option<usize>,
right: Option<usize>,
predictor: Option<usize>,
threshold: f64,
class_index: usize,
probabilities: Vec<f64>,
}
#[derive(Clone, Debug)]
struct TreeFit {
nodes: Vec<Node>,
num_splits: usize,
}
#[derive(Clone, Debug)]
struct SplitCandidate {
predictor: usize,
threshold: f64,
gain: f64,
left_rows: Vec<usize>,
right_rows: Vec<usize>,
}
#[runtime_builtin(
name = "fitctree",
category = "stats/ml",
summary = "Fit a binary decision tree for multiclass classification.",
keywords = "fitctree,classification tree,decision tree,machine learning,statistics",
type_resolver(fitctree_type),
descriptor(crate::builtins::stats::ml::classification_tree::FITCTREE_DESCRIPTOR),
builtin_path = "crate::builtins::stats::ml::classification_tree"
)]
async fn fitctree_builtin(first: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
let first = gather(first)
.await
.map_err(|err| fitctree_invalid(err.message))?;
let rest = gather_all(rest)
.await
.map_err(|err| fitctree_invalid(err.message))?;
let spec = build_fit_spec(first, rest)?;
let fit = fit_tree(&spec)?;
model_object(spec, fit).map(Value::Object)
}
async fn gather(value: Value) -> BuiltinResult<Value> {
gather_if_needed_async(&value)
.await
.map_err(|err| fitctree_invalid(format!("fitctree: {err}")))
}
async fn gather_all(values: Vec<Value>) -> BuiltinResult<Vec<Value>> {
let mut out = Vec::with_capacity(values.len());
for value in values {
out.push(gather(value).await?);
}
Ok(out)
}
fn build_fit_spec(first: Value, mut rest: Vec<Value>) -> BuiltinResult<FitSpec> {
if is_table_value(&first) {
build_table_fit_spec(first, rest)
} else {
if rest.is_empty() {
return Err(fitctree_invalid(
"fitctree: matrix input requires response vector Y",
));
}
let y = rest.remove(0);
build_matrix_fit_spec(first, y, rest)
}
}
fn is_table_value(value: &Value) -> bool {
matches!(value, Value::Object(object) if is_tabular_object(object))
}
fn build_matrix_fit_spec(
x_value: Value,
y_value: Value,
rest: Vec<Value>,
) -> BuiltinResult<FitSpec> {
let x = tensor::value_into_tensor_for(FITCTREE_NAME, x_value)
.map_err(|err| fitctree_invalid(format!("fitctree: {err}")))?;
if x.shape.len() > 2 {
return Err(fitctree_invalid("fitctree: X must be a 2-D numeric matrix"));
}
let labels = label_vector(y_value, "Y")?;
if labels.labels.len() != x.rows {
return Err(fitctree_invalid(
"fitctree: length(Y) must match the number of rows in X",
));
}
let mut options = FitOptions::default();
parse_fit_options(&rest, &mut options)?;
if let Some(WeightSpec::Variable(_)) = options.weights {
return Err(fitctree_invalid(
"fitctree: weights variable names require table input",
));
}
let predictor_names = if let Some(names) = options.predictor_names.clone() {
if names.len() != x.cols {
return Err(fitctree_invalid(format!(
"fitctree: PredictorNames must contain {} names",
x.cols
)));
}
names
} else {
(1..=x.cols).map(|idx| format!("x{idx}")).collect()
};
let response_name = options
.response_name
.clone()
.unwrap_or_else(|| "Y".to_string());
let weights = match &options.weights {
Some(WeightSpec::Values(values)) => values.clone(),
Some(WeightSpec::Variable(_)) => unreachable!("checked above"),
None => vec![1.0; x.rows],
};
finalize_fit_spec(x, labels, predictor_names, response_name, weights, options)
}
fn build_table_fit_spec(first: Value, rest: Vec<Value>) -> BuiltinResult<FitSpec> {
let object = match first {
Value::Object(object) => object,
_ => unreachable!("table checked"),
};
let variable_names = table_variable_names_from_object(&object)
.map_err(|err| fitctree_invalid(format!("fitctree: {err}")))?;
if variable_names.len() < 2 {
return Err(fitctree_invalid(
"fitctree: table input must contain at least one predictor and one response",
));
}
let variables =
table_variables(&object).map_err(|err| fitctree_invalid(format!("fitctree: {err}")))?;
let mut options = FitOptions::default();
let mut rest_start = 0usize;
let mut response_name = variable_names
.last()
.cloned()
.ok_or_else(|| fitctree_invalid("fitctree: response variable missing"))?;
let mut external_y = None;
let mut formula_predictors = None;
if let Some(first_arg) = rest.first() {
if !is_name_value_option(first_arg) {
if let Some(text) = scalar_text(first_arg) {
if text.contains('~') {
let parsed = parse_formula(&text, &variable_names)?;
response_name = parsed.0;
formula_predictors = Some(parsed.1);
} else {
response_name = text;
}
} else {
external_y = Some(label_vector(first_arg.clone(), "Y")?);
response_name = options
.response_name
.clone()
.unwrap_or_else(|| "Y".to_string());
}
rest_start = 1;
}
}
parse_fit_options(&rest[rest_start..], &mut options)?;
if external_y.is_none() {
response_name = options.response_name.clone().unwrap_or(response_name);
}
let has_external_y = external_y.is_some();
let labels = if let Some(labels) = external_y {
labels
} else {
let value = variables.fields.get(&response_name).ok_or_else(|| {
fitctree_invalid(format!(
"fitctree: response variable '{response_name}' is not in the table"
))
})?;
label_vector(value.clone(), &response_name)?
};
let predictor_names = if let Some(names) = options.predictor_names.clone() {
names
} else if let Some(names) = formula_predictors {
names
} else {
let weight_variable = match &options.weights {
Some(WeightSpec::Variable(name)) => Some(name.as_str()),
_ => None,
};
variable_names
.iter()
.filter(|name| {
(has_external_y || **name != response_name)
&& weight_variable != Some(name.as_str())
})
.cloned()
.collect()
};
if predictor_names.is_empty() {
return Err(fitctree_invalid(
"fitctree: at least one predictor variable is required",
));
}
let mut columns = Vec::with_capacity(predictor_names.len());
for name in &predictor_names {
let value = variables.fields.get(name).ok_or_else(|| {
fitctree_invalid(format!(
"fitctree: predictor variable '{name}' is not in the table"
))
})?;
let tensor = tensor::value_into_tensor_for(FITCTREE_NAME, value.clone()).map_err(|_| {
fitctree_invalid(format!(
"fitctree: predictor variable '{name}' must be numeric"
))
})?;
columns.push(vector_values(&tensor, name)?);
}
let x = columns_to_tensor(columns)?;
let weights = match &options.weights {
Some(WeightSpec::Values(values)) => values.clone(),
Some(WeightSpec::Variable(name)) => {
let value = variables.fields.get(name).ok_or_else(|| {
fitctree_invalid(format!(
"fitctree: weights variable '{name}' is not in the table"
))
})?;
numeric_vector(value, "Weights")?
}
None => vec![1.0; x.rows],
};
finalize_fit_spec(x, labels, predictor_names, response_name, weights, options)
}
fn finalize_fit_spec(
x: Tensor,
labels: LabelVector,
predictor_names: Vec<String>,
response_name: String,
weights: Vec<f64>,
mut options: FitOptions,
) -> BuiltinResult<FitSpec> {
if labels.labels.len() != x.rows {
return Err(fitctree_invalid(
"fitctree: response length must match the number of observations",
));
}
if weights.len() != x.rows {
return Err(fitctree_invalid(
"fitctree: Weights length must match the number of observations",
));
}
for weight in &weights {
if !weight.is_finite() || *weight < 0.0 {
return Err(fitctree_invalid(
"fitctree: weights must be nonnegative finite values",
));
}
}
let class_names = if let Some(names) = options.class_names.clone() {
validate_class_names(names, labels.kind)?
} else {
unique_labels(&labels.labels, labels.kind)
};
if class_names.len() < 2 {
return Err(fitctree_invalid(
"fitctree: at least two classes are required",
));
}
let mut clean_rows = Vec::new();
let mut clean_classes = Vec::new();
let mut clean_weights = Vec::new();
for row in 0..x.rows {
let weight = weights[row];
if weight == 0.0 {
continue;
}
if is_missing_label(&labels.labels[row]) {
continue;
}
let class_index = class_names
.iter()
.position(|label| same_label(label, &labels.labels[row]))
.ok_or_else(|| {
fitctree_invalid("fitctree: response contains a class outside ClassNames")
})?;
let mut row_values = Vec::with_capacity(x.cols);
let mut complete = true;
for col in 0..x.cols {
let value = x_value(&x, row, col);
if value.is_infinite() {
return Err(fitctree_invalid(
"fitctree: predictor values must not contain Inf",
));
}
if value.is_nan() {
complete = false;
break;
}
row_values.push(value);
}
if complete {
clean_rows.push(row_values);
clean_classes.push(class_index);
clean_weights.push(weight);
}
}
if clean_classes.is_empty() {
return Err(fitctree_invalid(
"fitctree: at least one complete weighted observation is required",
));
}
let mut observed_classes = vec![false; class_names.len()];
for class in &clean_classes {
observed_classes[*class] = true;
}
if observed_classes
.iter()
.filter(|observed| **observed)
.count()
< 2
{
return Err(fitctree_invalid(
"fitctree: at least two observed classes are required",
));
}
let mut clean_data = Vec::with_capacity(clean_rows.len() * x.cols);
for col in 0..x.cols {
for row in &clean_rows {
clean_data.push(row[col]);
}
}
let clean_x = Tensor::new(clean_data, vec![clean_rows.len(), x.cols])
.map_err(|err| fitctree_internal(format!("fitctree: {err}")))?;
if options.max_num_splits.is_none() {
options.max_num_splits = Some(clean_classes.len().saturating_sub(1));
}
Ok(FitSpec {
x: clean_x,
class_indices: clean_classes,
class_names,
class_kind: labels.kind,
predictor_names,
response_name,
weights: clean_weights,
options,
})
}
fn parse_fit_options(values: &[Value], options: &mut FitOptions) -> BuiltinResult<()> {
if !values.len().is_multiple_of(2) {
return Err(fitctree_invalid(
"fitctree: name-value options must be paired",
));
}
let mut index = 0usize;
while index < values.len() {
let name = scalar_text(&values[index])
.ok_or_else(|| fitctree_invalid("fitctree: option name must be text"))?;
let value = &values[index + 1];
match name.to_ascii_lowercase().as_str() {
"classnames" => {
let labels = label_vector(value.clone(), "ClassNames")?;
options.class_names = Some(labels.labels);
}
"predictornames" => {
options.predictor_names = Some(text_list(value, "PredictorNames")?);
}
"responsename" => {
options.response_name = Some(
scalar_text(value)
.ok_or_else(|| fitctree_invalid("fitctree: ResponseName must be text"))?,
);
}
"maxnumsplits" => {
options.max_num_splits = Some(nonnegative_integer(value, "MaxNumSplits")?);
}
"minleafsize" => {
options.min_leaf_size = positive_integer(value, "MinLeafSize")?;
}
"minparentsize" => {
options.min_parent_size = positive_integer(value, "MinParentSize")?;
}
"splitcriterion" => {
let text = scalar_text(value)
.ok_or_else(|| fitctree_invalid("fitctree: SplitCriterion must be text"))?;
options.split_criterion = match text.to_ascii_lowercase().as_str() {
"gdi" => SplitCriterion::Gdi,
"deviance" => SplitCriterion::Deviance,
other => {
return Err(fitctree_invalid(format!(
"fitctree: unsupported SplitCriterion '{other}'"
)));
}
};
}
"weights" => {
options.weights = Some(if let Some(name) = scalar_text(value) {
WeightSpec::Variable(name)
} else {
WeightSpec::Values(numeric_vector(value, "Weights")?)
});
}
"scoretransform" => {
let text = scalar_text(value)
.ok_or_else(|| fitctree_invalid("fitctree: ScoreTransform must be text"))?;
if !text.eq_ignore_ascii_case("none") {
return Err(fitctree_invalid(
"fitctree: only ScoreTransform='none' is supported",
));
}
}
"algorithmforcategorical"
| "categoricalpredictors"
| "cost"
| "crossval"
| "cvpartition"
| "holdout"
| "kfold"
| "leaveout"
| "maxnumcategories"
| "mergeleaves"
| "numbins"
| "numvariablestosample"
| "optimizehyperparameters"
| "hyperparameteroptimizationoptions"
| "predictorselection"
| "prior"
| "prune"
| "prunecriterion"
| "reproducible"
| "surrogate" => {
return Err(fitctree_invalid(format!(
"fitctree: option '{name}' is not supported yet"
)));
}
other => {
return Err(fitctree_invalid(format!(
"fitctree: unknown option '{other}'"
)))
}
}
index += 2;
}
Ok(())
}
fn parse_formula(text: &str, variables: &[String]) -> BuiltinResult<(String, Vec<String>)> {
let Some((lhs, rhs)) = text.split_once('~') else {
return Err(fitctree_invalid("fitctree: formula must contain '~'"));
};
let response = lhs.trim().to_string();
if response.is_empty() {
return Err(fitctree_invalid("fitctree: formula response is empty"));
}
let mut predictors = Vec::new();
for term in rhs.split('+') {
let term = term.trim();
if term.is_empty() {
continue;
}
if term == "1" {
continue;
}
if term.contains('*') || term.contains(':') || term.contains('^') {
return Err(fitctree_invalid(
"fitctree: formula interactions are not supported yet",
));
}
if !variables.iter().any(|name| name == term) {
return Err(fitctree_invalid(format!(
"fitctree: formula predictor '{term}' is not in the table"
)));
}
predictors.push(term.to_string());
}
if predictors.is_empty() {
return Err(fitctree_invalid(
"fitctree: formula must contain at least one predictor",
));
}
Ok((response, predictors))
}
fn fit_tree(spec: &FitSpec) -> BuiltinResult<TreeFit> {
let mut nodes = Vec::new();
let mut splits = 0usize;
let rows = (0..spec.x.rows).collect::<Vec<_>>();
build_node(spec, &rows, &mut nodes, &mut splits, 0)?;
Ok(TreeFit {
nodes,
num_splits: splits,
})
}
fn build_node(
spec: &FitSpec,
rows: &[usize],
nodes: &mut Vec<Node>,
split_count: &mut usize,
depth: usize,
) -> BuiltinResult<usize> {
if depth > MAX_TREE_DEPTH {
return Err(fitctree_invalid(
"fitctree: tree depth exceeds RunMat safety limit",
));
}
let counts = class_counts(
rows,
&spec.class_indices,
&spec.weights,
spec.class_names.len(),
);
let total_weight = counts.iter().sum::<f64>();
let probabilities = if total_weight > 0.0 {
counts.iter().map(|value| value / total_weight).collect()
} else {
vec![0.0; counts.len()]
};
let class_index = probabilities
.iter()
.enumerate()
.max_by(|left, right| {
left.1
.partial_cmp(right.1)
.unwrap_or(Ordering::Equal)
.then_with(|| right.0.cmp(&left.0))
})
.map(|(idx, _)| idx)
.unwrap_or(0);
let node_index = nodes.len();
nodes.push(Node {
left: None,
right: None,
predictor: None,
threshold: f64::NAN,
class_index,
probabilities,
});
let max_splits = spec.options.max_num_splits.unwrap_or(0);
if *split_count >= max_splits
|| rows.len() < spec.options.min_parent_size
|| rows.len() < spec.options.min_leaf_size.saturating_mul(2)
|| counts.iter().filter(|value| **value > 0.0).count() <= 1
{
return Ok(node_index);
}
if let Some(best) = best_split(spec, rows, &counts)? {
if best.gain <= EPS {
return Ok(node_index);
}
*split_count += 1;
let left = build_node(spec, &best.left_rows, nodes, split_count, depth + 1)?;
let right = build_node(spec, &best.right_rows, nodes, split_count, depth + 1)?;
nodes[node_index].left = Some(left);
nodes[node_index].right = Some(right);
nodes[node_index].predictor = Some(best.predictor);
nodes[node_index].threshold = best.threshold;
}
Ok(node_index)
}
fn best_split(
spec: &FitSpec,
rows: &[usize],
parent_counts: &[f64],
) -> BuiltinResult<Option<SplitCandidate>> {
let parent_weight = parent_counts.iter().sum::<f64>();
let parent_impurity = impurity(parent_counts, spec.options.split_criterion);
let mut best = None;
for col in 0..spec.x.cols {
let mut entries = rows
.iter()
.map(|row| {
(
x_value(&spec.x, *row, col),
spec.class_indices[*row],
spec.weights[*row],
*row,
)
})
.collect::<Vec<_>>();
entries.sort_by(|left, right| left.0.partial_cmp(&right.0).unwrap_or(Ordering::Equal));
if entries.len() < 2 {
continue;
}
let mut left_counts = vec![0.0; parent_counts.len()];
let mut right_counts = parent_counts.to_vec();
let mut left_weight = 0.0;
let mut right_weight = parent_weight;
let mut left_len = 0usize;
let mut right_len = entries.len();
for idx in 0..entries.len() - 1 {
let (_, class_idx, weight, _) = entries[idx];
left_counts[class_idx] += weight;
right_counts[class_idx] -= weight;
left_weight += weight;
right_weight -= weight;
left_len += 1;
right_len -= 1;
let current = entries[idx].0;
let next = entries[idx + 1].0;
if (current - next).abs() <= EPS {
continue;
}
if left_len < spec.options.min_leaf_size || right_len < spec.options.min_leaf_size {
continue;
}
if left_weight <= 0.0 || right_weight <= 0.0 {
continue;
}
let left_impurity = impurity(&left_counts, spec.options.split_criterion);
let right_impurity = impurity(&right_counts, spec.options.split_criterion);
let gain = parent_impurity
- (left_weight / parent_weight) * left_impurity
- (right_weight / parent_weight) * right_impurity;
if gain
> best
.as_ref()
.map(|candidate: &SplitCandidate| candidate.gain)
.unwrap_or(-1.0)
{
let threshold = current + (next - current) / 2.0;
let mut left_rows = Vec::with_capacity(left_len);
let mut right_rows = Vec::with_capacity(right_len);
for (value, _, _, row) in &entries {
if *value <= threshold {
left_rows.push(*row);
} else {
right_rows.push(*row);
}
}
best = Some(SplitCandidate {
predictor: col,
threshold,
gain,
left_rows,
right_rows,
});
}
}
}
Ok(best)
}
fn impurity(counts: &[f64], criterion: SplitCriterion) -> f64 {
let total = counts.iter().sum::<f64>();
if total <= 0.0 {
return 0.0;
}
match criterion {
SplitCriterion::Gdi => {
1.0 - counts
.iter()
.map(|count| {
let p = count / total;
p * p
})
.sum::<f64>()
}
SplitCriterion::Deviance => counts
.iter()
.filter(|count| **count > 0.0)
.map(|count| {
let p = count / total;
-p * p.ln()
})
.sum(),
}
}
fn class_counts(
rows: &[usize],
classes: &[usize],
weights: &[f64],
class_count: usize,
) -> Vec<f64> {
let mut counts = vec![0.0; class_count];
for row in rows {
counts[classes[*row]] += weights[*row];
}
counts
}
fn model_object(spec: FitSpec, fit: TreeFit) -> BuiltinResult<ObjectInstance> {
let mut object = ObjectInstance::new(CLASSIFICATION_TREE_CLASS.to_string());
object.properties.insert(
"ResponseName".to_string(),
Value::String(spec.response_name),
);
object.properties.insert(
"PredictorNames".to_string(),
string_column(&spec.predictor_names)?,
);
object.properties.insert(
"ClassNames".to_string(),
labels_value(&spec.class_names, spec.class_kind)?,
);
object.properties.insert(
"CategoricalPredictors".to_string(),
Value::Tensor(Tensor::new(Vec::new(), vec![0, 1]).map_err(|err| {
fitctree_internal(format!(
"fitctree: categorical predictor shape error: {err}"
))
})?),
);
object.properties.insert(
"ScoreTransform".to_string(),
Value::String("none".to_string()),
);
object.properties.insert(
"NumObservations".to_string(),
Value::Num(spec.x.rows as f64),
);
object
.properties
.insert("NumPredictors".to_string(), Value::Num(spec.x.cols as f64));
object
.properties
.insert("NumNodes".to_string(), Value::Num(fit.nodes.len() as f64));
object
.properties
.insert("NumSplits".to_string(), Value::Num(fit.num_splits as f64));
let mut params = StructValue::new();
params.insert(
"MaxNumSplits",
Value::Num(spec.options.max_num_splits.unwrap_or(0) as f64),
);
params.insert("MinLeaf", Value::Num(spec.options.min_leaf_size as f64));
params.insert("MinParent", Value::Num(spec.options.min_parent_size as f64));
params.insert(
"SplitCriterion",
Value::String(
match spec.options.split_criterion {
SplitCriterion::Gdi => "gdi",
SplitCriterion::Deviance => "deviance",
}
.to_string(),
),
);
object
.properties
.insert("ModelParameters".to_string(), Value::Struct(params));
object.properties.insert(
"__RunMatTreeClassKind".to_string(),
Value::String(
match spec.class_kind {
LabelKind::Numeric => "numeric",
LabelKind::Text => "text",
LabelKind::Logical => "logical",
}
.to_string(),
),
);
object.properties.insert(
"__RunMatTreeClassNames".to_string(),
labels_value(&spec.class_names, spec.class_kind)?,
);
object.properties.insert(
"__RunMatTreeChildren".to_string(),
tree_children_tensor(&fit.nodes)?,
);
object.properties.insert(
"__RunMatTreePredictor".to_string(),
tree_predictor_tensor(&fit.nodes)?,
);
object.properties.insert(
"__RunMatTreeThreshold".to_string(),
tree_threshold_tensor(&fit.nodes)?,
);
object.properties.insert(
"__RunMatTreeNodeClass".to_string(),
tree_class_tensor(&fit.nodes)?,
);
object.properties.insert(
"__RunMatTreeNodeProb".to_string(),
tree_prob_tensor(&fit.nodes, spec.class_names.len())?,
);
Ok(object)
}
pub(crate) fn predict_classification_tree_object(
object: ObjectInstance,
xnew: Value,
rest: Vec<Value>,
) -> BuiltinResult<Vec<Value>> {
if !rest.is_empty() {
return Err(predict_invalid(
"predict: ClassificationTree Subtrees and name-value options are not supported yet",
));
}
let predictor_names = string_property(&object, "PredictorNames")?;
let x = predictors_for_prediction(xnew, &predictor_names)?;
let children = matrix_property(&object, "__RunMatTreeChildren")?;
let predictors = numeric_property(&object, "__RunMatTreePredictor")?;
let thresholds = numeric_property(&object, "__RunMatTreeThreshold")?;
let node_classes = numeric_property(&object, "__RunMatTreeNodeClass")?;
let node_probs = matrix_property(&object, "__RunMatTreeNodeProb")?;
let class_names_value = object
.properties
.get("__RunMatTreeClassNames")
.cloned()
.ok_or_else(|| predict_invalid("predict: tree is missing class names"))?;
let class_kind = class_kind_property(&object)?;
let class_names = labels_from_value(class_names_value, "ClassNames")?;
let class_count = class_names.labels.len();
let node_count = predictors.len();
if children.len() != node_count * 2
|| thresholds.len() != node_count
|| node_classes.len() != node_count
|| node_probs.len() != node_count * class_count
{
return Err(predict_invalid("predict: tree metadata is inconsistent"));
}
let mut label_indices = Vec::with_capacity(x.rows);
let mut score_data = vec![0.0; x.rows * class_count];
let mut node_data = Vec::with_capacity(x.rows);
let mut class_number_data = Vec::with_capacity(x.rows);
for row in 0..x.rows {
let node = traverse_tree(&x, row, &children, &predictors, &thresholds, node_count)?;
let class_index = checked_one_based_index(node_classes[node], class_count, "class number")?;
label_indices.push(class_index);
node_data.push(node as f64 + 1.0);
class_number_data.push(class_index as f64 + 1.0);
for class in 0..class_count {
score_data[class * x.rows + row] = node_probs[class * node_count + node];
}
}
let labels = predicted_labels_value(&class_names.labels, class_kind, &label_indices)?;
let scores = Value::Tensor(
Tensor::new(score_data, vec![x.rows, class_count])
.map_err(|err| predict_internal(format!("predict: {err}")))?,
);
let node = Value::Tensor(
Tensor::new(node_data, vec![x.rows, 1])
.map_err(|err| predict_internal(format!("predict: {err}")))?,
);
let cnum = Value::Tensor(
Tensor::new(class_number_data, vec![x.rows, 1])
.map_err(|err| predict_internal(format!("predict: {err}")))?,
);
Ok(vec![labels, scores, node, cnum])
}
fn traverse_tree(
x: &Tensor,
row: usize,
children: &[f64],
predictors: &[f64],
thresholds: &[f64],
node_count: usize,
) -> BuiltinResult<usize> {
let mut node = 0usize;
for _ in 0..node_count {
let predictor = predictors[node];
if predictor <= 0.0 {
return Ok(node);
}
let predictor_idx = checked_one_based_index(predictor, x.cols, "predictor")?;
let value = x_value(x, row, predictor_idx);
if value.is_nan() {
return Ok(node);
}
let child_col = if value <= thresholds[node] { 0 } else { 1 };
let child = children[child_col * node_count + node];
if child <= 0.0 {
return Ok(node);
}
node = checked_one_based_index(child, node_count, "child node")?;
}
Err(predict_invalid("predict: tree traversal did not terminate"))
}
fn checked_one_based_index(value: f64, len: usize, label: &str) -> BuiltinResult<usize> {
if !value.is_finite() || value < 1.0 || value.fract() != 0.0 {
return Err(predict_invalid(format!("predict: invalid {label}")));
}
let index = value as usize - 1;
if index >= len {
return Err(predict_invalid(format!("predict: {label} out of range")));
}
Ok(index)
}
fn predictors_for_prediction(value: Value, predictor_names: &[String]) -> BuiltinResult<Tensor> {
if let Value::Object(object) = &value {
if is_tabular_object(object) {
let variables = table_variables(object)
.map_err(|err| predict_invalid(format!("predict: {err}")))?;
let mut columns = Vec::with_capacity(predictor_names.len());
for name in predictor_names {
let raw = variables.fields.get(name).ok_or_else(|| {
predict_invalid(format!("predict: missing predictor '{name}'"))
})?;
let tensor =
tensor::value_into_tensor_for(PREDICT_NAME, raw.clone()).map_err(|_| {
predict_invalid(format!("predict: predictor '{name}' must be numeric"))
})?;
columns.push(vector_values_predict(&tensor, name)?);
}
return columns_to_tensor_predict(columns);
}
}
let tensor = tensor::value_into_tensor_for(PREDICT_NAME, value)
.map_err(|err| predict_invalid(format!("predict: {err}")))?;
if tensor.shape.len() > 2 || tensor.cols != predictor_names.len() {
return Err(predict_invalid(format!(
"predict: X must have {} predictor columns",
predictor_names.len()
)));
}
Ok(tensor)
}
fn label_vector(value: Value, name: &str) -> BuiltinResult<LabelVector> {
labels_from_value(value, name).map_err(|err| fitctree_invalid(err.message))
}
fn labels_from_value(value: Value, name: &str) -> BuiltinResult<LabelVector> {
match value {
Value::Tensor(tensor) => Ok(LabelVector {
labels: vector_values_predict(&tensor, name)?
.into_iter()
.map(ClassLabel::Numeric)
.collect(),
kind: LabelKind::Numeric,
}),
Value::Num(value) => Ok(LabelVector {
labels: vec![ClassLabel::Numeric(value)],
kind: LabelKind::Numeric,
}),
Value::String(text) => Ok(LabelVector {
labels: vec![ClassLabel::Text(text)],
kind: LabelKind::Text,
}),
Value::CharArray(array) => Ok(LabelVector {
labels: char_rows(&array).into_iter().map(ClassLabel::Text).collect(),
kind: LabelKind::Text,
}),
Value::StringArray(array) => Ok(LabelVector {
labels: array.data.into_iter().map(ClassLabel::Text).collect(),
kind: LabelKind::Text,
}),
Value::Cell(cell) => {
let mut labels = Vec::with_capacity(cell.data.len());
for item in cell.data {
labels.push(ClassLabel::Text(
scalar_text(&item).ok_or_else(|| {
predict_invalid(format!("predict: {name} cell entries must be text"))
})?,
));
}
Ok(LabelVector {
labels,
kind: LabelKind::Text,
})
}
Value::Bool(value) => Ok(LabelVector {
labels: vec![ClassLabel::Logical(value)],
kind: LabelKind::Logical,
}),
Value::LogicalArray(array) => Ok(LabelVector {
labels: array
.data
.into_iter()
.map(|value| ClassLabel::Logical(value != 0))
.collect(),
kind: LabelKind::Logical,
}),
other => Err(predict_invalid(format!(
"predict: {name} must be numeric, logical, string, character, or cellstr labels; got {other:?}"
))),
}
}
fn validate_class_names(names: Vec<ClassLabel>, kind: LabelKind) -> BuiltinResult<Vec<ClassLabel>> {
if names.is_empty() {
return Err(fitctree_invalid("fitctree: ClassNames must not be empty"));
}
for name in &names {
if is_missing_label(name) {
return Err(fitctree_invalid(
"fitctree: ClassNames must not contain missing values",
));
}
if label_kind(name) != kind {
return Err(fitctree_invalid(
"fitctree: ClassNames must have the same type as the response",
));
}
}
let mut out = Vec::new();
for name in names {
if out.iter().any(|existing| same_label(existing, &name)) {
return Err(fitctree_invalid(
"fitctree: ClassNames must contain unique values",
));
}
out.push(name);
}
Ok(out)
}
fn unique_labels(labels: &[ClassLabel], kind: LabelKind) -> Vec<ClassLabel> {
let mut out = Vec::new();
for label in labels {
if is_missing_label(label) {
continue;
}
if !out.iter().any(|existing| same_label(existing, label)) {
out.push(label.clone());
}
}
match kind {
LabelKind::Numeric => out.sort_by(|left, right| match (left, right) {
(ClassLabel::Numeric(a), ClassLabel::Numeric(b)) => {
a.partial_cmp(b).unwrap_or(Ordering::Equal)
}
_ => Ordering::Equal,
}),
LabelKind::Text => out.sort_by(|left, right| match (left, right) {
(ClassLabel::Text(a), ClassLabel::Text(b)) => a.cmp(b),
_ => Ordering::Equal,
}),
LabelKind::Logical => out.sort_by_key(|label| match label {
ClassLabel::Logical(value) => *value,
_ => false,
}),
}
out
}
fn label_kind(label: &ClassLabel) -> LabelKind {
match label {
ClassLabel::Numeric(_) => LabelKind::Numeric,
ClassLabel::Text(_) => LabelKind::Text,
ClassLabel::Logical(_) => LabelKind::Logical,
}
}
fn same_label(left: &ClassLabel, right: &ClassLabel) -> bool {
match (left, right) {
(ClassLabel::Numeric(a), ClassLabel::Numeric(b)) => (a == b) || (a.is_nan() && b.is_nan()),
(ClassLabel::Text(a), ClassLabel::Text(b)) => a == b,
(ClassLabel::Logical(a), ClassLabel::Logical(b)) => a == b,
_ => false,
}
}
fn is_missing_label(label: &ClassLabel) -> bool {
matches!(label, ClassLabel::Numeric(value) if value.is_nan())
}
fn labels_value(labels: &[ClassLabel], kind: LabelKind) -> BuiltinResult<Value> {
predicted_labels_value(labels, kind, &(0..labels.len()).collect::<Vec<_>>())
}
fn predicted_labels_value(
labels: &[ClassLabel],
kind: LabelKind,
indices: &[usize],
) -> BuiltinResult<Value> {
match kind {
LabelKind::Numeric => {
let data = indices
.iter()
.map(|idx| match labels.get(*idx) {
Some(ClassLabel::Numeric(value)) => *value,
_ => f64::NAN,
})
.collect::<Vec<_>>();
Ok(Value::Tensor(
Tensor::new(data, vec![indices.len(), 1])
.map_err(|err| predict_internal(format!("predict: {err}")))?,
))
}
LabelKind::Text => {
let data = indices
.iter()
.map(|idx| match labels.get(*idx) {
Some(ClassLabel::Text(value)) => value.clone(),
_ => String::new(),
})
.collect::<Vec<_>>();
Ok(Value::StringArray(StringArray {
data,
shape: vec![indices.len(), 1],
rows: indices.len(),
cols: 1,
}))
}
LabelKind::Logical => {
let data = indices
.iter()
.map(|idx| match labels.get(*idx) {
Some(ClassLabel::Logical(value)) if *value => 1,
_ => 0,
})
.collect::<Vec<_>>();
Ok(Value::LogicalArray(
LogicalArray::new(data, vec![indices.len(), 1])
.map_err(|err| predict_internal(format!("predict: {err}")))?,
))
}
}
}
fn class_kind_property(object: &ObjectInstance) -> BuiltinResult<LabelKind> {
match object.properties.get("__RunMatTreeClassKind") {
Some(Value::String(text)) if text == "numeric" => Ok(LabelKind::Numeric),
Some(Value::String(text)) if text == "text" => Ok(LabelKind::Text),
Some(Value::String(text)) if text == "logical" => Ok(LabelKind::Logical),
_ => Err(predict_invalid(
"predict: tree is missing class type metadata",
)),
}
}
fn scalar_text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::CharArray(array) if array.rows == 1 => Some(array.data.iter().collect()),
Value::StringArray(array) if array.data.len() == 1 => Some(array.data[0].clone()),
_ => None,
}
}
fn text_list(value: &Value, name: &str) -> BuiltinResult<Vec<String>> {
match value {
Value::String(text) => Ok(vec![text.clone()]),
Value::CharArray(array) => Ok(char_rows(array)),
Value::StringArray(array) => Ok(array.data.clone()),
Value::Cell(cell) => cell
.data
.iter()
.map(|value| {
scalar_text(value).ok_or_else(|| {
fitctree_invalid(format!("fitctree: {name} entries must be text"))
})
})
.collect(),
_ => Err(fitctree_invalid(format!(
"fitctree: {name} must be text or a text collection"
))),
}
}
fn char_rows(array: &CharArray) -> Vec<String> {
let mut out = Vec::with_capacity(array.rows);
for row in 0..array.rows {
let start = row * array.cols;
out.push(array.data[start..start + array.cols].iter().collect());
}
out
}
fn numeric_vector(value: &Value, name: &str) -> BuiltinResult<Vec<f64>> {
match value {
Value::Num(value) => Ok(vec![*value]),
Value::Tensor(tensor) => vector_values(tensor, name),
Value::LogicalArray(array) => Ok(array
.data
.iter()
.map(|value| if *value == 0 { 0.0 } else { 1.0 })
.collect()),
Value::Bool(value) => Ok(vec![if *value { 1.0 } else { 0.0 }]),
_ => Err(fitctree_invalid(format!(
"fitctree: {name} must be numeric"
))),
}
}
fn vector_values(tensor: &Tensor, name: &str) -> BuiltinResult<Vec<f64>> {
vector_values_predict(tensor, name).map_err(|err| fitctree_invalid(err.message))
}
fn vector_values_predict(tensor: &Tensor, name: &str) -> BuiltinResult<Vec<f64>> {
if tensor.shape.len() > 2 || (tensor.rows != 1 && tensor.cols != 1) {
return Err(predict_invalid(format!("predict: {name} must be a vector")));
}
Ok(tensor.data.clone())
}
fn nonnegative_integer(value: &Value, name: &str) -> BuiltinResult<usize> {
let value = numeric_scalar(value, name)?;
if !value.is_finite() || value < 0.0 || value.fract() != 0.0 {
return Err(fitctree_invalid(format!(
"fitctree: {name} must be a nonnegative integer"
)));
}
Ok(value as usize)
}
fn positive_integer(value: &Value, name: &str) -> BuiltinResult<usize> {
let value = numeric_scalar(value, name)?;
if !value.is_finite() || value < 1.0 || value.fract() != 0.0 {
return Err(fitctree_invalid(format!(
"fitctree: {name} must be a positive integer"
)));
}
Ok(value as usize)
}
fn numeric_scalar(value: &Value, name: &str) -> BuiltinResult<f64> {
match value {
Value::Num(value) => Ok(*value),
Value::Tensor(tensor) if tensor.data.len() == 1 => Ok(tensor.data[0]),
Value::Bool(value) => Ok(if *value { 1.0 } else { 0.0 }),
_ => Err(fitctree_invalid(format!(
"fitctree: {name} must be a numeric scalar"
))),
}
}
fn is_name_value_option(value: &Value) -> bool {
scalar_text(value).is_some_and(|text| {
matches!(
text.to_ascii_lowercase().as_str(),
"algorithmforcategorical"
| "categoricalpredictors"
| "classnames"
| "cost"
| "crossval"
| "cvpartition"
| "holdout"
| "hyperparameteroptimizationoptions"
| "kfold"
| "leaveout"
| "maxnumcategories"
| "maxnumsplits"
| "mergeleaves"
| "minleafsize"
| "minparentsize"
| "numbins"
| "numvariablestosample"
| "optimizehyperparameters"
| "predictornames"
| "predictorselection"
| "prior"
| "prune"
| "prunecriterion"
| "reproducible"
| "responsename"
| "scoretransform"
| "splitcriterion"
| "surrogate"
| "weights"
)
})
}
fn columns_to_tensor(columns: Vec<Vec<f64>>) -> BuiltinResult<Tensor> {
if columns.is_empty() {
return Err(fitctree_invalid("fitctree: no predictor columns"));
}
let rows = columns[0].len();
let cols = columns.len();
let mut data = Vec::with_capacity(rows * cols);
for column in columns {
if column.len() != rows {
return Err(fitctree_invalid(
"fitctree: predictor variables must have the same height",
));
}
data.extend(column);
}
Tensor::new(data, vec![rows, cols]).map_err(|err| fitctree_internal(format!("fitctree: {err}")))
}
fn columns_to_tensor_predict(columns: Vec<Vec<f64>>) -> BuiltinResult<Tensor> {
if columns.is_empty() {
return Err(predict_invalid("predict: no predictor columns"));
}
let rows = columns[0].len();
let cols = columns.len();
let mut data = Vec::with_capacity(rows * cols);
for column in columns {
if column.len() != rows {
return Err(predict_invalid(
"predict: predictor table variables must have the same height",
));
}
data.extend(column);
}
Tensor::new(data, vec![rows, cols]).map_err(|err| predict_internal(format!("predict: {err}")))
}
fn x_value(tensor: &Tensor, row: usize, col: usize) -> f64 {
tensor.data[col * tensor.rows + row]
}
fn string_column(values: &[String]) -> BuiltinResult<Value> {
Ok(Value::StringArray(StringArray {
data: values.to_vec(),
shape: vec![values.len(), 1],
rows: values.len(),
cols: 1,
}))
}
fn numeric_property(object: &ObjectInstance, name: &str) -> BuiltinResult<Vec<f64>> {
match object.properties.get(name) {
Some(Value::Tensor(tensor)) => Ok(tensor.data.clone()),
_ => Err(predict_invalid(format!("predict: tree is missing {name}"))),
}
}
fn matrix_property(object: &ObjectInstance, name: &str) -> BuiltinResult<Vec<f64>> {
numeric_property(object, name)
}
fn string_property(object: &ObjectInstance, name: &str) -> BuiltinResult<Vec<String>> {
match object.properties.get(name) {
Some(Value::StringArray(array)) => Ok(array.data.clone()),
Some(Value::String(text)) => Ok(vec![text.clone()]),
_ => Err(predict_invalid(format!("predict: tree is missing {name}"))),
}
}
fn tree_children_tensor(nodes: &[Node]) -> BuiltinResult<Value> {
let rows = nodes.len();
let mut data = Vec::with_capacity(rows * 2);
for node in nodes {
data.push(node.left.map(|idx| idx as f64 + 1.0).unwrap_or(0.0));
}
for node in nodes {
data.push(node.right.map(|idx| idx as f64 + 1.0).unwrap_or(0.0));
}
Ok(Value::Tensor(Tensor::new(data, vec![rows, 2]).map_err(
|err| fitctree_internal(format!("fitctree: {err}")),
)?))
}
fn tree_predictor_tensor(nodes: &[Node]) -> BuiltinResult<Value> {
Ok(Value::Tensor(
Tensor::new(
nodes
.iter()
.map(|node| node.predictor.map(|idx| idx as f64 + 1.0).unwrap_or(0.0))
.collect(),
vec![nodes.len(), 1],
)
.map_err(|err| fitctree_internal(format!("fitctree: {err}")))?,
))
}
fn tree_threshold_tensor(nodes: &[Node]) -> BuiltinResult<Value> {
Ok(Value::Tensor(
Tensor::new(
nodes.iter().map(|node| node.threshold).collect(),
vec![nodes.len(), 1],
)
.map_err(|err| fitctree_internal(format!("fitctree: {err}")))?,
))
}
fn tree_class_tensor(nodes: &[Node]) -> BuiltinResult<Value> {
Ok(Value::Tensor(
Tensor::new(
nodes
.iter()
.map(|node| node.class_index as f64 + 1.0)
.collect(),
vec![nodes.len(), 1],
)
.map_err(|err| fitctree_internal(format!("fitctree: {err}")))?,
))
}
fn tree_prob_tensor(nodes: &[Node], class_count: usize) -> BuiltinResult<Value> {
let rows = nodes.len();
let mut data = Vec::with_capacity(rows * class_count);
for class in 0..class_count {
for node in nodes {
data.push(node.probabilities.get(class).copied().unwrap_or(0.0));
}
}
Ok(Value::Tensor(
Tensor::new(data, vec![rows, class_count])
.map_err(|err| fitctree_internal(format!("fitctree: {err}")))?,
))
}
#[cfg(test)]
mod tests {
use super::*;
use futures::executor::block_on;
fn tensor_value(data: Vec<f64>, rows: usize, cols: usize) -> Value {
Value::Tensor(Tensor::new(data, vec![rows, cols]).unwrap())
}
fn object(value: Value) -> ObjectInstance {
match value {
Value::Object(object) => object,
other => panic!("expected object, got {other:?}"),
}
}
fn tensor(value: &Value) -> &Tensor {
match value {
Value::Tensor(tensor) => tensor,
other => panic!("expected tensor, got {other:?}"),
}
}
#[test]
fn fitctree_matrix_predicts_numeric_labels_and_scores() {
let x = tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1);
let y = tensor_value(vec![0.0, 0.0, 1.0, 1.0], 4, 1);
let model = block_on(fitctree_builtin(
x,
vec![
y,
Value::from("MaxNumSplits"),
Value::Num(1.0),
Value::from("MinParentSize"),
Value::Num(2.0),
],
))
.unwrap();
let output_count = crate::output_count::push_output_count(Some(4));
let predicted = crate::builtins::stats::ml::linear_model::predict_builtin(
model,
tensor_value(vec![0.5, 2.5], 2, 1),
Vec::new(),
);
let predicted = block_on(predicted).unwrap();
drop(output_count);
let Value::OutputList(values) = predicted else {
panic!("expected output list");
};
assert_eq!(tensor(&values[0]).data, vec![0.0, 1.0]);
assert_eq!(tensor(&values[1]).shape, vec![2, 2]);
assert_eq!(tensor(&values[2]).shape, vec![2, 1]);
assert_eq!(tensor(&values[3]).data, vec![1.0, 2.0]);
}
#[test]
fn fitctree_table_response_name_predicts_string_labels() {
let table = crate::builtins::table::table_from_columns(
vec!["A".into(), "Y".into()],
vec![
tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1),
Value::StringArray(StringArray {
data: vec!["low".into(), "low".into(), "high".into(), "high".into()],
shape: vec![4, 1],
rows: 4,
cols: 1,
}),
],
)
.unwrap();
let model = block_on(fitctree_builtin(
table.clone(),
vec![
Value::String("Y".into()),
Value::from("MaxNumSplits"),
Value::Num(1.0),
Value::from("MinParentSize"),
Value::Num(2.0),
],
))
.unwrap();
let predicted = block_on(crate::builtins::stats::ml::linear_model::predict_builtin(
model,
table,
Vec::new(),
))
.unwrap();
let Value::StringArray(labels) = predicted else {
panic!("expected string labels");
};
assert_eq!(labels.shape, vec![4, 1]);
assert_eq!(labels.data[0], "low");
assert_eq!(labels.data[3], "high");
}
#[test]
fn fitctree_respects_class_names_and_depth_controls() {
let x = tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1);
let y = tensor_value(vec![2.0, 2.0, 5.0, 5.0], 4, 1);
let model = object(
block_on(fitctree_builtin(
x,
vec![
y,
Value::from("ClassNames"),
tensor_value(vec![5.0, 2.0], 2, 1),
Value::from("MaxNumSplits"),
Value::Num(0.0),
],
))
.unwrap(),
);
assert_eq!(model.properties.get("NumSplits"), Some(&Value::Num(0.0)));
let classes = tensor(model.properties.get("ClassNames").unwrap());
assert_eq!(classes.data, vec![5.0, 2.0]);
}
#[test]
fn fitctree_omits_nan_responses_and_requires_two_observed_classes() {
let model = object(
block_on(fitctree_builtin(
tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1),
vec![
tensor_value(vec![0.0, f64::NAN, 1.0, 1.0], 4, 1),
Value::from("MinParentSize"),
Value::Num(2.0),
],
))
.unwrap(),
);
assert_eq!(
model.properties.get("NumObservations"),
Some(&Value::Num(3.0))
);
assert_eq!(
tensor(model.properties.get("ClassNames").unwrap()).data,
vec![0.0, 1.0]
);
let err = block_on(fitctree_builtin(
tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1),
vec![
tensor_value(vec![0.0, f64::NAN, f64::NAN, 0.0], 4, 1),
Value::from("ClassNames"),
tensor_value(vec![0.0, 1.0], 2, 1),
],
))
.unwrap_err();
assert!(err.message().contains("at least two observed classes"));
}
#[test]
fn fitctree_excludes_table_weight_variable_from_default_predictors() {
let table = crate::builtins::table::table_from_columns(
vec!["A".into(), "W".into(), "Y".into()],
vec![
tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1),
tensor_value(vec![1.0, 2.0, 1.0, 2.0], 4, 1),
tensor_value(vec![0.0, 0.0, 1.0, 1.0], 4, 1),
],
)
.unwrap();
let model = object(
block_on(fitctree_builtin(
table,
vec![Value::from("Y"), Value::from("Weights"), Value::from("W")],
))
.unwrap(),
);
assert_eq!(
model.properties.get("NumPredictors"),
Some(&Value::Num(1.0))
);
let Value::StringArray(names) = model.properties.get("PredictorNames").unwrap() else {
panic!("expected predictor names");
};
assert_eq!(names.data, vec!["A"]);
}
#[test]
fn fitctree_predict_rejects_too_many_outputs() {
let model = block_on(fitctree_builtin(
tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1),
vec![
tensor_value(vec![0.0, 0.0, 1.0, 1.0], 4, 1),
Value::from("MaxNumSplits"),
Value::Num(1.0),
Value::from("MinParentSize"),
Value::Num(2.0),
],
))
.unwrap();
let output_count = crate::output_count::push_output_count(Some(5));
let err = block_on(crate::builtins::stats::ml::linear_model::predict_builtin(
model,
tensor_value(vec![0.5, 2.5], 2, 1),
Vec::new(),
))
.unwrap_err();
drop(output_count);
assert!(err.message().contains("too many output"));
}
#[test]
fn fitctree_rejects_unsupported_cross_validation_options() {
let x = tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1);
let y = tensor_value(vec![0.0, 0.0, 1.0, 1.0], 4, 1);
let err = block_on(fitctree_builtin(
x,
vec![y, Value::from("CrossVal"), Value::String("on".into())],
))
.unwrap_err();
assert!(err.message.contains("not supported"));
}
}