示例#1
0
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)
示例#2
0
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
示例#3
0
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
示例#5
0
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
示例#6
0
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)
示例#7
0
                                    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'
        )