differential-equations 0.6.1

//! Adaptive IRK for ODEs
use crate::{
    error::Error,
    interpolate::{Interpolation, cubic_hermite_interpolate},
    linalg::Matrix,
    methods::h_init::InitialStepSize,
    methods::{Adaptive, ImplicitRungeKutta, Ordinary},
    ode::{ODE, OrdinaryNumericalMethod},
    stats::Evals,
    status::Status,
    traits::{Real, State},
    utils::{constrain_step_size, validate_step_size_parameters},
};

impl<T: Real, Y: State<T>, const O: usize, const S: usize, const I: usize>
    OrdinaryNumericalMethod<T, Y> for ImplicitRungeKutta<Ordinary, Adaptive, T, Y, O, S, I>
{
    fn init<F>(&mut self, ode: &F, t0: T, tf: T, y0: &Y) -> Result<Evals, Error<T, Y>>
    where
        F: ODE<T, Y> + ?Sized,
    {
        let mut evals = Evals::new();

        // Compute h0 if not set
        if self.h0 == T::zero() {
            // Implicit initial step size heuristic
            self.h0 = InitialStepSize::<Ordinary>::compute(
                ode, t0, tf, y0, self.order, &self.rtol, &self.atol, self.h_min, self.h_max,
                &mut evals,
            );
        }

        // Validate step size bounds
        match validate_step_size_parameters::<T, Y>(self.h0, self.h_min, self.h_max, t0, tf) {
            Ok(h0) => self.h = (self.filter)(h0),
            Err(status) => return Err(status),
        }

        // Stats
        self.stiffness_counter = 0;
        self.newton_iterations = 0;
        self.jacobian_evaluations = 0;
        self.lu_decompositions = 0;

        // State
        self.t = t0;
        self.y = y0.clone();
        self.dydt = y0.zeros_like();
        self.y_prev = y0.clone();
        self.dydt_prev = y0.zeros_like();
        self.k = core::array::from_fn(|_| y0.zeros_like());
        self.z = core::array::from_fn(|_| y0.zeros_like());
        ode.diff(self.t, &self.y, &mut self.dydt);
        evals.function += 1;

        // Previous state
        self.t_prev = self.t;
        self.y_prev = self.y.clone();
        self.dydt_prev = self.dydt.clone();

        // Linear algebra workspace
        let dim = y0.len();
        let newton_system_size = self.stages * dim;
        self.stage_jacobians = core::array::from_fn(|_| Matrix::zeros(dim, dim));
        self.newton_matrix = Matrix::zeros(newton_system_size, newton_system_size);
        // Use State<T> storage for RHS and solution vectors
        self.rhs_newton = vec![T::zero(); newton_system_size];
        self.delta_k_vec = vec![T::zero(); newton_system_size];
        self.jacobian_age = 0;

        // Status
        self.status = Status::Initialized;

        Ok(evals)
    }

    fn step<F>(&mut self, ode: &F) -> Result<Evals, Error<T, Y>>
    where
        F: ODE<T, Y> + ?Sized,
    {
        let mut evals = Evals::new();

        // Step size guard
        if self.h.abs() < self.h_prev.abs() * T::from_f64(1e-14).unwrap() {
            self.status = Status::Error(Error::StepSize {
                t: self.t,
                y: self.y.clone(),
            });
            return Err(Error::StepSize {
                t: self.t,
                y: self.y.clone(),
            });
        }

        // Max steps guard
        if self.steps >= self.max_steps {
            self.status = Status::Error(Error::MaxSteps {
                t: self.t,
                y: self.y.clone(),
            });
            return Err(Error::MaxSteps {
                t: self.t,
                y: self.y.clone(),
            });
        }
        self.steps += 1;

        // Initial stage guesses: copy current state
        let dim = self.y.len();
        for i in 0..self.stages {
            self.z[i] = self.y.clone();
        }

        // Newton solve for F(z) = z - y_n - h*A*f(z) = 0
        let mut newton_converged = false;
        let mut newton_iter = 0;

        // Track increment norm
        let mut increment_norm = T::infinity();

        while !newton_converged && newton_iter < self.max_newton_iter {
            newton_iter += 1;
            self.newton_iterations += 1;
            evals.newton += 1;

            // Evaluate f at stage guesses
            for i in 0..self.stages {
                ode.diff(self.t + self.c[i] * self.h, &self.z[i], &mut self.k[i]);
            }
            evals.function += self.stages;

            // Residual and max-norm
            let mut residual_norm = T::zero();
            for i in 0..self.stages {
                // Start with z_i - y_n
                let mut residual = self.z[i].minus(&self.y);

                // Subtract h*sum(a_ij * f_j)
                for j in 0..self.stages {
                    residual.add_scaled(-(self.a[i][j] * self.h), &self.k[j]);
                }

                // Infinity norm and RHS
                for row_idx in 0..dim {
                    let res_val = residual.get_component(row_idx);
                    residual_norm = residual_norm.max(res_val.abs());
                    // Store residual in Newton RHS (negative for solving delta_z)
                    self.rhs_newton[i * dim + row_idx] = -res_val;
                }
            }

            // Converged by residual
            if residual_norm < self.newton_tol {
                newton_converged = true;
                break;
            }

            // Converged by increment
            if newton_iter > 1 && increment_norm < self.newton_tol {
                newton_converged = true;
                break;
            }

            // Refresh Jacobians if needed
            if newton_iter == 1 || self.jacobian_age > 3 {
                // Stage Jacobians
                for i in 0..self.stages {
                    ode.jacobian(
                        self.t + self.c[i] * self.h,
                        &self.z[i],
                        &mut self.stage_jacobians[i],
                    );
                    evals.jacobian += 1;
                }

                // Build Newton matrix: I - h*(A ⊗ J)
                // Zero the block Newton matrix (ensure Full storage)
                let nsys = self.stages * dim;
                let mut nm = Matrix::zeros(nsys, nsys);
                // Fill blocks
                for i in 0..self.stages {
                    for j in 0..self.stages {
                        let scale = -self.h * self.a[i][j];
                        for r in 0..dim {
                            for c_col in 0..dim {
                                nm[(i * dim + r, j * dim + c_col)] =
                                    self.stage_jacobians[j][(r, c_col)] * scale;
                            }
                        }
                    }
                    // Add identity on block diagonal
                    for d_idx in 0..dim {
                        let idx = i * dim + d_idx;
                        nm[(idx, idx)] += T::one();
                    }
                }
                self.newton_matrix = nm;

                self.jacobian_age = 0;
            }
            self.jacobian_age += 1;

            // Solve (I - h*A⊗J) Δz = -F(z) using our LU in-place over a flat slice
            for i in 0..self.delta_k_vec.len() {
                self.delta_k_vec[i] = self.rhs_newton[i];
            }
            self.newton_matrix
                .lin_solve_mut(&mut self.delta_k_vec[..])
                .map_err(|e| crate::error::Error::LinearAlgebra {
                    t: self.t,
                    y: self.y.clone(),
                    msg: e.to_string(),
                })?;
            self.lu_decompositions += 1;

            // Update z_i and increment norm
            increment_norm = T::zero();
            for i in 0..self.stages {
                for row_idx in 0..dim {
                    let delta_val = self.delta_k_vec[i * dim + row_idx];
                    let current_z = self.z[i].get_component(row_idx);
                    self.z[i].set_component(row_idx, current_z + delta_val);
                    // Calculate infinity norm of increment
                    increment_norm = increment_norm.max(delta_val.abs());
                }
            }

            // Next loop will re-check
        }

        // Newton failed to converge
        if !newton_converged {
            // Reduce h and retry later
            self.h *= T::from_f64(0.25).unwrap();
            self.h = constrain_step_size(self.h, self.h_min, self.h_max);
            self.h = (self.filter)(self.h);
            self.status = Status::RejectedStep;
            self.stiffness_counter += 1;

            if self.stiffness_counter >= self.max_rejects {
                self.status = Status::Error(Error::Stiffness {
                    t: self.t,
                    y: self.y.clone(),
                });
                return Err(Error::Stiffness {
                    t: self.t,
                    y: self.y.clone(),
                });
            }
            return Ok(evals);
        }

        // Final stage derivatives
        for i in 0..self.stages {
            ode.diff(self.t + self.c[i] * self.h, &self.z[i], &mut self.k[i]);
        }
        evals.function += self.stages;

        // y_{n+1} = y_n + h Σ b_i f_i
        let mut y_new = self.y.clone();
        for i in 0..self.stages {
            y_new.add_scaled(self.b[i] * self.h, &self.k[i]);
        }

        // Embedded error estimate (bh)
        let mut err_norm = T::zero();
        let bh = &self.bh.unwrap();

        // Lower-order solution
        let mut y_low = self.y.clone();
        for i in 0..self.stages {
            y_low.add_scaled(bh[i] * self.h, &self.k[i]);
        }

        // Weighted max-norm
        let dim = self.y.len();
        for n in 0..dim {
            let scale = self.atol[n]
                + self.rtol[n]
                    * self
                        .y
                        .get_component(n)
                        .abs()
                        .max(y_new.get_component(n).abs());
            if scale > T::zero() {
                err_norm =
                    err_norm.max(((y_new.get_component(n) - y_low.get_component(n)) / scale).abs());
            }
        }

        // Avoid vanishing error
        err_norm = err_norm.max(T::default_epsilon() * T::from_f64(100.0).unwrap());

        // Step scale factor
        let order = T::from_usize(self.order).unwrap();
        let error_exponent = T::one() / order;
        let mut scale = self.safety_factor * err_norm.powf(-error_exponent);

        // Clamp scale factor
        scale = scale.max(self.min_scale).min(self.max_scale);

        // Accept/reject
        if err_norm <= T::one() {
            // Accepted
            self.status = Status::Solving;

            // Log previous
            self.t_prev = self.t;
            self.y_prev = self.y.clone();
            self.dydt_prev = self.dydt.clone();
            self.h_prev = self.h;

            // Advance state
            self.t += self.h;
            self.y = y_new;

            // Next-step derivative
            ode.diff(self.t, &self.y, &mut self.dydt);
            evals.function += 1;

            // If we were rejecting, limit growth
            if let Status::RejectedStep = self.status {
                self.stiffness_counter = 0;

                // Avoid oscillations
                scale = scale.min(T::one());
            }
        } else {
            // Rejected
            self.status = Status::RejectedStep;
            self.stiffness_counter += 1;

            // Too many rejections
            if self.stiffness_counter >= self.max_rejects {
                self.status = Status::Error(Error::Stiffness {
                    t: self.t,
                    y: self.y.clone(),
                });
                return Err(Error::Stiffness {
                    t: self.t,
                    y: self.y.clone(),
                });
            }
        }

        // Update h
        self.h *= scale;

        // Constrain h
        self.h = constrain_step_size(self.h, self.h_min, self.h_max);

        // Apply step size filter
        self.h = (self.filter)(self.h);

        Ok(evals)
    }

    fn t(&self) -> T {
        self.t
    }
    fn y(&self) -> &Y {
        &self.y
    }
    fn t_prev(&self) -> T {
        self.t_prev
    }
    fn y_prev(&self) -> &Y {
        &self.y_prev
    }
    fn h(&self) -> T {
        self.h
    }
    fn set_h(&mut self, h: T) {
        self.h = (self.filter)(h);
    }
    fn status(&self) -> &Status<T, Y> {
        &self.status
    }
    fn set_status(&mut self, status: Status<T, Y>) {
        self.status = status;
    }
}

impl<T: Real, Y: State<T>, const O: usize, const S: usize, const I: usize> Interpolation<T, Y>
    for ImplicitRungeKutta<Ordinary, Adaptive, T, Y, O, S, I>
{
    fn interpolate(&mut self, t_interp: T) -> Result<Y, Error<T, Y>> {
        // Check if t is within bounds
        if t_interp < self.t_prev || t_interp > self.t {
            return Err(Error::OutOfBounds {
                t_interp,
                t_prev: self.t_prev,
                t_curr: self.t,
            });
        }

        // Use cubic Hermite interpolation
        let y_interp = cubic_hermite_interpolate(
            self.t_prev,
            self.t,
            &self.y_prev,
            &self.y,
            &self.dydt_prev,
            &self.dydt,
            t_interp,
        );

        Ok(y_interp)
    }
}