Drone in tube

opty/direct_collocation.py

@@ -436,7 +436,7 @@ def plot_trajectories(self, vector, axes=None):
         return axes
 
     @_optional_plt_dep
-    def plot_constraint_violations(self, vector, axes=None):
+    def plot_constraint_violations(self, vector, axes=None, detailed_eoms=False):
         """Returns an axis with the state constraint violations plotted versus
         node number and the instance constraints as a bar graph.
 
@@ -446,6 +446,12 @@ def plot_constraint_violations(self, vector, axes=None):
             The initial guess, solution, or any other vector that is in the
             canonical form.
 
+        detailed_eoms : boolean, optional. If True, the equations of motion
+            will be plotted in a separate plot for each state. Default is False.
+
+        axes : ndarray of AxesSubplot, optional. If given, it is the user's
+            responsibility to provide the correct number of axes.
+
         Returns
         =======
         axes : ndarray of AxesSubplot
@@ -523,16 +529,43 @@ def plot_constraint_violations(self, vector, axes=None):
         con_nodes = range(1, self.collocator.num_collocation_nodes)
 
         if axes is None:
-            fig, axes = plt.subplots(1 + num_plots, 1,
-                                     figsize=(6.4, 1.50*(1 + num_plots)),
-                                     layout='compressed')
+            if detailed_eoms == False or self.collocator.num_states == 1:
+                num_eom_plots = 1
+            else:
+                num_eom_plots = self.collocator.num_states
+            fig, axes = plt.subplots(num_eom_plots + num_plots, 1,
+                            figsize=(6.4, 1.75*(num_eom_plots + num_plots)),
+                            constrained_layout=True)
+
+        else:
+            num_eom_plots = len(axes) - num_plots
 
         axes = np.asarray(axes).ravel()
 
-        axes[0].plot(con_nodes, state_violations.T)
-        axes[0].set_title('Constraint violations')
-        axes[0].set_xlabel('Node Number')
-        axes[0].set_ylabel('EoM violation')
+        if detailed_eoms == False or self.collocator.num_states == 1:
+            axes[0].plot(con_nodes, state_violations.T)
+            axes[0].set_title('Constraint violations')
+            axes[0].set_xlabel('Node Number')
+            axes[0].set_ylabel('EoM violation')
+
+        else:
+            for i in range(self.collocator.num_states):
+                k = i + 1
+                if k in (11,12,13):
+                    msg = 'th'
+                elif k % 10 == 1:
+                    msg = 'st'
+                elif k % 10 == 2:
+                    msg = 'nd'
+                elif k % 10 == 3:
+                    msg = 'rd'
+                else:
+                    msg = 'th'
+
+                axes[i].plot(con_nodes, state_violations[i])
+                axes[i].set_ylabel(f'{str(k)}-{msg} EOM violation')
+            axes[num_eom_plots-1].set_xlabel('Node Number')
+            axes[0].set_title('Constraint violations')
 
         if self.collocator.instance_constraints is not None:
             # reduce the instance constrtaints to 2 digits after the decimal
@@ -575,11 +608,11 @@ def plot_constraint_violations(self, vector, axes=None):
                 inst_constr = instance_constr_plot[beginn: endd]
 
                 width = [0.06*num_ticks for _ in range(num_ticks)]
-                axes[i+1].bar(range(num_ticks), inst_viol,
+                axes[i+num_eom_plots].bar(range(num_ticks), inst_viol,
                               tick_label=[sm.latex(s, mode='inline') for s in
                                           inst_constr], width=width)
-                axes[i+1].set_ylabel('Instance')
-                axes[i+1].set_xticklabels(axes[i+1].get_xticklabels(),
+                axes[i+num_eom_plots].set_ylabel('Instance')
+                axes[i+num_eom_plots].set_xticklabels(axes[i+num_eom_plots].get_xticklabels(),
                                           rotation=rotation)
 
         return axes
@@ -1181,9 +1214,6 @@ def constraints(state_values, specified_values, constant_values,
 
             """
 
-            if state_values.shape[0] < 2:
-                raise ValueError('There should always be at least two states.')
-
             assert state_values.shape == (self.num_states,
                                           self.num_collocation_nodes)
             # n x N - 1
@@ -1562,8 +1592,6 @@ def constraints_jacobian(state_values, specified_values,
             - n*(N - 1) : number of constraints
 
             """
-            if state_values.shape[0] < 2:
-                raise ValueError('There should always be at least two states.')
 
             # Each of these arrays are shape(n, N - 1). The x_adjacent is
             # either the previous value of the state or the next value of

opty/tests/test_direct_collocation.py

@@ -1516,8 +1516,8 @@ def test_known_and_unknown_order():
 
 def test_for_algebraic_eoms():
     """
-    If algebraic equations of motion are given to Problem, a ValueError should
-    be raised. This a a test for this
+    If only algebraic equations of motion are given to Problem, a ValueError
+    should be raised. This a a test for this
     """
 
     target_angle = np.pi
@@ -1566,3 +1566,61 @@ def test_for_algebraic_eoms():
         )
 
     assert excinfo.type is ValueError
+
+def test_one_eom_only():
+    """
+    Only one differential equation should work. This tests for the corrrect
+    shape of the constraints and jacobian.
+
+    """
+    # Equations of motion.
+    t = mech.dynamicsymbols._t
+    y, u = mech.dynamicsymbols('y u')
+
+    eom = sym.Matrix([-y.diff(t) - y**3 + u])
+
+    t0, tf = 0.0, 10.0
+    num_nodes = 100
+    interval_value = (tf - t0)/(num_nodes - 1)
+
+    state_symbols = (y, )
+    specified_symbols = (u,)
+
+    # Specify the objective function and form the gradient.
+    obj_func = sym.Integral(y**2 + u**2, t)
+    obj, obj_grad = create_objective_function(
+        obj_func,
+        state_symbols,
+        specified_symbols,
+        tuple(),
+        num_nodes,
+        node_time_interval=interval_value
+    )
+
+    # Specify the symbolic instance constraints.
+    instance_constraints = (
+        y.func(t0) - 1,
+        y.func(tf) - 1.5,
+    )
+
+        # Create the optimization problem and set any options.
+    prob = Problem(
+        obj,
+        obj_grad,
+        eom,
+        state_symbols,
+        num_nodes,
+        interval_value,
+        instance_constraints=instance_constraints,
+        )
+
+    initial_guess = np.zeros(prob.num_free)
+    initial_guess[0] = 1.0
+    initial_guess[num_nodes-1] = 1.5
+
+    # assert that prob.constraints and prob.jacobian have the correct shape.
+    length = 1*(num_nodes-1) + 2
+    assert prob.constraints(initial_guess).shape == (length,)
+
+    length = (2*1 + 1 + 0 + 0) * (1*(num_nodes-1)) + 2
+    assert prob.jacobian(initial_guess).shape == (length,)