# An R Exercise

# Let's make a phase space distribution and explore its
# properties upon injection into the g-2 Storage Ring
# Use a uniform distribution in x and momentum, and a Gaussian
# distribution in y

rm(list=ls())  # clear any variables left over from last time...
Npar = 5000
x0  = runif(Npar, -9, 9)  # alfa_x x + beta_x x'
Px0 = runif(Npar,-50,50)
y0  = rnorm(Npar,0,15)
Py0 = rnorm(Npar,0,15)  # alfa_y y + beta_y y'

# offset the x coordinate by 77 mm (injection)
x0 = x0 + 77

plot(x0,Px0,asp=1,pch=".")

plot(y0,Py0,asp=1,pch=".")

# Give each particle a momentum (dp/p)
del = runif(Npar,-0.006,0.006)
hist(del,breaks=25)

# Rotate each particle horizontally about ITS closed orbit, to the kicker:
D = 8100   # dispersion [mm]
x  = -Px0 + D*del
Px = -(x0 - D*del)
y  =   y0
Py =  Py0
# Q:  why this is the right thing to do?

plot(x,Px,xlim=c(-60,60),asp=1,pch=".")

# Kick the distribution onto the central orbit
Ak    = 77
phik  = 0
Px = Px + Ak*cos(phik)
x  = x  + Ak*sin(phik)

plot(x,Px,xlim=c(-60,60),asp=1,pch=".")

# Q: why does it have this shape?

# Make a data frame of our final particle data
df = data.frame(x,Px,y,Py,del)

# Look at the first few lines of the data frame...
head(df)

##            x         Px           y         Py           del
## 1 -10.481481 -16.218133 -11.9642981 -19.362421 -0.0009076567
## 2  18.228174  -1.330378   0.1119496   9.174353  0.0009274216
## 3   4.230873 -22.741721 -10.8386293  -3.445242 -0.0025593979
## 4 -23.273731  -1.243135  -1.9065723  -8.571396  0.0002269044
## 5  15.724556  49.947384  14.2516737   0.273611  0.0055267592
## 6  -6.509867  31.055719  -4.8381024   4.462012  0.0041444486

# Let's add some color to the momentum
plot(y0,Py0,asp=1,pch=".",col=2+500*(del+0.006))

plot(x0,Px0,asp=1,pch=".",col=2+500*(del+0.006))

df$color = 2+500*(del+0.006)
plot(df$x,df$Px,xlim=c(-60,60),asp=1,pch=".",col=df$color)

# OK, now we can sort/subset this data

# Example:  which particles in our distribution are currently
#    inside a circular aperture of radius 45 mm?
rAp = 45
df2 = subset(df, subset= x^2 + y^2 < rAp^2 )

plot(df2$x,df2$y,xlim=c(-60,60),asp=1,pch=".",col=df2$color)

plot(df2$x,df2$Px,xlim=c(-60,60),asp=1,pch=".",col=df2$color)

plot(df2$y,df2$Py,xlim=c(-60,60),asp=1,pch=".",col=df2$color)

# number of particles in the subset:
length(df2$x)

## [1] 3366

# Example:  Suppose have a square aperture; which particles 
#   stay inside the aperture?
df3 = subset(df, subset = 
        D*abs(del) + sqrt((x-D*del)^2 + Px^2 ) < rAp &
          sqrt(y^2 + Py^2) <rAp )

plot(df3$x,df3$y,xlim=c(-60,60),asp=1,pch=".",col=df3$color)

hist(df3$del)

# Q:  why is this the correct subset?

# Q:  create df4, for the case of a round aperture

# Q:  compare the number of particles that survive each df above

# Q:  compare the rms momentum spread for each f above
#      also, make histograms of each momentum distribution

# Q:  change the kicker amplitude (Ak) and phase (phik) and 
#      look at the resulting momentum distribution for the 
#      case of a circular aperture