use nalgebra::{DMatrix, DVector};
use runmat_builtins::{
BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
CharArray, ObjectInstance, ResolveContext, StringArray, StructValue, Tensor, Type, Value,
};
use runmat_macros::runtime_builtin;
use crate::builtins::common::tensor;
use crate::builtins::stats::summary::distribution_math::{student_t_cdf_upper, student_t_inv};
use crate::builtins::table::{
is_tabular_object, table_from_columns, table_height, table_variable_names_from_object,
table_variables,
};
use crate::{build_runtime_error, gather_if_needed_async, BuiltinResult, RuntimeError};
const FITLM_NAME: &str = "fitlm";
const PREDICT_NAME: &str = "predict";
const LINEAR_MODEL_CLASS: &str = "LinearModel";
const EPS: f64 = 1.0e-12;
const OUTPUT_MDL: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
name: "mdl",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Required,
default: None,
description: "LinearModel object containing coefficients, fitted values, residuals, and summary statistics.",
}];
const OUTPUT_YPRED: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
name: "ypred",
ty: BuiltinParamType::NumericArray,
arity: BuiltinParamArity::Required,
default: None,
description: "Predicted response values.",
}];
const OUTPUT_YPRED_YCI: [BuiltinParamDescriptor; 2] = [
BuiltinParamDescriptor {
name: "ypred",
ty: BuiltinParamType::NumericArray,
arity: BuiltinParamArity::Required,
default: None,
description: "Predicted response values.",
},
BuiltinParamDescriptor {
name: "yci",
ty: BuiltinParamType::NumericArray,
arity: BuiltinParamArity::Required,
default: None,
description: "Pointwise confidence intervals for the predicted response.",
},
];
const OUTPUT_LABEL: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
name: "label",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Required,
default: None,
description: "Predicted class labels.",
}];
const OUTPUT_LABEL_SCORE_NODE_CNUM: [BuiltinParamDescriptor; 4] = [
BuiltinParamDescriptor {
name: "label",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Required,
default: None,
description: "Predicted class labels.",
},
BuiltinParamDescriptor {
name: "score",
ty: BuiltinParamType::NumericArray,
arity: BuiltinParamArity::Required,
default: None,
description: "Posterior probabilities for each class.",
},
BuiltinParamDescriptor {
name: "node",
ty: BuiltinParamType::NumericArray,
arity: BuiltinParamArity::Required,
default: None,
description: "One-based terminal node index for each prediction.",
},
BuiltinParamDescriptor {
name: "cnum",
ty: BuiltinParamType::NumericArray,
arity: BuiltinParamArity::Required,
default: None,
description: "One-based class number for each predicted label.",
},
];
const OUTPUT_LABEL_SCORE: [BuiltinParamDescriptor; 2] = [
BuiltinParamDescriptor {
name: "label",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Required,
default: None,
description: "Predicted class labels.",
},
BuiltinParamDescriptor {
name: "score",
ty: BuiltinParamType::NumericArray,
arity: BuiltinParamArity::Required,
default: None,
description: "Classification scores or posterior probabilities.",
},
];
const PARAM_TBL_OR_X: BuiltinParamDescriptor = BuiltinParamDescriptor {
name: "tblOrX",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Required,
default: None,
description: "Input table or predictor matrix.",
};
const PARAM_Y_OR_SPEC: BuiltinParamDescriptor = BuiltinParamDescriptor {
name: "yOrSpec",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Optional,
default: None,
description: "Response vector, response variable name, model formula, or model specification.",
};
const PARAM_REST: BuiltinParamDescriptor = BuiltinParamDescriptor {
name: "options",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Variadic,
default: None,
description: "Model specification and name-value options such as Intercept, ResponseVar, PredictorVars, VarNames, Exclude, and Weights.",
};
const PARAM_PREDICT_OPTIONS: BuiltinParamDescriptor = BuiltinParamDescriptor {
name: "options",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Variadic,
default: None,
description: "Prediction name-value options such as Alpha, Prediction, and Simultaneous.",
};
const PARAM_MDL: BuiltinParamDescriptor = BuiltinParamDescriptor {
name: "mdl",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Required,
default: None,
description: "Supported fitted model object, such as LinearModel from fitlm, ClassificationTree from fitctree, or ClassificationLinear from fitclinear.",
};
const PARAM_XNEW: BuiltinParamDescriptor = BuiltinParamDescriptor {
name: "Xnew",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Required,
default: None,
description: "New predictor table or matrix.",
};
const FITLM_INPUTS_X_Y: [BuiltinParamDescriptor; 2] = [PARAM_TBL_OR_X, PARAM_Y_OR_SPEC];
const FITLM_INPUTS_FULL: [BuiltinParamDescriptor; 3] =
[PARAM_TBL_OR_X, PARAM_Y_OR_SPEC, PARAM_REST];
const PREDICT_INPUTS: [BuiltinParamDescriptor; 2] = [PARAM_MDL, PARAM_XNEW];
const PREDICT_INPUTS_OPTIONS: [BuiltinParamDescriptor; 3] =
[PARAM_MDL, PARAM_XNEW, PARAM_PREDICT_OPTIONS];
const FITLM_SIGNATURES: [BuiltinSignatureDescriptor; 4] = [
BuiltinSignatureDescriptor {
label: "mdl = fitlm(tbl)",
inputs: &[PARAM_TBL_OR_X],
outputs: &OUTPUT_MDL,
},
BuiltinSignatureDescriptor {
label: "mdl = fitlm(tbl, modelspec)",
inputs: &FITLM_INPUTS_X_Y,
outputs: &OUTPUT_MDL,
},
BuiltinSignatureDescriptor {
label: "mdl = fitlm(X, y)",
inputs: &FITLM_INPUTS_X_Y,
outputs: &OUTPUT_MDL,
},
BuiltinSignatureDescriptor {
label: "mdl = fitlm(___, Name, Value)",
inputs: &FITLM_INPUTS_FULL,
outputs: &OUTPUT_MDL,
},
];
const PREDICT_SIGNATURES: [BuiltinSignatureDescriptor; 6] = [
BuiltinSignatureDescriptor {
label: "ypred = predict(mdl, Xnew)",
inputs: &PREDICT_INPUTS,
outputs: &OUTPUT_YPRED,
},
BuiltinSignatureDescriptor {
label: "[ypred, yci] = predict(mdl, Xnew)",
inputs: &PREDICT_INPUTS,
outputs: &OUTPUT_YPRED_YCI,
},
BuiltinSignatureDescriptor {
label: "[ypred, yci] = predict(mdl, Xnew, Name, Value)",
inputs: &PREDICT_INPUTS_OPTIONS,
outputs: &OUTPUT_YPRED_YCI,
},
BuiltinSignatureDescriptor {
label: "label = predict(tree, Xnew)",
inputs: &PREDICT_INPUTS,
outputs: &OUTPUT_LABEL,
},
BuiltinSignatureDescriptor {
label: "[label,score,node,cnum] = predict(tree, Xnew)",
inputs: &PREDICT_INPUTS,
outputs: &OUTPUT_LABEL_SCORE_NODE_CNUM,
},
BuiltinSignatureDescriptor {
label: "[label,score] = predict(linearClassifier, Xnew)",
inputs: &PREDICT_INPUTS_OPTIONS,
outputs: &OUTPUT_LABEL_SCORE,
},
];
const ERROR_FITLM_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.FITLM.INVALID_ARGUMENT",
identifier: Some("RunMat:fitlm:InvalidArgument"),
when: "Inputs, model specifications, dimensions, or name-value options are malformed or unsupported.",
message: "fitlm: invalid argument",
};
const ERROR_FITLM_NUMERICAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.FITLM.NUMERICAL",
identifier: Some("RunMat:fitlm:Numerical"),
when: "The linear model cannot be fit numerically.",
message: "fitlm: numerical failure",
};
const ERROR_FITLM_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.FITLM.INTERNAL",
identifier: Some("RunMat:fitlm:Internal"),
when: "RunMat cannot construct a LinearModel result.",
message: "fitlm: internal error",
};
const ERROR_PREDICT_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.PREDICT.INVALID_ARGUMENT",
identifier: Some("RunMat:predict:InvalidArgument"),
when: "The model, predictor data, or prediction options are malformed or unsupported.",
message: "predict: invalid argument",
};
const ERROR_PREDICT_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.PREDICT.INTERNAL",
identifier: Some("RunMat:predict:Internal"),
when: "RunMat cannot construct prediction outputs.",
message: "predict: internal error",
};
const FITLM_ERRORS: [BuiltinErrorDescriptor; 3] = [
ERROR_FITLM_INVALID_ARGUMENT,
ERROR_FITLM_NUMERICAL,
ERROR_FITLM_INTERNAL,
];
const PREDICT_ERRORS: [BuiltinErrorDescriptor; 2] =
[ERROR_PREDICT_INVALID_ARGUMENT, ERROR_PREDICT_INTERNAL];
pub const FITLM_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
signatures: &FITLM_SIGNATURES,
output_mode: BuiltinOutputMode::Fixed,
completion_policy: BuiltinCompletionPolicy::Public,
errors: &FITLM_ERRORS,
};
pub const PREDICT_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
signatures: &PREDICT_SIGNATURES,
output_mode: BuiltinOutputMode::ByRequestedOutputCount,
completion_policy: BuiltinCompletionPolicy::Public,
errors: &PREDICT_ERRORS,
};
fn fitlm_type(_args: &[Type], _ctx: &ResolveContext) -> Type {
Type::Unknown
}
pub(crate) fn predict_type(_args: &[Type], _ctx: &ResolveContext) -> Type {
Type::Unknown
}
fn fitlm_error(
message: impl Into<String>,
descriptor: &'static BuiltinErrorDescriptor,
) -> RuntimeError {
let mut builder = build_runtime_error(message).with_builtin(FITLM_NAME);
if let Some(identifier) = descriptor.identifier {
builder = builder.with_identifier(identifier);
}
builder.build()
}
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()
}
fn fitlm_invalid(message: impl Into<String>) -> RuntimeError {
fitlm_error(message, &ERROR_FITLM_INVALID_ARGUMENT)
}
fn fitlm_numerical(message: impl Into<String>) -> RuntimeError {
fitlm_error(message, &ERROR_FITLM_NUMERICAL)
}
fn fitlm_internal(message: impl Into<String>) -> RuntimeError {
fitlm_error(message, &ERROR_FITLM_INTERNAL)
}
pub(crate) 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)
}
#[derive(Clone, Debug)]
struct FitOptions {
intercept: bool,
predictor_vars: Option<Vec<String>>,
response_var: Option<String>,
var_names: Option<Vec<String>>,
exclude: Option<ExcludeSpec>,
weights: Option<WeightSpec>,
}
impl Default for FitOptions {
fn default() -> Self {
Self {
intercept: true,
predictor_vars: None,
response_var: None,
var_names: None,
exclude: None,
weights: None,
}
}
}
#[derive(Clone, Debug)]
enum ExcludeSpec {
Mask(Vec<bool>),
Indices(Vec<usize>),
}
#[derive(Clone, Debug)]
enum WeightSpec {
Values(Vec<f64>),
Variable(String),
}
#[derive(Clone, Debug)]
struct PredictOptions {
alpha: f64,
prediction: PredictionKind,
}
impl Default for PredictOptions {
fn default() -> Self {
Self {
alpha: 0.05,
prediction: PredictionKind::Curve,
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum PredictionKind {
Curve,
Observation,
}
#[derive(Clone, Debug)]
struct FitSpec {
x: Tensor,
y: Vec<f64>,
predictor_names: Vec<String>,
response_name: String,
terms: Vec<Vec<u32>>,
coefficient_names: Vec<String>,
model_formula: String,
exclude: Option<ExcludeSpec>,
weights: Option<Vec<f64>>,
}
#[derive(Clone, Debug)]
struct CleanData {
design: DMatrix<f64>,
weighted_design: DMatrix<f64>,
y: DVector<f64>,
weighted_y: DVector<f64>,
weights: Vec<f64>,
}
#[derive(Clone, Debug)]
struct OlsFit {
beta: Vec<f64>,
fitted: Vec<f64>,
residuals: Vec<f64>,
covariance: Vec<f64>,
se: Vec<f64>,
tstat: Vec<f64>,
pvalue: Vec<f64>,
sse: f64,
mse: f64,
rmse: f64,
r_squared: f64,
adjusted_r_squared: f64,
dfe: f64,
rank: usize,
observations: usize,
}
#[runtime_builtin(
name = "fitlm",
category = "stats/ml",
summary = "Fit an ordinary least-squares linear regression model.",
keywords = "fitlm,linear model,regression,least squares,statistics,machine learning",
type_resolver(fitlm_type),
descriptor(crate::builtins::stats::ml::linear_model::FITLM_DESCRIPTOR),
builtin_path = "crate::builtins::stats::ml::linear_model"
)]
async fn fitlm_builtin(first: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
let first = gather(first)
.await
.map_err(|err| fitlm_invalid(err.message))?;
let rest = gather_all(rest)
.await
.map_err(|err| fitlm_invalid(err.message))?;
let spec = build_fit_spec(first, rest)?;
let fit = fit_ols(&spec)?;
model_object(spec, fit).map(Value::Object)
}
#[cfg(test)]
pub(crate) async fn predict_builtin(
model: Value,
xnew: Value,
rest: Vec<Value>,
) -> BuiltinResult<Value> {
super::predict::predict_builtin(model, xnew, rest).await
}
pub(crate) fn predict_linear_model_dispatch(
model: Value,
xnew: Value,
rest: Vec<Value>,
) -> BuiltinResult<Vec<Value>> {
let options = parse_predict_options(rest)?;
let output = predict_linear_model(model, xnew, options)?;
Ok(vec![output.0, output.1])
}
async fn gather(value: Value) -> BuiltinResult<Value> {
gather_if_needed_async(&value)
.await
.map_err(|err| fitlm_invalid(format!("fitlm: {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(fitlm_invalid(
"fitlm: 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(FITLM_NAME, x_value)
.map_err(|err| fitlm_invalid(format!("fitlm: {err}")))?;
if x.shape.len() > 2 {
return Err(fitlm_invalid("fitlm: X must be a 2-D numeric matrix"));
}
let y_tensor = tensor::value_into_tensor_for(FITLM_NAME, y_value)
.map_err(|err| fitlm_invalid(format!("fitlm: {err}")))?;
let y = vector_values(&y_tensor, "y")?;
if y.len() != x.rows {
return Err(fitlm_invalid(
"fitlm: length(y) must match the number of rows in X",
));
}
let cols = x.cols;
let mut options = FitOptions::default();
let mut model_spec = ModelSpec::Linear;
let mut index = 0;
if index < rest.len() && !is_name_value_option(&rest[index]) {
model_spec = parse_model_spec(&rest[index], cols)?;
index += 1;
}
parse_fit_name_values(&rest[index..], &mut options)?;
if options.response_var.is_some() {
return Err(fitlm_invalid(
"fitlm: ResponseVar is only valid for table input",
));
}
if options.predictor_vars.is_some() {
return Err(fitlm_invalid(
"fitlm: PredictorVars is only valid for table input",
));
}
let (predictor_names, response_name) = if let Some(names) = options.var_names.clone() {
if names.len() != cols + 1 {
return Err(fitlm_invalid(format!(
"fitlm: VarNames must contain {} predictor names and one response name",
cols
)));
}
(names[..cols].to_vec(), names[cols].clone())
} else {
(
(1..=cols).map(|idx| format!("x{idx}")).collect(),
"y".to_string(),
)
};
let terms = model_terms(model_spec, cols, options.intercept);
let coefficient_names = coefficient_names(&terms, &predictor_names);
let model_formula = formula_text(&response_name, &coefficient_names);
let weights = match options.weights {
Some(WeightSpec::Values(values)) => Some(values),
Some(WeightSpec::Variable(_)) => {
return Err(fitlm_invalid(
"fitlm: weights variable names require table input",
));
}
None => None,
};
Ok(FitSpec {
x,
y,
predictor_names,
response_name,
terms,
coefficient_names,
model_formula,
exclude: options.exclude,
weights,
})
}
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| fitlm_invalid(format!("fitlm: {err}")))?;
if variable_names.len() < 2 {
return Err(fitlm_invalid(
"fitlm: table input must contain at least one predictor and one response",
));
}
let variables =
table_variables(&object).map_err(|err| fitlm_invalid(format!("fitlm: {err}")))?;
let mut options = FitOptions::default();
let mut model_spec = ModelSpec::Linear;
let mut formula: Option<String> = None;
let mut response_arg: Option<String> = None;
let mut index = 0;
if index < rest.len() && !is_name_value_option(&rest[index]) {
if let Some(text) = scalar_text(&rest[index]) {
if text.contains('~') {
formula = Some(text);
} else if is_known_model_spec(&text) {
model_spec = parse_model_spec_text(&text, variable_names.len() - 1)?;
} else {
response_arg = Some(text);
}
} else {
model_spec = parse_model_spec(&rest[index], variable_names.len() - 1)?;
}
index += 1;
}
parse_fit_name_values(&rest[index..], &mut options)?;
let response_name = formula
.as_ref()
.and_then(|text| text.split_once('~').map(|(lhs, _)| lhs.trim().to_string()))
.or(options.response_var.clone())
.or(response_arg)
.unwrap_or_else(|| variable_names.last().cloned().unwrap());
if !variable_names.iter().any(|name| name == &response_name) {
return Err(fitlm_invalid(format!(
"fitlm: response variable '{response_name}' is not in the table"
)));
}
let mut predictor_names = if let Some(text) = &formula {
parse_formula_predictors(text, &variable_names, &response_name, &mut options)?
} else if let Some(names) = options.predictor_vars.clone() {
names
} else {
variable_names
.iter()
.filter(|name| *name != &response_name)
.cloned()
.collect()
};
let weights = match options.weights.clone() {
Some(WeightSpec::Variable(name)) => {
let value = variables.fields.get(&name).ok_or_else(|| {
fitlm_invalid(format!(
"fitlm: weights variable '{name}' is not in the table"
))
})?;
predictor_names.retain(|predictor| predictor != &name);
Some(numeric_vector(value, "Weights")?)
}
Some(WeightSpec::Values(values)) => Some(values),
None => None,
};
if predictor_names.is_empty() && !options.intercept {
return Err(fitlm_invalid("fitlm: model must contain at least one term"));
}
for name in &predictor_names {
if !variable_names.iter().any(|candidate| candidate == name) {
return Err(fitlm_invalid(format!(
"fitlm: predictor variable '{name}' is not in the table"
)));
}
}
let y_value = variables
.fields
.get(&response_name)
.ok_or_else(|| fitlm_invalid("fitlm: response variable missing from table"))?;
let y_tensor = tensor::value_into_tensor_for(FITLM_NAME, y_value.clone())
.map_err(|err| fitlm_invalid(format!("fitlm: {err}")))?;
let y = vector_values(&y_tensor, &response_name)?;
let x = if predictor_names.is_empty() {
Tensor::new(Vec::new(), vec![y.len(), 0])
.map_err(|err| fitlm_internal(format!("fitlm: {err}")))?
} else {
let mut columns = Vec::with_capacity(predictor_names.len());
for name in &predictor_names {
let value = variables
.fields
.get(name)
.ok_or_else(|| fitlm_invalid(format!("fitlm: predictor '{name}' missing")))?;
let tensor = tensor::value_into_tensor_for(FITLM_NAME, value.clone())
.map_err(|_| fitlm_invalid(format!("fitlm: predictor '{name}' must be numeric")))?;
let values = vector_values(&tensor, name)?;
if values.len() != y.len() {
return Err(fitlm_invalid(format!(
"fitlm: predictor '{name}' height does not match response"
)));
}
columns.push(values);
}
columns_to_tensor(columns)?
};
let terms = if let Some(text) = &formula {
parse_formula_terms(text, &predictor_names, options.intercept)?
} else {
model_terms(model_spec, predictor_names.len(), options.intercept)
};
let coefficient_names = coefficient_names(&terms, &predictor_names);
let model_formula = formula
.as_ref()
.cloned()
.unwrap_or_else(|| formula_text(&response_name, &coefficient_names));
Ok(FitSpec {
x,
y,
predictor_names,
response_name,
terms,
coefficient_names,
model_formula,
exclude: options.exclude,
weights,
})
}
fn parse_fit_name_values(values: &[Value], options: &mut FitOptions) -> BuiltinResult<()> {
if !values.len().is_multiple_of(2) {
return Err(fitlm_invalid("fitlm: name-value options must be paired"));
}
let mut index = 0;
while index < values.len() {
let name = scalar_text(&values[index])
.ok_or_else(|| fitlm_invalid("fitlm: option name must be text"))?;
let value = &values[index + 1];
if name.eq_ignore_ascii_case("Intercept") {
options.intercept = bool_scalar(value, "Intercept")?;
} else if name.eq_ignore_ascii_case("ResponseVar") {
options.response_var = Some(
scalar_text(value)
.ok_or_else(|| fitlm_invalid("fitlm: ResponseVar must be text"))?,
);
} else if name.eq_ignore_ascii_case("PredictorVars") {
options.predictor_vars = Some(text_list(value, "PredictorVars")?);
} else if name.eq_ignore_ascii_case("VarNames") {
options.var_names = Some(text_list(value, "VarNames")?);
} else if name.eq_ignore_ascii_case("Exclude") {
options.exclude = Some(exclude_spec(value)?);
} else if name.eq_ignore_ascii_case("Weights") {
options.weights = Some(if let Some(name) = scalar_text(value) {
WeightSpec::Variable(name)
} else {
WeightSpec::Values(numeric_vector(value, "Weights")?)
});
} else if name.eq_ignore_ascii_case("RobustOpts")
|| name.eq_ignore_ascii_case("CategoricalVars")
{
return Err(fitlm_invalid(format!(
"fitlm: option '{name}' is not supported yet"
)));
} else {
return Err(fitlm_invalid(format!("fitlm: unknown option '{name}'")));
}
index += 2;
}
Ok(())
}
fn parse_predict_options(values: Vec<Value>) -> BuiltinResult<PredictOptions> {
if !values.len().is_multiple_of(2) {
return Err(predict_invalid(
"predict: name-value options must be paired",
));
}
let mut options = PredictOptions::default();
let mut index = 0;
while index < values.len() {
let name = scalar_text(&values[index])
.ok_or_else(|| predict_invalid("predict: option name must be text"))?;
let value = &values[index + 1];
if name.eq_ignore_ascii_case("Alpha") {
let alpha =
numeric_scalar(value, "Alpha").map_err(|err| predict_invalid(err.message))?;
if !(0.0..1.0).contains(&alpha) {
return Err(predict_invalid("predict: Alpha must be between 0 and 1"));
}
options.alpha = alpha;
} else if name.eq_ignore_ascii_case("Prediction") {
let text = scalar_text(value)
.ok_or_else(|| predict_invalid("predict: Prediction must be text"))?;
if text.eq_ignore_ascii_case("curve") {
options.prediction = PredictionKind::Curve;
} else if text.eq_ignore_ascii_case("observation") {
options.prediction = PredictionKind::Observation;
} else {
return Err(predict_invalid(
"predict: Prediction must be 'curve' or 'observation'",
));
}
} else if name.eq_ignore_ascii_case("Simultaneous") {
let simultaneous =
bool_scalar(value, "Simultaneous").map_err(|err| predict_invalid(err.message))?;
if simultaneous {
return Err(predict_invalid(
"predict: Simultaneous=true intervals are not supported yet",
));
}
} else {
return Err(predict_invalid(format!("predict: unknown option '{name}'")));
}
index += 2;
}
Ok(options)
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ModelSpec {
Constant,
Linear,
Interactions,
PureQuadratic,
Quadratic,
}
fn parse_model_spec(value: &Value, predictors: usize) -> BuiltinResult<ModelSpec> {
if let Some(text) = scalar_text(value) {
parse_model_spec_text(&text, predictors)
} else {
Err(fitlm_invalid(
"fitlm: numeric terms matrices are not supported yet",
))
}
}
fn parse_model_spec_text(text: &str, _predictors: usize) -> BuiltinResult<ModelSpec> {
if text.eq_ignore_ascii_case("constant") {
Ok(ModelSpec::Constant)
} else if text.eq_ignore_ascii_case("linear") {
Ok(ModelSpec::Linear)
} else if text.eq_ignore_ascii_case("interactions") {
Ok(ModelSpec::Interactions)
} else if text.eq_ignore_ascii_case("purequadratic") {
Ok(ModelSpec::PureQuadratic)
} else if text.eq_ignore_ascii_case("quadratic") {
Ok(ModelSpec::Quadratic)
} else {
Err(fitlm_invalid(format!(
"fitlm: unsupported model specification '{text}'"
)))
}
}
fn is_known_model_spec(text: &str) -> bool {
matches!(
text.to_ascii_lowercase().as_str(),
"constant" | "linear" | "interactions" | "purequadratic" | "quadratic"
)
}
fn model_terms(spec: ModelSpec, predictors: usize, intercept: bool) -> Vec<Vec<u32>> {
let mut terms = Vec::new();
if intercept {
terms.push(vec![0; predictors]);
}
if matches!(
spec,
ModelSpec::Linear
| ModelSpec::Interactions
| ModelSpec::PureQuadratic
| ModelSpec::Quadratic
) {
for col in 0..predictors {
let mut term = vec![0; predictors];
term[col] = 1;
terms.push(term);
}
}
if matches!(spec, ModelSpec::Interactions | ModelSpec::Quadratic) {
for left in 0..predictors {
for right in left + 1..predictors {
let mut term = vec![0; predictors];
term[left] = 1;
term[right] = 1;
terms.push(term);
}
}
}
if matches!(spec, ModelSpec::PureQuadratic | ModelSpec::Quadratic) {
for col in 0..predictors {
let mut term = vec![0; predictors];
term[col] = 2;
terms.push(term);
}
}
terms
}
fn parse_formula_predictors(
formula: &str,
table_names: &[String],
response_name: &str,
options: &mut FitOptions,
) -> BuiltinResult<Vec<String>> {
let (_, rhs) = formula
.split_once('~')
.ok_or_else(|| fitlm_invalid("fitlm: formula must contain '~'"))?;
let mut out = Vec::new();
for raw in formula_rhs_tokens(rhs) {
let term = raw.trim();
if term.is_empty() || term == "1" {
continue;
}
if term == "0" || term == "-1" {
options.intercept = false;
continue;
}
if term.contains(':') || term.contains('*') || term.contains('^') {
return Err(fitlm_invalid(
"fitlm: formula interactions and powers are not supported yet",
));
}
if term == response_name {
return Err(fitlm_invalid("fitlm: response cannot also be a predictor"));
}
if !table_names.iter().any(|name| name == term) {
return Err(fitlm_invalid(format!(
"fitlm: formula predictor '{term}' is not in the table"
)));
}
if !out.iter().any(|name| name == term) {
out.push(term.to_string());
}
}
Ok(out)
}
fn parse_formula_terms(
formula: &str,
predictor_names: &[String],
intercept: bool,
) -> BuiltinResult<Vec<Vec<u32>>> {
let (_, rhs) = formula
.split_once('~')
.ok_or_else(|| fitlm_invalid("fitlm: formula must contain '~'"))?;
let mut terms = Vec::new();
if intercept {
terms.push(vec![0; predictor_names.len()]);
}
for raw in formula_rhs_tokens(rhs) {
let term = raw.trim();
if term.is_empty() || term == "1" || term == "0" || term == "-1" {
continue;
}
let position = predictor_names
.iter()
.position(|name| name == term)
.ok_or_else(|| fitlm_invalid(format!("fitlm: unknown formula term '{term}'")))?;
let mut exponents = vec![0; predictor_names.len()];
exponents[position] = 1;
terms.push(exponents);
}
Ok(terms)
}
fn formula_rhs_tokens(rhs: &str) -> Vec<String> {
let mut out = Vec::new();
let mut current = String::new();
let mut sign = 1_i8;
let mut saw_token = false;
for ch in rhs.chars() {
match ch {
'+' | '-' => {
if saw_token || !current.trim().is_empty() {
let token = current.trim();
if !token.is_empty() {
out.push(if sign < 0 {
format!("-{token}")
} else {
token.to_string()
});
}
current.clear();
saw_token = false;
}
sign = if ch == '-' { -1 } else { 1 };
}
_ => {
if !ch.is_whitespace() {
saw_token = true;
}
current.push(ch);
}
}
}
let token = current.trim();
if !token.is_empty() {
out.push(if sign < 0 {
format!("-{token}")
} else {
token.to_string()
});
}
out
}
fn coefficient_names(terms: &[Vec<u32>], predictor_names: &[String]) -> Vec<String> {
terms
.iter()
.map(|term| {
if term.iter().all(|power| *power == 0) {
"(Intercept)".to_string()
} else {
let mut factors = Vec::new();
for (name, power) in predictor_names.iter().zip(term) {
match *power {
0 => {}
1 => factors.push(name.clone()),
n => factors.push(format!("{name}^{n}")),
}
}
factors.join(":")
}
})
.collect()
}
fn formula_text(response_name: &str, coefficient_names: &[String]) -> String {
let rhs = if coefficient_names.is_empty() {
"0".to_string()
} else {
coefficient_names
.iter()
.map(|name| {
if name == "(Intercept)" {
"1".to_string()
} else {
name.clone()
}
})
.collect::<Vec<_>>()
.join(" + ")
};
format!("{response_name} ~ {rhs}")
}
fn fit_ols(spec: &FitSpec) -> BuiltinResult<OlsFit> {
let clean = clean_data(spec)?;
let observations = clean.y.len();
if observations == 0 {
return Err(fitlm_invalid(
"fitlm: at least one complete observation is required",
));
}
let svd = clean.weighted_design.clone().svd(true, true);
let u = svd
.u
.ok_or_else(|| fitlm_numerical("fitlm: SVD did not return left singular vectors"))?;
let v_t = svd
.v_t
.ok_or_else(|| fitlm_numerical("fitlm: SVD did not return right singular vectors"))?;
let singular_values = svd.singular_values.iter().copied().collect::<Vec<_>>();
let largest = singular_values.iter().copied().fold(0.0_f64, f64::max);
let tolerance = (observations.max(spec.terms.len()) as f64) * f64::EPSILON * largest.max(1.0);
let rank = singular_values
.iter()
.filter(|value| value.abs() > tolerance)
.count();
let mut beta = DVector::zeros(spec.terms.len());
let mut inv_xtx = DMatrix::<f64>::zeros(spec.terms.len(), spec.terms.len());
for (index, singular_value) in singular_values.iter().copied().enumerate() {
if singular_value.abs() <= tolerance {
continue;
}
let uty = u.column(index).dot(&clean.weighted_y);
let contribution = uty / singular_value;
for row in 0..spec.terms.len() {
beta[row] += v_t[(index, row)] * contribution;
}
let inv_s2 = 1.0 / (singular_value * singular_value);
for row in 0..spec.terms.len() {
for col in 0..spec.terms.len() {
inv_xtx[(row, col)] += v_t[(index, row)] * v_t[(index, col)] * inv_s2;
}
}
}
let fitted_vec = &clean.design * β
let residual_vec = &clean.y - &fitted_vec;
let sse = residual_vec
.iter()
.zip(clean.weights.iter())
.map(|(residual, weight)| residual * residual * weight)
.sum::<f64>();
let weight_sum = clean.weights.iter().sum::<f64>();
let y_mean = if weight_sum > 0.0 {
clean
.y
.iter()
.zip(clean.weights.iter())
.map(|(y, weight)| y * weight)
.sum::<f64>()
/ weight_sum
} else {
f64::NAN
};
let sst = clean
.y
.iter()
.zip(clean.weights.iter())
.map(|(y, weight)| {
let diff = y - y_mean;
diff * diff * weight
})
.sum::<f64>();
let dfe = observations as f64 - rank as f64;
let mse = if dfe > 0.0 { sse / dfe } else { f64::NAN };
let rmse = mse.sqrt();
let r_squared = if sst > EPS {
1.0 - sse / sst
} else if sse <= EPS {
1.0
} else {
f64::NAN
};
let adjusted_r_squared = if dfe > 0.0 && observations > 1 && r_squared.is_finite() {
1.0 - (1.0 - r_squared) * ((observations - 1) as f64) / dfe
} else {
f64::NAN
};
let covariance = inv_xtx.iter().map(|value| value * mse).collect::<Vec<_>>();
let mut se = Vec::with_capacity(spec.terms.len());
let mut tstat = Vec::with_capacity(spec.terms.len());
let mut pvalue = Vec::with_capacity(spec.terms.len());
for i in 0..spec.terms.len() {
let variance = covariance[i + i * spec.terms.len()];
let std_error = if variance >= 0.0 {
variance.sqrt()
} else {
f64::NAN
};
let t = beta[i] / std_error;
let p = if dfe > 0.0 {
(2.0 * student_t_cdf_upper(t.abs(), dfe)).min(1.0)
} else {
f64::NAN
};
se.push(std_error);
tstat.push(t);
pvalue.push(p);
}
Ok(OlsFit {
beta: beta.iter().copied().collect(),
fitted: fitted_vec.iter().copied().collect(),
residuals: residual_vec.iter().copied().collect(),
covariance,
se,
tstat,
pvalue,
sse,
mse,
rmse,
r_squared,
adjusted_r_squared,
dfe,
rank,
observations,
})
}
fn clean_data(spec: &FitSpec) -> BuiltinResult<CleanData> {
let rows = spec.x.rows;
let cols = spec.x.cols;
let exclude_mask = exclude_mask(spec.exclude.as_ref(), rows)?;
let weights = spec.weights.as_ref();
if let Some(weights) = weights {
if weights.len() != rows {
return Err(fitlm_invalid(
"fitlm: Weights length must match the number of observations",
));
}
}
let mut design_rows = Vec::new();
let mut y_values = Vec::new();
let mut clean_weights = Vec::new();
for row in 0..rows {
if exclude_mask.as_ref().is_some_and(|mask| mask[row]) {
continue;
}
let y = spec.y[row];
let weight = weights.map(|w| w[row]).unwrap_or(1.0);
if weight < 0.0 || !weight.is_finite() {
return Err(fitlm_invalid(
"fitlm: weights must be nonnegative finite values",
));
}
let mut raw = Vec::with_capacity(cols);
let mut has_nan = y.is_nan();
let mut has_inf = y.is_infinite();
for col in 0..cols {
let value = x_value(&spec.x, row, col);
has_nan |= value.is_nan();
has_inf |= value.is_infinite();
raw.push(value);
}
if has_inf {
return Err(fitlm_invalid("fitlm: X and y must not contain Inf values"));
}
if has_nan || weight == 0.0 {
continue;
}
for term in &spec.terms {
design_rows.push(evaluate_term(&raw, term));
}
y_values.push(y);
clean_weights.push(weight);
}
let observations = y_values.len();
let predictors = spec.terms.len();
let design = DMatrix::from_row_slice(observations, predictors, &design_rows);
let mut weighted_rows = design_rows;
let mut weighted_y = y_values.clone();
for row in 0..observations {
let scale = clean_weights[row].sqrt();
weighted_y[row] *= scale;
for col in 0..predictors {
weighted_rows[row * predictors + col] *= scale;
}
}
Ok(CleanData {
design,
weighted_design: DMatrix::from_row_slice(observations, predictors, &weighted_rows),
y: DVector::from_column_slice(&y_values),
weighted_y: DVector::from_column_slice(&weighted_y),
weights: clean_weights,
})
}
fn evaluate_term(row: &[f64], term: &[u32]) -> f64 {
let mut value = 1.0;
for (x, power) in row.iter().zip(term) {
match *power {
0 => {}
1 => value *= *x,
2 => value *= x * x,
n => value *= x.powi(n as i32),
}
}
value
}
fn model_object(spec: FitSpec, fit: OlsFit) -> BuiltinResult<ObjectInstance> {
let mut object = ObjectInstance::new(LINEAR_MODEL_CLASS.to_string());
object.properties.insert(
"Coefficients".to_string(),
coefficients_table(&spec.coefficient_names, &fit)?,
);
object.properties.insert(
"CoefficientNames".to_string(),
string_row(spec.coefficient_names.clone())?,
);
object
.properties
.insert("Formula".to_string(), Value::String(spec.model_formula));
object.properties.insert(
"NumObservations".to_string(),
Value::Num(fit.observations as f64),
);
object
.properties
.insert("DFE".to_string(), Value::Num(fit.dfe));
object
.properties
.insert("RMSE".to_string(), Value::Num(fit.rmse));
let mut rsquared = StructValue::new();
rsquared.insert("Ordinary", Value::Num(fit.r_squared));
rsquared.insert("Adjusted", Value::Num(fit.adjusted_r_squared));
object
.properties
.insert("Rsquared".to_string(), Value::Struct(rsquared));
let mut model_criterion = StructValue::new();
model_criterion.insert("MSE", Value::Num(fit.mse));
model_criterion.insert("SSE", Value::Num(fit.sse));
object
.properties
.insert("ModelCriterion".to_string(), Value::Struct(model_criterion));
object.properties.insert(
"Fitted".to_string(),
Value::Tensor(
Tensor::new(fit.fitted.clone(), vec![fit.fitted.len(), 1]).map_err(|err| {
fitlm_internal(format!("fitlm: failed to construct fitted values: {err}"))
})?,
),
);
let mut residuals = StructValue::new();
residuals.insert(
"Raw",
Value::Tensor(
Tensor::new(fit.residuals.clone(), vec![fit.residuals.len(), 1]).map_err(|err| {
fitlm_internal(format!("fitlm: failed to construct residuals: {err}"))
})?,
),
);
object
.properties
.insert("Residuals".to_string(), Value::Struct(residuals));
object.properties.insert(
"PredictorNames".to_string(),
string_row(spec.predictor_names.clone())?,
);
object.properties.insert(
"ResponseName".to_string(),
Value::String(spec.response_name),
);
object.properties.insert(
"NumPredictors".to_string(),
Value::Num(spec.predictor_names.len() as f64),
);
object.properties.insert(
"HasIntercept".to_string(),
Value::Bool(has_intercept(&spec.terms)),
);
object
.properties
.insert("DesignRank".to_string(), Value::Num(fit.rank as f64));
object
.properties
.insert("__RunMatTerms".to_string(), terms_tensor(&spec.terms)?);
object
.properties
.insert("__RunMatBeta".to_string(), numeric_column(&fit.beta)?);
object.properties.insert(
"__RunMatCovariance".to_string(),
Value::Tensor(
Tensor::new(
fit.covariance.clone(),
vec![spec.terms.len(), spec.terms.len()],
)
.map_err(|err| fitlm_internal(format!("fitlm: covariance shape error: {err}")))?,
),
);
object
.properties
.insert("__RunMatMSE".to_string(), Value::Num(fit.mse));
object
.properties
.insert("__RunMatDFE".to_string(), Value::Num(fit.dfe));
Ok(object)
}
fn coefficients_table(names: &[String], fit: &OlsFit) -> BuiltinResult<Value> {
let columns = vec![
numeric_column(&fit.beta)?,
numeric_column(&fit.se)?,
numeric_column(&fit.tstat)?,
numeric_column(&fit.pvalue)?,
];
let table = table_from_columns(
vec![
"Estimate".to_string(),
"SE".to_string(),
"tStat".to_string(),
"pValue".to_string(),
],
columns,
)
.map_err(|err| fitlm_internal(format!("fitlm: {err}")))?;
let mut object = match table {
Value::Object(object) => object,
_ => unreachable!("table_from_columns returns object"),
};
let row_names = string_column(names.to_vec())?;
if let Some(Value::Struct(props)) = object.properties.get_mut("Properties") {
props.insert("RowNames", row_names.clone());
}
if let Some(Value::Struct(props)) = object.properties.get_mut("__table_properties") {
props.insert("RowNames", row_names.clone());
}
object
.properties
.insert("__CoefficientNames".to_string(), row_names);
Ok(Value::Object(object))
}
fn predict_linear_model(
model: Value,
xnew: Value,
options: PredictOptions,
) -> BuiltinResult<(Value, Value)> {
let object = match model {
Value::Object(object) if object.class_name == LINEAR_MODEL_CLASS => object,
other => {
return Err(predict_invalid(format!(
"predict: expected LinearModel object, got {other:?}"
)));
}
};
let terms = terms_from_object(&object)?;
let beta = numeric_property(&object, "__RunMatBeta")?;
let covariance = matrix_property(&object, "__RunMatCovariance")?;
let predictor_names = string_property(&object, "PredictorNames")?;
let x = predictors_for_prediction(xnew, &predictor_names)?;
let design = design_matrix(&x, &terms);
let rows = x.rows;
let cols = terms.len();
if beta.len() != cols {
return Err(predict_invalid(
"predict: model coefficient length is invalid",
));
}
let mut ypred = Vec::with_capacity(rows);
let mut ci_low = Vec::with_capacity(rows);
let mut ci_high = Vec::with_capacity(rows);
let mse = numeric_scalar_property(&object, "__RunMatMSE")?;
let dfe = numeric_scalar_property(&object, "__RunMatDFE")?;
let crit = if dfe > 0.0 {
student_t_inv(1.0 - options.alpha / 2.0, dfe)
} else {
f64::NAN
};
for row in 0..rows {
let mut y = 0.0;
for col in 0..cols {
y += design[row * cols + col] * beta[col];
}
ypred.push(y);
let leverage = leverage_for_row(&design[row * cols..row * cols + cols], &covariance, cols);
let variance = match options.prediction {
PredictionKind::Curve => leverage,
PredictionKind::Observation => leverage + mse,
};
let delta = if variance >= 0.0 {
crit * variance.sqrt()
} else {
f64::NAN
};
ci_low.push(y - delta);
ci_high.push(y + delta);
}
let mut intervals = ci_low;
intervals.extend(ci_high);
let y_value = Value::Tensor(
Tensor::new(ypred, vec![rows, 1])
.map_err(|err| predict_internal(format!("predict: {err}")))?,
);
let ci_value = Value::Tensor(
Tensor::new(intervals, vec![rows, 2])
.map_err(|err| predict_internal(format!("predict: {err}")))?,
);
Ok((y_value, ci_value))
}
fn predictors_for_prediction(value: Value, predictor_names: &[String]) -> BuiltinResult<Tensor> {
if let Value::Object(object) = &value {
if is_tabular_object(object) {
if predictor_names.is_empty() {
let rows = table_height(object)
.map_err(|err| predict_invalid(format!("predict: {err}")))?;
return Tensor::new(Vec::new(), vec![rows, 0])
.map_err(|err| predict_internal(format!("predict: {err}")));
}
let variables = table_variables(object)
.map_err(|err| predict_invalid(format!("predict: {err}")))?;
let mut columns = Vec::with_capacity(predictor_names.len());
let mut height = None;
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"))
})?;
let values = vector_values_predict(&tensor, name)?;
if let Some(expected) = height {
if values.len() != expected {
return Err(predict_invalid(
"predict: predictor table variables must have the same height",
));
}
} else {
height = Some(values.len());
}
columns.push(values);
}
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: Xnew must have {} predictor columns",
predictor_names.len()
)));
}
Ok(tensor)
}
fn design_matrix(x: &Tensor, terms: &[Vec<u32>]) -> Vec<f64> {
let mut out = Vec::with_capacity(x.rows * terms.len());
for row in 0..x.rows {
let raw = (0..x.cols)
.map(|col| x_value(x, row, col))
.collect::<Vec<_>>();
for term in terms {
out.push(evaluate_term(&raw, term));
}
}
out
}
fn leverage_for_row(row: &[f64], covariance: &[f64], cols: usize) -> f64 {
let mut value = 0.0;
for i in 0..cols {
for j in 0..cols {
value += row[i] * covariance[i + j * cols] * row[j];
}
}
value
}
fn terms_from_object(object: &ObjectInstance) -> BuiltinResult<Vec<Vec<u32>>> {
let tensor = match object.properties.get("__RunMatTerms") {
Some(Value::Tensor(tensor)) => tensor,
_ => return Err(predict_invalid("predict: model is missing term metadata")),
};
let mut terms = Vec::with_capacity(tensor.rows);
for row in 0..tensor.rows {
let mut term = Vec::with_capacity(tensor.cols);
for col in 0..tensor.cols {
let value = x_value(tensor, row, col);
if value < 0.0 || value.fract().abs() > EPS {
return Err(predict_invalid("predict: model term metadata is invalid"));
}
term.push(value as u32);
}
terms.push(term);
}
Ok(terms)
}
fn numeric_property(object: &ObjectInstance, name: &str) -> BuiltinResult<Vec<f64>> {
match object.properties.get(name) {
Some(Value::Tensor(tensor)) => Ok(tensor.data.clone()),
Some(Value::Num(value)) => Ok(vec![*value]),
_ => Err(predict_invalid(format!(
"predict: model is missing numeric property '{name}'"
))),
}
}
fn matrix_property(object: &ObjectInstance, name: &str) -> BuiltinResult<Vec<f64>> {
match object.properties.get(name) {
Some(Value::Tensor(tensor)) if tensor.rows == tensor.cols => Ok(tensor.data.clone()),
_ => Err(predict_invalid(format!(
"predict: model is missing matrix property '{name}'"
))),
}
}
fn numeric_scalar_property(object: &ObjectInstance, name: &str) -> BuiltinResult<f64> {
match object.properties.get(name) {
Some(Value::Num(value)) => Ok(*value),
Some(Value::Tensor(tensor)) if tensor.data.len() == 1 => Ok(tensor.data[0]),
_ => Err(predict_invalid(format!(
"predict: model is missing scalar property '{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::Cell(cell)) => cell
.data
.iter()
.map(|value| {
scalar_text(value).ok_or_else(|| {
predict_invalid(format!("predict: property '{name}' is invalid"))
})
})
.collect(),
_ => Err(predict_invalid(format!(
"predict: model is missing string property '{name}'"
))),
}
}
fn vector_values(tensor: &Tensor, label: &str) -> BuiltinResult<Vec<f64>> {
if tensor.shape.len() > 2 || !(tensor.rows == 1 || tensor.cols == 1) {
return Err(fitlm_invalid(format!("fitlm: {label} must be a vector")));
}
Ok(tensor.data.clone())
}
fn vector_values_predict(tensor: &Tensor, label: &str) -> BuiltinResult<Vec<f64>> {
if tensor.shape.len() > 2 || !(tensor.rows == 1 || tensor.cols == 1) {
return Err(predict_invalid(format!(
"predict: {label} must be a vector"
)));
}
Ok(tensor.data.clone())
}
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),
_ => Err(fitlm_invalid(format!("fitlm: {name} must be numeric"))),
}
}
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]),
_ => Err(fitlm_invalid(format!("fitlm: {name} must be a scalar"))),
}
}
fn exclude_spec(value: &Value) -> BuiltinResult<ExcludeSpec> {
match value {
Value::Bool(value) => Ok(ExcludeSpec::Mask(vec![*value])),
Value::LogicalArray(array) => Ok(ExcludeSpec::Mask(
array.data.iter().map(|value| *value != 0).collect(),
)),
Value::Tensor(tensor) => {
let mut indices = Vec::with_capacity(tensor.data.len());
for value in &tensor.data {
if *value < 1.0 || value.fract().abs() > EPS {
return Err(fitlm_invalid(
"fitlm: Exclude indices must be positive integers",
));
}
indices.push(*value as usize);
}
Ok(ExcludeSpec::Indices(indices))
}
_ => Err(fitlm_invalid(
"fitlm: Exclude must be logical mask or row indices",
)),
}
}
fn exclude_mask(spec: Option<&ExcludeSpec>, rows: usize) -> BuiltinResult<Option<Vec<bool>>> {
let Some(spec) = spec else {
return Ok(None);
};
let mut mask = vec![false; rows];
match spec {
ExcludeSpec::Mask(values) => {
if values.len() == 1 {
mask.fill(values[0]);
} else if values.len() == rows {
mask.clone_from(values);
} else {
return Err(fitlm_invalid(
"fitlm: Exclude mask length must match the number of observations",
));
}
}
ExcludeSpec::Indices(indices) => {
for index in indices {
if *index == 0 || *index > rows {
return Err(fitlm_invalid("fitlm: Exclude index exceeds observations"));
}
mask[*index - 1] = true;
}
}
}
Ok(Some(mask))
}
fn bool_scalar(value: &Value, name: &str) -> BuiltinResult<bool> {
match value {
Value::Bool(value) => Ok(*value),
Value::Num(value) => Ok(*value != 0.0),
Value::Tensor(tensor) if tensor.data.len() == 1 => Ok(tensor.data[0] != 0.0),
_ => Err(fitlm_invalid(format!(
"fitlm: {name} must be logical scalar"
))),
}
}
fn scalar_text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::CharArray(chars) => Some(chars_to_string(chars)),
Value::StringArray(array) if array.data.len() == 1 => array.data.first().cloned(),
_ => None,
}
}
fn text_list(value: &Value, name: &str) -> BuiltinResult<Vec<String>> {
match value {
Value::String(text) => Ok(vec![text.clone()]),
Value::CharArray(chars) => Ok(vec![chars_to_string(chars)]),
Value::StringArray(array) => Ok(array.data.clone()),
Value::Cell(cell) => cell
.data
.iter()
.map(|value| {
scalar_text(value)
.ok_or_else(|| fitlm_invalid(format!("fitlm: {name} entries must be text")))
})
.collect(),
_ => Err(fitlm_invalid(format!(
"fitlm: {name} must be text or a text collection"
))),
}
}
fn chars_to_string(chars: &CharArray) -> String {
chars.data.iter().collect()
}
fn is_name_value_option(value: &Value) -> bool {
scalar_text(value).is_some_and(|text| {
matches!(
text.to_ascii_lowercase().as_str(),
"intercept"
| "responsevar"
| "predictorvars"
| "varnames"
| "exclude"
| "weights"
| "robustopts"
| "categoricalvars"
)
})
}
fn columns_to_tensor(columns: Vec<Vec<f64>>) -> BuiltinResult<Tensor> {
if columns.is_empty() {
return Tensor::new(Vec::new(), vec![0, 0])
.map_err(|err| fitlm_internal(format!("fitlm: {err}")));
}
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(fitlm_invalid(
"fitlm: predictor variables must have the same height",
));
}
data.extend(column);
}
Tensor::new(data, vec![rows, cols]).map_err(|err| fitlm_internal(format!("fitlm: {err}")))
}
fn columns_to_tensor_predict(columns: Vec<Vec<f64>>) -> BuiltinResult<Tensor> {
if columns.is_empty() {
return Tensor::new(Vec::new(), vec![0, 0])
.map_err(|err| predict_internal(format!("predict: {err}")));
}
let rows = columns[0].len();
let cols = columns.len();
let mut data = Vec::with_capacity(rows * cols);
for column in columns {
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 terms_tensor(terms: &[Vec<u32>]) -> BuiltinResult<Value> {
let rows = terms.len();
let cols = terms.first().map(|term| term.len()).unwrap_or(0);
let mut data = Vec::with_capacity(rows * cols);
for col in 0..cols {
for term in terms {
data.push(term[col] as f64);
}
}
Ok(Value::Tensor(Tensor::new(data, vec![rows, cols]).map_err(
|err| fitlm_internal(format!("fitlm: terms metadata error: {err}")),
)?))
}
fn numeric_column(values: &[f64]) -> BuiltinResult<Value> {
Ok(Value::Tensor(
Tensor::new(values.to_vec(), vec![values.len(), 1])
.map_err(|err| fitlm_internal(format!("fitlm: numeric column error: {err}")))?,
))
}
fn string_row(values: Vec<String>) -> BuiltinResult<Value> {
let len = values.len();
Ok(Value::StringArray(
StringArray::new(values, vec![1, len])
.map_err(|err| fitlm_internal(format!("fitlm: string row error: {err}")))?,
))
}
fn string_column(values: Vec<String>) -> BuiltinResult<Value> {
let len = values.len();
Ok(Value::StringArray(
StringArray::new(values, vec![len, 1])
.map_err(|err| fitlm_internal(format!("fitlm: string column error: {err}")))?,
))
}
fn has_intercept(terms: &[Vec<u32>]) -> bool {
terms
.iter()
.any(|term| term.iter().all(|power| *power == 0))
}
#[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 fitlm_matrix_fits_coefficients_and_predicts() {
let x = tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1);
let y = tensor_value(vec![1.0, 3.0, 5.0, 7.0], 4, 1);
let model = block_on(fitlm_builtin(x, vec![y])).unwrap();
let model_object = object(model.clone());
assert_eq!(
model_object.properties.get("NumObservations"),
Some(&Value::Num(4.0))
);
let coeffs = match model_object.properties.get("Coefficients").unwrap() {
Value::Object(object) => table_variables(object).unwrap(),
other => panic!("coefficients {other:?}"),
};
let estimates = tensor(coeffs.fields.get("Estimate").unwrap());
assert!((estimates.data[0] - 1.0).abs() < 1.0e-10);
assert!((estimates.data[1] - 2.0).abs() < 1.0e-10);
let predicted = block_on(predict_builtin(
model,
tensor_value(vec![4.0, 5.0], 2, 1),
Vec::new(),
))
.unwrap();
let ypred = tensor(&predicted);
assert_eq!(ypred.shape, vec![2, 1]);
assert!((ypred.data[0] - 9.0).abs() < 1.0e-10);
assert!((ypred.data[1] - 11.0).abs() < 1.0e-10);
}
#[test]
fn fitlm_table_formula_uses_named_predictors() {
let table = table_from_columns(
vec!["A".into(), "B".into(), "Y".into()],
vec![
tensor_value(vec![1.0, 2.0, 3.0, 4.0], 4, 1),
tensor_value(vec![0.0, 1.0, 0.0, 1.0], 4, 1),
tensor_value(vec![3.0, 7.0, 9.0, 13.0], 4, 1),
],
)
.unwrap();
let model = block_on(fitlm_builtin(
table,
vec![Value::String("Y ~ A + B".into())],
))
.unwrap();
let object = object(model);
assert_eq!(
object.properties.get("Formula"),
Some(&Value::String("Y ~ A + B".into()))
);
let names = match object.properties.get("PredictorNames").unwrap() {
Value::StringArray(array) => array.data.clone(),
other => panic!("names {other:?}"),
};
assert_eq!(names, vec!["A", "B"]);
}
#[test]
fn fitlm_quadratic_adds_square_term() {
let x = tensor_value(vec![-1.0, 0.0, 1.0, 2.0], 4, 1);
let y = tensor_value(vec![3.0, 1.0, 3.0, 9.0], 4, 1);
let model =
object(block_on(fitlm_builtin(x, vec![y, Value::String("quadratic".into())])).unwrap());
let coeffs = match model.properties.get("Coefficients").unwrap() {
Value::Object(object) => table_variables(object).unwrap(),
other => panic!("coefficients {other:?}"),
};
let estimates = tensor(coeffs.fields.get("Estimate").unwrap());
assert_eq!(estimates.shape, vec![3, 1]);
assert!((estimates.data[0] - 1.0).abs() < 1.0e-10);
assert!(estimates.data[1].abs() < 1.0e-10);
assert!((estimates.data[2] - 2.0).abs() < 1.0e-10);
}
#[test]
fn fitlm_omits_nan_rows() {
let x = tensor_value(vec![0.0, 1.0, f64::NAN, 3.0], 4, 1);
let y = tensor_value(vec![1.0, 3.0, 5.0, 7.0], 4, 1);
let model = object(block_on(fitlm_builtin(x, vec![y])).unwrap());
assert_eq!(
model.properties.get("NumObservations"),
Some(&Value::Num(3.0))
);
}
#[test]
fn predict_confidence_intervals_are_column_major() {
let x = tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1);
let y = tensor_value(vec![1.0, 3.1, 4.9, 7.0], 4, 1);
let model = block_on(fitlm_builtin(x, vec![y])).unwrap();
let output_count = crate::output_count::push_output_count(Some(2));
let predicted = block_on(predict_builtin(
model,
tensor_value(vec![4.0, 5.0], 2, 1),
Vec::new(),
))
.unwrap();
drop(output_count);
let yci = match predicted {
Value::OutputList(list) => tensor(&list[1]).clone(),
other => panic!("expected tensor or output list, got {other:?}"),
};
assert_eq!(yci.shape, vec![2, 2]);
assert!(yci.data[0] < yci.data[1]);
assert!(yci.data[2] < yci.data[3]);
assert!(yci.data[0] < yci.data[2]);
assert!(yci.data[1] < yci.data[3]);
}
#[test]
fn fitlm_formula_no_intercept_and_intercept_only_are_supported() {
let table = table_from_columns(
vec!["A".into(), "Y".into()],
vec![
tensor_value(vec![1.0, 2.0, 3.0], 3, 1),
tensor_value(vec![2.0, 4.0, 6.0], 3, 1),
],
)
.unwrap();
let no_intercept = object(
block_on(fitlm_builtin(
table.clone(),
vec![Value::String("Y ~ A - 1".into())],
))
.unwrap(),
);
let coeffs = match no_intercept.properties.get("Coefficients").unwrap() {
Value::Object(object) => table_variables(object).unwrap(),
other => panic!("coefficients {other:?}"),
};
let estimates = tensor(coeffs.fields.get("Estimate").unwrap());
assert_eq!(estimates.shape, vec![1, 1]);
assert!((estimates.data[0] - 2.0).abs() < 1.0e-10);
let intercept_only =
object(block_on(fitlm_builtin(table, vec![Value::String("Y ~ 1".into())])).unwrap());
let coeffs = match intercept_only.properties.get("Coefficients").unwrap() {
Value::Object(object) => table_variables(object).unwrap(),
other => panic!("coefficients {other:?}"),
};
let estimates = tensor(coeffs.fields.get("Estimate").unwrap());
assert_eq!(estimates.shape, vec![1, 1]);
assert!((estimates.data[0] - 4.0).abs() < 1.0e-10);
}
#[test]
fn intercept_only_table_model_predicts_one_value_per_table_row() {
let train = table_from_columns(
vec!["A".into(), "Y".into()],
vec![
tensor_value(vec![1.0, 2.0, 3.0], 3, 1),
tensor_value(vec![2.0, 4.0, 6.0], 3, 1),
],
)
.unwrap();
let model = block_on(fitlm_builtin(train, vec![Value::String("Y ~ 1".into())])).unwrap();
let xnew = table_from_columns(
vec!["A".into(), "Y".into()],
vec![
tensor_value(vec![10.0, 20.0], 2, 1),
tensor_value(vec![0.0, 0.0], 2, 1),
],
)
.unwrap();
let predicted = block_on(predict_builtin(model, xnew, Vec::new())).unwrap();
let yhat = tensor(&predicted);
assert_eq!(yhat.shape, vec![2, 1]);
assert_eq!(yhat.data, vec![4.0, 4.0]);
}
#[test]
fn predict_rejects_unsupported_simultaneous_intervals() {
let x = tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1);
let y = tensor_value(vec![1.0, 3.0, 5.0, 7.0], 4, 1);
let model = block_on(fitlm_builtin(x, vec![y])).unwrap();
let err = block_on(predict_builtin(
model,
tensor_value(vec![4.0], 1, 1),
vec![Value::from("Simultaneous"), Value::Bool(true)],
))
.unwrap_err();
assert!(err.message.contains("Simultaneous=true"));
}
#[test]
fn fitlm_matrix_rejects_table_only_options_and_bad_varnames() {
let x = tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1);
let y = tensor_value(vec![1.0, 3.0, 5.0, 7.0], 4, 1);
let err = block_on(fitlm_builtin(
x.clone(),
vec![
y.clone(),
Value::from("ResponseVar"),
Value::String("Y".into()),
],
))
.unwrap_err();
assert!(err.message.contains("ResponseVar"));
let err = block_on(fitlm_builtin(
x.clone(),
vec![
y.clone(),
Value::from("PredictorVars"),
Value::String("X".into()),
],
))
.unwrap_err();
assert!(err.message.contains("PredictorVars"));
let err = block_on(fitlm_builtin(
x,
vec![
y,
Value::from("VarNames"),
Value::StringArray(StringArray {
data: vec!["only_x".into()],
shape: vec![1, 1],
rows: 1,
cols: 1,
}),
],
))
.unwrap_err();
assert!(err.message.contains("VarNames"));
}
#[test]
fn coefficients_table_exposes_row_names() {
let x = tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1);
let y = tensor_value(vec![1.0, 3.0, 5.0, 7.0], 4, 1);
let model = object(block_on(fitlm_builtin(x, vec![y])).unwrap());
let coeffs = match model.properties.get("Coefficients").unwrap() {
Value::Object(object) => object,
other => panic!("coefficients {other:?}"),
};
let props = match coeffs.properties.get("Properties").unwrap() {
Value::Struct(props) => props,
other => panic!("properties {other:?}"),
};
match props.fields.get("RowNames").unwrap() {
Value::StringArray(array) => {
assert_eq!(array.data, vec!["(Intercept)", "x1"]);
assert_eq!(array.shape, vec![2, 1]);
}
other => panic!("row names {other:?}"),
}
}
#[test]
fn fitlm_accepts_exclude_indices_and_table_weight_variable() {
let x = tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1);
let y = tensor_value(vec![100.0, 3.0, 5.0, 7.0], 4, 1);
let model = object(
block_on(fitlm_builtin(
x,
vec![y, Value::from("Exclude"), tensor_value(vec![1.0], 1, 1)],
))
.unwrap(),
);
assert_eq!(
model.properties.get("NumObservations"),
Some(&Value::Num(3.0))
);
let table = table_from_columns(
vec!["A".into(), "Y".into(), "W".into()],
vec![
tensor_value(vec![0.0, 1.0, 2.0, 3.0], 4, 1),
tensor_value(vec![1.0, 3.0, 5.0, 7.0], 4, 1),
tensor_value(vec![1.0, 1.0, 1.0, 1.0], 4, 1),
],
)
.unwrap();
let weighted = object(
block_on(fitlm_builtin(
table,
vec![
Value::String("Y ~ A".into()),
Value::from("Weights"),
Value::String("W".into()),
],
))
.unwrap(),
);
assert_eq!(
weighted.properties.get("NumPredictors"),
Some(&Value::Num(1.0))
);
}
}