newton.rs

use std::{
    ops::{
        Neg,
        Add,
        AddAssign,
        Sub,
        SubAssign,
        Mul,
        MulAssign,
        Div,
        DivAssign,
    },
};
use ndarray::{
    self as nd,
    array,
};
use thiserror::Error;

#[derive(Error, Debug)]
pub enum NewtonError {
    #[error("rka: error bound could not be satisfied")]
    RKAErrorBound,
}
pub type NewtonResult<T> = Result<T, NewtonError>;

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct ThreeVector(pub f64, pub f64, pub f64);

impl ThreeVector {
    /// Generate from spherical coordinates (r, theta, phi) where phi is the
    /// azimuthal angle.
    pub fn from_angles(r: f64, theta: f64, phi: f64) -> Self {
        return r * Self(
            phi.cos() * theta.sin(),
            phi.sin() * theta.sin(),
            theta.cos(),
        );
    }

    /// Generate from cylindrical coordinates (r, phi, z).
    pub fn from_angles_cylindrical(r: f64, phi: f64, z: f64) -> Self {
        return Self(
            r * phi.cos(),
            r * phi.sin(),
            z
        );
    }

    /// Generate along a Cartesian axis.
    pub fn from_axis(r: f64, axis: Axis) -> Self {
        return match axis {
            Axis::X => Self(r, 0.0, 0.0),
            Axis::Y => Self(0.0, r, 0.0),
            Axis::Z => Self(0.0, 0.0, r),
        };
    }

    pub fn norm(&self) -> f64 {
        return (self.0.powi(2) + self.1.powi(2) + self.2.powi(2)).sqrt();
    }

    pub fn normalized(&self) -> Self {
        return *self / self.norm();
    }

    pub fn abs(&self) -> Self {
        return Self(self.0.abs(), self.1.abs(), self.2.abs());
    }

    pub fn to_array(&self) -> nd::Array1<f64> {
        return array![self.0, self.1, self.2];
    }

    pub fn into_array(self) -> nd::Array1<f64> {
        return array![self.0, self.1, self.2];
    }

    pub fn get_component<A>(&self, axis: A) -> f64
    where A: Into<Axis>
    {
        return match axis.into() {
            Axis::X => self.0,
            Axis::Y => self.1,
            Axis::Z => self.2,
        };
    }

    pub fn get_components_except<A>(&self, axis: A) -> (f64, f64)
    where A: Into<Axis>
    {
        return match axis.into() {
            Axis::X => (self.1, self.2),
            Axis::Y => (self.0, self.2),
            Axis::Z => (self.0, self.1),
        };
    }
}

impl Neg for ThreeVector {
    type Output = Self;
    fn neg(self) -> Self { Self(-self.0, -self.1, -self.2) }
}

impl Add<ThreeVector> for ThreeVector {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        return Self(self.0 + rhs.0, self.1 + rhs.1, self.2 + rhs.2);
    }
}

impl AddAssign<ThreeVector> for ThreeVector {
    fn add_assign(&mut self, rhs: Self) {
        self.0 += rhs.0;
        self.1 += rhs.1;
        self.2 += rhs.2;
    }
}

impl Sub<ThreeVector> for ThreeVector {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        return Self(self.0 - rhs.0, self.1 - rhs.1, self.2 - rhs.2);
    }
}

impl SubAssign<ThreeVector> for ThreeVector {
    fn sub_assign(&mut self, rhs: Self) {
        self.0 -= rhs.0;
        self.1 -= rhs.1;
        self.2 -= rhs.2;
    }
}

impl Mul<f64> for ThreeVector {
    type Output = Self;
    fn mul(self, rhs: f64) -> Self {
        return Self(self.0 * rhs, self.1 * rhs, self.2 * rhs);
    }
}

impl MulAssign<f64> for ThreeVector {
    fn mul_assign(&mut self, rhs: f64) {
        self.0 *= rhs;
        self.1 *= rhs;
        self.2 *= rhs;
    }
}

impl Mul<ThreeVector> for f64 {
    type Output = ThreeVector;
    fn mul(self, rhs: ThreeVector) -> ThreeVector {
        return ThreeVector(self * rhs.0, self * rhs.1, self * rhs.2);
    }
}

impl Div<f64> for ThreeVector {
    type Output = Self;
    fn div(self, rhs: f64) -> Self {
        return ThreeVector(self.0 / rhs, self.1 / rhs, self.2 / rhs);
    }
}

