# 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)

# Let's add some color to the momentum
df$color = 2+500*(del+0.006)
df$color

plot(y0,Py0,asp=1,pch=".",col=2+500*(del+0.006))
plot(x0,Px0,asp=1,pch=".",col=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=df$color)
plot(df2$x,df2$Px,xlim=c(-60,60),asp=1,pch=".",col=df$color)
plot(df2$y,df2$Py,xlim=c(-60,60),asp=1,pch=".",col=df$color)

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

# 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 )
# 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