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