#!/usr/bin/env python
"""Try out odeint for a second-order differential equation, namely the equation of motion:

a = (q/m)*(v cross B)

where 

B is constant and given. 

The idea is to have a way of testing the accuracy of the solver odeint. 

How can this be done? Try to implement it.

"""



from numpy import *
from math import *
from scipy.integrate import odeint
from numpy.linalg import norm

B_field = [0.,0.,1.]


def func(y,dt):
    x,v = hsplit(y,2)
    #print 'x, v = ', x,v
    xp = v
    vp = multiply(-1.,cross(v,B_field))
    #print 'vp = ', vp
    #print 'radius = ', norm(x)
    z = hstack([xp, vp])
    return z


def main():
    vel = [0.,1.,0.]
    x   = [1.,0.,0.]
    t = [0.]
    dt = 1.e-1
    print 'initial radius = ', sqrt(sum(power(x,2)))
    #t = [0.1,0.2,0.4,0.6,0.8,1.]
    for i in xrange(0,2000):
        #t = [0.1*i]
        t = t[-1]
        t = append(t,dt*i)
        y = hstack([x,vel])
        #print 'y: ', y
        z = odeint(func,y,t)
        x,vel  = hsplit(z[-1],2)
        #print 'new position: ', x_new
        #print i, norm(x)
        if i%100 == 0: print i, 'at t = ', t[-1],' new radius = ', norm(x), 'x = ', x
    #print t

main()