def readpickle():
files = sorted(os.listdir(newdir))
files[0], files[1], files[2], files[3] = files[2], files[3], files[0], files[1]
phinom = []
phidis = []
nomdata = []
disdata = []
nom = {}
dis = {}
for file in files:
if file.endswith('nom.pkl'):
e_solvernom = trajectory_solver.from_pickle(newdir+'/'+file)
phinom.append(e_solvernom.init_conds.phi0)
nomdata.append(e_solvernom.dataframe)
#nom.update({str(e_solvernom.init_conds.phi0), e_solvernom.dataframe})
if file.endswith('dis.pkl'):
e_solverdis = trajectory_solver.from_pickle(newdir+'/'+file)
phidis.append(e_solverdis.init_conds.phi0)
disdata.append(e_solverdis.dataframe)
#dis.update({str(e_solverdis.init_conds.phi0), e_solverdis.dataframe})
if phinom == phidis:
for i in range(0, len(phinom), 1):
nom.update({phinom[i]:nomdata[i]})
dis.update({phidis[i]:disdata[i]})
if phinom == phidis:
return nom, dis, phinom, len(phinom)
def analyze_particle_momentum(particle_num, Bfunc, name, outdir):
# load track (pickle)
fname = outdir + name + f'.{particle_num:03d}.pkl.nom.pkl'
e = trajectory_solver.from_pickle(fname)
# analyze
e.B_func = Bfunc
e.mom_LHelix = None
e.analyze_trajectory_LHelix(step=step, stride=1)
return e.mom_LHelix
def closest_residual(e_file):
e = trajectory_solver.from_pickle(EMTRACK_RUN_DIR + e_file)
df_e = e.dataframe.copy()[::ESTRIDE]
df_e.reset_index(drop=True, inplace=True)
for res in ['B_res', 'Br_res', 'Bphi_res', 'Bz_res']:
df_e.loc[:, res] = griddata(df[['X', 'Y', 'Z']].values,
df[res],
df_e[["x", "y", "z"]].values,
method='nearest')
return df_e.values
def analyze_particle_momentum(filename, datadir=datadir):
# load track (pickle)
# fname = datadir+f'{scale}_{theta}_{phi}_0.054_.pkl'
e = trajectory_solver.from_pickle(datadir+filename)
# scale from filename
scale = float(filename[:4])
# get full theta, phi from B
theta0 = e.init_conds.theta0
phi0 = e.init_conds.phi0
# analyze in nominal only
# nominal
e.B_func = B_Mu2e_nom
e.analyze_trajectory_LHelix(step=50, stride=1)
mom_nom = e.mom_LHelix
status = e.LHelix_success
return scale, theta0, phi0, mom_nom, status
def analyze_particle_momentum(particle_num, name, outdir):
# load track (pickle)
fname = outdir + name + f'.{particle_num:03d}.pkl.nom.pkl'
e_Mu2e = trajectory_solver.from_pickle(fname)
# analyze
# nominal
e_Mu2e.B_func = B_Mu2e_nom
e_Mu2e.analyze_trajectory_LHelix(step=100, stride=1)
mom_nom = e_Mu2e.mom_LHelix
# distorted
mom_dis_list = []
for B_Mu2e_dis in B_Mu2e_dis_list:
e_Mu2e.B_func = B_Mu2e_dis
e_Mu2e.analyze_trajectory_LHelix(step=100, stride=1)
mom_dis_list.append(e_Mu2e.mom_LHelix)
return mom_nom, mom_dis_list
def track_residual(
particle_num,
Rrange=None,
name='run_04',
outdir='/home/ckampa/data/pickles/distortions/linear_gradient/run_04/'
):
# tracker radius range is about 40cm < r < 70cm
query = 'z >= 8.41 & z <= 11.66' # tracker query
fname = outdir + name + f'.{particle_num:03d}.pkl.nom.pkl'
e = trajectory_solver.from_pickle(fname)
if not Rrange is None:
query += f'& r >= {Rrange[0]} & r <= {Rrange[1]}'
e.dataframe.eval('r = (x**2+y**2)**(1/2)', inplace=True)
xyz_track = e.dataframe.query(query)[['x', 'y', 'z']].values[::ESTRIDE]
xyz_res_track = np.array([xyz_res_func(xyz) for xyz in xyz_track])
mag_res_track = np.array([mag_res_func(xyz)
for xyz in xyz_track])[:, 0].reshape(-1, 1)
# return np.concatenate([xyz_res_track, mag_res_track], axis=1)
return np.concatenate([xyz_res_track, mag_res_track], axis=1)
df_e[["x", "y", "z"]].values,
method='nearest')
return df_e.values
# this takes awhile (~20 min for 1000 tracks, 100 locations per track)...most time spent in griddata part. maybe switch to mapinterp. FIX ME!
if E_ANALYSIS:
num_cpu = multiprocessing.cpu_count()
print(f"CPU Cores: {num_cpu}")
print("Starting e track analysis")
# get run files
e_files = [
f for f in os.listdir(EMTRACK_RUN_DIR) if ('run' in f) & ('nom' in f)
]
e = trajectory_solver.from_pickle(EMTRACK_RUN_DIR + e_files[0])
cols = list(
e.dataframe.columns) + ['B_res', 'Br_res', 'Bphi_res', 'Bz_res']
df_nps = Parallel(n_jobs=num_cpu)(
delayed(closest_residual)(ef)
for ef in tqdm(e_files, file=sys.stdout, desc='trajectory #'))
df_nps = np.vstack(df_nps)
df_results = pd.DataFrame(df_nps, columns=cols)
df_results.to_pickle(mu2e_ext_path + 'Bmaps/Mau13_etracks_residuals.p')
else:
if 'Mau13_etracks_residuals.p' not in os.listdir(mu2e_ext_path + 'Bmaps'):
raise Exception(
'Please run with E_ANALYSIS = True to generate the residuals dataframe'
)