<s> from pylab import * <EOL> import matplotlib . animation <EOL> import scipy . interpolate <EOL> import scipy . integrate <EOL> import sklearn <EOL> import sklearn . preprocessing <EOL> import control <EOL> import transformations <EOL> import sympy <EOL> import sympy . matrices <EOL> import sympy . physics <EOL> import sympy . physics . mechanics <EOL> import sympy . physics . mechanics . functions <EOL> import trigutils <EOL> import sympyutils <EOL> import pathutils <EOL> m = <NUM_LIT:1.0> <EOL> g = <NUM_LIT> <EOL> d = <NUM_LIT:1.0> <EOL> I = m * d ** <NUM_LIT> <EOL> f_gravity_world = matrix ( [ - m * g , <NUM_LIT:0> ] ) . T <EOL> f_external_world = f_gravity_world <EOL> construct_sympy_expressions = False <EOL> if construct_sympy_expressions : <EOL> print "<STR_LIT>" <EOL> d_expr = sympy . Symbol ( "<STR_LIT:d>" ) <EOL> m_expr = sympy . Symbol ( "<STR_LIT:m>" ) <EOL> I_expr = sympy . Symbol ( "<STR_LIT:I>" ) <EOL> f_external_expr , f_external_expr_entries = sympyutils . construct_matrix_and_entries ( "<STR_LIT>" , ( <NUM_LIT:2> , <NUM_LIT:1> ) ) <EOL> t_expr = sympy . Symbol ( "<STR_LIT:t>" ) <EOL> p_y_expr = sympy . physics . mechanics . dynamicsymbols ( "<STR_LIT>" ) <EOL> p_x_expr = sympy . physics . mechanics . dynamicsymbols ( "<STR_LIT>" ) <EOL> theta_expr = sympy . physics . mechanics . dynamicsymbols ( "<STR_LIT>" ) <EOL> q_expr = sympy . Matrix ( [ p_y_expr , p_x_expr , theta_expr ] ) <EOL> q_dot_expr = q_expr . diff ( t_expr ) <EOL> u_expr , u_expr_entries = sympyutils . construct_matrix_and_entries ( "<STR_LIT:u>" , ( <NUM_LIT:2> , <NUM_LIT:1> ) ) <EOL> print "<STR_LIT>" <EOL> H_00_expr = sympy . Matrix ( m_expr * sympy . eye ( <NUM_LIT:2> ) ) <EOL> H_01_expr = sympy . Matrix . zeros ( <NUM_LIT:2> , <NUM_LIT:1> ) <EOL> H_10_expr = sympy . Matrix . zeros ( <NUM_LIT:1> , <NUM_LIT:2> ) <EOL> H_11_expr = sympy . Matrix ( [ I_expr ] ) <EOL> C_00_expr = sympy . Matrix . zeros ( <NUM_LIT:2> , <NUM_LIT:2> ) <EOL> C_01_expr = sympy . Matrix . zeros ( <NUM_LIT:2> , <NUM_LIT:1> ) <EOL> C_10_expr = sympy . Matrix . zeros ( <NUM_LIT:1> , <NUM_LIT:2> ) <EOL> C_11_expr = sympy . Matrix . zeros ( <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> G_0_expr = - f_external_expr <EOL> G_1_expr = sympy . Matrix . zeros ( <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> B_0_expr = sympy . Matrix ( [ [ sympy . cos ( theta_expr ) , sympy . cos ( theta_expr ) ] , [ - sympy . sin ( theta_expr ) , - sympy . sin ( theta_expr ) ] ] ) <EOL> B_1_expr = sympy . Matrix ( [ [ - d_expr , d_expr ] ] ) <EOL> H_expr = sympyutils . construct_matrix_from_block_matrix ( sympy . Matrix ( [ [ H_00_expr , H_01_expr ] , [ H_10_expr , H_11_expr ] ] ) ) <EOL> C_expr = sympyutils . construct_matrix_from_block_matrix ( sympy . Matrix ( [ [ C_00_expr , C_01_expr ] , [ C_10_expr , C_11_expr ] ] ) ) <EOL> G_expr = sympyutils . construct_matrix_from_block_matrix ( sympy . Matrix ( [ [ G_0_expr ] , [ G_1_expr ] ] ) ) <EOL> B_expr = sympyutils . construct_matrix_from_block_matrix ( sympy . Matrix ( [ [ B_0_expr ] , [ B_1_expr ] ] ) ) <EOL> q_dot_dot_expr = H_expr . inv ( ) * ( B_expr * u_expr - ( C_expr * q_dot_expr + G_expr ) ) <EOL> dqdotdot_dq_expr = q_dot_dot_expr . jacobian ( q_expr ) <EOL> dqdotdot_dqdot_expr = q_dot_dot_expr . jacobian ( q_dot_expr ) <EOL> const_syms = hstack ( [ d_expr , m_expr , I_expr , matrix ( f_external_expr ) . A1 ] ) <EOL> x_and_u_syms = hstack ( [ matrix ( q_expr ) . A1 , matrix ( q_dot_expr ) . A1 , matrix ( u_expr ) . A1 ] ) <EOL> print "<STR_LIT>" <EOL> const_vals = hstack ( [ d , m , I , f_external_world . A1 ] ) <EOL> const_subs = dict ( zip ( const_syms , const_vals ) ) <EOL> dqdotdot_dq_expr = sympyutils . nsimplify_matrix ( dqdotdot_dq_expr . subs ( const_subs ) ) <EOL> dqdotdot_dqdot_expr = sympyutils . nsimplify_matrix ( dqdotdot_dqdot_expr . subs ( const_subs ) ) <EOL> print "<STR_LIT>" <EOL> dqdotdot_dq_expr_dummy , dqdotdot_dq_expr_dummy_syms = sympyutils . dummify ( dqdotdot_dq_expr , x_and_u_syms ) <EOL> dqdotdot_dqdot_expr_dummy , dqdotdot_dqdot_expr_dummy_syms = sympyutils . dummify ( dqdotdot_dqdot_expr , x_and_u_syms ) <EOL> print "<STR_LIT>" <EOL> current_source_file_path = pathutils . get_current_source_file_path ( ) <EOL> with open ( current_source_file_path + "<STR_LIT>" , "<STR_LIT:w>" ) as f : f . write ( sympy . srepr ( dqdotdot_dq_expr_dummy ) ) <EOL> with open ( current_source_file_path + "<STR_LIT>" , "<STR_LIT:w>" ) as f : f . write ( sympy . srepr ( dqdotdot_dqdot_expr_dummy ) ) <EOL> with open ( current_source_file_path + "<STR_LIT>" , "<STR_LIT:w>" ) as f : f . write ( sympy . srepr ( sympy . Matrix ( dqdotdot_dq_expr_dummy_syms ) ) ) <EOL> with open ( current_source_file_path + "<STR_LIT>" , "<STR_LIT:w>" ) as f : f . write ( sympy . srepr ( sympy . Matrix ( dqdotdot_dqdot_expr_dummy_syms ) ) ) <EOL> else : <EOL> print "<STR_LIT>" <EOL> current_source_file_path = pathutils . get_current_source_file_path ( ) <EOL> with open ( current_source_file_path + "<STR_LIT>" , "<STR_LIT:r>" ) as f : dqdotdot_dq_expr_dummy = sympy . sympify ( f . read ( ) ) <EOL> with open ( current_source_file_path + "<STR_LIT>" , "<STR_LIT:r>" ) as f : dqdotdot_dqdot_expr_dummy = sympy . sympify ( f . read ( ) ) <EOL> with open ( current_source_file_path + "<STR_LIT>" , "<STR_LIT:r>" ) as f : dqdotdot_dq_expr_dummy_syms = array ( sympy . sympify ( f . read ( ) ) ) . squeeze ( ) <EOL> with open ( current_source_file_path + "<STR_LIT>" , "<STR_LIT:r>" ) as f : dqdotdot_dqdot_expr_dummy_syms = array ( sympy . sympify ( f . read ( ) ) ) . squeeze ( ) <EOL> print "<STR_LIT>" <EOL> dqdotdot_dq_anon_funcs_ufuncify = sympyutils . construct_matrix_anon_funcs_ufuncify ( dqdotdot_dq_expr_dummy , dqdotdot_dq_expr_dummy_syms ) <EOL> dqdotdot_dqdot_anon_funcs_ufuncify = sympyutils . construct_matrix_anon_funcs_ufuncify ( dqdotdot_dqdot_expr_dummy , dqdotdot_dqdot_expr_dummy_syms ) <EOL> print "<STR_LIT>" <EOL> def unpack_state ( x ) : <EOL> p = matrix ( [ x . item ( <NUM_LIT:0> ) , x . item ( <NUM_LIT:1> ) ] ) . T <EOL> theta = x . item ( <NUM_LIT:2> ) <EOL> p_dot = matrix ( [ x . item ( <NUM_LIT:3> ) , x . item ( <NUM_LIT:4> ) ] ) . T <EOL> theta_dot = x . item ( <NUM_LIT:5> ) <EOL> q = matrix ( [ p . item ( <NUM_LIT:0> ) , p . item ( <NUM_LIT:1> ) , theta ] ) . T <EOL> q_dot = matrix ( [ p_dot . item ( <NUM_LIT:0> ) , p_dot . item ( <NUM_LIT:1> ) , theta_dot ] ) . T <EOL> return p , theta , p_dot , theta_dot , q , q_dot <EOL> def compute_manipulator_matrices ( x ) : <EOL> p , theta , p_dot , theta_dot , q , q_dot = unpack_state ( x ) <EOL> H_00 = matrix ( m * identity ( <NUM_LIT:2> ) ) <EOL> H_01 = matrix ( zeros ( ( <NUM_LIT:2> , <NUM_LIT:1> ) ) ) <EOL> H_10 = matrix ( zeros ( ( <NUM_LIT:1> , <NUM_LIT:2> ) ) ) <EOL> H_11 = matrix ( I ) <EOL> C_00 = matrix ( zeros ( ( <NUM_LIT:2> , <NUM_LIT:2> ) ) ) <EOL> C_01 = matrix ( zeros ( ( <NUM_LIT:2> , <NUM_LIT:1> ) ) ) <EOL> C_10 = matrix ( zeros ( ( <NUM_LIT:1> , <NUM_LIT:2> ) ) ) <EOL> C_11 = matrix ( zeros ( ( <NUM_LIT:1> , <NUM_LIT:1> ) ) ) <EOL> G_0 = - f_external_world <EOL> G_1 = matrix ( zeros ( ( <NUM_LIT:1> , <NUM_LIT:1> ) ) ) <EOL> B_0 = matrix ( [ [ cos ( theta ) , cos ( theta ) ] , [ - sin ( theta ) , - sin ( theta ) ] ] ) <EOL> B_1 = matrix ( [ [ - d , d ] ] ) <EOL> H = bmat ( [ [ H_00 , H_01 ] , [ H_10 , H_11 ] ] ) <EOL> C = bmat ( [ [ C_00 , C_01 ] , [ C_10 , C_11 ] ] ) <EOL> G = bmat ( [ [ G_0 ] , [ G_1 ] ] ) <EOL> B = bmat ( [ [ B_0 ] , [ B_1 ] ] ) <EOL> return H , C , G , B <EOL> def compute_manipulator_matrix_derivatives ( x , u ) : <EOL> p , theta , p_dot , theta_dot , q , q_dot = unpack_state ( x ) <EOL> x_and_u_vals = hstack ( [ q . A1 , q_dot . A1 , u . A1 ] ) <EOL> dqdotdot_dq = matrix ( sympyutils . evaluate_matrix_anon_funcs ( dqdotdot_dq_anon_funcs_ufuncify , x_and_u_vals [ newaxis , : ] ) ) <EOL> dqdotdot_dqdot = matrix ( sympyutils . evaluate_matrix_anon_funcs ( dqdotdot_dqdot_anon_funcs_ufuncify , x_and_u_vals [ newaxis , : ] ) ) <EOL> return dqdotdot_dq , dqdotdot_dqdot <EOL> def compute_df_dx_and_df_du ( x , u ) : <EOL> H , C , G , B = compute_manipulator_matrices ( x ) <EOL> dqdotdot_dq , dqdotdot_dqdot = compute_manipulator_matrix_derivatives ( x , u ) <EOL> dqdot_dq = matrix ( zeros ( ( <NUM_LIT:3> , <NUM_LIT:3> ) ) ) <EOL> dqdot_dqdot = matrix ( identity ( <NUM_LIT:3> ) ) <EOL> dqdotdot_dq = dqdotdot_dq <EOL> dqdotdot_dqdot = dqdotdot_dqdot <EOL> dqdot_du = matrix ( zeros ( ( <NUM_LIT:3> , <NUM_LIT:2> ) ) ) <EOL> dqdotdot_du = H . I * B <EOL> df_dx = bmat ( [ [ dqdot_dq , dqdot_dqdot ] , [ dqdotdot_dq , dqdotdot_dqdot ] ] ) <EOL> df_du = bmat ( [ [ dqdot_du ] , [ dqdotdot_du ] ] ) <EOL> return df_dx , df_du <EOL> def compute_x_dot ( x , u ) : <EOL> p , theta , p_dot , theta_dot , q , q_dot = unpack_state ( x ) <EOL> H , C , G , B = compute_manipulator_matrices ( x ) <EOL> f1 = H . I * ( - G - C * q_dot ) <EOL> f2 = H . I * B <EOL> q_dot_dot = f1 + f2 * u <EOL> x_dot = bmat ( [ [ q_dot ] , [ q_dot_dot ] ] ) <EOL> return x_dot <EOL> def compute_state_space_trajectory_and_derivatives ( p , dt ) : <EOL> f_thrust_world_norm_threshold = <NUM_LIT> <EOL> num_timesteps = len ( p ) <EOL> p_dot = c_ [ gradient ( p [ : , <NUM_LIT:0> ] , dt ) , gradient ( p [ : , <NUM_LIT:1> ] , dt ) ] <EOL> p_dot_dot = c_ [ gradient ( p_dot [ : , <NUM_LIT:0> ] , dt ) , gradient ( p_dot [ : , <NUM_LIT:1> ] , dt ) ] <EOL> f_thrust_world = m * p_dot_dot - f_external_world . T . A <EOL> f_thrust_world_norm = linalg . norm ( f_thrust_world , axis = <NUM_LIT:1> ) <EOL> theta_raw = arctan2 ( f_thrust_world [ : , <NUM_LIT:0> ] , f_thrust_world [ : , <NUM_LIT:1> ] ) - ( pi / <NUM_LIT> ) <EOL> t_tmp = arange ( num_timesteps ) <EOL> theta_func = scipy . interpolate . interp1d ( t_tmp [ f_thrust_world_norm > f_thrust_world_norm_threshold ] , theta_raw [ f_thrust_world_norm > f_thrust_world_norm_threshold ] , kind = "<STR_LIT>" ) <EOL> theta = theta_func ( t_tmp ) <EOL> theta = trigutils . compute_continuous_angle_array ( theta ) <EOL> theta_dot = gradient ( theta , dt ) <EOL> theta_dot_dot = gradient ( theta_dot , dt ) <EOL> return p , p_dot , p_dot_dot , theta , theta_dot , theta_dot_dot <EOL> def compute_control_trajectory ( q_q_dot_q_dot_dot ) : <EOL> p , p_dot , p_dot_dot , theta , theta_dot , theta_dot_dot = q_q_dot_q_dot_dot <EOL> num_timesteps = len ( p ) <EOL> u = zeros ( ( num_timesteps , <NUM_LIT:2> ) ) <EOL> for ti in range ( int ( num_timesteps ) ) : <EOL> q_ti = matrix ( [ p [ ti , <NUM_LIT:0> ] , p [ ti , <NUM_LIT:1> ] , theta [ ti ] ] ) . T <EOL> q_dot_ti = matrix ( [ p_dot [ ti , <NUM_LIT:0> ] , p_dot [ ti , <NUM_LIT:1> ] , theta_dot [ ti ] ] ) . T <EOL> q_dot_dot_ti = matrix ( [ p_dot_dot [ ti , <NUM_LIT:0> ] , p_dot_dot [ ti , <NUM_LIT:1> ] , theta_dot_dot [ ti ] ] ) . T <EOL> x_ti = bmat ( [ [ q_ti ] , [ q_dot_ti ] ] ) <EOL> H , C , G , B = compute_manipulator_matrices ( x_ti ) <EOL> f1 = H . I * ( - G - C * q_dot_ti ) <EOL> f2 = H . I * B <EOL> f2_pinv = linalg . pinv ( f2 ) <EOL> u_ti = f2_pinv * ( q_dot_dot_ti - f1 ) <EOL> u [ ti , : ] = u_ti . T <EOL> return u </s>
