differential_equations/methods/irk/
mod.rs

1//! Implicit Runge–Kutta (IRK) methods
2
3mod adaptive;
4mod fixed;
5mod radau;
6
7use std::marker::PhantomData;
8
9use crate::{
10    linalg::Matrix,
11    status::Status,
12    tolerance::Tolerance,
13    traits::{Real, State},
14};
15
16/// IRK solver core. Supports fixed/adaptive stepping and common IRK families
17/// (Gauss, Radau, Lobatto).
18///
19/// Type params: E (equation), F (family), T (scalar), Y (state), D (callback),
20/// O (order), S (stages), I (dense output terms).
21pub struct ImplicitRungeKutta<
22    E,
23    F,
24    T: Real,
25    Y: State<T>,
26    const O: usize,
27    const S: usize,
28    const I: usize,
29> {
30    // Initial step
31    pub h0: T,
32
33    // Current step
34    h: T,
35
36    // Current state
37    t: T,
38    y: Y,
39    dydt: Y,
40
41    // Previous state
42    h_prev: T,
43    t_prev: T,
44    y_prev: Y,
45    dydt_prev: Y,
46
47    // Stage data
48    k: [Y; I], // Stage derivatives
49    z: [Y; S], // Stage states (z_i)
50
51    // Butcher tableau
52    c: [T; S],          // Nodes c
53    a: [[T; S]; S],     // Matrix a
54    b: [T; S],          // Weights b
55    bh: Option<[T; S]>, // Embedded weights
56
57    // Newton settings
58    pub newton_tol: T,          // Convergence tol
59    pub max_newton_iter: usize, // Max iterations per solve
60
61    // Adaptive settings
62    pub rtol: Tolerance<T>,
63    pub atol: Tolerance<T>,
64    pub h_max: T,
65    pub h_min: T,
66    pub max_steps: usize,
67    pub max_rejects: usize,
68    pub safety_factor: T,
69    pub min_scale: T,
70    pub max_scale: T,
71
72    // Iteration tracking
73    stage_jacobians: [Matrix<T>; S], // J_i at each stage
74    newton_matrix: Matrix<T>,        // I - h*(A⊗J)
75    rhs_newton: Vec<T>,              // Newton RHS
76    delta_k_vec: Vec<T>,             // Newton solution
77    jacobian_age: usize,             // Reuse counter
78    stiffness_counter: usize,
79    steps: usize,
80    newton_iterations: usize,    // Total Newton iterations
81    jacobian_evaluations: usize, // Total jacobian evaluations
82    lu_decompositions: usize,    // Total LU decompositions
83
84    // Status
85    status: Status<T, Y>,
86
87    // Method info
88    order: usize,
89    stages: usize,
90    dense_stages: usize,
91
92    // Family classification
93    family: PhantomData<F>,
94
95    // Equation type
96    equation: PhantomData<E>,
97}
98
99impl<E, F, T: Real, Y: State<T>, const O: usize, const S: usize, const I: usize> Default
100    for ImplicitRungeKutta<E, F, T, Y, O, S, I>
101{
102    fn default() -> Self {
103        Self {
104            h0: T::zero(),
105            h: T::zero(),
106            t: T::zero(),
107            y: Y::zeros(),
108            dydt: Y::zeros(),
109            h_prev: T::zero(),
110            t_prev: T::zero(),
111            y_prev: Y::zeros(),
112            dydt_prev: Y::zeros(),
113            k: [Y::zeros(); I],
114            z: [Y::zeros(); S],
115            c: [T::zero(); S],
116            a: [[T::zero(); S]; S],
117            b: [T::zero(); S],
118            bh: None,
119            newton_tol: T::from_f64(1.0e-10).unwrap(),
120            max_newton_iter: 50,
121            rtol: Tolerance::Scalar(T::from_f64(1.0e-6).unwrap()),
122            atol: Tolerance::Scalar(T::from_f64(1.0e-6).unwrap()),
123            h_max: T::infinity(),
124            h_min: T::zero(),
125            max_steps: 10_000,
126            max_rejects: 100,
127            safety_factor: T::from_f64(0.9).unwrap(),
128            min_scale: T::from_f64(0.2).unwrap(),
129            max_scale: T::from_f64(10.0).unwrap(),
130            stiffness_counter: 0,
131            steps: 0,
132            newton_iterations: 0,
133            jacobian_evaluations: 0,
134            lu_decompositions: 0,
135            status: Status::Uninitialized,
136            order: O,
137            stages: S,
138            dense_stages: I,
139            family: PhantomData,
140            equation: PhantomData,
141            stage_jacobians: core::array::from_fn(|_| Matrix::zeros(0, 0)),
142            newton_matrix: Matrix::zeros(0, 0),
143            rhs_newton: Vec::new(),
144            delta_k_vec: Vec::new(),
145            jacobian_age: 0,
146        }
147    }
148}
149
150impl<E, F, T: Real, Y: State<T>, const O: usize, const S: usize, const I: usize>
151    ImplicitRungeKutta<E, F, T, Y, O, S, I>
152{
153    /// Set relative tolerance
154    pub fn rtol<V: Into<Tolerance<T>>>(mut self, rtol: V) -> Self {
155        self.rtol = rtol.into();
156        self
157    }
158
159    /// Set absolute tolerance
160    pub fn atol<V: Into<Tolerance<T>>>(mut self, atol: V) -> Self {
161        self.atol = atol.into();
162        self
163    }
164
165    /// Set initial step size
166    pub fn h0(mut self, h0: T) -> Self {
167        self.h0 = h0;
168        self
169    }
170
171    /// Set minimum step size
172    pub fn h_min(mut self, h_min: T) -> Self {
173        self.h_min = h_min;
174        self
175    }
176
177    /// Set maximum step size
178    pub fn h_max(mut self, h_max: T) -> Self {
179        self.h_max = h_max;
180        self
181    }
182
183    /// Set max steps
184    pub fn max_steps(mut self, max_steps: usize) -> Self {
185        self.max_steps = max_steps;
186        self
187    }
188
189    /// Set max consecutive rejections
190    pub fn max_rejects(mut self, max_rejects: usize) -> Self {
191        self.max_rejects = max_rejects;
192        self
193    }
194
195    /// Set step size safety factor (default: 0.9)
196    pub fn safety_factor(mut self, safety_factor: T) -> Self {
197        self.safety_factor = safety_factor;
198        self
199    }
200
201    /// Set minimum scale for step changes (default: 0.2)
202    pub fn min_scale(mut self, min_scale: T) -> Self {
203        self.min_scale = min_scale;
204        self
205    }
206
207    /// Set maximum scale for step changes (default: 10.0)
208    pub fn max_scale(mut self, max_scale: T) -> Self {
209        self.max_scale = max_scale;
210        self
211    }
212
213    /// Set Newton tolerance (default: 1e-10)
214    pub fn newton_tol(mut self, newton_tol: T) -> Self {
215        self.newton_tol = newton_tol;
216        self
217    }
218
219    /// Set max Newton iterations per stage (default: 50)
220    pub fn max_newton_iter(mut self, max_newton_iter: usize) -> Self {
221        self.max_newton_iter = max_newton_iter;
222        self
223    }
224
225    /// Get method order
226    pub fn order(&self) -> usize {
227        self.order
228    }
229
230    /// Get number of stages
231    pub fn stages(&self) -> usize {
232        self.stages
233    }
234
235    /// Get dense output terms
236    pub fn dense_stages(&self) -> usize {
237        self.dense_stages
238    }
239}