1use 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
285pub async fn evaluate(
287 x: Value,
288 y: Value,
289 degree: Value,
290 rest: &[Value],
291) -> BuiltinResult<PolyfitEval> {
292 let deg = parse_degree(°ree)?;
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#[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 pub fn coefficients(&self) -> Value {
448 self.coefficients.clone()
449 }
450
451 pub fn stats(&self) -> Value {
453 self.stats.clone()
454 }
455
456 pub fn mu(&self) -> Value {
458 self.mu.clone()
459 }
460
461 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}