nyx-space 2.4.0

/*
    Nyx, blazing fast astrodynamics
    Copyright (C) 2018-onwards Christopher Rabotin <christopher.rabotin@gmail.com>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published
    by the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

use anise::analysis::prelude::OrbitalElement;
use anise::astro::orbit_gradient::{OrbitGrad, OrbitalElementPartials};
use anise::prelude::{Frame, Orbit};

use super::{AstroError, AstroPhysicsSnafu};
use crate::cosmic::NotHyperbolicSnafu;
use crate::linalg::{Matrix2, Matrix3, Vector2, Vector3};
use crate::md::objective::Objective;
use crate::md::{AstroSnafu, StateParameter, TargetingError};
use crate::time::{Duration, Epoch, Unit};
use crate::utils::between_pm_180;
use hyperdual::linalg::norm;
use hyperdual::{Float, OHyperdual};

use snafu::{ResultExt, ensure};
use std::convert::From;
use std::fmt;

/// Stores a B-Plane
#[derive(Copy, Clone, Debug)]

pub struct BPlane {
    /// The $B_T$ component, in kilometers
    pub b_t_km: OrbitalElementPartials,
    /// The $B_R$ component, in kilometers
    pub b_r_km: OrbitalElementPartials,
    /// The Linearized Time of Flight
    pub ltof_s: OrbitalElementPartials,
    /// The B-Plane rotation matrix
    pub str_dcm: Matrix3<f64>,
    /// The frame in which this B Plane was computed
    pub frame: Frame,
    /// The time of computation
    pub epoch: Epoch,
}

impl BPlane {
    /// Returns a newly define B-Plane if the orbit is hyperbolic and already in Dual form
    pub fn from_dual(orbit: OrbitGrad) -> Result<Self, AstroError> {
        ensure!(
            orbit.ecc().context(AstroPhysicsSnafu)?.real() > 1.0,
            NotHyperbolicSnafu
        );

        let one = OHyperdual::from(1.0);
        let zero = OHyperdual::from(0.0);

        let e_hat =
            orbit.evec().context(AstroPhysicsSnafu)? / orbit.ecc().context(AstroPhysicsSnafu)?.dual;
        let h_hat = orbit.hvec() / orbit.hmag().dual;
        let n_hat = h_hat.cross(&e_hat);

        // The reals implementation (which was initially validated) was:
        // let s = e_hat / orbit.ecc() + (1.0 - (1.0 / orbit.ecc()).powi(2)).sqrt() * n_hat;
        // let s_hat = s / s.norm();

        let ecc = orbit.ecc().context(AstroPhysicsSnafu)?;

        let incoming_asymptote_fact = (one - (one / ecc.dual).powi(2)).sqrt();

        let s = Vector3::new(
            e_hat[0] / ecc.dual + incoming_asymptote_fact * n_hat[0],
            e_hat[1] / ecc.dual + incoming_asymptote_fact * n_hat[1],
            e_hat[2] / ecc.dual + incoming_asymptote_fact * n_hat[2],
        );

        let s_hat = s / norm(&s); // Just to make sure to renormalize everything

        // The reals implementation (which was initially validated) was:
        // let b_vec = orbit.semi_minor_axis()
        //     * ((1.0 - (1.0 / orbit.ecc()).powi(2)).sqrt() * e_hat
        //         - (1.0 / orbit.ecc() * n_hat));
        let semi_minor_axis = orbit.semi_minor_axis_km().context(AstroPhysicsSnafu)?;

        let b_vec = Vector3::new(
            semi_minor_axis.dual
                * (incoming_asymptote_fact * e_hat[0] - ((one / ecc.dual) * n_hat[0])),
            semi_minor_axis.dual
                * (incoming_asymptote_fact * e_hat[1] - ((one / ecc.dual) * n_hat[1])),
            semi_minor_axis.dual
                * (incoming_asymptote_fact * e_hat[2] - ((one / ecc.dual) * n_hat[2])),
        );

        let t = s_hat.cross(&Vector3::new(zero, zero, one));
        let t_hat = t / norm(&t);
        let r_hat = s_hat.cross(&t_hat);

        // Build the rotation matrix from inertial to B Plane
        let str_rot = Matrix3::new(
            s_hat[0].real(),
            s_hat[1].real(),
            s_hat[2].real(),
            t_hat[0].real(),
            t_hat[1].real(),
            t_hat[2].real(),
            r_hat[0].real(),
            r_hat[1].real(),
            r_hat[2].real(),
        );

        Ok(BPlane {
            b_r_km: OrbitalElementPartials {
                dual: b_vec.dot(&r_hat),
                oe: OrbitalElement::Custom,
            },
            b_t_km: OrbitalElementPartials {
                dual: b_vec.dot(&t_hat),
                oe: OrbitalElement::Custom,
            },
            ltof_s: OrbitalElementPartials {
                dual: b_vec.dot(&s_hat) / orbit.vmag_km_s().dual,
                oe: OrbitalElement::Custom,
            },
            str_dcm: str_rot,
            frame: orbit.frame,
            epoch: orbit.epoch,
        })
    }

    /// Returns a newly defined B-Plane if the orbit is hyperbolic.
    pub fn new(orbit: Orbit) -> Result<Self, AstroError> {
        // Convert to OrbitGrad so we can target it
        Self::from_dual(OrbitGrad::from(orbit))
    }

    /// Returns the DCM to convert to the B Plane from the inertial frame
    pub fn inertial_to_bplane(&self) -> Matrix3<f64> {
        self.str_dcm
    }

    /// Returns the Jacobian of the B plane (BR, BT, LTOF) with respect to the velocity
    pub fn jacobian(&self) -> Matrix3<f64> {
        Matrix3::new(
            self.b_r_km.wrt_vx(),
            self.b_r_km.wrt_vy(),
            self.b_r_km.wrt_vz(),
            self.b_t_km.wrt_vx(),
            self.b_t_km.wrt_vy(),
            self.b_t_km.wrt_vz(),
            self.ltof_s.wrt_vx(),
            self.ltof_s.wrt_vy(),
            self.ltof_s.wrt_vz(),
        )
    }

    /// Returns the Jacobian of the B plane (BR, BT) with respect to two of the velocity components
    pub fn jacobian2(&self, invariant: StateParameter) -> Result<Matrix2<f64>, AstroError> {
        match invariant {
            StateParameter::Element(OrbitalElement::VX) => Ok(Matrix2::new(
                self.b_r_km.wrt_vy(),
                self.b_r_km.wrt_vz(),
                self.b_t_km.wrt_vy(),
                self.b_t_km.wrt_vz(),
            )),
            StateParameter::Element(OrbitalElement::VY) => Ok(Matrix2::new(
                self.b_r_km.wrt_vx(),
                self.b_r_km.wrt_vz(),
                self.b_t_km.wrt_vx(),
                self.b_t_km.wrt_vz(),
            )),
            StateParameter::Element(OrbitalElement::VZ) => Ok(Matrix2::new(
                self.b_r_km.wrt_vx(),
                self.b_r_km.wrt_vy(),
                self.b_t_km.wrt_vx(),
                self.b_t_km.wrt_vy(),
            )),
            _ => Err(AstroError::BPlaneInvariant),
        }
    }
}

impl BPlane {
    /// Returns the B dot T, in km
    pub fn b_dot_t_km(&self) -> f64 {
        self.b_t_km.real()
    }

    /// Returns the B dot R, in km
    pub fn b_dot_r_km(&self) -> f64 {
        self.b_r_km.real()
    }

    /// Returns the linearized time of flight, in seconds
    ///
    /// NOTE:
    /// Although the LTOF should allow to build a 3x3 Jacobian for a differential corrector
    /// it's historically super finicky and one will typically have better results by cascading
    /// DCs which target different parameters instead of trying to use the LTOF.
    pub fn ltof(&self) -> Duration {
        self.ltof_s.real() * Unit::Second
    }

    /// Returns the B plane angle in degrees between -180 and 180
    pub fn angle_deg(&self) -> f64 {
        between_pm_180(self.b_dot_r_km().atan2(self.b_dot_t_km()).to_degrees())
    }

    /// Returns the B plane vector magnitude, in kilometers
    pub fn magnitude_km(&self) -> f64 {
        (self.b_dot_t_km().powi(2) + self.b_dot_r_km().powi(2)).sqrt()
    }
}

impl fmt::Display for BPlane {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "[{}] {} B-Plane: B∙R = {:.3} km\tB∙T = {:.3} km\tAngle = {:.3} deg",
            self.frame,
            self.epoch,
            self.b_dot_r_km(),
            self.b_dot_t_km(),
            self.angle_deg()
        )
    }
}

#[derive(Copy, Clone, Debug)]
pub struct BPlaneTarget {
    /// The $B_T$ component, in kilometers
    pub b_t_km: f64,
    /// The $B_R$ component, in kilometers
    pub b_r_km: f64,
    /// The Linearized Time of Flight, in seconds
    pub ltof_s: f64,

    /// The tolerance on the $B_T$ component, in kilometers
    pub tol_b_t_km: f64,
    /// The tolerance on the $B_R$ component, in kilometers
    pub tol_b_r_km: f64,
    /// The tolerance on the Linearized Time of Flight, in seconds
    pub tol_ltof_s: f64,
}

impl BPlaneTarget {
    /// Initializes a new B Plane target with only the targets and the default tolerances.
    /// Default tolerances are 1 millimeter in positions and 1 second in LTOF
    pub fn from_targets(b_t_km: f64, b_r_km: f64, ltof: Duration) -> Self {
        let tol_ltof: Duration = 6.0 * Unit::Hour;
        Self {
            b_t_km,
            b_r_km,
            ltof_s: ltof.to_seconds(),
            tol_b_t_km: 1e-6,
            tol_b_r_km: 1e-6,
            tol_ltof_s: tol_ltof.to_seconds(),
        }
    }

    /// Initializes a new B Plane target with only the B Plane targets (not LTOF constraint) and the default tolerances.
    /// Default tolerances are 1 millimeter in positions. Here, the LTOF tolerance is set to 100 days.
    pub fn from_bt_br(b_t_km: f64, b_r_km: f64) -> Self {
        let ltof_tol: Duration = 100 * Unit::Day;
        Self {
            b_t_km,
            b_r_km,
            ltof_s: 0.0,
            tol_b_t_km: 1e-6,
            tol_b_r_km: 1e-6,
            tol_ltof_s: ltof_tol.to_seconds(),
        }
    }

    pub fn ltof_target_set(&self) -> bool {
        self.ltof_s.abs() > 1e-10
    }

    pub fn to_objectives(self) -> [Objective; 2] {
        self.to_objectives_with_tolerance(1.0)
    }

    pub fn to_objectives_with_tolerance(self, tol_km: f64) -> [Objective; 2] {
        [
            Objective::within_tolerance(StateParameter::BdotR(), self.b_r_km, tol_km),
            Objective::within_tolerance(StateParameter::BdotT(), self.b_t_km, tol_km),
        ]
    }

    /// Includes the linearized time of flight as an objective
    pub fn to_all_objectives_with_tolerance(self, tol_km: f64) -> [Objective; 3] {
        [
            Objective::within_tolerance(StateParameter::BdotR(), self.b_r_km, tol_km),
            Objective::within_tolerance(StateParameter::BdotT(), self.b_t_km, tol_km),
            Objective::within_tolerance(
                StateParameter::BLTOF(),
                self.ltof_s,
                self.tol_ltof_s * 1e5,
            ),
        ]
    }
}

impl fmt::Display for BPlaneTarget {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "B-Plane target: B∙R = {:.3} km (+/- {:.1} m)\tB∙T = {:.3} km (+/- {:.1} m)",
            self.b_r_km,
            self.tol_b_r_km * 1e-3,
            self.b_t_km,
            self.tol_b_t_km * 1e-3,
        )
    }
}

/// Returns the Delta V (in km/s) needed to achieve the B Plane specified by B dot R and B dot T.
/// If no LTOF target is set, this method will fix VX, VY and VZ successively and use the minimum of those as a seed for the LTOF variation finding.
/// If the 3x3 search is worse than any of the 2x2s, then a 2x2 will be returned.
/// This uses the hyperdual formulation of the Jacobian and will also vary the linearize time of flight (LTOF).
pub fn try_achieve_b_plane(
    orbit: Orbit,
    target: BPlaneTarget,
) -> Result<(Vector3<f64>, BPlane), TargetingError> {
    let mut total_dv = Vector3::zeros();
    let mut attempt_no = 0;
    let max_iter = 10;

    let mut real_orbit = orbit;
    let mut prev_b_plane_err = f64::INFINITY;

    if !target.ltof_target_set() {
        // If no LTOF is targeted, we'll solve this with a least squared approach.
        loop {
            if attempt_no > max_iter {
                return Err(TargetingError::TooManyIterations);
            }

            // Build current B Plane
            let b_plane = BPlane::new(real_orbit).context(AstroSnafu)?;

            // Check convergence
            let br_err = target.b_r_km - b_plane.b_dot_r_km();
            let bt_err = target.b_t_km - b_plane.b_dot_t_km();

            if br_err.abs() < target.tol_b_r_km && bt_err.abs() < target.tol_b_t_km {
                return Ok((total_dv, b_plane));
            }

            // Build the error vector
            let b_plane_err = Vector2::new(br_err, bt_err);

            if b_plane_err.norm() >= prev_b_plane_err {
                // If the error is not going down, we'll raise an error
                return Err(TargetingError::CorrectionIneffective {
                    prev_val: prev_b_plane_err,
                    cur_val: b_plane_err.norm(),
                    action: "Delta-V correction is ineffective at reducing the B-Plane error",
                });
            }
            prev_b_plane_err = b_plane_err.norm();

            // Grab the first two rows of the Jacobian (discard the rest).
            let full_jac = b_plane.jacobian();
            let jac = full_jac.fixed_rows::<2>(0);
            // Solve the Least Squares / compute the delta-v
            let dv = jac.transpose() * (jac * jac.transpose()).try_inverse().unwrap() * b_plane_err;

            total_dv[0] += dv[0];
            total_dv[1] += dv[1];
            total_dv[2] += dv[2];

            // Rebuild a new orbit
            real_orbit.velocity_km_s.x += dv[0];
            real_orbit.velocity_km_s.y += dv[1];
            real_orbit.velocity_km_s.z += dv[2];

            attempt_no += 1;
        }
    } else {
        // The LTOF targeting seems to break often, but it's still implemented
        loop {
            if attempt_no > max_iter {
                return Err(TargetingError::TooManyIterations);
            }

            // Build current B Plane
            let b_plane = BPlane::new(real_orbit).context(AstroSnafu)?;

            // Check convergence
            let br_err = target.b_r_km - b_plane.b_dot_r_km();
            let bt_err = target.b_t_km - b_plane.b_dot_t_km();
            let ltof_err = target.ltof_s - b_plane.ltof_s.real();

            if br_err.abs() < target.tol_b_r_km
                && bt_err.abs() < target.tol_b_t_km
                && ltof_err.abs() < target.tol_ltof_s
            {
                return Ok((total_dv, b_plane));
            }

            // Build the error vector
            let b_plane_err = Vector3::new(br_err, bt_err, ltof_err);

            if b_plane_err.norm() >= prev_b_plane_err {
                return Err(TargetingError::CorrectionIneffective {
                    prev_val: prev_b_plane_err,
                    cur_val: b_plane_err.norm(),
                    action: "LTOF enabled correction is failing. Try to not set an LTOF target.",
                });
            }
            prev_b_plane_err = b_plane_err.norm();

            // Compute the delta-v
            let dv = b_plane.jacobian() * b_plane_err;

            total_dv[0] += dv[0];
            total_dv[1] += dv[1];
            total_dv[2] += dv[2];

            // Rebuild a new orbit
            real_orbit.velocity_km_s.x += dv[0];
            real_orbit.velocity_km_s.y += dv[1];
            real_orbit.velocity_km_s.z += dv[2];

            attempt_no += 1;
        }
    }
}