Skip to main content

runmat_runtime/builtins/math/poly/
polyfit.rs

1//! MATLAB-compatible `polyfit` builtin with GPU-aware semantics for RunMat.
2
3use log::{trace, warn};
4use num_complex::Complex64;
5use runmat_accelerate_api::ProviderPolyfitResult;
6use runmat_builtins::{
7    BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
8    BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
9    ComplexTensor, StructValue, Tensor, Value,
10};
11use runmat_macros::runtime_builtin;
12
13use crate::builtins::common::tensor;
14use crate::dispatcher;
15use crate::{build_runtime_error, BuiltinResult, RuntimeError};
16
17use crate::builtins::common::spec::{
18    BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
19    ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
20};
21use crate::builtins::math::poly::type_resolvers::polyfit_type;
22
23const EPS: f64 = 1.0e-12;
24const EPS_NAN: f64 = 1.0e-12;
25const BUILTIN_NAME: &str = "polyfit";
26
27const POLYFIT_OUTPUT_P: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
28    name: "p",
29    ty: BuiltinParamType::Any,
30    arity: BuiltinParamArity::Required,
31    default: None,
32    description: "Polynomial coefficient vector ordered from highest power to constant term.",
33}];
34
35const POLYFIT_OUTPUT_PS: [BuiltinParamDescriptor; 2] = [
36    BuiltinParamDescriptor {
37        name: "p",
38        ty: BuiltinParamType::Any,
39        arity: BuiltinParamArity::Required,
40        default: None,
41        description: "Polynomial coefficient vector ordered from highest power to constant term.",
42    },
43    BuiltinParamDescriptor {
44        name: "S",
45        ty: BuiltinParamType::Any,
46        arity: BuiltinParamArity::Required,
47        default: None,
48        description: "Fit statistics structure with fields R, df, and normr.",
49    },
50];
51
52const POLYFIT_OUTPUT_PSMU: [BuiltinParamDescriptor; 3] = [
53    BuiltinParamDescriptor {
54        name: "p",
55        ty: BuiltinParamType::Any,
56        arity: BuiltinParamArity::Required,
57        default: None,
58        description: "Polynomial coefficient vector ordered from highest power to constant term.",
59    },
60    BuiltinParamDescriptor {
61        name: "S",
62        ty: BuiltinParamType::Any,
63        arity: BuiltinParamArity::Required,
64        default: None,
65        description: "Fit statistics structure with fields R, df, and normr.",
66    },
67    BuiltinParamDescriptor {
68        name: "mu",
69        ty: BuiltinParamType::NumericArray,
70        arity: BuiltinParamArity::Required,
71        default: None,
72        description: "Centering and scaling vector [mean(x), std(x)].",
73    },
74];
75
76const POLYFIT_INPUTS: [BuiltinParamDescriptor; 3] = [
77    BuiltinParamDescriptor {
78        name: "X",
79        ty: BuiltinParamType::Any,
80        arity: BuiltinParamArity::Required,
81        default: None,
82        description: "Sample x-values as a numeric vector.",
83    },
84    BuiltinParamDescriptor {
85        name: "Y",
86        ty: BuiltinParamType::Any,
87        arity: BuiltinParamArity::Required,
88        default: None,
89        description: "Sample y-values as a numeric vector.",
90    },
91    BuiltinParamDescriptor {
92        name: "n",
93        ty: BuiltinParamType::IntegerScalar,
94        arity: BuiltinParamArity::Required,
95        default: None,
96        description: "Polynomial degree.",
97    },
98];
99
100const POLYFIT_INPUTS_WEIGHTS: [BuiltinParamDescriptor; 4] = [
101    BuiltinParamDescriptor {
102        name: "X",
103        ty: BuiltinParamType::Any,
104        arity: BuiltinParamArity::Required,
105        default: None,
106        description: "Sample x-values as a numeric vector.",
107    },
108    BuiltinParamDescriptor {
109        name: "Y",
110        ty: BuiltinParamType::Any,
111        arity: BuiltinParamArity::Required,
112        default: None,
113        description: "Sample y-values as a numeric vector.",
114    },
115    BuiltinParamDescriptor {
116        name: "n",
117        ty: BuiltinParamType::IntegerScalar,
118        arity: BuiltinParamArity::Required,
119        default: None,
120        description: "Polynomial degree.",
121    },
122    BuiltinParamDescriptor {
123        name: "weights",
124        ty: BuiltinParamType::Any,
125        arity: BuiltinParamArity::Optional,
126        default: None,
127        description: "Optional nonnegative weight vector matching X and Y length.",
128    },
129];
130
131const POLYFIT_SIGNATURES: [BuiltinSignatureDescriptor; 6] = [
132    BuiltinSignatureDescriptor {
133        label: "p = polyfit(X, Y, n)",
134        inputs: &POLYFIT_INPUTS,
135        outputs: &POLYFIT_OUTPUT_P,
136    },
137    BuiltinSignatureDescriptor {
138        label: "p = polyfit(X, Y, n, weights)",
139        inputs: &POLYFIT_INPUTS_WEIGHTS,
140        outputs: &POLYFIT_OUTPUT_P,
141    },
142    BuiltinSignatureDescriptor {
143        label: "[p, S] = polyfit(X, Y, n)",
144        inputs: &POLYFIT_INPUTS,
145        outputs: &POLYFIT_OUTPUT_PS,
146    },
147    BuiltinSignatureDescriptor {
148        label: "[p, S] = polyfit(X, Y, n, weights)",
149        inputs: &POLYFIT_INPUTS_WEIGHTS,
150        outputs: &POLYFIT_OUTPUT_PS,
151    },
152    BuiltinSignatureDescriptor {
153        label: "[p, S, mu] = polyfit(X, Y, n)",
154        inputs: &POLYFIT_INPUTS,
155        outputs: &POLYFIT_OUTPUT_PSMU,
156    },
157    BuiltinSignatureDescriptor {
158        label: "[p, S, mu] = polyfit(X, Y, n, weights)",
159        inputs: &POLYFIT_INPUTS_WEIGHTS,
160        outputs: &POLYFIT_OUTPUT_PSMU,
161    },
162];
163
164const POLYFIT_ERROR_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
165    code: "RM.POLYFIT.INVALID_ARGUMENT",
166    identifier: Some("RunMat:polyfit:InvalidArgument"),
167    when: "Degree/weight arguments are malformed or unsupported.",
168    message: "polyfit: invalid argument",
169};
170
171const POLYFIT_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
172    code: "RM.POLYFIT.INVALID_INPUT",
173    identifier: Some("RunMat:polyfit:InvalidInput"),
174    when: "Input vectors or values cannot be processed for polynomial fitting.",
175    message: "polyfit: invalid input",
176};
177
178const POLYFIT_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
179    code: "RM.POLYFIT.INTERNAL",
180    identifier: Some("RunMat:polyfit:Internal"),
181    when: "Runtime fails while constructing fit outputs or provider payloads.",
182    message: "polyfit: internal runtime failure",
183};
184
185const POLYFIT_ERRORS: [BuiltinErrorDescriptor; 3] = [
186    POLYFIT_ERROR_INVALID_ARGUMENT,
187    POLYFIT_ERROR_INVALID_INPUT,
188    POLYFIT_ERROR_INTERNAL,
189];
190
191pub const POLYFIT_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
192    signatures: &POLYFIT_SIGNATURES,
193    output_mode: BuiltinOutputMode::ByRequestedOutputCount,
194    completion_policy: BuiltinCompletionPolicy::Public,
195    errors: &POLYFIT_ERRORS,
196};
197
198#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::math::poly::polyfit")]
199pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
200    name: "polyfit",
201    op_kind: GpuOpKind::Custom("polyfit"),
202    supported_precisions: &[ScalarType::F32, ScalarType::F64],
203    broadcast: BroadcastSemantics::Matlab,
204    provider_hooks: &[ProviderHook::Custom("polyfit")],
205    constant_strategy: ConstantStrategy::UniformBuffer,
206    residency: ResidencyPolicy::GatherImmediately,
207    nan_mode: ReductionNaN::Include,
208    two_pass_threshold: None,
209    workgroup_size: None,
210    accepts_nan_mode: false,
211    notes:
212        "Providers may gather to the host and invoke the shared Householder QR solver; WGPU implements this path today.",
213};
214
215fn polyfit_error_with(
216    message: impl Into<String>,
217    error: &'static BuiltinErrorDescriptor,
218) -> RuntimeError {
219    let mut builder = build_runtime_error(message).with_builtin(BUILTIN_NAME);
220    if let Some(identifier) = error.identifier {
221        builder = builder.with_identifier(identifier);
222    }
223    builder.build()
224}
225
226fn polyfit_error(message: impl Into<String>) -> RuntimeError {
227    polyfit_error_with(message, &POLYFIT_ERROR_INVALID_INPUT)
228}
229
230fn polyfit_argument_error(message: impl Into<String>) -> RuntimeError {
231    polyfit_error_with(message, &POLYFIT_ERROR_INVALID_ARGUMENT)
232}
233
234fn polyfit_internal_error(message: impl Into<String>) -> RuntimeError {
235    polyfit_error_with(message, &POLYFIT_ERROR_INTERNAL)
236}
237
238#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::math::poly::polyfit")]
239pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
240    name: "polyfit",
241    shape: ShapeRequirements::Any,
242    constant_strategy: ConstantStrategy::UniformBuffer,
243    elementwise: None,
244    reduction: None,
245    emits_nan: false,
246    notes: "Acts as a sink node—polynomial fitting materialises results eagerly and terminates fusion graphs.",
247};
248
249#[runtime_builtin(
250    name = "polyfit",
251    category = "math/poly",
252    summary = "Fit polynomials to data using least squares.",
253    keywords = "polyfit,polynomial,least-squares,gpu",
254    accel = "sink",
255    sink = true,
256    type_resolver(polyfit_type),
257    descriptor(crate::builtins::math::poly::polyfit::POLYFIT_DESCRIPTOR),
258    builtin_path = "crate::builtins::math::poly::polyfit"
259)]
260async fn polyfit_builtin(
261    x: Value,
262    y: Value,
263    degree: Value,
264    rest: Vec<Value>,
265) -> crate::BuiltinResult<Value> {
266    let eval = evaluate(x, y, degree, &rest).await?;
267    if let Some(out_count) = crate::output_count::current_output_count() {
268        if out_count == 0 {
269            return Ok(Value::OutputList(Vec::new()));
270        }
271        let mut outputs = vec![eval.coefficients()];
272        if out_count >= 2 {
273            outputs.push(eval.stats());
274        }
275        if out_count >= 3 {
276            outputs.push(eval.mu());
277        }
278        return Ok(crate::output_count::output_list_with_padding(
279            out_count, outputs,
280        ));
281    }
282    Ok(eval.coefficients())
283}
284
285/// Evaluate `polyfit`, returning the multi-output envelope used by the VM.
286pub async fn evaluate(
287    x: Value,
288    y: Value,
289    degree: Value,
290    rest: &[Value],
291) -> BuiltinResult<PolyfitEval> {
292    let deg = parse_degree(&degree)?;
293
294    if let Some(eval) = try_gpu_polyfit(&x, &y, deg, rest).await? {
295        return Ok(eval);
296    }
297
298    let x_host = dispatcher::gather_if_needed_async(&x).await?;
299    let y_host = dispatcher::gather_if_needed_async(&y).await?;
300
301    let x_data = real_vector("polyfit", "X", x_host).await?;
302    let (y_data, is_complex_input) = complex_vector("polyfit", "Y", y_host).await?;
303
304    if x_data.len() != y_data.len() {
305        return Err(polyfit_error(
306            "polyfit: X and Y vectors must be the same length",
307        ));
308    }
309    if x_data.is_empty() {
310        return Err(polyfit_error(
311            "polyfit: X and Y must contain at least one sample",
312        ));
313    }
314    if deg + 1 > x_data.len() && x_data.len() > 1 {
315        warn!(
316            "polyfit: polynomial degree {} is ill-conditioned for {} data points; results may be inaccurate",
317            deg,
318            x_data.len()
319        );
320    }
321
322    let weights = parse_weights(rest, x_data.len()).await?;
323    let mut solution = solve_polyfit(&x_data, &y_data, deg, weights.as_deref())?;
324    if is_complex_input {
325        solution.is_complex = true;
326    }
327
328    PolyfitEval::from_solution(solution)
329}
330
331async fn try_gpu_polyfit(
332    x: &Value,
333    y: &Value,
334    degree: usize,
335    rest: &[Value],
336) -> BuiltinResult<Option<PolyfitEval>> {
337    let provider = match runmat_accelerate_api::provider() {
338        Some(p) => p,
339        None => return Ok(None),
340    };
341
342    let x_handle = match x {
343        Value::GpuTensor(handle) => handle,
344        _ => return Ok(None),
345    };
346    let y_handle = match y {
347        Value::GpuTensor(handle) => handle,
348        _ => return Ok(None),
349    };
350
351    if rest.len() > 1 {
352        return Ok(None);
353    }
354
355    let weight_handle = match rest.first() {
356        Some(Value::GpuTensor(handle)) => Some(handle),
357        Some(_) => return Ok(None),
358        None => None,
359    };
360
361    let result = match provider
362        .polyfit(x_handle, y_handle, degree, weight_handle)
363        .await
364    {
365        Ok(res) => res,
366        Err(err) => {
367            trace!("polyfit: provider path unavailable ({err}); falling back to host");
368            return Ok(None);
369        }
370    };
371
372    let solution = PolyfitSolution::from_provider(result)?;
373    PolyfitEval::from_solution(solution).map(Some)
374}
375
376#[derive(Clone, Debug)]
377struct PolyfitSolution {
378    coeffs: Vec<Complex64>,
379    r_matrix: Vec<f64>,
380    mu_mean: f64,
381    mu_scale: f64,
382    normr: f64,
383    df: f64,
384    cols: usize,
385    is_complex: bool,
386}
387
388impl PolyfitSolution {
389    fn from_provider(result: ProviderPolyfitResult) -> BuiltinResult<Self> {
390        let cols = result.coefficients.len();
391        if cols == 0 {
392            return Err(polyfit_internal_error(
393                "polyfit: provider returned empty coefficient vector",
394            ));
395        }
396        if result.r_matrix.len() != cols * cols {
397            return Err(polyfit_internal_error(
398                "polyfit: provider returned malformed R matrix",
399            ));
400        }
401        let [mu_mean, mu_scale] = result.mu;
402        Ok(Self {
403            coeffs: result
404                .coefficients
405                .into_iter()
406                .map(|re| Complex64::new(re, 0.0))
407                .collect(),
408            r_matrix: result.r_matrix,
409            mu_mean,
410            mu_scale,
411            normr: result.normr,
412            df: result.df,
413            cols,
414            is_complex: false,
415        })
416    }
417}
418
419/// Multi-output envelope for `polyfit`, mirroring MATLAB semantics.
420#[derive(Debug)]
421pub struct PolyfitEval {
422    coefficients: Value,
423    stats: Value,
424    mu: Value,
425    is_complex: bool,
426}
427
428impl PolyfitEval {
429    fn from_solution(solution: PolyfitSolution) -> BuiltinResult<Self> {
430        let coefficients = coefficients_to_value(&solution.coeffs)?;
431        let stats = build_stats(
432            &solution.r_matrix,
433            solution.cols,
434            solution.normr,
435            solution.df,
436        )?;
437        let mu = build_mu(solution.mu_mean, solution.mu_scale)?;
438        Ok(Self {
439            coefficients,
440            stats,
441            mu,
442            is_complex: solution.is_complex,
443        })
444    }
445
446    /// Polynomial coefficients ordered from highest power to constant term.
447    pub fn coefficients(&self) -> Value {
448        self.coefficients.clone()
449    }
450
451    /// Structure `S` containing fields `R`, `df`, and `normr`.
452    pub fn stats(&self) -> Value {
453        self.stats.clone()
454    }
455
456    /// Centering and scaling vector `[mean(x), std(x)]`.
457    pub fn mu(&self) -> Value {
458        self.mu.clone()
459    }
460
461    /// Returns `true` if the fitted polynomial contains a complex coefficient.
462    pub fn is_complex(&self) -> bool {
463        self.is_complex
464    }
465}
466
467fn parse_degree(value: &Value) -> BuiltinResult<usize> {
468    match value {
469        Value::Int(i) => {
470            let raw = i.to_i64();
471            if raw < 0 {
472                return Err(polyfit_argument_error(
473                    "polyfit: degree must be a non-negative integer",
474                ));
475            }
476            Ok(raw as usize)
477        }
478        Value::Num(n) => {
479            if !n.is_finite() {
480                return Err(polyfit_argument_error("polyfit: degree must be finite"));
481            }
482            let rounded = n.round();
483            if (rounded - n).abs() > EPS {
484                return Err(polyfit_argument_error("polyfit: degree must be an integer"));
485            }
486            if rounded < 0.0 {
487                return Err(polyfit_argument_error(
488                    "polyfit: degree must be a non-negative integer",
489                ));
490            }
491            Ok(rounded as usize)
492        }
493        Value::Tensor(t) if tensor::is_scalar_tensor(t) => parse_degree(&Value::Num(t.data[0])),
494        Value::LogicalArray(l) if l.len() == 1 => {
495            parse_degree(&Value::Num(if l.data[0] != 0 { 1.0 } else { 0.0 }))
496        }
497        other => Err(polyfit_argument_error(format!(
498            "polyfit: degree must be a scalar numeric value, got {other:?}"
499        ))),
500    }
501}
502
503#[async_recursion::async_recursion(?Send)]
504async fn real_vector(context: &str, label: &str, value: Value) -> BuiltinResult<Vec<f64>> {
505    match value {
506        Value::Tensor(mut tensor) => {
507            ensure_vector_shape(context, label, &tensor.shape)?;
508            Ok(tensor.data.drain(..).collect())
509        }
510        Value::LogicalArray(logical) => {
511            let tensor = tensor::logical_to_tensor(&logical).map_err(polyfit_error)?;
512            ensure_vector_shape(context, label, &tensor.shape)?;
513            Ok(tensor.data)
514        }
515        Value::Num(n) => Ok(vec![n]),
516        Value::Int(i) => Ok(vec![i.to_f64()]),
517        Value::Bool(b) => Ok(vec![if b { 1.0 } else { 0.0 }]),
518        Value::GpuTensor(handle) => {
519            let gathered =
520                crate::builtins::common::gpu_helpers::gather_tensor_async(&handle).await?;
521            real_vector(context, label, Value::Tensor(gathered)).await
522        }
523        Value::Complex(_, _) | Value::ComplexTensor(_) => Err(polyfit_error(format!(
524            "{context}: {label} must be real-valued; complex inputs are not supported"
525        ))),
526        other => Err(polyfit_error(format!(
527            "{context}: expected {label} to be a numeric vector, got {other:?}"
528        ))),
529    }
530}
531
532#[async_recursion::async_recursion(?Send)]
533async fn complex_vector(
534    context: &str,
535    label: &str,
536    value: Value,
537) -> BuiltinResult<(Vec<Complex64>, bool)> {
538    match value {
539        Value::Tensor(mut tensor) => {
540            ensure_vector_shape(context, label, &tensor.shape)?;
541            let all_real = true;
542            let data = tensor
543                .data
544                .drain(..)
545                .map(|x| Complex64::new(x, 0.0))
546                .collect();
547            Ok((data, all_real))
548        }
549        Value::ComplexTensor(tensor) => {
550            ensure_vector_shape(context, label, &tensor.shape)?;
551            let is_complex = tensor.data.iter().any(|&(_, im)| im.abs() > EPS);
552            let data = tensor
553                .data
554                .into_iter()
555                .map(|(re, im)| Complex64::new(re, im))
556                .collect::<Vec<_>>();
557            Ok((data, is_complex))
558        }
559        Value::LogicalArray(logical) => {
560            let tensor = tensor::logical_to_tensor(&logical).map_err(polyfit_error)?;
561            ensure_vector_shape(context, label, &tensor.shape)?;
562            Ok((
563                tensor
564                    .data
565                    .iter()
566                    .map(|&x| Complex64::new(x, 0.0))
567                    .collect(),
568                false,
569            ))
570        }
571        Value::Num(n) => Ok((vec![Complex64::new(n, 0.0)], false)),
572        Value::Int(i) => Ok((vec![Complex64::new(i.to_f64(), 0.0)], false)),
573        Value::Bool(b) => Ok((vec![Complex64::new(if b { 1.0 } else { 0.0 }, 0.0)], false)),
574        Value::Complex(re, im) => Ok((vec![Complex64::new(re, im)], im.abs() > EPS)),
575        Value::GpuTensor(handle) => {
576            let gathered =
577                crate::builtins::common::gpu_helpers::gather_tensor_async(&handle).await?;
578            complex_vector(context, label, Value::Tensor(gathered)).await
579        }
580        other => Err(polyfit_error(format!(
581            "{context}: expected {label} to be a numeric vector, got {other:?}"
582        ))),
583    }
584}
585
586async fn parse_weights(rest: &[Value], len: usize) -> BuiltinResult<Option<Vec<f64>>> {
587    match rest.len() {
588        0 => Ok(None),
589        1 => {
590            let gathered = dispatcher::gather_if_needed_async(&rest[0]).await?;
591            let data = real_vector("polyfit", "weights", gathered).await?;
592            if data.len() != len {
593                return Err(polyfit_argument_error(
594                    "polyfit: weight vector must match the size of X",
595                ));
596            }
597            validate_weights(&data)?;
598            Ok(Some(data))
599        }
600        _ => Err(polyfit_argument_error("polyfit: too many input arguments")),
601    }
602}
603
604fn validate_weights(weights: &[f64]) -> BuiltinResult<()> {
605    for (idx, w) in weights.iter().enumerate() {
606        if !w.is_finite() {
607            return Err(polyfit_argument_error(format!(
608                "polyfit: weight at position {} must be finite",
609                idx + 1
610            )));
611        }
612        if *w < 0.0 {
613            return Err(polyfit_argument_error(
614                "polyfit: weights must be non-negative",
615            ));
616        }
617    }
618    Ok(())
619}
620
621fn solve_polyfit(
622    x_data: &[f64],
623    y_data: &[Complex64],
624    degree: usize,
625    weights: Option<&[f64]>,
626) -> BuiltinResult<PolyfitSolution> {
627    if x_data.len() != y_data.len() {
628        return Err(polyfit_error(
629            "polyfit: X and Y vectors must be the same length",
630        ));
631    }
632    if x_data.is_empty() {
633        return Err(polyfit_error(
634            "polyfit: X and Y must contain at least one sample",
635        ));
636    }
637    if let Some(w) = weights {
638        if w.len() != x_data.len() {
639            return Err(polyfit_argument_error(
640                "polyfit: weight vector must match the size of X",
641            ));
642        }
643        validate_weights(w)?;
644    }
645
646    let mean = x_data.iter().sum::<f64>() / x_data.len() as f64;
647    if !mean.is_finite() {
648        return Err(polyfit_error("polyfit: mean of X must be finite"));
649    }
650    let scale = compute_scale(x_data, mean)?;
651    let scaled: Vec<f64> = x_data.iter().map(|&v| (v - mean) / scale).collect();
652
653    let mut rhs = y_data.to_vec();
654    for (idx, value) in rhs.iter().enumerate() {
655        if !value.re.is_finite() || !value.im.is_finite() {
656            return Err(polyfit_error(format!(
657                "polyfit: Y must contain finite values (encountered NaN/Inf at position {})",
658                idx + 1
659            )));
660        }
661    }
662    if let Some(w) = weights {
663        apply_weights_rhs(&mut rhs, w)?;
664    }
665
666    let rows = scaled.len();
667    let cols = degree + 1;
668    let mut vandermonde = build_vandermonde(&scaled, cols);
669    if let Some(w) = weights {
670        apply_weights_matrix(&mut vandermonde, rows, cols, w)?;
671    }
672
673    let mut transformed_rhs = rhs.clone();
674    householder_qr(&mut vandermonde, rows, cols, &mut transformed_rhs)?;
675    let coeff_scaled = solve_upper(&vandermonde, rows, cols, &transformed_rhs)?;
676    let coeff_original = transform_coefficients(&coeff_scaled, mean, scale);
677
678    let normr = residual_norm(&transformed_rhs, rows, cols);
679    let df = if rows > cols {
680        (rows - cols) as f64
681    } else {
682        0.0
683    };
684    let r_matrix = extract_upper(&vandermonde, rows, cols);
685    let is_complex = coeff_original.iter().any(|c| c.im.abs() > EPS_NAN);
686
687    Ok(PolyfitSolution {
688        coeffs: coeff_original,
689        r_matrix,
690        mu_mean: mean,
691        mu_scale: scale,
692        normr,
693        df,
694        cols,
695        is_complex,
696    })
697}
698
699fn compute_scale(data: &[f64], mean: f64) -> BuiltinResult<f64> {
700    if data.len() <= 1 {
701        return Ok(1.0);
702    }
703    let mut acc = 0.0;
704    for &value in data {
705        if !value.is_finite() {
706            return Err(polyfit_error("polyfit: X must contain finite values"));
707        }
708        let diff = value - mean;
709        acc += diff * diff;
710    }
711    let denom = (data.len() as f64 - 1.0).max(1.0);
712    let std = (acc / denom).sqrt();
713    let scale = if std.abs() <= EPS { 1.0 } else { std };
714    if !scale.is_finite() {
715        return Err(polyfit_error(
716            "polyfit: failed to compute a stable scaling factor",
717        ));
718    }
719    Ok(scale)
720}
721
722fn build_vandermonde(u: &[f64], cols: usize) -> Vec<f64> {
723    let rows = u.len();
724    let mut matrix = vec![0.0; rows * cols];
725    if cols == 0 {
726        return matrix;
727    }
728    for (row_idx, &value) in u.iter().enumerate() {
729        let mut powers = vec![0.0; cols];
730        powers[cols - 1] = 1.0;
731        for idx in (0..cols - 1).rev() {
732            powers[idx] = powers[idx + 1] * value;
733        }
734        for col_idx in 0..cols {
735            matrix[row_idx + col_idx * rows] = powers[col_idx];
736        }
737    }
738    matrix
739}
740
741fn apply_weights_matrix(
742    matrix: &mut [f64],
743    rows: usize,
744    cols: usize,
745    weights: &[f64],
746) -> BuiltinResult<()> {
747    for (row, weight) in weights.iter().enumerate().take(rows) {
748        let sqrt_w = weight.sqrt();
749        if !sqrt_w.is_finite() {
750            return Err(polyfit_error(format!(
751                "polyfit: weight at position {} must be finite",
752                row + 1
753            )));
754        }
755        for col in 0..cols {
756            let idx = row + col * rows;
757            matrix[idx] *= sqrt_w;
758        }
759    }
760    Ok(())
761}
762
763fn apply_weights_rhs(rhs: &mut [Complex64], weights: &[f64]) -> BuiltinResult<()> {
764    for (idx, (value, weight)) in rhs.iter_mut().zip(weights.iter()).enumerate() {
765        let sqrt_w = weight.sqrt();
766        if !sqrt_w.is_finite() {
767            return Err(polyfit_error(format!(
768                "polyfit: weight at position {} must be finite",
769                idx + 1
770            )));
771        }
772        *value *= sqrt_w;
773    }
774    Ok(())
775}
776
777fn ensure_vector_shape(context: &str, label: &str, shape: &[usize]) -> BuiltinResult<()> {
778    if !is_vector_shape(shape) {
779        return Err(polyfit_error(format!(
780            "{context}: {label} must be a vector"
781        )));
782    }
783    Ok(())
784}
785
786fn is_vector_shape(shape: &[usize]) -> bool {
787    shape.iter().copied().filter(|&dim| dim > 1).count() <= 1
788}
789
790fn householder_qr(
791    matrix: &mut [f64],
792    rows: usize,
793    cols: usize,
794    rhs: &mut [Complex64],
795) -> BuiltinResult<()> {
796    let min_dim = rows.min(cols);
797    for k in 0..min_dim {
798        let mut norm_sq = 0.0;
799        for row in k..rows {
800            let val = matrix[row + k * rows];
801            norm_sq += val * val;
802        }
803        if norm_sq <= EPS {
804            continue;
805        }
806        let norm = norm_sq.sqrt();
807        let x0 = matrix[k + k * rows];
808        let alpha = if x0 >= 0.0 { -norm } else { norm };
809        let mut v = vec![0.0; rows - k];
810        v[0] = x0 - alpha;
811        for row in (k + 1)..rows {
812            v[row - k] = matrix[row + k * rows];
813        }
814        let v_norm_sq: f64 = v.iter().map(|&x| x * x).sum();
815        if v_norm_sq <= EPS {
816            continue;
817        }
818        let beta = 2.0 / v_norm_sq;
819        matrix[k + k * rows] = alpha;
820        for row in (k + 1)..rows {
821            matrix[row + k * rows] = 0.0;
822        }
823        for col in (k + 1)..cols {
824            let mut dot = 0.0;
825            for (idx, &vi) in v.iter().enumerate() {
826                let row_idx = k + idx;
827                dot += vi * matrix[row_idx + col * rows];
828            }
829            let factor = beta * dot;
830            for (idx, &vi) in v.iter().enumerate() {
831                let row_idx = k + idx;
832                matrix[row_idx + col * rows] -= factor * vi;
833            }
834        }
835        let mut dot = Complex64::new(0.0, 0.0);
836        for (idx, &vi) in v.iter().enumerate() {
837            let row_idx = k + idx;
838            dot += rhs[row_idx] * vi;
839        }
840        let factor = Complex64::new(beta, 0.0) * dot;
841        for (idx, &vi) in v.iter().enumerate() {
842            let row_idx = k + idx;
843            rhs[row_idx] -= factor * vi;
844        }
845    }
846    Ok(())
847}
848
849fn solve_upper(
850    matrix: &[f64],
851    rows: usize,
852    cols: usize,
853    rhs: &[Complex64],
854) -> BuiltinResult<Vec<Complex64>> {
855    if rhs.len() < rows {
856        return Err(polyfit_error(
857            "polyfit internal error: RHS dimension mismatch",
858        ));
859    }
860    let mut coeffs = vec![Complex64::new(0.0, 0.0); cols];
861    for col in (0..cols).rev() {
862        let diag = if col < rows {
863            matrix[col + col * rows]
864        } else {
865            0.0
866        };
867        if diag.abs() <= EPS {
868            coeffs[col] = Complex64::new(0.0, 0.0);
869            continue;
870        }
871        let mut acc = if col < rows {
872            rhs[col]
873        } else {
874            Complex64::new(0.0, 0.0)
875        };
876        for next in (col + 1)..cols {
877            let idx = if col < rows {
878                matrix[col + next * rows]
879            } else {
880                0.0
881            };
882            acc -= Complex64::new(idx, 0.0) * coeffs[next];
883        }
884        coeffs[col] = acc / Complex64::new(diag, 0.0);
885    }
886    Ok(coeffs)
887}
888
889fn residual_norm(rhs: &[Complex64], rows: usize, cols: usize) -> f64 {
890    let tail_start = rows.min(cols);
891    let mut acc = 0.0;
892    for value in rhs.iter().skip(tail_start) {
893        acc += value.norm_sqr();
894    }
895    acc.sqrt()
896}
897
898fn extract_upper(matrix: &[f64], rows: usize, cols: usize) -> Vec<f64> {
899    let mut output = vec![0.0; cols * cols];
900    for col in 0..cols {
901        for row in 0..=col {
902            if row < rows {
903                output[row + col * cols] = matrix[row + col * rows];
904            }
905        }
906    }
907    output
908}
909
910fn transform_coefficients(coeffs: &[Complex64], mean: f64, scale: f64) -> Vec<Complex64> {
911    let mut poly: Vec<Complex64> = Vec::new();
912    for &coeff in coeffs {
913        let mut next = vec![Complex64::new(0.0, 0.0); poly.len() + 1];
914        for (idx, &value) in poly.iter().enumerate() {
915            next[idx + 1] += value / scale;
916            next[idx] -= value * (mean / scale);
917        }
918        next[0] += coeff;
919        poly = next;
920    }
921    poly.reverse();
922    poly
923}
924
925fn coefficients_to_value(coeffs: &[Complex64]) -> BuiltinResult<Value> {
926    let all_real = coeffs
927        .iter()
928        .all(|c| c.im.abs() <= EPS_NAN && c.re.is_finite());
929    if all_real {
930        let data: Vec<f64> = coeffs.iter().map(|c| c.re).collect();
931        let tensor = Tensor::new(data, vec![1, coeffs.len()])
932            .map_err(|e| polyfit_error(format!("polyfit: {e}")))?;
933        Ok(Value::Tensor(tensor))
934    } else {
935        let data: Vec<(f64, f64)> = coeffs.iter().map(|c| (c.re, c.im)).collect();
936        let tensor = ComplexTensor::new(data, vec![1, coeffs.len()])
937            .map_err(|e| polyfit_error(format!("polyfit: {e}")))?;
938        Ok(Value::ComplexTensor(tensor))
939    }
940}
941
942fn build_stats(r: &[f64], n: usize, normr: f64, df: f64) -> BuiltinResult<Value> {
943    let tensor =
944        Tensor::new(r.to_vec(), vec![n, n]).map_err(|e| polyfit_error(format!("polyfit: {e}")))?;
945    let mut st = StructValue::new();
946    st.fields.insert("R".to_string(), Value::Tensor(tensor));
947    st.fields.insert("df".to_string(), Value::Num(df));
948    st.fields.insert("normr".to_string(), Value::Num(normr));
949    Ok(Value::Struct(st))
950}
951
952fn build_mu(mean: f64, scale: f64) -> BuiltinResult<Value> {
953    if !scale.is_finite() || scale.abs() <= EPS {
954        return Err(polyfit_error("polyfit: mu(2) must be non-zero and finite"));
955    }
956    let tensor = Tensor::new(vec![mean, scale], vec![1, 2])
957        .map_err(|e| polyfit_error(format!("polyfit: {e}")))?;
958    Ok(Value::Tensor(tensor))
959}
960
961#[derive(Debug, Clone)]
962pub struct PolyfitHostRealResult {
963    pub coefficients: Vec<f64>,
964    pub r_matrix: Vec<f64>,
965    pub mu: [f64; 2],
966    pub normr: f64,
967    pub df: f64,
968}
969
970pub fn polyfit_host_real_for_provider(
971    x: &[f64],
972    y: &[f64],
973    degree: usize,
974    weights: Option<&[f64]>,
975) -> BuiltinResult<PolyfitHostRealResult> {
976    if x.len() != y.len() {
977        return Err(polyfit_error(
978            "polyfit: X and Y vectors must be the same length",
979        ));
980    }
981    if let Some(w) = weights {
982        if w.len() != x.len() {
983            return Err(polyfit_error(
984                "polyfit: weight vector must match the size of X",
985            ));
986        }
987        validate_weights(w)?;
988    }
989    let complex_y: Vec<Complex64> = y.iter().copied().map(|v| Complex64::new(v, 0.0)).collect();
990    let solution = solve_polyfit(x, &complex_y, degree, weights)?;
991    let PolyfitSolution {
992        coeffs,
993        r_matrix,
994        mu_mean,
995        mu_scale,
996        normr,
997        df,
998        cols: _,
999        is_complex,
1000    } = solution;
1001    if is_complex {
1002        return Err(polyfit_error(
1003            "polyfit: provider fallback produced complex coefficients for real data",
1004        ));
1005    }
1006    let coeffs: Vec<f64> = coeffs.into_iter().map(|c| c.re).collect();
1007    let mu = [mu_mean, mu_scale];
1008    Ok(PolyfitHostRealResult {
1009        coefficients: coeffs,
1010        r_matrix,
1011        mu,
1012        normr,
1013        df,
1014    })
1015}
1016
1017#[cfg(test)]
1018pub(crate) mod tests {
1019    use super::*;
1020    use crate::builtins::common::test_support;
1021    use futures::executor::block_on;
1022
1023    fn assert_error_contains(err: crate::RuntimeError, needle: &str) {
1024        assert!(
1025            err.message().contains(needle),
1026            "expected error containing '{needle}', got '{}'",
1027            err.message()
1028        );
1029    }
1030
1031    fn evaluate(
1032        x: Value,
1033        y: Value,
1034        degree: Value,
1035        rest: &[Value],
1036    ) -> Result<PolyfitEval, RuntimeError> {
1037        block_on(super::evaluate(x, y, degree, rest))
1038    }
1039
1040    #[test]
1041    fn polyfit_descriptor_signatures_cover_core_forms() {
1042        let labels: Vec<&str> = POLYFIT_DESCRIPTOR
1043            .signatures
1044            .iter()
1045            .map(|signature| signature.label)
1046            .collect();
1047        assert!(labels.contains(&"p = polyfit(X, Y, n)"));
1048        assert!(labels.contains(&"p = polyfit(X, Y, n, weights)"));
1049        assert!(labels.contains(&"[p, S] = polyfit(X, Y, n)"));
1050        assert!(labels.contains(&"[p, S, mu] = polyfit(X, Y, n, weights)"));
1051    }
1052
1053    #[test]
1054    fn polyfit_descriptor_errors_have_stable_codes() {
1055        let codes: Vec<&str> = POLYFIT_DESCRIPTOR
1056            .errors
1057            .iter()
1058            .map(|error| error.code)
1059            .collect();
1060        assert!(codes.contains(&"RM.POLYFIT.INVALID_ARGUMENT"));
1061        assert!(codes.contains(&"RM.POLYFIT.INVALID_INPUT"));
1062        assert!(codes.contains(&"RM.POLYFIT.INTERNAL"));
1063    }
1064
1065    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1066    #[test]
1067    fn fits_linear_data() {
1068        let x = Tensor::new(vec![0.0, 1.0, 2.0, 3.0], vec![4, 1]).unwrap();
1069        let mut y_vals = Vec::new();
1070        for i in 0..4 {
1071            y_vals.push(1.5 * i as f64 + 2.0);
1072        }
1073        let y = Tensor::new(y_vals, vec![4, 1]).unwrap();
1074        let eval = evaluate(
1075            Value::Tensor(x),
1076            Value::Tensor(y),
1077            Value::Int(runmat_builtins::IntValue::I32(1)),
1078            &[],
1079        )
1080        .expect("polyfit");
1081        match eval.coefficients() {
1082            Value::Tensor(t) => {
1083                assert_eq!(t.shape, vec![1, 2]);
1084                assert!((t.data[0] - 1.5).abs() < 1e-10);
1085                assert!((t.data[1] - 2.0).abs() < 1e-10);
1086            }
1087            other => panic!("expected tensor coefficients, got {other:?}"),
1088        }
1089    }
1090
1091    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1092    #[test]
1093    fn returns_struct_and_mu() {
1094        let x = Tensor::new(vec![-1.0, 0.0, 1.0], vec![3, 1]).unwrap();
1095        let y = Tensor::new(vec![1.0, 0.0, 1.0], vec![3, 1]).unwrap();
1096        let eval = evaluate(
1097            Value::Tensor(x),
1098            Value::Tensor(y),
1099            Value::Int(runmat_builtins::IntValue::I32(2)),
1100            &[],
1101        )
1102        .expect("polyfit");
1103        match eval.stats() {
1104            Value::Struct(s) => {
1105                assert!(s.fields.contains_key("R"));
1106                assert!(s.fields.contains_key("df"));
1107                assert!(s.fields.contains_key("normr"));
1108            }
1109            other => panic!("expected struct, got {other:?}"),
1110        }
1111        match eval.mu() {
1112            Value::Tensor(t) => {
1113                assert_eq!(t.shape, vec![1, 2]);
1114                assert!((t.data[0]).abs() < 1e-10);
1115                assert!(t.data[1].abs() > 0.0);
1116            }
1117            other => panic!("expected tensor mu, got {other:?}"),
1118        }
1119    }
1120
1121    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1122    #[test]
1123    fn weighted_fit_matches_unweighted_when_weights_equal() {
1124        let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
1125        let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
1126        let weights = Tensor::new(vec![1.0, 1.0, 1.0], vec![3, 1]).unwrap();
1127        let eval_unweighted = evaluate(
1128            Value::Tensor(x.clone()),
1129            Value::Tensor(y.clone()),
1130            Value::Int(runmat_builtins::IntValue::I32(2)),
1131            &[],
1132        )
1133        .expect("polyfit");
1134        let eval_weighted = evaluate(
1135            Value::Tensor(x),
1136            Value::Tensor(y),
1137            Value::Int(runmat_builtins::IntValue::I32(2)),
1138            &[Value::Tensor(weights)],
1139        )
1140        .expect("polyfit");
1141        assert_eq!(eval_unweighted.coefficients(), eval_weighted.coefficients());
1142    }
1143
1144    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1145    #[test]
1146    fn accepts_logical_degree_scalar() {
1147        let x = Tensor::new(vec![0.0, 1.0], vec![2, 1]).unwrap();
1148        let y = Tensor::new(vec![1.0, 3.0], vec![2, 1]).unwrap();
1149        let logical = runmat_builtins::LogicalArray::new(vec![1], vec![1, 1]).unwrap();
1150        let eval = evaluate(
1151            Value::Tensor(x),
1152            Value::Tensor(y),
1153            Value::LogicalArray(logical),
1154            &[],
1155        )
1156        .expect("polyfit");
1157        assert!(matches!(eval.coefficients(), Value::Tensor(_)));
1158    }
1159
1160    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1161    #[test]
1162    fn rejects_non_integer_degree() {
1163        let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
1164        let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
1165        let err = evaluate(Value::Tensor(x), Value::Tensor(y), Value::Num(1.5), &[])
1166            .expect_err("polyfit should reject non-integer degree");
1167        assert_eq!(err.identifier(), POLYFIT_ERROR_INVALID_ARGUMENT.identifier);
1168        assert_error_contains(err, "integer");
1169    }
1170
1171    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1172    #[test]
1173    fn rejects_infinite_weights() {
1174        let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
1175        let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
1176        let weights = Tensor::new(vec![1.0, f64::INFINITY, 1.0], vec![3, 1]).unwrap();
1177        let err = evaluate(
1178            Value::Tensor(x),
1179            Value::Tensor(y),
1180            Value::Int(runmat_builtins::IntValue::I32(2)),
1181            &[Value::Tensor(weights)],
1182        )
1183        .expect_err("polyfit should reject infinite weights");
1184        assert_error_contains(err, "weight at position 2");
1185    }
1186
1187    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1188    #[test]
1189    fn gpu_inputs_are_gathered() {
1190        test_support::with_test_provider(|provider| {
1191            let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
1192            let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
1193            let view = runmat_accelerate_api::HostTensorView {
1194                data: &x.data,
1195                shape: &x.shape,
1196            };
1197            let x_handle = provider.upload(&view).expect("upload");
1198            let view_y = runmat_accelerate_api::HostTensorView {
1199                data: &y.data,
1200                shape: &y.shape,
1201            };
1202            let y_handle = provider.upload(&view_y).expect("upload");
1203            let eval = evaluate(
1204                Value::GpuTensor(x_handle),
1205                Value::GpuTensor(y_handle),
1206                Value::Int(runmat_builtins::IntValue::I32(2)),
1207                &[],
1208            )
1209            .expect("polyfit");
1210            assert!(matches!(eval.coefficients(), Value::Tensor(_)));
1211        });
1212    }
1213
1214    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1215    #[test]
1216    fn gpu_weights_are_gathered() {
1217        test_support::with_test_provider(|provider| {
1218            let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
1219            let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
1220            let weights = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
1221
1222            let x_view = runmat_accelerate_api::HostTensorView {
1223                data: &x.data,
1224                shape: &x.shape,
1225            };
1226            let y_view = runmat_accelerate_api::HostTensorView {
1227                data: &y.data,
1228                shape: &y.shape,
1229            };
1230            let w_view = runmat_accelerate_api::HostTensorView {
1231                data: &weights.data,
1232                shape: &weights.shape,
1233            };
1234
1235            let x_handle = provider.upload(&x_view).expect("upload x");
1236            let y_handle = provider.upload(&y_view).expect("upload y");
1237            let w_handle = provider.upload(&w_view).expect("upload weights");
1238
1239            let cpu_eval = evaluate(
1240                Value::Tensor(x.clone()),
1241                Value::Tensor(y.clone()),
1242                Value::Int(runmat_builtins::IntValue::I32(2)),
1243                &[Value::Tensor(weights.clone())],
1244            )
1245            .expect("cpu polyfit");
1246
1247            let gpu_eval = evaluate(
1248                Value::GpuTensor(x_handle.clone()),
1249                Value::GpuTensor(y_handle.clone()),
1250                Value::Int(runmat_builtins::IntValue::I32(2)),
1251                &[Value::GpuTensor(w_handle.clone())],
1252            )
1253            .expect("gpu polyfit with weights");
1254
1255            assert_eq!(cpu_eval.coefficients(), gpu_eval.coefficients());
1256            assert_eq!(cpu_eval.mu(), gpu_eval.mu());
1257
1258            let _ = provider.free(&x_handle);
1259            let _ = provider.free(&y_handle);
1260            let _ = provider.free(&w_handle);
1261        });
1262    }
1263
1264    #[cfg(feature = "wgpu")]
1265    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1266    #[test]
1267    fn polyfit_wgpu_matches_cpu() {
1268        let options = runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default();
1269        let _provider =
1270            match runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(options) {
1271                Ok(p) => p,
1272                Err(err) => {
1273                    warn!("polyfit_wgpu_matches_cpu: skipping test ({err})");
1274                    return;
1275                }
1276            };
1277        let x = Tensor::new(vec![0.0, 1.0, 2.0, 3.0], vec![4, 1]).unwrap();
1278        let y = Tensor::new(vec![1.0, 3.0, 7.0, 13.0], vec![4, 1]).unwrap();
1279
1280        let cpu_eval = evaluate(
1281            Value::Tensor(x.clone()),
1282            Value::Tensor(y.clone()),
1283            Value::Int(runmat_builtins::IntValue::I32(2)),
1284            &[],
1285        )
1286        .expect("cpu polyfit");
1287
1288        let trait_provider = runmat_accelerate_api::provider().expect("wgpu provider registered");
1289        let x_view = runmat_accelerate_api::HostTensorView {
1290            data: &x.data,
1291            shape: &x.shape,
1292        };
1293        let y_view = runmat_accelerate_api::HostTensorView {
1294            data: &y.data,
1295            shape: &y.shape,
1296        };
1297        let x_handle = trait_provider.upload(&x_view).expect("upload x");
1298        let y_handle = trait_provider.upload(&y_view).expect("upload y");
1299
1300        let gpu_eval = evaluate(
1301            Value::GpuTensor(x_handle.clone()),
1302            Value::GpuTensor(y_handle.clone()),
1303            Value::Int(runmat_builtins::IntValue::I32(2)),
1304            &[],
1305        )
1306        .expect("gpu polyfit");
1307
1308        let _ = trait_provider.free(&x_handle);
1309        let _ = trait_provider.free(&y_handle);
1310
1311        let cpu_coeff = match cpu_eval.coefficients() {
1312            Value::Tensor(t) => t,
1313            other => panic!("expected tensor coefficients, got {other:?}"),
1314        };
1315        let gpu_coeff = match gpu_eval.coefficients() {
1316            Value::Tensor(t) => t,
1317            other => panic!("expected tensor coefficients, got {other:?}"),
1318        };
1319        assert_eq!(cpu_coeff.shape, gpu_coeff.shape);
1320        for (a, b) in cpu_coeff.data.iter().zip(gpu_coeff.data.iter()) {
1321            assert!((a - b).abs() < 1e-9, "coeff mismatch {a} vs {b}");
1322        }
1323
1324        let cpu_mu = match cpu_eval.mu() {
1325            Value::Tensor(t) => t,
1326            other => panic!("expected tensor mu, got {other:?}"),
1327        };
1328        let gpu_mu = match gpu_eval.mu() {
1329            Value::Tensor(t) => t,
1330            other => panic!("expected tensor mu, got {other:?}"),
1331        };
1332        assert_eq!(cpu_mu.shape, gpu_mu.shape);
1333        for (a, b) in cpu_mu.data.iter().zip(gpu_mu.data.iter()) {
1334            assert!((a - b).abs() < 1e-9, "mu mismatch {a} vs {b}");
1335        }
1336
1337        let cpu_stats = match cpu_eval.stats() {
1338            Value::Struct(s) => s,
1339            other => panic!("expected struct stats, got {other:?}"),
1340        };
1341        let gpu_stats = match gpu_eval.stats() {
1342            Value::Struct(s) => s,
1343            other => panic!("expected struct stats, got {other:?}"),
1344        };
1345        let cpu_r = match cpu_stats.fields.get("R").expect("R present") {
1346            Value::Tensor(t) => t.clone(),
1347            other => panic!("expected tensor R, got {other:?}"),
1348        };
1349        let gpu_r = match gpu_stats.fields.get("R").expect("R present") {
1350            Value::Tensor(t) => t.clone(),
1351            other => panic!("expected tensor R, got {other:?}"),
1352        };
1353        assert_eq!(cpu_r.shape, gpu_r.shape);
1354        for (a, b) in cpu_r.data.iter().zip(gpu_r.data.iter()) {
1355            assert!((a - b).abs() < 1e-9, "R mismatch {a} vs {b}");
1356        }
1357        let cpu_df = match cpu_stats.fields.get("df").expect("df present") {
1358            Value::Num(n) => *n,
1359            other => panic!("expected numeric df, got {other:?}"),
1360        };
1361        let gpu_df = match gpu_stats.fields.get("df").expect("df present") {
1362            Value::Num(n) => *n,
1363            other => panic!("expected numeric df, got {other:?}"),
1364        };
1365        assert!((cpu_df - gpu_df).abs() < 1e-9);
1366        let cpu_normr = match cpu_stats.fields.get("normr").expect("normr present") {
1367            Value::Num(n) => *n,
1368            other => panic!("expected numeric normr, got {other:?}"),
1369        };
1370        let gpu_normr = match gpu_stats.fields.get("normr").expect("normr present") {
1371            Value::Num(n) => *n,
1372            other => panic!("expected numeric normr, got {other:?}"),
1373        };
1374        assert!((cpu_normr - gpu_normr).abs() < 1e-9);
1375    }
1376
1377    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1378    #[test]
1379    fn rejects_mismatched_lengths() {
1380        let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
1381        let y = Tensor::new(vec![1.0, 2.0], vec![2, 1]).unwrap();
1382        let err = evaluate(
1383            Value::Tensor(x),
1384            Value::Tensor(y),
1385            Value::Int(runmat_builtins::IntValue::I32(1)),
1386            &[],
1387        )
1388        .expect_err("polyfit should reject mismatched vector lengths");
1389        assert_error_contains(err, "same length");
1390    }
1391
1392    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1393    #[test]
1394    fn rejects_non_vector_inputs() {
1395        let x = Tensor::new(vec![0.0, 1.0, 2.0, 3.0], vec![2, 2]).unwrap();
1396        let y = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();
1397        let err = evaluate(
1398            Value::Tensor(x),
1399            Value::Tensor(y),
1400            Value::Int(runmat_builtins::IntValue::I32(1)),
1401            &[],
1402        )
1403        .expect_err("polyfit should reject non-vector X");
1404        assert_error_contains(err, "vector");
1405    }
1406
1407    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1408    #[test]
1409    fn rejects_weight_length_mismatch() {
1410        let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
1411        let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
1412        let weights = Tensor::new(vec![1.0, 1.0], vec![2, 1]).unwrap();
1413        let err = evaluate(
1414            Value::Tensor(x),
1415            Value::Tensor(y),
1416            Value::Int(runmat_builtins::IntValue::I32(2)),
1417            &[Value::Tensor(weights)],
1418        )
1419        .expect_err("polyfit should reject mismatched weights");
1420        assert_error_contains(err, "weight vector must match");
1421    }
1422
1423    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1424    #[test]
1425    fn rejects_negative_weights() {
1426        let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
1427        let y = Tensor::new(vec![1.0, 3.0, 7.0], vec![3, 1]).unwrap();
1428        let weights = Tensor::new(vec![1.0, -1.0, 1.0], vec![3, 1]).unwrap();
1429        let err = evaluate(
1430            Value::Tensor(x),
1431            Value::Tensor(y),
1432            Value::Int(runmat_builtins::IntValue::I32(2)),
1433            &[Value::Tensor(weights)],
1434        )
1435        .expect_err("polyfit should reject negative weights");
1436        assert_error_contains(err, "non-negative");
1437    }
1438
1439    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1440    #[test]
1441    fn fits_complex_data() {
1442        let x = Tensor::new(vec![0.0, 1.0, 2.0], vec![3, 1]).unwrap();
1443        let complex_values =
1444            ComplexTensor::new(vec![(0.0, 1.0), (1.0, 0.5), (4.0, -0.25)], vec![3, 1]).unwrap();
1445        let eval = evaluate(
1446            Value::Tensor(x),
1447            Value::ComplexTensor(complex_values),
1448            Value::Int(runmat_builtins::IntValue::I32(2)),
1449            &[],
1450        )
1451        .expect("polyfit complex");
1452        match eval.coefficients() {
1453            Value::ComplexTensor(t) => {
1454                assert_eq!(t.shape, vec![1, 3]);
1455            }
1456            other => panic!("expected complex tensor coefficients, got {other:?}"),
1457        }
1458    }
1459}