Trajectory optimisation example

This example shows how to use pyoptgra to optimise a trajectory using the GODOT cosmos module.

It was created by Gabor Varga as part of the Godot tutorials, he generously gave permission to reproduce it here.

First we load some python packages for plotting and linear algebra.

In [1]:
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

Then we load godot specific libraries and the pygmo optimisation package

In [2]:
from godot.core import tempo, util, constants
from godot.core import autodif as ad
from godot.core.autodif import bridge as br
from godot.model import interface
from godot import cosmos
import pygmo
import pygmo_plugins_nonfree
In [3]:
import pyoptgra

Load the universe configuration and create the universe object

In [4]:
util.suppressLogger()
uni_config = cosmos.util.load_yaml('universe.yml')
uni = cosmos.Universe(uni_config)

Load the trajectory configuration and create the trajectory object using the universe object

In [5]:
tra_config = cosmos.util.load_yaml('trajectory.yml')
tra = cosmos.Trajectory(uni, tra_config)

We generate a plotting function to create visualization for our trajectory in the solar system.

In [6]:
def man_data(name):
    start = tra.point(f"{name}_start")
    end = tra.point(f"{name}_end")
    ran = tempo.EpochRange(start, end)
    grid = ran.createGrid(1.0 * tempo.SecondsInDay)
    t = [e.mjd() for e in grid]
    r = np.asarray([uni.frames.vector3('Sun', 'SC_center', 'EMC', e) / constants.AU for e in grid])
    thr = np.asarray([uni.evaluables["SC_sep_thrust"].eval(e) for e in grid])
    return t, r, thr

def plot():
    ran = tempo.EpochRange(tempo.Epoch("2020-01-01 TDB"), tempo.Epoch("2022-01-01 TDB"))
    earth = np.asarray([uni.frames.vector3('Sun', 'Earth', 'EMC', e) / constants.AU for e in ran.createGrid(1.0 * tempo.SecondsInDay)])
    mars = np.asarray([uni.frames.vector3('Sun', 'Mars', 'EMC', e) / constants.AU for e in ran.createGrid(1.0 * tempo.SecondsInDay)])

    def pos(e):
        try:
            return uni.frames.vector3('Sun', 'SC_center', 'EMC', e) / constants.AU
        except:
            return [np.nan] * 3

    grid = tra.range().createGrid(1.0 * tempo.SecondsInDay)
    t = [e.mjd() for e in grid]
    x = np.asarray([pos(e) for e in grid])

    man1_t, man1_r, man1_thr = man_data("man1")
    man2_t, man2_r, man2_thr = man_data("man2")

    plt.figure(figsize=(8, 8))
    
    plt.xlabel('EMC X (AU)')
    plt.ylabel('EMC Y (AU)')
    plt.plot(earth[:, 0], earth[:, 1], '--', label="Earth")
    plt.plot(mars[:, 0], mars[:, 1], '--', label="Mars")
    plt.plot(x[:, 0], x[:, 1], '-k', linewidth=2, label="Transfer")
    plt.plot(man1_r[:, 0], man1_r[:, 1], '-r', linewidth=4, label="SEP 1")
    plt.plot(man2_r[:, 0], man2_r[:, 1], '-g', linewidth=4, label="SEP 2")
    plt.plot(x[0, 0], x[0, 1], 'ok')
    plt.plot(x[-1, 0], x[-1, 1], 'ok')
    plt.xlim([-2, 2])
    plt.ylim([-2, 2])
    plt.legend()
    plt.grid()

    plt.figure(figsize=(8, 4))
    plt.xlabel('Time (MJD)')
    plt.ylabel('SEP thrust (mN)')
    plt.plot(man1_t, 1e6 * man1_thr, '-r', linewidth=3, label="SEP 1")
    plt.plot(man2_t, 1e6 * man2_thr, '-g', linewidth=3, label="SEP 2")
    plt.xlim([9750, 10150])
    plt.ylim([0, 1000])
    plt.legend()
    plt.grid()

We compute the initial guess of the transfer trajectory and visualize it.

Immediately we can see that there is a large discontinuity at the matching point and this optimisation problem needs to be solved.

In [7]:
tra.compute(False)
plot()

Load the problem configuration and create the problem object using the universe, a list of considered trajectories

In [8]:
pro_config = cosmos.util.load_yaml('problem.yml')
pro = cosmos.Problem(uni, [tra], pro_config)

Now it is time to create the pygmo problem class that simply takes ours, and define some constraint tolerances

In [9]:
problem = pygmo.problem(pro)
conTol = 1e-6
problem.c_tol = [conTol] * problem.get_nc()

We create the pygmo population to be optimised by grabbing the initial value of the free parameters

In [10]:
x0 = pro.get_x()
pop = pygmo.population(problem, 0)
pop.push_back(x0)

Finally we use pyoptgra to optimise the problem

In [11]:
algo = pygmo.algorithm(pyoptgra.optgra())
algo.set_verbosity(1)
pop = algo.evolve(pop)

We can visualize the final converged trajectory.

In [12]:
plot()

Finally, we can save the optimised and updated trajectory configuration

In [13]:
up = tra.applyParameterChanges()
tra_config = cosmos.util.deep_update(tra_config, up)
cosmos.util.save_yaml(tra_config, 'trajectory_out.yml')