! Muon Trajectory version 11

! Uses the Numerical Recipes routine odeint.f90 to track muons through a specified EM field

! To use, you must:
!      Enter the desired E and B fields in the subroutine derivs.f90
!         The kicker & quadrupole fields are calculated by functions contained within derivs.f90
!      Check that the parameters in parameters.f90 are as desired
!         This is where the size of the storage ring, strength of the dipole field, location of the inflector
!         end, and other physical dimensions and locations are stored. It also has the values of some physical
!         constants, and conversion factors for the dimensionless units I use within the program
!      Enter the number of turns, desired accuracy, estimated first step size, minimum stepsize, saving interval,
!         number of muons, minimum stepsize, initial coordinate means & standard deviations,
!         and kicker timing & strength into an input file
!             Check the read statements below to make sure you are using the right units for each

! Beware: the differential equations in derivs.f90 are written in dimensionless form. This requires:
!       position measured in units of (c*timeunit)
!       momentum measured in units of (mc)
!       electric field measured in units of (mc/q*timeunit)
!       magnetic field measured in units of (m/q*timeunit)
!    where
!           c = speed of light
!           m = mass of muon
!           q = charge of muon
!           timeunit = whatever time unit you are using (probably nanoseconds)

! Date: 8/15/2011

! changes from version 10:
!       includes scattering (but not energy loss) when muons pass through quadrupole plate



! Record of previous versions:

! version 10: tracks more than one muon

! version 9: stop and start integration at the boundary of each element of the ring

! version 8: added electrostatic quadrupole fields (simple version)

! version 7: change names of coordinates to match normal accelerator conventions
!                  x = radial, y = vertical, z = along beam path
!            changed independent coordinate from time to angle

! version 6: switched to LEPP versions of Numerical Recipes routines (double precision)

! version 5: RLC circuit time dependence for kicker field

! version 4: adjusted diff eqs to work with small deviations from ideal quantities for some variables
!            to decrease roundoff error

! version 3: changed from Cartesian coordinates to cylindrical coordinates
!            placed the origin at the center of the storage ring
!            added some more parameters

! version 2: added parameters module to hold values of constants and coordinates
!            added a function to calculate the kicker field (in the derivs subroutine) -- ideal step-function

! version 1: new program



PROGRAM muon_trajectory_11
USE nrtype
USE precision_def
USE parameters ! module with coordinates and constants
USE ode_path ! module used for data storage
USE nr, ONLY : odeint,rkqs,gasdev

IMPLICIT NONE


! Interface blocks for the derivs subroutine (other interfaces are in nr module):
INTERFACE
   SUBROUTINE derivs(x,y,dydx)
   USE nrtype
   USE precision_def
   USE parameters
   IMPLICIT NONE
   REAL(RP), INTENT(IN) :: x
   REAL(RP), DIMENSION(:), INTENT(IN) :: y
   REAL(RP), DIMENSION(:), INTENT(OUT) :: dydx
   END SUBROUTINE derivs
END INTERFACE


INTEGER(I4B) :: i,j,n,nturns,nmuons
REAL(RP) :: theta1,theta2,eps,h1,hmin,x_input,y_input,kickermagnitude_T
REAL(RP) :: xmean_m,xsigma_m,ymean_m,ysigma_m
REAL(RP) :: xmean,xsigma,ymean,ysigma,tmean,tsigma,pxmean,pxsigma,pymean,pysigma,pzmean,pzsigma
CHARACTER(len=6) :: number
CHARACTER(len=25) :: directory
REAL(RP), DIMENSION(:), ALLOCATABLE :: array,xarray,yarray,tarray,pxarray,pyarray,pzarray
REAL(RP), DIMENSION(6) :: coords
REAL(RP), DIMENSION(2) :: circlecenter
REAL(RP) :: horizontalmomentum,circleradius,Q1plateradius,thetahit
REAL(RP) :: pathlength,theta0,dev1,dev2
REAL(RP) :: yprime,py,time,theta3,px1,py1,px2,py2

! Create namelist for various input parameters

namelist/parameterinputs/nturns,eps,h1,hmin,dxsav,nmuons,xmean_m,xsigma_m,ymean_m,ysigma_m,tmean,tsigma,pxmean,pxsigma,pymean,pysigma,pzmean,pzsigma,kickermaxtime,kickermagnitude_T,directory
save_steps = .true.

