Linkage units - Fixes Issue 274 (#275)

* fix for linkage variable units when variables are not states, controls, or time in the first phase in the linkage. * corrected placement of use_tempdirs in the two burn orbit raise tests
OpenMDAO · Feb 5, 2020 · 51cb73c · 51cb73c
1 parent 2bdb9a3
commit 51cb73c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/dymos/examples/finite_burn_orbit_raise/test/test_ex_two_burn_orbit_raise.py b/dymos/examples/finite_burn_orbit_raise/test/test_ex_two_burn_orbit_raise.py
@@ -263,9 +263,9 @@ def two_burn_orbit_raise_problem(transcription='gauss-lobatto', optimizer='SLSQP
     return p
 
 
+@use_tempdirs
 class TestExampleTwoBurnOrbitRaise(unittest.TestCase):
 
-    @use_tempdirs
     def test_ex_two_burn_orbit_raise(self):
         _, optimizer = set_pyoptsparse_opt('SNOPT', fallback=False)
 
@@ -279,9 +279,9 @@ def test_ex_two_burn_orbit_raise(self):
 
 
 # This test is separate because connected phases aren't directly parallelizable.
+@use_tempdirs
 class TestExampleTwoBurnOrbitRaiseConnected(unittest.TestCase):
 
-    @use_tempdirs
     def test_ex_two_burn_orbit_raise_connected(self):
         _, optimizer = set_pyoptsparse_opt('SNOPT', fallback=False)
 

diff --git a/dymos/trajectory/trajectory.py b/dymos/trajectory/trajectory.py
@@ -372,12 +372,21 @@ def _setup_linkages(self):
                         p2.set_time_options(input_initial=True)
                 else:
                     vars_to_constrain.append(var)
-                    if var in p1_states:
+                    if var in p1.state_options:
                         units_map[var] = p1.state_options[var]['units']
                         shape_map[var] = p1.state_options[var]['shape']
-                    elif var in p1_controls:
+                    elif var in p1.control_options:
                         units_map[var] = p1.control_options[var]['units']
                         shape_map[var] = p1.control_options[var]['shape']
+                    elif var in p1.polynomial_control_options:
+                        units_map[var] = p1.polynomial_control_options[var]['units']
+                        shape_map[var] = p1.polynomial_control_options[var]['shape']
+                    elif var in p1.design_parameter_options:
+                        units_map[var] = p1.design_parameter_options[var]['units']
+                        shape_map[var] = p1.design_parameter_options[var]['shape']
+                    elif var in p1.input_parameter_options:
+                        units_map[var] = p1.input_parameter_options[var]['units']
+                        shape_map[var] = p1.input_parameter_options[var]['shape']
                     elif var == 'time':
                         units_map[var] = p1.time_options['units']
                         shape_map[var] = (1,)