KSP 2D Ascent Trajectory Optimization with PSOPT

gistfile1.cpp

// KSP 2D Ascent Trajectory Optimization Problem
#include "psopt.h"

//////////////////////////////////////////////////////////////////////////
///////////////////  Define the end point (Mayer) cost function //////////
//////////////////////////////////////////////////////////////////////////

adouble endpoint_cost(adouble* initial_states, adouble* final_states, 
                      adouble* parameters,adouble& t0, adouble& tf, 
                      adouble* xad, int iphase)
{
    adouble mf = final_states[3];
    return -mf;
} 

//////////////////////////////////////////////////////////////////////////
///////////////////  Define the integrand (Lagrange) cost function  //////
//////////////////////////////////////////////////////////////////////////

adouble integrand_cost(adouble* states, adouble* controls, 
                       adouble* parameters, adouble& time, adouble* xad, 
                       int iphase)
{
    return  0.0;
} 

//////////////////////////////////////////////////////////////////////////
///////////////////  Define the DAE's ////////////////////////////////////
//////////////////////////////////////////////////////////////////////////

void dae(adouble* derivatives, adouble* path, adouble* states, 
         adouble* controls, adouble* parameters, adouble& time, 
         adouble* xad, int iphase)
{
   adouble Dr, Drdot, Dthetadot, Dm;

   adouble r = states[0];
   adouble rdot = states[1];
   adouble thetadot = states[2];
   adouble m = states[3];

   adouble Tr = controls[0];
   adouble Ttheta = controls[1];

   double Cd = 0.2;
   double Kp = 1.2230948554874*0.008;
   double P0 = 1.0;
   double rmin = 600000.0;
   double g0 = 9.8;
   double mu = 3.5316e12;
   double H0 = 5000.0;
   double Isp0 = 320.0;
   double Ispvac = 370.0;
   double omegarot = 2.0*pi/21600.0;

   adouble P = P0*exp(-(r-rmin)/H0);
   adouble Isp_lerp = P*Isp0+(1.0-P)*Ispvac;
   
   adouble drag_multiplier = -0.5*Kp*P*Cd;
   adouble a_grav = -mu/(r*r);
   adouble vrelsq = pow(rdot, 2) + pow(r*(thetadot-omegarot),2);
   adouble Fdtangent = drag_multiplier*sqrt(vrelsq)*r*(thetadot-omegarot);
   adouble Fdradial = drag_multiplier*sqrt(vrelsq)*rdot;
   
   adouble Tmag = sqrt(Tr*Tr+Ttheta*Ttheta);
   
   Dr = rdot;
   Drdot = Tr/m+a_grav+Fdradial;
   Dthetadot = (1.0/r)*(Ttheta/m+Fdtangent-2.0*rdot*thetadot);
   Dm = -Tmag/(g0*Isp_lerp);

   derivatives[0] = Dr;
   derivatives[1] = Drdot;
   derivatives[2] = Dthetadot;
   derivatives[3] = Dm;
   path[0] = Tmag;
}

////////////////////////////////////////////////////////////////////////////
///////////////////  Define the events function ////////////////////////////
////////////////////////////////////////////////////////////////////////////

void events(adouble* e, adouble* initial_states, adouble* final_states, 
            adouble* parameters,adouble& t0, adouble& tf, adouble* xad, 
            int iphase) 
{
   adouble ri = initial_states[0];
   adouble rdoti = initial_states[1];
   adouble thetadoti = initial_states[2];
   adouble mi = initial_states[3];
   adouble rf = final_states[0];
   adouble rdotf = final_states[1];
   adouble thetadotf = final_states[2];
  

   e[0] = ri;
   e[1] = rdoti;
   e[2] = thetadoti;
   e[3] = mi;
   e[4] = rf;
   e[5] = rdotf;
   e[6] = thetadotf;

}


///////////////////////////////////////////////////////////////////////////
///////////////////  Define the phase linkages function ///////////////////
///////////////////////////////////////////////////////////////////////////

void linkages( adouble* linkages, adouble* xad)
{
  // No linkages as this is a single phase problem
}


////////////////////////////////////////////////////////////////////////////
///////////////////  Define the main routine ///////////////////////////////
////////////////////////////////////////////////////////////////////////////

