conspire/math/integrate/field/adaptive/
mod.rs

1#[cfg(test)]
2mod test;
3
4use super::{
5    Integrable,
6    hermite::{HermiteSegment, hermite_at, interpolate_hermite},
7    reconstruct_or_err,
8    rkmk::{rkmk_dae_stage_slopes_into, rkmk_stage_slopes_into, weight},
9};
10use crate::math::{
11    Derivative, Differentiable, Quantity, Scalar, Tensor, TensorVec,
12    integrate::{ButcherTableau, EmbeddedTableau, IntegrationError, Times},
13    optimize::{EqualityConstraint, FirstOrderRootFinding, SecondOrderOptimization},
14    sparse::SparseSolver,
15};
16use std::ops::Mul;
17
18const DT_CUT: Scalar = 0.2;
19
20/// Adaptive [`super::rkmk_dae_step`]: embedded local-error control from the
21/// tableau's `D` weights over the span `[time[0], time[last]]`, with the same
22/// controller as [`integrate_rkmk_adaptive`]. A rejected step costs no
23/// endpoint constraint solve. Returns the accepted times, the state history,
24/// and the matching algebraic history.
25///
26/// Dense output follows the convention of the flat DAE loop: `time` of length
27/// two supplies only the span and the accepted steps are reported, while a
28/// longer `time` is a list of requested report times. Each of those is served by
29/// the geodesic [`HermiteSegment`] of the accepted step containing it, so the
30/// reported state is on the manifold at every requested time and not only at the
31/// accepted ones; the algebraic unknown is then re-solved from its constraint
32/// there, warm-started along the grid. Building the segments costs one extra
33/// rate evaluation per accepted step, so it is skipped when not requested.
34#[allow(clippy::type_complexity)]
35pub fn integrate_rkmk_dae_adaptive<Field, Tab, Z, U, V, T>(
36    mut rate: impl FnMut(
37        Quantity<T>,
38        &Field::Point,
39        &Z,
40    ) -> Result<Derivative<Field::Increment, T>, String>,
41    mut solve: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<Z, String>,
42    time: &[Quantity<T>],
43    initial_condition: (Field::Point, Z),
44    abs_tol: Scalar,
45    rel_tol: Scalar,
46) -> Result<(Times<T>, U, V), IntegrationError>
47where
48    Field: Integrable,
49    Tab: EmbeddedTableau,
50    Field::Point: Clone,
51    Field::Increment: Clone + Differentiable<T>,
52    Z: Clone,
53    T: Copy,
54    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
55    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
56    U: TensorVec<Item = Field::Point>,
57    V: TensorVec<Item = Z>,
58{
59    if time.len() < 2 {
60        return Err(IntegrationError::LengthTimeLessThanTwo);
61    }
62    let t_0 = time[0];
63    let t_f = time[time.len() - 1];
64    if t_0 >= t_f {
65        return Err(IntegrationError::InitialTimeNotLessThanFinalTime);
66    }
67    let exponent = 1.0 / Tab::ORDER;
68    let dt_min = (t_f - t_0) * 1e-10;
69    let mut t = t_0;
70    let mut dt = t_f - t_0;
71    let dense = time.len() > 2;
72    let (mut point, mut z) = initial_condition;
73    let z_0 = z.clone();
74    let mut points = U::new();
75    let mut algebraics = V::new();
76    let mut times = Times::new();
77    let mut slopes = Vec::new();
78    let mut segments = Vec::new();
79    let mut carry: Option<Derivative<Field::Increment, T>> = None;
80    points.push(point.clone());
81    algebraics.push(z.clone());
82    times.push(t_0);
83    while t_f - t > dt_min {
84        dt = dt.min(t_f - t);
85        let stage = rkmk_dae_stage_slopes_into::<Field, Tab, Z, T>(
86            &mut rate,
87            &mut solve,
88            &point,
89            &z,
90            t,
91            dt,
92            &mut slopes,
93            carry.as_ref(),
94        );
95        let (z_stage, next_carry) = match stage {
96            Ok(stage) => stage,
97            Err(error) => {
98                if dt <= dt_min {
99                    return Err(error);
100                }
101                dt *= DT_CUT;
102                continue;
103            }
104        };
105        let sigma = weight(&slopes, Tab::B);
106        let trial = match reconstruct_or_err::<Field>(&point, &sigma) {
107            Ok(trial) => trial,
108            Err(error) => {
109                if dt <= dt_min {
110                    return Err(error);
111                }
112                dt *= DT_CUT;
113                continue;
114            }
115        };
116        let error = weight(&slopes, Tab::D).norm().value().abs();
117        let tolerance = abs_tol + rel_tol * trial.norm().value();
118        let accept = error <= tolerance;
119        if accept {
120            let t_previous = t;
121            let t_next = t + dt;
122            match solve(t_next, &trial, &z_stage) {
123                Ok(z_next) => {
124                    t = t_next;
125                    z = z_next;
126                    carry = next_carry;
127                    if dense {
128                        let slope_1 = Field::dexpinv(&sigma, &rate(t, &trial, &z)? * dt);
129                        segments.push(HermiteSegment::new(
130                            t_previous,
131                            dt,
132                            point.clone(),
133                            sigma,
134                            slopes[0].clone(),
135                            slope_1,
136                        ));
137                    }
138                    point = trial;
139                    points.push(point.clone());
140                    algebraics.push(z.clone());
141                    times.push(t);
142                }
143                Err(error) => {
144                    if dt <= dt_min {
145                        return Err(IntegrationError::from(error));
146                    }
147                    dt *= DT_CUT;
148                    continue;
149                }
150            }
151        }
152        let scale = if error > 0.0 {
153            (0.9 * (tolerance / error).powf(exponent)).clamp(DT_CUT, 5.0)
154        } else {
155            5.0
156        };
157        dt *= scale;
158        if !accept && dt <= dt_min {
159            return Err(IntegrationError::from(
160                "the adaptive RKMK-DAE step fell below the floor".to_string(),
161            ));
162        }
163    }
164    if dense {
165        let mut points = U::new();
166        let mut algebraics = V::new();
167        let mut guess = z_0;
168        for time_k in time {
169            let point = hermite_at::<Field, T>(&segments, *time_k)?;
170            guess = solve(*time_k, &point, &guess)?;
171            points.push(point);
172            algebraics.push(guess.clone());
173        }
174        Ok((Times::from(time), points, algebraics))
175    } else {
176        Ok((times, points, algebraics))
177    }
178}
179
180/// [`integrate_rkmk_dae_adaptive`] with the algebraic unknown resolved by
181/// first-order root-finding at every stage abscissa, built from
182/// `function`/`jacobian`/`solver` the same way
183/// [`super::rkmk_dae_step_first_order_root`] builds it for a single step.
184#[allow(clippy::too_many_arguments, clippy::type_complexity)]
185pub fn integrate_rkmk_dae_adaptive_first_order_root<Field, Tab, F, J, Z, U, V, T>(
186    rate: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<Derivative<Field::Increment, T>, String>,
187    mut function: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<F, String>,
188    mut jacobian: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<J, String>,
189    solver: &impl FirstOrderRootFinding<F, J, Z>,
190    time: &[Quantity<T>],
191    initial_condition: (Field::Point, Z),
192    abs_tol: Scalar,
193    rel_tol: Scalar,
194    mut equality_constraint: impl FnMut(Quantity<T>) -> EqualityConstraint,
195) -> Result<(Times<T>, U, V), IntegrationError>
196where
197    Field: Integrable,
198    Tab: EmbeddedTableau,
199    Field::Point: Clone,
200    Field::Increment: Clone + Differentiable<T>,
201    Z: Clone,
202    T: Copy,
203    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
204    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
205    U: TensorVec<Item = Field::Point>,
206    V: TensorVec<Item = Z>,
207{
208    let solve = |t: Quantity<T>, point: &Field::Point, z_guess: &Z| -> Result<Z, String> {
209        Ok(solver.root(
210            |z| function(t, point, z),
211            |z| jacobian(t, point, z),
212            z_guess.clone(),
213            equality_constraint(t),
214            None,
215        )?)
216    };
217    integrate_rkmk_dae_adaptive::<Field, Tab, Z, U, V, T>(
218        rate,
219        solve,
220        time,
221        initial_condition,
222        abs_tol,
223        rel_tol,
224    )
225}
226
227/// [`integrate_rkmk_dae_adaptive`] with the algebraic unknown resolved by
228/// second-order minimization at every stage abscissa, built from
229/// `function`/`jacobian`/`hessian`/`solver` the same way
230/// [`super::rkmk_dae_step_second_order_minimize`] builds it for a single step.
231#[allow(clippy::too_many_arguments, clippy::type_complexity)]
232pub fn integrate_rkmk_dae_adaptive_second_order_minimize<Field, Tab, F, J, H, Z, U, V, T>(
233    rate: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<Derivative<Field::Increment, T>, String>,
234    mut function: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<F, String>,
235    mut jacobian: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<J, String>,
236    mut hessian: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<H, String>,
237    solver: &impl SecondOrderOptimization<F, J, H, Z>,
238    time: &[Quantity<T>],
239    initial_condition: (Field::Point, Z),
240    abs_tol: Scalar,
241    rel_tol: Scalar,
242    mut equality_constraint: impl FnMut(Quantity<T>) -> EqualityConstraint,
243    sparse: Option<SparseSolver>,
244) -> Result<(Times<T>, U, V), IntegrationError>
245where
246    Field: Integrable,
247    Tab: EmbeddedTableau,
248    Field::Point: Clone,
249    Field::Increment: Clone + Differentiable<T>,
250    Z: Clone,
251    T: Copy,
252    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
253    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
254    U: TensorVec<Item = Field::Point>,
255    V: TensorVec<Item = Z>,
256{
257    let solve = |t: Quantity<T>, point: &Field::Point, z_guess: &Z| -> Result<Z, String> {
258        Ok(solver.minimize(
259            |z| function(t, point, z),
260            |z| jacobian(t, point, z),
261            |z| hessian(t, point, z),
262            z_guess.clone(),
263            equality_constraint(t),
264            sparse.clone(),
265        )?)
266    };
267    integrate_rkmk_dae_adaptive::<Field, Tab, Z, U, V, T>(
268        rate,
269        solve,
270        time,
271        initial_condition,
272        abs_tol,
273        rel_tol,
274    )
275}
276
277/// Runge–Kutta–Munthe-Kaas: a fixed-step [`ButcherTableau`] run in the field's
278/// Lie algebra, with the [`Integrable::dexpinv`] correction per stage and a
279/// single [`Integrable::reconstruct`] per step. Reduces to the plain tableau
280/// on a flat field. See [`super::rkmk_step`] for the allocation-free single
281/// step.
282pub fn integrate_rkmk<Field, Tab, U, T>(
283    mut rate: impl FnMut(Quantity<T>, &Field::Point) -> Result<Derivative<Field::Increment, T>, String>,
284    time: &[Quantity<T>],
285    initial_condition: Field::Point,
286) -> Result<(Times<T>, U), IntegrationError>
287where
288    Field: Integrable,
289    Tab: ButcherTableau,
290    Field::Point: Clone,
291    Field::Increment: Clone + Differentiable<T>,
292    T: Copy,
293    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
294    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
295    U: TensorVec<Item = Field::Point>,
296{
297    let mut point = initial_condition;
298    let mut points = U::new();
299    let mut times = Times::new();
300    let mut scratch = Vec::new();
301    let mut carry: Option<Derivative<Field::Increment, T>> = None;
302    points.push(point.clone());
303    times.push(time[0]);
304    for step in time.windows(2) {
305        carry = rkmk_stage_slopes_into::<Field, Tab, T>(
306            &mut rate,
307            &point,
308            step[0],
309            step[1] - step[0],
310            &mut scratch,
311            carry.as_ref(),
312        )?;
313        point = reconstruct_or_err::<Field>(&point, &weight(&scratch, Tab::B))?;
314        points.push(point.clone());
315        times.push(step[1]);
316    }
317    Ok((times, points))
318}
319
320/// Adaptive RKMK: [`integrate_rkmk`] with embedded local-error control from the
321/// tableau's `D` weights. The step is grown or shrunk by `0.9 (tol / e)^{1/p}`
322/// (clamped to `[DT_CUT, 5]`), and a step whose error `e` exceeds
323/// `abs_tol + rel_tol ‖x_{n+1}‖` is rejected. A rate-evaluation failure is
324/// retried with `dt *= DT_CUT`, the same as a rejected accuracy estimate.
325///
326/// Dense output follows the convention of [`integrate_rkmk_dae_adaptive`]: `time`
327/// of length two supplies only the span and the accepted steps are reported,
328/// while a longer `time` is a list of requested report times, each served by the
329/// geodesic [`HermiteSegment`] of the accepted step containing it. Building the
330/// segments costs one extra rate evaluation per accepted step, so it is skipped
331/// when not requested.
332pub fn integrate_rkmk_adaptive<Field, Tab, U, T>(
333    mut rate: impl FnMut(Quantity<T>, &Field::Point) -> Result<Derivative<Field::Increment, T>, String>,
334    time: &[Quantity<T>],
335    initial_condition: Field::Point,
336    abs_tol: Scalar,
337    rel_tol: Scalar,
338) -> Result<(Times<T>, U), IntegrationError>
339where
340    Field: Integrable,
341    Tab: EmbeddedTableau,
342    Field::Point: Clone,
343    Field::Increment: Clone + Differentiable<T>,
344    T: Copy,
345    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
346    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
347    U: TensorVec<Item = Field::Point>,
348{
349    if time.len() < 2 {
350        return Err(IntegrationError::LengthTimeLessThanTwo);
351    }
352    let t_0 = time[0];
353    let t_f = time[time.len() - 1];
354    if t_0 >= t_f {
355        return Err(IntegrationError::InitialTimeNotLessThanFinalTime);
356    }
357    let exponent = 1.0 / Tab::ORDER;
358    let dt_min = (t_f - t_0) * 1e-10;
359    let mut t = t_0;
360    let mut dt = t_f - t_0;
361    let dense = time.len() > 2;
362    let mut point = initial_condition;
363    let mut points = U::new();
364    let mut times = Times::new();
365    let mut slopes = Vec::new();
366    let mut segments = Vec::new();
367    let mut carry: Option<Derivative<Field::Increment, T>> = None;
368    points.push(point.clone());
369    times.push(t_0);
370    while t_f - t > dt_min {
371        dt = dt.min(t_f - t);
372        let stage = rkmk_stage_slopes_into::<Field, Tab, T>(
373            &mut rate,
374            &point,
375            t,
376            dt,
377            &mut slopes,
378            carry.as_ref(),
379        );
380        let next_carry = match stage {
381            Ok(next_carry) => next_carry,
382            Err(error) => {
383                if dt <= dt_min {
384                    return Err(error);
385                }
386                dt *= DT_CUT;
387                continue;
388            }
389        };
390        let sigma = weight(&slopes, Tab::B);
391        let trial = match reconstruct_or_err::<Field>(&point, &sigma) {
392            Ok(trial) => trial,
393            Err(error) => {
394                if dt <= dt_min {
395                    return Err(error);
396                }
397                dt *= DT_CUT;
398                continue;
399            }
400        };
401        let error = weight(&slopes, Tab::D).norm().value().abs();
402        let tolerance = abs_tol + rel_tol * trial.norm().value();
403        let accept = error <= tolerance;
404        if accept {
405            let t_previous = t;
406            let t_next = t + dt;
407            let slope_1 = if dense {
408                match rate(t_next, &trial) {
409                    Ok(raw) => Some(Field::dexpinv(&sigma, &raw * dt)),
410                    Err(error) => {
411                        if dt <= dt_min {
412                            return Err(IntegrationError::from(error));
413                        }
414                        dt *= DT_CUT;
415                        continue;
416                    }
417                }
418            } else {
419                None
420            };
421            t = t_next;
422            carry = next_carry;
423            if let Some(slope_1) = slope_1 {
424                segments.push(HermiteSegment::new(
425                    t_previous,
426                    dt,
427                    point.clone(),
428                    sigma,
429                    slopes[0].clone(),
430                    slope_1,
431                ));
432            }
433            point = trial;
434            points.push(point.clone());
435            times.push(t);
436        }
437        let scale = if error > 0.0 {
438            (0.9 * (tolerance / error).powf(exponent)).clamp(DT_CUT, 5.0)
439        } else {
440            5.0
441        };
442        dt *= scale;
443        if !accept && dt <= dt_min {
444            return Err(IntegrationError::from(
445                "the adaptive RKMK step fell below the floor".to_string(),
446            ));
447        }
448    }
449    if dense {
450        Ok((
451            Times::from(time),
452            interpolate_hermite::<Field, U, T>(&segments, time)?,
453        ))
454    } else {
455        Ok((times, points))
456    }
457}
conspire/math/integrate/field/adaptive/mod.rs

conspire/math/integrate/field/adaptive/
mod.rs
