title, 'Beta Track';

! Some parameters:
pc0 = 15 ;                 ! momentum [GeV/c]
brho = 10/clight*1e8*pc0 ; ! rigidity [T-m]

! The beam:
beam, particle= proton, pc= pc0 ;

value, clight, pmass, pc0, brho ;

! The magnets:
Lquad =  2.0 ;    ! length of quadrupole magnet [m]
Fquad = 25.0 ;    ! focal length [m]

QF: quadrupole, L= Lquad,  k1=  1.0/Fquad/Lquad ;
QD: quadrupole, L= Lquad,  k1= -1.0/Fquad/Lquad ;

! Drift spaces:
Lcell2 = 30.0 ;    ! half-cell length [m]

O  :   drift,  L= Lcell2 - QF->L ;   !  space between quads [m]

value, O->L ;

! The layout:
FODO   : line=( O, QD, O, QF ) ;
FODO27 : line=( 27*FODO, O, QD ) ;

! Calculate Courant-Snyder (Twiss) parameters:
use, period=FODO27 ;
twiss, betx=60, alfx=1.2, bety=20, file=FODOvals0.txt ; 

select, flag=twiss, column= name, s, betx, alfx, bety, alfy, mux, muy ;
write, TABLE=TWISS,FILE=FODOvals.txt ;

! Make a plot:
plot, haxis=s, vaxis=betx ;

! Re-plot, with some options:
plot, haxis=s, vaxis=betx,
	 hmin=0,hmax=150, vmax=200 ;

plot, haxis=s, vaxis=betx,
	 hmin=0,hmax=150, vmin=0,vmax=200,
	 colour=100, interpolate=TRUE ;

plot, haxis=s, vaxis=betx,bety,
	 hmin=0,hmax=150, vmin=0,vmax=200,
	 colour=100, interpolate=TRUE ;

!**********
!Stop ;
!**********

! Compare periodic functions to tracking arbitrary initial functions

use, period=FODO ;

twiss, betx=60, alfx=1.2, bety=20 ;
plot, haxis=s, vaxis=betx, bety,
	 hmin=0,hmax=60, vmin=0,vmax=200,
	 colour=100, interpolate=TRUE ;

value, table(summ,Q1)*360 ;  !          horiz. cell phase advance (degrees)

twiss ;      ! Note:  if no initial parameters, then uses periodic parameters
plot, haxis=s, vaxis=betx,bety,
	 hmin=0,hmax=60, vmin=0,vmax=200,
	 colour=100, interpolate=TRUE ;

value, table(summ,Q1)*360 ;  ! periodic horiz. cell phase advance (degrees)

!**********
!Stop ;
!**********

! Use MAD-X to do nonlinear optimization; here, match one FODO cell to another

! Create a longer FODO cell to be used "downstream":
QFb: quadrupole, L= Lquad,  k1=  1.0/Fquad/Lquad/1.8 ;
QDb: quadrupole, L= Lquad,  k1= -1.0/Fquad/Lquad/1.8 ;
Ob :  drift, L= (O->L)*1.5 ;

FODOb: line=( Ob, QDb, Ob, QFb ) ;

! Next, create a "matching section" with quads that we will vary in strength:
Q1: quadrupole, L= Lquad,  k1= -1.0/Fquad/Lquad/1.8 ;
Q2: quadrupole, L= Lquad,  k1=  1.0/Fquad/Lquad/1.8 ;
Q3: quadrupole, L= Lquad,  k1= -1.0/Fquad/Lquad/1.8 ;
Q4: quadrupole, L= Lquad,  k1=  1.0/Fquad/Lquad/1.8 ;
dd :  drift, L= 38 ;

mcell: line=( dd Q1 dd Q2 dd Q3 dd Q4 ) ;

! Make a total beam line that connects the old FODO cells with the new ones,
!     using the matching section in between:
totaLine: line=( 3*FODO, mcell, 3*FODOb ) ;

! Prepare to vary the 4 tuning quads:
! First, calculate the initial CS values to use:
use, period=FODO ;
savebeta, label=inbet, place=#e, sequence=FODO ;
twiss;
show, inbet ;

! Secondly, generate final desired CS values to match to:
use, period=FODOb ;
savebeta, label=outbet, place=#e, sequence=FODOb ;
twiss;
show, outbet ;

! Make a plot, to see where things start:
use,period=totaLine ;
twiss, beta0=inbet;
plot, haxis=s, vaxis=betx,bety,
	 vmin=0,vmax=400,
	 colour=100, interpolate=TRUE ;

! Perform the Match:
use,period=totaLine ;
match, sequence=totaLine, beta0=inbet ;
	constraint, sequence=totaLine, range=#e, 
	   betx=outbet->betx, alfx=outbet->alfx, bety=outbet->bety, alfy=outbet->alfy;
	vary, name=Q1->K1, step=1.0e-4  ;
	vary, name=Q2->K1, step=1.0e-4  ;
	vary, name=Q3->K1, step=1.0e-4  ;	
	vary, name=Q4->K1, step=1.0e-4  ;
	simplex, calls=1000, tolerance=1.0e-5 ;
endmatch;

! Plot the final result:
use,period=totaLine ;
twiss, beta0=inbet;
plot, haxis=s, vaxis=betx,bety,
	 vmin=0,vmax=400,
	 colour=100, interpolate=TRUE ;