bacon_sci/polynomial/
mod.rs

1use crate::roots::newton_polynomial;
2use nalgebra::{ComplexField, RealField};
3use num_complex::Complex;
4use num_traits::{FromPrimitive, One, Zero};
5use std::collections::VecDeque;
6use std::{any::TypeId, f64, iter::FromIterator, ops};
7
8/// Polynomial on a [`ComplexField`].
9#[derive(Debug, Clone)]
10pub struct Polynomial<N: ComplexField + FromPrimitive + Copy>
11where
12    <N as ComplexField>::RealField: FromPrimitive + Copy,
13{
14    // Index 0 is constant, 1 is linear, etc.
15    coefficients: Vec<N>,
16    tolerance: <N as ComplexField>::RealField,
17}
18
19#[macro_export]
20macro_rules! polynomial {
21  ( $( $x:expr ),* ) => {
22    $crate::polynomial::Polynomial::from_slice(&[$($x),*])
23  }
24}
25
26impl<N: ComplexField + FromPrimitive + Copy> Polynomial<N>
27where
28    <N as ComplexField>::RealField: FromPrimitive + Copy,
29{
30    /// Returns the zero polynomial on a given field
31    pub fn new() -> Self {
32        Polynomial {
33            coefficients: vec![N::from_f64(0.0).unwrap()],
34            tolerance: N::RealField::from_f64(1e-10).unwrap(),
35        }
36    }
37
38    pub fn with_tolerance(tolerance: <N as ComplexField>::RealField) -> Result<Self, String> {
39        if !tolerance.is_sign_positive() {
40            return Err("Polynomial with_tolerance: Tolerance must be positive".to_owned());
41        }
42        Ok(Polynomial {
43            coefficients: vec![N::from_f64(0.0).unwrap()],
44            tolerance,
45        })
46    }
47
48    /// Returns the zero polynomial on a given field with preallocated memory
49    pub fn with_capacity(capacity: usize) -> Self {
50        let mut coefficients = Vec::with_capacity(capacity);
51        coefficients.push(N::zero());
52        coefficients.iter().copied().collect()
53    }
54
55    /// Create a polynomial from a slice, with the first element of the slice being the highest power
56    pub fn from_slice(data: &[N]) -> Self {
57        if data.is_empty() {
58            return Polynomial {
59                coefficients: vec![N::zero()],
60                tolerance: N::RealField::from_f64(1e-10).unwrap(),
61            };
62        }
63        Polynomial {
64            coefficients: data.iter().rev().copied().collect(),
65            tolerance: N::RealField::from_f64(1e-10).unwrap(),
66        }
67    }
68
69    pub fn set_tolerance(
70        &mut self,
71        tolerance: <N as ComplexField>::RealField,
72    ) -> Result<(), String> {
73        if !tolerance.is_sign_positive() {
74            return Err("Polynomial set_tolerance: tolerance must be positive".to_owned());
75        }
76
77        self.tolerance = tolerance;
78        Ok(())
79    }
80
81    pub fn get_tolerance(&self) -> <N as ComplexField>::RealField {
82        self.tolerance
83    }
84
85    /// Get the order of the polynomial
86    pub fn order(&self) -> usize {
87        self.coefficients.len() - 1
88    }
89
90    /// Returns the coefficients in the correct order to recreate the polynomial with [`Polynomial::from_slice`]
91    pub fn get_coefficients(&self) -> Vec<N> {
92        let mut cln = self.coefficients.clone();
93        cln.reverse();
94        cln
95    }
96
97    /// Get the coefficient of a power
98    pub fn get_coefficient(&self, ind: usize) -> N {
99        if ind >= self.coefficients.len() {
100            N::zero()
101        } else {
102            self.coefficients[ind]
103        }
104    }
105
106    /// Make a polynomial complex
107    pub fn make_complex(&self) -> Polynomial<Complex<<N as ComplexField>::RealField>> {
108        let mut coefficients: Vec<Complex<N::RealField>> =
109            Vec::with_capacity(self.coefficients.len());
110        for val in &self.coefficients {
111            coefficients.push(Complex::<N::RealField>::new(val.real(), val.imaginary()));
112        }
113        Polynomial {
114            coefficients,
115            tolerance: self.tolerance,
116        }
117    }
118
119    /// Evaluate a polynomial at a value
120    pub fn evaluate(&self, x: N) -> N {
121        let mut acc = *self.coefficients.last().unwrap();
122        for val in self.coefficients.iter().rev().skip(1) {
123            acc *= x;
124            acc += *val;
125        }
126
127        acc
128    }
129
130    /// Evaluate a polynomial and its derivative at a value
131    pub fn evaluate_derivative(&self, x: N) -> (N, N) {
132        if self.coefficients.len() == 1 {
133            return (self.coefficients[0], N::zero());
134        }
135        // Start with biggest coefficients
136        let mut acc_eval = *self.coefficients.last().unwrap();
137        let mut acc_deriv = *self.coefficients.last().unwrap();
138        // For every coefficient except the constant and largest
139        for val in self.coefficients.iter().skip(1).rev().skip(1) {
140            acc_eval = acc_eval * x + *val;
141            acc_deriv = acc_deriv * x + acc_eval;
142        }
143        // Do the constant for the polynomial evaluation
144        acc_eval = x * acc_eval + self.coefficients[0];
145
146        (acc_eval, acc_deriv)
147    }
148
149    /// Set a coefficient of a power in the polynomial
150    pub fn set_coefficient(&mut self, power: u32, coefficient: N) {
151        while (power + 1) > self.coefficients.len() as u32 {
152            self.coefficients.push(N::from_f64(0.0).unwrap());
153        }
154        self.coefficients[power as usize] = coefficient;
155    }
156
157    /// Remove the coefficient of a power in the polynomial
158    pub fn purge_coefficient(&mut self, power: usize) {
159        match self.coefficients.len() {
160            len if len == power && len != 1 => {
161                self.coefficients.pop();
162            }
163            _ => {
164                self.coefficients[power] = N::from_f64(0.0).unwrap();
165            }
166        };
167    }
168
169    /// Remove all leading 0 coefficients
170    pub fn purge_leading(&mut self) {
171        while self.coefficients.len() > 1
172            && self.coefficients.last().unwrap().real().abs() <= self.tolerance
173            && self.coefficients.last().unwrap().imaginary().abs() <= self.tolerance
174        {
175            self.coefficients.pop();
176        }
177    }
178
179    /// Get the derivative of the polynomial
180    pub fn derivative(&self) -> Self {
181        if self.coefficients.len() == 1 {
182            return Polynomial {
183                coefficients: vec![N::from_f64(0.0).unwrap()],
184                tolerance: self.tolerance,
185            };
186        }
187
188        let mut deriv_coeff = Vec::with_capacity(self.coefficients.len() - 1);
189
190        for (i, val) in self.coefficients.iter().enumerate().skip(1) {
191            deriv_coeff.push(N::from_f64(i as f64).unwrap() * *val);
192        }
193
194        Polynomial {
195            coefficients: deriv_coeff,
196            tolerance: self.tolerance,
197        }
198    }
199
200    /// Get the antiderivative of the polynomial with specified constant
201    pub fn antiderivative(&self, constant: N) -> Self {
202        let mut coefficients = Vec::with_capacity(self.coefficients.len() + 1);
203        coefficients.push(constant);
204        for (ind, val) in self.coefficients.iter().enumerate() {
205            coefficients.push(*val * N::from_f64(1.0 / (ind + 1) as f64).unwrap());
206        }
207        Polynomial {
208            coefficients,
209            tolerance: self.tolerance,
210        }
211    }
212
213    /// Integrate this polynomial between to starting points
214    pub fn integrate(&self, lower: N, upper: N) -> N {
215        let poly_anti = self.antiderivative(N::zero());
216        poly_anti.evaluate(upper) - poly_anti.evaluate(lower)
217    }
218
219    /// Divide this polynomial by another, getting a quotient and remainder, using tol to check for 0
220    pub fn divide(&self, divisor: &Polynomial<N>) -> Result<(Self, Self), String> {
221        if divisor.coefficients.len() == 1
222            && divisor.coefficients[0].real().abs() < self.tolerance
223            && divisor.coefficients[0].imaginary().abs() < self.tolerance
224        {
225            return Err("Polynomial division: Can not divide by 0".to_owned());
226        }
227
228        let mut quotient = Polynomial::with_tolerance(self.tolerance)?;
229        let mut remainder = Polynomial::from_iter(self.coefficients.iter().copied());
230        remainder.tolerance = self.tolerance;
231        remainder.purge_leading();
232        let mut temp = Polynomial::new();
233
234        if divisor.coefficients.len() == 1 {
235            let idivisor = N::from_f64(1.0).unwrap() / divisor.coefficients[0];
236            return Ok((
237                remainder
238                    .coefficients
239                    .iter()
240                    .map(|c| *c * idivisor)
241                    .collect(),
242                Polynomial::new(),
243            ));
244        }
245
246        while remainder.coefficients.len() >= divisor.coefficients.len()
247            && !(remainder.coefficients.len() == 1
248                && remainder.coefficients[0].real().abs() < self.tolerance
249                && remainder.coefficients[0].imaginary().abs() < self.tolerance)
250        {
251            // Get the power left over from dividing lead terms
252            let order = remainder.coefficients.len() - divisor.coefficients.len();
253            // Make a vector that is just the lead coefficients divided at the right power
254            temp.coefficients = vec![N::zero(); order + 1];
255            temp.coefficients[order] =
256                *remainder.coefficients.last().unwrap() / *divisor.coefficients.last().unwrap();
257            // Add the division to the quotient
258            quotient += &temp;
259            // Get the amount to shift divisor by
260            let padding = temp.coefficients.len() - 1;
261            // Multiply every coefficient in divisor by temp's coefficient
262            temp = divisor
263                .coefficients
264                .iter()
265                .map(|c| *c * *temp.coefficients.last().unwrap())
266                .collect();
267            // Shift the coefficients to multiply by the right power of x
268            for _ in 0..padding {
269                temp.coefficients.insert(0, N::zero());
270            }
271            // remainder -= temp x d;
272            remainder -= &temp;
273            while remainder.coefficients.len() > 1
274                && remainder.coefficients.last().unwrap().real().abs() < self.tolerance
275                && remainder.coefficients.last().unwrap().imaginary().abs() < self.tolerance
276            {
277                remainder.coefficients.pop();
278            }
279        }
280
281        Ok((quotient, remainder))
282    }
283
284    /// Get the n (possibly including repeats) of the polynomial given n using Laguerre's method
285    pub fn roots(
286        &self,
287        tol: <N as ComplexField>::RealField,
288        n_max: usize,
289    ) -> Result<VecDeque<Complex<<N as ComplexField>::RealField>>, String> {
290        if self.coefficients.len() > 1
291            && self.coefficients.last().unwrap().real().abs() < tol
292            && self.coefficients.last().unwrap().imaginary().abs() < tol
293        {
294            return Err("Polynomial roots: Leading 0 coefficient!".to_owned());
295        }
296
297        match self.coefficients.len() {
298            1 => {
299                // Only constant, root only if constant is 0
300                if self.coefficients[0].real().abs() < tol
301                    && self.coefficients[0].imaginary().abs() < tol
302                {
303                    return Ok(VecDeque::from(vec![Complex::<N::RealField>::zero()]));
304                }
305                return Err("Polynomial roots: Non-zero constant has no root".to_owned());
306            }
307            2 => {
308                // Linear term, root easy
309                let division = -self.coefficients[0] / self.coefficients[1];
310                return Ok(VecDeque::from(vec![Complex::<N::RealField>::new(
311                    division.real(),
312                    division.imaginary(),
313                )]));
314            }
315            3 => {
316                // Use quadratic formula and return in right order
317                let determinant = self.coefficients[1].powi(2)
318                    - N::from_f64(4.0).unwrap() * self.coefficients[2] * self.coefficients[0];
319                let determinant =
320                    Complex::<N::RealField>::new(determinant.real(), determinant.imaginary())
321                        .sqrt();
322                let leading = self.coefficients[2];
323                let leading = Complex::<N::RealField>::new(leading.real(), leading.imaginary());
324                let leading = leading
325                    * Complex::<N::RealField>::new(
326                        N::from_f64(2.0).unwrap().real(),
327                        N::zero().real(),
328                    );
329                let secondary = self.coefficients[1];
330                let secondary =
331                    Complex::<N::RealField>::new(secondary.real(), secondary.imaginary());
332                let positive = (-secondary + determinant) / leading;
333                let negative = (-secondary - determinant) / leading;
334                return Ok(VecDeque::from(vec![positive, negative]));
335            }
336            _ => {}
337        }
338
339        let complex = self.make_complex();
340        let derivative = complex.derivative();
341
342        let mut guess = Complex::<N::RealField>::zero();
343        let mut k = 0;
344        'out: while k < n_max {
345            let val = complex.evaluate(guess);
346            if val.abs() < tol {
347                break 'out;
348            }
349            let (deriv, second_deriv) = derivative.evaluate_derivative(guess);
350            let deriv_quotient = deriv / val;
351            let g_sq = deriv_quotient.powi(2);
352            let second_deriv_quotient = g_sq - second_deriv / val;
353            let order = Complex::<N::RealField>::from_usize(self.coefficients.len() - 1).unwrap();
354            let sqrt = ((order - Complex::<N::RealField>::one())
355                * (order * second_deriv_quotient - g_sq))
356                .sqrt();
357            let plus = deriv_quotient + sqrt;
358            let minus = deriv_quotient - sqrt;
359            let a = if plus.abs() > minus.abs() {
360                order / plus
361            } else {
362                order / minus
363            };
364            guess -= a;
365            k += 1;
366        }
367        if k == n_max {
368            return Err("Polynomial roots: maximum iterations exceeded".to_owned());
369        }
370
371        let divisor = polynomial![Complex::<N::RealField>::one(), -guess];
372        let (quotient, _) = complex.divide(&divisor)?;
373        let mut roots = quotient.roots(tol, n_max)?;
374        roots.push_front(guess);
375
376        let mut corrected_roots = VecDeque::with_capacity(roots.len());
377        for root in roots.iter() {
378            corrected_roots.push_back(newton_polynomial(*root, &complex, tol, n_max)?);
379        }
380
381        Ok(corrected_roots)
382    }
383
384    // Pad to the smallest power of two less than or equal to size
385    fn pad_power_of_two(&mut self, size: usize) {
386        let mut power: usize = 1;
387        while power < size {
388            power <<= 1;
389        }
390        while self.coefficients.len() < power {
391            self.coefficients.push(N::zero());
392        }
393    }
394
395    /// Get the polynomial in point form evaluated at roots of unity at k points
396    /// where k is the smallest power of 2 greater than or equal to size
397    pub fn dft(&self, size: usize) -> Vec<Complex<<N as ComplexField>::RealField>> {
398        let mut poly = self.make_complex();
399        poly.pad_power_of_two(size);
400        let mut working = bit_reverse_copy(&poly.coefficients);
401        let len = working.len();
402        for s in 1..(len as f64).log2() as usize + 1 {
403            let m = 1 << s;
404            let angle = 2.0 * f64::consts::PI / m as f64;
405            let angle = N::RealField::from_f64(angle).unwrap();
406            let root_of_unity = Complex::<N::RealField>::new(angle.cos(), angle.sin());
407            let mut w = Complex::<N::RealField>::new(N::RealField::one(), N::RealField::zero());
408            for j in 0..m / 2 {
409                for k in (j..len).step_by(m) {
410                    let temp = w * working[k + m / 2];
411                    let u = working[k];
412                    working[k] = u + temp;
413                    working[k + m / 2] = u - temp;
414                }
415                w *= root_of_unity;
416            }
417        }
418        working
419    }
420
421    // Assumes power of 2
422    pub fn idft(
423        vec: &[Complex<<N as ComplexField>::RealField>],
424        tol: <N as ComplexField>::RealField,
425    ) -> Self {
426        let mut working = bit_reverse_copy(vec);
427        let len = working.len();
428        for s in 1..(len as f64).log2() as usize + 1 {
429            let m = 1 << s;
430            let angle = -2.0 * f64::consts::PI / m as f64;
431            let angle = N::RealField::from_f64(angle).unwrap();
432            let root_of_unity = Complex::<N::RealField>::new(angle.cos(), angle.sin());
433            let mut w = Complex::<N::RealField>::new(N::RealField::one(), N::RealField::zero());
434            for j in 0..m / 2 {
435                for k in (j..len).step_by(m) {
436                    let temp = w * working[k + m / 2];
437                    let u = working[k];
438                    working[k] = u + temp;
439                    working[k + m / 2] = u - temp;
440                }
441                w *= root_of_unity;
442            }
443        }
444        let ilen = Complex::<N::RealField>::new(
445            N::from_f64(1.0 / len as f64).unwrap().real(),
446            N::zero().real(),
447        );
448        for val in &mut working {
449            *val *= ilen;
450        }
451        let coefficients = if TypeId::of::<N::RealField>() == TypeId::of::<N>() {
452            working
453                .iter()
454                .map(|c| N::from_real(c.re))
455                .collect::<Vec<_>>()
456        } else {
457            working
458                .iter()
459                .map(|c| N::from_real(c.re) + (-N::one()).sqrt() * N::from_real(c.im))
460                .collect::<Vec<_>>()
461        };
462
463        let mut poly = Polynomial {
464            coefficients,
465            tolerance: tol,
466        };
467        poly.purge_leading();
468        poly
469    }
470}
471
472fn bit_reverse(mut k: usize, num_bits: usize) -> usize {
473    let mut result: usize = 0;
474    for _ in 0..num_bits {
475        result |= k & 1;
476        result <<= 1;
477        k >>= 1;
478    }
479    result >>= 1;
480    result
481}
482
483// Assumes vec is a power of 2 length
484fn bit_reverse_copy<N: RealField + Copy>(vec: &[Complex<N>]) -> Vec<Complex<N>> {
485    let len = vec.len();
486    let mut result = vec![Complex::new(N::zero(), N::zero()); len];
487    let num_bits = (len as f64).log2() as usize;
488    for k in 0..len {
489        result[bit_reverse(k, num_bits)] = vec[k];
490    }
491    result
492}
493
494impl<N: ComplexField + FromPrimitive + Copy> FromIterator<N> for Polynomial<N>
495where
496    <N as ComplexField>::RealField: FromPrimitive + Copy,
497{
498    fn from_iter<I: IntoIterator<Item = N>>(iter: I) -> Polynomial<N> {
499        Polynomial {
500            coefficients: Vec::from_iter(iter),
501            tolerance: N::RealField::from_f64(1e-10).unwrap(),
502        }
503    }
504}
505
506impl<N: ComplexField + FromPrimitive + Copy> Default for Polynomial<N>
507where
508    <N as ComplexField>::RealField: FromPrimitive + Copy,
509{
510    fn default() -> Self {
511        Self::new()
512    }
513}
514
515impl<N: ComplexField + FromPrimitive + Copy> Zero for Polynomial<N>
516where
517    <N as ComplexField>::RealField: FromPrimitive + Copy,
518{
519    fn zero() -> Polynomial<N> {
520        Polynomial::new()
521    }
522
523    fn is_zero(&self) -> bool {
524        for val in &self.coefficients {
525            if !val.is_zero() {
526                return false;
527            }
528        }
529        true
530    }
531}
532
533// Operator overloading
534
535impl<N: ComplexField + FromPrimitive + Copy> ops::Add<N> for Polynomial<N>
536where
537    <N as ComplexField>::RealField: FromPrimitive + Copy,
538{
539    type Output = Polynomial<N>;
540
541    fn add(mut self, rhs: N) -> Polynomial<N> {
542        self.coefficients[0] += rhs;
543        self
544    }
545}
546
547impl<N: ComplexField + FromPrimitive + Copy> ops::Add<N> for &Polynomial<N>
548where
549    <N as ComplexField>::RealField: FromPrimitive + Copy,
550{
551    type Output = Polynomial<N>;
552
553    fn add(self, rhs: N) -> Polynomial<N> {
554        let mut coefficients = Vec::from(self.coefficients.as_slice());
555        coefficients[0] += rhs;
556        Polynomial {
557            coefficients,
558            tolerance: self.tolerance,
559        }
560    }
561}
562
563impl<N: ComplexField + FromPrimitive + Copy> ops::Add<Polynomial<N>> for Polynomial<N>
564where
565    <N as ComplexField>::RealField: FromPrimitive + Copy,
566{
567    type Output = Polynomial<N>;
568
569    fn add(mut self, rhs: Polynomial<N>) -> Polynomial<N> {
570        let min_order = self.coefficients.len().min(rhs.coefficients.len());
571        for (ind, val) in self.coefficients.iter_mut().take(min_order).enumerate() {
572            *val += rhs.coefficients[ind];
573        }
574
575        for val in rhs.coefficients.iter().skip(min_order) {
576            self.coefficients.push(*val);
577        }
578
579        self
580    }
581}
582
583impl<N: ComplexField + FromPrimitive + Copy> ops::Add<&Polynomial<N>> for Polynomial<N>
584where
585    <N as ComplexField>::RealField: FromPrimitive + Copy,
586{
587    type Output = Polynomial<N>;
588
589    fn add(mut self, rhs: &Polynomial<N>) -> Polynomial<N> {
590        let min_order = self.coefficients.len().min(rhs.coefficients.len());
591        for (ind, val) in self.coefficients.iter_mut().take(min_order).enumerate() {
592            *val += rhs.coefficients[ind];
593        }
594
595        // Will only run if rhs has higher order
596        for val in rhs.coefficients.iter().skip(min_order) {
597            self.coefficients.push(*val);
598        }
599
600        self
601    }
602}
603
604impl<N: ComplexField + FromPrimitive + Copy> ops::Add<Polynomial<N>> for &Polynomial<N>
605where
606    <N as ComplexField>::RealField: FromPrimitive + Copy,
607{
608    type Output = Polynomial<N>;
609
610    fn add(self, rhs: Polynomial<N>) -> Polynomial<N> {
611        let min_order = self.coefficients.len().min(rhs.coefficients.len());
612        let mut coefficients =
613            Vec::with_capacity(self.coefficients.len().max(rhs.coefficients.len()));
614        for (ind, val) in self.coefficients.iter().take(min_order).enumerate() {
615            coefficients.push(*val + rhs.coefficients[ind]);
616        }
617
618        // Only one loop will run
619        for val in self.coefficients.iter().skip(min_order) {
620            coefficients.push(*val);
621        }
622
623        for val in rhs.coefficients.iter().skip(min_order) {
624            coefficients.push(*val);
625        }
626
627        Polynomial {
628            coefficients,
629            tolerance: self.tolerance,
630        }
631    }
632}
633
634impl<N: ComplexField + FromPrimitive + Copy> ops::Add<&Polynomial<N>> for &Polynomial<N>
635where
636    <N as ComplexField>::RealField: FromPrimitive + Copy,
637{
638    type Output = Polynomial<N>;
639
640    fn add(self, rhs: &Polynomial<N>) -> Polynomial<N> {
641        let min_order = self.coefficients.len().min(rhs.coefficients.len());
642        let mut coefficients =
643            Vec::with_capacity(self.coefficients.len().max(rhs.coefficients.len()));
644        for (ind, val) in self.coefficients.iter().take(min_order).enumerate() {
645            coefficients.push(*val + rhs.coefficients[ind]);
646        }
647
648        // Only one loop will run
649        for val in self.coefficients.iter().skip(min_order) {
650            coefficients.push(*val);
651        }
652
653        for val in rhs.coefficients.iter().skip(min_order) {
654            coefficients.push(*val);
655        }
656
657        Polynomial {
658            coefficients,
659            tolerance: self.tolerance,
660        }
661    }
662}
663
664impl<N: ComplexField + FromPrimitive + Copy> ops::AddAssign<N> for Polynomial<N>
665where
666    <N as ComplexField>::RealField: FromPrimitive + Copy,
667{
668    fn add_assign(&mut self, rhs: N) {
669        self.coefficients[0] += rhs;
670    }
671}
672
673impl<N: ComplexField + FromPrimitive + Copy> ops::AddAssign<Polynomial<N>> for Polynomial<N>
674where
675    <N as ComplexField>::RealField: FromPrimitive + Copy,
676{
677    fn add_assign(&mut self, rhs: Polynomial<N>) {
678        let min_order = self.coefficients.len().min(rhs.coefficients.len());
679        for (ind, val) in self.coefficients.iter_mut().take(min_order).enumerate() {
680            *val += rhs.coefficients[ind];
681        }
682
683        for val in rhs.coefficients.iter().skip(min_order) {
684            self.coefficients.push(*val);
685        }
686    }
687}
688
689impl<N: ComplexField + FromPrimitive + Copy> ops::AddAssign<&Polynomial<N>> for Polynomial<N>
690where
691    <N as ComplexField>::RealField: FromPrimitive + Copy,
692{
693    fn add_assign(&mut self, rhs: &Polynomial<N>) {
694        let min_order = self.coefficients.len().min(rhs.coefficients.len());
695        for (ind, val) in self.coefficients.iter_mut().take(min_order).enumerate() {
696            *val += rhs.coefficients[ind];
697        }
698
699        for val in rhs.coefficients.iter().skip(min_order) {
700            self.coefficients.push(*val);
701        }
702    }
703}
704
705impl<N: ComplexField + FromPrimitive + Copy> ops::Sub<N> for Polynomial<N>
706where
707    <N as ComplexField>::RealField: FromPrimitive + Copy,
708{
709    type Output = Polynomial<N>;
710
711    fn sub(mut self, rhs: N) -> Polynomial<N> {
712        self.coefficients[0] -= rhs;
713        self
714    }
715}
716
717impl<N: ComplexField + FromPrimitive + Copy> ops::Sub<N> for &Polynomial<N>
718where
719    <N as ComplexField>::RealField: FromPrimitive + Copy,
720{
721    type Output = Polynomial<N>;
722
723    fn sub(self, rhs: N) -> Polynomial<N> {
724        let mut coefficients = Vec::from(self.coefficients.as_slice());
725        coefficients[0] -= rhs;
726        Polynomial {
727            coefficients,
728            tolerance: self.tolerance,
729        }
730    }
731}
732
733impl<N: ComplexField + FromPrimitive + Copy> ops::Sub<Polynomial<N>> for Polynomial<N>
734where
735    <N as ComplexField>::RealField: FromPrimitive + Copy,
736{
737    type Output = Polynomial<N>;
738
739    fn sub(mut self, rhs: Polynomial<N>) -> Polynomial<N> {
740        let min_order = self.coefficients.len().min(rhs.coefficients.len());
741        for (ind, val) in self.coefficients.iter_mut().take(min_order).enumerate() {
742            *val -= rhs.coefficients[ind];
743        }
744
745        for val in rhs.coefficients.iter().skip(min_order) {
746            self.coefficients.push(-*val);
747        }
748
749        self
750    }
751}
752
753impl<N: ComplexField + FromPrimitive + Copy> ops::Sub<Polynomial<N>> for &Polynomial<N>
754where
755    <N as ComplexField>::RealField: FromPrimitive + Copy,
756{
757    type Output = Polynomial<N>;
758
759    fn sub(self, rhs: Polynomial<N>) -> Polynomial<N> {
760        let min_order = self.coefficients.len().min(rhs.coefficients.len());
761        let mut coefficients =
762            Vec::with_capacity(self.coefficients.len().max(rhs.coefficients.len()));
763        for (ind, val) in self.coefficients.iter().take(min_order).enumerate() {
764            coefficients.push(*val - rhs.coefficients[ind]);
765        }
766
767        // Only one for loop runs
768        for val in self.coefficients.iter().skip(min_order) {
769            coefficients.push(*val);
770        }
771
772        for val in rhs.coefficients.iter().skip(min_order) {
773            coefficients.push(-*val);
774        }
775
776        Polynomial {
777            coefficients,
778            tolerance: self.tolerance,
779        }
780    }
781}
782
783impl<N: ComplexField + FromPrimitive + Copy> ops::Sub<&Polynomial<N>> for Polynomial<N>
784where
785    <N as ComplexField>::RealField: FromPrimitive + Copy,
786{
787    type Output = Polynomial<N>;
788
789    fn sub(mut self, rhs: &Polynomial<N>) -> Polynomial<N> {
790        let min_order = self.coefficients.len().min(rhs.coefficients.len());
791        for (ind, val) in self.coefficients.iter_mut().take(min_order).enumerate() {
792            *val -= rhs.coefficients[ind];
793        }
794
795        for val in rhs.coefficients.iter().skip(min_order) {
796            self.coefficients.push(-*val);
797        }
798
799        self
800    }
801}
802
803impl<N: ComplexField + FromPrimitive + Copy> ops::Sub<&Polynomial<N>> for &Polynomial<N>
804where
805    <N as ComplexField>::RealField: FromPrimitive + Copy,
806{
807    type Output = Polynomial<N>;
808
809    fn sub(self, rhs: &Polynomial<N>) -> Polynomial<N> {
810        let min_order = self.coefficients.len().min(rhs.coefficients.len());
811        let mut coefficients =
812            Vec::with_capacity(self.coefficients.len().max(rhs.coefficients.len()));
813        for (ind, val) in self.coefficients.iter().take(min_order).enumerate() {
814            coefficients.push(*val - rhs.coefficients[ind]);
815        }
816
817        // Only one for loop runs
818        for val in self.coefficients.iter().skip(min_order) {
819            coefficients.push(*val);
820        }
821
822        for val in rhs.coefficients.iter().skip(min_order) {
823            coefficients.push(-*val);
824        }
825
826        Polynomial {
827            coefficients,
828            tolerance: self.tolerance,
829        }
830    }
831}
832
833impl<N: ComplexField + FromPrimitive + Copy> ops::SubAssign<N> for Polynomial<N>
834where
835    <N as ComplexField>::RealField: FromPrimitive + Copy,
836{
837    fn sub_assign(&mut self, rhs: N) {
838        self.coefficients[0] -= rhs;
839    }
840}
841
842impl<N: ComplexField + FromPrimitive + Copy> ops::SubAssign<Polynomial<N>> for Polynomial<N>
843where
844    <N as ComplexField>::RealField: FromPrimitive + Copy,
845{
846    fn sub_assign(&mut self, rhs: Polynomial<N>) {
847        let min_order = self.coefficients.len().min(rhs.coefficients.len());
848        for (ind, val) in self.coefficients.iter_mut().take(min_order).enumerate() {
849            *val -= rhs.coefficients[ind];
850        }
851
852        for val in rhs.coefficients.iter().skip(min_order) {
853            self.coefficients.push(-*val);
854        }
855    }
856}
857
858impl<N: ComplexField + FromPrimitive + Copy> ops::SubAssign<&Polynomial<N>> for Polynomial<N>
859where
860    <N as ComplexField>::RealField: FromPrimitive + Copy,
861{
862    fn sub_assign(&mut self, rhs: &Polynomial<N>) {
863        let min_order = self.coefficients.len().min(rhs.coefficients.len());
864        for (ind, val) in self.coefficients.iter_mut().take(min_order).enumerate() {
865            *val -= rhs.coefficients[ind];
866        }
867
868        for val in rhs.coefficients.iter().skip(min_order) {
869            self.coefficients.push(-*val);
870        }
871    }
872}
873
874impl<N: ComplexField + FromPrimitive + Copy> ops::Mul<N> for Polynomial<N>
875where
876    <N as ComplexField>::RealField: FromPrimitive + Copy,
877{
878    type Output = Polynomial<N>;
879
880    fn mul(mut self, rhs: N) -> Polynomial<N> {
881        for val in &mut self.coefficients {
882            *val *= rhs;
883        }
884        self
885    }
886}
887
888impl<N: ComplexField + FromPrimitive + Copy> ops::Mul<N> for &Polynomial<N>
889where
890    <N as ComplexField>::RealField: FromPrimitive + Copy,
891{
892    type Output = Polynomial<N>;
893
894    fn mul(self, rhs: N) -> Polynomial<N> {
895        let mut coefficients = Vec::with_capacity(self.coefficients.len());
896        for val in &self.coefficients {
897            coefficients.push(*val * rhs);
898        }
899        Polynomial {
900            coefficients,
901            tolerance: self.tolerance,
902        }
903    }
904}
905
906fn multiply<N: ComplexField + FromPrimitive + Copy>(
907    lhs: &Polynomial<N>,
908    rhs: &Polynomial<N>,
909) -> Polynomial<N>
910where
911    <N as ComplexField>::RealField: FromPrimitive + Copy,
912{
913    // Do scalar multiplication if one side has no powers
914    if rhs.coefficients.len() == 1 {
915        return lhs * rhs.coefficients[0];
916    }
917    if lhs.coefficients.len() == 1 {
918        return rhs * lhs.coefficients[0];
919    }
920
921    // Special case linear term multiplication
922    if rhs.coefficients.len() == 2 {
923        let mut shifted = lhs * rhs.coefficients[1];
924        shifted.coefficients.insert(0, N::zero());
925        return shifted + lhs * rhs.coefficients[0];
926    }
927    if lhs.coefficients.len() == 2 {
928        let mut shifted = rhs * lhs.coefficients[1];
929        shifted.coefficients.insert(0, N::zero());
930        return shifted + rhs * lhs.coefficients[0];
931    }
932
933    let bound = lhs.coefficients.len().max(rhs.coefficients.len()) * 2;
934    let left_points = lhs.dft(bound);
935    let right_points = rhs.dft(bound);
936    let product_points: Vec<_> = left_points
937        .iter()
938        .zip(right_points.iter())
939        .map(|(l_p, r_p)| *l_p * r_p)
940        .collect();
941    Polynomial::<N>::idft(&product_points, lhs.tolerance)
942}
943
944impl<N: ComplexField + FromPrimitive + Copy> ops::Mul<Polynomial<N>> for Polynomial<N>
945where
946    <N as ComplexField>::RealField: FromPrimitive + Copy,
947{
948    type Output = Polynomial<N>;
949
950    fn mul(self, rhs: Polynomial<N>) -> Polynomial<N> {
951        multiply(&self, &rhs)
952    }
953}
954
955impl<N: ComplexField + FromPrimitive + Copy> ops::Mul<&Polynomial<N>> for Polynomial<N>
956where
957    <N as ComplexField>::RealField: FromPrimitive + Copy,
958{
959    type Output = Polynomial<N>;
960
961    fn mul(self, rhs: &Polynomial<N>) -> Polynomial<N> {
962        multiply(&self, rhs)
963    }
964}
965
966impl<N: ComplexField + FromPrimitive + Copy> ops::Mul<Polynomial<N>> for &Polynomial<N>
967where
968    <N as ComplexField>::RealField: FromPrimitive + Copy,
969{
970    type Output = Polynomial<N>;
971
972    fn mul(self, rhs: Polynomial<N>) -> Polynomial<N> {
973        multiply(self, &rhs)
974    }
975}
976
977impl<N: ComplexField + FromPrimitive + Copy> ops::Mul<&Polynomial<N>> for &Polynomial<N>
978where
979    <N as ComplexField>::RealField: FromPrimitive + Copy,
980{
981    type Output = Polynomial<N>;
982
983    fn mul(self, rhs: &Polynomial<N>) -> Polynomial<N> {
984        multiply(self, rhs)
985    }
986}
987
988impl<N: ComplexField + FromPrimitive + Copy> ops::MulAssign<N> for Polynomial<N>
989where
990    <N as ComplexField>::RealField: FromPrimitive + Copy,
991{
992    fn mul_assign(&mut self, rhs: N) {
993        for val in self.coefficients.iter_mut() {
994            *val *= rhs;
995        }
996    }
997}
998
999impl<N: ComplexField + FromPrimitive + Copy> ops::MulAssign<Polynomial<N>> for Polynomial<N>
1000where
1001    <N as ComplexField>::RealField: FromPrimitive + Copy,
1002{
1003    fn mul_assign(&mut self, rhs: Polynomial<N>) {
1004        self.coefficients = multiply(self, &rhs).coefficients;
1005    }
1006}
1007
1008impl<N: ComplexField + FromPrimitive + Copy> ops::MulAssign<&Polynomial<N>> for Polynomial<N>
1009where
1010    <N as ComplexField>::RealField: FromPrimitive + Copy,
1011{
1012    fn mul_assign(&mut self, rhs: &Polynomial<N>) {
1013        self.coefficients = multiply(self, rhs).coefficients;
1014    }
1015}
1016
1017impl<N: ComplexField + FromPrimitive + Copy> ops::Div<N> for Polynomial<N>
1018where
1019    <N as ComplexField>::RealField: FromPrimitive + Copy,
1020{
1021    type Output = Polynomial<N>;
1022
1023    fn div(mut self, rhs: N) -> Polynomial<N> {
1024        for val in &mut self.coefficients {
1025            *val /= rhs;
1026        }
1027        self
1028    }
1029}
1030
1031impl<N: ComplexField + FromPrimitive + Copy> ops::Div<N> for &Polynomial<N>
1032where
1033    <N as ComplexField>::RealField: FromPrimitive + Copy,
1034{
1035    type Output = Polynomial<N>;
1036
1037    fn div(self, rhs: N) -> Polynomial<N> {
1038        let mut coefficients = Vec::from(self.coefficients.as_slice());
1039        for val in &mut coefficients {
1040            *val /= rhs;
1041        }
1042        Polynomial {
1043            coefficients,
1044            tolerance: self.tolerance,
1045        }
1046    }
1047}
1048
1049impl<N: ComplexField + FromPrimitive + Copy> ops::DivAssign<N> for Polynomial<N>
1050where
1051    <N as ComplexField>::RealField: FromPrimitive + Copy,
1052{
1053    fn div_assign(&mut self, rhs: N) {
1054        for val in &mut self.coefficients {
1055            *val /= rhs;
1056        }
1057    }
1058}
1059
1060impl<N: ComplexField + FromPrimitive + Copy> ops::Neg for Polynomial<N>
1061where
1062    <N as ComplexField>::RealField: FromPrimitive + Copy,
1063{
1064    type Output = Polynomial<N>;
1065
1066    fn neg(mut self) -> Polynomial<N> {
1067        for val in &mut self.coefficients {
1068            *val = -*val;
1069        }
1070        self
1071    }
1072}
1073
1074impl<N: ComplexField + FromPrimitive + Copy> ops::Neg for &Polynomial<N>
1075where
1076    <N as ComplexField>::RealField: FromPrimitive + Copy,
1077{
1078    type Output = Polynomial<N>;
1079
1080    fn neg(self) -> Polynomial<N> {
1081        Polynomial {
1082            coefficients: self.coefficients.iter().map(|c| -*c).collect(),
1083            tolerance: self.tolerance,
1084        }
1085    }
1086}
1087
1088impl<N: ComplexField + FromPrimitive + Copy> From<N> for Polynomial<N>
1089where
1090    <N as ComplexField>::RealField: FromPrimitive + Copy,
1091{
1092    fn from(n: N) -> Polynomial<N> {
1093        polynomial![n]
1094    }
1095}
1096
1097impl<N: RealField + FromPrimitive + Copy> From<Polynomial<N>> for Polynomial<Complex<N>> {
1098    fn from(poly: Polynomial<N>) -> Polynomial<Complex<N>> {
1099        poly.make_complex()
1100    }
1101}
bacon_sci/polynomial/mod.rs

bacon_sci/polynomial/
mod.rs