impl DivAssign<f64> for ThreeVector {
    fn div_assign(&mut self, rhs: f64) {
        self.0 /= rhs;
        self.1 /= rhs;
        self.2 /= rhs;
    }
}

#[derive(Copy, Clone, Debug)]
pub struct PhaseSpace {
    pub pos: ThreeVector,
    pub mom: ThreeVector,
}

impl PhaseSpace {
    pub fn norm(&self) -> f64 {
        return (self.pos.norm().powi(2) + self.mom.norm().powi(2)).sqrt();
    }

    pub fn normalized(&self) -> Self {
        return Self { pos: self.pos.normalized(), mom: self.mom.normalized() };
    }

    pub fn abs(&self) -> Self {
        return Self { pos: self.pos.abs(), mom: self.mom.abs() };
    }

    pub fn to_array(&self) -> nd::Array1<f64> {
        return array![
            self.pos.0, self.pos.1, self.pos.2,
            self.mom.0, self.mom.1, self.mom.2,
        ];
    }

    pub fn into_array(self) -> nd::Array1<f64> {
        return array![
            self.pos.0, self.pos.1, self.pos.2,
            self.mom.0, self.mom.1, self.mom.2,
        ];
    }

    pub fn get_component<A>(&self, axis: A) -> (f64, f64)
    where A: Into<Axis>
    {
        return match axis.into() {
            Axis::X => (self.pos.0, self.mom.0),
            Axis::Y => (self.pos.1, self.mom.1),
            Axis::Z => (self.pos.2, self.mom.2),
        };
    }

    pub fn get_components_except<A>(&self, axis: A) -> ((f64, f64), (f64, f64))
    where A: Into<Axis>
    {
        return match axis.into() {
            Axis::X => ((self.pos.1, self.pos.2), (self.mom.1, self.mom.2)),
            Axis::Y => ((self.pos.0, self.pos.2), (self.mom.0, self.mom.2)),
            Axis::Z => ((self.pos.0, self.pos.1), (self.mom.0, self.mom.1)),
        };
    }
}

impl Neg for PhaseSpace {
    type Output = Self;
    fn neg(self) -> Self { Self { pos: -self.pos, mom: -self.mom } }
}

impl Add<PhaseSpace> for PhaseSpace {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        return Self { pos: self.pos + rhs.pos, mom: self.mom + rhs.mom };
    }
}

impl Sub<PhaseSpace> for PhaseSpace {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        return Self { pos: self.pos - rhs.pos, mom: self.pos - rhs.mom };
    }
}

impl Mul<f64> for PhaseSpace {
    type Output = Self;
    fn mul(self, rhs: f64) -> Self {
        return Self { pos: self.pos * rhs, mom: self.mom * rhs };
    }
}

impl Mul<PhaseSpace> for f64 {
    type Output = PhaseSpace;
    fn mul(self, rhs: PhaseSpace) -> PhaseSpace {
        return PhaseSpace { pos: self * rhs.pos, mom: self * rhs.mom };
    }
}

impl Div<f64> for PhaseSpace {
    type Output = Self;
    fn div(self, rhs: f64) -> Self {
        return PhaseSpace { pos: self.pos / rhs, mom: self.mom / rhs };
    }
}

pub fn traj_as_array(traj: &[ThreeVector]) -> nd::Array2<f64> {
    let arrays: Vec<nd::Array1<f64>>
        = traj.iter().map(|vk| vk.to_array()).collect();
    return nd::stack(
        nd::Axis(1),
        &arrays.iter()
            .map(|ak| ak.view())
            .collect::<Vec<nd::ArrayView1<f64>>>(),
    ).unwrap();
}

