def test_integrate(self):
h = 1.
x0 = np.array([0, 0, 1.])
dt, vu0, uu = .01, .01, .5
self.s.electrode("rf").rf = uu*h**2
t, x, v = [], [], []
for ti, xi, vi in self.s.trajectory(
x0, np.array([0, 0, vu0*uu*h]), axis=(1, 2),
t1=20*2*np.pi, dt=dt*2*np.pi):
t.append(ti)
x.append(xi)
v.append(vi)
t = np.array(t)
x = np.array(x)
v = np.array(v)
self.assertEqual(np.alltrue(utils.norm(x, axis=1)<3), True)
self.assertEqual(np.alltrue(utils.norm(v, axis=1)<1), True)
avg = int(1/dt)
kin = (((x[:-avg]-x[avg:])/(2*np.pi))**2).sum(axis=-1)/2*4 # 4?
pot = self.s.potential(np.array([x0[0]+0*x[:,0], x[:,0], x[:,1]]).T)
pot = pot[avg/2:-avg/2]
t = t[avg/2:-avg/2]
do_avg = lambda ar: ar[:ar.size/avg*avg].reshape(
(-1, avg)).mean(axis=-1)
t, kin, pot = map(do_avg, (t, kin, pot))
self.assertEqual(np.alltrue(np.std(kin+pot)/np.mean(kin+pot)<.01),
True)
def test_norm(self):
self.assertEqual(utils.norm([1,2,3.]), 14**.5)
self.assertEqual(utils.norm([[1,2,3.]], 1), 14**.5)
