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

1#[cfg(test)]
2mod test;
3
4use super::{Integrable, reconstruct_or_err};
5use crate::math::{
6    Derivative, Differentiable, Quantity, Scalar,
7    integrate::{ButcherTableau, IntegrationError},
8    optimize::{EqualityConstraint, FirstOrderRootFinding, SecondOrderOptimization},
9    sparse::SparseSolver,
10};
11use std::ops::{AddAssign, Mul};
12
13/// Fills `slopes` with one RKMK step's corrected stage slopes `k̃ᵢ` in the
14/// field's Lie algebra: per stage combine the earlier `k̃ⱼ` by row `Aᵢ`,
15/// `reconstruct` the stage point, evaluate the rate, scale by `dt`, apply
16/// [`Integrable::dexpinv`] at the accumulated algebra element. `slopes` is
17/// cleared first and reused, so a caller that steps in a loop allocates nothing.
18/// The caller weights the entries by `B` (the step) and, for an embedded pair,
19/// by `D` (the error estimate).
20///
21/// FSAL: `first_rate` seeds stage 0 (`C[0] == 0`) with a rate carried from the
22/// previous step, skipping that evaluation; when `Tab::FSAL`, the raw rate at
23/// the final stage (whose point is the step solution) is returned for the next
24/// step to seed with.
25pub(super) fn rkmk_stage_slopes_into<Field, Tab, T>(
26    rate: &mut impl FnMut(Quantity<T>, &Field::Point) -> Result<Derivative<Field::Increment, T>, String>,
27    point: &Field::Point,
28    t: Quantity<T>,
29    dt: Quantity<T>,
30    slopes: &mut Vec<Field::Increment>,
31    first_rate: Option<&Derivative<Field::Increment, T>>,
32) -> Result<Option<Derivative<Field::Increment, T>>, IntegrationError>
33where
34    Field: Integrable,
35    Tab: ButcherTableau,
36    Field::Point: Clone,
37    Field::Increment: Clone + Differentiable<T>,
38    T: Copy,
39    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
40    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
41{
42    slopes.clear();
43    slopes.reserve(Tab::STAGES);
44    let mut carry = None;
45    for i in 0..Tab::STAGES {
46        let sigma = if i == 0 {
47            None
48        } else {
49            let mut accumulated = slopes[0].clone() * Tab::A[i][0];
50            for (j, slope) in slopes.iter().enumerate().take(i).skip(1) {
51                accumulated += slope.clone() * Tab::A[i][j];
52            }
53            Some(accumulated)
54        };
55        let stage_point = match &sigma {
56            Some(sigma) => reconstruct_or_err::<Field>(point, sigma)?,
57            None => point.clone(),
58        };
59        let increment = match (i, first_rate) {
60            (0, Some(seed)) => seed * dt,
61            _ => {
62                let raw = rate(t + dt * Tab::C[i], &stage_point)?;
63                let increment = &raw * dt;
64                if Tab::FSAL && i + 1 == Tab::STAGES {
65                    carry = Some(raw);
66                }
67                increment
68            }
69        };
70        slopes.push(match &sigma {
71            Some(sigma) => Field::dexpinv(sigma, increment),
72            None => increment,
73        });
74    }
75    Ok(carry)
76}
77
78pub(super) fn weight<P>(slopes: &[P], weights: &[Scalar]) -> P
79where
80    P: Clone + Mul<Scalar, Output = P> + AddAssign,
81{
82    let mut sum = slopes[0].clone() * weights[0];
83    for (i, slope) in slopes.iter().enumerate().skip(1) {
84        sum += slope.clone() * weights[i];
85    }
86    sum
87}
88
89/// Advances `point` one RKMK step from `t` to `t + dt` with the `Tab` tableau —
90/// [`super::integrate_rkmk`] without the history, and with the stage-slope buffer
91/// `scratch` passed in so a stepping loop allocates nothing per step. `scratch`
92/// may start empty; its contents are overwritten.
93pub fn rkmk_step<Field, Tab, T>(
94    rate: &mut impl FnMut(Quantity<T>, &Field::Point) -> Result<Derivative<Field::Increment, T>, String>,
95    point: &Field::Point,
96    t: Quantity<T>,
97    dt: Quantity<T>,
98    scratch: &mut Vec<Field::Increment>,
99) -> Result<Field::Point, IntegrationError>
100where
101    Field: Integrable,
102    Tab: ButcherTableau,
103    Field::Point: Clone,
104    Field::Increment: Clone + Differentiable<T>,
105    T: Copy,
106    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
107    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
108{
109    rkmk_stage_slopes_into::<Field, Tab, T>(rate, point, t, dt, scratch, None)?;
110    reconstruct_or_err::<Field>(point, &weight(scratch, Tab::B))
111}
112
113/// One RKMK step of a semi-explicit DAE: the differential field advances on its
114/// manifold while the algebraic unknown is re-solved from its constraint at
115/// every stage abscissa.
116///
117/// [`rkmk_step`] freezes the drive across the whole window, which caps the
118/// coupling at first order however the two legs are ordered. Here `solve`
119/// supplies `z` at each stage time `t + cᵢ Δt` from the stage point, so the
120/// drive is resolved *within* the window — the half-explicit RK treatment of an
121/// index-1 DAE, but with the state leg kept on its group.
122///
123/// `solve` is seeded with the previous stage's `z` and must return a `z`
124/// satisfying the constraint at the stage it is given; the returned `z` is the
125/// one consistent with the step's own endpoint.
126///
127/// FSAL: `first_rate` seeds stage 0 with a rate carried from the previous
128/// step, skipping both that rate evaluation and its constraint solve (`z`
129/// stays the `z` passed in, which is already consistent with `(t, point)`);
130/// when `Tab::FSAL`, the raw rate at the final stage is returned for the next
131/// step to seed with.
132#[allow(clippy::too_many_arguments, clippy::type_complexity)]
133pub fn rkmk_dae_step<Field, Tab, Z, T>(
134    rate: &mut impl FnMut(
135        Quantity<T>,
136        &Field::Point,
137        &Z,
138    ) -> Result<Derivative<Field::Increment, T>, String>,
139    solve: &mut impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<Z, String>,
140    point: &Field::Point,
141    z: &Z,
142    t: Quantity<T>,
143    dt: Quantity<T>,
144    scratch: &mut Vec<Field::Increment>,
145    first_rate: Option<&Derivative<Field::Increment, T>>,
146) -> Result<(Field::Point, Z, Option<Derivative<Field::Increment, T>>), IntegrationError>
147where
148    Field: Integrable,
149    Tab: ButcherTableau,
150    Field::Point: Clone,
151    Field::Increment: Clone + Differentiable<T>,
152    Z: Clone,
153    T: Copy,
154    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
155    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
156{
157    let (z_stage, carry) = rkmk_dae_stage_slopes_into::<Field, Tab, Z, T>(
158        rate, solve, point, z, t, dt, scratch, first_rate,
159    )?;
160    let advanced = reconstruct_or_err::<Field>(point, &weight(scratch, Tab::B))?;
161    let z_final = solve(t + dt, &advanced, &z_stage)?;
162    Ok((advanced, z_final, carry))
163}
164
165/// [`rkmk_dae_step`] with the algebraic unknown resolved by first-order
166/// root-finding at every stage abscissa, built from `function`/`jacobian`/
167/// `solver` exactly as `ExplicitDaeVariableStepExplicitFirstOrderRoot` builds
168/// its `solution` closure for the legacy flat DAE solver — the split between
169/// root-finding and minimization is orthogonal to which field the state lives
170/// on, so this is the one place that wrapping happens for the RKMK-DAE path.
171/// Any [`super::StateEvolution`] model that also supplies a residual and its
172/// Jacobian in terms of the *whole* field state gets the manifold-aware
173/// stage-equilibrium step for free, without hand-rolling this closure itself.
174#[allow(clippy::too_many_arguments, clippy::type_complexity)]
175pub fn rkmk_dae_step_first_order_root<Field, Tab, F, J, Z, T>(
176    rate: &mut impl FnMut(
177        Quantity<T>,
178        &Field::Point,
179        &Z,
180    ) -> Result<Derivative<Field::Increment, T>, String>,
181    mut function: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<F, String>,
182    mut jacobian: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<J, String>,
183    solver: &impl FirstOrderRootFinding<F, J, Z>,
184    point: &Field::Point,
185    z: &Z,
186    t: Quantity<T>,
187    dt: Quantity<T>,
188    scratch: &mut Vec<Field::Increment>,
189    first_rate: Option<&Derivative<Field::Increment, T>>,
190    mut equality_constraint: impl FnMut(Quantity<T>) -> EqualityConstraint,
191) -> Result<(Field::Point, Z, Option<Derivative<Field::Increment, T>>), IntegrationError>
192where
193    Field: Integrable,
194    Tab: ButcherTableau,
195    Field::Point: Clone,
196    Field::Increment: Clone + Differentiable<T>,
197    Z: Clone,
198    T: Copy,
199    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
200    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
201{
202    let mut solve = |t: Quantity<T>, point: &Field::Point, z_guess: &Z| -> Result<Z, String> {
203        Ok(solver.root(
204            |z| function(t, point, z),
205            |z| jacobian(t, point, z),
206            z_guess.clone(),
207            equality_constraint(t),
208            None,
209        )?)
210    };
211    rkmk_dae_step::<Field, Tab, Z, T>(rate, &mut solve, point, z, t, dt, scratch, first_rate)
212}
213
214/// [`rkmk_dae_step`] with the algebraic unknown resolved by second-order
215/// minimization at every stage abscissa, built from `function`/`jacobian`/
216/// `hessian`/`solver` the same way [`rkmk_dae_step_first_order_root`] builds
217/// it for root-finding — the two are siblings so a model whose equilibrium is
218/// naturally posed as a potential (rather than a residual) gets the same
219/// manifold-aware stage-equilibrium step.
220#[allow(clippy::too_many_arguments, clippy::type_complexity)]
221pub fn rkmk_dae_step_second_order_minimize<Field, Tab, F, J, H, Z, T>(
222    rate: &mut impl FnMut(
223        Quantity<T>,
224        &Field::Point,
225        &Z,
226    ) -> Result<Derivative<Field::Increment, T>, String>,
227    mut function: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<F, String>,
228    mut jacobian: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<J, String>,
229    mut hessian: impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<H, String>,
230    solver: &impl SecondOrderOptimization<F, J, H, Z>,
231    point: &Field::Point,
232    z: &Z,
233    t: Quantity<T>,
234    dt: Quantity<T>,
235    scratch: &mut Vec<Field::Increment>,
236    first_rate: Option<&Derivative<Field::Increment, T>>,
237    mut equality_constraint: impl FnMut(Quantity<T>) -> EqualityConstraint,
238    sparse: Option<SparseSolver>,
239) -> Result<(Field::Point, Z, Option<Derivative<Field::Increment, T>>), IntegrationError>
240where
241    Field: Integrable,
242    Tab: ButcherTableau,
243    Field::Point: Clone,
244    Field::Increment: Clone + Differentiable<T>,
245    Z: Clone,
246    T: Copy,
247    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
248    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
249{
250    let mut solve = |t: Quantity<T>, point: &Field::Point, z_guess: &Z| -> Result<Z, String> {
251        Ok(solver.minimize(
252            |z| function(t, point, z),
253            |z| jacobian(t, point, z),
254            |z| hessian(t, point, z),
255            z_guess.clone(),
256            equality_constraint(t),
257            sparse.clone(),
258        )?)
259    };
260    rkmk_dae_step::<Field, Tab, Z, T>(rate, &mut solve, point, z, t, dt, scratch, first_rate)
261}
262
263/// Fills `slopes` with one RKMK-DAE step's corrected stage slopes, resolving the
264/// algebraic unknown at each stage abscissa; returns the last stage's `z` as the
265/// seed for the caller's endpoint solve, and the FSAL carry (see
266/// [`rkmk_dae_step`]). [`rkmk_dae_step`] without the endpoint, so an adaptive
267/// driver can weight the slopes by `D` and reject a step before paying for that
268/// solve.
269#[allow(clippy::too_many_arguments, clippy::type_complexity)]
270pub(super) fn rkmk_dae_stage_slopes_into<Field, Tab, Z, T>(
271    rate: &mut impl FnMut(
272        Quantity<T>,
273        &Field::Point,
274        &Z,
275    ) -> Result<Derivative<Field::Increment, T>, String>,
276    solve: &mut impl FnMut(Quantity<T>, &Field::Point, &Z) -> Result<Z, String>,
277    point: &Field::Point,
278    z: &Z,
279    t: Quantity<T>,
280    dt: Quantity<T>,
281    slopes: &mut Vec<Field::Increment>,
282    first_rate: Option<&Derivative<Field::Increment, T>>,
283) -> Result<(Z, Option<Derivative<Field::Increment, T>>), IntegrationError>
284where
285    Field: Integrable,
286    Tab: ButcherTableau,
287    Field::Point: Clone,
288    Field::Increment: Clone + Differentiable<T>,
289    Z: Clone,
290    T: Copy,
291    Quantity<T>: Mul<Scalar, Output = Quantity<T>>,
292    for<'a> &'a Derivative<Field::Increment, T>: Mul<Quantity<T>, Output = Field::Increment>,
293{
294    slopes.clear();
295    slopes.reserve(Tab::STAGES);
296    let mut z_stage = z.clone();
297    let mut carry = None;
298    for i in 0..Tab::STAGES {
299        let sigma = if i == 0 {
300            None
301        } else {
302            let mut accumulated = slopes[0].clone() * Tab::A[i][0];
303            for (j, slope) in slopes.iter().enumerate().take(i).skip(1) {
304                accumulated += slope.clone() * Tab::A[i][j];
305            }
306            Some(accumulated)
307        };
308        let stage_point = match &sigma {
309            Some(sigma) => reconstruct_or_err::<Field>(point, sigma)?,
310            None => point.clone(),
311        };
312        let t_stage = t + dt * Tab::C[i];
313        let increment = match (i, first_rate) {
314            (0, Some(seed)) => seed * dt,
315            _ => {
316                z_stage = solve(t_stage, &stage_point, &z_stage)?;
317                let raw = rate(t_stage, &stage_point, &z_stage)?;
318                let increment = &raw * dt;
319                if Tab::FSAL && i + 1 == Tab::STAGES {
320                    carry = Some(raw);
321                }
322                increment
323            }
324        };
325        slopes.push(match &sigma {
326            Some(sigma) => Field::dexpinv(sigma, increment),
327            None => increment,
328        });
329    }
330    Ok((z_stage, carry))
331}
conspire/math/integrate/field/rkmk/mod.rs

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