I have two phases, one much shorter than the other, but it looks like I have to make the transcription the same for both- same number of nodes and polynomial order for each one, other wise I get a numpy mismatched array size / broadcasting error. Is there a way around this?
You do not need to have the same transcription for both.
Here is an example of setting up a very simple cannonball problem with two phases. The first phase is Radau, the second phase is GaussLabotto. Both use the same ODE, but different orders, numbers of segments, and compression setting.
import numpy as np
from scipy.interpolate import interp1d
import openmdao.api as om
from openmdao.components.interp_util.interp import InterpND
import dymos as dm
from dymos.models.atmosphere.atmos_1976 import USatm1976Data
# CREATE an atmosphere interpolant
english_to_metric_rho = om.unit_conversion('slug/ft**3', 'kg/m**3')[0]
english_to_metric_alt = om.unit_conversion('ft', 'm')[0]
rho_interp = interp1d(np.array(USatm1976Data.alt*english_to_metric_alt, dtype=complex),
np.array(USatm1976Data.rho*english_to_metric_rho, dtype=complex), kind='linear')
GRAVITY = 9.80665
class CannonballODE(om.ExplicitComponent):
def initialize(self):
self.options.declare('num_nodes', types=int)
def setup(self):
nn = self.options['num_nodes']
# static parameters
self.add_input('radius', units='m')
self.add_input('density', units='kg/m**3')
self.add_input('CD', units=None)
self.add_input('alt', units='m', shape=nn)
self.add_input('v', units='m/s', shape=nn)
self.add_input('gam', units='rad', shape=nn)
self.add_output('v_dot', shape=(nn,), units='m/s**2')
self.add_output('gam_dot', shape=(nn,), units='rad/s')
self.add_output('h_dot', shape=(nn,), units='m/s')
self.add_output('r_dot', shape=(nn,), units='m/s')
self.add_output('ke', shape=(nn,), units='J')
self.add_output('mass', shape=1, units='kg')
self.declare_partials('*', '*', method='cs')
def compute(self, inputs, outputs):
CD = inputs['CD']
gam = inputs['gam']
v = inputs['v']
alt = inputs['alt']
radius = inputs['radius']
dens = inputs['density']
m = (4/3.)*dens*np.pi*radius**3
S = np.pi*radius**2
if np.iscomplexobj(alt):
rho = rho_interp(inputs['alt'])
else:
rho = rho_interp(inputs['alt']).real
q = 0.5*rho*inputs['v']**2
qS = q * S
D = qS * CD
cgam = np.cos(gam)
sgam = np.sin(gam)
mv = m*v
outputs['v_dot'] = - D/m-GRAVITY*sgam
outputs['gam_dot'] = -(GRAVITY/v)*cgam
outputs['h_dot'] = v*sgam
outputs['r_dot'] = v*cgam
outputs['ke'] = 0.5*m*v**2
if __name__ == "__main__":
p = om.Problem()
p.driver = om.pyOptSparseDriver()
p.driver.options['optimizer'] = 'SLSQP'
p.driver.declare_coloring()
traj = p.model.add_subsystem('traj', dm.Trajectory())
# constants across the whole trajectory
traj.add_parameter('radius', units='m', val=0.01, dynamic=False)
traj.add_parameter('density', units='kg/m**3', dynamic=False)
p.model.add_design_var('traj.parameters:radius', lower=0.01, upper=0.10, ref0=0.01, ref=0.10)
transcription = dm.Radau(num_segments=5, order=3, compressed=True)
ascent = dm.Phase(transcription=transcription, ode_class=CannonballODE)
traj.add_phase('ascent', ascent)
ascent.add_state('r', units='m', rate_source='r_dot')
ascent.add_state('h', units='m', rate_source='h_dot')
ascent.add_state('gam', units='rad', rate_source='gam_dot')
ascent.add_state('v', units='m/s', rate_source='v_dot')
# All initial states except flight path angle are fixed
# Final flight path angle is fixed (we will set it to zero so that the phase ends at apogee)
ascent.set_time_options(fix_initial=True, duration_bounds=(1, 100), duration_ref=100, units='s')
ascent.set_state_options('r', fix_initial=True, fix_final=False)
ascent.set_state_options('h', fix_initial=True, fix_final=False)
ascent.set_state_options('gam', fix_initial=False, fix_final=True)
ascent.set_state_options('v', fix_initial=False, fix_final=False)
ascent.add_parameter('radius', units='m', dynamic=False)
ascent.add_parameter('density', units='kg/m**3', dynamic=False)
ascent.add_parameter('CD', val=0.5, dynamic=False)
# Limit the muzzle energy
ascent.add_boundary_constraint('ke', loc='initial', units='J',
upper=400000, lower=0, ref=100000)
# Second Phase (descent)
transcription = dm.GaussLobatto(num_segments=2, order=5, compressed=False)
descent = dm.Phase(transcription=transcription, ode_class=CannonballODE)
traj.add_phase('descent', descent )
# All initial states and time are free (they will be linked to the final states of ascent.
# Final altitude is fixed (we will set it to zero so that the phase ends at ground impact)
descent.set_time_options(initial_bounds=(.5, 100), duration_bounds=(.5, 100),
duration_ref=100, units='s')
descent.add_state('r', units='m', rate_source='r_dot')
descent.add_state('h', units='m', rate_source='h_dot')
descent.add_state('gam', units='rad', rate_source='gam_dot')
descent.add_state('v', units='m/s', rate_source='v_dot',)
descent.set_state_options('r', )
descent.set_state_options('h', fix_initial=False, fix_final=True)
descent.set_state_options('gam', fix_initial=False, fix_final=False)
descent.set_state_options('v', fix_initial=False, fix_final=False)
descent.add_parameter('radius', units='m', dynamic=False)
descent.add_parameter('density', units='kg/m**3', dynamic=False)
descent.add_parameter('CD', val=0.5, dynamic=False)
# Link Phases (link time and all state variables)
traj.link_phases(phases=['ascent', 'descent'], vars=['*'])
descent.add_objective('r', loc='final', scaler=-1.0)
# Finish Problem Setup
p.setup()
# Set Initial Guesses
p.set_val('traj.parameters:radius', 0.05, units='m')
p.set_val('traj.parameters:density', 7.87, units='g/cm**3')
# initial guesses
p.set_val('traj.ascent.t_initial', 0.0)
p.set_val('traj.ascent.t_duration', 10.0)
p.set_val('traj.ascent.states:r', ascent.interpolate(ys=[0, 100], nodes='state_input'))
p.set_val('traj.ascent.states:h', ascent.interpolate(ys=[0, 100], nodes='state_input'))
p.set_val('traj.ascent.states:v', ascent.interpolate(ys=[200, 150], nodes='state_input'))
p.set_val('traj.ascent.states:gam', ascent.interpolate(ys=[25, 0], nodes='state_input'),
units='deg')
# more intitial guesses
p.set_val('traj.descent.t_initial', 10.0)
p.set_val('traj.descent.t_duration', 10.0)
p.set_val('traj.descent.states:r', descent.interpolate(ys=[100, 200], nodes='state_input'))
p.set_val('traj.descent.states:h', descent.interpolate(ys=[100, 0], nodes='state_input'))
p.set_val('traj.descent.states:v', descent.interpolate(ys=[150, 200], nodes='state_input'))
p.set_val('traj.descent.states:gam', descent.interpolate(ys=[0, -45], nodes='state_input'),
units='deg')
dm.run_problem(p, simulate=True, make_plots=True)
# p.list_problem_vars(print_arrays=True)