import numpy as np
import matplotlib.pyplot as pl
from cpymad.madx import Madx
import pysixtracklib as pyst


mad=Madx()
mad.options.echo=False
mad.call(file="SPS_Q20_thin.seq")
mad.use(sequence='sps')
twiss=mad.twiss()
q1mad=twiss.summary['q1']
q2mad=twiss.summary['q2']

elements=pyst.Elements.from_mad(mad.sequence.sps)
nturns=2**14
elements.BeamMonitor(num_stores=nturns)

npart=10
particles=pyst.Particles.from_ref(npart,p0c=26e6)
particles.x += np.linspace(0,1e-6,npart)
job = pyst.TrackJob(elements, particles,until_turn_elem_by_elem=1)
job.track_elem_by_elem(1)
pl.plot(job.output.particles[0].x[1::10])

npart=10
particles=pyst.Particles.from_ref(npart,p0c=26e6)
particles.x += np.linspace(0,1e-6,npart)
job = pyst.TrackJob(elements, particles)
job.track(nturns)

ff=np.linspace(0,0.5,nturns//2+1)
x=job.output.particles[0].x[1::10]
xf=abs(np.fft.rfft(x))
plot(ff,xf)
q1st=ff[xf.argmax()]
print((q1mad-20)-q1st)


@@ -103,16 +103,31 @@
  pyst%0A%0A%0A
+# prepare madx%0A
 mad=Madx
@@ -281,16 +281,49 @@
 %5B'q2'%5D%0A%0A
+# Build elements for SixTrackLib%0A
 elements
@@ -377,16 +377,20 @@
 s=2**14%0A
+bpm=
 elements
@@ -422,16 +422,33 @@
 turns)%0A%0A
+# Track one turn%0A
 npart=10
@@ -675,16 +675,39 @@
 ::10%5D)%0A%0A
+# Track many turns CPU%0A
 npart=10
@@ -858,16 +858,29 @@
 turns)%0A%0A
+# Find tunes%0A
 ff=np.li
@@ -936,18 +936,21 @@
 0%5D.x%5B1::
-10
+npart
 %5D%0Axf=abs
@@ -1023,8 +1023,508 @@
 -q1st)%0A%0A
+# Track many particles few turns GPU%0Anpart=5000; nturns=256%0Abpm.num_stores=nturns%0Aparticles=pyst.Particles.from_ref(npart,p0c=26e6)%0Aparticles.x += np.linspace(0,1e-1,npart)%0Ajob = pyst.TrackJob(elements, particles,device=%22opencl:0.0%22)%0Ajob.track(nturns)%0Ajob.collect()%0A%0Ax=job.output.particles%5B0%5D.x.copy()%0Ax=x.reshape(nturns,npart)%0Apl.imshow(abs(np.fft.rfft(x,axis=0)),aspect='auto')%0A%0Apx=job.output.particles%5B0%5D.px.copy()%0Apx=px.reshape(nturns,npart)%0A%0Afor i in range(npart):%0A    plot(x%5B:,i%5D,px%5B:,i%5D,'.')%0A%0A

update test