Skip to main content

runmat_runtime/builtins/stats/ml/
predict.rs

1//! Generic model prediction dispatcher.
2
3use runmat_builtins::Value;
4use runmat_macros::runtime_builtin;
5
6use crate::{gather_if_needed_async, BuiltinResult};
7
8use super::classification_linear::{
9    predict_classification_linear_object, CLASSIFICATION_LINEAR_CLASS,
10};
11use super::classification_tree::{predict_classification_tree_object, CLASSIFICATION_TREE_CLASS};
12use super::linear_model::{predict_invalid, predict_linear_model_dispatch, predict_type};
13
14#[runtime_builtin(
15    name = "predict",
16    category = "stats/ml",
17    summary = "Predict responses from a fitted model.",
18    keywords = "predict,fitlm,linear model,classification,deep learning,prediction",
19    type_resolver(predict_type),
20    descriptor(crate::builtins::stats::ml::linear_model::PREDICT_DESCRIPTOR),
21    builtin_path = "crate::builtins::stats::ml::predict"
22)]
23pub(crate) async fn predict_builtin(
24    model: Value,
25    xnew: Value,
26    rest: Vec<Value>,
27) -> BuiltinResult<Value> {
28    let model = gather(model)
29        .await
30        .map_err(|err| predict_invalid(err.message))?;
31    let xnew = gather(xnew)
32        .await
33        .map_err(|err| predict_invalid(err.message))?;
34    let rest = gather_all(rest)
35        .await
36        .map_err(|err| predict_invalid(err.message))?;
37    let output = match model {
38        Value::Object(object)
39            if crate::builtins::deep_learning::model::is_deep_learning_network_object(&object) =>
40        {
41            crate::builtins::deep_learning::model::predict_deep_learning_object(object, xnew, rest)?
42        }
43        Value::Object(object) if object.class_name == CLASSIFICATION_TREE_CLASS => {
44            predict_classification_tree_object(object, xnew, rest)?
45        }
46        Value::Object(object) if object.class_name == CLASSIFICATION_LINEAR_CLASS => {
47            predict_classification_linear_object(object, xnew, rest)?
48        }
49        other => predict_linear_model_dispatch(other, xnew, rest)?,
50    };
51    match crate::output_count::current_output_count() {
52        Some(0) => Ok(Value::OutputList(Vec::new())),
53        Some(1) => Ok(Value::OutputList(vec![output[0].clone()])),
54        Some(out_count) if out_count > output.len() => {
55            Err(predict_invalid("predict: too many output arguments"))
56        }
57        Some(out_count) => Ok(crate::output_count::output_list_with_padding(
58            out_count, output,
59        )),
60        None => Ok(output
61            .into_iter()
62            .next()
63            .expect("predict dispatch always returns at least one output")),
64    }
65}
66
67async fn gather(value: Value) -> BuiltinResult<Value> {
68    gather_if_needed_async(&value).await
69}
70
71async fn gather_all(values: Vec<Value>) -> BuiltinResult<Vec<Value>> {
72    let mut out = Vec::with_capacity(values.len());
73    for value in values {
74        out.push(gather(value).await?);
75    }
76    Ok(out)
77}