Automatically differentiate state transition matrix via the use of dual numbers

[ChristopherRabotin]

(And probably write a paper on this method)

Summary

I determined how to compute the state transition matrix of dynamics using dual numbers. By semi-clever organization of the vector returned by the equation of motion, it is possible to generate the gradient of the EOM for the provided state.

The following proof of concept uses the dual space, and as such requires N calls to the EOMs, where N corresponds to the size of the state. However, if using a hyperdual space of size N, the EOMs only need to be computed once.

Proof of concept

Output

Note that the expected STM is computed through the analytical version, as done here

     Running `target/debug/dualgradtest`

  ┌                                                                                                                                                                         ┐
  │                           0                           0                           0                           1                           0                           0 │
  │                           0                           0                           0                           0                           1                           0 │
  │                           0                           0                           0                           0                           0                           1 │
  │ -0.000000018628398676538285   -0.0000000408977477510809  -0.00000001544439654966729                           0                           0                           0 │
  │ -0.000000040897747751080905    0.0000000452532717518738  0.000000031658392121967554                           0                           0                           0 │
  │  -0.00000001544439654966729   0.00000003165839212196755 -0.000000026624873075335535                           0                           0                           0 │
  └                                                                                                                                                                         ┘


|delta| = 5.909567257049796e-23

main.rs

extern crate dual_num;
extern crate nalgebra as na;

use dual_num::{differentiate, DualNumber, Float, FloatConst};
use na::{DimName, Matrix6, U3, U6, Vector3, Vector6};

fn diff<F>(x: Vector6<f64>, func: F) -> Matrix6<f64>
where
    F: Fn(&Vector6<DualNumber<f64>>) -> Vector6<DualNumber<f64>>,
{
    let mut grad_as_slice = Vec::with_capacity(U6::dim() * U6::dim());
    for coord in 0..U6::dim() {
        let x11 = Vector6::<DualNumber<f64>>::new(
            DualNumber::new(x[(0, 0)], if coord == 0 { 1.0 } else { 0.0 }),
            DualNumber::new(x[(1, 0)], if coord == 1 { 1.0 } else { 0.0 }),
            DualNumber::new(x[(2, 0)], if coord == 2 { 1.0 } else { 0.0 }),
            DualNumber::new(x[(3, 0)], if coord == 3 { 1.0 } else { 0.0 }),
            DualNumber::new(x[(4, 0)], if coord == 4 { 1.0 } else { 0.0 }),
            DualNumber::new(x[(5, 0)], if coord == 5 { 1.0 } else { 0.0 }),
        );
        let df_di = func(&x11);
        for i in 0..6 {
            grad_as_slice.push(df_di[(i, 0)].dual());
        }
    }

    let grad = Matrix6::from_column_slice(&grad_as_slice);
    grad
}

fn eom(state: &Vector6<DualNumber<f64>>) -> Vector6<DualNumber<f64>> {
    let radius = state.fixed_rows::<U3>(0).into_owned();
    let velocity = state.fixed_rows::<U3>(3).into_owned();
    let norm = (radius[(0, 0)].powi(2) + radius[(1, 0)].powi(2) + radius[(2, 0)].powi(2)).sqrt();
    let body_acceleration_f = DualNumber::from_real(-398_600.4415) / norm.powi(3);
    let body_acceleration = Vector3::new(
        radius[(0, 0)] * body_acceleration_f,
        radius[(1, 0)] * body_acceleration_f,
        radius[(2, 0)] * body_acceleration_f,
    );
    Vector6::from_iterator(velocity.iter().chain(body_acceleration.iter()).cloned())
}

fn main() {
    let state = Vector6::new(
        -9042.862233600335,
        18536.333069123244,
        6999.9570694864115,
        -3.28878900377057,
        -2.226285193102822,
        1.6467383807226765,
    );
    let grad = diff(state, eom);
    println!("{}", grad);

    let mut expected = Matrix6::zeros();
    expected[(0, 3)] = 1.0;
    expected[(1, 4)] = 1.0;
    expected[(2, 5)] = 1.0;
    expected[(3, 0)] = -0.000000018628398676538285;
    expected[(4, 0)] = -0.00000004089774775108092;
    expected[(5, 0)] = -0.0000000154443965496673;
    expected[(3, 1)] = -0.00000004089774775108092;
    expected[(4, 1)] = 0.000000045253271751873843;
    expected[(5, 1)] = 0.00000003165839212196757;
    expected[(3, 2)] = -0.0000000154443965496673;
    expected[(4, 2)] = 0.00000003165839212196757;
    expected[(5, 2)] = -0.000000026624873075335538;

    println!("|delta| = {:e}", (grad - expected).norm());
}

Cargo.toml

[package]
name = "dualgradtest"
version = "0.1.0"
authors = ["Christopher Rabotin <christopher.rabotin@gmail.com>"]

[dependencies]
dual_num = "0.1"
nalgebra = "0.15"

[ChristopherRabotin]

Despite novacrazy/dual_num#2 being merged, I still need to make Dual a na::Scalar in dual_num in order to support all the na::MatrixMN operations. This is important so that the definition of all the eoms are (almost?) identical whether autodifferentiation is used or not. The almost might be required for slicing the state vector.

Moreover, I would want the nabla function of dual_num to be using hyperdual space instead of "simulating" it so as to limit the number of calls to the eoms. Although I do not think this will reduce the mathematical complexity of the problem at hand (since the same number of operations is needed in both cases), it'll make the call stack smaller, and therefore probably provide some speed up of the gradient computation.