Efficient implementation of nonlinear dynamics and constraints

[afossa]

Hello,

I am using InfiniteOpt to solve optimal control problems in trajectory optimization. These type of problems frequently have complicated nonlinear expressions in both the dynamics and the constraints, and they rely a lot on operations like matrix-vector products, cross products, norms, etc.

I was able to successfully solve an optimal transfer between the Earth and Mars, but I am looking for suggestions on how I can improve my implementation, specifically how to avoid recomputing quantities that appear in more than one constraint, such as the distance from the Sun. Should I add a variable for each of them, or is there a better way to achieve this?

I am also experimenting with different problem formulations, in particular I want to express the control variables for the three components of the thrust vector in either the inertial frame or in a rotating frame defined by the spacecraft position and velocity vectors. In the second case, I need to compute the rotation matrix from the rotating to the inertial frame, but I am getting an error that seems to be related to the normalization of the angular momentum vector (computed as the cross product between position and velocity):

Unable to create a nonlinear expression because it did not contain any JuMP scalars. head = `:+`, args = `Any[0.0]`.

So I modified the problem formulation to avoid it, however, the accuracy of the solution with the thrust expressed in rotating frame is much lower than that of the solution with the thrust expressed directly in inertial frame. My guess is that the rotation matrix, as I am computing it now, is not perfectly orthogonal. This is the working code I came up with:

@variables(mdl, begin
    r[1:3], Infinite(t) # position vector
    v[1:3], Infinite(t) # velocity vector
    m_dry <= m <= m0, Infinite(t) # mass
    0.0 <= Γ, Infinite(t), (start=0.0) # energy integral
    0.0 <= T <= thr_max, Infinite(t), (start=0.0) # thrust magnitude
    u[1:3], Infinite(t), (start=0.0) # thrust components
end);

# zeroth-order hold control
constant_over_collocation(T, t);
[constant_over_collocation(u[i], t) for i in 1:3];

# initial guess
[set_start_value_function(r[i], t -> linear_interpolation(tv_l, rvv_l[i,:], extrapolation_bc=Line())(t)) for i in 1:3];
[set_start_value_function(v[i], t -> linear_interpolation(tv_l, rvv_l[i+3,:], extrapolation_bc=Line())(t)) for i in 1:3];
set_start_value_function(m, t -> linear_interpolation([t0, tf], [m0, m_dry], extrapolation_bc=Line())(t));

if frm_id != 0 # thrust components in local orbital frame

    # unit vectors along position, velocity, and angular momentum vectors
    @expression(mdl, ur, r ./ dot(r,r)^0.5);
    @expression(mdl, uv, v ./ dot(v,v)^0.5);
    @expression(mdl, uh, [ur[2]*uv[3] - ur[3]*uv[2]; ur[3]*uv[1] - ur[1]*uv[3]; ur[1]*uv[2] - ur[2]*uv[1]]);

    # unit vectors defining the local orbital frame
    if frm_id == 1 # qsw frame
        @expression(mdl, ux, ur);
    else # tnw frame
        @expression(mdl, ux, uv);
    end;
    @expression(mdl, uz, uh);
    @expression(mdl, uy, [uz[2]*ux[3] - uz[3]*ux[2]; uz[3]*ux[1] - uz[1]*ux[3]; uz[1]*ux[2] - uz[2]*ux[1]]);

    # thrust components in inertial frame
    @expression(mdl, ui, [ux[1]*u[1] + uy[1]*u[2] + uz[1]*u[3]; ux[2]*u[1] + uy[2]*u[2] + uz[2]*u[3]; ux[3]*u[1] + uy[3]*u[2] + uz[3]*u[3]]);

    # thrust acceleration in inertial frame
    @expression(mdl, at, ui ./ m);
else # thrust components in inertial frame
    @expression(mdl, at, u ./ m); # thrust acceleration in inertial frame
end;

# gravitational acceleration in inertial frame
@expression(mdl, ag, -gm / dot(r,r)^1.5 * r);

# dynamics
@constraint(mdl, [i=1:3], ∂(r[i], t) == v[i]);
@constraint(mdl, [i=1:3], ∂(v[i], t) == ag[i] + at[i]);
@constraint(mdl, ∂(m, t) == -T / vex);
@constraint(mdl, ∂(Γ, t) == 0.5 * T^2);

# control constraints
@constraint(mdl, dot(u,u) <= T^2); # thrust magnitude

Do you have any suggestion on how I can improve it, both in terms of accuracy and efficiency as discussed above?

The results I obtain in inertial frame are:

Problem type:   energy optimal
Thrust vector:  inertial frame
Remaining fuel: 56.51349789895097 kg
Errors between final BCs and explicit simulation:
Position: 1.667613622724852 km
Velocity: 6.277496265271775e-8 km/s

and in QSW frame:

Problem type:   energy optimal
Thrust vector:  qsw frame
Remaining fuel: 55.46058819518962 kg
Errors between final BCs and explicit simulation:
Position: 903149.9219073891 km
Velocity: 0.12782349401040421 km/s

finally, in TNW frame:

Problem type:   energy optimal
Thrust vector:  tnw frame
Remaining fuel: 54.81937906234515 kg
Errors between final BCs and explicit simulation:
Position: 492797.71709358005 km
Velocity: 0.09807832598412268 km/s

Attached is the full script and input data to reproduce these results. I am using the version of InfiniteOpt checked out from the master branch to use the constant_over_collocation method that enforces a piecewise-constant control.

earth_mars_input.json
earth_mars_script.txt