pub fn phasespace_traj_as_array(traj: &[PhaseSpace]) -> nd::Array2<f64> {
    let arrays: Vec<nd::Array1<f64>>
        = traj.iter().map(|qk| qk.to_array()).collect();
    return nd::stack(
        nd::Axis(1),
        &arrays.iter()
            .map(|ak| ak.view())
            .collect::<Vec<nd::ArrayView1<f64>>>(),
    ).unwrap();
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Axis {
    X = 0,
    Y = 1,
    Z = 2,
}

impl From<usize> for Axis {
    fn from(axis: usize) -> Self {
        return match axis {
            0 => Axis::X,
            1 => Axis::Y,
            _ => Axis::Z,
        };
    }
}

impl From<Axis> for usize {
    fn from(axis: Axis) -> Self { axis as usize }
}

fn rk4_step<F>(x: f64, y: PhaseSpace, dx: f64, rhs: &F) -> PhaseSpace
where F: Fn(f64, PhaseSpace) -> PhaseSpace
{
    let x_half: f64 = x + dx / 2.0;
    let k1: PhaseSpace = rhs(x, y);
    let k2: PhaseSpace = rhs(x_half, y + dx / 2.0 * k1);
    let k3: PhaseSpace = rhs(x_half, y + dx / 2.0 * k2);
    let k4: PhaseSpace = rhs(x + dx, y + dx * k3);
    return y + dx / 6.0 * (k1 + 2.0 * (k2 + k3) + k4);
}

pub fn rka_step<F>(x: f64, y: PhaseSpace, dx: f64, rhs: &F, err: f64)
    -> NewtonResult<(f64, PhaseSpace, f64)>
where F: Fn(f64, PhaseSpace) -> PhaseSpace
{
    // define safety numbers for choosing next step size -- particular to rk4
    let safe1: f64 = 0.9;
    let safe2: f64 = 4.0;

    let mut dx_old: f64;
    let mut dx_new: f64 = dx;
    let (mut dx_cond1, mut dx_cond2): (f64, f64);
    let (mut dx_half, mut x_half, mut x_full): (f64, f64, f64);
    let (mut y_half, mut y_half2, mut y_full):
        (PhaseSpace, PhaseSpace, PhaseSpace);
    let (mut scale, mut diff, mut error_ratio): (PhaseSpace, PhaseSpace, f64);
    for _ in 0..100 {
        // take two half-sized steps
        dx_half = dx / 2.0;
        x_half = x + dx_half;
        y_half = rk4_step(x, y, dx_half, rhs);
        y_half2 = rk4_step(x_half, y_half, dx_half, rhs);

        // take one full-sized step
        x_full = x + dx;
        y_full = rk4_step(x, y, dx, rhs);

        // compute the estimated local truncation error
        scale = err * (y_half2.abs() + y_full.abs()) / 2.0;
        diff = (y_half2 - y_full).abs();
        error_ratio
            = scale.into_array().into_iter().zip(diff.into_array().into_iter())
            .map(|(s, d)| d / (s + f64::EPSILON))
            .max_by(|l, r| {
                match l.partial_cmp(r) {
                    Some(ord) => ord,
                    None => std::cmp::Ordering::Less,
                }
            }).unwrap();

        // estimate new step size (with safety factors)
        dx_old = dx_new;
        if error_ratio == 0.0 {
            dx_new = dx_old / safe2;
            continue;
        }
        dx_new = dx_old * error_ratio.powf(-0.2) * safe1;
        dx_cond1 = dx_old / safe2;
        dx_cond2 = dx_old * safe2;
        dx_new = if dx_cond1 > dx_new { dx_cond1 } else { dx_new };
        dx_new = if dx_cond2 < dx_new { dx_cond2 } else { dx_new };

        if error_ratio < 1.0 {
            return Ok((x_full, y_half2, dx_new));
        }
    }
    return Err(NewtonError::RKAErrorBound);
}

pub fn rka<F>(
    x_bounds: (f64, f64),
    y0: PhaseSpace,
    dx0: f64,
    rhs: F,
    epsilon: f64
) -> NewtonResult<(Vec<f64>, Vec<PhaseSpace>)>
where F: Fn(f64, PhaseSpace) -> PhaseSpace
{
    let mut x: Vec<f64> = vec![x_bounds.0];
    let mut x_prev: f64 = x_bounds.0;
    let mut x_next: f64;
    let mut dx: f64 = dx0;
    let mut y: Vec<PhaseSpace> = vec![y0];
    let mut y_prev: &PhaseSpace = y.last().unwrap();
    let mut y_next: PhaseSpace;
    let mut step: (f64, PhaseSpace, f64);
    while x_prev < x_bounds.1 {
        dx = dx.min(x_bounds.1 - x_prev);
        
        step = rka_step(x_prev, *y_prev, dx, &rhs, epsilon)?;
        x_next = step.0;
        y_next = step.1;
        dx = step.2;

        x.push(x_next);
        y.push(y_next);

        x_prev = x_next;
        y_prev = y.last().unwrap();
    }
    return Ok((x, y));
}