differential_equations/methods/irk/
mod.rs

1//! Implicit Runge-Kutta (IRK) methods
2
3mod adaptive;
4mod fixed;
5
6use crate::{
7    Status,
8    traits::{CallBackData, Real, State},
9};
10use std::{
11    marker::PhantomData
12};
13
14/// Implicit Runge-Kutta solver that can handle:
15/// - Fixed-step methods with Newton iteration for stage equations
16/// - Adaptive step methods with embedded error estimation
17/// - Gauss methods (A-stable, symplectic for Hamiltonian systems)
18/// - Radau methods (L-stable, good for stiff problems)
19/// - Lobatto methods (A-stable, good for constrained systems)
20///
21/// # Type Parameters
22///
23/// * `E`: Equation type (e.g., Ordinary, Delay, Stochastic)
24/// * `F`: Family type (e.g., Adaptive, Fixed, Gauss, Radau, Lobatto)
25/// * `T`: Real number type (f32, f64)
26/// * `V`: State vector type
27/// * `D`: Callback data type
28/// * `const O`: Order of the method
29/// * `const S`: Number of stages in the method
30/// * `const I`: Total number of stages including interpolation (equal to S for methods without dense output)
31pub struct ImplicitRungeKutta<E, F, T: Real, V: State<T>, D: CallBackData, const O: usize, const S: usize, const I: usize> {
32    // Initial Step Size
33    pub h0: T,
34
35    // Current Step Size
36    h: T,
37
38    // Current State
39    t: T,
40    y: V,
41    dydt: V,
42
43    // Previous State
44    h_prev: T,
45    t_prev: T,
46    y_prev: V,
47    dydt_prev: V,
48
49    // Stage values
50    k: [V; I],           // Stage derivatives
51    y_stages: [V; S],    // Stage values (Y_i = y_n + h * sum(a_ij * k_j))
52
53    // Constants from Butcher tableau
54    c: [T; S],                    // Stage time coefficients
55    a: [[T; S]; S],              // Runge-Kutta matrix (typically dense for implicit methods)
56    b: [T; S],                   // Weights for final solution
57    bh: Option<[T; S]>,          // Lower order coefficients for embedded methods
58
59    // Newton iteration settings
60    pub newton_tol: T,           // Tolerance for Newton iteration convergence
61    pub max_newton_iter: usize,  // Maximum Newton iterations per stage
62
63    // Settings
64    pub rtol: T,
65    pub atol: T,
66    pub h_max: T,
67    pub h_min: T,
68    pub max_steps: usize,
69    pub max_rejects: usize,
70    pub safety_factor: T,
71    pub min_scale: T,
72    pub max_scale: T,
73
74    // Iteration tracking
75    jacobian_matrix: nalgebra::DMatrix<T>,    // Jacobian matrix J = df/dy
76    newton_matrix: nalgebra::DMatrix<T>,      // Newton system matrix M = I - h*(A⊗J)
77    rhs_newton: nalgebra::DVector<T>,         // Right-hand side vector for Newton system
78    delta_k_vec: nalgebra::DVector<T>,        // Solution vector for Newton system
79    jacobian_age: usize,                      // Age of current Jacobian (for reuse)
80    stiffness_counter: usize,
81    steps: usize,
82    newton_iterations: usize,    // Total Newton iterations
83    jacobian_evaluations: usize, // Total Jacobian evaluations
84    lu_decompositions: usize,    // Total LU decompositions
85
86    // Status
87    status: Status<T, V, D>,
88    
89    // Method info
90    order: usize,
91    stages: usize,
92    dense_stages: usize,
93
94    // Family classification
95    family: PhantomData<F>,
96
97    // Equation type
98    equation: PhantomData<E>,
99}
100
101impl<E, F, T: Real, V: State<T>, D: CallBackData, const O: usize, const S: usize, const I: usize> Default for ImplicitRungeKutta<E, F, T, V, D, O, S, I> {
102    fn default() -> Self {
103        Self {
104            h0: T::zero(),
105            h: T::zero(),
106            t: T::zero(),
107            y: V::zeros(),
108            dydt: V::zeros(),
109            h_prev: T::zero(),
110            t_prev: T::zero(),
111            y_prev: V::zeros(),
112            dydt_prev: V::zeros(),
113            k: [V::zeros(); I],
114            y_stages: [V::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: T::from_f64(1.0e-6).unwrap(),
122            atol: 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            jacobian_matrix: nalgebra::DMatrix::zeros(0, 0),
142            newton_matrix: nalgebra::DMatrix::zeros(0, 0),
143            rhs_newton: nalgebra::DVector::zeros(0),
144            delta_k_vec: nalgebra::DVector::zeros(0),
145            jacobian_age: 0,
146        }
147    }
148}
149
150impl<E, F, T: Real, V: State<T>, D: CallBackData, const O: usize, const S: usize, const I: usize> ImplicitRungeKutta<E, F, T, V, D, O, S, I> {
151    /// Set the relative tolerance for error control
152    pub fn rtol(mut self, rtol: T) -> Self {
153        self.rtol = rtol;
154        self
155    }
156
157    /// Set the absolute tolerance for error control
158    pub fn atol(mut self, atol: T) -> Self {
159        self.atol = atol;
160        self
161    }
162
163    /// Set the initial step size
164    pub fn h0(mut self, h0: T) -> Self {
165        self.h0 = h0;
166        self
167    }
168
169    /// Set the minimum allowed step size
170    pub fn h_min(mut self, h_min: T) -> Self {
171        self.h_min = h_min;
172        self
173    }
174
175    /// Set the maximum allowed step size
176    pub fn h_max(mut self, h_max: T) -> Self {
177        self.h_max = h_max;
178        self
179    }
180
181    /// Set the maximum number of steps allowed
182    pub fn max_steps(mut self, max_steps: usize) -> Self {
183        self.max_steps = max_steps;
184        self
185    }
186
187    /// Set the maximum number of consecutive rejected steps before declaring stiffness
188    pub fn max_rejects(mut self, max_rejects: usize) -> Self {
189        self.max_rejects = max_rejects;
190        self
191    }
192
193    /// Set the safety factor for step size control (default: 0.9)
194    pub fn safety_factor(mut self, safety_factor: T) -> Self {
195        self.safety_factor = safety_factor;
196        self
197    }
198
199    /// Set the minimum scale factor for step size changes (default: 0.2)
200    pub fn min_scale(mut self, min_scale: T) -> Self {
201        self.min_scale = min_scale;
202        self
203    }
204
205    /// Set the maximum scale factor for step size changes (default: 10.0)
206    pub fn max_scale(mut self, max_scale: T) -> Self {
207        self.max_scale = max_scale;
208        self
209    }
210
211    /// Set the Newton iteration tolerance (default: 1e-10)
212    pub fn newton_tol(mut self, newton_tol: T) -> Self {
213        self.newton_tol = newton_tol;
214        self
215    }
216
217    /// Set the maximum number of Newton iterations per stage (default: 50)
218    pub fn max_newton_iter(mut self, max_newton_iter: usize) -> Self {
219        self.max_newton_iter = max_newton_iter;
220        self
221    }
222
223    /// Get the order of the method
224    pub fn order(&self) -> usize {
225        self.order
226    }
227
228    /// Get the number of stages in the method
229    pub fn stages(&self) -> usize {
230        self.stages
231    }
232
233    /// Get the number of terms in the dense output interpolation polynomial
234    pub fn dense_stages(&self) -> usize {
235        self.dense_stages
236    }
237}
differential_equations/methods/irk/mod.rs

differential_equations/methods/irk/
mod.rs