! Read in the various input parameters

open (unit=1, file = 'parameterinputs.in', status ='old')
read (1, nml=parameterinputs)

! Convert inputs to dimensionless units, as needed:
xmean = xmean_m/lengthunit_m
xsigma = xsigma_m/lengthunit_m
ymean = ymean_m/lengthunit_m
ysigma = ysigma_m/lengthunit_m
kickermagnitude = kickermagnitude_T/Bfieldunit_T

! Generate gaussian distributions of input parameters
allocate(array(nmuons))
allocate(xarray(nmuons))
allocate(yarray(nmuons))
allocate(tarray(nmuons))
allocate(pxarray(nmuons))
allocate(pyarray(nmuons))
allocate(pzarray(nmuons))
call gasdev(array)
xarray = xmean + xsigma*array
call gasdev(array)
yarray = ymean + ysigma*array
call gasdev(array)
tarray = tmean + tsigma*array
call gasdev(array)
pxarray = pxmean + pxsigma*array
call gasdev(array)
pyarray = pymean + pysigma*array
call gasdev(array)
pzarray = pzmean + pzsigma*array



do time = -100, 100, 1
j=1
do n = 1,nmuons ! track each muon in turn

write(11, '(/,a1,/)')'#'

! initial coordinates for this muon:
coords = [xarray(n),yarray(n),time,pxarray(n),pyarray(n),pzarray(n)]

