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