Import GPT particles

(*Commonly used constants in SI units*)
cLightN=299792458.;
mec2N=0.510998928 (10^6);
plotopts = {Frame->True, BaseStyle->{FontSize->15}, ImageSize->500};
SetOptions[{Plot, ListPlot}, plotopts];

Clear[ImprortGPTScreen];
ImprortGPTScreen[gptparticlefile_,screennumber_Integer]:=
(*Imports gpt_particles.txt style file

General format of gpt_particles.txt:
@logo B&M-General Particle Tracer
		numderivs 0
		cputime 0.000e+00
		time 0.000e+00
		x y z G Bx By Bz rxy m q nmacro rmacro ID fEx fEy fEz fBx fBy fBz
		-1.000e-03-1.000e-03 0.000e+00 1.000e+00 0.000e+00 0.000e+00 0.000e+00 1.414e-03 9.110e-31-1.602e-19 2.312e+08 0.000e+00 1.000e+00 0.000e+00 0.000e+00 2.730e+06 0.000e+00 0.000e+00 0.000e+00
		-1.000e-03-1.000e-03 0.000e+00 1.000e+00 0.000e+00 0.000e+00 0.000e+00 1.414e-03 9.110e-31-1.602e-19 2.312e+08 0.000e+00 2.000e+00 0.000e+00 0.000e+00 2.730e+06 0.000e+00 0.000e+00 0.000e+00
		...


*)
Module[{screens, screen,nextscreen,nparticles, dat, header,hcol,  pdat, xC, yC, zC, GC, BxC, ByC, BzC, qC, nmacroC,  timebmadline, opalline, time},
Check[dat = Import[gptparticlefile, "Table"];, Print["Problem with : "<>gptparticlefile];Abort[]];
screens=Position[dat, "time"][[All,1 ]];
screen = screens[[screennumber]];

nparticles = If[Length[screens]>1, nparticles  = screens[[2]]  - screens[[1]]-3, Length[dat]-screens[[1]]+1];

time=dat[[screen]][[2]];
header=dat[[screen+1]];

pdat = dat[[screen+2;;screen+2 + nparticles-1]];
Print["Found "<>ToString[nparticles]<>" particles", Length[pdat]];
Print["Time: "<>ToString[FortranForm[time]]];


Do[hcol[header[[i]]] = i, {i,1,Length[header]}];
(*G et column indices *)
(*"x","y","z","rxy","Bx","By","Bz","G","t","m","q","nmacro","rmacro","ID"*)
xC = hcol["x"];  
yC= hcol["y"];
zC= hcol["z"];
GC = hcol["G"];
BxC = hcol["Bx"];
ByC = hcol["By"];
BzC = hcol["Bz"];
qC = hcol["q"];
nmacroC = hcol["nmacro"];
(*Convert to x, y, z, cpx, cpy, cpz, t, Q*)
{#[[xC]], #[[yC]], #[[zC]], mec2N #[[GC]] #[[BxC]],  mec2N #[[GC]] #[[ByC]],  mec2N #[[GC]] #[[BzC]], time,  #[[nmacroC ]] #[[qC ]]}&/@pdat

];

Convert to Bmad format

BmadParticleFromGPTParticle[{x_, y_, z_, cpx_, cpy_, cpz_, t_, Q_}, z0_,cp0_, t0_, mc2_:mec2N]:=Module[{Δz,Δt, β, cp, energy},
(*Drift by Δz to place particle exactly at z0*)
Δz = z0-z;
cp = Sqrt[(cpx^2)+(cpy^2)+(cpz^2)];
energy = Sqrt[(cpx^2)+(cpy^2)+(cpz^2) + (mc2^2)];
Δt = Δz energy / cpz / cLightN;

(* x, px/p0, y, py/p0, -βc(t-t0), δ = p/p0 -1, charge *)
{x + Δz (cpx/cpz), (cpx/cp0), y + Δz (cpy/cpz),  (cpy/cp0), ((-cp)/energy)cLightN( (*ignore particle's time: t + *)Δt),(cp/cp0)-1, Abs[Q], 1(* Status *)}
];

BmadBunchFromGPTBunch[gptbunch_, mc2_:mec2N]:=Module[{ref,z0, cp0, t0},
ref = Mean/@Transpose[gptbunch];
z0 = ref[[3]];
cp0 = Sqrt[((ref[[4]])^2) + ((ref[[5]])^2)+((ref[[6]])^2)];
t0 = ref[[7]];

BmadParticleFromGPTParticle[#[[1;;8]], z0, cp0, t0, mc2]&/@gptbunch


];

Export for Tao

TaoBunchHeader[nparticle_Integer, totalcharge_]:=Module[{},
{ "0  ! ix_ele",
"1  ! n_bunch",
ToString[nparticle]<>"  ! n_particle",
"BEGIN_BUNCH",
"ELECTRON",
ToString[Abs[totalcharge], FortranForm], 
"0  ! z_center",
"0  ! t_center"
}
];
TaoBunchFooter = {"END_BUNCH"};

Clear[ExportTaoBunch];
ExportTaoBunch[outfile_, bmadbunch_, nparticle_:∞]:=Module[{ilist, outbunch, resample },
If[nparticle < Length[bmadbunch],
Print["Resampling to ", nparticle, " particles"]; 
(*Resample, assuming all particles represent the same charge*)
ilist = RandomSample[Range[ Length[bmadbunch]], nparticle] ;
outbunch = bmadbunch[[ilist]];
outbunch[[All,7]] = Total[bmadbunch[[All,7]]]/nparticle,
(*else*)
(*Print["No resample"];*)
outbunch = bmadbunch;


];

Export[outfile , Join[TaoBunchHeader[Length[outbunch], Total[outbunch[[All,7]]]] , outbunch, TaoBunchFooter], "Table"]
];

Test e_gun

testdir = NotebookDirectory[]<>"e_gun/";

gptparticlefile = testdir<>"gpt_particles.txt";
gptbunch = ImprortGPTScreen[gptparticlefile,-1]

ListPlot[{#[[1]], #[[2]]}&/@gptbunch, FrameLabel->{"x (m)", "y (m)", "GPT Particles"}]

bmadbunch = BmadBunchFromGPTBunch[gptbunch ]

ListPlot[{#[[1]], #[[3]]}&/@bmadbunch, FrameLabel->{"x (m)", "y (m)", "Bmad Particles"}]

(*Export for use with Tao*)
ExportTaoBunch[testdir<>"converted.beam0", bmadbunch]