def write_image(k):
	q,p,mu,q1 = jpf.state_to_weinstein_darboux( y_data[k] )
	f = gf.display_velocity(q,p,mu)
	time_s = str(time_data[k])
	plt.suptitle('t = '+ time_s[0:4] , fontsize=16 , x = 0.75 , y = 0.25 )
	fname = './movie_frames/frame_%03i.png' % k
	f.savefig( fname )
	plt.close(f)
	sys.stdout.write(' '+str(k))
	sys.stdout.flush()
print 'initial momentum: %.3f,%.3f  %.3f' % ( pi[0], pi[1], Li[0][1] )

print 'initial J_R^1 momenta:'
Ki = np.zeros([N,DIM,DIM])
for i in range(0,N):
    Ki[i] = jpf.Jr1_momentum(q,p,mu,q1,particles=[i])
    print Ki[i]

state =  jpf.weinstein_darboux_to_state( q , p , mu , q1 )
step_max = 406
t_span = np.linspace( 0. , T , step_max )
y_span = odeint( jpf.ode_function , state , t_span , rtol=0.0000001 )
np.save('state_data',y_span)
np.save('time_data',t_span)

q,p,mu,q1 = jpf.state_to_weinstein_darboux( y_span[step_max-1] )

Ef = jpf.Hamiltonian(q,p,mu)
pf = jpf.lin_momentum(q,p,mu)
Lf = jpf.ang_momentum(q,p,mu)
print '  final energy: %.3f  diff = %.3e' % ( Ef, Ef-Ei )
print '  final momentum: %.3f,%.3f  %.3f  diff = %.3e %.3e' % \
    ( pf[0], pf[1], Lf[0][1], np.linalg.norm(pf-pi), np.fabs(Lf[0][1]-Li[0][1]) )
print '  final position: %.2f,%.2f' % ( q[0][0], q[0][1] )
print 'dist = %.3e   p_0 = %.3e' % ( np.linalg.norm(q[1]-q[0]), np.linalg.norm(p[0]) )

print '  final J_R^1 momenta:'
Kf = np.zeros([N,DIM,DIM])
for i in range(0,N):
    Kf[i] = jpf.Jr1_momentum(q,p,mu,q1,particles=[i])
    print Kf[i]
import numpy as np
import sys

N = jpf.N
DIM = jpf.DIM
SIGMA = jpf.SIGMA

y_data = np.load('./state_data.npy')
times  = np.load('./time_data.npy')

N_timestep = y_data.shape[0]

momenta = np.zeros([N,3,N_timestep])

for i in range(N_timestep):
    q,p,mu,q1 = jpf.state_to_weinstein_darboux( y_data[i] )

    for j in range(N):
        tmp = jpf.lin_momentum(q,p,mu,[j])
        momenta[j][0][i] = tmp[0]
        momenta[j][1][i] = tmp[1]
        momenta[j][2][i] = jpf.ang_momentum(q,p,mu,[j])[0][1]

plt.figure()
plt.xlabel('t')
plt.plot(times,momenta[0][2],'b-',
         times,momenta[1][2],'r-',
         times,momenta[2][2],'g-')
#plt.axis([0,times[N_timestep-1],0,1])
plt.show()