def write_bad_spills_to_table(server, schema, bad_spill_set, clear=False):
try:
db = mdb.connect(read_default_file='../.my.cnf',
read_default_group='guest',
db=schema,
host=server,
port=server_dict[server]['port'])
cur = db.cursor()
if not table_exists(server, schema, "bad_spills"):
cur.execute("CREATE TABLE bad_spills (spillID INT PRIMARY KEY)")
if clear:
cur.execute("DELETE FROM bad_spills")
if len(bad_spill_set) > 0:
query = "INSERT INTO bad_spills (spillID) VALUES "
for bad_spill in bad_spill_set:
query += "(" + str(bad_spill) + "), "
query = query[:-2]
cur.execute(query)
except mdb.Error, e:
print "Error %d: %s" % (e.args[0], e.args[1])
return 1
def mc_truth_get_analysis_data(server='seaquel.physics.illinois.edu',
analysis_schema='user_dannowitz_mc',
mc_schema_list=['mc_drellyan_LD2_M001_S001'],
analysis_table='kDimuon_mc_truth',
fresh_start=False,
*wargs):
if not schema_exists(server, analysis_schema):
make_analysis_schema(server, analysis_schema)
if not table_exists(server, analysis_schema, analysis_table):
mc_truth_make_mDimuon_table(server, analysis_schema, analysis_table,
mc_schema_list[0])
if fresh_start:
drop_analysis_table(server, analysis_schema, analysis_table)
mc_truth_make_mDimuon_table(server, analysis_schema, analysis_table,
mc_schema_list[0])
if fresh_start:
for merged_schema in mc_schema_list:
mc_truth_fill_dimuon_table(server, analysis_schema, analysis_table,
merged_schema)
dimuon_df = mc_truth_get_dimuon_df(server, analysis_schema, analysis_table)
return dimuon_df
def get_dimuon_df(server='e906-db3.fnal.gov',
analysis_schema='user_dannowitz_analysis',
source_schema_list=(),
analysis_table='kDimuon_analysis',
fresh_start=False,
truth_mc=False,
tracked_mc=False,
likesign=None):
if len(source_schema_list) == 0:
print ("Please pass one or many source schemas.")
if not schema_exists(server, analysis_schema):
make_analysis_schema(server, analysis_schema)
if not fresh_start and not table_exists(server, analysis_schema, analysis_table):
make_analysis_table(server, analysis_schema,
analysis_table, source_schema_list[0],
truth_mc=truth_mc, tracked_mc=tracked_mc)
if fresh_start:
drop_analysis_table(server, analysis_schema, analysis_table)
make_analysis_table(server, analysis_schema,
analysis_table,
source_schema_list[0],
truth_mc=truth_mc,
tracked_mc=tracked_mc,
likesign=likesign)
for source_schema in source_schema_list:
print ('>>> Retrieving and Cleaning Dimuon Data for %s' %
source_schema)
fill_analysis_table(server, analysis_schema,
analysis_table, source_schema,
tracked_mc=tracked_mc,
truth_mc=truth_mc,
likesign=likesign)
dimuon_df = fetch_analysis_table(server, analysis_schema, analysis_table,
truth_mc=truth_mc,
tracked_mc=tracked_mc)
if all(field in dimuon_df.columns for field in ('dpx', 'dpy')):
dimuon_df['dpt'] = np.sqrt(dimuon_df['dpx'] ** 2 + dimuon_df['dpy'] ** 2)
if all(field in dimuon_df.columns for field in ('px1', 'px2', 'py1', 'py2')):
dimuon_df['pt1'] = np.sqrt(dimuon_df['px1'] ** 2 + dimuon_df['py1'] ** 2)
dimuon_df['pt2'] = np.sqrt(dimuon_df['px2'] ** 2 + dimuon_df['py2'] ** 2)
dimuon_df = add_target(dimuon_df)
if not truth_mc and not tracked_mc:
dimuon_df = add_roadset(dimuon_df)
dimuon_df = add_weight(dimuon_df)
if 'weight' in dimuon_df.columns:
dimuon_df['weight_sq'] = dimuon_df['weight'] ** 2
return dimuon_df
def main():
print 'Hello World!'
print table_exists('e906-db3.fnal.gov', 'user_dannowitz_analysis_Mar30',
"bad_spills")