0
0

compton_combiner.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614
  1. """Script to combine compton measurements with runs and process these data
  2. """
  3. import argparse
  4. from configparser import ConfigParser
  5. from datetime import datetime, timedelta, timezone
  6. import logging
  7. import os
  8. import sys
  9. import sqlite3
  10. from typing import Union, Tuple, Optional
  11. from compton_filter import CalibrdbHandler
  12. from iminuit import Minuit
  13. import matplotlib.dates as mdates
  14. import matplotlib.pyplot as plt
  15. from mysql.connector import connect, Error
  16. import numpy as np
  17. import pandas as pd
  18. from tqdm import tqdm
  19. SEASONS = {
  20. 'name': ['PHIOMEGA2012', 'RHO2013', 'BRK2013/16', 'HIGH2017', 'RHO2018', 'HIGH2019', 'LOW2020', 'HIGH2020', 'HIGH2021', 'NNBAR2022', 'HIGH2023'],
  21. 'start_run': [17405, 18809, 32076, 36872, 48938, 70014, 85224, 89973, 98116, 107342, 131913, None],
  22. }
  23. class RunsDBHandler():
  24. def __init__(self, host: str = 'cmddb', database: str = 'online', user: str = None, password: str = None):
  25. self.conn = connect(host = host, database = database, user = user, password = password)
  26. self.cur = self.conn.cursor()
  27. self.cur.execute("SET time_zone = '+07:00';")
  28. @property
  29. def fields(self) -> list:
  30. """Returns a list of available columns in the RunsDB
  31. """
  32. self.cur.execute("""DESCRIBE Runlog""")
  33. return self.cur.fetchall()
  34. def load_tables(self, range: Union[Tuple[int, Optional[int]], Tuple[datetime, datetime]], energy_point: Optional[float] = None, select_bad_runs: bool = False):
  35. """
  36. Returns a slice of the table with following fields: run, starttime, stoptime, energy, luminosity
  37. Parameters
  38. ----------
  39. range : Union[Tuple[int, Optional[int]], Tuple[datetime, datetime]]
  40. selection range
  41. int range defines an interval in runs (first included, last excluded)
  42. datetime range defines a time interval (NSK: +7:00 time)
  43. energy_point : Optional[float]
  44. energy point name, MeV (default is None)
  45. select_bad_runs : bool
  46. select runs with labels except (Y) (default is False)
  47. """
  48. cond = ""
  49. if isinstance(range[0], int):
  50. cond = f" AND run >= {range[0]} "
  51. if range[1] is not None:
  52. cond += f" AND run < {range[1]} "
  53. elif isinstance(range[0], datetime):
  54. cond = f" AND starttime >= %s "
  55. if range[1] is not None:
  56. cond += " AND stoptime <= %s"
  57. else:
  58. range = (range[0], )
  59. energy_cond = ""
  60. if energy_point is not None:
  61. energy_cond = f" AND energy = {energy_point}"
  62. quality_cond = ' quality = "Y" '
  63. if select_bad_runs:
  64. quality_cond = ' quality <> "Y" '
  65. sql_query = f"""
  66. SELECT
  67. run,
  68. starttime,
  69. stoptime,
  70. energy,
  71. luminosity
  72. FROM Runlog
  73. WHERE
  74. {quality_cond}
  75. {cond}
  76. {energy_cond}
  77. AND luminosity > 0
  78. AND stoptime > starttime
  79. AND nevent > 0
  80. ORDER BY run DESC"""
  81. if isinstance(range[0], datetime):
  82. self.cur.execute(sql_query, range)
  83. else:
  84. self.cur.execute(sql_query)
  85. field_names = [i[0] for i in self.cur.description]
  86. res = self.cur.fetchall()
  87. return res, field_names
  88. def __del__(self):
  89. self.conn.close()
  90. class Combiner():
  91. """Combines a dataframe with runs and a dataframe with compton measurements together
  92. """
  93. def __init__(self, runsdb: Tuple[list, list], clbrdb: Tuple[list, list]):
  94. """
  95. Parameters
  96. ----------
  97. runsdb : Tuple[list, list]
  98. table of runs (rows and field names)
  99. clbrdb : Tuple[list, list]
  100. table of compton measurements (rows and field names)
  101. """
  102. rdb_rows, r_fld = runsdb
  103. cdb_rows, c_fld = clbrdb
  104. self.conn = sqlite3.connect(":memory:", detect_types=sqlite3.PARSE_DECLTYPES|sqlite3.PARSE_COLNAMES)
  105. self.cur = self.conn.cursor()
  106. self.cur.execute(f"CREATE table runs (run, elabel, starttime timestamp, stoptime timestamp, luminosity)")
  107. self.cur.execute(f"CREATE table compton (begintime timestamp, endtime timestamp, e_mean, e_std, spread_mean, spread_std)")
  108. run_row_generator = map(lambda x: (x[r_fld.index("run")], x[r_fld.index("energy")],
  109. x[r_fld.index("starttime")], x[r_fld.index("stoptime")],
  110. x[r_fld.index("luminosity")]), rdb_rows)
  111. c_data_idx = c_fld.index("data")
  112. compton_row_generator = map(lambda x: (x[c_fld.index("begintime")], x[c_fld.index("endtime")],
  113. float(x[c_data_idx][0]), float(x[c_data_idx][1]),
  114. float(x[c_data_idx][2]), float(x[c_data_idx][3])), cdb_rows)
  115. self.cur.executemany(f"""INSERT into runs VALUES ({','.join(['?']*5)})""", run_row_generator)
  116. self.cur.executemany(f"""INSERT into compton VALUES ({','.join(['?']*6)})""", compton_row_generator)
  117. self.__create_combined_table()
  118. def __create_combined_table(self):
  119. create_combined_query = """
  120. CREATE TABLE combined_table AS
  121. SELECT
  122. runs.run AS run,
  123. runs.elabel AS elabel,
  124. runs.starttime as "run_start [timestamp]",
  125. runs.stoptime AS "run_stop [timestamp]",
  126. compton.begintime AS "compton_start [timestamp]",
  127. compton.endtime AS "compton_stop [timestamp]",
  128. runs.luminosity, compton.e_mean, compton.e_std, compton.spread_mean, compton.spread_std
  129. FROM runs, compton
  130. WHERE
  131. (runs.starttime BETWEEN compton.begintime AND compton.endtime)
  132. OR (runs.stoptime BETWEEN compton.begintime AND compton.endtime)
  133. OR (compton.begintime BETWEEN runs.starttime AND runs.stoptime)
  134. OR (compton.endtime BETWEEN runs.starttime AND runs.stoptime);
  135. """
  136. self.cur.execute(create_combined_query)
  137. return
  138. def combined_table(self) -> pd.DataFrame:
  139. """Returns combined dataframe
  140. """
  141. sql_query = """
  142. SELECT * FROM combined_table;
  143. """
  144. df = pd.read_sql(sql_query, self.conn)
  145. df['common_duration'] = df[['run_stop', 'compton_stop']].min(axis=1) - df[['run_start', 'compton_start']].max(axis=1)
  146. df['run_duration'] = df['run_stop'] - df['run_start']
  147. df['run_in_measurement'] = df['common_duration']/df['run_duration']
  148. df = df.sort_values(by='run_in_measurement', ascending=False).drop_duplicates(subset='run').sort_values(by='run')
  149. df = df.drop(['run_duration', 'common_duration'], axis=1) #, 'run_start', 'run_stop'
  150. return df
  151. def __del__(self):
  152. self.conn.close()
  153. class Likelihood():
  154. """
  155. Likelihood function
  156. """
  157. def __init__(self, means: np.array, sigmas: np.array, weights: np.array):
  158. """
  159. Parameters
  160. ----------
  161. means : np.array
  162. array of means, [MeV]
  163. sigmas : np.array
  164. array of standard deviations, [MeV]
  165. weights : np.array
  166. array of luminosities
  167. """
  168. self.means = means
  169. self.sigmas = sigmas
  170. self.weights = weights/weights.mean()
  171. # w_norm = (weights**2).sum()/(weights.sum())
  172. # self.weights = weights/w_norm
  173. def __call__(self, mean: float, sigma: float):
  174. """
  175. Calls likelihood calculation
  176. Parameters
  177. ----------
  178. mean : float
  179. expected mean
  180. sigma : float
  181. expected standard deviation
  182. """
  183. sigma_total = np.sqrt(sigma**2 + self.sigmas**2)
  184. ln_L = -np.sum( self.weights*( ((mean - self.means)**2)/(2*(sigma_total**2)) + np.log(sigma_total) ) )
  185. return -ln_L
  186. def __estimate_point_with_closest(comb_df: pd.DataFrame, runs_df: pd.DataFrame, compton_df: pd.DataFrame):
  187. # estimate energy by the nearest points
  188. min_run_time = runs_df[runs_df.run == comb_df.iloc[0].at['run_first']].iloc[0].at['starttime']
  189. max_run_time = runs_df[runs_df.run == comb_df.iloc[0].at['run_last']].iloc[0].at['stoptime']
  190. nearest_row_before = compton_df.iloc[pd.Index(compton_df.endtime).get_loc(min_run_time, 'nearest')]
  191. nearest_row_after = compton_df.iloc[pd.Index(compton_df.begintime).get_loc(max_run_time, 'nearest')]
  192. # regulatization
  193. nearest_row_before['data'][1] = max(nearest_row_before['data'][3], 1e-3)
  194. nearest_row_after['data'][3] = max(nearest_row_after['data'][3], 1e-3)
  195. nearest_row_before['data'][1] = max(nearest_row_before['data'][1], 1e-3)
  196. nearest_row_after['data'][3] = max(nearest_row_after['data'][3], 1e-3)
  197. mean_energy = (nearest_row_before['data'][0] + nearest_row_after['data'][0])/2
  198. mean_spread = (nearest_row_before['data'][2] + nearest_row_after['data'][2])/2
  199. std_energy = np.sqrt(1/(1/(nearest_row_before['data'][1])**2 + 1/(nearest_row_after['data'][1])**2))
  200. std_spread = np.sqrt(1/(1/(nearest_row_before['data'][3])**2 + 1/(nearest_row_after['data'][3])**2))
  201. sys_energy = np.std([nearest_row_before['data'][0], nearest_row_after['data'][0]])
  202. return {
  203. 'energy_point': comb_df.elabel.min(),
  204. 'first_run': comb_df.run_first.min(),
  205. 'last_run': comb_df.run_last.max(),
  206. 'mean_energy': mean_energy,
  207. 'mean_energy_stat_err': std_energy,
  208. 'mean_energy_sys_err': sys_energy,
  209. 'mean_spread': mean_spread,
  210. 'mean_spread_stat_err': std_spread,
  211. 'used_lum': 0,
  212. 'comment': 'indirect measurement #2',
  213. }, pd.DataFrame([])
  214. def averager_on_lums(df: pd.DataFrame) -> dict:
  215. """Averaging as <E> = \frac{\sum{L_i E_i}{\sum{L_i}},
  216. \deltaE^2 = \frac{\sum{(L_i \delta E_i)^2}}{(\sum L_i)^2}
  217. Attention: I think it's incorrect way of avaraging.
  218. Parameters
  219. ----------
  220. df : pd.DataFrame
  221. input dataframe containg means and spreads
  222. Returns
  223. -------
  224. dict
  225. averaged mean and spread
  226. """
  227. mean_en = (df.e_mean * df.luminosity).sum() / df.luminosity.sum()
  228. sys_err = df.e_mean.std()
  229. stat_err = np.sqrt( np.sum((df.luminosity * df.e_std)**2) ) / df.luminosity.sum()
  230. mean_spread = (df.spread_mean * df.luminosity).sum() / df.luminosity.sum()
  231. std_spread = np.sqrt( np.sum((df.luminosity * df.spread_std)**2) ) / df.luminosity.sum()
  232. return {
  233. 'mean_energy': mean_en,
  234. 'mean_energy_stat_err': stat_err,
  235. 'mean_energy_sys_err': sys_err,
  236. 'mean_spread': mean_spread,
  237. 'mean_spread_stat_err': std_spread,
  238. }
  239. def ultimate_averager(df: pd.DataFrame) -> dict:
  240. """Complete averager for estimation of mean energy and energy spread
  241. Parameters
  242. ----------
  243. df : pd.DataFrame
  244. input dataframe containing means and spreads
  245. Returns
  246. -------
  247. dict
  248. averaged mean and spread
  249. """
  250. m = Minuit(Likelihood(df.e_mean, df.e_std, df.luminosity), mean=df.e_mean.mean(), sigma=df.e_mean.std(ddof=0))
  251. m.errordef = 0.5
  252. m.limits['sigma'] = (0, None)
  253. m.migrad();
  254. # print(m.migrad())
  255. sys_err = m.values['sigma']
  256. mean_en = m.values['mean']
  257. mean_spread = np.sum(df.spread_mean*df.luminosity/(df.spread_std**2))/np.sum(df.luminosity/(df.spread_std**2))
  258. std_spread = np.sqrt(1/np.sum((df.luminosity/df.luminosity.mean())/df.spread_std**2))
  259. return {
  260. 'mean_energy': mean_en,
  261. 'mean_energy_stat_err': m.errors['mean'],
  262. 'mean_energy_sys_err': sys_err,
  263. 'mean_spread': mean_spread,
  264. 'mean_spread_stat_err': std_spread,
  265. }
  266. def calculate_point(comb_df: pd.DataFrame, runs_df: pd.DataFrame, compton_df: pd.DataFrame, rdb, averager: callable = ultimate_averager) -> dict:
  267. """Calculates parameters of the energy (mean, std, spread) in this dataFrame
  268. Parameters
  269. ----------
  270. comb_df : pd.DataFrame
  271. table of the measurements linked with runs
  272. runs_df : pd.DataFrame
  273. table of the runs
  274. compton_df : pd.DataFrame
  275. table of the comptons
  276. averager : callable
  277. function for averaging (ultimate_averager or averager_on_lums)
  278. Returns
  279. -------
  280. dict, pd.DataFrame
  281. average parameters on this DataFrame, clean dataFrame
  282. """
  283. if (len(comb_df) == 1) and pd.isnull(comb_df.iloc[0].at['compton_start']):
  284. # no direct measurements of the compton during data runs
  285. min_Yruntime = runs_df[runs_df.run == comb_df.iloc[0].at['run_first']].iloc[0].at['starttime']
  286. max_Yruntime = runs_df[runs_df.run == comb_df.iloc[0].at['run_last']].iloc[0].at['stoptime']
  287. dlt0 = timedelta(days=1)
  288. # assymetric time because energy can be stable only after
  289. runs_df_with_bads = rdb.load_tables((min_Yruntime, max_Yruntime + dlt0), energy_point = comb_df.iloc[0].at['elabel'], select_bad_runs = True)
  290. if len(runs_df_with_bads[0]) == 0:
  291. return __estimate_point_with_closest(comb_df, runs_df, compton_df)
  292. runs_df_with_bads_df = pd.DataFrame(runs_df_with_bads[0], columns = runs_df_with_bads[1])
  293. min_run_time, max_run_time = min(min_Yruntime, runs_df_with_bads_df.starttime.min()), max(max_Yruntime, runs_df_with_bads_df.stoptime.max())
  294. compton_meas = compton_df.query('((begintime>=@min_run_time)&(begintime<=@max_run_time))|((endtime>=@min_run_time)&(endtime<=@max_run_time))').copy()
  295. if len(compton_meas) == 0:
  296. # no compton measurements
  297. raise Exception("No measurement in this point. Pass it.")
  298. res_df = pd.DataFrame(list(map(lambda x: {
  299. 'compton_start': x[1]['begintime'],
  300. 'compton_stop': x[1]['endtime'],
  301. 'e_mean': float(x[1]['data'][0]),
  302. 'e_std': float(x[1]['data'][1]),
  303. 'spread_mean': float(x[1]['data'][2]),
  304. 'spread_std': float(x[1]['data'][3]),
  305. }, compton_meas.iterrows())))
  306. res_df = res_df.query(f'abs(e_mean -{comb_df.iloc[0].at["elabel"]})<5')
  307. if len(res_df) == 0:
  308. return __estimate_point_with_closest(comb_df, runs_df, compton_df)
  309. return {
  310. 'energy_point': comb_df.elabel.min(),
  311. 'first_run': comb_df.run_first.min(),
  312. 'last_run': comb_df.run_last.max(),
  313. 'mean_energy': res_df.e_mean.mean(),
  314. 'mean_energy_stat_err': np.sqrt(1/np.sum(1/(res_df.e_std)**2)),
  315. 'mean_energy_sys_err': np.abs(comb_df.iloc[0].at['elabel'] - res_df.e_mean.mean()),
  316. 'mean_spread': res_df.spread_mean.mean(),
  317. 'mean_spread_stat_err':np.sqrt(1/np.sum(1/(res_df.spread_std)**2)),
  318. 'used_lum': 0,
  319. 'comment': 'indirect measurement #1',
  320. }, res_df
  321. comb_df = comb_df.reset_index()
  322. df = comb_df.loc[~comb_df.compton_start.isna()].copy()
  323. df.spread_std = np.where(df.spread_std < 1e-4, 1e-4, df.spread_std)
  324. df = df[df.e_std > 0]
  325. mean_energy = np.sum(df.e_mean*df.luminosity/(df.e_std**2))/np.sum(df.luminosity/(df.e_std**2))
  326. good_criterion = np.abs((df.e_mean - mean_energy)/np.sqrt(df.e_mean.std(ddof=0)**2 + df.e_std**2)) < 5
  327. # print('WTF:', df[~good_criterion].index)
  328. df = df[good_criterion]
  329. df['accepted'] = 1
  330. averages = averager(df)
  331. res_dict = {
  332. 'energy_point': comb_df.elabel.min(),
  333. 'first_run': comb_df.run_first.min(),
  334. 'last_run': comb_df.run_last.max(),
  335. 'mean_energy': averages['mean_energy'],
  336. 'mean_energy_stat_err': averages['mean_energy_stat_err'],
  337. 'mean_energy_sys_err': averages['mean_energy_sys_err'],
  338. 'mean_spread': averages['mean_spread'],
  339. 'mean_spread_stat_err': averages['mean_spread_stat_err'],
  340. 'used_lum': df.luminosity.sum()/comb_df.luminosity_total.sum(),
  341. 'comment': '',
  342. }
  343. comb_df['accepted'] = 0
  344. comb_df.loc[df.index, 'accepted'] = 1
  345. return res_dict, comb_df.set_index('point_idx')
  346. def process_intersected_compton_meas(combined_df: pd.DataFrame) -> pd.DataFrame:
  347. """Replaces compton measurements writed on the border of two energy points on NaNs
  348. """
  349. energy_point_borders = combined_df[['point_idx', 'elabel', 'run_start', 'run_stop']].groupby(['point_idx'], dropna=True).agg(
  350. elabel_start_time=('run_start', 'min'), elabel_stop_time=('run_stop', 'max'),
  351. )
  352. df_comb = combined_df.set_index('point_idx').join(energy_point_borders, how='left')
  353. df_comb['comptonmeas_in_elabel'] = (df_comb[['elabel_stop_time', 'compton_stop']].min(axis=1) - df_comb[['elabel_start_time', 'compton_start']].max(axis=1))/(df_comb['compton_stop'] - df_comb['compton_start'])
  354. df_comb = df_comb.query('comptonmeas_in_elabel < 0.7')
  355. border_comptons = df_comb.compton_start.values
  356. combined_df.loc[combined_df.compton_start.isin(border_comptons),
  357. ['compton_start', 'compton_stop', 'e_mean', 'e_std', 'spread_mean', 'spread_std', 'luminosity']] = np.nan
  358. return combined_df
  359. def process_combined(combined_df: pd.DataFrame, runs_df: pd.DataFrame, compton_df: pd.DataFrame, pics_folder: Optional[str] = None, rdb: Optional[RunsDBHandler] = None, old_averager: bool = False, energy_point_csv_folder: Optional[str] = None) -> pd.DataFrame:
  360. if pics_folder is not None:
  361. plt.ioff()
  362. plt.style.use('ggplot')
  363. locator = mdates.AutoDateLocator(minticks=5)
  364. formatter = mdates.ConciseDateFormatter(locator)
  365. formatter.formats = ['%y', '%b', '%d', '%H:%M', '%H:%M', '%S.%f', ]
  366. formatter.zero_formats = [''] + formatter.formats[:-1]
  367. formatter.zero_formats[3] = '%d-%b'
  368. formatter.offset_formats = ['', '%Y', '%b %Y', '%d %b %Y', '%d %b %Y', '%d %b %Y %H:%M', ]
  369. runs_df = runs_df.rename({'luminosity': 'luminosity_full', 'energy': 'elabel'}, axis=1)
  370. combined_df = pd.merge(combined_df.drop(['elabel'], axis=1), runs_df[['run', 'elabel', 'luminosity_full']], how='outer')
  371. combined_df = combined_df.sort_values(by='run')
  372. combined_df['point_idx'] = np.cumsum(~np.isclose(combined_df.elabel, combined_df.elabel.shift(1), atol=1e-4))
  373. combined_df = process_intersected_compton_meas(combined_df)
  374. combined_df['luminosity'] = combined_df['luminosity'].fillna(0)
  375. # combined_df.to_csv('file.csv')
  376. combined_df = combined_df.groupby(['point_idx', 'compton_start'], dropna=False).agg(
  377. elabel=('elabel', 'min'), elabel_test=('elabel', 'max'),
  378. run_first=('run', 'min'), run_last=('run', 'max'),
  379. luminosity=('luminosity', 'sum'), luminosity_total=('luminosity_full', 'sum'),
  380. compton_stop=('compton_stop', 'min'), compton_stop_test=('compton_stop', 'max'),
  381. e_mean=('e_mean', 'min'), e_mean_test=('e_mean', 'max'),
  382. e_std=('e_std', 'min'), e_std_test=('e_std', 'max'),
  383. spread_mean=('spread_mean', 'min'), spread_mean_test=('spread_mean', 'max'),
  384. spread_std=('spread_std', 'min'), spread_std_test=('spread_std', 'max'),
  385. ).reset_index().set_index('point_idx')
  386. # return combined_df
  387. result_df = pd.DataFrame(columns=['energy_point', 'first_run', 'last_run', 'mean_energy', 'mean_energy_stat_err', 'mean_energy_sys_err', 'mean_spread', 'mean_spread_stat_err', 'used_lum', 'comment'])
  388. for i, table in tqdm(combined_df.groupby('point_idx', dropna=False)):
  389. try:
  390. res_dict, good_df = calculate_point(table, runs_df, compton_df, rdb, averager_on_lums if old_averager else ultimate_averager)
  391. if energy_point_csv_folder is not None:
  392. save_columns = ['elabel', 'run_first', 'run_last', 'luminosity', 'compton_start', 'compton_stop', 'e_mean', 'e_std', 'spread_mean', 'spread_std', 'accepted']
  393. save_csv(good_df[save_columns].dropna(), f'{energy_point_csv_folder}/{res_dict["energy_point"]}_{res_dict["first_run"]}.csv', update_current=False)
  394. good_df = good_df.query('accepted==1')
  395. except Exception:
  396. continue
  397. # result_df = result_df.append(res_dict, ignore_index=True)
  398. result_df = pd.concat([result_df, pd.Series(res_dict).to_frame().T], ignore_index=True)
  399. if pics_folder is not None:
  400. plt_table = good_df.dropna()
  401. if len(plt_table) == 0:
  402. continue
  403. total_error = np.sqrt(res_dict["mean_energy_stat_err"]**2 + res_dict["mean_energy_sys_err"]**2)
  404. half_timedelta = (plt_table.compton_stop - plt_table.compton_start)/2
  405. time = plt_table.compton_start + half_timedelta
  406. dlt0, total_time = timedelta(days=1), plt_table.compton_stop.max() - plt_table.compton_stop.min()
  407. timelim = [plt_table.compton_start.min() - 0.05*total_time, plt_table.compton_stop.max() + 0.05*total_time]
  408. fig, ax = plt.subplots(1, 1, dpi=120, tight_layout=True)
  409. ax.errorbar(time, plt_table.e_mean, xerr=half_timedelta, yerr=plt_table.e_std, fmt='.')
  410. ax.axhline(res_dict['mean_energy'], color='black', zorder=3, label='Mean')
  411. ax.fill_between(timelim,
  412. [res_dict['mean_energy'] - total_error]*2,
  413. [res_dict['mean_energy'] + total_error]*2, color='green', zorder=1, alpha=0.4)
  414. ax.tick_params(axis='x', labelrotation=45)
  415. ax.xaxis.set_major_locator(locator)
  416. ax.xaxis.set_major_formatter(formatter)
  417. ax.set(title=f'{res_dict["energy_point"]}, E = {res_dict["mean_energy"]:.3f} ± {res_dict["mean_energy_stat_err"]:.3f} ± {res_dict["mean_energy_sys_err"]:.3f} MeV',
  418. xlabel='Time, NSK', ylabel='Energy, [MeV]', xlim=timelim)
  419. plt.savefig(f'{pics_folder}/{res_dict["first_run"]}_{res_dict["energy_point"]}.png', transparent=True)
  420. plt.close()
  421. types_dict = {ftype : float for ftype in ['energy_point', 'mean_energy', 'mean_energy_stat_err', 'mean_energy_sys_err', 'mean_spread', 'mean_spread_stat_err', 'used_lum']}
  422. types_dict.update({itype : int for itype in ['first_run', 'last_run']})
  423. result_df = result_df.astype(types_dict)
  424. return result_df
  425. def final_table_to_clbrdb(df: pd.DataFrame, clbrdb: CalibrdbHandler, runs_df: pd.DataFrame, season: str):
  426. """Write good values from the averaged table into clbrdb
  427. """
  428. good_values = (df.comment=='')|((df.comment!='')&((df.mean_energy.astype(float) - df.energy_point).abs()<5))
  429. df_clbrdb = df.loc[good_values].drop(['comment', 'used_lum'], axis=1)
  430. df_clbrdb = pd.merge(df_clbrdb, runs_df[['run', 'starttime']], how='left', left_on='first_run', right_on='run').drop(['run'], axis=1)
  431. df_clbrdb = pd.merge(df_clbrdb, runs_df[['run', 'stoptime']], how='left', left_on='last_run', right_on='run').drop(['run'], axis=1)
  432. df_clbrdb = df_clbrdb.assign(writetime=lambda df: df['stoptime'])
  433. df_clbrdb = df_clbrdb[['writetime', 'starttime', 'stoptime',
  434. 'energy_point', 'first_run', 'last_run', 'mean_energy',
  435. 'mean_energy_stat_err', 'mean_energy_sys_err', 'mean_spread', 'mean_spread_stat_err']].values.tolist()
  436. clbrdb.insert(df_clbrdb, 'Misc', 'RunHeader', 'Compton_run_avg', 'Default', comment = season)
  437. clbrdb.commit()
  438. def save_csv(df: pd.DataFrame, filepath: str, update_current: bool = True):
  439. """Saves csv file. Updates current file in filepath if exists"""
  440. if (os.path.isfile(filepath) and update_current):
  441. df_current = pd.read_csv(filepath)
  442. # df_current = df_current.append(df, ignore_index=True)
  443. df_current = pd.concat([df, df_current], ignore_index=True)
  444. df_current = df_current.drop_duplicates(subset=['energy_point', 'first_run'], keep='last')
  445. df = df_current
  446. df.to_csv(filepath, index=False, float_format='%g')
  447. return
  448. # python scripts/compton_combiner.py -s NNBAR2021 -c database.ini --csv_dir . --clbrdb
  449. def main():
  450. log_format = '[%(asctime)s] %(levelname)s: %(message)s'
  451. logging.basicConfig(stream=sys.stdout, format=log_format, level=logging.INFO) #"filename=compton_combiner.log"
  452. logging.info("compton_combiner is started")
  453. parser = argparse.ArgumentParser(description = 'Mean compton energy measurements from clbrdb')
  454. parser.add_argument('-s', '--season', help = 'Name of the season')
  455. parser.add_argument('-c', '--config', help = 'Config file containing information for access to databases')
  456. parser.add_argument('--csv_dir', help = 'Save csv file with data in the folder or not if skip it')
  457. parser.add_argument('--clbrdb', action = 'store_true', help = 'Update Compton_run_avg clbrdb or not')
  458. parser.add_argument('--pics_folder', help = 'Path to the directory for saving the pictures')
  459. parser.add_argument('--energy_point_csv_folder', help = 'Path to the directory for saving the result in detail for each energy point')
  460. parser.add_argument('--only_last', action = 'store_true', help = 'Compute values of the last (in Compton_run_avg clbrdb) and new points only')
  461. parser.add_argument('--old_averaging', action = 'store_true', help = 'Use old incomplete <E> = \frac{\sum{L_i E_i}{\sum{L_i}} averaging')
  462. args = parser.parse_args()
  463. logging.info(f"""Arguments: season {args.season}, config {args.config}, csv_dir {args.csv_dir}, save_to_clbrdb {args.clbrdb},
  464. pics_folder {args.pics_folder}, detailed_csv_folder {args.energy_point_csv_folder}, only_last {args.only_last}, old_average: {args.old_averaging}""")
  465. parser = ConfigParser()
  466. parser.read(args.config);
  467. rdb = RunsDBHandler(**parser['cmdruns'])
  468. clbrdb = CalibrdbHandler(**parser['clbrDB'])
  469. idx = SEASONS['name'].index(args.season)
  470. runs_range = (SEASONS['start_run'][idx], SEASONS['start_run'][idx+1])
  471. if args.only_last:
  472. res_avg = clbrdb.load_table('Misc', 'RunHeader', 'Compton_run_avg', num_last_rows = 1)
  473. if len(res_avg[0]) != 0:
  474. begintime = res_avg[0][0][res_avg[1].index("begintime")]
  475. runs_range = (begintime, None)
  476. logging.info(f"only_last flag enabled. Time range for runs database is used: {runs_range}")
  477. res_rdb = rdb.load_tables(runs_range)
  478. runs_df = pd.DataFrame(res_rdb[0], columns=res_rdb[1])
  479. if args.only_last:
  480. # Remain the studied season only (important for the first averaging in the season with turned on only last flag)
  481. start_season_run, stop_season_run = SEASONS['start_run'][idx], SEASONS['start_run'][idx+1] if SEASONS['start_run'][idx+1] else np.inf
  482. runs_df = runs_df.query(f'(run>=@start_season_run)&(run<@stop_season_run)')
  483. tdlt0 = timedelta(days=2)
  484. time_range = (runs_df.starttime.min() - tdlt0, runs_df.stoptime.max() + tdlt0)
  485. res_clbrdb = clbrdb.load_table('Misc', 'RunHeader', 'Compton_run', num_last_rows = None, timerange = time_range)
  486. cb = Combiner(res_rdb, res_clbrdb)
  487. comb_df = cb.combined_table()
  488. compton_df = pd.DataFrame(res_clbrdb[0], columns=res_clbrdb[1])
  489. cdf = process_combined(comb_df, runs_df, compton_df, args.pics_folder, rdb, args.old_averaging, args.energy_point_csv_folder)
  490. if args.csv_dir is not None:
  491. csv_path = os.path.join(args.csv_dir, f'{args.season}.csv')
  492. save_csv(cdf, csv_path)
  493. # cdf.to_csv(f'{args.season}.csv', index=False, float_format='%g')
  494. if args.clbrdb:
  495. final_table_to_clbrdb(cdf, clbrdb, runs_df, args.season)
  496. return
  497. if __name__ == "__main__":
  498. main()