! create file to save trajectory data for this muon:
write(number,"(I6)") n  ! change n into a character (to be part of the file name)
!write(number,"(I6)") time
open(5,file=trim(adjustl(directory))//'/muon'//trim(adjustl(number)),status='unknown',form='formatted')


! calculate location in ring where muon hits the Q1 quadrupole plate:
circlecenter = [coords(4)/dipoleBfield(2), &
                    storageradius + coords(1) - (coords(6)+magicmomentum)/dipoleBfield(2)]
     ! these are coordinates of the center of circular trajectory
horizontalmomentum = sqrt(coords(4)**2 + (coords(6)+magicmomentum)**2)
circleradius = horizontalmomentum/dipoleBfield(2)
Q1plateradius = storageradius + 0.5*Qplatesep
thetahit = acos((Q1plateradius**2 + circlecenter(1)**2 + circlecenter(2)**2 - circleradius**2)&
               /(2*Q1plateradius*sqrt(circlecenter(1)**2 + circlecenter(2)**2))) &
        + atan(circlecenter(1)/circlecenter(2))


! track muon from inflector end to where it hits the plate:
theta1 = 0
theta2 = thetahit
quad = .false.
kick = .false.
call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
print '(8f12.4)',theta1, theta2,coords
write (5,'(3F12.7, F12.2, 3F12.7)') &
   (xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)


! apply a scattering angle:
pathlength = -Q1platethickness*(coords(6)+magicmomentum)/coords(4)
               ! pathlength = distance traveled through plate (small angle approximation)
theta0 = 0 !((13.6/momentumunit_MeVperc)/magicmomentum)*sqrt(pathlength/radlength)*(1+0.038*log(pathlength/radlength))
call gasdev(dev1)
call gasdev(dev2)
coords(4) = coords(4) + (coords(6)+magicmomentum)*theta0*dev1
coords(5) = coords(5) + (coords(6)+magicmomentum)*theta0*dev2
! approximation: leave tangential momentum unchanged

!track muon to end of quadrupole region:
theta1 = thetahit
theta2 = Q1end
quad = .true.
kick = .false.
call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
print '(8f12.4)',theta1, theta2,coords
write (6,'(3F12.7, F12.2, 3F12.7)') &
     (xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)


! track muon around ring the designated number of times
do i = 0,nturns-1
 
   theta1 = Q1end
   theta2 = kickerstart
   quad = .false.
   kick = .false.
   call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
   print '(8f12.4)',theta1, theta2,coords
  ! if (i >= 10) then
  ! write(11, '(11f12.4)') i+xp(j)/twopi, time, theta1, theta2, theta3, coords
  ! end if
   write (6,'(3F12.7, F12.2, 3F12.7)') &
      (i+xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)

   theta1 = kickerstart
   call transform(theta1, coords(4), coords(6)+magicmomentum, px1, py1)
   theta2 = kickerend
   quad = .false.
   kick = .true.
   call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
   call transform(theta2, coords(4), coords(6)+magicmomentum, px2, py2)
   theta3 = acos((px1*px2 + py1*py2)/(sqrt((px1*px1+py1*py1))*sqrt((px2*px2+py2*py2))))
   print '(8f12.4)',theta1, theta2,coords
   !if (i <= 15) then
      write(11, '(11f12.4)') i+xp(j)/twopi, time, theta1, theta2, theta3, coords
      write(12, '(2f12.4)') i, theta3
   !end if
   write (6,'(3F12.7, F12.2, 3F12.7)') &
      (i+xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)

   theta1 = kickerend
   theta2 = Q2start
   quad = .false.
   kick = .false.
   call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
   print '(8f12.4)',theta1, theta2,coords
   !if (i <= 15) then
      write(11, '(11f12.4)') i+xp(j)/twopi, time, theta1, theta2, theta3, coords
   !end if
   write (6,'(3F12.7, F12.2, 3F12.7)') &
      (i+xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)

   theta1 = Q2start
   theta2 = Q2end
   quad = .true.
   kick = .false.
   call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
   print '(8f12.4)',theta1, theta2,coords
   !if (i <= 15) then
      write(11, '(11f12.4)') i+xp(j)/twopi, time, theta1, theta2, theta3, coords
   !end if
   write (6,'(3F12.7, F12.2, 3F12.7)') &
      (i+xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)

   theta1 = Q2end
   theta2 = Q3start
   quad = .false.
   kick = .false.
   call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
   print '(8f12.4)',theta1, theta2,coords
   !if (i <= 15) then
      write(11, '(11f12.4)') i+xp(j)/twopi, time, theta1, theta2, theta3, coords
   !end if
   write (6,'(3F12.7, F12.2, 3F12.7)') &
      (i+xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)

   theta1 = Q3start
   theta2 = Q3end
   quad = .true.
   kick = .false.
   call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
   print '(8f12.4)',theta1, theta2,coords
   !if (i <= 15) then
      write(11, '(11f12.4)') i+xp(j)/twopi, time, theta1, theta2, theta3, coords
   !end if
   write (6,'(3F12.7, F12.2, 3F12.7)') &
      (i+xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)

   theta1 = Q3end
   theta2 = Q4start
   quad = .false.
   kick = .false.
   call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
   print '(8f12.4)',theta1, theta2,coords
   !if (i <= 15) then
      write(11, '(11f12.4)') i+xp(j)/twopi, time, theta1, theta2, theta3, coords
   !end if
   write (6,'(3F12.7, F12.2, 3F12.7)') &
      (i+xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)

   theta1 = Q4start
   theta2 = Q4end
   quad = .true.
   kick = .false.
   call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
   print '(8f12.4)',theta1, theta2,coords
   !if (i <= 15) then
      write(11, '(11f12.4)') i+xp(j)/twopi, time, theta1, theta2, theta3, coords
   !end if
   write (6,'(3F12.7, F12.2, 3F12.7)') &
      (i+xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)

   theta1 = Q4end
   theta2 = twopi + Q1start
   quad = .false.
   kick = .false.
   call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
   print '(8f12.4)',theta1, theta2,coords
   !if (i <= 15) then
      write(11, '(11f12.4)') i+xp(j)/twopi, time, theta1, theta2, theta3, coords
   !end if
   write (6,'(3F12.7, F12.2, 3F12.7)') &
      (i+xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)

   theta1 = twopi + Q1start
   theta2 = twopi + Q1end
   quad = .true.
   kick = .false.
   call odeint(coords,theta1,theta2,eps,h1,hmin,derivs,rkqs)
   print '(8f12.4)',theta1, theta2, coords
   !if (i <= 15) then
      write(11, '(11f12.4)') i+xp(j)/twopi, time, theta1, theta2, theta3, coords
   !end if
   write (6,'(3F12.7, F12.2, 3F12.7)') &
      (i+xp(j)/twopi, yp(1:1,j)*lengthunit_m, yp(2:2,j)*lengthunit_m, yp(3:3,j), yp(4:6,j),j=1,kount)


end do ! loop over turns

!close(5) ! close the output file for that muon

end do ! loop over muons

end do ! loop over angles

! output file has: turn #, x, y, t, px, py, pzdev


END PROGRAM muon_trajectory_11