runmat_runtime/builtins/math/poly/
polyval.rs

1//! MATLAB-compatible `polyval` builtin with GPU-aware semantics for RunMat.
2
3use log::debug;
4use num_complex::Complex64;
5use runmat_accelerate_api::{HostTensorView, ProviderPolyvalMu, ProviderPolyvalOptions};
6use runmat_builtins::{ComplexTensor, LogicalArray, Tensor, Value};
7use runmat_macros::runtime_builtin;
8
9use crate::builtins::common::spec::{
10    BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
11    ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
12};
13use crate::builtins::common::{gpu_helpers, tensor};
14#[cfg(feature = "doc_export")]
15use crate::register_builtin_doc_text;
16use crate::{register_builtin_fusion_spec, register_builtin_gpu_spec};
17
18const EPS: f64 = 1.0e-12;
19
20#[cfg(feature = "doc_export")]
21pub const DOC_MD: &str = r#"---
22title: "polyval"
23category: "math/poly"
24keywords: ["polyval", "polynomial evaluation", "prediction interval", "polyfit", "gpu"]
25summary: "Evaluate a polynomial at given points and optionally compute prediction intervals using polyfit output."
26references:
27  - title: "MATLAB polyval documentation"
28    url: "https://www.mathworks.com/help/matlab/ref/polyval.html"
29  - title: "Numerical Recipes: Polynomial fitting and evaluation"
30    url: "https://numerical.recipes/"
31gpu_support:
32  elementwise: false
33  reduction: false
34  precisions: ["f32", "f64"]
35  broadcasting: "matlab"
36  notes: "Providers with Horner kernels keep real-valued workloads on the device; complex inputs and prediction intervals fall back to the CPU implementation."
37fusion:
38  elementwise: false
39  reduction: false
40  max_inputs: 2
41  constants: "uniform"
42requires_feature: null
43tested:
44  unit: "builtins::math::poly::polyval::tests"
45  integration: "builtins::math::poly::polyval::tests::polyval_gpu_roundtrip"
46---
47
48# What does the `polyval` function do in MATLAB / RunMat?
49`polyval(p, x)` evaluates the polynomial defined by the coefficient vector `p` at every element of
50`x`. Coefficients follow MATLAB’s convention: `p(1)` is the highest-order term, `p(end)` is the
51constant term. Both real and complex inputs are supported, and the output matches the shape of `x`.
52
53## How does the `polyval` function behave in MATLAB / RunMat?
54- Accepts scalar, vector, or N-D coefficient inputs that have at most one non-singleton dimension.
55- Logical and integer coefficients are promoted to double precision; complex coefficients are kept
56  exactly.
57- When `p` and `x` are real-valued and a provider is registered, RunMat issues a Horner-series GPU
58  kernel via RunMat Accelerate. Mixed or complex inputs fall back to the reference CPU
59  implementation, with purely real outputs re-uploaded to the device when residency makes sense.
60- `polyval(p, x, [], mu)` applies centering and scaling parameters from `polyfit`, evaluating the
61  polynomial at `(x - mu(1)) / mu(2)`.
62- `[y, delta] = polyval(p, x, S, mu)` computes prediction intervals using the structure `S`
63  produced by `polyfit`. RunMat mirrors MATLAB rules: `S` must contain the fields `R`, `normr`,
64  and `df`, and the interval collapses to zeros when `df <= 0` or `normr == 0`.
65- Empty inputs yield empty outputs; an empty coefficient vector represents the zero polynomial.
66- MATLAB-compatible errors are raised for non-numeric inputs, invalid `mu` vectors, or malformed
67  `S` structures.
68
69## `polyval` Function GPU Execution Behaviour
70When a GPU provider is active, RunMat first attempts to evaluate the polynomial in device memory
71using a dedicated Horner kernel. Coefficients and inputs are uploaded automatically when required.
72If the call requests complex arithmetic, prediction intervals, or otherwise falls outside the GPU
73kernel’s contract, RunMat gathers to the host, executes the CPU implementation, and (for real-valued
74results) pushes the output back to the GPU so downstream kernels retain residency.
75
76## Examples of using the `polyval` function in MATLAB / RunMat
77
78### Evaluating a polynomial at scalar points
79
80```matlab
81p = [2 -3 5];   % 2x^2 - 3x + 5
82y = polyval(p, 4);
83```
84
85Expected output:
86
87```matlab
88y = 21;
89```
90
91### Evaluating across a vector of inputs
92
93```matlab
94p = [1 0 -2 1];
95x = linspace(-2, 2, 5);
96y = polyval(p, x);
97```
98
99Expected output:
100
101```matlab
102y = [ -3   2   1   0   5 ];
103```
104
105### Evaluating a polynomial over a matrix grid
106
107```matlab
108[X, Y] = meshgrid(-1:1);
109Z = polyval([1 -3 2], X + Y);
110```
111
112Expected output (matching the shape of `X` and `Y`):
113
114```matlab
115Z =
116    12     6     2
117     6     2     0
118     2     0     0
119```
120
121### Using centering and scaling parameters from `polyfit`
122
123```matlab
124x = -2:2;
125noise = 0.05 * randn(size(x));
126[p, S, mu] = polyfit(x, sin(x) + noise, 3);
127y = polyval(p, x, [], mu);
128```
129
130Expected output:
131
132```matlab
133% y closely matches sin(x) + noise with polynomial smoothing
134```
135
136### Computing prediction intervals with polyfit output
137
138```matlab
139[p, S, mu] = polyfit((0:10)', exp((0:10)'/10), 3);
140[y, delta] = polyval(p, 5, S, mu);
141```
142
143Expected output:
144
145```matlab
146% y is the fitted value at x = 5
147% delta is the 1σ prediction interval (standard error)
148```
149
150### Handling complex coefficients and inputs
151
152```matlab
153p = [1+2i, -3, 4i];
154z = [-1+1i, 0, 1-2i];
155y = polyval(p, z);
156```
157
158Expected output:
159
160```matlab
161% Complex results that agree with MATLAB's polyval
162```
163
164### Evaluating on a gpuArray input
165
166```matlab
167x = gpuArray.linspace(-1, 1, 2048);
168p = [1 0 1];
169y = polyval(p, x);    % Runs on the GPU for real-valued inputs
170```
171
172Expected behaviour:
173
174```matlab
175y is a gpuArray because the result is real-valued.
176```
177
178## FAQ
179
180### Do coefficients need to be a row vector?
181No. They can be row or column vectors (or any N-D shape with a single non-singleton dimension). The
182output always matches the shape of `x`.
183
184### What kinds of inputs are accepted?
185Numeric scalars, vectors, matrices, N-D arrays, logical arrays, and complex data are all accepted.
186Logical and integer inputs are promoted to double precision automatically.
187
188### How do centering (`mu`) parameters work?
189RunMat mirrors MATLAB: the polynomial is evaluated at `(x - mu(1)) / mu(2)`. The `mu` vector must
190contain at least two finite values, and the scale `mu(2)` must be non-zero.
191
192### Why does `[y, delta] = polyval(...)` require the structure `S`?
193The prediction interval comes from the QR factorization stored in `S` by `polyfit`. The structure
194must include the fields `R`, `df`, and `normr`; without them the interval cannot be computed.
195
196### What happens when `df <= 0` or `normr == 0`?
197The prediction interval collapses to zeros (RunMat matches MATLAB and Octave here). This typically
198occurs when the fit is exact or when there are insufficient degrees of freedom.
199
200### Can I keep results on the GPU?
201Yes. RunMat re-uploads real-valued results to the active provider after the host evaluation. Complex
202outputs stay on the host until providers add complex buffer uploads.
203
204### Does `polyval` support sparse inputs?
205Not yet. Dense inputs (including gpuArray tensors) are supported today. Sparse support will arrive
206once RunMat's sparse infrastructure stabilises.
207
208## See Also
209[polyfit](./polyfit), [conv](../signal/conv), [deconv](../signal/deconv), [poly](./poly), [polyder](./polyder), [gpuArray](../../acceleration/gpu/gpuArray), [gather](../../acceleration/gpu/gather)
210
211## Source & Feedback
212- Source: `crates/runmat-runtime/src/builtins/math/poly/polyval.rs`
213- Found an issue or behavioural difference? [Open an issue](https://github.com/runmat-org/runmat/issues/new/choose) with a minimal reproduction.
214"#;
215
216pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
217    name: "polyval",
218    op_kind: GpuOpKind::Custom("polyval"),
219    supported_precisions: &[ScalarType::F32, ScalarType::F64],
220    broadcast: BroadcastSemantics::Matlab,
221    provider_hooks: &[ProviderHook::Custom("polyval")],
222    constant_strategy: ConstantStrategy::UniformBuffer,
223    residency: ResidencyPolicy::NewHandle,
224    nan_mode: ReductionNaN::Include,
225    two_pass_threshold: None,
226    workgroup_size: None,
227    accepts_nan_mode: false,
228    notes:
229        "Uses provider-level Horner kernels for real coefficients/inputs; falls back to host evaluation (with upload) for complex or prediction-interval paths.",
230};
231
232register_builtin_gpu_spec!(GPU_SPEC);
233
234pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
235    name: "polyval",
236    shape: ShapeRequirements::Any,
237    constant_strategy: ConstantStrategy::UniformBuffer,
238    elementwise: None,
239    reduction: None,
240    emits_nan: true,
241    notes: "Acts as a fusion sink; real-valued workloads stay on device, while complex/delta paths gather to the host.",
242};
243
244register_builtin_fusion_spec!(FUSION_SPEC);
245
246#[cfg(feature = "doc_export")]
247register_builtin_doc_text!("polyval", DOC_MD);
248
249#[runtime_builtin(
250    name = "polyval",
251    category = "math/poly",
252    summary = "Evaluate a polynomial at given points with MATLAB-compatible options.",
253    keywords = "polyval,polynomial,polyfit,delta,gpu",
254    accel = "sink",
255    sink = true
256)]
257fn polyval_builtin(p: Value, x: Value, rest: Vec<Value>) -> Result<Value, String> {
258    let eval = evaluate(p, x, &rest, false)?;
259    Ok(eval.value())
260}
261
262/// Evaluate `polyval`, optionally computing the prediction interval.
263pub fn evaluate(
264    coefficients: Value,
265    points: Value,
266    rest: &[Value],
267    want_delta: bool,
268) -> Result<PolyvalEval, String> {
269    let options = parse_option_values(rest)?;
270
271    let coeff_clone = coefficients.clone();
272    let points_clone = points.clone();
273
274    let coeff_was_gpu = matches!(coefficients, Value::GpuTensor(_));
275    let (coeffs, coeff_real) = convert_coefficients(coeff_clone)?;
276
277    let (mut inputs, prefer_gpu_points) = convert_points(points_clone)?;
278    let prefer_gpu_output = prefer_gpu_points || coeff_was_gpu;
279
280    let mu = match options.mu.clone() {
281        Some(mu_value) => Some(parse_mu(mu_value)?),
282        None => None,
283    };
284
285    if prefer_gpu_output && !want_delta && options.s.is_none() {
286        if let Some(value) = try_gpu_polyval(&coeffs, coeff_real, &inputs, mu, prefer_gpu_output)? {
287            return Ok(PolyvalEval::new(value, None));
288        }
289    }
290
291    if let Some(mu_val) = mu {
292        apply_mu(&mut inputs.data, mu_val)?;
293    }
294
295    let stats = if let Some(s_value) = options.s {
296        parse_stats(s_value, coeffs.len())?
297    } else {
298        None
299    };
300
301    if want_delta && stats.is_none() {
302        return Err(
303            "polyval: S input (structure returned by polyfit) is required for delta output"
304                .to_string(),
305        );
306    }
307
308    if inputs.data.is_empty() {
309        let y = zeros_like(&inputs.shape, prefer_gpu_output)?;
310        let delta = if want_delta {
311            Some(zeros_like(&inputs.shape, prefer_gpu_output)?)
312        } else {
313            None
314        };
315        return Ok(PolyvalEval::new(y, delta));
316    }
317
318    if coeffs.is_empty() {
319        let zeros = zeros_like(&inputs.shape, prefer_gpu_output)?;
320        let delta = if want_delta {
321            Some(zeros_like(&inputs.shape, prefer_gpu_output)?)
322        } else {
323            None
324        };
325        return Ok(PolyvalEval::new(zeros, delta));
326    }
327
328    let output_real = coeff_real && inputs.all_real;
329    let values = evaluate_polynomial(&coeffs, &inputs.data);
330    let result_value = finalize_values(
331        &values,
332        &inputs.shape,
333        prefer_gpu_output,
334        output_real && values_are_real(&values),
335    )?;
336
337    let delta_value = if want_delta {
338        let stats = stats.expect("delta requires stats");
339        let delta = compute_prediction_interval(&coeffs, &inputs.data, &stats)?;
340        let prefer = prefer_gpu_output && stats.is_real;
341        Some(finalize_delta(delta, &inputs.shape, prefer)?)
342    } else {
343        None
344    };
345
346    Ok(PolyvalEval::new(result_value, delta_value))
347}
348
349fn try_gpu_polyval(
350    coeffs: &[Complex64],
351    coeff_real: bool,
352    inputs: &NumericArray,
353    mu: Option<Mu>,
354    prefer_gpu_output: bool,
355) -> Result<Option<Value>, String> {
356    if !coeff_real || !inputs.all_real {
357        return Ok(None);
358    }
359    if coeffs.is_empty() || inputs.data.is_empty() {
360        return Ok(None);
361    }
362    let Some(provider) = runmat_accelerate_api::provider() else {
363        return Ok(None);
364    };
365
366    let coeff_data: Vec<f64> = coeffs.iter().map(|c| c.re).collect();
367    let coeff_shape = vec![1usize, coeffs.len()];
368    let coeff_view = HostTensorView {
369        data: &coeff_data,
370        shape: &coeff_shape,
371    };
372    let coeff_handle = match provider.upload(&coeff_view) {
373        Ok(handle) => handle,
374        Err(err) => {
375            debug!("polyval: GPU upload of coefficients failed, falling back: {err}");
376            return Ok(None);
377        }
378    };
379
380    let input_data: Vec<f64> = inputs.data.iter().map(|c| c.re).collect();
381    let input_shape = inputs.shape.clone();
382    let input_view = HostTensorView {
383        data: &input_data,
384        shape: &input_shape,
385    };
386    let input_handle = match provider.upload(&input_view) {
387        Ok(handle) => handle,
388        Err(err) => {
389            debug!("polyval: GPU upload of evaluation points failed, falling back: {err}");
390            let _ = provider.free(&coeff_handle);
391            return Ok(None);
392        }
393    };
394
395    let options = ProviderPolyvalOptions {
396        mu: mu.map(|m| ProviderPolyvalMu {
397            mean: m.mean,
398            scale: m.scale,
399        }),
400    };
401
402    let result_handle = match provider.polyval(&coeff_handle, &input_handle, &options) {
403        Ok(handle) => handle,
404        Err(err) => {
405            debug!("polyval: GPU kernel execution failed, falling back: {err}");
406            let _ = provider.free(&coeff_handle);
407            let _ = provider.free(&input_handle);
408            return Ok(None);
409        }
410    };
411
412    let _ = provider.free(&coeff_handle);
413    let _ = provider.free(&input_handle);
414
415    if prefer_gpu_output {
416        return Ok(Some(Value::GpuTensor(result_handle)));
417    }
418
419    let host = match provider.download(&result_handle) {
420        Ok(host) => host,
421        Err(err) => {
422            debug!("polyval: GPU download failed, falling back: {err}");
423            let _ = provider.free(&result_handle);
424            return Ok(None);
425        }
426    };
427    let _ = provider.free(&result_handle);
428
429    let tensor = Tensor::new(host.data, host.shape).map_err(|e| format!("polyval: {e}"))?;
430    Ok(Some(tensor::tensor_into_value(tensor)))
431}
432
433/// Result object for polyval evaluation.
434#[derive(Debug)]
435pub struct PolyvalEval {
436    value: Value,
437    delta: Option<Value>,
438}
439
440impl PolyvalEval {
441    fn new(value: Value, delta: Option<Value>) -> Self {
442        Self { value, delta }
443    }
444
445    /// Primary output (`y`).
446    pub fn value(&self) -> Value {
447        self.value.clone()
448    }
449
450    /// Optional prediction interval (`delta`).
451    pub fn delta(&self) -> Result<Value, String> {
452        self.delta
453            .clone()
454            .ok_or_else(|| "polyval: delta output not computed".to_string())
455    }
456
457    /// Consume into the main value.
458    pub fn into_value(self) -> Value {
459        self.value
460    }
461
462    /// Consume into `(value, delta)` pair.
463    pub fn into_pair(self) -> Result<(Value, Value), String> {
464        match self.delta {
465            Some(delta) => Ok((self.value, delta)),
466            None => Err("polyval: delta output not computed".to_string()),
467        }
468    }
469}
470
471#[derive(Clone, Copy)]
472struct Mu {
473    mean: f64,
474    scale: f64,
475}
476
477impl Mu {
478    fn new(mean: f64, scale: f64) -> Result<Self, String> {
479        if !mean.is_finite() || !scale.is_finite() {
480            return Err("polyval: mu values must be finite".to_string());
481        }
482        if scale.abs() <= EPS {
483            return Err("polyval: mu(2) must be non-zero".to_string());
484        }
485        Ok(Self { mean, scale })
486    }
487}
488
489#[derive(Clone)]
490struct NumericArray {
491    data: Vec<Complex64>,
492    shape: Vec<usize>,
493    all_real: bool,
494}
495
496#[derive(Clone)]
497struct PolyfitStats {
498    r: Matrix,
499    df: f64,
500    normr: f64,
501    is_real: bool,
502}
503
504impl PolyfitStats {
505    fn is_effective(&self) -> bool {
506        self.r.len() > 0 && self.df > 0.0 && self.normr.is_finite()
507    }
508}
509
510#[derive(Clone)]
511struct Matrix {
512    rows: usize,
513    cols: usize,
514    data: Vec<Complex64>,
515}
516
517impl Matrix {
518    fn get(&self, row: usize, col: usize) -> Complex64 {
519        self.data[row + col * self.rows]
520    }
521
522    fn len(&self) -> usize {
523        self.rows * self.cols
524    }
525}
526
527struct ParsedOptions {
528    s: Option<Value>,
529    mu: Option<Value>,
530}
531
532fn parse_option_values(rest: &[Value]) -> Result<ParsedOptions, String> {
533    match rest.len() {
534        0 => Ok(ParsedOptions { s: None, mu: None }),
535        1 => Ok(ParsedOptions {
536            s: if is_empty_value(&rest[0]) {
537                None
538            } else {
539                Some(rest[0].clone())
540            },
541            mu: None,
542        }),
543        2 => Ok(ParsedOptions {
544            s: if is_empty_value(&rest[0]) {
545                None
546            } else {
547                Some(rest[0].clone())
548            },
549            mu: Some(rest[1].clone()),
550        }),
551        _ => Err("polyval: too many input arguments".to_string()),
552    }
553}
554
555fn convert_coefficients(value: Value) -> Result<(Vec<Complex64>, bool), String> {
556    match value {
557        Value::GpuTensor(handle) => {
558            let gathered = gpu_helpers::gather_value(&Value::GpuTensor(handle.clone()))?;
559            convert_coefficients(gathered)
560        }
561        Value::Tensor(mut tensor) => {
562            ensure_vector_shape("polyval", &tensor.shape)?;
563            let data = tensor
564                .data
565                .drain(..)
566                .map(|re| Complex64::new(re, 0.0))
567                .collect();
568            Ok((data, true))
569        }
570        Value::ComplexTensor(mut tensor) => {
571            ensure_vector_shape("polyval", &tensor.shape)?;
572            let all_real = tensor.data.iter().all(|&(_, im)| im.abs() <= EPS);
573            let data = tensor
574                .data
575                .drain(..)
576                .map(|(re, im)| Complex64::new(re, im))
577                .collect();
578            Ok((data, all_real))
579        }
580        Value::LogicalArray(mut array) => {
581            ensure_vector_data_shape("polyval", &array.shape)?;
582            let data = array
583                .data
584                .drain(..)
585                .map(|bit| Complex64::new(if bit != 0 { 1.0 } else { 0.0 }, 0.0))
586                .collect();
587            Ok((data, true))
588        }
589        Value::Num(n) => Ok((vec![Complex64::new(n, 0.0)], true)),
590        Value::Int(i) => Ok((vec![Complex64::new(i.to_f64(), 0.0)], true)),
591        Value::Bool(flag) => Ok((
592            vec![Complex64::new(if flag { 1.0 } else { 0.0 }, 0.0)],
593            true,
594        )),
595        Value::Complex(re, im) => Ok((vec![Complex64::new(re, im)], im.abs() <= EPS)),
596        other => Err(format!(
597            "polyval: coefficients must be numeric, got {other:?}"
598        )),
599    }
600}
601
602fn convert_points(value: Value) -> Result<(NumericArray, bool), String> {
603    match value {
604        Value::GpuTensor(handle) => {
605            let tensor = gpu_helpers::gather_tensor(&handle)?;
606            let array = NumericArray {
607                data: tensor
608                    .data
609                    .iter()
610                    .map(|&re| Complex64::new(re, 0.0))
611                    .collect(),
612                shape: tensor.shape.clone(),
613                all_real: true,
614            };
615            Ok((array, true))
616        }
617        Value::Tensor(tensor) => Ok((
618            NumericArray {
619                data: tensor
620                    .data
621                    .iter()
622                    .map(|&re| Complex64::new(re, 0.0))
623                    .collect(),
624                shape: tensor.shape.clone(),
625                all_real: true,
626            },
627            false,
628        )),
629        Value::ComplexTensor(tensor) => Ok((
630            NumericArray {
631                data: tensor
632                    .data
633                    .iter()
634                    .map(|&(re, im)| Complex64::new(re, im))
635                    .collect(),
636                shape: tensor.shape.clone(),
637                all_real: tensor.data.iter().all(|&(_, im)| im.abs() <= EPS),
638            },
639            false,
640        )),
641        Value::LogicalArray(array) => Ok((
642            NumericArray {
643                data: array
644                    .data
645                    .iter()
646                    .map(|&bit| Complex64::new(if bit != 0 { 1.0 } else { 0.0 }, 0.0))
647                    .collect(),
648                shape: array.shape.clone(),
649                all_real: true,
650            },
651            false,
652        )),
653        Value::Num(n) => Ok((
654            NumericArray {
655                data: vec![Complex64::new(n, 0.0)],
656                shape: vec![1, 1],
657                all_real: true,
658            },
659            false,
660        )),
661        Value::Int(i) => Ok((
662            NumericArray {
663                data: vec![Complex64::new(i.to_f64(), 0.0)],
664                shape: vec![1, 1],
665                all_real: true,
666            },
667            false,
668        )),
669        Value::Bool(flag) => Ok((
670            NumericArray {
671                data: vec![Complex64::new(if flag { 1.0 } else { 0.0 }, 0.0)],
672                shape: vec![1, 1],
673                all_real: true,
674            },
675            false,
676        )),
677        Value::Complex(re, im) => Ok((
678            NumericArray {
679                data: vec![Complex64::new(re, im)],
680                shape: vec![1, 1],
681                all_real: im.abs() <= EPS,
682            },
683            false,
684        )),
685        other => Err(format!("polyval: X must be numeric, got {other:?}")),
686    }
687}
688
689fn parse_mu(value: Value) -> Result<Mu, String> {
690    match value {
691        Value::GpuTensor(handle) => {
692            let gathered = gpu_helpers::gather_tensor(&handle)?;
693            parse_mu(Value::Tensor(gathered))
694        }
695        Value::Tensor(tensor) => {
696            if tensor.data.len() < 2 {
697                return Err("polyval: mu must contain at least two elements".to_string());
698            }
699            Mu::new(tensor.data[0], tensor.data[1])
700        }
701        Value::LogicalArray(array) => {
702            if array.data.len() < 2 {
703                return Err("polyval: mu must contain at least two elements".to_string());
704            }
705            let mean = if array.data[0] != 0 { 1.0 } else { 0.0 };
706            let scale = if array.data[1] != 0 { 1.0 } else { 0.0 };
707            Mu::new(mean, scale)
708        }
709        Value::Num(_) | Value::Int(_) | Value::Bool(_) | Value::Complex(_, _) => {
710            Err("polyval: mu must be a numeric vector with at least two values".to_string())
711        }
712        Value::ComplexTensor(tensor) => {
713            if tensor.data.len() < 2 {
714                return Err("polyval: mu must contain at least two elements".to_string());
715            }
716            let (mean_re, mean_im) = tensor.data[0];
717            let (scale_re, scale_im) = tensor.data[1];
718            if mean_im.abs() > EPS || scale_im.abs() > EPS {
719                return Err("polyval: mu values must be real".to_string());
720            }
721            Mu::new(mean_re, scale_re)
722        }
723        _ => Err("polyval: mu must be a numeric vector with at least two values".to_string()),
724    }
725}
726
727fn parse_stats(value: Value, coeff_len: usize) -> Result<Option<PolyfitStats>, String> {
728    if is_empty_value(&value) {
729        return Ok(None);
730    }
731    let struct_value = match value {
732        Value::Struct(s) => s,
733        Value::GpuTensor(handle) => {
734            let gathered = gpu_helpers::gather_value(&Value::GpuTensor(handle))?;
735            return parse_stats(gathered, coeff_len);
736        }
737        other => {
738            return Err(format!(
739                "polyval: S input must be the structure returned by polyfit, got {other:?}"
740            ))
741        }
742    };
743    let r_value = struct_value
744        .fields
745        .get("R")
746        .cloned()
747        .ok_or_else(|| "polyval: S input is missing the field 'R'".to_string())?;
748    let df_value = struct_value
749        .fields
750        .get("df")
751        .cloned()
752        .ok_or_else(|| "polyval: S input is missing the field 'df'".to_string())?;
753    let normr_value = struct_value
754        .fields
755        .get("normr")
756        .cloned()
757        .ok_or_else(|| "polyval: S input is missing the field 'normr'".to_string())?;
758
759    let (matrix, is_real) = convert_matrix(r_value, coeff_len)?;
760    let df = scalar_to_f64(df_value, "polyval: S.df")?;
761    let normr = scalar_to_f64(normr_value, "polyval: S.normr")?;
762
763    Ok(Some(PolyfitStats {
764        r: matrix,
765        df,
766        normr,
767        is_real,
768    }))
769}
770
771fn convert_matrix(value: Value, coeff_len: usize) -> Result<(Matrix, bool), String> {
772    match value {
773        Value::GpuTensor(handle) => {
774            let tensor = gpu_helpers::gather_tensor(&handle)?;
775            convert_matrix(Value::Tensor(tensor), coeff_len)
776        }
777        Value::Tensor(tensor) => {
778            let Tensor {
779                data, rows, cols, ..
780            } = tensor;
781            if rows != coeff_len || cols != coeff_len {
782                return Err("polyval: size of S.R must match the coefficient vector".to_string());
783            }
784            let data = data.into_iter().map(|re| Complex64::new(re, 0.0)).collect();
785            Ok((Matrix { rows, cols, data }, true))
786        }
787        Value::ComplexTensor(tensor) => {
788            let ComplexTensor {
789                data, rows, cols, ..
790            } = tensor;
791            if rows != coeff_len || cols != coeff_len {
792                return Err("polyval: size of S.R must match the coefficient vector".to_string());
793            }
794            let imag_small = data.iter().all(|&(_, im)| im.abs() <= EPS);
795            let data = data
796                .into_iter()
797                .map(|(re, im)| Complex64::new(re, im))
798                .collect();
799            Ok((Matrix { rows, cols, data }, imag_small))
800        }
801        Value::LogicalArray(array) => {
802            let LogicalArray { data, shape } = array;
803            let rows = shape.first().copied().unwrap_or(0);
804            let cols = shape.get(1).copied().unwrap_or(0);
805            if rows != coeff_len || cols != coeff_len {
806                return Err("polyval: size of S.R must match the coefficient vector".to_string());
807            }
808            let data = data
809                .into_iter()
810                .map(|bit| Complex64::new(if bit != 0 { 1.0 } else { 0.0 }, 0.0))
811                .collect();
812            Ok((Matrix { rows, cols, data }, true))
813        }
814        Value::Num(_) | Value::Int(_) | Value::Bool(_) | Value::Complex(_, _) => Err(
815            "polyval: S.R must be a square numeric matrix matching the coefficient vector length"
816                .to_string(),
817        ),
818        Value::Struct(_)
819        | Value::Cell(_)
820        | Value::String(_)
821        | Value::StringArray(_)
822        | Value::CharArray(_) => Err(
823            "polyval: S.R must be a square numeric matrix matching the coefficient vector length"
824                .to_string(),
825        ),
826        _ => Err(
827            "polyval: S.R must be a square numeric matrix matching the coefficient vector length"
828                .to_string(),
829        ),
830    }
831}
832
833fn scalar_to_f64(value: Value, context: &str) -> Result<f64, String> {
834    match value {
835        Value::Num(n) => Ok(n),
836        Value::Int(i) => Ok(i.to_f64()),
837        Value::Bool(flag) => Ok(if flag { 1.0 } else { 0.0 }),
838        Value::Tensor(tensor) => {
839            if tensor.data.len() != 1 {
840                return Err(format!("{context} must be a scalar"));
841            }
842            Ok(tensor.data[0])
843        }
844        Value::LogicalArray(array) => {
845            if array.data.len() != 1 {
846                return Err(format!("{context} must be a scalar"));
847            }
848            Ok(if array.data[0] != 0 { 1.0 } else { 0.0 })
849        }
850        Value::GpuTensor(handle) => {
851            let tensor = gpu_helpers::gather_tensor(&handle)?;
852            scalar_to_f64(Value::Tensor(tensor), context)
853        }
854        Value::Complex(_, _) | Value::ComplexTensor(_) => {
855            Err(format!("{context} must be real-valued"))
856        }
857        other => Err(format!("{context} must be a scalar, got {other:?}")),
858    }
859}
860
861fn apply_mu(values: &mut [Complex64], mu: Mu) -> Result<(), String> {
862    let mean = Complex64::new(mu.mean, 0.0);
863    let scale = Complex64::new(mu.scale, 0.0);
864    for v in values.iter_mut() {
865        *v = (*v - mean) / scale;
866    }
867    Ok(())
868}
869
870fn evaluate_polynomial(coeffs: &[Complex64], inputs: &[Complex64]) -> Vec<Complex64> {
871    let mut outputs = Vec::with_capacity(inputs.len());
872    for &x in inputs {
873        let mut acc = Complex64::new(0.0, 0.0);
874        for &c in coeffs {
875            acc = acc * x + c;
876        }
877        outputs.push(acc);
878    }
879    outputs
880}
881
882fn compute_prediction_interval(
883    coeffs: &[Complex64],
884    inputs: &[Complex64],
885    stats: &PolyfitStats,
886) -> Result<Vec<f64>, String> {
887    if !stats.is_effective() {
888        return Ok(vec![0.0; inputs.len()]);
889    }
890    let n = coeffs.len();
891    let mut delta = Vec::with_capacity(inputs.len());
892    for &x in inputs {
893        let row = vandermonde_row(x, n);
894        let solved = solve_row_against_upper(&row, &stats.r)?;
895        let sum_sq: f64 = solved.iter().map(|c| c.norm_sqr()).sum();
896        let interval = (1.0 + sum_sq).sqrt() * (stats.normr / stats.df.sqrt());
897        delta.push(interval);
898    }
899    Ok(delta)
900}
901
902fn vandermonde_row(x: Complex64, len: usize) -> Vec<Complex64> {
903    if len == 0 {
904        return vec![Complex64::new(1.0, 0.0)];
905    }
906    let degree = len - 1;
907    let mut powers = vec![Complex64::new(1.0, 0.0); degree + 1];
908    for idx in 1..=degree {
909        powers[idx] = powers[idx - 1] * x;
910    }
911    let mut row = vec![Complex64::new(0.0, 0.0); degree + 1];
912    for (i, value) in powers.into_iter().enumerate() {
913        row[degree - i] = value;
914    }
915    row
916}
917
918fn solve_row_against_upper(row: &[Complex64], matrix: &Matrix) -> Result<Vec<Complex64>, String> {
919    let n = row.len();
920    if matrix.rows != n || matrix.cols != n {
921        return Err("polyval: size of S.R must match the coefficient vector".to_string());
922    }
923    let mut result = vec![Complex64::new(0.0, 0.0); n];
924    for j in (0..n).rev() {
925        let mut acc = row[j];
926        for (k, value) in result.iter().enumerate().skip(j + 1) {
927            acc -= *value * matrix.get(k, j);
928        }
929        let diag = matrix.get(j, j);
930        if diag.norm() <= EPS {
931            return Err("polyval: S.R is singular".to_string());
932        }
933        result[j] = acc / diag;
934    }
935    Ok(result)
936}
937
938fn finalize_values(
939    data: &[Complex64],
940    shape: &[usize],
941    prefer_gpu: bool,
942    real_only: bool,
943) -> Result<Value, String> {
944    if real_only {
945        let real_data: Vec<f64> = data.iter().map(|c| c.re).collect();
946        finalize_real(real_data, shape, prefer_gpu)
947    } else if data.len() == 1 {
948        let value = data[0];
949        Ok(Value::Complex(value.re, value.im))
950    } else {
951        let complex_data: Vec<(f64, f64)> = data.iter().map(|c| (c.re, c.im)).collect();
952        let tensor = ComplexTensor::new(complex_data, shape.to_vec())
953            .map_err(|e| format!("polyval: failed to build complex tensor: {e}"))?;
954        Ok(Value::ComplexTensor(tensor))
955    }
956}
957
958fn finalize_delta(data: Vec<f64>, shape: &[usize], prefer_gpu: bool) -> Result<Value, String> {
959    finalize_real(data, shape, prefer_gpu)
960}
961
962fn finalize_real(data: Vec<f64>, shape: &[usize], prefer_gpu: bool) -> Result<Value, String> {
963    let tensor = Tensor::new(data, shape.to_vec())
964        .map_err(|e| format!("polyval: failed to build tensor: {e}"))?;
965    if prefer_gpu {
966        if let Some(provider) = runmat_accelerate_api::provider() {
967            let view = HostTensorView {
968                data: &tensor.data,
969                shape: &tensor.shape,
970            };
971            if let Ok(handle) = provider.upload(&view) {
972                return Ok(Value::GpuTensor(handle));
973            }
974        }
975    }
976    Ok(tensor::tensor_into_value(tensor))
977}
978
979fn zeros_like(shape: &[usize], prefer_gpu: bool) -> Result<Value, String> {
980    let len = shape.iter().product();
981    finalize_real(vec![0.0; len], shape, prefer_gpu)
982}
983
984fn ensure_vector_shape(name: &str, shape: &[usize]) -> Result<(), String> {
985    if !is_vector_shape(shape) {
986        Err(format!(
987            "{name}: coefficients must be a scalar, row vector, or column vector"
988        ))
989    } else {
990        Ok(())
991    }
992}
993
994fn ensure_vector_data_shape(name: &str, shape: &[usize]) -> Result<(), String> {
995    if !is_vector_shape(shape) {
996        Err(format!("{name}: inputs must be vectors or scalars"))
997    } else {
998        Ok(())
999    }
1000}
1001
1002fn is_vector_shape(shape: &[usize]) -> bool {
1003    shape.iter().filter(|&&dim| dim > 1).count() <= 1
1004}
1005
1006fn is_empty_value(value: &Value) -> bool {
1007    match value {
1008        Value::Tensor(t) => t.data.is_empty(),
1009        Value::LogicalArray(l) => l.data.is_empty(),
1010        Value::Cell(ca) => ca.data.is_empty(),
1011        Value::GpuTensor(handle) => {
1012            match gpu_helpers::gather_value(&Value::GpuTensor(handle.clone())) {
1013                Ok(gathered) => is_empty_value(&gathered),
1014                Err(_) => false,
1015            }
1016        }
1017        _ => false,
1018    }
1019}
1020
1021fn values_are_real(values: &[Complex64]) -> bool {
1022    values.iter().all(|c| c.im.abs() <= EPS)
1023}
1024
1025#[cfg(test)]
1026mod tests {
1027    use super::*;
1028    use crate::builtins::common::test_support;
1029    use runmat_builtins::StructValue;
1030
1031    #[test]
1032    fn polyval_scalar() {
1033        let coeffs = Tensor::new(vec![2.0, -3.0, 5.0], vec![1, 3]).unwrap();
1034        let value =
1035            polyval_builtin(Value::Tensor(coeffs), Value::Num(4.0), Vec::new()).expect("polyval");
1036        match value {
1037            Value::Num(n) => assert!((n - (2.0 * 16.0 - 12.0 + 5.0)).abs() < 1e-12),
1038            other => panic!("expected scalar, got {other:?}"),
1039        }
1040    }
1041
1042    #[test]
1043    fn polyval_matrix_input() {
1044        let coeffs = Tensor::new(vec![1.0, 0.0, -2.0, 1.0], vec![1, 4]).unwrap();
1045        let points = Tensor::new(vec![-2.0, -1.0, 0.0, 1.0, 2.0], vec![5, 1]).unwrap();
1046        let value = polyval_builtin(
1047            Value::Tensor(coeffs),
1048            Value::Tensor(points.clone()),
1049            Vec::new(),
1050        )
1051        .expect("polyval");
1052        match value {
1053            Value::Tensor(tensor) => {
1054                assert_eq!(tensor.shape, points.shape);
1055                let expected = vec![-3.0, 2.0, 1.0, 0.0, 5.0];
1056                assert_eq!(tensor.data, expected);
1057            }
1058            other => panic!("expected tensor output, got {other:?}"),
1059        }
1060    }
1061
1062    #[test]
1063    fn polyval_complex_inputs() {
1064        let coeffs =
1065            ComplexTensor::new(vec![(1.0, 2.0), (-3.0, 0.0), (0.0, 4.0)], vec![1, 3]).unwrap();
1066        let points =
1067            ComplexTensor::new(vec![(-1.0, 1.0), (0.0, 0.0), (1.0, -2.0)], vec![1, 3]).unwrap();
1068        let value = polyval_builtin(
1069            Value::ComplexTensor(coeffs),
1070            Value::ComplexTensor(points.clone()),
1071            Vec::new(),
1072        )
1073        .expect("polyval");
1074        match value {
1075            Value::ComplexTensor(tensor) => {
1076                assert_eq!(tensor.shape, points.shape);
1077                assert_eq!(tensor.data.len(), 3);
1078            }
1079            other => panic!("expected complex tensor, got {other:?}"),
1080        }
1081    }
1082
1083    #[test]
1084    fn polyval_with_mu() {
1085        let coeffs = Tensor::new(vec![1.0, 0.0, 0.0], vec![1, 3]).unwrap();
1086        let points = Tensor::new(vec![0.0, 1.0, 2.0], vec![1, 3]).unwrap();
1087        let mu = Tensor::new(vec![1.0, 2.0], vec![1, 2]).unwrap();
1088        let value = polyval_builtin(
1089            Value::Tensor(coeffs),
1090            Value::Tensor(points),
1091            vec![
1092                Value::Tensor(Tensor::new(vec![], vec![0, 0]).unwrap()),
1093                Value::Tensor(mu),
1094            ],
1095        )
1096        .expect("polyval");
1097        match value {
1098            Value::Tensor(tensor) => {
1099                assert_eq!(tensor.data, vec![0.25, 0.0, 0.25]);
1100            }
1101            other => panic!("expected tensor output, got {other:?}"),
1102        }
1103    }
1104
1105    #[test]
1106    fn polyval_delta_computation() {
1107        let coeffs = Tensor::new(vec![1.0, -3.0, 2.0], vec![1, 3]).unwrap();
1108        let points = Tensor::new(vec![0.0, 1.0, 2.0], vec![1, 3]).unwrap();
1109        let mut st = StructValue::new();
1110        let r = Tensor::new(
1111            vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0],
1112            vec![3, 3],
1113        )
1114        .unwrap();
1115        st.fields.insert("R".to_string(), Value::Tensor(r));
1116        st.fields.insert("df".to_string(), Value::Num(4.0));
1117        st.fields.insert("normr".to_string(), Value::Num(2.0));
1118        let stats = Value::Struct(st);
1119        let eval = evaluate(Value::Tensor(coeffs), Value::Tensor(points), &[stats], true)
1120            .expect("polyval");
1121        let (_, delta) = eval.into_pair().expect("delta available");
1122        match delta {
1123            Value::Tensor(tensor) => {
1124                assert_eq!(tensor.shape, vec![1, 3]);
1125                assert_eq!(tensor.data.len(), 3);
1126            }
1127            other => panic!("expected tensor delta, got {other:?}"),
1128        }
1129    }
1130
1131    #[test]
1132    fn polyval_delta_requires_stats() {
1133        let coeffs = Tensor::new(vec![1.0, 0.0], vec![1, 2]).unwrap();
1134        let points = Tensor::new(vec![1.0], vec![1, 1]).unwrap();
1135        let err = evaluate(Value::Tensor(coeffs), Value::Tensor(points), &[], true)
1136            .expect_err("expected error");
1137        assert!(err.contains("S input"));
1138    }
1139
1140    #[test]
1141    fn polyval_invalid_mu_length_errors() {
1142        let coeffs = Tensor::new(vec![1.0, 0.0], vec![1, 2]).unwrap();
1143        let points = Tensor::new(vec![0.0], vec![1, 1]).unwrap();
1144        let mu = Tensor::new(vec![1.0], vec![1, 1]).unwrap();
1145        let placeholder = Tensor::new(vec![], vec![0, 0]).unwrap();
1146        let err = polyval_builtin(
1147            Value::Tensor(coeffs),
1148            Value::Tensor(points),
1149            vec![Value::Tensor(placeholder), Value::Tensor(mu)],
1150        )
1151        .expect_err("expected mu length error");
1152        assert!(err.contains("mu must contain at least two elements"));
1153    }
1154
1155    #[test]
1156    fn polyval_complex_mu_rejected() {
1157        let coeffs = Tensor::new(vec![1.0, 0.0], vec![1, 2]).unwrap();
1158        let points = Tensor::new(vec![0.0], vec![1, 1]).unwrap();
1159        let complex_mu =
1160            ComplexTensor::new(vec![(0.0, 0.0), (1.0, 0.5)], vec![1, 2]).expect("complex mu");
1161        let placeholder = Tensor::new(vec![], vec![0, 0]).unwrap();
1162        let err = polyval_builtin(
1163            Value::Tensor(coeffs),
1164            Value::Tensor(points),
1165            vec![Value::Tensor(placeholder), Value::ComplexTensor(complex_mu)],
1166        )
1167        .expect_err("expected complex mu error");
1168        assert!(err.contains("mu values must be real"));
1169    }
1170
1171    #[test]
1172    fn polyval_invalid_stats_missing_r() {
1173        let coeffs = Tensor::new(vec![1.0, -3.0, 2.0], vec![1, 3]).unwrap();
1174        let points = Tensor::new(vec![0.0], vec![1, 1]).unwrap();
1175        let mut st = StructValue::new();
1176        st.fields.insert("df".to_string(), Value::Num(1.0));
1177        st.fields.insert("normr".to_string(), Value::Num(1.0));
1178        let stats = Value::Struct(st);
1179        let err = polyval_builtin(Value::Tensor(coeffs), Value::Tensor(points), vec![stats])
1180            .expect_err("expected missing R error");
1181        assert!(err.contains("missing the field 'R'"));
1182    }
1183
1184    #[test]
1185    fn polyval_gpu_roundtrip() {
1186        test_support::with_test_provider(|provider| {
1187            let coeffs = Tensor::new(vec![1.0, 0.0, 1.0], vec![1, 3]).unwrap();
1188            let points = Tensor::new(vec![-1.0, 0.0, 1.0], vec![3, 1]).unwrap();
1189            let coeff_handle = provider
1190                .upload(&HostTensorView {
1191                    data: &coeffs.data,
1192                    shape: &coeffs.shape,
1193                })
1194                .expect("upload coeff");
1195            let point_handle = provider
1196                .upload(&HostTensorView {
1197                    data: &points.data,
1198                    shape: &points.shape,
1199                })
1200                .expect("upload points");
1201            let value = polyval_builtin(
1202                Value::GpuTensor(coeff_handle),
1203                Value::GpuTensor(point_handle),
1204                Vec::new(),
1205            )
1206            .expect("polyval");
1207            match value {
1208                Value::GpuTensor(handle) => {
1209                    let gathered = test_support::gather(Value::GpuTensor(handle)).expect("gather");
1210                    assert_eq!(gathered.data, vec![2.0, 1.0, 2.0]);
1211                }
1212                other => panic!("expected gpu tensor, got {other:?}"),
1213            }
1214        });
1215    }
1216
1217    #[test]
1218    #[cfg(feature = "wgpu")]
1219    fn polyval_wgpu_matches_cpu_real_inputs() {
1220        let _ = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
1221            runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
1222        );
1223        let coeffs = Tensor::new(vec![1.0, -3.0, 2.0], vec![1, 3]).unwrap();
1224        let points = Tensor::new(vec![-2.0, -1.0, 0.5, 2.5], vec![4, 1]).unwrap();
1225
1226        let provider = runmat_accelerate_api::provider().expect("wgpu provider");
1227        let coeff_handle = provider
1228            .upload(&HostTensorView {
1229                data: &coeffs.data,
1230                shape: &coeffs.shape,
1231            })
1232            .expect("upload coeffs");
1233        let point_handle = provider
1234            .upload(&HostTensorView {
1235                data: &points.data,
1236                shape: &points.shape,
1237            })
1238            .expect("upload points");
1239
1240        let gpu_value = polyval_builtin(
1241            Value::GpuTensor(coeff_handle.clone()),
1242            Value::GpuTensor(point_handle.clone()),
1243            Vec::new(),
1244        )
1245        .expect("polyval gpu");
1246
1247        let _ = provider.free(&coeff_handle);
1248        let _ = provider.free(&point_handle);
1249
1250        let gathered = test_support::gather(gpu_value).expect("gather");
1251
1252        let coeff_complex: Vec<Complex64> = coeffs
1253            .data
1254            .iter()
1255            .map(|&c| Complex64::new(c, 0.0))
1256            .collect();
1257        let point_complex: Vec<Complex64> = points
1258            .data
1259            .iter()
1260            .map(|&x| Complex64::new(x, 0.0))
1261            .collect();
1262        let expected_vals = evaluate_polynomial(&coeff_complex, &point_complex);
1263        let expected: Vec<f64> = expected_vals.iter().map(|c| c.re).collect();
1264
1265        assert_eq!(gathered.shape, vec![4, 1]);
1266        assert_eq!(gathered.data, expected);
1267    }
1268
1269    #[test]
1270    #[cfg(feature = "doc_export")]
1271    fn doc_examples_present() {
1272        let blocks = test_support::doc_examples(DOC_MD);
1273        assert!(!blocks.is_empty());
1274    }
1275}