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