differential-equations 0.6.1

//! Defines system of differential equations for numerical solvers.
//! The NumericalMethods use this trait to take a input system from the user and solve
//! Includes a differential equation. Event handling is provided via the separate `Event` trait.

use crate::{
    linalg::Matrix,
    traits::{DefaultState, Real, State},
};

/// ODE Trait for Differential Equations
///
/// ODE trait defines the differential equation dydt = f(t, y) for the solver.
/// The differential equation is used to solve the ordinary differential equation.
///
/// # Impl
/// * `diff`    - Differential Equation dydt = f(t, y) in form f(t, &y, &mut dydt).
/// * `jacobian` - Jacobian matrix J = df/dy for the system of equations.
///
/// Note that the jacobian function is optional and can be left out when implementing.
pub trait ODE<T = f64, Y = DefaultState<T>>
where
    T: Real,
    Y: State<T>,
{
    /// Differential Equation dydt = f(t, y)
    ///
    /// An ordinary differential equation (ODE) takes a independent variable
    /// which in this case is 't' as it is typically time and a dependent variable
    /// which is a vector of values 'y'. The ODE returns the derivative of the
    /// dependent variable 'y' with respect to the independent variable 't' as
    /// dydt = f(t, y).
    ///
    /// For efficiency and ergonomics the derivative is calculated from an argument
    /// of a mutable reference to the derivative vector dydt. This allows for a
    /// derivatives to be calculated in place which is more efficient as iterative
    /// ODE solvers require the derivative to be calculated at each step without
    /// regard to the previous value.
    ///
    /// # Arguments
    /// * `t`    - Independent variable point.
    /// * `y`    - Dependent variable point.
    /// * `dydt` - Derivative point.
    ///
    fn diff(&self, t: T, y: &Y, dydt: &mut Y);

    /// jacobian matrix J = df/dy
    ///
    /// The jacobian matrix is a matrix of partial derivatives of a vector-valued function.
    /// It describes the local behavior of the system of equations and can be used to improve
    /// the efficiency of certain solvers by providing information about the local behavior
    /// of the system of equations.
    ///
    /// By default, this method uses a finite difference approximation.
    /// Users can override this with an analytical implementation for better efficiency.
    ///
    /// # Arguments
    /// * `t` - Independent variable grid point.
    /// * `y` - Dependent variable vector.
    /// * `j` - jacobian matrix. This matrix should be pre-sized by the caller to `dim x dim` where `dim = y.len()`.
    ///
    fn jacobian(&self, t: T, y: &Y, j: &mut Matrix<T>) {
        // Default implementation using forward finite differences
        let dim = y.len();
        let mut y_perturbed = y.clone();
        let mut f_perturbed = y.zeros_like();
        let mut f_origin = y.zeros_like();

        // Compute the unperturbed derivative
        self.diff(t, y, &mut f_origin);

        // Use sqrt of machine epsilon for finite differences
        let eps = T::default_epsilon().sqrt();

        // For each column of the jacobian
        for j_col in 0..dim {
            // Get the original value
            let y_original_j = y.get_component(j_col);

            // Calculate perturbation size (max of component magnitude or 1.0)
            let perturbation = eps * y_original_j.abs().max(T::one());

            // Perturb the component
            y_perturbed.copy_from_state(y);
            y_perturbed.set_component(j_col, y_original_j + perturbation);

            // Evaluate function with perturbed value
            self.diff(t, &y_perturbed, &mut f_perturbed);

            // Compute finite difference approximation for this column
            for i_row in 0..dim {
                j[(i_row, j_col)] = (f_perturbed.get_component(i_row)
                    - f_origin.get_component(i_row))
                    / perturbation;
            }
        }
    }
}

impl<EqType, T: Real, Y: State<T>> ODE<T, Y> for &EqType
where
    EqType: ODE<T, Y> + ?Sized,
{
    fn diff(&self, t: T, y: &Y, dydt: &mut Y) {
        (*self).diff(t, y, dydt);
    }

    fn jacobian(&self, t: T, y: &Y, j: &mut Matrix<T>) {
        (*self).jacobian(t, y, j);
    }
}