Tests the thin fixed order multipole implementations (ThinQuadrupole
, ThinSkewQuadrupole
, ThinSextupole
and ThinOctupole
) against the multipole kick formula and times the results.
from __future__ import division, print_function
import numpy as np
from scipy.constants import m_p, c, e
# sets the PyHEADTAIL directory etc.
try:
from settings import *
except:
pass
import PyHEADTAIL.multipoles.multipoles as multip
PyHEADTAIL v1.10.5.271
class B(object):
def __init__(self, x, y, xp, yp):
self.x = np.array(x, dtype=float)
self.y = np.array(y, dtype=float)
self.xp = np.array(xp, dtype=float)
self.yp = np.array(yp, dtype=float)
def make_beam(n_particles):
np.random.seed(42)
return B(np.random.random(n_particles),
np.random.random(n_particles),
np.random.random(n_particles),
np.random.random(n_particles))
quadrupole = multip.ThinQuadrupole(np.pi)
beam = make_beam(int(1e6))
xp0, yp0 = beam.xp.copy(), beam.yp.copy()
quadrupole.track(beam)
res_quad_xp = beam.xp - xp0
res_quad_yp = beam.yp - yp0
multipole = multip.ThinMultipole([0, np.pi])
beam = make_beam(int(1e6))
xp0, yp0 = beam.xp.copy(), beam.yp.copy()
multipole.track(beam)
print ("xp kicks correct:", np.allclose(beam.xp - xp0, res_quad_xp))
print ("yp kicks correct:", np.allclose(beam.yp - yp0, res_quad_yp))
xp kicks correct: True yp kicks correct: True
%timeit make_beam(int(1e6))
10 loops, best of 3: 77.2 ms per loop
%timeit quadrupole.track(make_beam(int(1e6)))
10 loops, best of 3: 82.6 ms per loop
%timeit multipole.track(make_beam(int(1e6)))
10 loops, best of 3: 97 ms per loop
# using ztaylor
proper_ctaylor = multipole.ctaylor
multipole.ctaylor = multipole.ztaylor
%timeit multipole.track(make_beam(int(1e6)))
multipole.ctaylor = proper_ctaylor
10 loops, best of 3: 134 ms per loop
squadrupole = multip.ThinSkewQuadrupole(np.pi)
beam = make_beam(int(1e6))
xp0, yp0 = beam.xp.copy(), beam.yp.copy()
squadrupole.track(beam)
res_squad_xp = beam.xp - xp0
res_squad_yp = beam.yp - yp0
multipole = multip.ThinMultipole([0, 0], [0, np.pi])
beam = make_beam(int(1e6))
xp0, yp0 = beam.xp.copy(), beam.yp.copy()
multipole.track(beam)
print ("xp kicks correct:", np.allclose(beam.xp - xp0, res_squad_xp))
print ("yp kicks correct:", np.allclose(beam.yp - yp0, res_squad_yp))
xp kicks correct: True yp kicks correct: True
sextupole = multip.ThinSextupole(np.pi)
beam = make_beam(int(1e6))
xp0, yp0 = beam.xp.copy(), beam.yp.copy()
sextupole.track(beam)
res_sext_xp = beam.xp - xp0
res_sext_yp = beam.yp - yp0
multipole = multip.ThinMultipole([0, 0, np.pi])
beam = make_beam(int(1e6))
xp0, yp0 = beam.xp.copy(), beam.yp.copy()
multipole.track(beam)
print ("xp kicks correct:", np.allclose(beam.xp - xp0, res_sext_xp))
print ("yp kicks correct:", np.allclose(beam.yp - yp0, res_sext_yp))
xp kicks correct: True yp kicks correct: True
octupole = multip.ThinOctupole(np.pi)
beam = make_beam(int(1e6))
xp0, yp0 = beam.xp.copy(), beam.yp.copy()
octupole.track(beam)
res_octu_xp = beam.xp - xp0
res_octu_yp = beam.yp - yp0
multipole = multip.ThinMultipole([0, 0, 0, np.pi])
beam = make_beam(int(1e6))
xp0, yp0 = beam.xp.copy(), beam.yp.copy()
multipole.track(beam)
print ("xp kicks correct:", np.allclose(beam.xp - xp0, res_octu_xp))
print ("yp kicks correct:", np.allclose(beam.yp - yp0, res_octu_yp))
xp kicks correct: True yp kicks correct: True