def get_residuals_interp_func(mapfile): df_fit = pd.read_pickle(mapfile) # calculate |B| and all residuals df_fit.eval('B = (Bx**2+By**2+Bz**2)**(1/2)', inplace=True) df_fit.eval('B_fit = (Bx_fit**2+By_fit**2+Bz_fit**2)**(1/2)', inplace=True) df_fit.eval('B_res = B - B_fit', inplace=True) df_fit.eval('Bx_res = Bx - Bx_fit', inplace=True) df_fit.eval('By_res = By - By_fit', inplace=True) df_fit.eval('Bz_res = Bz - Bz_fit', inplace=True) # setup interp functions xyz_res_func = get_df_interp_func(df=df_fit, Blabels=['Bx_res', 'By_res', 'Bz_res']) mag_res_func = get_df_interp_func(df=df_fit, Blabels=['B_res', 'By_res', 'Bz_res']) return xyz_res_func, mag_res_func
def map_run(ic_df, B_file=fBnom, name='Mau13'): # load Bmap B_ = get_df_interp_func(filename=B_file, gauss=False) print("Start of emtracks Sim in Pion Degrader (tan(dip))") print("-----------------------------------------") # print(f"Directory: {outdir},\nFilename: {name},\nNumber: {N}") print(f"Using BField File: {B_file}") start = time.time() num_cpu = multiprocessing.cpu_count() print(f"CPU Cores: {num_cpu}") reco_tuples = Parallel(n_jobs=num_cpu)( delayed(run_trajectory)(row, B_) for i, row in tqdm(enumerate(ic_df.itertuples()), file=sys.stdout, desc='particle #', total=len(ic_df))) tand0s = np.array([i[0] for i in reco_tuples]) tands = np.array([i[1] for i in reco_tuples]) Rmaxs = np.array([i[2] for i in reco_tuples]) # store results in IC dataframe ic_df.loc[:, f"tand0_{name}"] = tand0s ic_df.loc[:, f"tand_{name}"] = tands ic_df.loc[:, f"Rmax_{name}"] = Rmaxs stop = time.time() dt = stop - start N = len(ic_df) print("Calculations Complete") print(f"Runtime: {dt} s, {dt/60.} min, {dt/60./60.} hr") print(f"Speed: {dt / (2*N)} s / trajectory\n") return ic_df
from emtracks.particle import trajectory_solver from emtracks.tools import InitConds, Bounds, ic_Mu2e, bounds_Mu2e from emtracks.mapinterp import get_df_interp_func from emtracks.Bdist import get_B_df_distorted ddir = "/home/ckampa/data/" forig = ddir + "root/Mu2e_ConversionElectron_MC.pkl" fBnom = ddir + "Bmaps/Mu2e_DSMap_V13.p" df_Mu2e_nom = pd.read_pickle(fBnom) df_Mu2e_dis = get_B_df_distorted(df_Mu2e_nom, v="0") # df_Mu2e_dis = get_B_df_distorted(df_Mu2e_nom, v="0", Bz0=1000.) # 10x gradient # B functions B_Mu2e_nom = get_df_interp_func(df=df_Mu2e_nom, gauss=False) #, bounds=bounds_Mu2e) B_Mu2e_dis = get_df_interp_func(df=df_Mu2e_dis, gauss=False) #, bounds=bounds_Mu2e) del (df_Mu2e_nom) del (df_Mu2e_dis) def load_origins(filename=forig, N=10): df = pd.read_pickle(filename) df = df.sample(N) df.reset_index(inplace=True) return df def create_init_conds(origin_tuple):
from joblib import Parallel, delayed import multiprocessing from emtracks.particle import trajectory_solver from emtracks.mapinterp import get_df_interp_func # generated momentum p0 = 104.96 # MeV/c # directories datadir = '/home/shared_data/mao10,mao13_analysis/data/mao13no_nom/' # set up B interpolator ddir = '/home/shared_data/' fBnom = ddir+"Bmaps/Mu2e_DSMap_V13.p" B_Mu2e_nom = get_df_interp_func(fBnom, gauss=False) # analyze a given track w nominal and distorted field 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)
scales_finer = np.linspace(.995, 1.005, 11) scales = np.concatenate([scales_coarse, scales_fine, scales_finer]) # course fields + new fields scales_str = [ f'{scale:.2f}' if int(round(scale * 1000 % 10)) == 0 else f'{scale:.3f}' for scale in scales ] #### Solenoid off # fields = [f'{n:0.2f}' for n in scales] # which fields to analyze fields = scales_str TSs = [ '1.00' if int(round(scale * 1000 % 10)) == 0 else '1.000' for scale in scales ] B_Mu2e_nom = get_df_interp_func(fBnom, gauss=False) B_Mu2e_dis_list = [ get_df_interp_func(fBdis.format(field, TS), gauss=False) for field, TS in zip(fields, TSs) ] #### #### Linear Gradient # df_Mu2e_nom = pd.read_pickle(fBnom) # df_Mu2e_dis = get_B_df_distorted(df_Mu2e_nom, v="0") # # B functions # B_Mu2e_nom = get_df_interp_func(df=df_Mu2e_nom, gauss=False)#, bounds=bounds_Mu2e) # B_Mu2e_dis = get_df_interp_func(df=df_Mu2e_dis, gauss=False)#, bounds=bounds_Mu2e) # del(df_Mu2e_nom)
ddir = '/home/ckampa/data/' # ddir = '/home/shared_data/' # Test of Brian Fit # fBnom = ddir+"Bmaps/Mu2e_DSMap_V13.p" # fBdis = ddir+"Bmaps/Mau13_standard_tracker_fit_df.p" # Test of Hall Probe biased up fBnom = ddir + "Bmaps/Mau13_standard_tracker_fit_df.p" # fBdis = ddir+"Bmaps/Mau13_middle_hp_bias_up_fit_df.p" nums = [0, 1, 2, 3, 4] fBdis_list = [ ddir + f"Bmaps/hp_bias/Mau13_hp_{num}_bias_up_fit_df.p" for num in nums ] # get interp funcs B_Mu2e_nom = get_df_interp_func(fBnom, gauss=False, Blabels=['Bx_fit', 'By_fit', 'Bz_fit']) # B_Mu2e_dis = get_df_interp_func(fBdis, gauss=False, Blabels=['Bx_fit', 'By_fit', 'Bz_fit']) B_Mu2e_dis_list = [ get_df_interp_func(fBdis, gauss=False, Blabels=['Bx_fit', 'By_fit', 'Bz_fit']) for fBdis in fBdis_list ] step = 100 # 100 default # analyze a given track w nominal and distorted field def analyze_particle_momentum(particle_num, name, outdir): # load track (pickle)
end_point = 14 initial_B = 50 #(rougly 1% distortion at z = 3.0, 0% at z = 14) final_B = 0 m = (final_B - initial_B) / (end_point - start_point) n = 50 step = (end_point - start_point) / n t = np.arange(start_point, end_point, step) x = plt.plot(t, ((t - start_point) * m) + initial_B) plt.title("Distortion") plt.xlabel("Z (meters)") plt.ylabel("B (gauss)") #MU2E FIELD df_Mu2e = pd.read_pickle(mapdir + "Mu2e_DSMap_V13.p") B_Mu2e_func = get_df_interp_func(mapdir + "Mu2e_DSMap_V13.p", gauss=False) #MU2E FIELD + DIS df_Mu2e_dis = get_B_df_distorted(df_Mu2e, v="0", Bz0=initial_B, Bzf=0, z0=start_point, zf=end_point) B_Mu2e_dis = get_df_interp_func(df=df_Mu2e_dis, gauss=False) #input N, return N random values between 0 and 2pi def get_random_phi(N): phis = np.random.uniform(0, 2 * math.pi, N) return phis
from emtracks.mapinterp import get_df_interp_func from emtracks.Bdist import get_B_df_distorted ddir = "/home/ckampa/data/" forig = ddir + "root/Mu2e_ConversionElectron_MC.pkl" # fBnom = ddir+"Bmaps/Mu2e_DSMap_V13.p" # fBnom = ddir+"Bmaps/Mau13_middle_hp_bias_up_fit_df.p" # bias middle Hall probe # BAD fBnom = ddir + "Bmaps/Mau13_standard_fit_df.p" # bias middle Hall probe # df_Mu2e_nom = pd.read_pickle(fBnom) # df_Mu2e_dis = get_B_df_distorted(df_Mu2e_nom, v="0") # # df_Mu2e_dis = get_B_df_distorted(df_Mu2e_nom, v="0", Bz0=1000.) # 10x gradient # B functions B_Mu2e_nom = get_df_interp_func(filename=fBnom, gauss=False, Blabels=['Bx_fit', 'By_fit', 'Bz_fit']) #, bounds=bounds_Mu2e) # B_Mu2e_nom = get_df_interp_func(df=df_Mu2e_nom, gauss=False)#, bounds=bounds_Mu2e) # B_Mu2e_dis = get_df_interp_func(df=df_Mu2e_dis, gauss=False)#, bounds=bounds_Mu2e) # del(df_Mu2e_nom) # del(df_Mu2e_dis) def load_origins(filename=forig, N=10): df = pd.read_pickle(filename) df = df.sample(N) df.reset_index(inplace=True) return df
args = parser.parse_args() # fill defaults where needed if args.map_name is None: args.map_name = mu2e_ext_path + 'Bmaps/Mu2e_DSMap_V13.p' else: args.map_name = mu2e_ext_path + 'Bmaps/' + args.map_name if args.pkl_dir is None: args.pkl_dir = mu2e_ext_path + 'pickles/Bfit_CE_reco/default/' else: args.pkl_dir = mu2e_ext_path + 'pickles/Bfit_CE_reco/' + args.pkl_dir + '/' # short Bfield name index = [i.start() for i in re.finditer('/', args.map_name)][-1] Bshort = args.map_name[index + 1:-9] outname = args.pkl_dir + Bshort + '_deltaP.p' # check pickle output directory and create if doesn't exist check_dir(args.pkl_dir) # get df func print(f'Using Bmap File: {args.map_name}') Bfunc = get_df_interp_func(filename=args.map_name, gauss=False, Blabels=['Bx_fit', 'By_fit', 'Bz_fit']) # run analysis # N_lim = 64 # testing N_lim = None run_analysis( Bfunc, outname, name="run_04", outdir="/home/ckampa/data/pickles/distortions/linear_gradient/run_04/", N_lim=N_lim)
# load and prep tracker fit dataframe # df_fit = pd.read_pickle('/home/ckampa/data/Bmaps/Mau13_standard_tracker_fit_df.p') # df_fit = pd.read_pickle('/home/ckampa/data/Bmaps/Mau13_middle_hp_bias_up_fit_df.p') # nums = [0, 1, 2, 3, 4] num = 4 df_fit = pd.read_pickle( f'/home/ckampa/data/Bmaps/hp_bias/Mau13_hp_{num}_bias_up_fit_df.p') df_fit.eval('B = (Bx**2+By**2+Bz**2)**(1/2)', inplace=True) df_fit.eval('B_fit = (Bx_fit**2+By_fit**2+Bz_fit**2)**(1/2)', inplace=True) df_fit.eval('B_res = B - B_fit', inplace=True) df_fit.eval('Bx_res = Bx - Bx_fit', inplace=True) df_fit.eval('By_res = By - By_fit', inplace=True) df_fit.eval('Bz_res = Bz - Bz_fit', inplace=True) # setup interp functions xyz_res_func = get_df_interp_func(df=df_fit, Blabels=['Bx_res', 'By_res', 'Bz_res']) mag_res_func = get_df_interp_func(df=df_fit, Blabels=['B_res', 'By_res', 'Bz_res']) # RUN = "run_04" # E_ANALYSIS = False # EMTRACK_RUN_DIR = f"/home/ckampa/data/pickles/distortions/linear_gradient/{RUN}/" # datadir = f"/home/ckampa/data/pickles/distortions/linear_gradient/" # ESTRIDE = 10 # testing (3 cm) ESTRIDE = 1 # real (3 mm) def track_residual( particle_num, Rrange=None, name='run_04',