Skip to content

Commit

Permalink
Cleaned up assert_cases_equal test utility a bit. (#779)
Browse files Browse the repository at this point in the history
* cleaned up assert_cases_equal test utility a bit.

* pep8 fixes
  • Loading branch information
robfalck authored Jul 11, 2022
1 parent a5f9187 commit bb1cadc
Show file tree
Hide file tree
Showing 4 changed files with 24 additions and 273 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@ def make_traj(transcription='gauss-lobatto', transcription_order=3, compressed=F
traj : dm.Trajectory
The Trajectory object.
"""

t = {'gauss-lobatto': dm.GaussLobatto(num_segments=5, order=transcription_order, compressed=compressed),
'radau': dm.Radau(num_segments=5, order=transcription_order, compressed=compressed)}

Expand Down Expand Up @@ -81,266 +80,6 @@ def make_traj(transcription='gauss-lobatto', transcription_order=3, compressed=F
# Third Phase (burn)
burn2 = dm.Phase(ode_class=FiniteBurnODE, transcription=t[transcription])

if connected:
traj.add_phase('coast', coast)
traj.add_phase('burn2', burn2)

burn2.set_time_options(initial_bounds=(1.0, 60), duration_bounds=(-10.0, -0.5),
initial_ref=10, units='TU')
burn2.add_state('r', fix_initial=True, fix_final=False, defect_scaler=100.0,
rate_source='r_dot', units='DU')
burn2.add_state('theta', fix_initial=False, fix_final=False, defect_scaler=100.0,
rate_source='theta_dot', units='rad')
burn2.add_state('vr', fix_initial=True, fix_final=False, defect_scaler=1000.0,
rate_source='vr_dot', units='DU/TU')
burn2.add_state('vt', fix_initial=True, fix_final=False, defect_scaler=1000.0,
rate_source='vt_dot', units='DU/TU')
burn2.add_state('accel', fix_initial=False, fix_final=False, defect_scaler=1.0,
rate_source='at_dot', units='DU/TU**2')
burn2.add_state('deltav', fix_initial=False, fix_final=False, defect_scaler=1.0,
rate_source='deltav_dot', units='DU/TU')

burn2.add_objective('deltav', loc='initial', scaler=100.0)

burn2.add_control('u1', rate_continuity=True, rate2_continuity=True, units='deg',
scaler=0.01, lower=-180, upper=180)
else:
traj.add_phase('coast', coast)
traj.add_phase('burn2', burn2)

burn2.set_time_options(initial_bounds=(0.5, 50), duration_bounds=(.5, 10), initial_ref=10, units='TU')
burn2.add_state('r', fix_initial=False, fix_final=True, defect_scaler=100.0,
rate_source='r_dot', targets=['r'], units='DU')
burn2.add_state('theta', fix_initial=False, fix_final=False, defect_scaler=100.0,
rate_source='theta_dot', targets=['theta'], units='rad')
burn2.add_state('vr', fix_initial=False, fix_final=True, defect_scaler=1000.0,
rate_source='vr_dot', targets=['vr'], units='DU/TU')
burn2.add_state('vt', fix_initial=False, fix_final=True, defect_scaler=1000.0,
rate_source='vt_dot', targets=['vt'], units='DU/TU')
burn2.add_state('accel', fix_initial=False, fix_final=False, defect_scaler=1.0,
rate_source='at_dot', targets=['accel'], units='DU/TU**2')
burn2.add_state('deltav', fix_initial=False, fix_final=False, defect_scaler=1.0,
rate_source='deltav_dot', units='DU/TU')

burn2.add_objective('deltav', loc='final', scaler=100.0)

burn2.add_control('u1', rate_continuity=True, rate2_continuity=True, units='deg',
scaler=0.01, lower=-180, upper=180, targets=['u1'])

burn1.add_timeseries_output('pos_x')
coast.add_timeseries_output('pos_x')
burn2.add_timeseries_output('pos_x')

burn1.add_timeseries_output('pos_y')
coast.add_timeseries_output('pos_y')
burn2.add_timeseries_output('pos_y')

# Link Phases
if connected:
traj.link_phases(phases=['burn1', 'coast'],
vars=['time', 'r', 'theta', 'vr', 'vt', 'deltav'],
connected=True)

# No direct connections to the end of a phase.
traj.link_phases(phases=['burn2', 'coast'],
vars=['r', 'theta', 'vr', 'vt', 'deltav'],
locs=('final', 'final'))
traj.link_phases(phases=['burn2', 'coast'],
vars=['time'], locs=('final', 'final'))

traj.link_phases(phases=['burn1', 'burn2'], vars=['accel'],
locs=('final', 'final'))

else:
traj.link_phases(phases=['burn1', 'coast', 'burn2'],
vars=['time', 'r', 'theta', 'vr', 'vt', 'deltav'])

traj.link_phases(phases=['burn1', 'burn2'], vars=['accel'])

return traj


def two_burn_orbit_raise_problem(transcription='gauss-lobatto', optimizer='SLSQP', r_target=3.0,
transcription_order=3, compressed=False,
show_output=True, connected=False, restart=None):
"""
Build and run the finite burn orbit raise problem.
Parameters
----------
transcription : str
One of 'radau' or 'gauss-lobatto'.
optimizer : str
The pyoptsparse optimizer to use.
r_target : float
The final radius target, in AU.
transcription_order : int
The order of the state transcription.
compressed : bool
If True, use a compressed transcription.
show_output : bool
If True, display the optimizer output.
connected : bool
If True, connect phases for continuity. Otherwise, enforce continuity via linkage constraints.
restart : str or None
The restart file to use, if available.
Returns
-------
p : om.Problem
The OpenMDAO Problem instance.
"""

p = om.Problem(model=om.Group())

p.driver = om.pyOptSparseDriver()
p.driver.options['optimizer'] = optimizer
p.driver.declare_coloring()
if optimizer == 'SNOPT':
p.driver.opt_settings['Major iterations limit'] = 100
p.driver.opt_settings['Major feasibility tolerance'] = 1.0E-6
p.driver.opt_settings['Major optimality tolerance'] = 1.0E-6
if show_output:
p.driver.opt_settings['iSumm'] = 6
elif optimizer == 'IPOPT':
p.driver.opt_settings['mu_init'] = 1e-3
p.driver.opt_settings['max_iter'] = 500
p.driver.opt_settings['acceptable_tol'] = 1e-3
p.driver.opt_settings['constr_viol_tol'] = 1e-3
p.driver.opt_settings['compl_inf_tol'] = 1e-3
p.driver.opt_settings['acceptable_iter'] = 0
p.driver.opt_settings['tol'] = 1e-3
p.driver.opt_settings['print_level'] = 5
p.driver.opt_settings['nlp_scaling_method'] = 'gradient-based' # for faster convergence
p.driver.opt_settings['alpha_for_y'] = 'safer-min-dual-infeas'
p.driver.opt_settings['mu_strategy'] = 'monotone'

traj = make_traj(transcription=transcription, transcription_order=transcription_order,
compressed=compressed, connected=connected)
p.model.add_subsystem('traj', subsys=traj)

# Needed to move the direct solver down into the phases for use with MPI.
# - After moving down, used fewer iterations (about 30 less)

p.setup(check=True)

# Set Initial Guesses
p.set_val('traj.parameters:c', val=1.5, units='DU/TU')

burn1 = p.model.traj.phases.burn1
burn2 = p.model.traj.phases.burn2
coast = p.model.traj.phases.coast

if burn1 in p.model.traj.phases._subsystems_myproc:
p.set_val('traj.burn1.t_initial', val=0.0)
p.set_val('traj.burn1.t_duration', val=2.25)
p.set_val('traj.burn1.states:r', val=burn1.interp('r', [1, 1.5]))
p.set_val('traj.burn1.states:theta', val=burn1.interp('theta', [0, 1.7]))
p.set_val('traj.burn1.states:vr', val=burn1.interp('vr', [0, 0]))
p.set_val('traj.burn1.states:vt', val=burn1.interp('vt', [1, 1]))
p.set_val('traj.burn1.states:accel', val=burn1.interp('accel', [0.1, 0]))
p.set_val('traj.burn1.states:deltav', val=burn1.interp('deltav', [0, 0.1]))
p.set_val('traj.burn1.controls:u1', val=burn1.interp('u1', [-3.5, 13.0]))

if coast in p.model.traj.phases._subsystems_myproc:
p.set_val('traj.coast.t_initial', val=2.25)
p.set_val('traj.coast.t_duration', val=3.0)

p.set_val('traj.coast.states:r', val=coast.interp('r', [1.3, 1.5]))
p.set_val('traj.coast.states:theta', val=coast.interp('theta', [2.1767, 1.7]))
p.set_val('traj.coast.states:vr', val=coast.interp('vr', [0.3285, 0]))
p.set_val('traj.coast.states:vt', val=coast.interp('vt', [0.97, 1]))
p.set_val('traj.coast.states:accel', val=coast.interp('accel', [0, 0]))

if burn2 in p.model.traj.phases._subsystems_myproc:
if connected:
p.set_val('traj.burn2.t_initial', val=7.0)
p.set_val('traj.burn2.t_duration', val=-1.75)

p.set_val('traj.burn2.states:r', val=burn2.interp('r', [r_target, 1]))
p.set_val('traj.burn2.states:theta', val=burn2.interp('theta', [4.0, 0.0]))
p.set_val('traj.burn2.states:vr', val=burn2.interp('vr', [0, 0]))
p.set_val('traj.burn2.states:vt', val=burn2.interp('vt', [np.sqrt(1 / r_target), 1]))
p.set_val('traj.burn2.states:deltav', val=burn2.interp('deltav', [0.2, 0.1]))
p.set_val('traj.burn2.states:accel', val=burn2.interp('accel', [0., 0.1]))

else:
p.set_val('traj.burn2.t_initial', val=5.25)
p.set_val('traj.burn2.t_duration', val=1.75)

p.set_val('traj.burn2.states:r', val=burn2.interp('r', [1, r_target]))
p.set_val('traj.burn2.states:theta', val=burn2.interp('theta', [0, 4.0]))
p.set_val('traj.burn2.states:vr', val=burn2.interp('vr', [0, 0]))
p.set_val('traj.burn2.states:vt', val=burn2.interp('vt', [1, np.sqrt(1 / r_target)]))
p.set_val('traj.burn2.states:deltav', val=burn2.interp('deltav', [0.1, 0.2]))
p.set_val('traj.burn2.states:accel', val=burn2.interp('accel', [0.1, 0]))
p.set_val('traj.burn2.controls:u1', val=burn2.interp('u1', [0, 0]))

# p.run_driver()
#
# return p

dm.run_problem(p, run_driver=True, simulate=True, restart=restart)

return p


def make_traj(transcription='gauss-lobatto', transcription_order=3, compressed=False,
connected=False):

t = {'gauss-lobatto': dm.GaussLobatto(num_segments=5, order=transcription_order, compressed=compressed),
'radau': dm.Radau(num_segments=20, order=transcription_order, compressed=compressed)}

traj = dm.Trajectory()

traj.add_parameter('c', opt=False, val=1.5, units='DU/TU',
targets={'burn1': ['c'], 'burn2': ['c']})

# First Phase (burn)

burn1 = dm.Phase(ode_class=FiniteBurnODE, transcription=t[transcription])

burn1 = traj.add_phase('burn1', burn1)

burn1.set_time_options(fix_initial=True, duration_bounds=(.5, 10), units='TU')
burn1.add_state('r', fix_initial=True, fix_final=False, defect_scaler=100.0,
rate_source='r_dot', units='DU')
burn1.add_state('theta', fix_initial=True, fix_final=False, defect_scaler=100.0,
rate_source='theta_dot', units='rad')
burn1.add_state('vr', fix_initial=True, fix_final=False, defect_scaler=100.0,
rate_source='vr_dot', units='DU/TU')
burn1.add_state('vt', fix_initial=True, fix_final=False, defect_scaler=100.0,
rate_source='vt_dot', units='DU/TU')
burn1.add_state('accel', fix_initial=True, fix_final=False,
rate_source='at_dot', units='DU/TU**2')
burn1.add_state('deltav', fix_initial=True, fix_final=False,
rate_source='deltav_dot', units='DU/TU')
burn1.add_control('u1', rate_continuity=True, rate2_continuity=True, units='deg', scaler=0.01,
rate_continuity_scaler=0.001, rate2_continuity_scaler=0.001,
lower=-30, upper=30)
# Second Phase (Coast)
coast = dm.Phase(ode_class=FiniteBurnODE, transcription=t[transcription])

coast.set_time_options(initial_bounds=(0.5, 20), duration_bounds=(.5, 50), duration_ref=50, units='TU')
coast.add_state('r', fix_initial=False, fix_final=False, defect_scaler=100.0,
rate_source='r_dot', units='DU')
coast.add_state('theta', fix_initial=False, fix_final=False, defect_scaler=100.0,
rate_source='theta_dot', units='rad')
coast.add_state('vr', fix_initial=False, fix_final=False, defect_scaler=100.0,
rate_source='vr_dot', units='DU/TU')
coast.add_state('vt', fix_initial=False, fix_final=False, defect_scaler=100.0,
rate_source='vt_dot', units='DU/TU')
coast.add_state('accel', fix_initial=True, fix_final=True,
rate_source='at_dot', units='DU/TU**2')
coast.add_state('deltav', fix_initial=False, fix_final=False,
rate_source='deltav_dot', units='DU/TU')

coast.add_parameter('u1', opt=False, val=0.0, units='deg')
coast.add_parameter('c', val=0.0, units='DU/TU')

# Third Phase (burn)
burn2 = dm.Phase(ode_class=FiniteBurnODE, transcription=t[transcription])

if connected:
traj.add_phase('burn2', burn2)
traj.add_phase('coast', coast)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ def test_restart_from_solution_radau_to_connected(self):

# Verify that the second case has the same inputs and outputs
assert_cases_equal(case1, p, tol=1.0E-9, require_same_vars=False)
assert_cases_equal(sim_case1, sim_case2, tol=1.0E-9)
assert_cases_equal(sim_case1, sim_case2, tol=1.0E-8)


if __name__ == '__main__': # pragma: no cover
Expand Down
13 changes: 8 additions & 5 deletions dymos/test/test_testing_utils.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import difflib
import os
import unittest

Expand All @@ -16,9 +17,8 @@ def tearDown(self):
for file in ('p1.db', 'p2.db'):
try:
os.remove(file)
print('removed', file)
except:
print(f'no file named {file}')
pass

def test_different_variables(self):

Expand Down Expand Up @@ -161,9 +161,12 @@ def test_different_values(self):
c1 = om.CaseReader('p1.db').get_case('final')
c2 = om.CaseReader('p2.db').get_case('final')

expected = "\nThe following variables contain different values:\nvar: " \
"error\na: [0. 2. 4.]\nb: [[0. 1. 2.]\n [3. 4. 5.]]\n" \
"c: [[4. 0. 0.]\n [0. 4. 0.]\n [0. 0. 4.]]"
expected = "\nThe following variables contain different values:\n" \
"var max error mean error\n" \
"--- ---------------- ----------------\n" \
" a 4.000000000e+00 2.000000000e+00\n" \
" b 5.000000000e+00 2.500000000e+00\n" \
" c 4.000000000e+00 1.333333333e+00\n"

with self.assertRaises(AssertionError) as e:
assert_cases_equal(c1, c2)
Expand Down
21 changes: 15 additions & 6 deletions dymos/utils/testing_utils.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import io

import numpy as np

from scipy.interpolate import interp1d
Expand Down Expand Up @@ -86,9 +88,9 @@ def assert_cases_equal(case1, case2, tol=1.0E-12, require_same_vars=True):
diff_err_msg += f'\nVariables in case2 but not in case1: {sorted(case2_minus_case1)}'

shape_errors = set()
val_errors = set()
val_errors = {}
shape_err_msg = '\nThe following variables have different shapes/sizes:'
val_err_msg = '\nThe following variables contain different values:\nvar: error'
val_err_msg = io.StringIO()

for var in sorted(set(case1_vars.keys()).intersection(case2_vars.keys())):
a = case1_vars[var]['val']
Expand All @@ -98,17 +100,24 @@ def assert_cases_equal(case1, case2, tol=1.0E-12, require_same_vars=True):
shape_err_msg += f'\n{var} has shape {a.shape} in case1 but shape {b.shape} in case2'
continue
err = np.abs(a - b)
if np.any(err > tol):
val_errors.add(var)
val_err_msg += f'\n{var}: {err}'
max_err = np.max(err)
mean_err = np.mean(err)
if np.any(max_err > tol):
val_errors[var] = (max_err, mean_err)

err_msg = ''
if diff_err_msg:
err_msg += diff_err_msg
if shape_errors:
err_msg += shape_err_msg
if val_errors:
err_msg += val_err_msg
val_err_msg.write('\nThe following variables contain different values:\n')
max_var_len = max(3, max([len(s) for s in val_errors.keys()]))
val_err_msg.write(f"{'var'.rjust(max_var_len)} {'max error'.rjust(16)} {'mean error'.rjust(16)}\n")
val_err_msg.write(max_var_len * '-' + ' ' + 16 * '-' + ' ' + 16 * '-' + '\n')
for varname, (max_err, mean_err) in val_errors.items():
val_err_msg.write(f"{varname.rjust(max_var_len)} {max_err:16.9e} {mean_err:16.9e}\n")
err_msg += val_err_msg.getvalue()

if err_msg:
raise AssertionError(err_msg)
Expand Down

0 comments on commit bb1cadc

Please sign in to comment.