1use alloc::vec::Vec;
2use core::ops::{
3 Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Shr, ShrAssign, Sub, SubAssign,
4};
5
6use num::{Complex, One, Zero};
7
8use crate::numerics::{Abs, AbsSqrt, IsNegativeOne, IsPositive, TryFromUsizeExact};
9use crate::polynomial::find_roots::{discriminant_trinomial, trinomial_roots};
10use crate::polynomial::polynomial::term_with_deg;
11use crate::{
12 Degree, Derivable, Evaluable, Integrable, Integral, LinearBinomial, MutablePolynomial,
13 Polynomial, Roots, SizedPolynomial, Term, TryAddError,
14};
15
16#[derive(Debug, Clone)]
17pub struct QuadraticTrinomial<N> {
21 pub coefficients: [N; 3],
22}
23
24impl<N: Sized> QuadraticTrinomial<N> {
25 pub fn new(coefficients: [N; 3]) -> QuadraticTrinomial<N> {
36 QuadraticTrinomial { coefficients }
37 }
38}
39
40impl<N: Zero + Copy> QuadraticTrinomial<N> {
41 pub fn ordered_term_iter(&self) -> impl Iterator<Item = (N, usize)> + '_ {
42 self.coefficients
43 .iter()
44 .enumerate()
45 .filter_map(|(index, &coeff)| {
46 if coeff.is_zero() {
47 None
48 } else {
49 Some((coeff, 2 - index))
50 }
51 })
52 }
53}
54
55impl<N> QuadraticTrinomial<N>
56where
57 N: Copy
58 + Zero
59 + Mul<Output = N>
60 + Neg<Output = N>
61 + Sub<Output = N>
62 + From<u8>
63 + Div<Output = N>
64 + AbsSqrt
65 + IsPositive
66 + One,
67{
68 pub fn discriminant(&self) -> N {
69 let [a, b, c] = self.coefficients;
70 discriminant_trinomial(a, b, c)
71 }
72
73 pub fn roots(&self) -> Roots<N> {
76 let [a, b, c] = self.coefficients;
77 trinomial_roots(a, b, c)
78 }
79
80 pub fn complex_factors(&self) -> (N, LinearBinomial<Complex<N>>, LinearBinomial<Complex<N>>) {
81 match self.roots() {
82 Roots::TwoComplexRoots(root_a, root_b) => (
83 self.coefficients[0],
84 LinearBinomial::new([Complex::new(N::one(), -N::zero()), root_a]),
85 LinearBinomial::new([Complex::new(N::one(), -N::zero()), root_b]),
86 ),
87 Roots::TwoRealRoots(a, b) => (
88 self.coefficients[0],
89 LinearBinomial::new([
90 Complex::new(N::one(), N::zero()),
91 Complex::new(-a, N::zero()),
92 ]),
93 LinearBinomial::new([
94 Complex::new(N::one(), N::zero()),
95 Complex::new(-b, N::zero()),
96 ]),
97 ),
98 _ => unreachable!(),
99 }
100 }
101
102 pub fn real_factors(&self) -> Option<(N, LinearBinomial<N>, LinearBinomial<N>)> {
103 if let Roots::TwoRealRoots(root_a, root_b) = self.roots() {
104 Some((
105 self.coefficients[0],
106 LinearBinomial::new([N::one(), -root_a]),
107 LinearBinomial::new([N::one(), -root_b]),
108 ))
109 } else {
110 None
111 }
112 }
113}
114
115impl<N: Copy + Zero> SizedPolynomial<N> for QuadraticTrinomial<N> {
116 fn term_with_degree(&self, degree: usize) -> Term<N> {
128 term_with_deg(&self.coefficients, degree)
129 }
130
131 fn terms_as_vec(&self) -> Vec<(N, usize)> {
132 self.ordered_term_iter().collect()
133 }
134
135 fn degree(&self) -> Degree {
151 if !self.coefficients[0].is_zero() {
152 Degree::Num(2)
153 } else if !self.coefficients[1].is_zero() {
154 Degree::Num(1)
155 } else if !self.coefficients[2].is_zero() {
156 Degree::Num(0)
157 } else {
158 Degree::NegInf
159 }
160 }
161
162 fn zero() -> Self {
171 QuadraticTrinomial::new([N::zero(); 3])
172 }
173
174 fn set_to_zero(&mut self) {
186 self.coefficients = [N::zero(); 3];
187 }
188}
189
190impl<N> MutablePolynomial<N> for QuadraticTrinomial<N>
191where
192 N: Zero + SubAssign + AddAssign + Copy,
193{
194 fn try_add_term(&mut self, coeff: N, degree: usize) -> Result<(), TryAddError> {
195 if degree <= 2 {
196 self.coefficients[2 - degree] += coeff;
197 Ok(())
198 } else {
199 Err(TryAddError::DegreeOutOfBounds)
200 }
201 }
202
203 fn try_sub_term(&mut self, coeff: N, degree: usize) -> Result<(), TryAddError> {
204 if degree <= 2 {
205 self.coefficients[2 - degree] -= coeff;
206 Ok(())
207 } else {
208 Err(TryAddError::DegreeOutOfBounds)
209 }
210 }
211}
212
213impl<N> Evaluable<N> for QuadraticTrinomial<N>
214where
215 N: Add<Output = N> + Mul<Output = N> + Copy,
216{
217 fn eval(&self, point: N) -> N {
228 point * (self.coefficients[0] * point + self.coefficients[1]) + self.coefficients[2]
229 }
230}
231
232impl<N> Derivable<N> for QuadraticTrinomial<N>
233where
234 N: Zero + One + Copy + Mul<Output = N> + TryFromUsizeExact,
235{
236 fn derivative(&self) -> QuadraticTrinomial<N> {
246 QuadraticTrinomial::new([
247 N::zero(),
248 self.coefficients[0]
249 * N::try_from_usize_exact(2).expect("Failed to convert 2usize to N."),
250 self.coefficients[1],
251 ])
252 }
253}
254
255impl<N> Integrable<N, Polynomial<N>> for QuadraticTrinomial<N>
256where
257 N: Zero
258 + TryFromUsizeExact
259 + Copy
260 + DivAssign
261 + Mul<Output = N>
262 + MulAssign
263 + AddAssign
264 + Div<Output = N>,
265{
266 fn integral(&self) -> Integral<N, Polynomial<N>> {
280 Integral::new(Polynomial::new(vec![
281 self.coefficients[0]
282 / N::try_from_usize_exact(3).expect("Failed to convert 3usize to N."),
283 self.coefficients[1]
284 / N::try_from_usize_exact(2).expect("Failed to convert 2usize to N."),
285 self.coefficients[2],
286 N::zero(),
287 ]))
288 }
289}
290
291impl<N> PartialEq for QuadraticTrinomial<N>
312where
313 N: Zero + PartialEq + Copy,
314{
315 fn eq(&self, other: &Self) -> bool {
317 self.coefficients == other.coefficients
318 }
319}
320
321macro_rules! from_trinomial_a_to_b {
322 ($A:ty, $B:ty) => {
323 impl From<QuadraticTrinomial<$A>> for QuadraticTrinomial<$B> {
324 fn from(item: QuadraticTrinomial<$A>) -> Self {
325 QuadraticTrinomial::new([
326 item.coefficients[0] as $B,
327 item.coefficients[1] as $B,
328 item.coefficients[2] as $B,
329 ])
330 }
331 }
332 };
333}
334
335upcast!(from_trinomial_a_to_b);
336poly_from_str!(QuadraticTrinomial);
337fmt_poly!(QuadraticTrinomial);
338
339impl<N: Copy + Neg<Output = N>> Neg for QuadraticTrinomial<N> {
340 type Output = QuadraticTrinomial<N>;
341
342 fn neg(self) -> QuadraticTrinomial<N> {
343 QuadraticTrinomial::new([
344 -self.coefficients[0],
345 -self.coefficients[1],
346 -self.coefficients[2],
347 ])
348 }
349}
350
351impl<N> Sub<QuadraticTrinomial<N>> for QuadraticTrinomial<N>
352where
353 N: Copy + Sub<Output = N>,
354{
355 type Output = QuadraticTrinomial<N>;
356
357 fn sub(self, rhs: QuadraticTrinomial<N>) -> QuadraticTrinomial<N> {
358 QuadraticTrinomial::new([
359 self.coefficients[0] - rhs.coefficients[0],
360 self.coefficients[1] - rhs.coefficients[1],
361 self.coefficients[2] - rhs.coefficients[2],
362 ])
363 }
364}
365
366impl<N> SubAssign<QuadraticTrinomial<N>> for QuadraticTrinomial<N>
367where
368 N: SubAssign + Copy,
369{
370 fn sub_assign(&mut self, rhs: QuadraticTrinomial<N>) {
371 self.coefficients[0] -= rhs.coefficients[0];
372 self.coefficients[1] -= rhs.coefficients[1];
373 self.coefficients[2] -= rhs.coefficients[2];
374 }
375}
376
377impl<N> Add<QuadraticTrinomial<N>> for QuadraticTrinomial<N>
378where
379 N: Add<Output = N> + Copy,
380{
381 type Output = QuadraticTrinomial<N>;
382
383 fn add(self, rhs: QuadraticTrinomial<N>) -> QuadraticTrinomial<N> {
384 QuadraticTrinomial::new([
385 self.coefficients[0] + rhs.coefficients[0],
386 self.coefficients[1] + rhs.coefficients[1],
387 self.coefficients[2] + rhs.coefficients[2],
388 ])
389 }
390}
391
392impl<N: Copy + AddAssign> AddAssign<QuadraticTrinomial<N>> for QuadraticTrinomial<N> {
393 fn add_assign(&mut self, rhs: QuadraticTrinomial<N>) {
394 self.coefficients[0] += rhs.coefficients[0];
395 self.coefficients[1] += rhs.coefficients[1];
396 self.coefficients[2] += rhs.coefficients[2];
397 }
398}
399
400impl<N: Mul<Output = N> + Copy> Mul<N> for QuadraticTrinomial<N> {
434 type Output = QuadraticTrinomial<N>;
435
436 fn mul(self, rhs: N) -> QuadraticTrinomial<N> {
437 QuadraticTrinomial::new([
438 self.coefficients[0] * rhs,
439 self.coefficients[1] * rhs,
440 self.coefficients[2] * rhs,
441 ])
442 }
443}
444
445impl<N: MulAssign + Copy> MulAssign<N> for QuadraticTrinomial<N> {
446 fn mul_assign(&mut self, rhs: N) {
447 self.coefficients[0] *= rhs;
448 self.coefficients[1] *= rhs;
449 self.coefficients[2] *= rhs;
450 }
451}
452
453impl<N: Div<Output = N> + Copy> Div<N> for QuadraticTrinomial<N> {
454 type Output = QuadraticTrinomial<N>;
455
456 fn div(self, rhs: N) -> QuadraticTrinomial<N> {
457 QuadraticTrinomial::new([
458 self.coefficients[0] / rhs,
459 self.coefficients[1] / rhs,
460 self.coefficients[2] / rhs,
461 ])
462 }
463}
464
465impl<N: DivAssign + Copy> DivAssign<N> for QuadraticTrinomial<N> {
466 fn div_assign(&mut self, rhs: N) {
467 self.coefficients[0] /= rhs;
468 self.coefficients[1] /= rhs;
469 self.coefficients[2] /= rhs;
470 }
471}
472
473impl<N: Zero + Copy> Shr<u32> for QuadraticTrinomial<N> {
525 type Output = QuadraticTrinomial<N>;
526
527 fn shr(self, rhs: u32) -> QuadraticTrinomial<N> {
528 match rhs {
529 0 => QuadraticTrinomial::new(self.coefficients),
530 1 => QuadraticTrinomial::new([N::zero(), self.coefficients[0], self.coefficients[1]]),
531 2 => QuadraticTrinomial::new([N::zero(), N::zero(), self.coefficients[0]]),
532 _ => QuadraticTrinomial::zero(),
533 }
534 }
535}
536
537impl<N: Zero + Copy> ShrAssign<u32> for QuadraticTrinomial<N> {
538 fn shr_assign(&mut self, rhs: u32) {
539 match rhs {
540 0 => {}
541 1 => {
542 self.coefficients[2] = self.coefficients[1];
543 self.coefficients[1] = self.coefficients[0];
544 self.coefficients[0] = N::zero();
545 }
546 2 => {
547 self.coefficients[2] = self.coefficients[0];
548 self.coefficients[1] = N::zero();
549 self.coefficients[0] = N::zero();
550 }
551 _ => {
552 self.coefficients[0] = N::zero();
553 self.coefficients[1] = N::zero();
554 self.coefficients[2] = N::zero();
555 }
556 }
557 }
558}
559
560#[cfg(test)]
561mod test {
562 use crate::{Derivable, Evaluable, QuadraticTrinomial, Roots, SizedPolynomial};
563 use num::Complex;
564
565 #[test]
566 fn test_eval() {
567 let a = QuadraticTrinomial::new([5, 0, 0]);
568 assert_eq!(125, a.eval(5));
569 }
570 #[test]
602 fn test_shr_pos() {
603 let a = QuadraticTrinomial::new([1, 0, 0]);
604 let c = QuadraticTrinomial::new([0, 0, 1]);
605 assert_eq!(c, a >> 2);
606 }
607
608 #[test]
609 fn test_shr_assign_pos() {
610 let mut a = QuadraticTrinomial::new([1, 0, 0]);
611 let c = QuadraticTrinomial::new([0, 0, 1]);
612 a >>= 2;
613 assert_eq!(c, a);
614 }
615
616 #[test]
632 fn test_shr_to_zero() {
633 let a = QuadraticTrinomial::new([1, 2, 3]);
634 assert_eq!(QuadraticTrinomial::zero(), a >> 5);
635 }
636
637 #[test]
638 fn test_shr_assign_to_zero() {
639 let mut a = QuadraticTrinomial::new([1, 2, 3]);
640 a >>= 5;
641 assert_eq!(QuadraticTrinomial::zero(), a);
642 }
643
644 #[test]
645 fn test_derivative_of_zero() {
646 let a: QuadraticTrinomial<i32> = QuadraticTrinomial::zero();
647 assert_eq!(QuadraticTrinomial::zero(), a.derivative());
648 }
649
650 #[test]
651 fn test_derivative_of_degree_zero() {
652 let a = QuadraticTrinomial::new([0, 0, 1]);
653 assert_eq!(QuadraticTrinomial::zero(), a.derivative());
654 }
655
656 #[test]
657 fn test_derivative() {
658 let a = QuadraticTrinomial::new([1, 2, 3]);
659 assert_eq!(QuadraticTrinomial::new([0, 2, 2]), a.derivative());
660 }
661
662 #[test]
663 fn test_roots_pos() {
664 let a = QuadraticTrinomial::new([1, 4, 4]);
665 let c = Roots::TwoRealRoots(-2i16, -2i16);
666 assert_eq!(c, a.roots());
667 }
668
669 #[test]
670 fn test_complex_roots_neg() {
671 let a = QuadraticTrinomial::new([1, 0, 4]);
672 let c = Roots::TwoComplexRoots(Complex::new(0, 2i16), Complex::new(0, -2i16));
673 assert_eq!(c, a.roots());
674 }
675}