int main(void)
{
////////////////////////////////////////////////////////////////////////////
///////////////////  Declare key structures ////////////////////////////////
////////////////////////////////////////////////////////////////////////////

    Alg  algorithm;
    Sol  solution;
    Prob problem;

////////////////////////////////////////////////////////////////////////////
///////////////////  Register problem name  ////////////////////////////////
////////////////////////////////////////////////////////////////////////////

    problem.name          	= "KSP Maximum Profile";
    problem.outfilename                 = "ksp.txt";
    int nodecount = 60;

////////////////////////////////////////////////////////////////////////////
////////////  Declare problem level constants & setup phases  //////////////
////////////////////////////////////////////////////////////////////////////

    problem.nphases   			= 1;
    problem.nlinkages                   = 0;

    psopt_level1_setup(problem);


/////////////////////////////////////////////////////////////////////////////
/////////   Define phase related information &  setup PSOPT  /////////////////
/////////////////////////////////////////////////////////////////////////////


    problem.phases(1).nstates   		= 4;
    problem.phases(1).ncontrols 		= 2;
    problem.phases(1).nevents   		= 7;
    problem.phases(1).npath     		= 1;
    problem.phases(1).nodes                     = nodecount; 

    psopt_level2_setup(problem, algorithm);

////////////////////////////////////////////////////////////////////////////
///////////////////  Declare DMatrix objects to store results //////////////
////////////////////////////////////////////////////////////////////////////

    DMatrix x, u, t;

////////////////////////////////////////////////////////////////////////////
///////////////////  Enter problem bounds information //////////////////////
////////////////////////////////////////////////////////////////////////////


    double r_L = 600000.0;
    double rdot_L = -INF;
    double thetadot_L = 0.0;
    double m_L = 0.0;
    
    double r_U = INF;
    double rdot_U = INF;
    double thetadot_U = INF;
    double m_U = 11.15;

    double Tr_L = 0.0;
    double Tr_U = 215.0;
    
    double Ttheta_L = 0.0;
    double Ttheta_U = 215.0;
    
    double r_i = 600000.0+70.0;
    double rdot_i = 0.0;
    double thetadot_i = 2.9089e-4;
    double m_i = 11.15;
    double r_f = 674000.0;
    double rdot_f = 0.0;
    double thetadot_f = 2289.053/674000.0;

    problem.phases(1).bounds.lower.states(1) = r_L;
    problem.phases(1).bounds.lower.states(2) = rdot_L;
    problem.phases(1).bounds.lower.states(3) = thetadot_L;
    problem.phases(1).bounds.lower.states(4) = m_L;

    problem.phases(1).bounds.upper.states(1) = r_U;
    problem.phases(1).bounds.upper.states(2) = rdot_U;
    problem.phases(1).bounds.upper.states(3) = thetadot_U;
    problem.phases(1).bounds.upper.states(4) = m_U;

    problem.phases(1).bounds.lower.controls(1) = Tr_L;
    problem.phases(1).bounds.upper.controls(1) = Tr_U;
    problem.phases(1).bounds.lower.controls(2) = Ttheta_L;
    problem.phases(1).bounds.upper.controls(2) = Ttheta_U;

    problem.phases(1).bounds.lower.events(1) = r_i;
    problem.phases(1).bounds.lower.events(2) = rdot_i;
    problem.phases(1).bounds.lower.events(3) = thetadot_i;
    problem.phases(1).bounds.lower.events(4) = m_i;
    problem.phases(1).bounds.lower.events(5) = r_f;
    problem.phases(1).bounds.lower.events(6) = rdot_f;
    problem.phases(1).bounds.lower.events(7) = thetadot_f;
   

    problem.phases(1).bounds.upper.events(1) = r_i;
    problem.phases(1).bounds.upper.events(2) = rdot_i;
    problem.phases(1).bounds.upper.events(3) = thetadot_i;
    problem.phases(1).bounds.upper.events(4) = m_i;
    problem.phases(1).bounds.upper.events(5) = r_f;
    problem.phases(1).bounds.upper.events(6) = rdot_f;
    problem.phases(1).bounds.upper.events(7) = thetadot_f;
    
    problem.phases(1).bounds.lower.path(1) = 0.0;
    problem.phases(1).bounds.upper.path(1) = 215.0;


    problem.phases(1).bounds.lower.StartTime    = 0.0;
    problem.phases(1).bounds.upper.StartTime    = 0.0;

    problem.phases(1).bounds.lower.EndTime      = 10.0;
    problem.phases(1).bounds.upper.EndTime      = 1000.0;


////////////////////////////////////////////////////////////////////////////
///////////////////  Register problem functions  ///////////////////////////
////////////////////////////////////////////////////////////////////////////


    problem.integrand_cost 	= &integrand_cost;
    problem.endpoint_cost 	= &endpoint_cost;
    problem.dae 		= &dae;
    problem.events 		= &events;
    problem.linkages		= &linkages;

////////////////////////////////////////////////////////////////////////////
///////////////////  Define & register initial guess ///////////////////////
////////////////////////////////////////////////////////////////////////////


    DMatrix x0(4,nodecount);

    x0(1,colon()) = linspace(600000.0,670000.0, nodecount);
    x0(2,colon()) = linspace(100.0,0.0, nodecount);
    x0(3,colon()) = linspace(1e-4,1e-3, nodecount);
    x0(4,colon()) = linspace(m_U,m_L, nodecount);
    
    DMatrix u0(2,nodecount);
    u0(1,colon()) = linspace(Tr_U, Tr_L, nodecount);
    u0(2,colon()) = linspace(Ttheta_L,Ttheta_U, nodecount);
    
    DMatrix tt(1,nodecount);
    tt(1,colon()) = linspace(0.0,300.0,nodecount);
    
    DMatrix initcond(4,1);
    initcond(1) = r_i;
    initcond(2) = rdot_i;
    initcond(3) = thetadot_i;
    initcond(4) = m_i;
    tra(initcond).Print("initcond:");
    DMatrix params(0);
    
    //rk4_propagate(dae, u0, tt, initcond, params, problem, 1, x0);
    
    tra(x0).Print("x0:");
    //return 0;
    
    problem.phases(1).guess.controls       = u0;
    problem.phases(1).guess.states         = x0;
    problem.phases(1).guess.time           = tt;


////////////////////////////////////////////////////////////////////////////
///////////////////  Enter algorithm options  //////////////////////////////
////////////////////////////////////////////////////////////////////////////
    

    algorithm.nlp_method                  = "IPOPT";
    algorithm.scaling                     = "automatic";
    algorithm.derivatives                 = "automatic";
    //algorithm.mesh_refinement             = "automatic";
    //algorithm.mr_max_iterations = 3;
    algorithm.nlp_iter_max                = 50000;
    algorithm.nlp_tolerance               = 1.e-6;
    //algorithm.collocation_method = "Hermite-Simpson";


////////////////////////////////////////////////////////////////////////////
///////////////////  Now call PSOPT to solve the problem   //////////////////
////////////////////////////////////////////////////////////////////////////

    psopt(solution, problem, algorithm);

////////////////////////////////////////////////////////////////////////////
///////////  Extract relevant variables from solution structure   //////////
////////////////////////////////////////////////////////////////////////////

    x = solution.get_states_in_phase(1);
    u = solution.get_controls_in_phase(1);
    t = solution.get_time_in_phase(1);

////////////////////////////////////////////////////////////////////////////
///////////  Save solution data to files if desired ////////////////////////
////////////////////////////////////////////////////////////////////////////

    x.Save("x.dat");
    u.Save("u.dat");
    t.Save("t.dat");

////////////////////////////////////////////////////////////////////////////
///////////  Plot some results if desired (requires gnuplot) ///////////////
////////////////////////////////////////////////////////////////////////////

    plot(t,x.Row(1),problem.name, "time (s)", "states", "r");
    plot(t,x.Row(2),problem.name, "time (s)", "states", "rdot");
    plot(t,x.Row(3),problem.name, "time (s)", "states", "thetadot");
    plot(t,x.Row(4),problem.name, "time (s)", "states", "m");

    plot(t,u.Row(1),problem.name, "time (s)", "control", "T");
    plot(t,u.Row(2),problem.name, "time (s)", "control", "Ttheta");

    //plot(t,x,problem.name, "time (s)", "states", "v h m",
                           //"pdf", "ksp_vert_states.pdf");

    //plot(t,u,problem.name, "time (s)", "control", "T",
                           //"pdf", "ksp_vert_control.pdf");

}