def plot_theta(df_path): ''' Load data apply cuts and return two data frames - one per station ''' print("Opening data...") data_hdf = pd.read_hdf(df_path) #open skimmed print("N before cuts", data_hdf.shape[0]) #apply cuts mom_cut = ( (data_hdf['trackMomentum'] > p_min) & (data_hdf['trackMomentum'] < p_max) ) # MeV time_cut =( (data_hdf['trackT0'] > t_min) & (data_hdf['trackT0'] < t_max) ) # MeV data_hdf=data_hdf[mom_cut & time_cut] data_hdf=data_hdf.reset_index() # reset index from 0 after cuts N=data_hdf.shape[0] print("Total tracks after cuts", round(N/1e6,2), "M") # calculate variables for plotting p=data_hdf['trackMomentum'] py=data_hdf['trackMomentumY'] theta_y_mrad = np.arctan2(py, p)*1e3 # rad -> mrad data_hdf['theta_y_mrad']=theta_y_mrad # add to the data frame # select all stations for simulation if(sim or len(stations)==1): data = [data_hdf] #split into two stations for data if(not sim and len(stations)==2): data = [ data_hdf[data_hdf['station'] == 12], data_hdf[data_hdf['station'] == 18] ]; for i_station, station in enumerate(stations): data_station=data[i_station] N=data_station.shape[0] print("Entries: ", N, " in S"+str(station)) ############# #Blinded (EDM) fit for B_Z ############ ### Resolve angle and times tmod_abs, weights=cu.get_abs_times_weights(data_station['trackT0'], g2period) ang=data_station['theta_y_mrad'] ### Digitise data with weights xy_bins=(bin_n, bin_n) h,xedges,yedges = np.histogram2d(tmod_abs, ang, weights=weights, bins=xy_bins); # expand (x_w, y_w), binsXY, dBinXY = ru.hist2np(h, (xedges,yedges)) print("Got XY bins", binsXY) #profile df_binned =cu.Profile(x_w, y_w, None, nbins=bin_n, xmin=np.min(x_w), xmax=np.max(x_w), mean=True, only_binned=True) x, y, y_e, x_e =df_binned['bincenters'], df_binned['ymean'], df_binned['yerr'], df_binned['xerr'] #Fit par, par_e, pcov, chi2_ndf, ndf = cu.fit_and_chi2(x, y, y_e, cu.thetaY_phase, p0_theta_blinded[i_station]) if (np.max(abs(par_e)) == np.Infinity ): raise Exception("\nOne of the fit parameters is infinity! Exiting...\n") if(args.scan==True): par_dump=np.array([[t_min], t_max, p_min, p_max, chi2_ndf, ndf, g2period, cu._LT, cu._phi, bin_w, bin_n, xy_bins[0], xy_bins[1], N, station, ds_name, *par, *par_e]) par_dump_keys = ["start", "stop", "p_min", "p_max", "chi2", "ndf", "g2period", "lt", "phase", "bin_w", "bin_n", "ndf_x", "ndf_y", "n", "station", "ds"] par_dump_keys.extend(par_names_theta) par_dump_keys.extend( [str(par)+"_e" for par in par_names_theta] ) dict_dump = dict(zip(par_dump_keys,par_dump)) df = pd.DataFrame.from_records(dict_dump, index='start') with open("../DATA/scans/edm_scan_"+keys[1]+".csv", 'a') as f: df.to_csv(f, mode='a', header=f.tell()==0)
def plot_theta(df_path): ''' Load data apply cuts and return two data frames - one per station ''' print("Opening data...") data_hdf = pd.read_hdf(df_path) #open skimmed print("N before cuts", data_hdf.shape[0]) #apply cuts mom_cut = ((data_hdf['trackMomentum'] > p_min) & (data_hdf['trackMomentum'] < p_max)) # MeV time_cut = ( (data_hdf['trackT0'] > t_min) & (data_hdf['trackT0'] < t_max)) # MeV data_hdf = data_hdf[mom_cut & time_cut] data_hdf = data_hdf.reset_index() # reset index from 0 after cuts N = data_hdf.shape[0] print("Total tracks after cuts", round(N / 1e6, 2), "M") # calculate variables for plotting p = data_hdf['trackMomentum'] py = data_hdf['trackMomentumY'] theta_y_mrad = np.arctan2(py, p) * 1e3 # rad -> mrad data_hdf['theta_y_mrad'] = theta_y_mrad # add to the data frame if (sim): t = data_hdf['trackT0'] mod_times = cu.get_g2_mod_time( t, g2period) # Module the g-2 oscillation time data_hdf['mod_times'] = mod_times # add to the data frame # select all stations for simulation if (sim or len(stations) == 1): data = [data_hdf] #split into two stations for data if (not sim and len(stations) == 2): data = [ data_hdf[data_hdf['station'] == 12], data_hdf[data_hdf['station'] == 18] ] for i_station, station in enumerate(stations): data_station = data[i_station] N = data_station.shape[0] print("Entries: ", N, " in S" + str(station)) ############# #Blinded (EDM) fit for B_Z ############ ### Resolve angle and times tmod_abs, weights = cu.get_abs_times_weights(data_station['trackT0'], g2period) ang = data_station['theta_y_mrad'] ### Digitise data with weights xy_bins = (bin_n, bin_n) h, xedges, yedges = np.histogram2d(tmod_abs, ang, weights=weights, bins=xy_bins) # expand (x_w, y_w), binsXY, dBinXY = ru.hist2np(h, (xedges, yedges)) print("Got XY bins", binsXY) #profile df_binned = cu.Profile(x_w, y_w, None, nbins=bin_n, xmin=np.min(x_w), xmax=np.max(x_w), mean=True, only_binned=True) x, y, y_e, x_e = df_binned['bincenters'], df_binned[ 'ymean'], df_binned['yerr'], df_binned['xerr'] #Fit par, par_e, pcov, chi2_ndf, ndf = cu.fit_and_chi2( x, y, y_e, cu.thetaY_phase, p0_theta_blinded[i_station]) if (np.max(abs(par_e)) == np.Infinity): raise Exception( "\nOne of the fit parameters is infinity! Exiting...\n") if (args.corr): print("Covariance matrix\n", pcov) np.save("../DATA/misc/pcov_theta_S" + str(station) + ".np", pcov) #Plot #Set legend title for the plot if (sim): legend = ds_name_official + " S" + str(station) else: legend = "Run-" + ds_name_official + " dataset S" + str(station) fig, ax, leg_data, leg_fit = cu.plot_edm( x, y, y_e, cu.thetaY_phase, par, par_e, chi2_ndf, ndf, bin_w, N, t_min, t_max, p_min, p_max, par_labels_theta, par_units_theta, legend_data=legend, legend_fit= r'Fit: $\langle \theta(t) \rangle = A_{\mathrm{B_z}}\cos(\omega_a t + \phi) + A_{\mathrm{EDM}}\sin(\omega_a t + \phi) + c$', ylabel=r"$\langle\theta_y\rangle$ [mrad] per " + str(int(bin_w * 1e3)) + " ns", font_size=font_size, prec=2, urad=urad_bool) ax.set_xlim(0, g2period) if (ds_name == "9D"): ax.set_ylim(-0.75, 0.40) elif (ds_name == "R1"): ax.set_ylim(-0.5, 0.13) if (p_min < 1800): ax.set_ylim(-1.0, -0.1) elif (ds_name == "EG"): ax.set_ylim(-0.75, 0.15) elif (ds_name == "HK"): ax.set_ylim(-0.60, 0.40) else: ax.set_ylim(-0.80, 0.55) if (sim): ax.set_ylim(-2.9, 2.5) cu.textL(ax, 0.75, 0.15, leg_data, fs=font_size) cu.textL(ax, 0.25, 0.17, leg_fit, fs=font_size, c="r") print("Fit in " + ds_name + " S:" + str(station), leg_fit) if (args.scan == False): fig.savefig("../fig/bz_" + ds_name + "_S" + str(station) + ".png", dpi=300) if (args.scan == True): par_dump = np.array([[t_min], t_max, p_min, p_max, chi2_ndf, ndf, g2period, cu._LT, cu._phi, bin_w, bin_n, xy_bins[0], xy_bins[1], N, station, ds_name, *par, *par_e]) par_dump_keys = [ "start", "stop", "p_min", "p_max", "chi2", "ndf", "g2period", "lt", "phase", "bin_w", "bin_n", "ndf_x", "ndf_y", "n", "station", "ds" ] par_dump_keys.extend(par_names_theta) par_dump_keys.extend([str(par) + "_e" for par in par_names_theta]) dict_dump = dict(zip(par_dump_keys, par_dump)) df = pd.DataFrame.from_records(dict_dump, index='start') with open("../DATA/scans/edm_scan_" + keys[1] + ".csv", 'a') as f: df.to_csv(f, mode='a', header=f.tell() == 0) plt.savefig("../fig/scans/bz_" + ds_name + "_S" + str(station) + scan_label + ".png", dpi=300) # get residuals for later plots residuals_theta[i_station] = cu.residuals(x, y, cu.thetaY_phase, par) times_theta[i_station] = x errors_theta[i_station] = y_e #make sanity plots if (args.hist): # if(not sim and args.hist): fig, _ = plt.subplots() bin_w_mom = 10 mom = data_station['trackMomentum'] n_bins_mom = int(round((max(mom) - min(mom)) / bin_w_mom, 2)) ax, _ = cu.plotHist(mom, n_bins=n_bins_mom, prec=3, units="MeV", label="Run-" + ds_name_official + " dataset S" + str(station)) legend = cu.legend5(*cu.stats5(mom), "MeV", prec=2) cu.textL(ax, 0.76, 0.85, str(legend), fs=14) ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[1] * 1.1) # ax.set_xlim(-50,50) ax.set_xlabel(r"$p$ [MeV]", fontsize=font_size) ax.set_ylabel("Entries per " + str(bin_w_mom) + " MeV", fontsize=font_size) ax.legend(fontsize=font_size, loc='upper center', bbox_to_anchor=(0.26, 1.0)) fig.savefig("../fig/mom_" + ds_name + "_S" + str(station) + ".png", dpi=300, bbox_inches='tight') fig, _ = plt.subplots() n_bins_ang = 400 * 2 ax, _ = cu.plotHist(ang, n_bins=n_bins_ang, prec=3, units="mrad", label="Run-" + ds_name_official + " dataset S" + str(station)) legend = cu.legend3_sd(*cu.stats3_sd(ang), "mrad", prec=3) cu.textL(ax, 0.8, 0.85, str(legend), fs=14) ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[1] * 1.1) ax.set_xlim(-50, 50) ax.set_xlabel(r"$\theta_y$ [mrad]", fontsize=font_size) ax.set_ylabel("Entries per " + str(round( (max(ang) - min(ang)) / n_bins_ang, 3)) + " mrad", fontsize=font_size) ax.legend(fontsize=font_size, loc='upper center', bbox_to_anchor=(0.3, 1.0)) fig.savefig("../fig/theta_" + ds_name + "_S" + str(station) + ".png", dpi=300, bbox_inches='tight') # fig, _ = plt.subplots() # n_binsXY_ang=(192,575) # jg, cb, legendX, legendY = cu.plotHist2D(data_station['trackT0'], ang, n_binsXY=n_binsXY_ang, prec=2, unitsXY=(r"[$\rm{\mu}$s]", "mrad"), label="S"+str(station), cmin=0) # jg.ax_joint.set_xlim(0, 100) # jg.ax_joint.set_ylim(-60, 60) # jg.ax_joint.set_ylabel(r"$\theta_y$ [mrad]", fontsize=font_size+2); # jg.ax_joint.set_xlabel(r"t [$\rm{\mu}$s]", fontsize=font_size+2); # plt.savefig("../fig/theta2D_"+ds_name+"_S"+str(station)+".png", dpi=300, bbox_inches='tight') # fig, _ = plt.subplots() # jg, cb, legendX, legendY = cu.plotHist2D(data_station['mod_times'], ang, n_binsXY=n_binsXY_ang, prec=3, unitsXY=(r"[$\rm{\mu}$s]", "mrad"), label="S"+str(station), cmin=0 ) # jg.ax_joint.set_xlim(0.0, g2period) # jg.ax_joint.set_ylim(-60, 60) # jg.ax_joint.set_ylabel(r"$\theta_y$ [mrad]", fontsize=font_size+2); # jg.ax_joint.set_xlabel(r"$t^{mod}_{g-2}$"+r"[$\rm{\mu}$s]", fontsize=font_size+2); # plt.savefig("../fig/theta2D_mod_"+ds_name+"_S"+str(station)+".png", dpi=300, bbox_inches='tight') ############# # Make truth (un-blinded fits) if simulation ############# if (sim): print("Making truth plots in simulation") # Bin df_binned = cu.Profile(data_station['mod_times'], data_station['theta_y_mrad'], None, nbins=bin_n, xmin=np.min(data_station['mod_times']), xmax=np.max(data_station['mod_times']), mean=True, only_binned=True) x, y, y_e, x_e = df_binned['bincenters'], df_binned[ 'ymean'], df_binned['yerr'], df_binned['xerr'] # Fit par, par_e, pcov, chi2_ndf, ndf = cu.fit_and_chi2( x, y, y_e, cu.thetaY_phase, p0_theta_truth[i_station]) if (np.max(abs(par_e)) == np.Infinity): raise Exception( "\nOne of the fit parameters is infinity! Exiting...\n") if (args.corr): print("Covariance matrix", pcov) np.save("../DATA/misc/pcov_truth_S" + str(station) + ".np", pcov) #Plot fig, ax, leg_data, leg_fit = cu.plot_edm( x, y, y_e, cu.thetaY_phase, par, par_e, chi2_ndf, ndf, bin_w, N, t_min, t_max, p_min, p_max, par_labels_truth, par_units_theta, legend_data=legend, legend_fit= r'Fit: $\langle \theta(t) \rangle = A_{\mathrm{B_z}}\cos(\omega_a t + \phi) + A_{\mathrm{EDM}}\sin(\omega_a t + \phi) + c$', ylabel=r"$\langle\theta_y\rangle$ [mrad] per " + str(int(bin_w * 1e3)) + " ns", font_size=font_size, prec=2, urad=urad_bool) cu.textL(ax, 0.74, 0.15, leg_data, fs=font_size) cu.textL(ax, 0.23, 0.15, leg_fit, fs=font_size, c="r") ax.set_xlim(0, g2period) ax.set_ylim(-0.80, 0.55) if (sim): ax.set_ylim(-2.9, 2.5) if (args.scan == False): fig.savefig("../fig/bz_truth_fit_S" + str(station) + ".png", dpi=300) if (args.scan == True): par_dump = np.array([[t_min], t_max, p_min, p_max, chi2_ndf, ndf, g2period, bin_w, N, station, ds_name, *par, *par_e]) par_dump_keys = [ "start", "stop", "p_min", "p_max", "chi2", "ndf", "g2period", "bin_w", "n", "station", "ds" ] par_dump_keys.extend(par_names_theta_truth) par_dump_keys.extend( [str(par) + "_e" for par in par_names_theta_truth]) dict_dump = dict(zip(par_dump_keys, par_dump)) df = pd.DataFrame.from_records(dict_dump, index='start') with open("../DATA/scans/edm_scan_" + keys[2] + ".csv", 'a') as f: df.to_csv(f, mode='a', header=f.tell() == 0) plt.savefig("../fig/scans/bz_truth_fit_S" + str(station) + scan_label + ".png", dpi=300) #-------end of looping over stations ## now if not scanning - get FFTs for both stations ## FFTs if (not args.scan and not args.count and args.corr): print("Plotting residuals and FFTs...") cu.residual_plots(times_counts, residuals_counts, sim=sim, eL="count", file_label=file_label)
# Profile Plot if (args.profile): # load the data dataXY=np.load("../DATA/misc/dataXY.npy") x=dataXY[0] y=dataXY[1] #convert 𝛉 into um y=y*1e3 # rad -> mrad print("Plotting a profile...") fig,ax=plt.subplots() ax, df_binned, df_input =cu.Profile(x, y, ax, nbins=15, xmin=np.min(x),xmax=np.max(x), mean=True) ax.set_ylabel(r"$\langle\theta_y\rangle$ [mrad]", fontsize=16) ax.set_xlabel(r"$t^{mod}_{g-2} \ \mathrm{[\mu}$s]", fontsize=16) N=cu.sci_notation(len(x)) # format as a cu.textL(ax, 0.88, 0.9, "N: "+N, fs=14) plt.tight_layout() plt.savefig("../fig/profile.png") # can save the profile points with errors to a file df_binned.to_csv("../DATA/misc/df_binned.csv") # iterative fits over many profiles if(args.iter and not args.bins): print("Plotting iterative profiles...") # df = pd.DataFrame(columns=['cut', 'A_mu','A_edm','c','w','chi2']) fu.iter_plots(n_prof_bins=15, gauss=args.gauss, df=df)
def plot_theta(df_path): ''' Load data apply cuts and return two data frames - one per station ''' print("Opening data...") data_hdf = pd.read_hdf(df_path, "trackerNTup/tracker") #open skimmed print("N before cuts", data_hdf.shape[0]) #apply cuts mom_cut = ((data_hdf['trackMomentum'] > p_min) & (data_hdf['trackMomentum'] < p_max)) # MeV time_cut = ( (data_hdf['trackT0'] > t_min) & (data_hdf['trackT0'] < t_max)) # MeV data_hdf = data_hdf[mom_cut & time_cut] data_hdf = data_hdf.reset_index() # reset index from 0 after cuts N = data_hdf.shape[0] print("Total tracks after cuts", round(N / 1e6, 2), "M") # calculate variables for plotting p = data_hdf['trackMomentum'] py = data_hdf['trackMomentumY'] t = data_hdf['trackT0'] theta_y_mrad = np.arctan2(py, p) * 1e3 # rad -> mrad mod_times = cu.get_g2_mod_time(t, g2period) # Module the g-2 oscillation time ############# # Make truth (un-blinded fits) if simulation ############# print("Making truth plots in simulation") # Bin df_binned = cu.Profile(mod_times, theta_y_mrad, None, nbins=bin_n, xmin=np.min(mod_times), xmax=np.max(mod_times), mean=True, only_binned=True) x, y, y_e, x_e = df_binned['bincenters'], df_binned['ymean'], df_binned[ 'yerr'], df_binned['xerr'] # Fit par, par_e, pcov, chi2_ndf, ndf = cu.fit_and_chi2(x, y, y_e, cu.thetaY_phase, p0_theta) if (np.max(abs(par_e)) == np.Infinity): raise Exception( "\nOne of the fit parameters is infinity! Exiting...\n") #Plot legend = ds_name + " S" + str(station) fig, ax, leg_data, leg_fit = cu.plot_edm( x, y, y_e, cu.thetaY_phase, par, par_e, chi2_ndf, ndf, bin_w, N, t_min, t_max, p_min, p_max, par_labels_theta, par_units_theta, legend_data=legend, legend_fit= r'Fit: $\langle \theta(t) \rangle = A_{\mathrm{B_z}}\cos(\omega_a t + \phi) + A_{\mathrm{EDM}}\sin(\omega_a t + \phi) + c$', ylabel=r"$\langle\theta_y\rangle$ [mrad] per " + str(int(bin_w * 1e3)) + " ns", font_size=font_size, prec=2) cu.textL(ax, 0.74, 0.15, leg_data, fs=font_size) cu.textL(ax, 0.23, 0.15, leg_fit, fs=font_size, c="r") ax.set_xlim(0, g2period) ax.set_ylim(-2.9, 2.5) if (args.scan == False): fig.savefig("../fig/bz_truth_fit_S" + str(station) + ".png", dpi=300) scan_label = "_" + str(t_min) + "_" + str(t_max) + "_" + str( p_min) + "_" + str(p_max) + "_" + str(ndf) if (args.scan == True): par_dump = np.array([[t_min], t_max, p_min, p_max, chi2_ndf, ndf, g2period, bin_w, N, station, ds_name, *par, *par_e]) par_dump_keys = [ "start", "stop", "p_min", "p_max", "chi2", "ndf", "g2period", "bin_w", "n", "station", "ds" ] par_dump_keys.extend(par_names_theta) par_dump_keys.extend([str(par) + "_e" for par in par_names_theta]) dict_dump = dict(zip(par_dump_keys, par_dump)) df = pd.DataFrame.from_records(dict_dump, index='start') with open("../DATA/scans/edm_scan_" + keys[1] + ".csv", 'a') as f: df.to_csv(f, mode='a', header=f.tell() == 0) plt.savefig("../fig/scans/bz_truth_fit_S" + str(station) + scan_label + ".png", dpi=300)
modulog2TimeReco = (g2fracTimeReco - g2fracTimeRecoInt) * g2period # resolve data into easy variable names time, theta, p = modulog2TimeReco, theta_y, data["decayVertexMom"] #sanity plot the momentum cut itself ax, legend = cu.plotHist(p, n_bins=45) plt.savefig("../fig/cuts/mom_" + cutLabel[i] + ".png") plt.clf() ### Profile (bin), fit, plot # return binned dataframe only (only_binned=True) df_binned = cu.Profile(time, theta, None, nbins=15, xmin=np.min(time), xmax=np.max(time), mean=True, only_binned=True) # resole data for fit x = df_binned['bincenters'] y = df_binned['ymean'] x_err = df_binned['xerr'] y_err = df_binned['yerr'] ''' Now perfrom a 4-parameter fit in simulation with the constants phase of 6.240(8) rad $$\theta(t) = A_{B_z}\cos(\omega t + \phi) + A_{\mathrm{EDM}}\sin(\omega t + \phi) + c$$ where [0] $A_{\mathrm{B_z}}$ is the $B_z$ amplitude [1] $A_{\mathrm{EDM}}$ is the EDM amplitude
def plot_counts_theta(data): for i_station, station in enumerate(stations): data_station = data[i_station] N = data_station.shape[0] print("Entries: ", N, " in S" + str(station)) if (1 == 1): ############# #Blinded (EDM) fit for B_Z ############ ang = data_station['theta_y_mrad'] tmod_abs = data_station['mod_times'] ### Digitise data with weights xy_bins = (bin_n, bin_n) h, xedges, yedges = np.histogram2d(tmod_abs, ang, bins=xy_bins) # expand (x_w, y_w), binsXY, dBinXY = ru.hist2np(h, (xedges, yedges)) print("Got XY bins", binsXY) #profile df_binned = cu.Profile(x_w, y_w, None, nbins=bin_n, xmin=np.min(x_w), xmax=np.max(x_w), mean=True, only_binned=True) x, y, y_e, x_e = df_binned['bincenters'], df_binned[ 'ymean'], df_binned['yerr'], df_binned['xerr'] #Fit par, par_e, pcov, chi2_ndf, ndf = cu.fit_and_chi2( x, y, y_e, cu.edm_sim, p0_theta) if (np.max(abs(par_e)) == np.Infinity): raise Exception( "\nOne of the fit parameters is infinity! Exiting...\n") #Plot if (sim): legend = ds_name + " S" + str(station) fig, ax, leg_data, leg_fit = cu.plot_edm( x, y, y_e, cu.edm_sim, par, par_e, chi2_ndf, ndf, bin_w, N, t_min, t_max, p_min, p_max, par_labels_theta, par_units_theta, legend_data=legend, legend_fit= r'Fit: $\langle \theta(t) \rangle = A_{\mathrm{EDM}}\sin(\omega_a t) + c$', ylabel=r"$\langle\theta_y\rangle$ [mrad] per " + str(int(bin_w * 1e3)) + " ns", font_size=font_size, prec=2, urad=False) ax.set_xlim(0, g2period) ax.set_ylim(ax.get_ylim()[0] * 2.0, ax.get_ylim()[1] * 1.6) if (not sim): ax.set_ylim(ax.get_ylim()[0] * 1.23, ax.get_ylim()[1] * 1.4) cu.textL(ax, 0.75, 0.15, leg_data, fs=font_size) cu.textL(ax, 0.25, 0.17, leg_fit, fs=font_size, c="r") print("Fit in " + ds_name + " S:" + str(station), leg_fit) plt.tight_layout() fig.savefig("../fig/sim_" + ds_name + "_S" + str(station) + ".png", dpi=300)
def iter_plots(n_prof_bins=15, extraLabel="", outdir="profFits", gauss=False, best=False, vertex=False, df=None): #loop over common plots # fileLabel=("LargeEDM", "noEDM") # fileName=("DATA/VLEDM.root", "DATA/noEDM.root") fileLabel = [("LargeEDM")] fileName = [("../DATA/VLEDM.root")] plotLabel = ("Tracks", "Vertices") plotName = ["TrackFit", "VertexExt"] if (best): plotLabel = [("Vertices")] plotName = [("VertexExt")] # qLabel=("QT", "NQT") # qName=("AllStations", "AllStationsNoTQ")\ qLabel = [("NQT")] qName = [("AllStationsNoTQ")] cutLabel = ("0_p_3600", "400_p_2700", "700_p_2400", "1500_p_3600", "1600_p_3600", "1700_p_3600", "1800_p_3600") cutName = ("t>0/0<p<3600", "t>0/400<p<2700", "t>0/700<p<2400", "t>0/1500<p<3600", "t>0/1600<p<3600", "t>0/1700<p<3600", "t>0/1800<p<3600") if (best): cutLabel = [("700_p_2400")] cutName = [("t>0/700<p<2400")] for i_file, i_fileName in enumerate(fileName): # change of name on noEDM if (i_file == 1): plotName[1] = "VertexExtap" #old art code had a typo... for i_cut, i_cutName in enumerate(cutName): for i_plot, i_plotName in enumerate(plotName): for i_q, i_qName in enumerate(qName): # form full paths and labels fullPath = i_qName + "/" + i_plotName + "/" + i_cutName + "/thetay_vs_time_modg2" fullLabel = extraLabel + fileLabel[ i_file] + "_" + plotLabel[i_plot] + "_" + qLabel[ i_q] + "_" + cutLabel[i_cut] print("Plotting a profile for", fullPath) #extract data from the histogram dataXY, n_binsXY, dBinsXY = ru.hist2np( file_path=i_fileName, hist_path=fullPath) #bin data into a profile if (gauss): df_data = cu.Profile(dataXY[0], dataXY[1], False, nbins=n_prof_bins, xmin=np.min(dataXY[0]), xmax=np.max(dataXY[0]), full_y=True, only_binned=True) y = df_data['y'] else: df_data = cu.Profile(dataXY[0], dataXY[1], False, nbins=n_prof_bins, xmin=np.min(dataXY[0]), xmax=np.max(dataXY[0]), mean=True, only_binned=True) y = df_data['ymean'] x = df_data['bincenters'] x_err = df_data['xerr'] y_err = df_data['yerr'] y = y * 1e3 # rad -> mrad y_err = y_err * 1e3 # rad -> mrad # extract info edm_setting = fullPath.split("/")[1].split(".")[0] Q_cut = fullPath.split("/")[0] data_type = fullPath.split("/")[1] time_cut = fullPath.split("/")[2] + r" $\mathrm{\mu}$s" p_cut = fullPath.split("/")[3] + " MeV" y_label = fullPath.split("/")[4].split("_")[0] x_label = fullPath.split("/")[4].split("_")[2:] N = len(dataXY[0]) #some specific string transforms if (edm_setting == "VLEDM"): edm_setting = r"$d_{\mu} = 5.4\times10^{-18} \ e\cdot{\mathrm{cm}}$" if (edm_setting == "noEDM"): edm_setting = r"$d_{\mu} = 0 \ e\cdot{\mathrm{cm}}$" if (y_label == "thetay"): y_label = r"$\langle\theta_y\rangle$ [mrad]" if (x_label[0] == "time" and x_label[1] == "modg2"): x_label = r"$t^{mod}_{g-2} \ \mathrm{[\mu}$s]" # if extracting y and delta(y) from a Gaussian fit if (gauss): means = [] means_errors = [] for i_point in range(0, len(df_data)): # fit for a range of data n_bins = 25 y_min, y_max = -25, +25 # all data y_hist, range is y y_hist = y[i_point] y_select = y_hist[np.logical_and( y_hist >= y_min, y_hist <= y_max)] #bin the data in range hist, bin_edges = np.histogram(y_select, bins=n_bins, density=False) bin_centres = (bin_edges[:-1] + bin_edges[1:]) / 2 bin_width = bin_edges[1] - bin_edges[0] #find the right number of bins for. all data n_bins_hist = int( (np.max(y_hist) - np.min(y_hist)) / bin_width) y_err = np.sqrt(hist) # sqrt(N) per bin # fit in range p0 = [1, 1, 1] par, pcov = optimize.curve_fit( cu.gauss, bin_centres, hist, p0=p0, sigma=y_err, absolute_sigma=False, method='trf') par_e = np.sqrt(np.diag(pcov)) chi2ndf = cu.chi2_ndf(bin_centres, hist, y_err, cu.gauss, par) #append the fit parameters means.append(par[1]) means_errors.append(par_e[1]) #plot and stats + legend units = "mrad" legend_fit = cu.legend4_fit(chi2ndf[0], par[1], par_e[1], par[2], par_e[2], units, prec=2) ax, legend = cu.plotHist(y_hist, n_bins=n_bins_hist, units=units, prec=2) ax.plot(bin_centres, cu.gauss(bin_centres, *par), color="red", linewidth=2, label='Fit') cu.textL(ax, 0.8, 0.78, r"$\theta_y$:" + "\n" + str(legend), font_size=15) cu.textL(ax, 0.2, 0.78, "Fit:" + "\n" + str(legend_fit), font_size=15, color="red") ax.set_xlabel(r"$\theta_y$ [mrad]", fontsize=18) ax.set_ylabel(r"$\theta_y$ / " + str(round(bin_width)) + " mrad", fontsize=18) plt.tight_layout() plt.savefig("../fig/Gauss/Gauss_" + fullLabel + "_" + str(i_point) + ".png", dpi=300) plt.clf() #done looping over y bins #reassign data "pointer names" y = means y_err = means_errors #fit a function and get pars par, pcov = optimize.curve_fit(cu.thetaY_unblinded, x, y, sigma=y_err, p0=[0.0, 0.18, -0.06, 1.4], absolute_sigma=False, method='lm') par_e = np.sqrt(np.diag(pcov)) chi2_n = cu.chi2_ndf(x, y, y_err, cu.thetaY_unblinded, par) # plot the fit and data fig, ax = plt.subplots() # data ax.errorbar(x, y, xerr=x_err, yerr=y_err, linewidth=0, elinewidth=2, color="green", marker="o", label="Sim.") # fit ax.plot(x, cu.thetaY_unblinded(x, par[0], par[1], par[2], par[3]), color="red", label='Fit') # deal with fitted parameters (to display nicely) parNames = [ r"$ A_{\mu}$", r"$ A_{\rm{EDM}}$", "c", r"$\omega$" ] units = ["mrad", "mrad", "mrad", "MhZ"] prec = 2 # set custom precision #form complex legends legend1_chi2 = cu.legend1_fit(chi2_n[0]) legned1_par = "" legned1_par = cu.legend_par(legned1_par, parNames, par, par_e, units) legend1 = legend1_chi2 + "\n" + legned1_par print(legend1) legend2 = data_type + "\n" + p_cut + "\n N=" + cu.sci_notation( N) #place on the plot and save y1, y2, x1, x2 = 0.15, 0.85, 0.25, 0.70 cu.textL(ax, x1, y1, legend1, font_size=16, color="red") cu.textL(ax, x2, y2, legend2, font_size=16) ax.legend(loc='center right', fontsize=16) ax.set_ylabel(y_label, fontsize=18) ax.set_xlabel(x_label, fontsize=18) plt.tight_layout() plt.savefig("../fig/" + outdir + "/" + fullLabel + ".png") plt.clf()