use log::{trace, warn};
use num_complex::Complex64;
use runmat_accelerate_api::ProviderPolyfitResult;
use runmat_builtins::{ComplexTensor, StructValue, Tensor, Value};
use runmat_macros::runtime_builtin;
use crate::builtins::common::tensor;
use crate::dispatcher;
use crate::{build_runtime_error, BuiltinResult, RuntimeError};
use crate::builtins::common::spec::{
BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
};
use crate::builtins::math::poly::type_resolvers::polyfit_type;
const EPS: f64 = 1.0e-12;
const EPS_NAN: f64 = 1.0e-12;
const BUILTIN_NAME: &str = "polyfit";
#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::math::poly::polyfit")]
pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
name: "polyfit",
op_kind: GpuOpKind::Custom("polyfit"),
supported_precisions: &[ScalarType::F32, ScalarType::F64],
broadcast: BroadcastSemantics::Matlab,
provider_hooks: &[ProviderHook::Custom("polyfit")],
constant_strategy: ConstantStrategy::UniformBuffer,
residency: ResidencyPolicy::GatherImmediately,
nan_mode: ReductionNaN::Include,
two_pass_threshold: None,
workgroup_size: None,
accepts_nan_mode: false,
notes:
"Providers may gather to the host and invoke the shared Householder QR solver; WGPU implements this path today.",
};
fn polyfit_error(message: impl Into<String>) -> RuntimeError {
build_runtime_error(message)
.with_builtin(BUILTIN_NAME)
.build()
}
#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::math::poly::polyfit")]
pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
name: "polyfit",
shape: ShapeRequirements::Any,
constant_strategy: ConstantStrategy::UniformBuffer,
elementwise: None,
reduction: None,
emits_nan: false,
notes: "Acts as a sink node—polynomial fitting materialises results eagerly and terminates fusion graphs.",
};
#[runtime_builtin(
name = "polyfit",
category = "math/poly",
summary = "Fit an n-th degree polynomial to data points with MATLAB-compatible outputs.",
keywords = "polyfit,polynomial,least-squares,gpu",
accel = "sink",
sink = true,
type_resolver(polyfit_type),
builtin_path = "crate::builtins::math::poly::polyfit"
)]
async fn polyfit_builtin(
x: Value,
y: Value,
degree: Value,
rest: Vec<Value>,
) -> crate::BuiltinResult<Value> {
let eval = evaluate(x, y, degree, &rest).await?;
if let Some(out_count) = crate::output_count::current_output_count() {
if out_count == 0 {
return Ok(Value::OutputList(Vec::new()));
}
let mut outputs = vec![eval.coefficients()];
if out_count >= 2 {
outputs.push(eval.stats());
}
if out_count >= 3 {
outputs.push(eval.mu());
}
return Ok(crate::output_count::output_list_with_padding(
out_count, outputs,
));
}
Ok(eval.coefficients())
}
pub async fn evaluate(
x: Value,
y: Value,
degree: Value,
rest: &[Value],
) -> BuiltinResult<PolyfitEval> {
let deg = parse_degree(°ree)?;
if let Some(eval) = try_gpu_polyfit(&x, &y, deg, rest).await? {
return Ok(eval);
}
let x_host = dispatcher::gather_if_needed_async(&x).await?;
let y_host = dispatcher::gather_if_needed_async(&y).await?;
let x_data = real_vector("polyfit", "X", x_host).await?;
let (y_data, is_complex_input) = complex_vector("polyfit", "Y", y_host).await?;
if x_data.len() != y_data.len() {
return Err(polyfit_error(
"polyfit: X and Y vectors must be the same length",
));
}
if x_data.is_empty() {
return Err(polyfit_error(
"polyfit: X and Y must contain at least one sample",
));
}
if deg + 1 > x_data.len() && x_data.len() > 1 {
warn!(
"polyfit: polynomial degree {} is ill-conditioned for {} data points; results may be inaccurate",
deg,
x_data.len()
);
}
let weights = parse_weights(rest, x_data.len()).await?;
let mut solution = solve_polyfit(&x_data, &y_data, deg, weights.as_deref())?;
if is_complex_input {
solution.is_complex = true;
}
PolyfitEval::from_solution(solution)
}
async fn try_gpu_polyfit(
x: &Value,
y: &Value,
degree: usize,
rest: &[Value],
) -> BuiltinResult<Option<PolyfitEval>> {
let provider = match runmat_accelerate_api::provider() {
Some(p) => p,
None => return Ok(None),
};
let x_handle = match x {
Value::GpuTensor(handle) => handle,
_ => return Ok(None),
};
let y_handle = match y {
Value::GpuTensor(handle) => handle,
_ => return Ok(None),
};
if rest.len() > 1 {
return Ok(None);
}
let weight_handle = match rest.first() {
Some(Value::GpuTensor(handle)) => Some(handle),
Some(_) => return Ok(None),
None => None,
};
let result = match provider
.polyfit(x_handle, y_handle, degree, weight_handle)
.await
{
Ok(res) => res,
Err(err) => {
trace!("polyfit: provider path unavailable ({err}); falling back to host");
return Ok(None);
}
};
let solution = PolyfitSolution::from_provider(result)?;
PolyfitEval::from_solution(solution).map(Some)
}
#[derive(Clone, Debug)]
struct PolyfitSolution {
coeffs: Vec<Complex64>,
r_matrix: Vec<f64>,
mu_mean: f64,
mu_scale: f64,
normr: f64,
df: f64,
cols: usize,
is_complex: bool,
}
impl PolyfitSolution {
fn from_provider(result: ProviderPolyfitResult) -> BuiltinResult<Self> {
let cols = result.coefficients.len();
if cols == 0 {
return Err(polyfit_error(
"polyfit: provider returned empty coefficient vector",
));
}
if result.r_matrix.len() != cols * cols {
return Err(polyfit_error(
"polyfit: provider returned malformed R matrix",
));
}
let [mu_mean, mu_scale] = result.mu;
Ok(Self {
coeffs: result
.coefficients
.into_iter()
.map(|re| Complex64::new(re, 0.0))
.collect(),
r_matrix: result.r_matrix,
mu_mean,
mu_scale,
normr: result.normr,
df: result.df,
cols,
is_complex: false,
})
}
}
#[derive(Debug)]
pub struct PolyfitEval {
coefficients: Value,
stats: Value,
mu: Value,
is_complex: bool,
}
impl PolyfitEval {
fn from_solution(solution: PolyfitSolution) -> BuiltinResult<Self> {
let coefficients = coefficients_to_value(&solution.coeffs)?;
let stats = build_stats(
&solution.r_matrix,
solution.cols,
solution.normr,
solution.df,
)?;
let mu = build_mu(solution.mu_mean, solution.mu_scale)?;
Ok(Self {
coefficients,
stats,
mu,
is_complex: solution.is_complex,
})
}
pub fn coefficients(&self) -> Value {
self.coefficients.clone()
}
pub fn stats(&self) -> Value {
self.stats.clone()
}
pub fn mu(&self) -> Value {
self.mu.clone()
}
pub fn is_complex(&self) -> bool {
self.is_complex
}
}
fn parse_degree(value: &Value) -> BuiltinResult<usize> {
match value {
Value::Int(i) => {
let raw = i.to_i64();
if raw < 0 {
return Err(polyfit_error(
"polyfit: degree must be a non-negative integer",
));
}
Ok(raw as usize)
}
Value::Num(n) => {
if !n.is_finite() {
return Err(polyfit_error("polyfit: degree must be finite"));
}
let rounded = n.round();
if (rounded - n).abs() > EPS {
return Err(polyfit_error("polyfit: degree must be an integer"));
}
if rounded < 0.0 {
return Err(polyfit_error(
"polyfit: degree must be a non-negative integer",
));
}
Ok(rounded as usize)
}
Value::Tensor(t) if tensor::is_scalar_tensor(t) => parse_degree(&Value::Num(t.data[0])),
Value::LogicalArray(l) if l.len() == 1 => {
parse_degree(&Value::Num(if l.data[0] != 0 { 1.0 } else { 0.0 }))
}
other => Err(polyfit_error(format!(
"polyfit: degree must be a scalar numeric value, got {other:?}"
))),
}
}
#[async_recursion::async_recursion(?Send)]
async fn real_vector(context: &str, label: &str, value: Value) -> BuiltinResult<Vec<f64>> {
match value {
Value::Tensor(mut tensor) => {
ensure_vector_shape(context, label, &tensor.shape)?;
Ok(tensor.data.drain(..).collect())
}
Value::LogicalArray(logical) => {
let tensor = tensor::logical_to_tensor(&logical).map_err(polyfit_error)?;
ensure_vector_shape(context, label, &tensor.shape)?;
Ok(tensor.data)
}
Value::Num(n) => Ok(vec![n]),
Value::Int(i) => Ok(vec![i.to_f64()]),
Value::Bool(b) => Ok(vec![if b { 1.0 } else { 0.0 }]),
Value::GpuTensor(handle) => {
let gathered =
crate::builtins::common::gpu_helpers::gather_tensor_async(&handle).await?;
real_vector(context, label, Value::Tensor(gathered)).await
}
Value::Complex(_, _) | Value::ComplexTensor(_) => Err(polyfit_error(format!(
"{context}: {label} must be real-valued; complex inputs are not supported"
))),
other => Err(polyfit_error(format!(
"{context}: expected {label} to be a numeric vector, got {other:?}"
))),
}
}
#[async_recursion::async_recursion(?Send)]
async fn complex_vector(
context: &str,
label: &str,
value: Value,
) -> BuiltinResult<(Vec<Complex64>, bool)> {
match value {
Value::Tensor(mut tensor) => {
ensure_vector_shape(context, label, &tensor.shape)?;
let all_real = true;
let data = tensor
.data
.drain(..)
.map(|x| Complex64::new(x, 0.0))
.collect();
Ok((data, all_real))
}
Value::ComplexTensor(tensor) => {
ensure_vector_shape(context, label, &tensor.shape)?;
let is_complex = tensor.data.iter().any(|&(_, im)| im.abs() > EPS);
let data = tensor
.data
.into_iter()
.map(|(re, im)| Complex64::new(re, im))
.collect::<Vec<_>>();
Ok((data, is_complex))
}
Value::LogicalArray(logical) => {
let tensor = tensor::logical_to_tensor(&logical).map_err(polyfit_error)?;
ensure_vector_shape(context, label, &tensor.shape)?;
Ok((
tensor
.data
.iter()
.map(|&x| Complex64::new(x, 0.0))
.collect(),
false,
))
}
Value::Num(n) => Ok((vec![Complex64::new(n, 0.0)], false)),
Value::Int(i) => Ok((vec![Complex64::new(i.to_f64(), 0.0)], false)),
Value::Bool(b) => Ok((vec![Complex64::new(if b { 1.0 } else { 0.0 }, 0.0)], false)),
Value::Complex(re, im) => Ok((vec![Complex64::new(re, im)], im.abs() > EPS)),
Value::GpuTensor(handle) => {
let gathered =
crate::builtins::common::gpu_helpers::gather_tensor_async(&handle).await?;
complex_vector(context, label, Value::Tensor(gathered)).await
}
other => Err(polyfit_error(format!(
"{context}: expected {label} to be a numeric vector, got {other:?}"
))),
}
}
async fn parse_weights(rest: &[Value], len: usize) -> BuiltinResult<Option<Vec<f64>>> {
match rest.len() {
0 => Ok(None),
1 => {
let gathered = dispatcher::gather_if_needed_async(&rest[0]).await?;
let data = real_vector("polyfit", "weights", gathered).await?;
if data.len() != len {
return Err(polyfit_error(
"polyfit: weight vector must match the size of X",
));
}
validate_weights(&data)?;
Ok(Some(data))
}
_ => Err(polyfit_error("polyfit: too many input arguments")),
}
}
fn validate_weights(weights: &[f64]) -> BuiltinResult<()> {
for (idx, w) in weights.iter().enumerate() {
if !w.is_finite() {
return Err(polyfit_error(format!(
"polyfit: weight at position {} must be finite",
idx + 1
)));
}
if *w < 0.0 {
return Err(polyfit_error("polyfit: weights must be non-negative"));
}
}
Ok(())
}
fn solve_polyfit(
x_data: &[f64],
y_data: &[Complex64],
degree: usize,
weights: Option<&[f64]>,
) -> BuiltinResult<PolyfitSolution> {
if x_data.len() != y_data.len() {
return Err(polyfit_error(
"polyfit: X and Y vectors must be the same length",
));
}
if x_data.is_empty() {
return Err(polyfit_error(
"polyfit: X and Y must contain at least one sample",
));
}
if let Some(w) = weights {
if w.len() != x_data.len() {
return Err(polyfit_error(
"polyfit: weight vector must match the size of X",
));
}
validate_weights(w)?;
}
let mean = x_data.iter().sum::<f64>() / x_data.len() as f64;
if !mean.is_finite() {
return Err(polyfit_error("polyfit: mean of X must be finite"));
}
let scale = compute_scale(x_data, mean)?;
let scaled: Vec<f64> = x_data.iter().map(|&v| (v - mean) / scale).collect();
let mut rhs = y_data.to_vec();
for (idx, value) in rhs.iter().enumerate() {
if !value.re.is_finite() || !value.im.is_finite() {
return Err(polyfit_error(format!(
"polyfit: Y must contain finite values (encountered NaN/Inf at position {})",
idx + 1
)));
}
}
if let Some(w) = weights {
apply_weights_rhs(&mut rhs, w)?;
}
let rows = scaled.len();
let cols = degree + 1;
let mut vandermonde = build_vandermonde(&scaled, cols);
if let Some(w) = weights {
apply_weights_matrix(&mut vandermonde, rows, cols, w)?;
}
let mut transformed_rhs = rhs.clone();
householder_qr(&mut vandermonde, rows, cols, &mut transformed_rhs)?;
let coeff_scaled = solve_upper(&vandermonde, rows, cols, &transformed_rhs)?;
let coeff_original = transform_coefficients(&coeff_scaled, mean, scale);
let normr = residual_norm(&transformed_rhs, rows, cols);
let df = if rows > cols {
(rows - cols) as f64
} else {
0.0
};
let r_matrix = extract_upper(&vandermonde, rows, cols);
let is_complex = coeff_original.iter().any(|c| c.im.abs() > EPS_NAN);
Ok(PolyfitSolution {
coeffs: coeff_original,
r_matrix,
mu_mean: mean,
mu_scale: scale,
normr,
df,
cols,
is_complex,
})
}
fn compute_scale(data: &[f64], mean: f64) -> BuiltinResult<f64> {
if data.len() <= 1 {
return Ok(1.0);
}
let mut acc = 0.0;
for &value in data {
if !value.is_finite() {
return Err(polyfit_error("polyfit: X must contain finite values"));
}
let diff = value - mean;
acc += diff * diff;
}
let denom = (data.len() as f64 - 1.0).max(1.0);
let std = (acc / denom).sqrt();
let scale = if std.abs() <= EPS { 1.0 } else { std };
if !scale.is_finite() {
return Err(polyfit_error(
"polyfit: failed to compute a stable scaling factor",
));
}
Ok(scale)
}
fn build_vandermonde(u: &[f64], cols: usize) -> Vec<f64> {
let rows = u.len();
let mut matrix = vec![0.0; rows * cols];
if cols == 0 {
return matrix;
}
for (row_idx, &value) in u.iter().enumerate() {
let mut powers = vec![0.0; cols];
powers[cols - 1] = 1.0;
for idx in (0..cols - 1).rev() {
powers[idx] = powers[idx + 1] * value;
}
for col_idx in 0..cols {
matrix[row_idx + col_idx * rows] = powers[col_idx];
}
}
matrix
}
fn apply_weights_matrix(
matrix: &mut [f64],
rows: usize,
cols: usize,
weights: &[f64],
) -> BuiltinResult<()> {
for (row, weight) in weights.iter().enumerate().take(rows) {
let sqrt_w = weight.sqrt();
if !sqrt_w.is_finite() {
return Err(polyfit_error(format!(
"polyfit: weight at position {} must be finite",
row + 1
)));
}
for col in 0..cols {
let idx = row + col * rows;
matrix[idx] *= sqrt_w;
}
}
Ok(())
}
fn apply_weights_rhs(rhs: &mut [Complex64], weights: &[f64]) -> BuiltinResult<()> {
for (idx, (value, weight)) in rhs.iter_mut().zip(weights.iter()).enumerate() {
let sqrt_w = weight.sqrt();
if !sqrt_w.is_finite() {
return Err(polyfit_error(format!(
"polyfit: weight at position {} must be finite",
idx + 1
)));
}
*value *= sqrt_w;
}
Ok(())
}
fn ensure_vector_shape(context: &str, label: &str, shape: &[usize]) -> BuiltinResult<()> {
if !is_vector_shape(shape) {
return Err(polyfit_error(format!(
"{context}: {label} must be a vector"
)));
}
Ok(())
}
fn is_vector_shape(shape: &[usize]) -> bool {
shape.iter().copied().filter(|&dim| dim > 1).count() <= 1
}
fn householder_qr(
matrix: &mut [f64],
rows: usize,
cols: usize,
rhs: &mut [Complex64],
) -> BuiltinResult<()> {
let min_dim = rows.min(cols);
for k in 0..min_dim {
let mut norm_sq = 0.0;
for row in k..rows {
let val = matrix[row + k * rows];
norm_sq += val * val;
}
if norm_sq <= EPS {
continue;
}
let norm = norm_sq.sqrt();
let x0 = matrix[k + k * rows];
let alpha = if x0 >= 0.0 { -norm } else { norm };
let mut v = vec![0.0; rows - k];
v[0] = x0 - alpha;
for row in (k + 1)..rows {
v[row - k] = matrix[row + k * rows];
}
let v_norm_sq: f64 = v.iter().map(|&x| x * x).sum();
if v_norm_sq <= EPS {
continue;
}
let beta = 2.0 / v_norm_sq;
matrix[k + k * rows] = alpha;
for row in (k + 1)..rows {
matrix[row + k * rows] = 0.0;
}
for col in (k + 1)..cols {
let mut dot = 0.0;
for (idx, &vi) in v.iter().enumerate() {
let row_idx = k + idx;
dot += vi * matrix[row_idx + col * rows];
}
let factor = beta * dot;
for (idx, &vi) in v.iter().enumerate() {
let row_idx = k + idx;
matrix[row_idx + col * rows] -= factor * vi;
}
}
let mut dot = Complex64::new(0.0, 0.0);
for (idx, &vi) in v.iter().enumerate() {
let row_idx = k + idx;
dot += rhs[row_idx] * vi;
}
let factor = Complex64::new(beta, 0.0) * dot;
for (idx, &vi) in v.iter().enumerate() {
let row_idx = k + idx;
rhs[row_idx] -= factor * vi;
}
}
Ok(())
}
fn solve_upper(
matrix: &[f64],
rows: usize,
cols: usize,
rhs: &[Complex64],
) -> BuiltinResult<Vec<Complex64>> {
if rhs.len() < rows {
return Err(polyfit_error(
"polyfit internal error: RHS dimension mismatch",
));
}
let mut coeffs = vec![Complex64::new(0.0, 0.0); cols];
for col in (0..cols).rev() {
let diag = if col < rows {
matrix[col + col * rows]
} else {
0.0
};
if diag.abs() <= EPS {
coeffs[col] = Complex64::new(0.0, 0.0);
continue;
}
let mut acc = if col < rows {
rhs[col]
} else {
Complex64::new(0.0, 0.0)
};
for next in (col + 1)..cols {
let idx = if col < rows {
matrix[col + next * rows]
} else {
0.0
};
acc -= Complex64::new(idx, 0.0) * coeffs[next];
}
coeffs[col] = acc / Complex64::new(diag, 0.0);
}
Ok(coeffs)
}
fn residual_norm(rhs: &[Complex64], rows: usize, cols: usize) -> f64 {
let tail_start = rows.min(cols);
let mut acc = 0.0;
for value in rhs.iter().skip(tail_start) {
acc += value.norm_sqr();
}
acc.sqrt()
}
fn extract_upper(matrix: &[f64], rows: usize, cols: usize) -> Vec<f64> {
let mut output = vec![0.0; cols * cols];
for col in 0..cols {
for row in 0..=col {
if row < rows {
output[row + col * cols] = matrix[row + col * rows];
}
}
}
output
}
fn transform_coefficients(coeffs: &[Complex64], mean: f64, scale: f64) -> Vec<Complex64> {
let mut poly: Vec<Complex64> = Vec::new();
for &coeff in coeffs {
let mut next = vec![Complex64::new(0.0, 0.0); poly.len() + 1];
for (idx, &value) in poly.iter().enumerate() {
next[idx + 1] += value / scale;
next[idx] -= value * (mean / scale);
}
next[0] += coeff;
poly = next;
}
poly.reverse();
poly
}
fn coefficients_to_value(coeffs: &[Complex64]) -> BuiltinResult<Value> {
let all_real = coeffs
.iter()
.all(|c| c.im.abs() <= EPS_NAN && c.re.is_finite());
if all_real {
let data: Vec<f64> = coeffs.iter().map(|c| c.re).collect();
let tensor = Tensor::new(data, vec![1, coeffs.len()])
.map_err(|e| polyfit_error(format!("polyfit: {e}")))?;
Ok(Value::Tensor(tensor))
} else {
let data: Vec<(f64, f64)> = coeffs.iter().map(|c| (c.re, c.im)).collect();
let tensor = ComplexTensor::new(data, vec![1, coeffs.len()])
.map_err(|e| polyfit_error(format!("polyfit: {e}")))?;
Ok(Value::ComplexTensor(tensor))
}
}
fn build_stats(r: &[f64], n: usize, normr: f64, df: f64) -> BuiltinResult<Value> {
let tensor =
Tensor::new(r.to_vec(), vec![n, n]).map_err(|e| polyfit_error(format!("polyfit: {e}")))?;
let mut st = StructValue::new();
st.fields.insert("R".to_string(), Value::Tensor(tensor));
st.fields.insert("df".to_string(), Value::Num(df));
st.fields.insert("normr".to_string(), Value::Num(normr));
Ok(Value::Struct(st))
}
fn build_mu(mean: f64, scale: f64) -> BuiltinResult<Value> {
if !scale.is_finite() || scale.abs() <= EPS {
return Err(polyfit_error("polyfit: mu(2) must be non-zero and finite"));
}
let tensor = Tensor::new(vec![mean, scale], vec![1, 2])
.map_err(|e| polyfit_error(format!("polyfit: {e}")))?;
Ok(Value::Tensor(tensor))
}
#[derive(Debug, Clone)]
pub struct PolyfitHostRealResult {
pub coefficients: Vec<f64>,
pub r_matrix: Vec<f64>,
pub mu: [f64; 2],
pub normr: f64,
pub df: f64,
}
pub fn polyfit_host_real_for_provider(
x: &[f64],
y: &[f64],
degree: usize,
weights: Option<&[f64]>,
) -> BuiltinResult<PolyfitHostRealResult> {
if x.len() != y.len() {
return Err(polyfit_error(
"polyfit: X and Y vectors must be the same length",
));
}
if let Some(w) = weights {
if w.len() != x.len() {
return Err(polyfit_error(
"polyfit: weight vector must match the size of X",
));
}
validate_weights(w)?;
}
let complex_y: Vec<Complex64> = y.iter().copied().map(|v| Complex64::new(v, 0.0)).collect();
let solution = solve_polyfit(x, &complex_y, degree, weights)?;
let PolyfitSolution {
coeffs,
r_matrix,
mu_mean,
mu_scale,
normr,
df,
cols: _,
is_complex,
} = solution;
if is_complex {
return Err(polyfit_error(
"polyfit: provider fallback produced complex coefficients for real data",
));
}
let coeffs: Vec<f64> = coeffs.into_iter().map(|c| c.re).collect();
let mu = [mu_mean, mu_scale];
Ok(PolyfitHostRealResult {
coefficients: coeffs,
r_matrix,
mu,
normr,
df,
})
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use crate::builtins::common::test_support;
use futures::executor::block_on;
fn assert_error_contains(err: crate::RuntimeError, needle: &str) {
assert!(
err.message().contains(needle),
"expected error containing '{needle}', got '{}'",
err.message()
);
}
fn evaluate(
x: Value,
y: Value,
degree: Value,
rest: &[Value],
) -> Result<PolyfitEval, RuntimeError> {
block_on(super::evaluate(x, y, degree, rest))
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn fits_linear_data() {
let x = Tensor::new(vec![0.0, 1.0, 2.0, 3.0], vec![4, 1]).unwrap();
let mut y_vals = Vec::new();
for i in 0..4 {
y_vals.push(1.5 * i as f64 + 2.0);
}
let y = Tensor::new(y_vals, vec![4, 1]).unwrap();
let eval = evaluate(
Value::Tensor(x),
Value::Tensor(y),
Value::Int(runmat_builtins::IntValue::I32(1)),
&[],
)
.expect("polyfit");
match eval.coefficients() {
Value::Tensor(t) => {
assert_eq!(t.shape, vec![1, 2]);
assert!((t.data[0] - 1.5).abs() < 1e-10);
assert!((t.data[1] - 2.0).abs() < 1e-10);
}
other => panic!("expected tensor coefficients, got {other:?}"),
}
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn returns_struct_and_mu() {
let x = Tensor::new(vec![-1.0, 0.0, 1.0], vec![3, 1]).unwrap();
let y = Tensor::new(vec![1.0, 0.0, 1.0], vec![3, 1]).unwrap();
let eval = evaluate(
Value::Tensor(x),
Value::Tensor(y),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[],
)
.expect("polyfit");
match eval.stats() {
Value::Struct(s) => {
assert!(s.fields.contains_key("R"));
assert!(s.fields.contains_key("df"));
assert!(s.fields.contains_key("normr"));
}
other => panic!("expected struct, got {other:?}"),
}
match eval.mu() {
Value::Tensor(t) => {
assert_eq!(t.shape, vec![1, 2]);
assert!((t.data[0]).abs() < 1e-10);
assert!(t.data[1].abs() > 0.0);
}
other => panic!("expected tensor mu, got {other:?}"),
}
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn weighted_fit_matches_unweighted_when_weights_equal() {
let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
let weights = Tensor::new(vec![1.0, 1.0, 1.0], vec![3, 1]).unwrap();
let eval_unweighted = evaluate(
Value::Tensor(x.clone()),
Value::Tensor(y.clone()),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[],
)
.expect("polyfit");
let eval_weighted = evaluate(
Value::Tensor(x),
Value::Tensor(y),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[Value::Tensor(weights)],
)
.expect("polyfit");
assert_eq!(eval_unweighted.coefficients(), eval_weighted.coefficients());
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn accepts_logical_degree_scalar() {
let x = Tensor::new(vec![0.0, 1.0], vec![2, 1]).unwrap();
let y = Tensor::new(vec![1.0, 3.0], vec![2, 1]).unwrap();
let logical = runmat_builtins::LogicalArray::new(vec![1], vec![1, 1]).unwrap();
let eval = evaluate(
Value::Tensor(x),
Value::Tensor(y),
Value::LogicalArray(logical),
&[],
)
.expect("polyfit");
assert!(matches!(eval.coefficients(), Value::Tensor(_)));
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn rejects_non_integer_degree() {
let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
let err = evaluate(Value::Tensor(x), Value::Tensor(y), Value::Num(1.5), &[])
.expect_err("polyfit should reject non-integer degree");
assert_error_contains(err, "integer");
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn rejects_infinite_weights() {
let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
let weights = Tensor::new(vec![1.0, f64::INFINITY, 1.0], vec![3, 1]).unwrap();
let err = evaluate(
Value::Tensor(x),
Value::Tensor(y),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[Value::Tensor(weights)],
)
.expect_err("polyfit should reject infinite weights");
assert_error_contains(err, "weight at position 2");
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn gpu_inputs_are_gathered() {
test_support::with_test_provider(|provider| {
let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
let view = runmat_accelerate_api::HostTensorView {
data: &x.data,
shape: &x.shape,
};
let x_handle = provider.upload(&view).expect("upload");
let view_y = runmat_accelerate_api::HostTensorView {
data: &y.data,
shape: &y.shape,
};
let y_handle = provider.upload(&view_y).expect("upload");
let eval = evaluate(
Value::GpuTensor(x_handle),
Value::GpuTensor(y_handle),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[],
)
.expect("polyfit");
assert!(matches!(eval.coefficients(), Value::Tensor(_)));
});
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn gpu_weights_are_gathered() {
test_support::with_test_provider(|provider| {
let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
let weights = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
let x_view = runmat_accelerate_api::HostTensorView {
data: &x.data,
shape: &x.shape,
};
let y_view = runmat_accelerate_api::HostTensorView {
data: &y.data,
shape: &y.shape,
};
let w_view = runmat_accelerate_api::HostTensorView {
data: &weights.data,
shape: &weights.shape,
};
let x_handle = provider.upload(&x_view).expect("upload x");
let y_handle = provider.upload(&y_view).expect("upload y");
let w_handle = provider.upload(&w_view).expect("upload weights");
let cpu_eval = evaluate(
Value::Tensor(x.clone()),
Value::Tensor(y.clone()),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[Value::Tensor(weights.clone())],
)
.expect("cpu polyfit");
let gpu_eval = evaluate(
Value::GpuTensor(x_handle.clone()),
Value::GpuTensor(y_handle.clone()),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[Value::GpuTensor(w_handle.clone())],
)
.expect("gpu polyfit with weights");
assert_eq!(cpu_eval.coefficients(), gpu_eval.coefficients());
assert_eq!(cpu_eval.mu(), gpu_eval.mu());
let _ = provider.free(&x_handle);
let _ = provider.free(&y_handle);
let _ = provider.free(&w_handle);
});
}
#[cfg(feature = "wgpu")]
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn polyfit_wgpu_matches_cpu() {
let options = runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default();
let _provider =
match runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(options) {
Ok(p) => p,
Err(err) => {
warn!("polyfit_wgpu_matches_cpu: skipping test ({err})");
return;
}
};
let x = Tensor::new(vec![0.0, 1.0, 2.0, 3.0], vec![4, 1]).unwrap();
let y = Tensor::new(vec![1.0, 3.0, 7.0, 13.0], vec![4, 1]).unwrap();
let cpu_eval = evaluate(
Value::Tensor(x.clone()),
Value::Tensor(y.clone()),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[],
)
.expect("cpu polyfit");
let trait_provider = runmat_accelerate_api::provider().expect("wgpu provider registered");
let x_view = runmat_accelerate_api::HostTensorView {
data: &x.data,
shape: &x.shape,
};
let y_view = runmat_accelerate_api::HostTensorView {
data: &y.data,
shape: &y.shape,
};
let x_handle = trait_provider.upload(&x_view).expect("upload x");
let y_handle = trait_provider.upload(&y_view).expect("upload y");
let gpu_eval = evaluate(
Value::GpuTensor(x_handle.clone()),
Value::GpuTensor(y_handle.clone()),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[],
)
.expect("gpu polyfit");
let _ = trait_provider.free(&x_handle);
let _ = trait_provider.free(&y_handle);
let cpu_coeff = match cpu_eval.coefficients() {
Value::Tensor(t) => t,
other => panic!("expected tensor coefficients, got {other:?}"),
};
let gpu_coeff = match gpu_eval.coefficients() {
Value::Tensor(t) => t,
other => panic!("expected tensor coefficients, got {other:?}"),
};
assert_eq!(cpu_coeff.shape, gpu_coeff.shape);
for (a, b) in cpu_coeff.data.iter().zip(gpu_coeff.data.iter()) {
assert!((a - b).abs() < 1e-9, "coeff mismatch {a} vs {b}");
}
let cpu_mu = match cpu_eval.mu() {
Value::Tensor(t) => t,
other => panic!("expected tensor mu, got {other:?}"),
};
let gpu_mu = match gpu_eval.mu() {
Value::Tensor(t) => t,
other => panic!("expected tensor mu, got {other:?}"),
};
assert_eq!(cpu_mu.shape, gpu_mu.shape);
for (a, b) in cpu_mu.data.iter().zip(gpu_mu.data.iter()) {
assert!((a - b).abs() < 1e-9, "mu mismatch {a} vs {b}");
}
let cpu_stats = match cpu_eval.stats() {
Value::Struct(s) => s,
other => panic!("expected struct stats, got {other:?}"),
};
let gpu_stats = match gpu_eval.stats() {
Value::Struct(s) => s,
other => panic!("expected struct stats, got {other:?}"),
};
let cpu_r = match cpu_stats.fields.get("R").expect("R present") {
Value::Tensor(t) => t.clone(),
other => panic!("expected tensor R, got {other:?}"),
};
let gpu_r = match gpu_stats.fields.get("R").expect("R present") {
Value::Tensor(t) => t.clone(),
other => panic!("expected tensor R, got {other:?}"),
};
assert_eq!(cpu_r.shape, gpu_r.shape);
for (a, b) in cpu_r.data.iter().zip(gpu_r.data.iter()) {
assert!((a - b).abs() < 1e-9, "R mismatch {a} vs {b}");
}
let cpu_df = match cpu_stats.fields.get("df").expect("df present") {
Value::Num(n) => *n,
other => panic!("expected numeric df, got {other:?}"),
};
let gpu_df = match gpu_stats.fields.get("df").expect("df present") {
Value::Num(n) => *n,
other => panic!("expected numeric df, got {other:?}"),
};
assert!((cpu_df - gpu_df).abs() < 1e-9);
let cpu_normr = match cpu_stats.fields.get("normr").expect("normr present") {
Value::Num(n) => *n,
other => panic!("expected numeric normr, got {other:?}"),
};
let gpu_normr = match gpu_stats.fields.get("normr").expect("normr present") {
Value::Num(n) => *n,
other => panic!("expected numeric normr, got {other:?}"),
};
assert!((cpu_normr - gpu_normr).abs() < 1e-9);
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn rejects_mismatched_lengths() {
let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
let y = Tensor::new(vec![1.0, 2.0], vec![2, 1]).unwrap();
let err = evaluate(
Value::Tensor(x),
Value::Tensor(y),
Value::Int(runmat_builtins::IntValue::I32(1)),
&[],
)
.expect_err("polyfit should reject mismatched vector lengths");
assert_error_contains(err, "same length");
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn rejects_non_vector_inputs() {
let x = Tensor::new(vec![0.0, 1.0, 2.0, 3.0], vec![2, 2]).unwrap();
let y = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();
let err = evaluate(
Value::Tensor(x),
Value::Tensor(y),
Value::Int(runmat_builtins::IntValue::I32(1)),
&[],
)
.expect_err("polyfit should reject non-vector X");
assert_error_contains(err, "vector");
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn rejects_weight_length_mismatch() {
let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
let weights = Tensor::new(vec![1.0, 1.0], vec![2, 1]).unwrap();
let err = evaluate(
Value::Tensor(x),
Value::Tensor(y),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[Value::Tensor(weights)],
)
.expect_err("polyfit should reject mismatched weights");
assert_error_contains(err, "weight vector must match");
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn rejects_negative_weights() {
let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
let weights = Tensor::new(vec![1.0, -1.0, 1.0], vec![3, 1]).unwrap();
let err = evaluate(
Value::Tensor(x),
Value::Tensor(y),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[Value::Tensor(weights)],
)
.expect_err("polyfit should reject negative weights");
assert_error_contains(err, "non-negative");
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn fits_complex_data() {
let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
let complex_values =
ComplexTensor::new(vec![(0.0, 1.0), (1.0, 0.5), (4.0, -0.25)], vec![3, 1]).unwrap();
let eval = evaluate(
Value::Tensor(x),
Value::ComplexTensor(complex_values),
Value::Int(runmat_builtins::IntValue::I32(2)),
&[],
)
.expect("polyfit complex");
match eval.coefficients() {
Value::ComplexTensor(t) => {
assert_eq!(t.shape, vec![1, 3]);
}
other => panic!("expected complex tensor coefficients, got {other:?}"),
}
}
}