def __init__(self, db=RUtils.RDateBasePool().begin(), token=None):
self.config = RUtils.RConfig()
self.info = UserInfo()
self.db = db
self.session = RUtils.RMemorySessionStore()
if token:
self.info.token = token
self.login_by_token()
def check_login(*args, **kwargs):
if 'token' not in args[1].params.keys():
raise RUtils.RError(2)
data = user.User(db=args[1].context['sql'],
token=args[1].params['token'])
data.login_by_token()
if not data.info.ifLogin:
raise RUtils.RError(3)
func(*args, user=user, **kwargs)
def reload(self):
result = RUtils.RDateBasePool().execute(
'select * from users where uuid = %s;', (self.uuid,)
)
if not result:
raise RUtils.RError(0)
result = result[0]
self.admin = result['admin']
self.status = result['status']
self.username = result['username']
self.email = result['email']
def on_get(self, req, resp, comic_id):
comic_data = self.db.comic_list.find_one({'_id': ObjectId(comic_id)})
if not comic_data:
raise RUtils.RError(404)
data = self.db.comic.find({'name': comic_data['name']})
result = {}
result['comic'] = {
'id': str(comic_data['_id']),
'name': comic_data['name'],
'update_time': comic_data['update_time'],
'mobi': comic_data['mobi'],
'mobi_size': comic_data['mobi_size']
}
result['pic'] = []
for i in data:
if 'update_time' not in i.keys():
i['update_time'] = 0
result['pic'].append({
'chapter': i['chapter'],
'pic_num': len(i['pic']),
'update_time': i['update_time'],
'next': i['next'],
'mobi': i['mobi'],
'mobi_size': i['mobi_size']
})
req.context['result'] = result
def on_get(self, req, resp):
if 'password' not in req.params.keys():
raise RUtils.RError(2)
if req.params['password'] != RUtils.RConfig().password:
raise RUtils.RError(2)
data = self.db.students.find({})
result = []
for i in data:
result.append({
'id': str(i['_id']),
'name': i['name'],
'phone': i['phone'],
'email': i['email'],
'academy': i['academy'],
'wechat': i['wechat'],
'arrive_date': i['arrive_date'],
'leave_date': i['leave_date'],
'job': i['job'],
'special': i['special'],
'register_time': i['register_time']
})
req.context['result'] = result
def read_excel(name, sheetname='Detail Data', date_cols=[], preselection=''):
import RUtils, os
Rcode = "library(readxl)"
Rcode += '\n' + "df=read_excel('%s', sheet='%s')" % (name, sheetname)
Rcode += '\n' + "write.csv(df,'/tmp/%s/my_csv_from_excel.csv') " % os.getenv(
'USER')
RUtils.Rscript(code=Rcode)
df = pd.read_csv('/tmp/%s/my_csv_from_excel.csv' % os.getenv('USER'))
try:
robust_column_names(df)
convert_str_to_dates(df, date_cols)
except:
convert_str_to_dates(df, date_cols)
robust_column_names(df)
pass
if len(preselection): df = df[df.eval(preselection)]
return df
def blinded_10_par(x, *pars):
'''
CBO (4 par) + lost muons (1 par) with constant LT (5 par, 4/5 var 1 const LT)
'''
import RUtils as ru
time = x
ds = _DS
if (ds == -1): raise Exception("DS not set via cu._DS=x")
K = pars[-1]
'''
The code and the LM spectra courtesy of N. Kinnaird (Muon spin precession frequency extraction and decay positron track fitting in Run-1 of the Fermilab Muon g − 2 experiment, PhD thesis, Boston University (2020).)
'''
# use a pre-made histograms of muon loss spectrum to get this integral
#the times in the histogram are in ns
L = 1.0 - K * ru.LM_integral(
time * 1e3,
ds) * -3.75e-7 # the 1e-10 is just an arbitrary scaling factor
return blinded_wiggle_function_cbo(time, *pars[0:10]) * L
def login_by_password(self, username=None, password=None):
if not username or not password:
return False
password = user_utils.hash_password(password)
result = self.db.execute(
'select * from users where username=%s and password=%s and status > 0;',
(username, password)
)
if not result:
return False
result = result[0]
self.info.uuid = result['uuid']
self.info.token = RUtils.generate_code(64)
self.info.ifLogin = True
self.info.login_expired_date = time.time() + self.config.login_expired_time
self.info.reload()
self.session.push(self.info.token, self.info)
return True
def on_post(self, req, resp):
if 'request' not in req.context.keys():
raise RUtils.RError(6)
request = req.context['request']
if 'password' not in request.keys():
raise RUtils.RError(8)
if request['password'] != self.config.password:
raise RUtils.RError(9)
if 'url' not in request.keys():
raise RUtils.RError(10)
if not re.match(
'((http|ftp|https):\/\/)?[\w\-_]+(\.[\w\-_]+)+([\w\-\.,@?^=%&:/~\+#]*[\w\-\@?^=%&/~\+#])?',
request['url']):
raise RUtils.RError(10)
if not self.db.comic_list.find_one({'url': request['url']}):
self.db.comic_list.insert({'url': request['url']})
else:
raise RUtils.RError(11)
def on_post(self, req, resp):
if 'request' not in req.context.keys():
raise RUtils.RError(6)
request = req.context['request']
if 'password' not in request.keys():
raise RUtils.RError(8)
if request['password'] != self.config.password:
raise RUtils.RError(9)
if 'id' not in request.keys():
raise RUtils.RError(13)
chapter = self.db.comic.find_one({"_id": request['id']})
if not chapter:
raise RUtils.RError(12)
if 'method' not in request.keys():
raise RUtils.RError(13)
if request['method'] == 'reset_mobi':
self.db.comic.update_one(
{"_id": chapter["_id"]},
{"$set": {
"mobi": False,
"mobi_failed": 0
}})
elif request['method'] == 'reset_download':
self.db.comic.update_one(
{"_id": chapter["_id"]},
{"$set": {
"flag": 0,
"download_failed": 0
}})
elif request['method'] == 'refetch':
self.db.comic.update_one({"_id": chapter["_id"]},
{"$set": {
"flag": -1
}})
else:
raise RUtils.RError(13)
arg_parser.add_argument("--gauss", action='store_true', default=False)
arg_parser.add_argument("--bins", action='store_true', default=False)
args=arg_parser.parse_args()
#TODO impliment check of at least 1 arg
# args_v = vars(arg_parser.parse_args())
# if all(not args_v):
# arg_parser.error('No arguments provided.')
#set some global variables
info=[]
names=[]
# get data from 2D hist
if(args.read):
dataXY, n_binsXY, dBinsXY = ru.hist2np(None, None, file_path=args.file_path, hist_path=args.hist_path, from_root=True)
# store that data
np.save("../DATA/misc/dataXY.npy", dataXY)
# and info
# edm_setting=args.file_path.split("/")[1].split(".")[0]
# print(args.hist_path)
# Q_cut=args.hist_path.split("/")[0]
# data_type=args.hist_path.split("/")[1]
# time_cut=args.hist_path.split("/")[2]+r" $\mathrm{\mu}$s"
# p_cut=args.hist_path.split("/")[3]+" MeV"
# y_label=args.hist_path.split("/")[4].split("_")[0]
# x_label=args.hist_path.split("/")[4].split("_")[2:]
# print(x_label)
# # N=len(dataXY[0])
# #some specific string transforms
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 __init__(self, client):
self.client = client
self.db = self.client.comic
self.config = RUtils.RConfig()
def on_post(self, req, resp):
if 'request' not in req.context.keys():
raise RUtils.RError(6)
request = req.context['request']
# Check Name
if 'name' not in request.keys():
raise RUtils.RError(3)
name = request['name']
if len(name) > 10:
raise RUtils.RError(3)
# Check Phone
if 'phone' not in request.keys():
raise RUtils.RError(4)
phone = request['phone']
if not re.match('0?(13|14|15|16|17|18)[0-9]{9}', phone):
raise RUtils.RError(4)
if self.db.students.find_one({'phone': phone}):
raise RUtils.RError(8)
# Check Email
if 'email' not in request.keys():
raise RUtils.RError(5)
email = request['email']
if not re.match('^[a-z0-9](\w|\.|-)*@([a-z0-9]+-?[a-z0-9]+\.){1,3}[a-z]{2,10}$', email):
raise RUtils.RError(5)
if self.db.students.find_one({'email': email}):
raise RUtils.RError(9)
# Check School
if 'academy' not in request.keys():
raise RUtils.RError(6)
school = request['academy']
# Wechat
wechat = ""
if 'wechat' in request.keys():
wechat = request['wechat']
# Arrive_date
arrive_date = ""
if 'arrive_date' in request.keys():
arrive_date = request['arrive_date']
# Leave_date
leave_date = ""
if 'leave_date' in request.keys():
leave_date = request['leave_date']
# Job
job = ""
if 'job' in request.keys():
job = request['job']
# Special
special = ""
if 'special' in request.keys():
special = request['special']
if len(self.db.students.find({})) > 100:
raise RUtils.RError(10)
register_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
self.db.students.insert({
'name': name,
'phone': phone,
'email': email,
'academy': school,
'wechat': wechat,
'arrive_date': arrive_date,
'leave_date': leave_date,
'job': job,
'special': special,
'register_time': register_time
})
def hash_password(password):
return hashlib.sha512(password +
RUtils.RConfig().password_salt).hexdigest()[0:64]
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 error_handle(ex, req, resp, params):
if isinstance(ex, falcon.HTTPError):
raise ex
else:
traceback.print_exc()
raise RUtils.RError(0)
def process_request(self, req, resp):
req.context['sql'] = RUtils.RDateBasePool().begin()
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)
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()