better way (slicing) to get fewer frames

examples-gallery/plot_two_link_pendulum_on_a_cart.py

@@ -1,4 +1,3 @@
-# %%
 """
 Upright a Double Pendulum
 =========================
@@ -43,7 +42,7 @@
 from matplotlib import patches
 
 # %%
-# Generate the nonholonomic equations of motion of the system.
+# Generate the equations of motion of the system.
 N, A1, A2 = sm.symbols('N A1, A2', cls=me.ReferenceFrame)
 t = me.dynamicsymbols._t
 O, P1, P2, P3 = sm.symbols('O P1 P2 P3', cls=me.Point)
@@ -78,10 +77,15 @@
 q_ind = [q1, q2, q3]
 u_ind = [u1, u2, u3]
 
-KM = me.KanesMethod(N, q_ind=q_ind, u_ind=u_ind, kd_eqs=kd)
-(fr, frstar) = KM.kanes_equations(bodies, loads=loads)
-EOM = kd.col_join(fr + frstar)
-sm.pprint(sm.trigsimp(EOM))
+KM = me.KanesMethod(
+    N,
+    q_ind=q_ind,
+    u_ind=u_ind,
+    kd_eqs=kd
+)
+fr, frstar = KM.kanes_equations(bodies, loads=loads)
+eom = kd.col_join(fr + frstar)
+sm.pprint(sm.trigsimp(eom))
 
 # %%
 # Define various objects to be use in the optimization problem.
@@ -143,25 +147,28 @@ def obj_grad(free):
                               initial_state_constraints.items()) +
                         tuple(xi.subs({t: duration}) - xi_val for xi, xi_val in
                               final_state_constraints.items()))
-print(instance_constraints)
 
 # %%
 # Bounding h > 0 helps avoid 'solutions' with h < 0.
-bounds = {F: (-150.0, 150.0), q1: (-5.0, 5.0), h: (1.e-5, 1.0)}
+bounds = {F: (-150.0, 150.0),
+          q1: (-5.0, 5.0),
+          h: (0.0, 1.0)
+}
 
 # %%
 # Create an optimization problem and solve it.
 prob = Problem(
     obj,
     obj_grad,
-    EOM,
+    eom,
     state_symbols,
     num_nodes,
     interval_value,
     known_parameter_map=par_map,
     instance_constraints=instance_constraints,
     time_symbol=t,
-    bounds=bounds)
+    bounds=bounds
+)
 
 # Initial guess.
 initial_guess = np.zeros(prob.num_free)
@@ -179,11 +186,18 @@ def obj_grad(free):
 prob.plot_trajectories(initial_guess)
 
 # Find the optimal solution.
+# As initial guess the solution of a previous run, stored in
+# 'two_link_pendulum_on_a_cart_solution.npy' is used.
+initial_guess = np.load('two_link_pendulum_on_a_cart_solution.npy')
 solution, info = prob.solve(initial_guess)
 print('Message from optimizer:', info['status_msg'])
 print('Iterations needed', len(prob.obj_value))
 print(f"Objective value {solution[-1]: .3e}")
 
+# %%
+#This is where the solution is saved to give better initial conditions
+# ```np.save('two_link_pendulum_on_a_cart_solution.npy', solution)```
+
 # %%
 # Plot the evolution of the objective function.
 prob.plot_objective_value()
@@ -199,18 +213,12 @@ def obj_grad(free):
 # %%
 # animate the solution.
 
-solution2 = []
 state_sol, _, _, h_var = parse_free(solution, len(state_symbols),
         len(specified_symbols),num_nodes, variable_duration=True)
-for i in range(len(state_symbols)):
-    solution1 = []
-    for j in range(num_nodes):
-        if j % 4 == 0:
-            solution1.append(state_sol[i][j])
-    solution2 += solution1
-num_nodes = len(solution1)
-solution = solution2 + [h_var]
-
+state_sol1 = state_sol.T[::4, :]
+num_nodes = state_sol1.shape[0]
+print('num nodes', num_nodes)
+solution = list(state_sol1.T.flatten()) + [h_var]
 
 P1_x = np.empty(num_nodes)
 P1_y = np.empty(num_nodes)