rust_poly/poly/
impl_num.rs

1#![allow(clippy::op_ref)]
2
3// Implementation of traits related to numeric operations, operators and number theory
4
5use itertools::Itertools;
6use num::{traits::CheckedRem, CheckedDiv, Complex, One, Zero};
7use std::ops::{Add, Div, Mul, Neg, Rem, Sub};
8
9use crate::{
10    util::{casting::usize_to_u32, linalg::convolve_1d},
11    Poly, RealScalar,
12};
13
14impl<T: RealScalar> Zero for Poly<T> {
15    fn zero() -> Self {
16        Self::from_real_slice(&[T::zero()])
17    }
18
19    fn is_zero(&self) -> bool {
20        debug_assert!(self.is_normalized());
21        if self.len_raw() != 1 {
22            return false;
23        }
24        self.0[0].is_zero()
25    }
26}
27
28impl<T: RealScalar> Poly<T> {
29    /// Calculate the quotient and remainder using long division. More efficient than
30    /// calculating them separately.
31    ///
32    /// # Panics
33    /// Panics if a division by zero is attempted
34    ///
35    /// # Examples
36    /// ```
37    /// use rust_poly::{Poly, poly};
38    /// use num::{Complex, One};
39    ///
40    /// let c1 = poly![1.0, 2.0, 3.0];
41    /// let c2 = poly![3.0, 2.0, 1.0];
42    /// let expected1 = (poly![3.0], poly![-8.0, -4.0]);
43    /// assert_eq!(c1.clone().div_rem(&c2).unwrap(), expected1);
44    /// ```
45    #[allow(clippy::cast_sign_loss)]
46    #[allow(clippy::cast_possible_wrap)]
47    #[must_use]
48    pub fn div_rem(self, other: &Self) -> Option<(Self, Self)> {
49        debug_assert!(self.is_normalized());
50        debug_assert!(other.is_normalized());
51
52        if other.is_zero() {
53            // bail!("Attempted to divide a polynomial by zero");
54            return None;
55        }
56
57        if other.is_one() {
58            return Some((self, Self::zero()));
59        }
60
61        let expected_degree = self.degree_raw() - other.degree_raw();
62
63        let den_c = other
64            .as_slice()
65            .last()
66            .expect("polynomials should always have at least one coefficient");
67        let den_k = other.degree_raw();
68        let mut this = self;
69        let mut div = Self::zero();
70        'for_else: {
71            // TODO: this should fail faster, what's the upper bound?
72            for _ in 0..u32::MAX {
73                if this.degree_raw() < other.degree_raw() {
74                    break 'for_else;
75                }
76                let num_c = this
77                    .as_slice()
78                    .last()
79                    .expect("polynomials should always have at least one coefficient");
80                let num_k = this.degree_raw();
81                let c = num_c / den_c;
82                let k = num_k - den_k;
83                let new_term = Self::term(c, usize_to_u32(k));
84                this = this.clone() - new_term.clone() * other;
85                div = div + new_term;
86            }
87            // else: did not converge
88            return None;
89        }
90        let rem = this;
91
92        // sanity check: result has the expected degree
93        debug_assert_eq!(div.degree_raw(), expected_degree);
94        Some((div, rem))
95    }
96}
97
98impl<T: RealScalar> One for Poly<T> {
99    fn one() -> Self {
100        Self(vec![Complex::<T>::one()])
101    }
102}
103
104impl<T: RealScalar> Add<Self> for Poly<T> {
105    type Output = Self;
106
107    fn add(self, rhs: Self) -> Self::Output {
108        // invariant: polynomials are normalized
109        debug_assert!(self.is_normalized());
110        debug_assert!(rhs.is_normalized());
111
112        let (mut longest, shortest) = if self.len_raw() >= rhs.len_raw() {
113            (self.0, rhs.0)
114        } else {
115            (rhs.0, self.0)
116        };
117        longest
118            .as_mut_slice()
119            .iter_mut()
120            .zip_longest(shortest.iter())
121            .for_each(|p| {
122                if let itertools::EitherOrBoth::Both(l, r) = p {
123                    *l += r;
124                }
125            });
126        Self(longest).normalize()
127    }
128}
129
130impl<T: RealScalar> Add<&Self> for Poly<T> {
131    type Output = Self;
132
133    fn add(self, rhs: &Self) -> Self::Output {
134        self + rhs.clone()
135    }
136}
137
138impl<T: RealScalar> Add<Poly<T>> for &Poly<T> {
139    type Output = Poly<T>;
140
141    fn add(self, rhs: Poly<T>) -> Self::Output {
142        self.clone() + rhs
143    }
144}
145
146impl<T: RealScalar> Add<&Poly<T>> for &Poly<T> {
147    type Output = Poly<T>;
148
149    fn add(self, rhs: &Poly<T>) -> Self::Output {
150        self.clone() + rhs.clone()
151    }
152}
153
154impl<T: RealScalar> Mul<Self> for Poly<T> {
155    type Output = Self;
156
157    fn mul(self, rhs: Self) -> Self::Output {
158        // invariant: polynomials are normalized
159        debug_assert!(self.is_normalized());
160        debug_assert!(rhs.is_normalized());
161
162        if self.is_zero() || rhs.is_zero() {
163            return Self::zero();
164        }
165        if self.is_one() {
166            return rhs;
167        }
168        if rhs.is_one() {
169            return self;
170        }
171
172        let ret = convolve_1d(&self.0, &rhs.0);
173        Self(ret).normalize()
174    }
175}
176
177impl<T: RealScalar> Mul<&Self> for Poly<T> {
178    type Output = Self;
179
180    fn mul(self, rhs: &Self) -> Self::Output {
181        self * rhs.clone()
182    }
183}
184
185impl<T: RealScalar> Mul<Poly<T>> for &Poly<T> {
186    type Output = Poly<T>;
187
188    fn mul(self, rhs: Poly<T>) -> Self::Output {
189        self.clone() * rhs
190    }
191}
192
193impl<T: RealScalar> Mul<&Poly<T>> for &Poly<T> {
194    type Output = Poly<T>;
195
196    fn mul(self, rhs: &Poly<T>) -> Self::Output {
197        self.clone() * rhs.clone()
198    }
199}
200
201impl<T: RealScalar> Mul<Complex<T>> for Poly<T> {
202    type Output = Self;
203
204    fn mul(self, rhs: Complex<T>) -> Self::Output {
205        self * &rhs
206    }
207}
208
209impl<T: RealScalar> Mul<&Complex<T>> for Poly<T> {
210    type Output = Self;
211
212    fn mul(self, rhs: &Complex<T>) -> Self::Output {
213        let mut lhs = self;
214        lhs.0.iter_mut().for_each(|c| *c *= rhs);
215        lhs.normalize()
216    }
217}
218
219impl<T: RealScalar> Mul<Complex<T>> for &Poly<T> {
220    type Output = Poly<T>;
221
222    fn mul(self, rhs: Complex<T>) -> Self::Output {
223        self.clone() * rhs
224    }
225}
226
227impl<T: RealScalar> Mul<&Complex<T>> for &Poly<T> {
228    type Output = Poly<T>;
229
230    fn mul(self, rhs: &Complex<T>) -> Self::Output {
231        self.clone() * rhs
232    }
233}
234
235impl<T: RealScalar> Sub<Self> for Poly<T> {
236    type Output = Self;
237
238    fn sub(self, rhs: Self) -> Self::Output {
239        // invariant: polynomials are normalized
240        debug_assert!(self.is_normalized());
241        debug_assert!(rhs.is_normalized());
242
243        let (mut longest, shortest) = if self.len_raw() >= rhs.len_raw() {
244            (self.0, rhs.0)
245        } else {
246            (rhs.0, self.0)
247        };
248        longest
249            .as_mut_slice()
250            .iter_mut()
251            .zip_longest(shortest.iter())
252            .for_each(|p| {
253                if let itertools::EitherOrBoth::Both(l, r) = p {
254                    *l -= r;
255                }
256            });
257        Self(longest).normalize()
258    }
259}
260
261impl<T: RealScalar> Sub<&Self> for Poly<T> {
262    type Output = Self;
263
264    fn sub(self, rhs: &Self) -> Self::Output {
265        self - rhs.clone()
266    }
267}
268
269impl<T: RealScalar> Sub<Poly<T>> for &Poly<T> {
270    type Output = Poly<T>;
271
272    fn sub(self, rhs: Poly<T>) -> Self::Output {
273        self.clone() - rhs
274    }
275}
276
277impl<T: RealScalar> Sub<&Poly<T>> for &Poly<T> {
278    type Output = Poly<T>;
279
280    fn sub(self, rhs: &Poly<T>) -> Self::Output {
281        self.clone() - rhs.clone()
282    }
283}
284
285impl<T: RealScalar> CheckedDiv for Poly<T> {
286    fn checked_div(&self, rhs: &Self) -> Option<Self> {
287        self.clone().checked_div_impl(rhs)
288    }
289}
290
291impl<T: RealScalar> CheckedRem for Poly<T> {
292    fn checked_rem(&self, rhs: &Self) -> Option<Self> {
293        self.clone().checked_rem_impl(rhs)
294    }
295}
296
297impl<T: RealScalar> Div<Self> for Poly<T> {
298    type Output = Self;
299
300    fn div(self, rhs: Self) -> Self::Output {
301        self / &rhs
302    }
303}
304
305impl<T: RealScalar> Div<&Self> for Poly<T> {
306    type Output = Self;
307
308    fn div(self, rhs: &Self) -> Self::Output {
309        self.checked_div_impl(rhs).expect("Division by zero")
310    }
311}
312
313impl<T: RealScalar> Div<Poly<T>> for &Poly<T> {
314    type Output = Poly<T>;
315
316    fn div(self, rhs: Poly<T>) -> Self::Output {
317        self.clone() / &rhs
318    }
319}
320
321impl<T: RealScalar> Div<&Poly<T>> for &Poly<T> {
322    type Output = Poly<T>;
323
324    fn div(self, rhs: &Poly<T>) -> Self::Output {
325        self.clone() / rhs
326    }
327}
328
329impl<T: RealScalar> Div<Complex<T>> for Poly<T> {
330    type Output = Self;
331
332    fn div(self, rhs: Complex<T>) -> Self::Output {
333        self / &rhs
334    }
335}
336
337impl<T: RealScalar> Div<&Complex<T>> for Poly<T> {
338    type Output = Self;
339
340    fn div(self, rhs: &Complex<T>) -> Self::Output {
341        let mut lhs = self;
342        lhs.0.iter_mut().for_each(|c| *c /= rhs);
343        lhs.normalize()
344    }
345}
346
347impl<T: RealScalar> Div<Complex<T>> for &Poly<T> {
348    type Output = Poly<T>;
349
350    fn div(self, rhs: Complex<T>) -> Self::Output {
351        self.clone() / rhs
352    }
353}
354
355impl<T: RealScalar> Div<&Complex<T>> for &Poly<T> {
356    type Output = Poly<T>;
357
358    fn div(self, rhs: &Complex<T>) -> Self::Output {
359        self.clone() / rhs
360    }
361}
362
363impl<T: RealScalar> Rem<Self> for Poly<T> {
364    type Output = Self;
365
366    fn rem(self, rhs: Self) -> Self::Output {
367        self % &rhs
368    }
369}
370
371impl<T: RealScalar> Rem<&Self> for Poly<T> {
372    type Output = Self;
373
374    fn rem(self, rhs: &Self) -> Self::Output {
375        self.checked_rem_impl(rhs).expect("Division by zero")
376    }
377}
378
379impl<T: RealScalar> Rem<Poly<T>> for &Poly<T> {
380    type Output = Poly<T>;
381
382    fn rem(self, rhs: Poly<T>) -> Self::Output {
383        self.clone() % &rhs
384    }
385}
386
387impl<T: RealScalar> Rem<&Poly<T>> for &Poly<T> {
388    type Output = Poly<T>;
389
390    fn rem(self, rhs: &Poly<T>) -> Self::Output {
391        self.clone() % rhs
392    }
393}
394
395impl<T: RealScalar> Neg for Poly<T> {
396    type Output = Self;
397
398    fn neg(self) -> Self::Output {
399        Self::zero() - self
400    }
401}
402
403impl<T: RealScalar> Neg for &Poly<T> {
404    type Output = Poly<T>;
405
406    fn neg(self) -> Self::Output {
407        -self.clone()
408    }
409}
410
411impl<T: RealScalar> std::iter::Sum for Poly<T> {
412    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
413        iter.fold(Self::zero(), |acc, x| acc + x).normalize()
414    }
415}
416
417impl<T: RealScalar> Poly<T> {
418    pub(crate) fn checked_div_impl(self, rhs: &Self) -> Option<Self> {
419        Some(self.div_rem(rhs)?.0)
420    }
421
422    pub(crate) fn checked_rem_impl(self, rhs: &Self) -> Option<Self> {
423        Some(self.div_rem(rhs)?.1)
424    }
425}
426
427#[cfg(test)]
428mod test {
429    #[test]
430    fn div() {
431        let dividend = poly![-4.0, 0.0, -2.0, 1.0];
432        let divisor = poly![-3.0, 1.0];
433        let (q, r) = dividend.div_rem(&divisor).unwrap();
434        assert_eq!(q, poly![3.0, 1.0, 1.0]);
435        assert_eq!(r, poly![5.0]);
436    }
437}