poly_cool/poly.rs
1//! Polynomials of dynamic (run-time) degree.
2
3use crate::{Cubic, InputError, TerminationCondition, different_signs};
4
5/// A polynomial of dynamic degree.
6///
7/// It would be nice to have polynomials of type-level degree,
8/// but that's a bit awkward without const generic expressions
9/// (e.g. to express the type of the derivative). It could be
10/// done with `typenum` and `generic_array`...
11#[derive(Clone, Debug)]
12pub struct Poly {
13 /// Coefficients in increasing order of degree.
14 ///
15 /// For example, `coeffs[0]` is the constant term.
16 coeffs: Vec<f64>,
17}
18
19impl<'a> std::ops::Mul<&'a Poly> for &'a Poly {
20 type Output = Poly;
21
22 fn mul(self, rhs: &Poly) -> Poly {
23 let mut coeffs = vec![0.0; (self.coeffs.len() + rhs.coeffs.len()).saturating_sub(1)];
24
25 for (i, c) in self.coeffs.iter().enumerate() {
26 for (j, d) in rhs.coeffs.iter().enumerate() {
27 coeffs[i + j] += c * d;
28 }
29 }
30 Poly { coeffs }
31 }
32}
33
34impl std::ops::Mul<&Poly> for Poly {
35 type Output = Poly;
36
37 fn mul(self, rhs: &Poly) -> Poly {
38 (&self) * rhs
39 }
40}
41
42impl Poly {
43 /// Constructs a new polynomial from coefficients.
44 ///
45 /// The first coefficient provided will be the constant term, the second will
46 /// be the linear term, and so on.
47 pub fn new(coeffs: impl IntoIterator<Item = f64>) -> Self {
48 Poly {
49 coeffs: coeffs.into_iter().collect(),
50 }
51 }
52
53 fn is_finite(&self) -> bool {
54 self.coeffs.iter().all(|c| c.is_finite())
55 }
56
57 /// Returns the polynomial that's the derivative of this polynomial.
58 pub fn deriv(&self) -> Poly {
59 let mut coeffs = Vec::with_capacity(self.coeffs.len() - 1);
60 // If we're empty (meaning that we're the constant zero polynomial),
61 // this will just return the zero polynomial again: no need for a
62 // special case.
63 for (i, c) in self.coeffs.iter().enumerate().skip(1) {
64 coeffs.push(c * (i as f64));
65 }
66 Poly { coeffs }
67 }
68
69 /// Evaluates this polynomial at a point.
70 pub fn eval(&self, x: f64) -> f64 {
71 let mut ret = 0.0;
72 let mut x_pow = 1.0;
73 for &c in &self.coeffs {
74 ret += c * x_pow;
75 x_pow *= x;
76 }
77 ret
78 }
79
80 /// The degree of this polynomial.
81 ///
82 /// This function only looks at the *presence* of coefficients, not their
83 /// value. If you construct a polynomial with three coefficients, this
84 /// method will say that it has degree 2 even if all of those coefficients
85 /// are zero.
86 ///
87 /// A polynomial with no coefficients will give zero as its degree, as will
88 /// a polynomial with one coefficient.
89 pub fn degree(&self) -> usize {
90 self.coeffs.len().saturating_sub(1)
91 }
92
93 /// If this polynomial has degree 3 or less, converts it to a [cubic](crate::Cubic).
94 fn to_cubic(&self) -> Option<Cubic> {
95 if self.degree() <= 3 {
96 Some(Cubic {
97 c0: self.coeffs.first().copied().unwrap_or(0.0),
98 c1: self.coeffs.get(1).copied().unwrap_or(0.0),
99 c2: self.coeffs.get(2).copied().unwrap_or(0.0),
100 c3: self.coeffs.get(3).copied().unwrap_or(0.0),
101 })
102 } else {
103 None
104 }
105 }
106
107 fn one_root<Term: TerminationCondition>(
108 &self,
109 deriv: &Poly,
110 mut lower: f64,
111 mut upper: f64,
112 val_lower: f64,
113 val_upper: f64,
114 term: Term,
115 ) -> f64 {
116 if !val_lower.is_finite() || !val_upper.is_finite() || !deriv.is_finite() {
117 return f64::NAN;
118 }
119 debug_assert!(different_signs(val_lower, val_upper));
120
121 let mut x = lower + (upper - lower) / 2.0;
122 let mut val_x = self.eval(x);
123 let mut step = (upper - lower) / 2.0;
124
125 while x.is_finite() && !term.stop(step, val_x) {
126 let root_in_first_half = different_signs(val_lower, val_x);
127 if root_in_first_half {
128 upper = x;
129 } else {
130 lower = x;
131 }
132
133 let deriv_x = self.deriv().eval(x);
134 debug_assert!(deriv_x.is_finite());
135 debug_assert!(val_x.is_finite());
136
137 step = -val_x / deriv_x;
138 let mut new_x = x + step;
139
140 if new_x <= lower || new_x >= upper {
141 new_x = lower + (upper - lower) / 2.0;
142
143 if new_x == upper || new_x == lower {
144 // This should be rare, but it happens if they ask for more
145 // accuracy than is reasonable. For example, suppse (because
146 // of large coefficients) the output value jumps from -1.0
147 // to 1.0 between adjacent floats and they ask for an output
148 // error of smaller than 0.5. Then we'll eventually shrink
149 // the search interval to a pair of adjacent floats and hit
150 // this case.
151 return new_x;
152 }
153 }
154 step = new_x - x;
155 x = new_x;
156 val_x = self.eval(x);
157 }
158 x
159 }
160
161 /// Finds all the roots in an interval, using Yuksel's algorithm.
162 ///
163 /// This is a numerical, iterative method. It first constructs critical
164 /// points to find bracketing intervals (intervals `[x0, x1]` where
165 /// `self.eval(x0)` and `self.eval(x1)` have different signs). Then it uses
166 /// a kind of modified Newton method to find a root on each bracketing
167 /// interval. It has a few limitations:
168 ///
169 /// - if there is only a small interval where the polynomial changes sign,
170 /// it can miss roots. For example, when two roots are very close together
171 /// it can miss them both.
172 /// - run time is quadratic in the degree. However, it is often very fast
173 /// in practice for polynomials of low degree, especially if the interval
174 /// `[lower, upper]` contains few roots.
175 pub fn roots_between(&self, lower: f64, upper: f64, x_error: f64) -> Vec<f64> {
176 let mut ret = Vec::new();
177 let mut scratch = Vec::new();
178 self.roots_between_with_buffer(lower, upper, x_error, &mut ret, &mut scratch);
179 ret
180 }
181
182 /// Finds all the roots in an interval, using Yuksel's algorithm.
183 ///
184 /// See [`Poly::roots_between`] for more details. This method differs from that
185 /// one in that it performs fewer allocations: you provide an `out` buffer
186 /// for the result and a `scratch` buffer for intermediate computations.
187 pub fn roots_between_with_buffer(
188 &self,
189 lower: f64,
190 upper: f64,
191 x_error: f64,
192 out: &mut Vec<f64>,
193 scratch: &mut Vec<f64>,
194 ) {
195 out.clear();
196 scratch.clear();
197
198 if let Some(c) = self.to_cubic() {
199 out.extend(c.roots_between(lower, upper, x_error));
200 return;
201 }
202
203 let deriv = self.deriv();
204 deriv.roots_between_with_buffer(lower, upper, x_error, scratch, out);
205 scratch.push(upper);
206 out.clear();
207 let mut last = lower;
208 let mut last_val = self.eval(last);
209
210 // `scratch` now contains all the critical points (in increasing order)
211 // and the upper endpoint of the interval. If we throw away all the
212 // critical points that are outside of (lower, upper), the things
213 // remaining in `scratch` are the endpoints of the potential bracketing
214 // intervals of our polynomial. So by filtering out uninteresting
215 // critical points, this loop is iterating over potential bracketing
216 // intervals.
217 for &mut x in scratch {
218 if x > last && x <= upper {
219 let val = self.eval(x);
220 if different_signs(last_val, val) {
221 out.push(self.one_root(&deriv, last, x, last_val, val, InputError(x_error)));
222 }
223
224 last = x;
225 last_val = val;
226 }
227 }
228 }
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234
235 #[test]
236 fn smoke() {
237 let x_minus_1 = Poly::new([-1.0, 1.0]);
238 let x_minus_2 = Poly::new([-2.0, 1.0]);
239 let x_minus_3 = Poly::new([-3.0, 1.0]);
240 let x_minus_4 = Poly::new([-4.0, 1.0]);
241
242 let p = &x_minus_1 * &x_minus_2 * &x_minus_3 * &x_minus_4;
243
244 let roots = p.roots_between(0.0, 5.0, 1e-6);
245 assert_eq!(roots.len(), 4);
246 assert!((roots[0] - 1.0).abs() <= 1e-6);
247 assert!((roots[1] - 2.0).abs() <= 1e-6);
248 assert!((roots[2] - 3.0).abs() <= 1e-6);
249 assert!((roots[3] - 4.0).abs() <= 1e-6);
250 }
251}