scirs2_integrate/ode/methods/implicit.rs
1//! Implicit ODE solver methods
2//!
3//! This module implements implicit methods for solving ODEs,
4//! including the Backward Differentiation Formula (BDF) method
5//! and the Radau IIA method.
6
7use crate::error::{IntegrateError, IntegrateResult};
8use crate::ode::types::{ODEMethod, ODEOptions, ODEResult};
9use crate::IntegrateFloat;
10use scirs2_core::ndarray::{Array1, Array2, ArrayView1};
11
12/// Solve ODE using the Backward Differentiation Formula (BDF) method
13///
14/// BDF is an implicit multistep method particularly suited for stiff problems.
15/// It is more computationally expensive per step than explicit methods but
16/// can take much larger steps for stiff problems, resulting in overall better
17/// performance for such systems.
18///
19/// # Arguments
20///
21/// * `f` - ODE function dy/dt = f(t, y)
22/// * `t_span` - Time span [t_start, t_end]
23/// * `y0` - Initial condition
24/// * `opts` - Solver options
25///
26/// # Returns
27///
28/// The solution as an ODEResult or an error
29#[allow(dead_code)]
30pub fn bdf_method<F, Func>(
31 f: Func,
32 t_span: [F; 2],
33 y0: Array1<F>,
34 opts: ODEOptions<F>,
35) -> IntegrateResult<ODEResult<F>>
36where
37 F: IntegrateFloat,
38 Func: Fn(F, ArrayView1<F>) -> Array1<F>,
39{
40 // Check BDF order is valid (1-5 supported)
41 let order = opts.max_order.unwrap_or(2);
42 if !(1..=5).contains(&order) {
43 return Err(IntegrateError::ValueError(
44 "BDF order must be between 1 and 5".to_string(),
45 ));
46 }
47
48 // Initialize
49 let [t_start, t_end] = t_span;
50 let n_dim = y0.len();
51
52 // Determine initial step size if not provided
53 let h0 = opts.h0.unwrap_or_else(|| {
54 // Simple heuristic for initial step size
55 let _span = t_end - t_start;
56 _span / F::from_usize(100).unwrap() * F::from_f64(0.1).unwrap() // 0.1% of interval
57 });
58
59 // Determine minimum and maximum step sizes
60 let min_step = opts.min_step.unwrap_or_else(|| {
61 let _span = t_end - t_start;
62 _span * F::from_f64(1e-10).unwrap() // Minimal step size
63 });
64
65 let max_step = opts.max_step.unwrap_or_else(|| {
66 t_end - t_start // Maximum step can be the whole interval
67 });
68
69 // For BDF methods we need previous steps
70 // We'll use another method (RK4) to bootstrap the initial steps
71 // We need order previous points for a BDF method of order 'order'
72
73 // First run RK4 to get the first 'order' points
74 let mut h = h0;
75 let mut t_values = vec![t_start];
76 let mut y_values = vec![y0.clone()];
77 let mut func_evals = 0;
78 let mut step_count = 0;
79 let mut accepted_steps = 0;
80 let mut rejected_steps = 0;
81 let mut newton_iters = F::zero();
82 let mut n_lu = 0;
83 let mut n_jac = 0;
84
85 // Generate initial points using RK4 (more accurate than Euler)
86 if order > 1 {
87 let two = F::from_f64(2.0).unwrap();
88 let mut t = t_start;
89 let mut y = y0.clone();
90
91 for _ in 0..(order - 1) {
92 // Don't go beyond t_end
93 if t + h > t_end {
94 h = t_end - t;
95 }
96
97 // Standard RK4 step
98 let half_step = h / two;
99
100 let k1 = f(t, y.view());
101 let k2 = f(t + half_step, (y.clone() + k1.clone() * half_step).view());
102 let k3 = f(t + half_step, (y.clone() + k2.clone() * half_step).view());
103 let k4 = f(t + h, (y.clone() + k3.clone() * h).view());
104 func_evals += 4;
105
106 // Combine slopes with appropriate weights
107 let slope = (k1 + k2.clone() * two + k3.clone() * two + k4) / F::from_f64(6.0).unwrap();
108 y = y + slope * h;
109
110 // Update time
111 t += h;
112
113 // Store results
114 t_values.push(t);
115 y_values.push(y.clone());
116
117 step_count += 1;
118 accepted_steps += 1;
119
120 // Don't continue if we've reached t_end
121 if t >= t_end {
122 break;
123 }
124 }
125 }
126
127 // Now we have enough points to start BDF
128 let mut t = *t_values.last().unwrap();
129 let mut y = y_values.last().unwrap().clone();
130
131 // BDF coefficients for different orders
132 // These are the coefficients for the BDF formula
133 // For BDF of order p, we use p previous points
134
135 // Coefficients for BDF1 (Implicit Euler) through BDF5
136 let bdf_coefs: [Vec<F>; 5] = [
137 // BDF1 (Implicit Euler): y_{n+1} - y_n = h * f(t_{n+1}, y_{n+1})
138 vec![F::one(), F::from_f64(-1.0).unwrap()],
139 // BDF2: 3/2 * y_{n+1} - 2 * y_n + 1/2 * y_{n-1} = h * f(t_{n+1}, y_{n+1})
140 vec![
141 F::from_f64(3.0 / 2.0).unwrap(),
142 F::from_f64(-2.0).unwrap(),
143 F::from_f64(1.0 / 2.0).unwrap(),
144 ],
145 // BDF3
146 vec![
147 F::from_f64(11.0 / 6.0).unwrap(),
148 F::from_f64(-3.0).unwrap(),
149 F::from_f64(3.0 / 2.0).unwrap(),
150 F::from_f64(-1.0 / 3.0).unwrap(),
151 ],
152 // BDF4
153 vec![
154 F::from_f64(25.0 / 12.0).unwrap(),
155 F::from_f64(-4.0).unwrap(),
156 F::from_f64(3.0).unwrap(),
157 F::from_f64(-4.0 / 3.0).unwrap(),
158 F::from_f64(1.0 / 4.0).unwrap(),
159 ],
160 // BDF5
161 vec![
162 F::from_f64(137.0 / 60.0).unwrap(),
163 F::from_f64(-5.0).unwrap(),
164 F::from_f64(5.0).unwrap(),
165 F::from_f64(-10.0 / 3.0).unwrap(),
166 F::from_f64(5.0 / 4.0).unwrap(),
167 F::from_f64(-1.0 / 5.0).unwrap(),
168 ],
169 ];
170
171 // Only use coefficients for the requested order
172 let coeffs = &bdf_coefs[order - 1];
173
174 // Main integration loop
175 while t < t_end && step_count < opts.max_steps {
176 // Adjust step size for the last step if needed
177 if t + h > t_end {
178 h = t_end - t;
179 }
180
181 // Limit step size to bounds
182 h = h.min(max_step).max(min_step);
183
184 // Create the system matrix for the implicit equation
185 // For BDF, we need to solve the nonlinear system:
186 // c_0 * y_{n+1} - h * f(t_{n+1}, y_{n+1}) - sum_{j=1}^p c_j * y_{n+1-j} = 0
187
188 // Initialize the residual function
189 let next_t = t + h;
190
191 // Predict initial value using explicit extrapolation
192 // This provides a better starting guess for Newton iteration
193 let mut y_pred = y.clone();
194
195 // If we have multiple previous points, use them for prediction
196 if order > 1 && y_values.len() >= order {
197 // Simple linear extrapolation for BDF2
198 if order == 2 {
199 let y_nm1 = &y_values[y_values.len() - 2];
200 y_pred = y.clone() * F::from_f64(2.0).unwrap() - y_nm1 * F::one();
201 } else {
202 // For higher orders, use a quadratic or higher extrapolation
203 // This is approximate but sufficient for an initial guess
204 let y_curr = &y_values[y_values.len() - 1];
205 let y_prev = &y_values[y_values.len() - 2];
206 let y_prev2 = &y_values[y_values.len() - 3];
207
208 let a = F::one();
209 let b = F::from_f64(2.0).unwrap();
210
211 y_pred = y_curr * (a + b) - y_prev * b + y_prev2 * (F::one() - a);
212 }
213 }
214
215 // Newton's method for solving the nonlinear system
216 let mut y_next = y_pred.clone();
217 let mut converged = false;
218 let mut iter_count = 0;
219 let max_newton_iters = 10; // Maximum iterations for Newton's method
220
221 while iter_count < max_newton_iters {
222 // Evaluate function at the current iterate
223 let f_eval = f(next_t, y_next.view());
224 func_evals += 1;
225
226 // Compute residual: c_0 * y_{n+1} - h * f(t_{n+1}, y_{n+1}) - sum_{j=1}^p c_j * y_{n+1-j}
227 let mut residual = y_next.clone() * coeffs[0];
228
229 // Previous values contribution
230 for (j, coeff) in coeffs.iter().enumerate().skip(1).take(order) {
231 if j <= y_values.len() {
232 let idx = y_values.len() - j;
233 residual = residual - y_values[idx].clone() * *coeff;
234 }
235 }
236
237 // Subtract h * f(t_{n+1}, y_{n+1})
238 residual = residual - f_eval.clone() * h;
239
240 // Compute Newton step
241 // In a full implementation, we would compute the Jacobian:
242 // J = c_0 * I - h * df/dy
243 // However, computing the actual Jacobian is complex
244 // For simplicity, we'll use a finite difference approximation
245
246 // Create approximate Jacobian using finite differences
247 let eps = F::from_f64(1e-8).unwrap();
248 let mut jacobian = Array2::<F>::zeros((n_dim, n_dim));
249 n_jac += 1;
250
251 for i in 0..n_dim {
252 let mut y_perturbed = y_next.clone();
253 y_perturbed[i] += eps;
254
255 let f_perturbed = f(next_t, y_perturbed.view());
256 func_evals += 1;
257
258 for j in 0..n_dim {
259 // Finite difference approximation of df_j/dy_i
260 let df_dy = (f_perturbed[j] - f_eval[j]) / eps;
261
262 // J_{ji} = c_0 * δ_{ji} - h * df_j/dy_i
263 jacobian[[j, i]] = if i == j {
264 coeffs[0] - h * df_dy
265 } else {
266 -h * df_dy
267 };
268 }
269 }
270
271 // For small systems (n_dim typically <= 10), we can use a simple approach
272 // instead of full matrix solver to avoid dependency issues
273
274 // For a 1D system, we can directly solve without any matrix inversion
275 if n_dim == 1 {
276 // For scalar case, J is just a number, and delta_y = -residual / J
277 if jacobian[[0, 0]].abs() < F::from_f64(1e-10).unwrap() {
278 // Nearly singular, reduce step size and try again
279 h *= F::from_f64(0.5).unwrap();
280 if h < min_step {
281 return Err(IntegrateError::ConvergenceError(
282 "Newton iteration failed to converge with minimum step size"
283 .to_string(),
284 ));
285 }
286 iter_count = 0;
287 continue;
288 }
289
290 // Direct solution for scalar case
291 let delta_y0 = residual[0] / jacobian[[0, 0]];
292 y_next[0] -= delta_y0;
293 }
294 // For larger systems, use Gaussian elimination
295 else {
296 // Implement Gaussian elimination for larger systems
297 // Copy the matrix and right-hand side for manipulation
298 let mut aug = Array2::<F>::zeros((n_dim, n_dim + 1));
299 for i in 0..n_dim {
300 for j in 0..n_dim {
301 aug[[i, j]] = jacobian[[i, j]];
302 }
303 aug[[i, n_dim]] = residual[i];
304 }
305
306 // Gaussian elimination with partial pivoting
307 n_lu += 1;
308 for i in 0..n_dim {
309 // Find pivot
310 let mut max_idx = i;
311 let mut max_val = aug[[i, i]].abs();
312
313 for j in i + 1..n_dim {
314 if aug[[j, i]].abs() > max_val {
315 max_idx = j;
316 max_val = aug[[j, i]].abs();
317 }
318 }
319
320 // Check if the matrix is singular
321 if max_val < F::from_f64(1e-10).unwrap() {
322 // Nearly singular matrix, reduce step size and try again
323 h *= F::from_f64(0.5).unwrap();
324 if h < min_step {
325 return Err(IntegrateError::ConvergenceError(
326 "Newton iteration failed to converge with minimum step size"
327 .to_string(),
328 ));
329 }
330 iter_count = 0;
331 continue;
332 }
333
334 // Swap rows if necessary
335 if max_idx != i {
336 for j in 0..n_dim + 1 {
337 let temp = aug[[i, j]];
338 aug[[i, j]] = aug[[max_idx, j]];
339 aug[[max_idx, j]] = temp;
340 }
341 }
342
343 // Eliminate below
344 for j in i + 1..n_dim {
345 let factor = aug[[j, i]] / aug[[i, i]];
346 for k in i..n_dim + 1 {
347 aug[[j, k]] = aug[[j, k]] - factor * aug[[i, k]];
348 }
349 }
350 }
351
352 // Back substitution
353 let mut delta_y = Array1::<F>::zeros(n_dim);
354 for i in (0..n_dim).rev() {
355 let mut sum = aug[[i, n_dim]];
356 for j in i + 1..n_dim {
357 sum -= aug[[i, j]] * delta_y[j];
358 }
359 delta_y[i] = sum / aug[[i, i]];
360 }
361
362 // Update solution
363 for i in 0..n_dim {
364 y_next[i] -= delta_y[i];
365 }
366 }
367
368 // Check convergence
369 let newton_tol = F::from_f64(1e-8).unwrap();
370 let mut err = F::zero();
371 for i in 0..n_dim {
372 let sc = opts.atol + opts.rtol * y_next[i].abs();
373 let e: F = residual[i] / sc;
374 err = err.max(e.abs());
375 }
376
377 if err <= newton_tol {
378 converged = true;
379 break;
380 }
381
382 iter_count += 1;
383 }
384
385 newton_iters += F::from(iter_count).unwrap();
386
387 if converged {
388 // Step accepted
389 t = next_t;
390 y = y_next;
391
392 // Store results
393 t_values.push(t);
394 y_values.push(y.clone());
395
396 step_count += 1;
397 accepted_steps += 1;
398
399 // Adjust step size based on Newton convergence
400 // If we converged quickly, increase step size
401 if iter_count <= 2 {
402 h *= F::from_f64(1.2).unwrap().min(F::from_f64(5.0).unwrap());
403 }
404 // If we needed many iterations, decrease step size
405 else if iter_count >= max_newton_iters - 1 {
406 h *= F::from_f64(0.8).unwrap().max(F::from_f64(0.2).unwrap());
407 }
408 } else {
409 // Newton iteration failed to converge, reduce step size
410 h *= F::from_f64(0.5).unwrap();
411 if h < min_step {
412 return Err(IntegrateError::ConvergenceError(
413 "Newton iteration failed to converge with minimum step size".to_string(),
414 ));
415 }
416 rejected_steps += 1;
417 }
418 }
419
420 let success = t >= t_end;
421 let message = if !success {
422 Some(format!(
423 "Maximum number of steps ({}) reached",
424 opts.max_steps
425 ))
426 } else {
427 None
428 };
429
430 // Return the solution
431 Ok(ODEResult {
432 t: t_values,
433 y: y_values,
434 success,
435 message,
436 n_eval: func_evals,
437 n_steps: step_count,
438 n_accepted: accepted_steps,
439 n_rejected: rejected_steps,
440 n_lu,
441 n_jac,
442 method: ODEMethod::Bdf,
443 })
444}
445
446/// Solve ODE using the Radau IIA method
447///
448/// Radau IIA is an implicit Runge-Kutta method with high stability properties,
449/// making it suitable for stiff problems. It is an L-stable method with high
450/// order of accuracy.
451///
452/// # Arguments
453///
454/// * `f` - ODE function dy/dt = f(t, y)
455/// * `t_span` - Time span [t_start, t_end]
456/// * `y0` - Initial condition
457/// * `opts` - Solver options
458///
459/// # Returns
460///
461/// The solution as an ODEResult or an error
462#[allow(dead_code)]
463pub fn radau_method<F, Func>(
464 f: Func,
465 t_span: [F; 2],
466 y0: Array1<F>,
467 opts: ODEOptions<F>,
468) -> IntegrateResult<ODEResult<F>>
469where
470 F: IntegrateFloat,
471 Func: Fn(F, ArrayView1<F>) -> Array1<F>,
472{
473 // Initialize
474 let [t_start, t_end] = t_span;
475 let n_dim = y0.len();
476
477 // Determine initial step size if not provided
478 let h0 = opts.h0.unwrap_or_else(|| {
479 // Simple heuristic for initial step size
480 let _span = t_end - t_start;
481 _span / F::from_usize(100).unwrap() * F::from_f64(0.1).unwrap() // 0.1% of interval
482 });
483
484 // Determine minimum and maximum step sizes
485 let min_step = opts.min_step.unwrap_or_else(|| {
486 let _span = t_end - t_start;
487 _span * F::from_f64(1e-10).unwrap() // Minimal step size
488 });
489
490 let max_step = opts.max_step.unwrap_or_else(|| {
491 t_end - t_start // Maximum step can be the whole interval
492 });
493
494 // Radau IIA 3-stage method (5th order)
495 // Butcher tableau for Radau IIA (3 stages)
496 // c_i | a_ij
497 // ----------
498 // | b_j
499 //
500 // c = [4-sqrt(6))/10, (4+sqrt(6))/10, 1]
501 // Exact values would be irrational, so we use high precision approximations
502
503 let c1 = F::from_f64(0.1550510257).unwrap();
504 let c2 = F::from_f64(0.6449489743).unwrap();
505 let c3 = F::one();
506
507 // Runge-Kutta matrix A (coefficients a_ij)
508 // We're using a 3-stage Radau IIA method
509 let a11 = F::from_f64(0.1968154772).unwrap();
510 let a12 = F::from_f64(-0.0678338608).unwrap();
511 let a13 = F::from_f64(-0.0207959730).unwrap();
512
513 let a21 = F::from_f64(0.3944243147).unwrap();
514 let a22 = F::from_f64(0.2921005631).unwrap();
515 let a23 = F::from_f64(0.0416635118).unwrap();
516
517 let a31 = F::from_f64(0.3764030627).unwrap();
518 let a32 = F::from_f64(0.5124858261).unwrap();
519 let a33 = F::from_f64(0.1111111111).unwrap();
520
521 // Weight coefficients b_j (same as last row of A for Radau IIA)
522 let b1 = a31;
523 let b2 = a32;
524 let b3 = a33;
525
526 // Integration variables
527 let mut t = t_start;
528 let mut y = y0.clone();
529 let mut h = h0;
530
531 // Result storage
532 let mut t_values = vec![t];
533 let mut y_values = vec![y.clone()];
534
535 // Statistics
536 let mut func_evals = 0;
537 let mut step_count = 0;
538 let mut accepted_steps = 0;
539 let mut rejected_steps = 0;
540 let mut n_lu = 0;
541 let mut n_jac = 0;
542
543 // Newton iteration parameters
544 let newton_tol = F::from_f64(1e-8).unwrap();
545 let max_newton_iters = 10;
546
547 // Main integration loop
548 while t < t_end && step_count < opts.max_steps {
549 // Adjust step size for the last step if needed
550 if t + h > t_end {
551 h = t_end - t;
552 }
553
554 // Limit step size to bounds
555 h = h.min(max_step).max(min_step);
556
557 // Stage values
558 let t1 = t + c1 * h;
559 let t2 = t + c2 * h;
560 let t3 = t + c3 * h;
561
562 // Initial guess for stages (simple extrapolation from current state)
563 let mut k1 = y.clone();
564 let mut k2 = y.clone();
565 let mut k3 = y.clone();
566
567 // Newton's method to solve for stage values
568 let mut converged = false;
569 let mut iter_count = 0;
570
571 // Newton iteration loop
572 while iter_count < max_newton_iters {
573 // Evaluate function at each stage
574 let f1 = f(t1, k1.view());
575 let f2 = f(t2, k2.view());
576 let f3 = f(t3, k3.view());
577 func_evals += 3;
578
579 // Compute residuals for each stage
580 // r_i = k_i - y_n - h * sum_j(a_ij * f_j)
581 let mut r1 = k1.clone();
582 r1 -= &y;
583 r1 = r1 - (&f1 * (h * a11) + &f2 * (h * a12) + &f3 * (h * a13));
584
585 let mut r2 = k2.clone();
586 r2 -= &y;
587 r2 = r2 - (&f1 * (h * a21) + &f2 * (h * a22) + &f3 * (h * a23));
588
589 let mut r3 = k3.clone();
590 r3 -= &y;
591 r3 = r3 - (&f1 * (h * a31) + &f2 * (h * a32) + &f3 * (h * a33));
592
593 // Check for convergence
594 let mut max_res = F::zero();
595 for i in 0..n_dim {
596 let sc = opts.atol + opts.rtol * y[i].abs();
597 max_res = max_res.max(r1[i].abs() / sc);
598 max_res = max_res.max(r2[i].abs() / sc);
599 max_res = max_res.max(r3[i].abs() / sc);
600 }
601
602 if max_res <= newton_tol {
603 converged = true;
604 break;
605 }
606
607 // Construct Jacobian for Newton iteration
608 // For simplicity, we'll use a block-diagonal approximation
609 // Each block corresponds to a single component across all stages
610 n_jac += 1;
611
612 // For small system sizes, we can use a direct approach for each component
613 for i in 0..n_dim {
614 // Extract i-th component for all stages
615 let mut yi = Array1::<F>::zeros(3);
616 let mut ri = Array1::<F>::zeros(3);
617 yi[0] = k1[i];
618 yi[1] = k2[i];
619 yi[2] = k3[i];
620 ri[0] = r1[i];
621 ri[1] = r2[i];
622 ri[2] = r3[i];
623
624 // Approximate Jacobian for this component using finite differences
625 let eps = F::from_f64(1e-8).unwrap();
626 let mut jac = Array2::<F>::zeros((3, 3));
627
628 // Disturb each stage value for this component
629 for j in 0..3 {
630 let mut k1_perturbed = k1.clone();
631 let mut k2_perturbed = k2.clone();
632 let mut k3_perturbed = k3.clone();
633
634 if j == 0 {
635 k1_perturbed[i] += eps;
636 } else if j == 1 {
637 k2_perturbed[i] += eps;
638 } else {
639 k3_perturbed[i] += eps;
640 }
641
642 // Evaluate function with perturbed values
643 let f1_perturbed = f(t1, k1_perturbed.view());
644 let f2_perturbed = f(t2, k2_perturbed.view());
645 let f3_perturbed = f(t3, k3_perturbed.view());
646 func_evals += 3;
647
648 // Compute perturbed residuals
649 let mut r1_perturbed = k1_perturbed.clone();
650 r1_perturbed -= &y;
651 r1_perturbed = r1_perturbed
652 - (&f1_perturbed * (h * a11)
653 + &f2_perturbed * (h * a12)
654 + &f3_perturbed * (h * a13));
655
656 let mut r2_perturbed = k2_perturbed.clone();
657 r2_perturbed -= &y;
658 r2_perturbed = r2_perturbed
659 - (&f1_perturbed * (h * a21)
660 + &f2_perturbed * (h * a22)
661 + &f3_perturbed * (h * a23));
662
663 let mut r3_perturbed = k3_perturbed.clone();
664 r3_perturbed -= &y;
665 r3_perturbed = r3_perturbed
666 - (&f1_perturbed * (h * a31)
667 + &f2_perturbed * (h * a32)
668 + &f3_perturbed * (h * a33));
669
670 // Finite difference approximation of Jacobian
671 jac[[0, j]] = (r1_perturbed[i] - r1[i]) / eps;
672 jac[[1, j]] = (r2_perturbed[i] - r2[i]) / eps;
673 jac[[2, j]] = (r3_perturbed[i] - r3[i]) / eps;
674 }
675
676 // Solve 3x3 system for this component
677 n_lu += 1;
678
679 // Augmented matrix for Gaussian elimination
680 let mut aug = Array2::<F>::zeros((3, 4));
681 for j in 0..3 {
682 for k in 0..3 {
683 aug[[j, k]] = jac[[j, k]];
684 }
685 aug[[j, 3]] = ri[j];
686 }
687
688 // Gaussian elimination with partial pivoting (for 3x3 system)
689 for j in 0..3 {
690 // Find pivot
691 let mut max_idx = j;
692 let mut max_val = aug[[j, j]].abs();
693
694 for k in j + 1..3 {
695 if aug[[k, j]].abs() > max_val {
696 max_idx = k;
697 max_val = aug[[k, j]].abs();
698 }
699 }
700
701 // Check for singularity
702 if max_val < F::from_f64(1e-10).unwrap() {
703 // Nearly singular matrix, reduce step size and try again
704 h *= F::from_f64(0.5).unwrap();
705 if h < min_step {
706 return Err(IntegrateError::ConvergenceError(
707 "Newton iteration failed to converge with minimum step size"
708 .to_string(),
709 ));
710 }
711 iter_count = 0;
712 continue;
713 }
714
715 // Swap rows if necessary
716 if max_idx != j {
717 for k in 0..4 {
718 let temp = aug[[j, k]];
719 aug[[j, k]] = aug[[max_idx, k]];
720 aug[[max_idx, k]] = temp;
721 }
722 }
723
724 // Eliminate below
725 for k in j + 1..3 {
726 let factor = aug[[k, j]] / aug[[j, j]];
727 for l in j..4 {
728 aug[[k, l]] = aug[[k, l]] - factor * aug[[j, l]];
729 }
730 }
731 }
732
733 // Back substitution
734 let mut delta = Array1::<F>::zeros(3);
735 for j in (0..3).rev() {
736 let mut sum = aug[[j, 3]];
737 for k in j + 1..3 {
738 sum -= aug[[j, k]] * delta[k];
739 }
740 delta[j] = sum / aug[[j, j]];
741 }
742
743 // Update stage values
744 k1[i] -= delta[0];
745 k2[i] -= delta[1];
746 k3[i] -= delta[2];
747 }
748
749 iter_count += 1;
750 }
751
752 if converged {
753 // Step accepted
754 // Compute the next state using the Butcher tableau weights
755 let f1 = f(t1, k1.view());
756 let f2 = f(t2, k2.view());
757 let f3 = f(t3, k3.view());
758 func_evals += 3;
759
760 let y_next = &y + &(&f1 * (h * b1) + &f2 * (h * b2) + &f3 * (h * b3));
761
762 // Update state
763 t += h;
764 y = y_next;
765
766 // Store results
767 t_values.push(t);
768 y_values.push(y.clone());
769
770 step_count += 1;
771 accepted_steps += 1;
772
773 // Adjust step size based on convergence
774 if iter_count <= 2 {
775 // Fast convergence, increase step size
776 h *= F::from_f64(1.2).unwrap().min(F::from_f64(5.0).unwrap());
777 } else if iter_count >= max_newton_iters - 1 {
778 // Slow convergence, decrease step size
779 h *= F::from_f64(0.8).unwrap().max(F::from_f64(0.2).unwrap());
780 }
781 } else {
782 // Step rejected, reduce step size
783 h *= F::from_f64(0.5).unwrap();
784 if h < min_step {
785 return Err(IntegrateError::ConvergenceError(
786 "Newton iteration failed to converge with minimum step size".to_string(),
787 ));
788 }
789 rejected_steps += 1;
790 }
791 }
792
793 let success = t >= t_end;
794 let message = if !success {
795 Some(format!(
796 "Maximum number of steps ({}) reached",
797 opts.max_steps
798 ))
799 } else {
800 None
801 };
802
803 // Return the solution
804 Ok(ODEResult {
805 t: t_values,
806 y: y_values,
807 success,
808 message,
809 n_eval: func_evals,
810 n_steps: step_count,
811 n_accepted: accepted_steps,
812 n_rejected: rejected_steps,
813 n_lu,
814 n_jac,
815 method: ODEMethod::Radau,
816 })
817}