nyx_space/dynamics/

gravity_field.rs

/*
    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::errors::OrientationSnafu;
use anise::prelude::Almanac;
use snafu::ResultExt;

use crate::cosmic::{AstroPhysicsSnafu, Orbit};
use crate::dynamics::AccelModel;
use crate::io::gravity::GravityFieldData;
use crate::linalg::{DMatrix, Matrix3, U7, Vector3, Vector4};
use hyperdual::linalg::norm;
use hyperdual::{Float, OHyperdual, hyperspace_from_vector};
use std::cmp::min;
use std::fmt;
use std::sync::Arc;

use super::{DynamicsAlmanacSnafu, DynamicsAstroSnafu, DynamicsError};

#[derive(Clone)]
pub struct GravityField {
    grav_data: GravityFieldData,
    a_nm: DMatrix<f64>,
    b_nm: DMatrix<f64>,
    c_nm: DMatrix<f64>,
    vr01: DMatrix<f64>,
    vr11: DMatrix<f64>,
    a_nm_h: DMatrix<OHyperdual<f64, U7>>,
    b_nm_h: DMatrix<OHyperdual<f64, U7>>,
    c_nm_h: DMatrix<OHyperdual<f64, U7>>,
    vr01_h: DMatrix<OHyperdual<f64, U7>>,
    vr11_h: DMatrix<OHyperdual<f64, U7>>,
}

impl GravityField {
    /// Create a new Harmonics dynamical model from the provided gravity potential storage instance.
    pub fn new(stor: GravityFieldData) -> Arc<Self> {
        let degree_np2 = stor.max_degree_n() + 2;
        let mut a_nm = DMatrix::from_element(degree_np2 + 1, degree_np2 + 1, 0.0);
        let mut b_nm = DMatrix::from_element(degree_np2, degree_np2, 0.0);
        let mut c_nm = DMatrix::from_element(degree_np2, degree_np2, 0.0);
        let mut vr01 = DMatrix::from_element(degree_np2, degree_np2, 0.0);
        let mut vr11 = DMatrix::from_element(degree_np2, degree_np2, 0.0);

        // Initialize the diagonal elements (not a function of the input)
        a_nm[(0, 0)] = 1.0;
        for n in 1..=degree_np2 {
            let nf64 = n as f64;
            // Diagonal element
            a_nm[(n, n)] = (1.0 + 1.0 / (2.0 * nf64)).sqrt() * a_nm[(n - 1, n - 1)];
        }

        // Pre-compute the B_nm, C_nm, vr01 and vr11 storages
        for n in 0..degree_np2 {
            for m in 0..degree_np2 {
                let nf64 = n as f64;
                let mf64 = m as f64;
                // Compute c_nm, which is B_nm/B_(n-1,m) in Jones' dissertation
                c_nm[(n, m)] = (((2.0 * nf64 + 1.0) * (nf64 + mf64 - 1.0) * (nf64 - mf64 - 1.0))
                    / ((nf64 - mf64) * (nf64 + mf64) * (2.0 * nf64 - 3.0)))
                    .sqrt();

                b_nm[(n, m)] = (((2.0 * nf64 + 1.0) * (2.0 * nf64 - 1.0))
                    / ((nf64 + mf64) * (nf64 - mf64)))
                    .sqrt();

                vr01[(n, m)] = ((nf64 - mf64) * (nf64 + mf64 + 1.0)).sqrt();
                vr11[(n, m)] = (((2.0 * nf64 + 1.0) * (nf64 + mf64 + 2.0) * (nf64 + mf64 + 1.0))
                    / (2.0 * nf64 + 3.0))
                    .sqrt();

                if m == 0 {
                    vr01[(n, m)] /= 2.0_f64.sqrt();
                    vr11[(n, m)] /= 2.0_f64.sqrt();
                }
            }
        }

        // Repeat for the hyperdual part in case we need to super the partials
        let mut a_nm_h =
            DMatrix::from_element(degree_np2 + 1, degree_np2 + 1, OHyperdual::from(0.0));
        let mut b_nm_h = DMatrix::from_element(degree_np2, degree_np2, OHyperdual::from(0.0));
        let mut c_nm_h = DMatrix::from_element(degree_np2, degree_np2, OHyperdual::from(0.0));
        let mut vr01_h = DMatrix::from_element(degree_np2, degree_np2, OHyperdual::from(0.0));
        let mut vr11_h = DMatrix::from_element(degree_np2, degree_np2, OHyperdual::from(0.0));

        // initialize the diagonal elements (not a function of the input)
        a_nm_h[(0, 0)] = OHyperdual::from(1.0);
        for n in 1..=degree_np2 {
            // Diagonal element
            a_nm_h[(n, n)] = OHyperdual::from(a_nm[(n, n)]);
        }

        // Pre-compute the B_nm, C_nm, vr01 and vr11 storages
        for n in 0..degree_np2 {
            for m in 0..degree_np2 {
                vr01_h[(n, m)] = OHyperdual::from(vr01[(n, m)]);
                vr11_h[(n, m)] = OHyperdual::from(vr11[(n, m)]);
                b_nm_h[(n, m)] = OHyperdual::from(b_nm[(n, m)]);
                c_nm_h[(n, m)] = OHyperdual::from(c_nm[(n, m)]);
            }
        }

        Arc::new(Self {
            grav_data: stor,
            a_nm,
            b_nm,
            c_nm,
            vr01,
            vr11,
            a_nm_h,
            b_nm_h,
            c_nm_h,
            vr01_h,
            vr11_h,
        })
    }
}

impl fmt::Display for GravityField {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{} gravity field {}x{} (order x degree)",
            self.grav_data.frame,
            self.grav_data.max_order_m(),
            self.grav_data.max_degree_n(),
        )
    }
}

impl AccelModel for GravityField {
    fn eom(&self, osc: &Orbit, almanac: &Almanac) -> Result<Vector3<f64>, DynamicsError> {
        // Convert the osculating orbit to the correct frame (needed for multiple harmonic fields)
        let state = almanac
            .transform_to(*osc, self.grav_data.frame, None)
            .context(DynamicsAlmanacSnafu {
                action: "transforming into gravity field frame",
            })?;

        // Using the GMAT notation, with extra character for ease of highlight
        let r_ = state.rmag_km();
        let s_ = state.radius_km.x / r_;
        let t_ = state.radius_km.y / r_;
        let u_ = state.radius_km.z / r_;
        let max_degree = self.grav_data.max_degree_n(); // In GMAT, the degree is NN
        let max_order = self.grav_data.max_order_m(); // In GMAT, the order is MM

        // Create the associated Legendre polynomials. Note that we add three items as per GMAT (this may be useful for the STM)
        let mut a_nm = self.a_nm.clone();

        // Initialize the diagonal elements (not a function of the input)
        a_nm[(1, 0)] = u_ * 3.0f64.sqrt();
        for n in 1..=max_degree + 1 {
            let nf64 = n as f64;
            // Off diagonal
            a_nm[(n + 1, n)] = (2.0 * nf64 + 3.0).sqrt() * u_ * a_nm[(n, n)];
        }

        for m in 0..=max_order + 1 {
            for n in (m + 2)..=max_degree + 1 {
                let hm_idx = (n, m);
                a_nm[hm_idx] = u_ * self.b_nm[hm_idx] * a_nm[(n - 1, m)]
                    - self.c_nm[hm_idx] * a_nm[(n - 2, m)];
            }
        }

        // Generate r_m and i_m
        let mut r_m = Vec::with_capacity(min(max_degree, max_order) + 1);
        let mut i_m = Vec::with_capacity(min(max_degree, max_order) + 1);

        r_m.push(1.0);
        i_m.push(0.0);

        for m in 1..=min(max_degree, max_order) {
            r_m.push(s_ * r_m[m - 1] - t_ * i_m[m - 1]);
            i_m.push(s_ * i_m[m - 1] + t_ * r_m[m - 1]);
        }

        let eq_radius_km = self.grav_data.radius_km.unwrap_or(
            self.grav_data
                .frame
                .mean_equatorial_radius_km()
                .context(AstroPhysicsSnafu)
                .context(DynamicsAstroSnafu)?,
        );

        let mu_km3_s2 = self.grav_data.mu_km3_s2.unwrap_or(
            self.grav_data
                .frame
                .mu_km3_s2()
                .context(AstroPhysicsSnafu)
                .context(DynamicsAstroSnafu)?,
        );

        let rho = eq_radius_km / r_;
        let mut rho_np1 = mu_km3_s2 / r_ * rho;
        let mut accel4: Vector4<f64> = Vector4::zeros();

        for n in 1..=max_degree {
            let mut sum: Vector4<f64> = Vector4::zeros();
            rho_np1 *= rho;

            for m in 0..=min(n, max_order) {
                let (c_val, s_val) = self.grav_data.cs_nm(n, m);
                let d_ = unsafe {
                    (c_val * r_m.get_unchecked(m) + s_val * i_m.get_unchecked(m)) * 2.0.sqrt()
                };
                let e_ = if m == 0 {
                    0.0
                } else {
                    unsafe {
                        (c_val * r_m.get_unchecked(m - 1) + s_val * i_m.get_unchecked(m - 1))
                            * 2.0.sqrt()
                    }
                };
                let f_ = if m == 0 {
                    0.0
                } else {
                    unsafe {
                        (s_val * r_m.get_unchecked(m - 1) - c_val * i_m.get_unchecked(m - 1))
                            * 2.0.sqrt()
                    }
                };

                unsafe {
                    sum.x += (m as f64) * a_nm.get_unchecked((n, m)) * e_;
                    sum.y += (m as f64) * a_nm.get_unchecked((n, m)) * f_;
                    sum.z += self.vr01.get_unchecked((n, m)) * a_nm.get_unchecked((n, m + 1)) * d_;
                    sum.w -=
                        self.vr11.get_unchecked((n, m)) * a_nm.get_unchecked((n + 1, m + 1)) * d_;
                }
            }
            let rr = rho_np1 / eq_radius_km;
            accel4 += rr * sum;
        }
        let accel = Vector3::new(
            accel4.x + accel4.w * s_,
            accel4.y + accel4.w * t_,
            accel4.z + accel4.w * u_,
        );
        // Rotate this acceleration vector back into the integration frame (no center change needed, it's just a vector)
        // As discussed with Sai, if the Earth was spinning faster, would the acceleration due to the harmonics be any different?
        // No. Therefore, we do not need to account for the transport theorem here.
        let dcm = almanac
            .rotate(self.grav_data.frame, osc.frame, osc.epoch)
            .context(OrientationSnafu {
                action: "transform state dcm",
            })
            .context(DynamicsAlmanacSnafu {
                action: "transforming into gravity field frame",
            })?;

        Ok(dcm.rot_mat * accel)
    }

    /// The gradient computation of the gravity field skip bound checks on the matrix fetching via the get_unchecked.
    /// This approach allows the gradient to be calculated to full machine precision while, surprisingly, being slightly
    /// faster than a single forward different gradient approach.
    fn gradient(
        &self,
        osc: &Orbit,
        almanac: &Almanac,
    ) -> Result<(Vector3<f64>, Matrix3<f64>), DynamicsError> {
        // Convert the osculating orbit to the correct frame (needed for multiple harmonic fields)
        let state = almanac
            .transform_to(*osc, self.grav_data.frame, None)
            .context(DynamicsAlmanacSnafu {
                action: "transforming into gravity field frame",
            })?;

        let radius: Vector3<OHyperdual<f64, U7>> = hyperspace_from_vector(&state.radius_km);

        // Using the GMAT notation, with extra character for ease of highlight
        let r_ = norm(&radius);
        let s_ = radius[0] / r_;
        let t_ = radius[1] / r_;
        let u_ = radius[2] / r_;
        let max_degree = self.grav_data.max_degree_n(); // In GMAT, the order is NN
        let max_order = self.grav_data.max_order_m(); // In GMAT, the order is MM

        // Create the associated Legendre polynomials. Note that we add three items as per GMAT (this may be useful for the STM)
        let mut a_nm = DMatrix::from_element(max_degree + 3, max_degree + 3, OHyperdual::from(0.0));
        // Copy only the pre-computed diagonals from self.a_nm_h manually
        for i in 0..=max_degree + 1 {
            a_nm[(i, i)] = self.a_nm_h[(i, i)];
        }

        // Initialize the diagonal elements (not a function of the input)
        a_nm[(1, 0)] = u_ * 3.0f64.sqrt();
        for n in 1..=max_degree + 1 {
            let nf64 = n as f64;
            // Off diagonal
            a_nm[(n + 1, n)] = OHyperdual::from((2.0 * nf64 + 3.0).sqrt()) * u_ * a_nm[(n, n)];
        }

        for m in 0..=max_order + 1 {
            for n in (m + 2)..=max_degree + 1 {
                let hm_idx = (n, m);
                a_nm[hm_idx] = u_ * self.b_nm_h[hm_idx] * a_nm[(n - 1, m)]
                    - self.c_nm_h[hm_idx] * a_nm[(n - 2, m)];
            }
        }

        // Generate r_m and i_m
        let mut r_m = Vec::with_capacity(min(max_degree, max_order) + 1);
        let mut i_m = Vec::with_capacity(min(max_degree, max_order) + 1);

        r_m.push(OHyperdual::<f64, U7>::from(1.0));
        i_m.push(OHyperdual::<f64, U7>::from(0.0));

        for m in 1..=min(max_degree, max_order) {
            r_m.push(s_ * r_m[m - 1] - t_ * i_m[m - 1]);
            i_m.push(s_ * i_m[m - 1] + t_ * r_m[m - 1]);
        }

        let real_eq_radius_km = self.grav_data.radius_km.unwrap_or(
            self.grav_data
                .frame
                .mean_equatorial_radius_km()
                .context(AstroPhysicsSnafu)
                .context(DynamicsAstroSnafu)?,
        );

        let real_mu_km3_s2 = self.grav_data.mu_km3_s2.unwrap_or(
            self.grav_data
                .frame
                .mu_km3_s2()
                .context(AstroPhysicsSnafu)
                .context(DynamicsAstroSnafu)?,
        );

        let eq_radius = OHyperdual::<f64, U7>::from(real_eq_radius_km);
        let rho = eq_radius / r_;
        let mut rho_np1 = OHyperdual::<f64, U7>::from(real_mu_km3_s2) / r_ * rho;

        let mut a0 = OHyperdual::<f64, U7>::from(0.0);
        let mut a1 = OHyperdual::<f64, U7>::from(0.0);
        let mut a2 = OHyperdual::<f64, U7>::from(0.0);
        let mut a3 = OHyperdual::<f64, U7>::from(0.0);
        let sqrt2 = OHyperdual::<f64, U7>::from(2.0.sqrt());

        for n in 1..=max_degree {
            let mut sum0 = OHyperdual::from(0.0);
            let mut sum1 = OHyperdual::from(0.0);
            let mut sum2 = OHyperdual::from(0.0);
            let mut sum3 = OHyperdual::from(0.0);
            rho_np1 *= rho;

            for m in 0..=min(n, max_order) {
                let (c_valf64, s_valf64) = self.grav_data.cs_nm(n, m);
                let c_val = OHyperdual::<f64, U7>::from(c_valf64);
                let s_val = OHyperdual::<f64, U7>::from(s_valf64);

                let d_ = unsafe {
                    (c_val * r_m.get_unchecked(m) + s_val * i_m.get_unchecked(m)) * sqrt2
                };
                let e_ = if m == 0 {
                    OHyperdual::from(0.0)
                } else {
                    unsafe {
                        (c_val * r_m.get_unchecked(m - 1) + s_val * i_m.get_unchecked(m - 1))
                            * sqrt2
                    }
                };
                let f_ = if m == 0 {
                    OHyperdual::from(0.0)
                } else {
                    unsafe {
                        (s_val * r_m.get_unchecked(m - 1) - c_val * i_m.get_unchecked(m - 1))
                            * sqrt2
                    }
                };

                unsafe {
                    sum0 += OHyperdual::from(m as f64) * a_nm.get_unchecked((n, m)) * e_;
                    sum1 += OHyperdual::from(m as f64) * a_nm.get_unchecked((n, m)) * f_;
                    sum2 +=
                        *self.vr01_h.get_unchecked((n, m)) * *a_nm.get_unchecked((n, m + 1)) * d_;
                    sum3 += *self.vr11_h.get_unchecked((n, m))
                        * *a_nm.get_unchecked((n + 1, m + 1))
                        * d_;
                }
            }
            let rr = rho_np1 / eq_radius;
            a0 += rr * sum0;
            a1 += rr * sum1;
            a2 += rr * sum2;
            a3 -= rr * sum3;
        }

        let dcm = almanac
            .rotate(self.grav_data.frame, osc.frame, osc.epoch)
            .context(OrientationSnafu {
                action: "transform state dcm",
            })
            .context(DynamicsAlmanacSnafu {
                action: "transforming into gravity field frame",
            })?
            .rot_mat;

        let accel_local = Vector3::new(a0 + a3 * s_, a1 + a3 * t_, a2 + a3 * u_);
        let dx = dcm
            * Vector3::new(
                accel_local[0].real(),
                accel_local[1].real(),
                accel_local[2].real(),
            );

        let mut grad_local = Matrix3::zeros();
        for i in 0..3 {
            // For each acceleration component
            for j in 1..4 {
                // For each derivative wrt to the position
                grad_local[(i, j - 1)] += accel_local[i][j];
            }
        }
        let grad = dcm * grad_local * dcm.transpose();
        Ok((dx, grad))
    }
}