bacon_sci/interp/
spline.rs

1/* This file is part of bacon.
2 * Copyright (c) Wyatt Campbell.
3 *
4 * See repository LICENSE for information.
5 */
6
7use crate::polynomial::Polynomial;
8use nalgebra::{ComplexField, RealField};
9use num_traits::FromPrimitive;
10
11#[derive(Debug, Clone)]
12pub struct CubicSpline<N: ComplexField + FromPrimitive + Copy>
13where
14    <N as ComplexField>::RealField: FromPrimitive + Copy,
15{
16    cubics: Vec<Polynomial<N>>,
17    ranges: Vec<(N::RealField, N::RealField)>,
18}
19
20impl<N: ComplexField + FromPrimitive + Copy> CubicSpline<N>
21where
22    <N as ComplexField>::RealField: FromPrimitive + Copy,
23{
24    pub fn evaluate(&self, x: N::RealField) -> Result<N, String> {
25        if self.cubics.is_empty() {
26            return Err("CubicSpline evaluate: Empty spline".to_owned());
27        }
28
29        for (ind, range) in self.ranges.iter().enumerate() {
30            if x >= range.0 && x <= range.1 {
31                return Ok(self.cubics[ind].evaluate(N::from_real(x)));
32            }
33        }
34
35        Err(format!("CubicSpline evaluate: {} outside of range", x))
36    }
37
38    pub fn evaluate_derivative(&self, x: N::RealField) -> Result<(N, N), String> {
39        if self.cubics.is_empty() {
40            return Err("CubicSpline evaluate: Empty spline".to_owned());
41        }
42
43        for (ind, range) in self.ranges.iter().enumerate() {
44            if x >= range.0 && x <= range.1 {
45                return Ok(self.cubics[ind].evaluate_derivative(N::from_real(x)));
46            }
47        }
48
49        Err(format!("CubicSpline evaluate: {} outside of range", x))
50    }
51}
52
53/// Create a free cubic spline interpolating the given points.
54///
55/// Given a set of ordered points, produce a piecewise function of
56/// cubic polynomials that interpolate the points given the second derivative
57/// of the piecewise function at the end points is zero and the piecewise function
58/// is smooth.
59///
60/// # Params
61/// `xs` x points. Must be real because cubic splines keep track of ranges within
62/// which it interpolates. Must be sorted.
63///
64/// `ys` y points. Can be complex. ys\[i\] must match with xs\[i\].
65///
66/// `tol` the tolerance of the polynomials
67///
68/// # Examples
69/// ```
70/// use bacon_sci::interp::spline_free;
71/// fn example() {
72///     let xs: Vec<_> = (0..=10).map(|x| x as f64).collect();
73///     let ys: Vec<_> = xs.iter().map(|x| x.exp()).collect();
74///
75///     let spline = spline_free(&xs, &ys, 1e-8).unwrap();
76///     for i in 0..1000 {
77///         let i = i as f64 * 0.001;
78///         assert!((spline.evaluate(i).unwrap() - i.exp()).abs() / i.exp() < 0.25);
79///     }
80/// }
81/// ```
82pub fn spline_free<N: ComplexField + FromPrimitive + Copy>(
83    xs: &[N::RealField],
84    ys: &[N],
85    tol: N::RealField,
86) -> Result<CubicSpline<N>, String>
87where
88    <N as ComplexField>::RealField: FromPrimitive + Copy,
89{
90    if xs.len() != ys.len() {
91        return Err("spline_free: xs and ys must be same length".to_owned());
92    }
93    if xs.len() < 2 {
94        return Err("spline_free: need at least two points".to_owned());
95    }
96
97    let mut hs = Vec::with_capacity(xs.len() - 1);
98    for i in 0..xs.len() - 1 {
99        hs.push(xs[i + 1] - xs[i]);
100    }
101
102    if hs.iter().any(|h| !h.is_sign_positive()) {
103        return Err("spline_free: xs must be sorted".to_owned());
104    }
105
106    let third = N::from_f64(1.0 / 3.0).unwrap();
107    let two = N::from_i32(2).unwrap();
108    let three = N::from_i32(3).unwrap();
109
110    let mut alphas = Vec::with_capacity(xs.len() - 1);
111    alphas.push(N::zero());
112    for i in 1..xs.len() - 1 {
113        alphas.push(
114            (three / N::from_real(hs[i])) * (ys[i + 1] - ys[i])
115                - (three / N::from_real(hs[i - 1])) * (ys[i] - ys[i - 1]),
116        );
117    }
118
119    let mut l = Vec::with_capacity(xs.len() - 1);
120    let mut mu = Vec::with_capacity(xs.len() - 1);
121    let mut z = Vec::with_capacity(xs.len() - 1);
122
123    l.push(N::one());
124    mu.push(N::zero());
125    z.push(N::zero());
126
127    for i in 1..xs.len() - 1 {
128        l.push(two * N::from_real(xs[i + 1] - xs[i - 1]) - N::from_real(hs[i - 1]) * mu[i - 1]);
129        mu.push(N::from_real(hs[i]) / l[i]);
130        z.push((alphas[i] - N::from_real(hs[i - 1]) * z[i - 1]) / l[i]);
131    }
132
133    l.push(N::one());
134    z.push(N::zero());
135
136    let mut c_coefficient = vec![N::zero(); xs.len()];
137    let mut b_coefficient = vec![N::zero(); xs.len()];
138    let mut d_coefficient = vec![N::zero(); xs.len()];
139    for i in (0..xs.len() - 1).rev() {
140        c_coefficient[i] = z[i] - mu[i] * c_coefficient[i + 1];
141        b_coefficient[i] = (ys[i + 1] - ys[i]) / N::from_real(hs[i])
142            - N::from_real(hs[i]) * (c_coefficient[i + 1] + two * c_coefficient[i]) * third;
143        d_coefficient[i] =
144            (c_coefficient[i + 1] - c_coefficient[i]) / (three * N::from_real(hs[i]));
145    }
146
147    let mut polynomials = Vec::with_capacity(xs.len() - 1);
148    let mut ranges = Vec::with_capacity(xs.len() - 1);
149
150    for i in 0..xs.len() - 1 {
151        // Horner's method to build polynomial
152        let term = polynomial![N::one(), N::from_real(-xs[i])];
153        let mut poly = &term * d_coefficient[i];
154        poly.set_tolerance(tol)?;
155        poly += c_coefficient[i];
156        poly *= &term;
157        poly += b_coefficient[i];
158        poly *= term;
159        poly += ys[i];
160        polynomials.push(poly);
161        ranges.push((xs[i], xs[i + 1]));
162    }
163
164    Ok(CubicSpline {
165        cubics: polynomials,
166        ranges,
167    })
168}
169
170/// Create a clamped cubic spline interpolating the given points.
171///
172/// Given a set of ordered points, produce a piecewise function of
173/// cubic polynomials that interpolate the points given the first derivative
174/// of the piecewise function at the end points is the same as the given values
175/// and the piecewise function is smooth.
176///
177/// # Params
178/// `xs` x points. Must be real because cubic splines keep track of ranges within
179/// which it interpolates. Must be sorted.
180///
181/// `ys` y points. Can be complex. ys\[i\] must match with xs\[i\].
182///
183/// `(f_0, f_n)` The derivative values at the end points.
184///
185/// `tol` the tolerance of the polynomials
186///
187/// # Examples
188/// ```
189/// use bacon_sci::interp::spline_clamped;
190/// fn example() {
191///     let xs: Vec<_> = (0..=10).map(|x| x as f64).collect();
192///     let ys: Vec<_> = xs.iter().map(|x| x.exp()).collect();
193///
194///     let spline = spline_clamped(&xs, &ys, (1.0, (10.0f64).exp()), 1e-8).unwrap();
195///     for i in 0..1000 {
196///         let i = i as f64 * 0.001;
197///         assert!((spline.evaluate(i).unwrap() - i.exp()).abs() / i.exp() < 0.25);
198///     }
199/// }
200/// ```
201pub fn spline_clamped<N: ComplexField + FromPrimitive + Copy>(
202    xs: &[N::RealField],
203    ys: &[N],
204    (f_0, f_n): (N, N),
205    tol: N::RealField,
206) -> Result<CubicSpline<N>, String>
207where
208    <N as ComplexField>::RealField: FromPrimitive + Copy,
209{
210    if xs.len() != ys.len() {
211        return Err("spline_clamped: xs and ys must be same length".to_owned());
212    }
213    if xs.len() < 2 {
214        return Err("spline_clamped: need at least two points".to_owned());
215    }
216
217    let mut hs = Vec::with_capacity(xs.len() - 1);
218    for i in 0..xs.len() - 1 {
219        hs.push(xs[i + 1] - xs[i]);
220    }
221
222    if hs.iter().any(|h| !h.is_sign_positive()) {
223        return Err("spline_clamped: xs must be sorted".to_owned());
224    }
225
226    let third = N::from_f64(1.0 / 3.0).unwrap();
227    let half = N::from_f64(0.5).unwrap();
228    let two = N::from_i32(2).unwrap();
229    let three = N::from_i32(3).unwrap();
230
231    let mut alphas = vec![N::zero(); xs.len()];
232    alphas[0] = three * ((ys[1] - ys[0]) / N::from_real(hs[0]) - f_0);
233    alphas[xs.len() - 1] =
234        three * (f_n - (ys[xs.len() - 1] - ys[xs.len() - 2]) / N::from_real(hs[xs.len() - 2]));
235
236    for i in 1..xs.len() - 1 {
237        alphas[i] = three
238            * ((ys[i + 1] - ys[i]) / N::from_real(hs[i])
239                - (ys[i] - ys[i - 1]) / N::from_real(hs[i - 1]));
240    }
241
242    let mut l = Vec::with_capacity(xs.len() - 1);
243    let mut mu = Vec::with_capacity(xs.len() - 1);
244    let mut z = Vec::with_capacity(xs.len() - 1);
245
246    l.push(two * N::from_real(hs[0]));
247    mu.push(half);
248    z.push(alphas[0] / l[0]);
249
250    for i in 1..xs.len() - 1 {
251        l.push(two * N::from_real(xs[i + 1] - xs[i - 1]) - N::from_real(hs[i - 1]) * mu[i - 1]);
252        mu.push(N::from_real(hs[i]) / l[i]);
253        z.push((alphas[i] - N::from_real(hs[i - 1]) * z[i - 1]) / l[i]);
254    }
255
256    l.push(N::from_real(hs[xs.len() - 2]) * (two - mu[xs.len() - 2]));
257    z.push(
258        (alphas[xs.len() - 1] - N::from_real(hs[xs.len() - 2]) * z[xs.len() - 2]) / l[xs.len() - 1],
259    );
260
261    let mut b_coefficient = vec![N::zero(); xs.len()];
262    let mut c_coefficient = vec![N::zero(); xs.len()];
263    let mut d_coefficient = vec![N::zero(); xs.len()];
264
265    c_coefficient[xs.len() - 1] = z[xs.len() - 1];
266
267    for i in (0..xs.len() - 1).rev() {
268        c_coefficient[i] = z[i] - mu[i] * c_coefficient[i + 1];
269        b_coefficient[i] = (ys[i + 1] - ys[i]) / N::from_real(hs[i])
270            - N::from_real(hs[i]) * third * (c_coefficient[i + 1] + two * c_coefficient[i]);
271        d_coefficient[i] =
272            (c_coefficient[i + 1] - c_coefficient[i]) / (three * N::from_real(hs[i]));
273    }
274
275    let mut polynomials = Vec::with_capacity(xs.len() - 1);
276    let mut ranges = Vec::with_capacity(xs.len() - 1);
277
278    for i in 0..xs.len() - 1 {
279        // Horner's method to build polynomial
280        let term = polynomial![N::one(), N::from_real(-xs[i])];
281        let mut poly = &term * d_coefficient[i];
282        poly.set_tolerance(tol)?;
283        poly += c_coefficient[i];
284        poly *= &term;
285        poly += b_coefficient[i];
286        poly *= term;
287        poly += ys[i];
288        polynomials.push(poly);
289        ranges.push((xs[i], xs[i + 1]));
290    }
291
292    Ok(CubicSpline {
293        cubics: polynomials,
294        ranges,
295    })
296}
bacon_sci/interp/spline.rs

bacon_sci/interp/
spline.rs
