differential_equations/ode/methods/runge_kutta/implicit/
radau5.rs

1//! Adaptive step size implicit Radau IIA 5th order method for solving ODEs.
2
3use crate::{
4    Error, Status,
5    alias::Evals,
6    interpolate::Interpolation,
7    ode::{ODENumericalMethod, ODE, methods::h_init},
8    traits::{CallBackData, Real, State},
9    utils::{constrain_step_size, validate_step_size_parameters},
10};
11
12// Constants for Radau IIA 3-stage, 5th order method
13const SQRT6: f64 = 2.449489743; // sqrt(6.0)
14
15// Time points c_i
16const C0: f64 = (4.0 - SQRT6) / 10.0;
17const C1: f64 = (4.0 + SQRT6) / 10.0;
18const C2: f64 = 1.0;
19
20// Butcher tableau A matrix coefficients
21const A11: f64 = (88.0 - 7.0 * SQRT6) / 360.0;
22const A12: f64 = (296.0 - 169.0 * SQRT6) / 1800.0;
23const A13: f64 = (-2.0 + 3.0 * SQRT6) / 225.0;
24const A21: f64 = (296.0 + 169.0 * SQRT6) / 1800.0;
25const A22: f64 = (88.0 + 7.0 * SQRT6) / 360.0;
26const A23: f64 = (-2.0 - 3.0 * SQRT6) / 225.0;
27const A31: f64 = (16.0 - SQRT6) / 36.0;
28const A32: f64 = (16.0 + SQRT6) / 36.0;
29const A33: f64 = 1.0 / 9.0;
30
31// Primary weights b_i
32const B0: f64 = (16.0 - SQRT6) / 36.0;
33const B1: f64 = (16.0 + SQRT6) / 36.0;
34const B2: f64 = 1.0 / 9.0;
35
36// Embedded method weights b_hat_i for error estimation
37const BHAT0: f64 = (16.0 - SQRT6) / 36.0 - 0.01;
38const BHAT1: f64 = (16.0 + SQRT6) / 36.0 - 0.01;
39const BHAT2: f64 = 1.0 / 9.0 + 0.02;
40
41// TODO: Add matrix coefficients and different step size selection strategies used
42// in the Fortran implementation. Current version only implements core and dense output
43// functionality with newton iteration and jacobian calls.
44
45/// Radau IIA method of order 5 (3 stages).
46///
47/// This is a 3-stage, 5th order implicit Runge-Kutta method, A-stable, L-stable,
48/// suitable for stiff problems. It uses Newton iteration, embedded error estimation,
49/// and specialized Radau dense output for accurate interpolation.
50/// 
51/// Note currently this uses the same constants as the Radau IIA 3-stage method.
52/// Unlike Fortran implemnentation no matrix, ti, coefficients are used and thus
53/// this is a simplified version. 
54/// 
55pub struct Radau5<T: Real, V: State<T>, D: CallBackData> {
56    // Initial Step Size
57    pub h0: T,
58    // Current Step Size
59    h: T,
60    // Previous step size (used for interpolation)
61    h_prev_step: T,
62
63    // Current State (t_n, y_n)
64    t: T,
65    y: V,
66    dydt: V, // Derivative f(t_n, y_n)
67
68    // Previous State (t_{n-1}, y_{n-1})
69    t_prev: T,
70    y_prev: V,
71    dydt_prev: V, // Derivative f(t_{n-1}, y_{n-1})
72
73    // Stage derivatives k_i = f(t_n + c_i*h, Y_i)
74    k: [V; 3],
75    // Stage values Y_i = y_n + h * sum(a_ij * k_j)
76    y_stage: [V; 3],
77    // f(t_stage, y_stage) during Newton iteration
78    f_at_stages: [V; 3],
79
80    // Butcher tableau coefficients (typed)
81    a: [[T; 3]; 3],
82    b_higher: [T; 3], // Primary weights b
83    b_lower: [T; 3],  // Secondary weights b_hat for error estimation
84    c: [T; 3],        // Time points
85
86    // Dense output coefficients for polynomial interpolation
87    cont: [V; 4],
88
89    // --- Adaptive Step Settings ---
90    pub rtol: T,
91    pub atol: T,
92    pub h_max: T,
93    pub h_min: T,
94    pub max_steps: usize,
95    pub max_rejects: usize,
96    pub safety_factor: T,
97    pub min_scale: T,
98    pub max_scale: T,
99
100    // --- Implicit Solver Settings ---
101    pub max_iter: usize, // Max iterations for Newton solver
102    pub tol: T,          // Tolerance for Newton solver convergence
103    fd_epsilon_sqrt: T, // Stores sqrt(machine_epsilon) for FD
104
105    // Iteration tracking & Status
106    reject: bool,
107    n_stiff: usize,
108    steps: usize,
109    status: Status<T, V, D>,
110
111    // --- Jacobian and Newton Solver Data ---
112    jacobian_matrix: nalgebra::DMatrix<T>, // Jacobian of f: J(t,y)
113    newton_matrix: nalgebra::DMatrix<T>,   // Matrix for Newton system (M)
114    rhs_newton: nalgebra::DVector<T>,      // RHS vector for Newton system (-phi)
115    delta_k_vec: nalgebra::DVector<T>,     // Solution of Newton system (delta_k)
116}
117
118impl<T: Real, V: State<T>, D: CallBackData> Default for Radau5<T, V, D> {
119    fn default() -> Self {
120        // Convert Butcher tableau f64 constants to type T
121        let a_coeffs: [[f64; 3]; 3] = [
122            [A11, A12, A13],
123            [A21, A22, A23],
124            [A31, A32, A33]
125        ];
126        let b_coeffs: [f64; 3] = [B0, B1, B2];
127        let b_hat_coeffs: [f64; 3] = [BHAT0, BHAT1, BHAT2];
128        let c_coeffs: [f64; 3] = [C0, C1, C2];
129
130        let a_t: [[T; 3]; 3] = a_coeffs.map(|row| row.map(|x| T::from_f64(x).unwrap()));
131        let b_higher_t: [T; 3] = b_coeffs.map(|x| T::from_f64(x).unwrap());
132        let b_lower_t: [T; 3] = b_hat_coeffs.map(|x| T::from_f64(x).unwrap());
133        let c_t: [T; 3] = c_coeffs.map(|x| T::from_f64(x).unwrap());
134
135        Radau5 {
136            h0: T::zero(), 
137            h: T::zero(),
138            h_prev_step: T::zero(),
139            t: T::zero(),
140            y: V::zeros(),
141            dydt: V::zeros(),
142            t_prev: T::zero(),
143            y_prev: V::zeros(),
144            dydt_prev: V::zeros(),
145            k: [V::zeros(); 3],
146            y_stage: [V::zeros(); 3],
147            f_at_stages: [V::zeros(); 3],
148            a: a_t,
149            b_higher: b_higher_t,
150            b_lower: b_lower_t,
151            c: c_t,
152            cont: [V::zeros(); 4],
153            // Adaptive defaults
154            rtol: T::from_f64(1.0e-6).unwrap(),
155            atol: T::from_f64(1.0e-6).unwrap(),
156            h_max: T::infinity(),
157            h_min: T::zero(),
158            max_steps: 10000,
159            max_rejects: 100,
160            safety_factor: T::from_f64(0.9).unwrap(),
161            min_scale: T::from_f64(0.2).unwrap(),
162            max_scale: T::from_f64(10.0).unwrap(),
163            // Implicit defaults
164            max_iter: 50,
165            tol: T::from_f64(1e-8).unwrap(),
166            fd_epsilon_sqrt: T::zero(),
167            // Status
168            reject: false,
169            n_stiff: 0,
170            steps: 0,
171            status: Status::Uninitialized,
172            // Initialize nalgebra structures (empty, to be sized in init)
173            jacobian_matrix: nalgebra::DMatrix::zeros(0, 0),
174            newton_matrix: nalgebra::DMatrix::zeros(0, 0),
175            rhs_newton: nalgebra::DVector::zeros(0),
176            delta_k_vec: nalgebra::DVector::zeros(0),
177        }
178    }
179}
180
181impl<T: Real, V: State<T>, D: CallBackData> ODENumericalMethod<T, V, D> for Radau5<T, V, D> {
182    fn init<F>(&mut self, ode: &F, t0: T, tf: T, y0: &V) -> Result<Evals, Error<T, V>>
183    where
184        F: ODE<T, V, D>,
185    {
186        let mut evals = Evals::new();
187
188        // Calculate initial derivative f(t0, y0)
189        let mut initial_dydt = V::zeros();
190        ode.diff(t0, y0, &mut initial_dydt);
191        evals.fcn += 1;
192
193        // If h0 is zero calculate h0 using initial derivative
194        if self.h0 == T::zero() {
195            self.h0 = h_init(ode, t0, tf, y0, 5, self.rtol, self.atol, self.h_min, self.h_max); // Order 5
196        }
197
198        // Check bounds
199        self.h = validate_step_size_parameters::<T, V, D>(self.h0, self.h_min, self.h_max, t0, tf)?;
200        self.h_prev_step = self.h;
201
202        // Initialize Statistics
203        self.reject = false;
204        self.n_stiff = 0;
205        self.steps = 0;
206
207        // Initialize State
208        self.t = t0;
209        self.y = *y0;
210        self.dydt = initial_dydt; // Store f(t0, y0)
211
212        // Initialize previous state (same as current initially)
213        self.t_prev = t0;
214        self.y_prev = *y0;
215        self.dydt_prev = initial_dydt;
216
217        // Initialize fd_epsilon_sqrt
218        self.fd_epsilon_sqrt = T::default_epsilon().sqrt();
219
220        // Initialize Status
221        self.status = Status::Initialized;
222
223        // Initialize Jacobian and Newton-related matrices/vectors with correct dimensions
224        let dim = y0.len();
225        self.jacobian_matrix = nalgebra::DMatrix::zeros(dim, dim);
226        let newton_system_size = 3 * dim; // 3 stages for Radau5
227        self.newton_matrix = nalgebra::DMatrix::zeros(newton_system_size, newton_system_size);
228        self.rhs_newton = nalgebra::DVector::zeros(newton_system_size);
229        self.delta_k_vec = nalgebra::DVector::zeros(newton_system_size);
230        self.f_at_stages = [V::zeros(); 3]; 
231
232        // Initialize dense output coefficients
233        self.cont = [V::zeros(); 4];
234
235        Ok(evals)
236    }
237
238    fn step<F>(&mut self, ode: &F) -> Result<Evals, Error<T, V>>
239    where
240        F: ODE<T, V, D>,
241    {
242        let mut evals = Evals::new();
243
244        // Check step size validity
245        if self.h.abs() < self.h_min || self.h.abs() < T::default_epsilon() {
246            self.status = Status::Error(Error::StepSize { t: self.t, y: self.y });
247            return Err(Error::StepSize { t: self.t, y: self.y });
248        }
249
250        // Check max steps
251        if self.steps >= self.max_steps {
252            self.status = Status::Error(Error::MaxSteps { t: self.t, y: self.y });
253            return Err(Error::MaxSteps { t: self.t, y: self.y });
254        }
255        self.steps += 1;
256
257        // --- Newton Iteration for stage derivatives k_i ---
258        // Initial guess for k_i: k_i^(0) = f(t_n, y_n) (stored in self.dydt)
259        for i in 0..3 { 
260            self.k[i] = self.dydt;
261        }
262
263        // Calculate Jacobian J_n = df/dy(t_n, y_n) once per step attempt
264        ode.jacobian(self.t, &self.y, &mut self.jacobian_matrix);
265        evals.jac += 1;
266
267        // Form Newton iteration matrix: M = I - h * (A ⊗ J)
268        let newton_converged = self.newton_iteration(ode, &mut evals)?;
269        
270        if !newton_converged {
271            self.h *= T::from_f64(0.25).unwrap();
272            self.h = constrain_step_size(self.h, self.h_min, self.h_max);
273            self.reject = true;
274            self.n_stiff += 1;
275            evals.fcn += 1;
276
277            if self.n_stiff >= self.max_rejects {
278                self.status = Status::Error(Error::Stiffness { t: self.t, y: self.y });
279                return Err(Error::Stiffness { t: self.t, y: self.y });
280            }
281            return Ok(evals); // Step rejection
282        }
283
284        // --- Newton iteration converged, compute solutions and error ---
285        let mut delta_y_high = V::zeros();
286        for i in 0..3 { 
287            delta_y_high += self.k[i] * (self.b_higher[i] * self.h);
288        }
289        let y_high = self.y + delta_y_high;
290
291        let mut delta_y_low = V::zeros();
292        for i in 0..3 { 
293            delta_y_low += self.k[i] * (self.b_lower[i] * self.h);
294        }
295        let y_low = self.y + delta_y_low;
296
297        let err_vec = y_high - y_low;
298
299        // Calculate scaled error norm
300        let mut err_norm = T::zero();
301        for n in 0..self.y.len() {
302            let scale = self.atol + self.rtol * self.y.get(n).abs().max(y_high.get(n).abs());
303            if scale > T::zero() {
304                err_norm = err_norm.max((err_vec.get(n) / scale).abs());
305            }
306        }
307        err_norm = err_norm.max(T::default_epsilon() * T::from_f64(100.0).unwrap());
308
309        // Calculate new step size based on error
310        let order_p1 = T::from_usize(5 + 1).unwrap(); // Order 5 for Radau5
311        let mut scale = self.safety_factor * err_norm.powf(-T::one() / order_p1);
312        scale = scale.max(self.min_scale).min(self.max_scale);
313        let h_new = self.h * scale;
314
315        if err_norm <= T::one() {
316            // Step accepted
317            self.status = Status::Solving;
318            
319            // Save current step for interpolation
320            self.h_prev_step = self.h;
321            self.t_prev = self.t;
322            self.y_prev = self.y;
323            self.dydt_prev = self.dydt;
324
325            // Update state
326            self.t += self.h;
327            self.y = y_high;
328
329            // Compute new derivative for next step
330            ode.diff(self.t, &self.y, &mut self.dydt);
331            evals.fcn += 1;
332
333            // Calculate dense output coefficients
334            self.compute_dense_output_coeffs();
335
336            if self.reject {
337                self.n_stiff = 0;
338                self.reject = false;
339            }
340            self.h = constrain_step_size(h_new, self.h_min, self.h_max);
341        } else {
342            // Step rejected
343            self.status = Status::RejectedStep;
344            self.reject = true;
345            self.n_stiff += 1;
346
347            if self.n_stiff >= self.max_rejects {
348                self.status = Status::Error(Error::Stiffness { t: self.t, y: self.y });
349                return Err(Error::Stiffness { t: self.t, y: self.y });
350            }
351
352            self.h = constrain_step_size(h_new, self.h_min, self.h_max);
353            return Ok(evals); // Step rejection
354        }
355
356        Ok(evals)
357    }
358
359    fn t(&self) -> T { self.t }
360    fn y(&self) -> &V { &self.y }
361    fn t_prev(&self) -> T { self.t_prev }
362    fn y_prev(&self) -> &V { &self.y_prev }
363    fn h(&self) -> T { self.h }
364    fn set_h(&mut self, h: T) { self.h = h; }
365    fn status(&self) -> &Status<T, V, D> { &self.status }
366    fn set_status(&mut self, status: Status<T, V, D>) { self.status = status; }
367}
368
369impl<T: Real, V: State<T>, D: CallBackData> Radau5<T, V, D> {
370    /// Newton iteration for solving the implicit equations for stage derivatives k_i.
371    fn newton_iteration<F>(&mut self, ode: &F, evals: &mut Evals) -> Result<bool, Error<T, V>>
372    where
373        F: ODE<T, V, D>,
374    {
375        let dim = self.y.len();
376        
377        // Form Newton matrix M = I_block - h * (A_butcher ⊗ J_f)
378        // where I_block is a block identity matrix.
379        for i in 0..3 { // block row index
380            for l in 0..3 { // block column index
381                let scale_factor = -self.h * self.a[i][l];
382                for r_idx in 0..dim { // row index within the block
383                    for c_idx in 0..dim { // column index within the block
384                        self.newton_matrix[(i * dim + r_idx, l * dim + c_idx)] = 
385                            self.jacobian_matrix[(r_idx, c_idx)] * scale_factor;
386                    }
387                }
388                if i == l { // If it's a diagonal block, add Identity
389                    for d_idx in 0..dim {
390                        self.newton_matrix[(i * dim + d_idx, l * dim + d_idx)] += T::one();
391                    }
392                }
393            }
394        }
395
396        let mut converged = false;
397        
398        for _iter in 0..self.max_iter {
399            // 1. Compute residual phi(K_current) and store -phi in rhs_newton
400            //    y_stage[i] = y_n + h * sum_j(a[i][j] * k[j])
401            //    rhs_newton_i = f(t_n + c_i*h, y_stage[i]) - k[i] (this is phi_i)
402            for i in 0..3 { 
403                self.y_stage[i] = self.y; 
404                for j in 0..3 { 
405                    self.y_stage[i] += self.k[j] * (self.a[i][j] * self.h);
406                }
407
408                ode.diff(self.t + self.c[i] * self.h, &self.y_stage[i], &mut self.f_at_stages[i]);
409                evals.fcn += 1;
410
411                for row_idx in 0..dim {
412                    self.rhs_newton[i * dim + row_idx] = self.f_at_stages[i].get(row_idx) - self.k[i].get(row_idx);
413                }
414            }
415
416            // 2. Solve M * delta_k_vec = rhs_newton
417            let lu_decomp = nalgebra::LU::new(self.newton_matrix.clone());
418            if let Some(solution) = lu_decomp.solve(&self.rhs_newton) {
419                self.delta_k_vec.copy_from(&solution);
420            } else {
421                return Ok(false); // Singular matrix
422            }
423
424            // 3. Update K: self.k[i] += delta_k_vec_i
425            let mut norm_delta_k_sq = T::zero();
426            for i in 0..3 { 
427                for row_idx in 0..dim {
428                    let delta_val = self.delta_k_vec[i * dim + row_idx];
429                    let current_val = self.k[i].get(row_idx);
430                    self.k[i].set(row_idx, current_val + delta_val);
431                    norm_delta_k_sq += delta_val * delta_val;
432                }
433            }
434
435            // 4. Check convergence
436            let dyno = norm_delta_k_sq.sqrt();
437            if dyno < self.tol {
438                converged = true;
439                break;
440            }
441        }
442
443        Ok(converged)
444    }
445
446    /// Compute the coefficients for dense output after a successful step.
447    /// Follows the original RADAU5 Fortran implementation.
448    fn compute_dense_output_coeffs(&mut self) {
449        // self.cont[0] stores y at the current time t (y_{n+1})
450        self.cont[0] = self.y;
451        
452        // Time points from self.c array:
453        // self.c[0] (initialized from C0) corresponds to Fortran C1
454        // self.c[1] (initialized from C1) corresponds to Fortran C2
455        // self.c[2] (initialized from C2) is 1.0
456        let c1_f = self.c[0]; 
457        let c2_f = self.c[1]; 
458
459        let c1m1 = c1_f - T::one(); // C1 - 1
460        let c2m1 = c2_f - T::one(); // C2 - 1
461        let c1mc2 = c1_f - c2_f;   // C1 - C2
462        
463        // The Fortran code uses Z1I, Z2I, Z3I for dense output coefficient calculation.
464        // These Zs are y(x_n + c_i*h) - y_n.
465        // In our code, self.y_stage[i] = y_n + h * sum(a_ij * k_j) = Y_i.
466        // self.y_prev is y_n at the point this function is called.
467        // So, (self.y_stage[i] - self.y_prev) corresponds to Fortran's Z_i values.
468        let z1_val = self.y_stage[0] - self.y_prev; 
469        let z2_val = self.y_stage[1] - self.y_prev; 
470        let z3_val = self.y_stage[2] - self.y_prev; 
471        
472        // Fortran formulas for CONT coefficients (translated):
473        // cont[1] = (Z2-Z3)/C2M1
474        // ak = (Z1-Z2)/C1MC2
475        // acont3_temp = Z1/C1
476        // acont3_temp = (ak-acont3_temp)/C2
477        // cont[2] = (ak-cont[1])/C1M1
478        // cont[3] = cont[2]-acont3_temp
479        
480        // Note: self.cont[0] is y_n+1 (current y).
481        // self.cont[1] is the coefficient for s in the interpolation polynomial.
482        // self.cont[2] is the coefficient for s*(s-c2m1).
483        // self.cont[3] is the coefficient for s*(s-c2m1)*(s-c1m1).
484
485        self.cont[1] = (z2_val - z3_val) / c2m1;
486        
487        let ak = (z1_val - z2_val) / c1mc2;
488        
489        let mut acont3_temp = z1_val / c1_f;
490        
491        acont3_temp = (ak - acont3_temp) / c2_f;
492        
493        self.cont[2] = (ak - self.cont[1]) / c1m1;
494        
495        self.cont[3] = self.cont[2] - acont3_temp;
496    }
497}
498
499impl<T: Real, V: State<T>, D: CallBackData> Interpolation<T, V> for Radau5<T, V, D> {
500    fn interpolate(&mut self, t_interp: T) -> Result<V, Error<T, V>> {
501        if t_interp < self.t_prev || t_interp > self.t {
502            return Err(Error::OutOfBounds {
503                t_interp,
504                t_prev: self.t_prev,
505                t_curr: self.t,
506            });
507        }
508        
509        // Normalized time: s = (t_interp - t_curr) / h_prev_step
510        // This makes s range from -1 (at t_prev) to 0 (at t_curr), matching Fortran's S in CONTR5.
511        let s = (t_interp - self.t) / self.h_prev_step;
512        
513        // Polynomial from CONTR5 function in radau5.f:
514        // CONTR5 = CONT(I) + S*(CONT(I+NN)+(S-C2M1)*(CONT(I+NN2)+(S-C1M1)*CONT(I+NN3)))
515        // CONT(I) is self.cont[0] (which is y_curr = y_{n+1})
516        // S is our s.
517        // CONT(I+NN) is self.cont[1], etc.
518        // C1M1 and C2M1 are (self.c[0]-1) and (self.c[1]-1) respectively.
519        
520        let c1_f = self.c[0]; // Corresponds to Fortran C1
521        let c2_f = self.c[1]; // Corresponds to Fortran C2
522        let c1m1 = c1_f - T::one();
523        let c2m1 = c2_f - T::one();
524        
525        // Interpolation polynomial:
526        // y_interp = y_curr + s * (d1 + (s - (c2-1)) * (d2 + (s - (c1-1)) * d3))
527        // where d1=self.cont[1], d2=self.cont[2], d3=self.cont[3]
528        let y_interp = self.cont[0] + (self.cont[1] + (self.cont[2] + self.cont[3] * (s - c1m1)) * (s - c2m1)) * s;
529        
530        Ok(y_interp)
531    }
532}
533
534// Builder pattern methods
535impl<
536    T: Real,
537    V: State<T>,
538    D: CallBackData,
539> Radau5<T, V, D> {
540    pub fn new() -> Self {
541        Self::default()
542    }
543    pub fn h0(mut self, h0: T) -> Self { self.h0 = h0; self }
544    pub fn rtol(mut self, rtol: T) -> Self { self.rtol = rtol; self }
545    pub fn atol(mut self, atol: T) -> Self { self.atol = atol; self }
546    pub fn h_min(mut self, h_min: T) -> Self { self.h_min = h_min; self }
547    pub fn h_max(mut self, h_max: T) -> Self { self.h_max = h_max; self }
548    pub fn max_steps(mut self, max_steps: usize) -> Self { self.max_steps = max_steps; self }
549    pub fn max_rejects(mut self, max_rejects: usize) -> Self { self.max_rejects = max_rejects; self }
550    pub fn safety_factor(mut self, safety_factor: T) -> Self { self.safety_factor = safety_factor; self }
551    pub fn min_scale(mut self, min_scale: T) -> Self { self.min_scale = min_scale; self }
552    pub fn max_scale(mut self, max_scale: T) -> Self { self.max_scale = max_scale; self }
553    pub fn max_iter(mut self, iter: usize) -> Self { self.max_iter = iter; self }
554    pub fn tol(mut self, tol: T) -> Self { self.tol = tol; self }
555}
differential_equations/ode/methods/runge_kutta/implicit/radau5.rs

differential_equations/ode/methods/runge_kutta/implicit/
radau5.rs
