Add plotting to initial guess test function.

opty/tests/test_direct_collocation.py

@@ -1908,7 +1908,7 @@ def test_attributes_read_only():
         with raises(AttributeError):
             setattr(test, XX, 5)
 
-def test_linear_initial_guess():
+def test_linear_initial_guess(plot=False):
     """Test to check if the initial guess is created correctly."""
 
     x, y, ux, uy = mech.dynamicsymbols('x y ux uy')
@@ -2037,6 +2037,9 @@ def obj_grad(free):
     initial_guess = prob.create_linear_initial_guess()
     np.testing.assert_allclose(initial_guess, expected_guess)
 
+    if plot:
+        prob.plot_trajectories(initial_guess)
+
     # A2: np.inf, -np.inf in bounds
     bounds[a1] = (-np.inf, 10.0)
     bounds[a2] = (-10.0, np.inf)
@@ -2081,6 +2084,9 @@ def obj_grad(free):
     initial_guess = prob.create_linear_initial_guess()
     np.testing.assert_allclose(initial_guess, expected_guess)
 
+    if plot:
+        prob.plot_trajectories(initial_guess)
+
     # A4: state instances in instance_constraints
     instance_constraints = (
         x.func(t0) - 3.0 + ux.func(tf),
@@ -2118,6 +2124,9 @@ def obj_grad(free):
     initial_guess = prob.create_linear_initial_guess()
     np.testing.assert_allclose(initial_guess, expected_guess)
 
+    if plot:
+        prob.plot_trajectories(initial_guess)
+
     # B: VARIABLE TIME INTERVAL
     # B1: normal
     h = sym.symbols('h')
@@ -2229,6 +2238,9 @@ def obj_grad(free):
     initial_guess = prob.create_linear_initial_guess()
     np.testing.assert_allclose(initial_guess, expected_guess)
 
+    if plot:
+        prob.plot_trajectories(initial_guess)
+
     # B2: np.inf, -np.inf in bounds
     bounds[a1] = (-np.inf, 10.0)
     bounds[a2] = (-10.0, np.inf)
@@ -2252,6 +2264,9 @@ def obj_grad(free):
     initial_guess = prob.create_linear_initial_guess()
     np.testing.assert_allclose(initial_guess, expected_guess)
 
+    if plot:
+        prob.plot_trajectories(initial_guess)
+
     # B3: no bounds
     expected_guess[4*num_nodes: 5*num_nodes] = 0.0
     expected_guess[5*num_nodes: 6*num_nodes] = 0.0
@@ -2273,6 +2288,9 @@ def obj_grad(free):
     initial_guess = prob.create_linear_initial_guess()
     np.testing.assert_allclose(initial_guess, expected_guess)
 
+    if plot:
+        prob.plot_trajectories(initial_guess)
+
     # B4: state instances in instance_constraints
     instance_constraints = (
         x.func(t00) - 3.0 + ux.func(tf0),
@@ -2309,3 +2327,6 @@ def obj_grad(free):
 
     initial_guess = prob.create_linear_initial_guess()
     np.testing.assert_allclose(initial_guess, expected_guess)
+
+    if plot:
+        prob.plot_trajectories(initial_guess)