poly_cool/quadratic.rs
1use arrayvec::ArrayVec;
2
3use crate::Quadratic;
4
5impl Quadratic {
6 /// This is like [`Quadratic::eval`] but faster.
7 ///
8 /// It would be nice if we could just make `eval` like this, but I couldn't
9 /// figure out how, given the lack of specialization.
10 #[doc(hidden)]
11 pub fn eval_opt(&self, x: f64) -> f64 {
12 let [c0, c1, c2] = self.coeffs;
13 c0 + c1 * x + c2 * x * x
14 }
15
16 /// Returns the roots of this quadratic, in increasing order.
17 ///
18 /// Double-roots are only counted once.
19 pub fn roots(&self) -> ArrayVec<f64, 2> {
20 let &[c, b, a] = self.coeffs();
21 let disc = b * b - 4.0 * a * c;
22 if disc.is_finite() {
23 let mut ret = ArrayVec::new();
24 let mut push = |r: f64| {
25 if r.is_finite() {
26 ret.push(r)
27 }
28 };
29 if disc > 0.0 {
30 let q = -0.5 * (b + disc.sqrt().copysign(b));
31 let r0 = q / a;
32 let r1 = c / q;
33 push(r0.min(r1));
34 push(r0.max(r1));
35 } else if disc == 0.0 {
36 let root = -0.5 * b / a;
37 if root.is_finite() {
38 push(root);
39 } else if c == 0.0 {
40 // This is kurbo's behavior: the intention is that if the
41 // whole thing is zero, return zero as a single root. I'm
42 // not sure I love it.
43 //
44 // Bear in mind that this branch is not *only* for the
45 // identically zero case: if a == c == 0.0 and b * b
46 // underflows then we will end up here. In that case,
47 // zero is the only root.
48 push(0.0);
49 }
50 } else {
51 // No roots.
52 }
53 ret
54 } else {
55 // At least one of the coefficients was too large and triggered
56 // overflow.
57 //
58 // The exponent of f64 maxes out at 1023, so scaling down by
59 // 2^{-512} is enough to ensure that squaring doesn't overflow. We
60 // do an extra factor of 2^{-3} for some wiggle room. This can't
61 // completely destroy all the coefficients: because of the overflow,
62 // we know that at least one of them was big.
63 let scale = 2.0f64.powi(-515);
64 // If we're infinite, just give up. (Otherwise, we'd stack overflow
65 // by repeatedly trying to rescale.)
66 if self.is_finite() {
67 (*self * scale).roots()
68 } else {
69 ArrayVec::new()
70 }
71 }
72 }
73
74 /// Returns the two distinct roots of this quadratic, but only if the
75 /// discriminant is positive.
76 pub fn positive_discriminant_roots(&self) -> Option<(f64, f64)> {
77 let &[c, b, a] = self.coeffs();
78 let disc = b * b - 4.0 * a * c;
79 if disc.is_finite() {
80 if disc > 0.0 {
81 let q = -0.5 * (b + disc.sqrt().copysign(b));
82 let r0 = q / a;
83 let r1 = c / q;
84 Some((r0.min(r1), r0.max(r1)))
85 } else {
86 None
87 }
88 } else {
89 self.positive_discriminant_roots_scaled()
90 }
91 }
92
93 #[cold]
94 fn positive_discriminant_roots_scaled(&self) -> Option<(f64, f64)> {
95 if self.is_finite() {
96 let scale = 2.0f64.powi(-515);
97 (*self * scale).positive_discriminant_roots()
98 } else {
99 None
100 }
101 }
102}
103
104#[cfg(test)]
105mod tests {
106 #[test]
107 fn root_evaluation() {
108 arbtest::arbtest(|u| {
109 let q = crate::arbitrary::quadratic(u)?;
110 // Arbitrary quadratics can have coefficients with wild magnitudes,
111 // so we need to adjust our error expectations accordingly.
112 let magnitude = q.magnitude().max(1.0);
113
114 for r in q.roots() {
115 let y = q.eval(r);
116 // To evaluate the polynomial, we need to square r, so our error
117 // should be relative to the magnitude of r squared.
118 let r_magnitude = r.abs().max(1.0);
119 let threshold = r_magnitude * 1e-14 * r_magnitude * magnitude;
120 if y.is_finite() && threshold.is_finite() {
121 assert!(y.abs() <= threshold);
122 }
123 }
124 Ok(())
125 })
126 .budget_ms(5_000);
127 }
128}