diff --git a/fast_response/AlertFollowup.py b/fast_response/AlertFollowup.py index f574778..7fec44b 100644 --- a/fast_response/AlertFollowup.py +++ b/fast_response/AlertFollowup.py @@ -25,6 +25,15 @@ class AlertFollowup(PriorFollowup): _fix_index = True _float_index = not _fix_index _index = 2.5 + _bg_trial_dir = os.path.join( + '/data/ana/analyses/NuSources/', + '2021_v2_alert_stacking_FRA/fast_response/', + 'alert_precomputed_trials/' + ) + _sens_dir = '/data/ana/analyses/NuSources/2021_v2_alert_stacking_FRA/' \ + + 'fast_response/reference_sensitivity_curves/' + # These directories will need to changed for each AlertFollowup LLH configuration + # i.e. dataset combination and such def run_background_trials(self, ntrials = 1000): r"""For alert events with specific time windows, @@ -36,15 +45,11 @@ def run_background_trials(self, ntrials = 1000): test-statistic distribution with weighting from alert event spatial prior """ + current_rate = self.llh.nbackground / (self.duration * 86400.) * 1000. closest_rate = sensitivity_utils.find_nearest(np.linspace(6.2, 7.2, 6), current_rate) - - bg_trial_dir = '/data/ana/analyses/NuSources/' \ - + '2021_v2_alert_stacking_FRA/fast_response/' \ - + 'alert_precomputed_trials/' - pre_ts_array = sparse.load_npz( - bg_trial_dir + self._bg_trial_dir + 'precomputed_trials_delta_t_' + '{:.2e}_trials_rate_{:.1f}_low_stats.npz'.format( self.duration * 86400., closest_rate, self.duration * 86400.)) @@ -80,10 +85,9 @@ def ps_sens_range(self): highest sensitivity within the 90% contour of the skymap """ - sens_dir = '/data/ana/analyses/NuSources/2021_v2_alert_stacking_FRA/' \ - + 'fast_response/reference_sensitivity_curves/' + filename = os.path.join(self._sens_dir, f'ideal_ps_sensitivity_deltaT_{self.duration:.2e}_50CL.pkl') - with open(f'{sens_dir}ideal_ps_sensitivity_deltaT_{self.duration:.2e}_50CL.pkl', 'rb') as f: + with open(filename, 'rb') as f: ideal = pickle.load(f, encoding='bytes') delta_t = self.duration * 86400. src_theta, src_phi = hp.pix2ang(self.nside, self.ipix_90) @@ -101,10 +105,9 @@ def sens_range_plot(self): """ fig, ax = plt.subplots() - sens_dir = '/data/ana/analyses/NuSources/2021_v2_alert_stacking_FRA/' \ - + 'fast_response/reference_sensitivity_curves/' + filename = os.path.join(self._sens_dir, f'ideal_ps_sensitivity_deltaT_{self.duration:.2e}_50CL.pkl') - with open(f'{sens_dir}ideal_ps_sensitivity_deltaT_{self.duration:.2e}_50CL.pkl', 'rb') as f: + with open(filename, 'rb') as f: ideal = pickle.load(f, encoding='bytes') delta_t = self.duration * 86400. plt.plot(ideal[b'sinDec'], np.array(ideal[b'sensitivity'])*delta_t*1e6, lw=3, ls='-', @@ -115,7 +118,7 @@ def sens_range_plot(self): src_dec = np.unique(src_dec) src_dec = np.sin(src_dec) ax.axvspan(src_dec.min(), src_dec.max(), alpha=0.3, color=sns.xkcd_rgb['light navy blue'], - label='90\% contour region') + label='90% contour region') plt.text(0.05, 3e1, 'Min sens.: {:.1e}'.format(self.sens_range[0]) + r' GeV cm$^{-2}$') plt.text(0.05, 1.5e1, 'Max sens.: {:.1e}'.format(self.sens_range[1]) + r' GeV cm$^{-2}$') plt.grid(which='both', alpha=0.2, zorder=1) @@ -309,7 +312,7 @@ def ipixs_in_percentage(self, percentage): elif percentage == 0.5: msk = (skymap < 22.2) * (skymap > 0.) else: - raise ValueError('Must use 50\% or 90\% containment for alert events') + raise ValueError('Must use 50% or 90% containment for alert events') msk *= ~np.isnan(skymap) msk *= ~np.isinf(skymap) ipix = np.asarray(indices[msk], dtype=int) diff --git a/fast_response/FastResponseAnalysis.py b/fast_response/FastResponseAnalysis.py index 4c8be36..39b2f11 100644 --- a/fast_response/FastResponseAnalysis.py +++ b/fast_response/FastResponseAnalysis.py @@ -2,20 +2,23 @@ Author: Alex Pizzuto Date: 2021 - ''' +''' from abc import abstractmethod import os, sys, time, subprocess import pickle, dateutil.parser, logging, warnings +from pathlib import Path import h5py -import healpy as hp -import numpy as np -import seaborn as sns -import matplotlib as mpl -import matplotlib.pyplot as plt +import healpy as hp +import numpy as np +import seaborn as sns +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy.lib.recfunctions as rf from astropy.time import Time from scipy.special import erfinv +from scipy import sparse from matplotlib.lines import Line2D from skylab.datasets import Datasets @@ -57,13 +60,30 @@ class FastResponseAnalysis(object): _season_names = [f"IC86, 201{y}" for y in range(1, 10)] _nb_days = 10 _ncpu = 5 - - def __init__(self, name, tstart, tstop, skipped=None, seed=None, - outdir=None, save=True, extension=None): + _jitter = False + _background_days = 6 # if not using archival data, get_data() will prepend this duration to load + + def __init__(self, name, tstart, tstop, + skipped=None, seed=None, + outdir=None, save=True, + extension=None, + index=None, + fix_index=None, + dataset=None, + ): + logging.debug('FastResponseAnalysis.__init__') self.name = name + + if index is not None: + self._index = float(index) + if dataset is not None: + self._dataset = dataset + if fix_index is not None: + self._fix_index = fix_index + self._float_index = not self._fix_index if seed is not None: - self.llh_seed(seed) + self.llh_seed = seed if outdir is None: outdir = os.environ.get('FAST_RESPONSE_OUTPUT') if outdir is None: @@ -94,6 +114,7 @@ def __init__(self, name, tstart, tstop, skipped=None, seed=None, # sys.exit() elif self.save_output: subprocess.call(['mkdir', self.analysispath]) + #os.makedirs(self.analysispath, exist_ok=False) if creating parent directories is ok if 'test' in self.name.lower(): self.scramble = True @@ -143,6 +164,39 @@ def llh_seed(self): def llh_seed(self, x): self._llh_seed = x + def unify_exp_array(self, _exp, enum=0): + """Turns a rec array into same fields, order and precision. + """ + merged_dtype = np.dtype([('run', ' 58933.0: + # 1a) legacy behaviour, these are added to GFU archival + # TODO this relies on the assumption the archival GFU is equal to the default + if (self.stop > 58933.0) and self.dataset.startswith('GFUOnline_v001p02'): # Add local 2020 if need be # TODO: Need to figure out what to do for zenith_smoothed exp_new = np.load( @@ -184,13 +252,18 @@ def get_data(self, livestream_start=None, livestream_stop=None): self._floor) exps.append(exp_new) grls.append(grl) + # concatenate the rest exp = np.concatenate(exps) grl = np.concatenate(grls) + # TODO discuss: for new analyses, can replace with new GFU? + # TODO discuss: can replace this v001p02 + /data/user/ with a Skylab dataset definition? + # 2) use the livestream method to grab fresh events else: if self._verbose: print("Recent time: querying the i3live database") + # (default) retrieve a fixed off-time window before the analysis window if livestream_start is None or livestream_stop is None: - livestream_start = self.start - 6. + livestream_start = self.start - self._background_days livestream_stop = self.stop exp, mc, livetime, grl = dset.livestream( livestream_start, livestream_stop, @@ -200,8 +273,13 @@ def get_data(self, livestream_start=None, livestream_stop=None): grl.sort(order='run') livetime = grl['livetime'].sum() - sinDec_bins = dset.sinDec_bins("livestream") - energy_bins = dset.energy_bins("livestream") + # workaround while not all datasets have livestream yet + if 'livestream' in dset.season_names(): + reference_season = 'livestream' + else: + reference_season = self._season_names[0] + sinDec_bins = dset.sinDec_bins(reference_season) + energy_bins = dset.energy_bins(reference_season) self.exp = exp self.mc = mc @@ -273,6 +351,7 @@ def initialize_llh(self, skipped=None, scramble=False): ncpu=self._ncpu, # use 10 CPUs when computing trials scramble=scramble, # set to False for unblinding timescramble=True, # not just RA scrambling + jitter=self._jitter, # depends on sample llh_model=llh_model, # likelihood model temporal_model=box, # use box for temporal model nsource_bounds=(0., 1e3), # bounds on fitted ns @@ -507,14 +586,14 @@ def plot_ontime(self, with_contour=False, contour_files=None, label_events=False try: self.plot_skymap_zoom(with_contour=with_contour, contour_files=contour_files) except Exception as e: - print('Failed to make skymap zoom plot') + print(f'Failed to make skymap zoom plot: {e}') try: self.plot_skymap(with_contour=with_contour, contour_files=contour_files, label_events=label_events) except Exception as e: print('Failed to make FULL skymap plot') - def plot_skymap_zoom(self, with_contour=False, contour_files=None): + def plot_skymap_zoom(self, with_contour=False, contour_files=None, reso=3.): r"""Make a zoomed in portion of a skymap with all ontime neutrino events within a certain range Outputs a plot (in png and pdf formats) to the analysis path @@ -525,18 +604,22 @@ def plot_skymap_zoom(self, with_contour=False, contour_files=None): plots the 90% containment contour of a skymap (default False) contour_files: string text file containing skymap contours to be plotted (default None) + reso: float or 'auto' + resolution to zoom in, degrees (default 3.0) """ - events = self.llh.exp + + events = self.llh_exp # TODO ask if they prefer alternative: + # events as optional kwarg + # separate out a method that takes events of one sample + # and the rest of the plotting + # then this method needs to know nothing of enum events = events[(events['time'] < self.stop) & (events['time'] > self.start)] - col_num = 5000 - seq_palette = sns.color_palette("icefire", col_num) - lscmap = mpl.colors.ListedColormap(seq_palette) - - rel_t = np.array((events['time'] - self.start) * col_num / (self.stop - self.start), dtype = int) - cols = np.array([seq_palette[j] for j in rel_t]) + if reso == 'auto': + reso = plotting_utils.auto_reso(events) + # plot skymap if given: if self.skymap is not None: skymap = self.skymap ra = self.skymap_fit_ra @@ -547,32 +630,56 @@ def plot_skymap_zoom(self, with_contour=False, contour_files=None): skymap = np.zeros(hp.nside2npix(self._nside)) ra = self.ra dec = self.dec - label_str = self.name + label_str = self.name.replace('_', ' ') cmap = mpl.colors.ListedColormap([(1.,1.,1.)] * 50) - plotting_utils.plot_zoom(skymap, ra, dec, "", range = [0,10], reso=3., cmap = cmap) + plotting_utils.plot_zoom(skymap, ra, dec, "", range = [0,10], reso=reso, cmap = cmap) + + # remove skipped event: if self.skipped is not None: try: msk = events['run'] == int(self.skipped[0][0]) msk *= events['event'] == int(self.skipped[0][1]) + # TODO here we don't know which sample the skipped event came from... + # TODO and strictly speaking, supply subevent and check that too plotting_utils.plot_events(self.skipped_event['dec'], self.skipped_event['ra'], self.skipped_event['sigma']*self._angScale, ra, dec, 2*6, sigma_scale=1.0, constant_sigma=False, same_marker=True, energy_size=True, col = 'grey', with_dash=True) events = events[~msk] - cols = cols[~msk] except: print("Removed event not in GFU") + # TODO print the actual dataset name + # (once one implemention Multi code is merged) if (self.stop - self.start) <= 21.: - plotting_utils.plot_events(events['dec'], events['ra'], events['sigma']*self._angScale, ra, dec, 2*6, sigma_scale=1.0, - constant_sigma=False, same_marker=True, energy_size=True, col = cols) + sigma_scale = 1.0 else: #Long time windows means don't plot contours - plotting_utils.plot_events(events['dec'], events['ra'], events['sigma']*self._angScale, ra, dec, 2*6, sigma_scale=None, - constant_sigma=False, same_marker=True, energy_size=True, col = cols) + sigma_scale = None + # FIXME this is ignored? + + for enum in np.unique(events['enum']): # not adding pandas as dependency + _mask = events['enum'] == enum + _events = events[_mask] + _style = plotting_utils.skymap_style[enum] + _cols = cols[_mask] + if self._verbose: + print(f'Found {_events.size} on-time events from {self.datasets[enum]}') + plotting_utils.plot_events(_events['dec'], _events['ra'], _events['sigma']*self._angScale, + ra, dec, 2*6, # this reso positional arg is not used + sigma_scale=reso/3., + constant_sigma=False, same_marker=True, energy_size=True, + col = _cols, + kw_style=plotting_utils.skymap_style[enum], + ) + + # plotting_utils.plot_events(events['dec'], events['ra'], events['sigma']*self._angScale, + # ra, dec, 2*6, + # sigma_scale=sigma_scale, + # constant_sigma=False, same_marker=True, energy_size=True, col = cols) if contour_files is not None: cont_ls = ['solid', 'dashed'] @@ -581,7 +688,7 @@ def plot_skymap_zoom(self, with_contour=False, contour_files=None): cont = np.loadtxt(c_file, skiprows=1) cont_ra = cont.T[0] cont_dec = cont.T[1] - label = 'Millipede 50\%, 90\% (160427A syst.)' \ + label = 'Millipede 50%, 90% (160427A syst.)' \ if contour_counter == 0 else '' hp.projplot(np.pi/2. - cont_dec, cont_ra, linewidth=3., color='k', linestyle=cont_ls[contour_counter], coord='C', @@ -607,7 +714,10 @@ def plot_skymap_zoom(self, with_contour=False, contour_files=None): plt.savefig(self.analysispath + '/' + self.analysisid + 'unblinded_skymap_zoom.pdf',bbox_inches='tight', dpi=300) plt.close() - def plot_skymap(self, with_contour=False, contour_files=None, label_events=False): + def plot_skymap(self, with_contour=False, contour_files=None, label_events=False, + labels=['GFU Event'], distinct_colorbars=False, + show=False, + ): r""" Make skymap with event localization and all neutrino events on the sky within the given time window Outputs a plot in png format to the analysis path @@ -623,21 +733,20 @@ def plot_skymap(self, with_contour=False, contour_files=None, label_events=False """ - events = self.llh.exp + events = self.llh_exp events = events[(events['time'] < self.stop) & (events['time'] > self.start)] - col_num = 5000 - seq_palette = sns.color_palette("icefire", col_num) - lscmap = mpl.colors.ListedColormap(seq_palette) - - rel_t = np.array((events['time'] - self.start) * col_num / (self.stop - self.start), dtype = int) - cols = [seq_palette[j] for j in rel_t] - # Set color map and plot skymap pdf_palette = sns.color_palette("Blues", 500) cmap = mpl.colors.ListedColormap(pdf_palette) cmap.set_under("w") + # Obtain color maps for event times + if distinct_colorbars: + tcmap = plotting_utils.TimeColormap(self.start, self.stop, n_maps=events['enum'].max()+1) + else: + tcmap = plotting_utils.TimeColormap(self.start, self.stop, n_maps=1) + if self.skymap is None: skymap = np.zeros(hp.nside2npix(self._nside)) max_val = 1. @@ -682,25 +791,42 @@ def plot_skymap(self, with_contour=False, contour_files=None, label_events=False # plot events on sky with error contours handles=[] - hp.projscatter(theta,phi,c=cols,marker='x',label='GFU Event',coord='C', zorder=5) + # TODO is this implementation ok? + # or re-factor this method, so it can be used from MultiFRA, knowing about the samples before? + for enum in np.unique(events['enum']): + _mask = events['enum'] == enum + _style = plotting_utils.skymap_style[enum] + _label = labels[enum] + hp.projscatter(theta[_mask], phi[_mask], + c=tcmap(events['time'][_mask], enum), + marker=_style['marker'], + label=_label, + s=128, + coord='C', zorder=5) + handles.append(Line2D([0], [0], marker=_style['marker'], ls='None', label=_label)) + if label_events: for j in range(len(theta)): hp.projtext(theta[j], phi[j]-0.11, '{}'.format(j+1), color='red', fontsize=18, zorder=6) - handles.append(Line2D([0], [0], marker='x', ls='None', label='GFU Event')) - + if (self.stop - self.start) <= 0.5: #Only plot contours if less than 2 days for i in range(events['ra'].size): + _enum = events['enum'][i] + _style = plotting_utils.skymap_style[_enum] + _col = tcmap(events['time'][i], _enum)[0] my_contour = plotting_utils.contour(events['ra'][i], events['dec'][i],sigma_90[i], self._nside) hp.projplot(my_contour[0], my_contour[1], linewidth=2., - color=cols[i], linestyle="solid",coord='C', zorder=5) + color=_col, linestyle=_style['linestyle'], + coord='C', zorder=5) if self.skymap is None: + label_str = self.name.replace('_', ' ') src_theta = np.pi/2. - self.dec src_phi = self.ra hp.projscatter(src_theta, src_phi, c = 'k', marker = '*', - label = self.name, coord='C', s=350) - handles.append(Line2D([0], [0], marker='*', c='k', ls='None', label=self.name)) + label = label_str, coord='C', s=350) + handles.append(Line2D([0], [0], marker='*', c='k', ls='None', label=label_str)) if contour_files is not None: cont_ls = ['solid', 'dashed'] @@ -719,8 +845,8 @@ def plot_skymap(self, with_contour=False, contour_files=None, label_events=False ### plot 90% containment contour of PDF levels = [0.9] theta, phi = plotting_utils.plot_contours(levels, probs) - hp.projplot(theta[0], phi[0], linewidth=2., c='k', label='Skymap (90\% cont.)') - handles.append(Line2D([0], [0], lw=2, c='k', label=r"Skymap (90\% cont.)")) + hp.projplot(theta[0], phi[0], linewidth=2., c='k', label='Skymap (90% cont.)') + handles.append(Line2D([0], [0], lw=2, c='k', label=r"Skymap (90% cont.)")) for i in range(1, len(theta)): hp.projplot(theta[i], phi[i], linewidth=2., c='k') @@ -731,7 +857,8 @@ def plot_skymap(self, with_contour=False, contour_files=None, label_events=False except: plt.title('Fast Response Skymap') plt.savefig(self.analysispath + '/' + self.analysisid + 'unblinded_skymap.png',bbox_inches='tight') - plt.close() + if not show: + plt.close() def generate_report(self): r"""Generates report using class attributes @@ -751,10 +878,19 @@ class PriorFollowup(FastResponseAnalysis): _containment = 0.99 _allow_neg = False _nside = 256 + _bg_dir = './' + _bg_format = '_'.join([ + 'precomputed_trials_delta_t_{delta_t:.2e}', + 'nside_{nside}', + 'index_{index}', + '{lookup}', + ]) def __init__(self, name, skymap_path, tstart, tstop, skipped=None, seed=None, outdir=None, save=True, extension=None): + logging.debug('PriorFollowup.__init__') + super().__init__(name, tstart, tstop, skipped=skipped, seed=seed, outdir=outdir, save=save, extension=extension) @@ -872,13 +1008,80 @@ def run_background_trials(self, ntrials=1000): self.tsd = tsd self.save_items['tsd'] = tsd - def find_coincident_events(self): + def load_background_trials(self, ntrials=None, rate=None, month=None): + """Produce background trials based on precomputed all-sky scans + stored in sparse matrices produced by fast_response/precomputed_background/... + precompute_ts.py (or its variants) + glob_precomputed_trials.py (or its variants) + + + Parameters + ---------- + ntrials : int, optional + Number of trials to return, by default return as many as available. + rate : float, optional + rate in mHz to look up + month : int, optional + month to look up + + Raises + ------ + TypeError + if neither month nor rate are supplied + """ + if not ((rate is None) ^ (month is None)): + raise TypeError("Need to supply either rate or month") + + # Assemble variables for the background file + filename = self._bg_format.format( + delta_t = self.duration * 86400., + nside = self.nside, + index = self._index, + lookup = f"{rate:.2f}_mHz" if month is None else f"{month:02d}", + ) + bg_files = list(Path(self._bg_dir).glob(filename)) + if not bg_files: + raise FileNotFoundError(f"Did not find precomputed bg {filename} in {self._bg_dir}") + # TODO also concatenate in here; then the glob script becomes superfluous + # Load sparse matrix of background scans + pre_ts_array = sparse.load_npz(bg_files[0]) + if hp.npix2nside(pre_ts_array.shape[1]) != self.nside: + # Should be ensured by file name but better check + raise ValueError(f"Loaded precomputed bg has nside != {self.nside}") + # Combine with prior as in GWFollowup + ts_prior = pre_ts_array.copy() + ts_norm = np.log(np.amax(self.skymap)) + # skymap was already reduced to the analysis nside upon loading + # FIXME reduction vs. interpolation?? + # TODO better way than to introduce inf's by log-ging the skymap? + ts_prior.data += 2.*(np.log(self.skymap[pre_ts_array.indices]) - ts_norm) + ts_prior.data[~np.isfinite(ts_prior.data)] = 0. # TODO discuss whether this applies. Not sure why inconsistent. + ts_prior.data[ts_prior.data < 0] = 0. + # Take the maximum per entry + tsd = ts_prior.max(axis=1).toarray()[:,0] + # Explicitly skip + empty = np.array([_ts.size==0 for _ts in pre_ts_array]) + tsd[empty] = -np.inf # FIXME why does it set these to -inf instead of 0? + self.tsd = tsd + if ntrials is None: + return tsd + elif ntrials <= self.tsd.size: + return tsd[:ntrials] + else: + raise ValueError(f"Could not load {ntrials} precomputed trials, only have {self.tsd.size}") + + + + + + def find_coincident_events(self, exp=None): r"""Find coincident events for a skymap based analysis. These are ontime events that are also in the 90% contour of the skymap """ - t_mask=(self.llh.exp['time']<=self.stop)&(self.llh.exp['time']>=self.start) - events = self.llh.exp[t_mask] + t_mask=(self.llh_exp['time']<=self.stop)&(self.llh_exp['time']>=self.start) + events = self.llh_exp[t_mask] + # Using the llh_exp property -- TODO get feedback if they mind exp_theta = 0.5*np.pi - events['dec'] exp_phi = events['ra'] exp_pix = hp.ang2pix(self.nside, exp_theta, exp_phi) @@ -1085,7 +1288,7 @@ def make_dNdE(self): low_5_min_dec = np.interp(0.05, cdf, a[1][:-1]) median_min_dec = np.interp(0.5, cdf, a[1][:-1]) high_5_min_dec = np.interp(0.95, cdf, a[1][:-1]) - plt.axvspan(low_5_min_dec, high_5_min_dec, color = sns.xkcd_rgb['windows blue'], alpha = 0.25, label="Central 90\%") + plt.axvspan(low_5_min_dec, high_5_min_dec, color = sns.xkcd_rgb['windows blue'], alpha = 0.25, label="Central 90%") lab = 'Median (min dec.)' plt.axvline(median_min_dec, c = sns.xkcd_rgb['windows blue'], alpha = 0.75, label = lab) @@ -1123,10 +1326,12 @@ class PointSourceFollowup(FastResponseAnalysis): Class for point-source or extended source followup i.e. there is a fixed location on the sky, not a healpy skymap """ + logging.debug('PointSourceFollowup.__init__') + _nside = 256 def __init__(self, name, ra, dec, tstart, tstop, extension=None, skipped=None, outdir=None, save=True, seed=None): - + logging.debug('PointSourceFollowup.__init__') super().__init__(name, tstart, tstop, skipped=skipped, seed=seed, outdir=outdir, save=save, extension=extension) @@ -1212,21 +1417,47 @@ def unblind_TS(self): ns = ns['nsignal'] if self._verbose: print("TS = {}".format(ts)) - print("ns = {}\n\n".format(ns)) + print("ns = {}".format(ns)) + for par, val in params.items(): + if isinstance(val, float): + print(f"{par} = {val:.3f}") + else: + print(f"{par} = {val}") + print("\n\n") self.ts, self.ns = ts, ns self.save_items['ts'] = ts self.save_items['ns'] = ns + # need gamma for report + # TODO alternatively change ReportGenerator to report any ns_params + # (if they are not a fixed spectrum) + if 'gamma' in params: + self.gamma = params['gamma'] + # save all parameters besides nsignal + # (can be other spectral models) + for par, val in params.items(): + if par in self.save_items: + if self._verbose: + print(f'Warning, not saving {par} as save_items already has such a key') + self.save_items.setdefault(par, val) + return ts, ns - def find_coincident_events(self): + def find_coincident_events(self, ns_params=None): r"""Find "coincident events" for the analysis. These are ontime events that have a spatial times energy weight greater than 10 + + Parameters + ----------- + ns_params: dict + Fit parameters to use for weight calculation, e.g. if fit happened in MultiPointSourceFollowup """ + if ns_params is None: + ns_params = self.ns_params spatial_weights = self.llh.llh_model.signal( self.ra, self.dec, self.llh._events, src_extension=self.extension)[0] / self.llh._events['B'] energy_ratio, _ = self.llh.llh_model.weight( - self.llh._events, **self.ns_params) + self.llh._events, **ns_params) temporal_weights = self.llh.temporal_model.signal(self.llh._events) msk = spatial_weights * energy_ratio * temporal_weights > 10 self.coincident_events = [] @@ -1331,6 +1562,8 @@ def upper_limit(self, n_per_sig=100, p0=None): fits[best_fit_ind]['ls'] = '-' self.upperlimit = self.inj.mu2flux(fits[best_fit_ind]['sens']) self.upperlimit_ninj = fits[best_fit_ind]['sens'] + # E^2 dN/dE at self.inj.E0 + upperlimit_fluence = self.upperlimit * self.duration * 86400. * self.inj.E0**2 fig, ax = plt.subplots() for fit_dict in fits: @@ -1344,9 +1577,14 @@ def upper_limit(self, n_per_sig=100, p0=None): if fit_dict['ls'] == '-': ax.axhline(0.9, color = 'm', linewidth = 0.3, linestyle = '-.') ax.axvline(fit_dict['sens'], color = 'm', linewidth = 0.3, linestyle = '-.') - ax.text(3.5, 0.8, r'Sens. = {:.2f} events'.format(fit_dict['sens']), fontsize = 16) - ax.text(3.5, 0.7, r' = {:.1e}'.format(self.upperlimit * self.duration * 86400. * 1e6) + r' GeV cm$^{-2}$', fontsize = 16) - #ax.text(5, 0.5, r'Sens. = {:.2e}'.format(self.inj.mu2flux(fit_dict['sens'])) + ' GeV^-1 cm^-2 s^-1') + limit_annotation = r'Sens. = {:.2f} events'.format(fit_dict['sens']) + '\n' + limit_annotation += r' = {:.1e}'.format(upperlimit_fluence) + r' GeV cm$^{-2}$' + '\n' + if self.index != 2: + # E^2 F not constant in energy, state pivot energy in label + limit_annotation += f'at {self.inj.E0:.0f} GeV' + ax.annotate(limit_annotation, + (3.5, 0.8), ha = 'left', va = 'top', xycoords = 'data', + fontsize = 16) ax.errorbar(signal_fluxes, passing, yerr=errs, capsize = 3, linestyle='', marker = 's', markersize = 2) ax.legend(loc=4, fontsize = 14) ax.set_xlabel(r'$\langle n_{inj} \rangle$', fontsize = 14) @@ -1384,7 +1622,7 @@ def make_dNdE(self): high_5 = np.interp(0.95, cdf, a[1][:-1]) self.low5 = low_5 self.high5 = high_5 - plt.axvspan(low_5, high_5, color = sns.xkcd_rgb['windows blue'], alpha = 0.25, label="Central 90\%") + plt.axvspan(low_5, high_5, color = sns.xkcd_rgb['windows blue'], alpha = 0.25, label="Central 90%") lab = 'Median' plt.axvline(median, c = sns.xkcd_rgb['windows blue'], alpha = 0.75, label = lab) plt.xlim(1e1, 1e8) diff --git a/fast_response/FastResponseAnalysis_ifelse.py b/fast_response/FastResponseAnalysis_ifelse.py new file mode 100644 index 0000000..936703e --- /dev/null +++ b/fast_response/FastResponseAnalysis_ifelse.py @@ -0,0 +1,1291 @@ +r''' General Fast Response Analysis Class. + + Author: Alex Pizzuto + Date: 2021 + + One of two variants to include mutiple datasets + ''' + +from abc import abstractmethod +import os, sys, time, subprocess +import pickle, dateutil.parser, logging, warnings +from argparse import Namespace + +import h5py +import healpy as hp +import numpy as np +import seaborn as sns +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy.lib.recfunctions as rf +from astropy.time import Time +from scipy.special import erfinv +from matplotlib.lines import Line2D + +from skylab.datasets import Datasets +from skylab.llh_models import EnergyLLH +from skylab.priors import SpatialPrior +from skylab.ps_injector import PointSourceInjector +from skylab.ps_llh import PointSourceLLH, MultiPointSourceLLH +from skylab.ps_injector import PriorInjector +from skylab.spectral_models import PowerLaw +from skylab.temporal_models import BoxProfile, TemporalModel +# import meander + +from . import web_utils +from . import sensitivity_utils +from . import plotting_utils +from .reports import FastResponseReport + +mpl.use('agg') +current_palette = sns.color_palette('colorblind', 10) +logging.getLogger().setLevel(logging.ERROR) +warnings.simplefilter("ignore", UserWarning) +warnings.simplefilter("ignore", RuntimeWarning) + +class FastResponseAnalysis(object): + """ + Object to do realtime followup analyses of + astrophysical transients with arbitrary event + localization + """ + _dataset = None + _fix_index = True + _float_index = not _fix_index + _index_range = [1., 4.] + _index = 2.0 + _floor = np.radians(0.2) + _verbose = True + _angScale = 2.145966 + _llh_seed = 1 + _season_names = [f"IC86, 20{y:02d}" for y in range(11, 22+1)] + _nb_days = 10 + _ncpu = 5 + + def __init__(self, name, tstart, tstop, skipped=None, seed=None, + outdir=None, save=True, extension=None): + self.name = name + self.multi = type(self.dataset) is list + + if seed is not None: + self.llh_seed(seed) + if outdir is None: + outdir = os.environ.get('FAST_RESPONSE_OUTPUT') + if outdir is None: + outdir = os.getcwd() + self.outdir = outdir + + start = Time(dateutil.parser.parse(tstart)).mjd + stop = Time(dateutil.parser.parse(tstop)).mjd + + start_date = Time(dateutil.parser.parse(tstart)).datetime + start_str = f'{start_date.year:02d}_' \ + + f'{start_date.month:02d}_{start_date.day:02d}' + + self.start = start + self.stop = stop + self.duration = stop - start + self.centertime = (start + stop) / 2. + + self.analysisid = start_str + '_' + self.name.replace(' ', '_') + self.analysispath = os.path.join(self.outdir, self.analysisid) + self.save_output = save + + if os.path.isdir(self.analysispath) and self.save_output: + print("Directory {} already exists. Deleting it ...".format( + self.analysispath)) + subprocess.call(['rm', '-r', self.analysispath]) + subprocess.call(['mkdir', self.analysispath]) + # sys.exit() + elif self.save_output: + subprocess.call(['mkdir', self.analysispath]) + #os.makedirs(self.analysispath, exist_ok=False) if creating parent directories is ok + + if 'test' in self.name.lower(): + self.scramble = True + if self._verbose: + print('Working on scrambled data') + else: + self.scramble = False + if self._verbose: + print('Working on unscrambled (UNBLINDED) data') + + self.save_items = { + 'start':self.start, 'stop': self.stop, + 'name': self.name, 'analysisid': self.analysisid, + 'analysispath': self.analysispath + } + self.extension = extension + if self.extension is not None: + self.extension = np.deg2rad(self.extension) + self.llh_seed = seed if seed is not None else 1 + self.skipped = skipped + self.skipped_event = None + self.exp = None + self.spectrum = None + if self._fix_index: + self.spectrum = PowerLaw(A=1, gamma=self.index, E0=1000.) + self.time_profile = BoxProfile(self.start, self.stop) + self.llh = self.initialize_llh(skipped=skipped, scramble=self.scramble) + self.inj = None + + @property + def dataset(self): + """Returns the datasets used""" + return self._dataset + @dataset.setter + def dataset(self, x): + self._dataset = x + + @property + def datasets(self): + '''Dataset(s) used as a list''' + if isinstance(self.dataset, list): + return self.dataset + else: + return [self.dataset] + + @property + def index(self): + """Returns the spectral index""" + return self._index + @index.setter + def index(self, x): + self._index = x + + @property + def llh_seed(self): + """Returns the seed used in the LLH""" + return self._llh_seed + @llh_seed.setter + def llh_seed(self, x): + self._llh_seed = x + + @property + def llh_exp(self): + """Returns a flat array of experimental data loaded across the used dataset(s) loaded into the LLH, + limited to fields used for plotting and amended with a field for the dataset.""" + # NOTE can not use self.dset_container as that is before self.remove_event(skipped=skipped) + # TODO make this more efficient? + if self.multi: + exp = {} + for enum, _llh in self.llh._samples.items(): + exp[enum] = _llh.exp + else: + exp = {-1: self.llh.exp} + + merged_dtype = np.dtype([('run', ' 58933.0) and dataset.startswith('GFUOnline'): + if self._verbose: + print('adding 2020 GFU from /data/user/apizzuto') + # Add local 2020 if need be + # TODO: Need to figure out what to do for zenith_smoothed + exp_new = np.load( + '/data/user/apizzuto/fast_response_skylab/fast-response/' + + 'fast_response/2020_data/IC86_2020_data.npy') + grl = np.load( + '/data/user/apizzuto/fast_response_skylab/fast-response/' + + 'fast_response/2020_data/GRL/IC86_2020_data.npy') + exp_new.dtype.names = exp.dtype.names + exp_new['sigma'] = np.where( + exp_new['sigma'] > self._floor, + exp_new['sigma'], + self._floor) + exps.append(exp_new) + grls.append(grl) + exp = np.concatenate(exps) + grl = np.concatenate(grls) + # TODO add livestream season to Greco dataset - big task + else: + if self._verbose: + print("Recent time: querying the i3live database") + if livestream_start is None or livestream_stop is None: + livestream_start = self.start - 6. + livestream_stop = self.stop + exp, mc, livetime, grl = dset.livestream( + livestream_start, livestream_stop, + append=archival_seasons, + floor=self._floor) + exp.sort(order='time') + grl.sort(order='run') + livetime = grl['livetime'].sum() + # workaround while not all datasets have livestream yet + if 'livestream' in dset.season_names(): + reference_season = 'livestream' + else: + reference_season = archival_seasons[0] + sinDec_bins = dset.sinDec_bins(reference_season) + energy_bins = dset.energy_bins(reference_season) + + dset_container.exp = exp + dset_container.mc = mc + dset_container.grl = grl + dset_container.livetime = livetime + dset_container.dset = dset + dset_container.sinDec_bins = sinDec_bins + dset_container.energy_bins = energy_bins + return dset_container + + def get_data(self, **kwargs): + """ + Gets the skylab data and MC from querying the i3live livestream + + Parameters + ----------- + livestream_start: float + (optional) start time for getting data (MJD) from i3Live. Default is start time of the analysis - 5 days + livestream_stop: float + (optional) stop time for getting data (MJD). Default is stop time of the analysis + """ + if self.multi: + dset_multi = {enum:self.get_data_single(_dataset) for enum, _dataset in enumerate(self.dataset)} + self.dset_container = dset_multi + # pivot from dictionary of namespaces + # to dictionaries in the class namespace + for enum, dset in dset_multi.items(): + for attr in vars(dset): + if not hasattr(self, attr) or getattr(self, attr) is None: + setattr(self, attr, {}) + getattr(self, attr)[enum] = getattr(dset, attr) + else: + dset = self.get_data_single(self.dataset, **kwargs) + self.dset_container = {-1: dset} + for attr in vars(dset): + setattr(self, attr, getattr(dset, attr)) + + + + + def initialize_llh(self, skipped=None, scramble=False): + """ + Grab data and format it all into a skylab llh object + + Parameters + ----------- + skipped: array of tuples + (optional) event(s) to be removed in the analysis. + Format: [(run_id,event_id),...] + scramble: bool + (optional) run on scrambled data (default=False) + + Returns + ---------- + llh: skylab llh object + Likelihood for the analysis + """ + if self.exp is None: + self.get_data() + + if self._verbose: + print("Initializing Point Source LLH in Skylab") + + assert self._fix_index != self._float_index,\ + 'Must choose to either float or fix the index' + + llh_list = [] + # kwargs to be broadcast to the BaseLLH constructor + # not specific to dataset + base_kwargs = dict( + nsource=1., # seed for nsignal fit + nsource_bounds=(0., 1e3), # bounds on fitted ns + ncpu=self._ncpu, # use 10 CPUs when computing trials + seed=self.llh_seed, + ) + + for enum in self.dset_container: + dset = self.dset_container[enum] + twodim_bins = [dset.energy_bins, dset.sinDec_bins] + if self._fix_index: + llh_model = EnergyLLH( + twodim_bins=twodim_bins, + allow_empty=True, + spectrum=self.spectrum, + ncpu=self._ncpu) + elif self._float_index: + llh_model = EnergyLLH( + twodim_bins=twodim_bins, + allow_empty=True, + bounds=self._index_range, + seed = self.index, + ncpu=self._ncpu) + + box = TemporalModel( + grl=dset.grl, + poisson_llh=True, + days=self._nb_days, + signal=self.time_profile) + + if skipped is not None: + dset.exp = self.remove_event(dset.exp, dset.dset, skipped) + + if self.extension is not None: + src_extension = float(self.extension) + self.save_items['extension'] = src_extension + else: + src_extension = None + + llh = PointSourceLLH( + dset.exp, # array with data events + dset.mc, # array with Monte Carlo events + dset.livetime, # total livetime of the data events + scramble=scramble, # set to False for unblinding + timescramble=True, # not just RA scrambling + llh_model=llh_model, # likelihood model + temporal_model=box, # use box for temporal model + src_extension = src_extension, # Model symmetric extended source + **base_kwargs, + ) + llh_list.append(llh) + + if self.multi: + multi_llh = MultiPointSourceLLH(**base_kwargs) + for name, llh in zip(self.dataset, llh_list): + llh.do_trials_seed = 1 + multi_llh.add_sample(name, llh) + return multi_llh + + elif len(llh_list)==0: + return llh_list[0] + + else: + raise AssertionError('have multiple instances of PointSourceLLH but FRA is not multi') + + def remove_event(self, exp, dset, skipped): + """ + Remove a given event from the analysis, eg. for an alert event + + Parameters + ----------- + exp: skylab data object + Data events loaded from livestream + dset: skylab Dataset object + Dataset used in the analysis + skipped: array of tuples + Event(s) to be attempted to be removed in the analysis. + Format: [(run_id,event_id),...] + + Returns + ---------- + exp: skylab data object + Data events, minus the removed event (if successful) + """ + try: + event = skipped[0] + mjd_keys = exp['time'][(exp['run'] == int(event[0])) & (exp['event'] == int(event[1]))] + self.skipped_event = exp[exp['time'] == mjd_keys[0]][0] + exp = dset.remove_ev(exp, mjd_keys=mjd_keys[0]) + self.save_items['skipped_event'] = self.skipped_event + except: + print("Was not able to remove event {}".format(skipped)) + return exp + + @abstractmethod + def initialize_injector(self): + pass + + @abstractmethod + def unblind_TS(self): + pass + + @abstractmethod + def run_background_trials(self, ntrials=1000): + pass + + def calc_pvalue(self, ntrials=1000, run_anyway=False): + r""" Given an unblinded TS value, calculate the p value + + Parameters + ----------- + ntrials: int + Number of trials to run (default 1000) + run_anyway: bool + Choose to override and run background trials, even if TS=0 (default False) + + Returns + -------- + p: float + P value for this analysis + """ + if self.ts is None: + if self._verbose: + print("Need TS value to find p value") + self.unblind_TS() + if self.ts == 0.0 and not run_anyway: + if self._verbose: + print("TS=0, no need to run background trials") + p = 1.0 + self.tsd = None + else: + self.run_background_trials(ntrials=ntrials) + p = np.count_nonzero(self.tsd >= self.ts) / float(self.tsd.size) + sigma = self.significance(p) + if self._verbose: + print(f"p: {p:.4f}") + print("Significance: {:.3f}sigma\n\n".format(sigma)) + self.p = p + self.save_items['p'] = p + return p + + @abstractmethod + def find_coincident_events(self): + pass + + @abstractmethod + def upper_limit(self): + pass + + @abstractmethod + def write_circular(self): + pass + + @abstractmethod + def make_dNdE(self): + pass + + @abstractmethod + def make_all_report_plots(self): + pass + + @abstractmethod + def make_all_report_tables(self): + pass + + def plot_tsd(self, allow_neg=False): + r""" + Outputs a plot of the background TS distributions + as well as the observed TS + + Parameters + ----------- + allow_neg: bool + Choose to allow negative TS values (all negative values are then TS=0). Default False + """ + if self.tsd is not None: + fig, ax = plt.subplots() + if allow_neg: + if str(np.min(self.tsd))== '-inf': + lower=-500. + else: + lower = np.max([np.min(self.tsd), -500.]) + + #lowest bin as underflow bin + self.tsd[self.tsd < lower] = lower + self.tsd[np.isinf(self.tsd)]= lower + + pos_bins=np.linspace(0., 25., 30) #fine pos bins + neg_bins=np.linspace(lower, 0., 30) #coarse neg bins + + bins=np.concatenate([ neg_bins[:-1], pos_bins ]) + else: + if (np.min(self.tsd) < 0.0) or (str(np.min(self.tsd))== '-inf'): + self.tsd[self.tsd < 0.] = 0. + self.tsd[np.isinf(self.tsd)]= 0. + bins = np.linspace(0., 25., 30) + + plt.hist(self.tsd, bins= bins, + label="Background Scrambles", density=True) + if self.ts >= -500.: + plt.axvline(self.ts, color = 'k', label = "Observed TS") + else: + plt.annotate("Unblinded TS: {:.0f}".format(self.ts), (lower+20, 1e-3)) + plt.legend(loc=1) + + if bins[0]<0.: + plt.xlabel("TS (Note: lowest bin is an underflow bin)") + else: + plt.xlabel("TS") + plt.ylabel("Probability Density") + plt.yscale('log') + plt.savefig( + self.analysispath + '/TS_distribution.png', + bbox_inches='tight') + + def __str__(self): + r"""Print a message with info about the source once + the analysis is running""" + int_str = '*'*80 + int_str += '\n*' + ' '*78 + '*\n' + int_str += '*' + ' '*((78-len(self.name))//2) + self.name + ' '*((78-len(self.name))//2 + len(self.name)%2) + '*' + int_str += '\n*' + ' '*78 + '*\n' + int_str += '*'*80 + '\n' + int_str += ' '*5 + str(Time(self.start, format = 'mjd', scale = 'utc').iso) + int_str += ' '*5 + str(Time(self.stop, format = 'mjd', scale = 'utc').iso) + '\n' + return int_str + + def significance(self, p): + r"""Given p value, report significance + + Parameters + ------------ + p: float + P value + + Returns + -------- + sigma: float + Significance + """ + # sigma = stats.norm.isf(p) + sigma = np.sqrt(2)*erfinv(1-2*p) + return sigma + + def save_results(self, alt_path=None): + r"""Once analysis has been performed, save + all relevant info to a new directory. + Saves to a file called [analysisid]_results.pickle + + Parameters + ----------- + alt_path: string + optional specified directory to save results to + + Returns + ----------- + save_items: dict + dictionary of saved analysis parameters + """ + if alt_path is None: + print("Saving results to directory:\n\t{}".format(self.analysispath)) + with open(self.analysispath + '/' + self.analysisid + '_' + 'results.pickle', 'wb') as f: + pickle.dump(self.save_items, f, protocol=pickle.HIGHEST_PROTOCOL) + print("Results successfully saved") + return self.save_items + else: + print("Saving to specified directory") + with open(alt_path + '/' + self.analysisid + '_' + 'results.pickle', 'wb') as f: + pickle.dump(self.save_items, f, protocol=pickle.HIGHEST_PROTOCOL) + print("Results successfully saved") + return self.save_items + + def plot_ontime(self, with_contour=False, contour_files=None, label_events=False): + r"""Plots ontime events on the full skymap and a + zoomed in version near the scan best-fit + + Parameters + ----------- + with_contour: bool + plots the 90% containment contour of a skymap (default False) + contour_files: string + text file containing skymap contours to be plotted (default None) + label_events: bool + adds a number label to events on skymap (default False) + + """ + + try: + self.plot_skymap_zoom(with_contour=with_contour, contour_files=contour_files) + except Exception as e: + print('Failed to make skymap zoom plot') + + try: + self.plot_skymap(with_contour=with_contour, contour_files=contour_files, label_events=label_events) + except Exception as e: + print('Failed to make FULL skymap plot') + + def plot_skymap_zoom(self, events=None, with_contour=False, contour_files=None): + r"""Make a zoomed in portion of a skymap with + all ontime neutrino events within a certain range + Outputs a plot (in png and pdf formats) to the analysis path + + Parameters + ------------ + with_contour: bool + plots the 90% containment contour of a skymap (default False) + contour_files: string + text file containing skymap contours to be plotted (default None) + + """ + if events is None: + events = self.llh_exp # flattened array of exp per sample + events = events[(events['time'] < self.stop) & (events['time'] > self.start)] + + col_num = 5000 + seq_palette = sns.color_palette("icefire", col_num) + lscmap = mpl.colors.ListedColormap(seq_palette) + + rel_t = np.array((events['time'] - self.start) * col_num / (self.stop - self.start), dtype = int) + cols = np.array([seq_palette[j] for j in rel_t]) + + if self.skymap is not None: + skymap = self.skymap + ra = self.skymap_fit_ra + dec = self.skymap_fit_dec + label_str = "Scan Hot Spot" + cmap = None + else: + skymap = np.zeros(hp.nside2npix(self._nside)) + ra = self.ra + dec = self.dec + label_str = self.name + cmap = mpl.colors.ListedColormap([(1.,1.,1.)] * 50) + + plotting_utils.plot_zoom(skymap, ra, dec, "", range = [0,10], reso=3., cmap = cmap) + # TODO change xlim, ylim via the gnomview tools to accommodate datasets with larger resolution + + if self.skipped is not None: + try: + msk = events['run'] == int(self.skipped[0][0]) + msk *= events['event'] == int(self.skipped[0][1]) + plotting_utils.plot_events(self.skipped_event['dec'], self.skipped_event['ra'], + self.skipped_event['sigma']*self._angScale, + ra, dec, 2*6, sigma_scale=1.0, constant_sigma=False, + same_marker=True, energy_size=True, col = 'grey', + with_dash=True) + events = events[~msk] + cols = cols[~msk] + except: + print("Removed event not in GFU") + + if (self.stop - self.start) <= 21.: + plotting_utils.plot_events(events['dec'], events['ra'], events['sigma']*self._angScale, ra, dec, 2*6, sigma_scale=1.0, + constant_sigma=False, same_marker=True, energy_size=True, col = cols) + else: + #Long time windows means don't plot contours + plotting_utils.plot_events(events['dec'], events['ra'], events['sigma']*self._angScale, ra, dec, 2*6, sigma_scale=None, + constant_sigma=False, same_marker=True, energy_size=True, col = cols) + + if contour_files is not None: + cont_ls = ['solid', 'dashed'] + contour_counter = 0 + for c_file in contour_files: + cont = np.loadtxt(c_file, skiprows=1) + cont_ra = cont.T[0] + cont_dec = cont.T[1] + label = 'Millipede 50\%, 90\% (160427A syst.)' \ + if contour_counter == 0 else '' + hp.projplot(np.pi/2. - cont_dec, cont_ra, linewidth=3., + color='k', linestyle=cont_ls[contour_counter], coord='C', + label=label) + contour_counter += 1 + + if with_contour: + probs = hp.pixelfunc.ud_grade(self.skymap, 64, power=-2) + probs = probs/np.sum(probs) + ### plot 90% containment contour of PDF + levels = [0.9] + theta, phi = plotting_utils.plot_contours(levels, probs) + hp.projplot(theta[0], phi[0], linewidth=2., c='k') + for i in range(1, len(theta)): + hp.projplot(theta[i], phi[i], linewidth=2., c='k', label=None) + + plt.scatter(0,0, marker='*', c = 'k', s = 130, label = label_str) + + plt.legend(loc = 2, ncol=1, mode = 'expand', fontsize = 18.5, framealpha = 0.95) + plotting_utils.plot_color_bar(range=[0,6], cmap=lscmap, col_label=r"IceCube Event Time", + offset=-50, labels = [r'-$\Delta T \Bigg/ 2$', r'+$\Delta T \Bigg/ 2$']) + plt.savefig(self.analysispath + '/' + self.analysisid + 'unblinded_skymap_zoom.png',bbox_inches='tight') + plt.savefig(self.analysispath + '/' + self.analysisid + 'unblinded_skymap_zoom.pdf',bbox_inches='tight', dpi=300) + plt.close() + + def plot_skymap(self, events = None, + with_contour=False, contour_files=None, label_events=False): + r""" Make skymap with event localization and all + neutrino events on the sky within the given time window + Outputs a plot in png format to the analysis path + + Parameters + ----------- + with_contour: bool + plots the 90% containment contour of a skymap (default False) + contour_files: string + text file containing skymap contours to be plotted (default None) + label_events: bool + adds a number label to events on skymap (default False) + + """ + if events is None: + events = self.llh_exp + events = events[(events['time'] < self.stop) & (events['time'] > self.start)] + + col_num = 5000 + seq_palette = sns.color_palette("icefire", col_num) + lscmap = mpl.colors.ListedColormap(seq_palette) + + rel_t = np.array((events['time'] - self.start) * col_num / (self.stop - self.start), dtype = int) + cols = [seq_palette[j] for j in rel_t] + + # Set color map and plot skymap + pdf_palette = sns.color_palette("Blues", 500) + cmap = mpl.colors.ListedColormap(pdf_palette) + cmap.set_under("w") + + if self.skymap is None: + skymap = np.zeros(hp.nside2npix(self._nside)) + max_val = 1. + cmap = mpl.colors.ListedColormap([(1.,1.,1.)] * 50) + else: + skymap = self.skymap + max_val = max(skymap) + min_val = min(skymap) + moll_cbar = True if self.skymap is not None else None + + try: + hp.mollview(skymap, coord='C', cmap=cmap, rot=180, cbar=moll_cbar) #cbar=None) # + except Exception as e: + if min_val<1.0e-16: + #for some reason, sometimes have an underflow issue here + skymap[skymap<1.0e-16] = 0. + else: + print(e) + print('Failed to make all-sky unblinded skymap. Retry making full skymap! ') + return + + plt.text(2.0,0., r"$0^\circ$", ha="left", va="center") + plt.text(1.9,0.45, r"$30^\circ$", ha="left", va="center") + plt.text(1.4,0.8, r"$60^\circ$", ha="left", va="center") + plt.text(1.9,-0.45, r"$-30^\circ$", ha="left", va="center") + plt.text(1.4,-0.8, r"$-60^\circ$", ha="left", va="center") + plt.text(2.0, -0.15, r"$0^\circ$", ha="center", va="center") + plt.text(1.333, -0.15, r"$60^\circ$", ha="center", va="center") + plt.text(.666, -0.15, r"$120^\circ$", ha="center", va="center") + plt.text(0.0, -0.15, r"$180^\circ$", ha="center", va="center") + plt.text(-.666, -0.15, r"$240^\circ$", ha="center", va="center") + plt.text(-1.333, -0.15, r"$300^\circ$", ha="center", va="center") + plt.text(-2.0, -0.15, r"$360^\circ$", ha="center", va="center") + + hp.graticule(verbose=False) + + theta=np.pi/2 - events['dec'] + phi = events['ra'] + + #plot 90% containment errors + sigma_90 = events['sigma']*self._angScale + + # plot events on sky with error contours + handles=[] + hp.projscatter(theta,phi,c=cols,marker='x',label='Event',coord='C', zorder=5) + if label_events: + for j in range(len(theta)): + hp.projtext(theta[j], phi[j]-0.11, '{}'.format(j+1), color='red', fontsize=18, zorder=6) + handles.append(Line2D([0], [0], marker='x', ls='None', label='Event')) + + if (self.stop - self.start) <= 0.5: #Only plot contours if less than 2 days + for i in range(events['ra'].size): + my_contour = plotting_utils.contour(events['ra'][i], + events['dec'][i],sigma_90[i], self._nside) + hp.projplot(my_contour[0], my_contour[1], linewidth=2., + color=cols[i], linestyle="solid",coord='C', zorder=5) + + if self.skymap is None: + src_theta = np.pi/2. - self.dec + src_phi = self.ra + hp.projscatter(src_theta, src_phi, c = 'k', marker = '*', + label = self.name, coord='C', s=350) + handles.append(Line2D([0], [0], marker='*', c='k', ls='None', label=self.name)) + + if contour_files is not None: + cont_ls = ['solid', 'dashed'] + contour_counter = 0 + for c_file in contour_files: + cont = np.loadtxt(c_file, skiprows=1) + cont_ra = cont.T[0] + cont_dec = cont.T[1] + hp.projplot(np.pi/2. - cont_dec, cont_ra, linewidth=3., + color='k', linestyle=cont_ls[contour_counter], coord='C') + contour_counter += 1 + + if with_contour and self.skymap is not None: + probs = hp.pixelfunc.ud_grade(self.skymap, 64, power=-2) + probs = probs/np.sum(probs) + ### plot 90% containment contour of PDF + levels = [0.9] + theta, phi = plotting_utils.plot_contours(levels, probs) + hp.projplot(theta[0], phi[0], linewidth=2., c='k', label='Skymap (90\% cont.)') + handles.append(Line2D([0], [0], lw=2, c='k', label=r"Skymap (90\% cont.)")) + for i in range(1, len(theta)): + hp.projplot(theta[i], phi[i], linewidth=2., c='k') + + plt.title(self.name.replace('_', ' ')) + plt.legend(loc=1, handles=handles) + try: + plt.savefig(self.analysispath + '/' + self.analysisid + 'unblinded_skymap.png',bbox_inches='tight') + except: + plt.title('Fast Response Skymap') + plt.savefig(self.analysispath + '/' + self.analysisid + 'unblinded_skymap.png',bbox_inches='tight') + plt.close() + + def generate_report(self): + r"""Generates report using class attributes + and the ReportGenerator Class + Makes full report and outputs as pdf + """ + report = FastResponseReport(self) + report.generate_report() + report.make_pdf() + + + +class PointSourceFollowup(FastResponseAnalysis): + """ + Class for point-source or extended source followup + i.e. there is a fixed location on the sky, not a healpy skymap + + The arguments index and dataset override the defaults of the class. + """ + _nside = 256 + def __init__(self, name, ra, dec, tstart, tstop, extension=None, + index=None, dataset=None, fix_index=None, + skipped=None, outdir=None, save=True, seed=None): + + if index is not None: + self._index = float(index) + if dataset is not None: + self._dataset = dataset + if fix_index is not None: + self._fix_index = fix_index + self._float_index = not self._fix_index + + super().__init__(name, tstart, tstop, skipped=skipped, seed=seed, + outdir=outdir, save=save, extension=extension) + + + + + self.ra = np.deg2rad(ra) + self.dec = np.deg2rad(dec) + + if (self.ra < 0 or self.ra > 2*np.pi) or \ + (self.dec < -np.pi/2. or self.dec > np.pi/2.): + print("Right ascension and declination are not valid") + sys.exit() + else: + self.save_items['ra'] = self.ra + self.save_items['dec'] = self.dec + + self.skymap, self.nside, self.ipix_90 = None, None, None + + def __str__(self): + int_str = super().__str__() + int_str += ' '*10 + 'RA, dec.: {:.2f}, {:+.3f}'.format( + self.ra*180. / np.pi, self.dec*180. / np.pi) + exts = 0. if self.extension is None else self.extension + int_str += ' '*10 + 'Extension: {:.2f}'.format(exts * 180. / np.pi) + int_str += '\n\n' + return int_str + + def run_background_trials(self, ntrials=1000): + r""" Method to run background trials for point source analysis + + Parameters + ----------- + ntrials: int + number of background trials to run (default 1000) + """ + trials = self.llh.do_trials( + int(ntrials), + src_ra=self.ra, + src_dec=self.dec) + self.tsd = trials['TS'] + self.save_items['tsd'] = self.tsd + + def initialize_injector(self, e_range=(0., np.inf)): + r"""Method to make relevant injector in Skylab. + Used in calculating upper limit + + Parameters + ----------- + e_range: tuple + optional energy range allowed for injector. default (0., np.inf) + + Returns + -------- + inj: Skylab injector object + Point source injector using skylab + """ + inj = PointSourceInjector( + gamma = self.index, + E0 = 1000., + e_range=e_range) + if self.multi: + temporal_model = {enum:_llh.temporal_model for enum,_llh in self.llh._samples.items()} + else: + temporal_model = self.llh.temporal_model + inj.fill( + self.dec, + self.llh.exp, + self.llh.mc, + self.llh.livetime, + temporal_model=temporal_model) + self.inj = inj + self.save_items['E0'] = self.inj.E0 + + def unblind_TS(self): + r""" Unblind TS at one location for a point source + + Returns + -------- + ts: float + unblinded test statistic + ns: float + best fit ns + """ + # Fix the case of getting best-fit gamma + # TODO: What if gamma is floated + ts, ns = self.llh.fit_source(src_ra=self.ra, src_dec=self.dec) + params = ns.copy() + params.pop('nsignal') + self.ns_params = params + ns = ns['nsignal'] + if self._verbose: + print("TS = {}".format(ts)) + print("ns = {}".format(ns)) + for par, val in params.items(): + print(f"{par} = {val:.3f}") + print("\n\n") + self.ts, self.ns = ts, ns + self.save_items['ts'] = ts + self.save_items['ns'] = ns + # TODO alternatively change ReportGenerator to report ns_params generically + if 'gamma' in params: + self.gamma = params['gamma'] + for par, val in params.items(): + if par in self.save_items: + if self._verbose: + print(f'Warning, not saving {par} as save_items already has such a key') + self.save_items.setdefault(par, val) + return ts, ns + + def find_coincident_events_single(self, llh): + """Find ontime events with a spatial weight > 10. + + Parameters: + ----------- + llh: BaseLLH + A single-sample LLH, such as PointSourceLLH. + + + Returns: + coincident_events: [dict] + + """ + + spatial_weights = llh.llh_model.signal( + self.ra, self.dec, llh._events, + src_extension=self.extension)[0] / llh._events['B'] + energy_ratio, _ = llh.llh_model.weight( + llh._events, **self.ns_params) + # TODO or should that use a fixed index, for consistent plots, rather than what contributes to the LLH? + # the threshold of 10 is also tuned to GFU, gamma=2 + temporal_weights = llh.temporal_model.signal(llh._events) + msk = spatial_weights * energy_ratio * temporal_weights > 10 + coincident_events = [] + if len(spatial_weights[msk]) > 0: + coincident_events = [] + for ev, s_w, en_w in zip(llh._events[msk], + spatial_weights[msk], energy_ratio[msk]): + coincident_events.append({}) + for key in ['run', 'event', 'ra', 'dec', 'sigma', 'logE', 'time']: + coincident_events[-1][key] = ev[key] + del_psi = sensitivity_utils.deltaPsi([ev['dec']], [ev['ra']], [self.dec], [self.ra])[0] + coincident_events[-1]['delta_psi'] = del_psi + coincident_events[-1]['spatial_w'] = s_w + coincident_events[-1]['energy_w'] = en_w + return coincident_events + + def find_coincident_events(self): + r"""Find "coincident events" for the analysis. + These are ontime events that have a spatial times energy weight greater than 10. + These will be combined from all used samples, with a key 'enum' to distinguish. + """ + # TODO clarify the docstring: actually it includes the temporal weight! + coincident_events = [] + if self.multi: + for enum, _sam in self.llh._samples.items(): + _coincident_events = self.find_coincident_events_single(_sam) + for _ev in _coincident_events: + _ev['enum'] = enum + coincident_events += _coincident_events + + else: + coincident_events = self.find_coincident_events_single(self.llh) + self.coincident_events = coincident_events + self.save_items['coincident_events'] = coincident_events + + def ns_scan(self, params = None): + r""" Calculate the llh as a function of ns, to + see how shallow the likelihood space is. + + Parameters + ----------- + params: mappable + params dictionary for skylab.ps_llh object. Default is + that of this analysis (set below, using self.spectrum) + + Returns + --------- + xs, delta_llh: array-like + Values of ns scanned and corresponding -2*delta_llh values""" + + if params is None: + if self._fix_index: + params = {'spectrum': str(self.spectrum)} + else: + params = self.ns_params + # TODO or minimize per point? + + bounds = self.ns * 2.5 + if bounds == 0.0: + bounds = 1.0 + xs = np.linspace(0., bounds, 101) + llhs = [] + for n in xs: + llhs.append(self.llh.llh(n, **params)[0]) + llhs = np.array(llhs) + best_llh = self.llh.llh(self.ns, **params)[0] + delta_llh = -2.*(llhs - best_llh) + + if self.save_output: + fig, ax = plt.subplots() + plt.plot(xs, delta_llh, lw = 3) + plt.axvline(self.ns, ls = '--', color = 'gray', lw = 1) + plt.xlabel(r'$n_s$', fontsize = 18) + plt.ylabel(r'$-2\times \Delta \mathrm{LLH}$', fontsize = 18) + plt.savefig(self.analysispath + '/llh_ns_scan.png', bbox_inches='tight', dpi=200) + + self.ns_profile = (xs, delta_llh) + self.save_items['ns_scan'] = xs, delta_llh + return xs, delta_llh + + def upper_limit(self, n_per_sig=100, p0=None): + r"""After calculating TS, find upper limit + Assuming an E^-2 spectrum + + Parameters + ----------- + n_per_sig: int + number of trials per injected signal events to run + p0: None, scalar, or n-length sequence + Initial guess for the parameters passed to curve_fit (optional, default None) + + Returns + -------- + upperlimit: float + Value of E^2 dN / dE in units of TeV / cm^2 + """ + # TODO this docstring is not general enough, + # and actually the return value is from mu2flux but does not match its units + # because injector index is set by self._index, and for multisample we want softer indices. + if self.inj is None: + self.initialize_injector() + ninj = np.array([1., 1.5, 2., 2.5, 3., 4., 5., 6.]) + passing = [] + for n in ninj: + results = self.llh.do_trials( + n_per_sig, src_ra=self.ra, src_dec=self.dec, + injector=self.inj, mean_signal=n, poisson=True) + msk = results['TS'] > self.ts + npass = len(results['TS'][msk]) + passing.append((n, npass, n_per_sig)) + + signal_fluxes, passing, number = list(zip(*passing)) + signal_fluxes = np.array(signal_fluxes) + passing = np.array(passing, dtype=float) + number = np.array(number, dtype=float) + passing /= number + errs = sensitivity_utils.binomial_error(passing, number) + fits, plist = [], [] + for func, func_name in [(sensitivity_utils.chi2cdf, 'chi2'), + (sensitivity_utils.erfunc, 'Error function'), + (sensitivity_utils.incomplete_gamma, + 'Incomplete gamma')]: + try: + fits.append(sensitivity_utils.sensitivity_fit( + signal_fluxes, passing, errs, func, p0=p0)) + plist.append(fits[-1]['pval']) + except: + print(f"{func_name} fit failed in upper limit calculation") + #Find best fit of the three, make it look different in plot + plist = np.array(plist) + best_fit_ind = np.argmax(plist) + fits[best_fit_ind]['ls'] = '-' + self.upperlimit = self.inj.mu2flux(fits[best_fit_ind]['sens']) + self.upperlimit_ninj = fits[best_fit_ind]['sens'] + # E^2 dN/dE at self.inj.E0 + upperlimit_fluence = self.upperlimit * self.duration * 86400. * self.inj.E0**2 + + fig, ax = plt.subplots() + for fit_dict in fits: + lw = 1. if fit_dict['ls'] == '-' else 1.5 + ax.plot( + fit_dict['xfit'], fit_dict['yfit'], + label = r'{}: $\chi^2$ = {:.2f}, d.o.f. = {}'.format( + fit_dict['name'], fit_dict['chi2'], fit_dict['dof']), + ls = fit_dict['ls'], + lw=lw) + if fit_dict['ls'] == '-': + ax.axhline(0.9, color = 'm', linewidth = 0.3, linestyle = '-.') + ax.axvline(fit_dict['sens'], color = 'm', linewidth = 0.3, linestyle = '-.') + limit_annotation = r'Sens. = {:.2f} events'.format(fit_dict['sens']) + '\n' + limit_annotation += r' = {:.1e}'.format(upperlimit_fluence) + r' GeV cm$^{-2}$' + '\n' + #ax.text(3.5, 0.8, r'Sens. = {:.2f} events'.format(fit_dict['sens']), fontsize = 16) + #ax.text(3.5, 0.7, r' = {:.1e}'.format(upperlimit_fluence) + r' GeV cm$^{-2}$', fontsize = 16) + if self.index != 2: + # E^2 F not constant in energy, state pivot energy in label + #ax.text(3.5, 0.7, r' = {:.1e}'.format(upperlimit_fluence) + r' GeV cm$^{-2}$' + f'\nat {self.inj.E0:.0f} GeV', fontsize = 16) + limit_annotation += f'at {self.inj.E0:.0f} GeV' + #ax.text(5, 0.5, r'Sens. = {:.2e}'.format(self.inj.mu2flux(fit_dict['sens'])) + ' GeV^-1 cm^-2 s^-1') + ax.annotate(limit_annotation, (0.6, 0.8), ha = 'left', va = 'top', xycoords = 'axes fraction', fontsize = 16) + ax.errorbar(signal_fluxes, passing, yerr=errs, capsize = 3, linestyle='', marker = 's', markersize = 2) + ax.legend(loc=4, fontsize = 14) + ax.set_xlabel(r'$\langle n_{inj} \rangle$', fontsize = 14) + ax.set_ylabel(r'Fraction TS $>$ unblinded TS', fontsize = 14) + if self.save_output: + plt.savefig(self.analysispath + '/upper_limit_distribution.png', bbox_inches='tight', dpi=200) + print("Found upper limit of {:.2f} events".format(self.upperlimit_ninj)) + self.save_items['upper_limit'] = self.upperlimit + self.save_items['upper_limit_ninj'] = self.upperlimit_ninj + return self.upperlimit + + def make_dNdE_single(self): + r"""Make an E^-2 or E^-2.5 dNdE with the central 90% + for the most relevant declination band + (+/- 5 deg around source dec) + for a single dataset + """ + dec_mask_1 = self.llh.mc['dec'] > self.dec - (5. * np.pi / 180.) + dec_mask_2 = self.llh.mc['dec'] < self.dec + (5. * np.pi / 180.) + dec_mask_3, dec_mask_4 = None, None + dec_mask = dec_mask_1 * dec_mask_2 + fig, ax = plt.subplots(figsize = (8,5)) + fig.set_facecolor('white') + lab = None + delta_gamma = -1. * self.index + 1. + a = plt.hist(self.llh.mc['trueE'][dec_mask], bins = np.logspace(1., 8., 50), + weights = self.llh.mc['ow'][dec_mask] * np.power(self.llh.mc['trueE'][dec_mask], delta_gamma) / self.llh.mc['trueE'][dec_mask], + histtype = 'step', linewidth = 2., color = sns.xkcd_rgb['windows blue'], label = lab) + plt.yscale('log') + plt.xscale('log') + plt.grid(which = 'major', alpha = 0.25) + plt.xlabel('Energy (GeV)', fontsize = 24) + cdf = np.cumsum(a[0]) / np.sum(a[0]) + low_5 = np.interp(0.05, cdf, a[1][:-1]) + median = np.interp(0.5, cdf, a[1][:-1]) + high_5 = np.interp(0.95, cdf, a[1][:-1]) + self.low5 = low_5 + self.high5 = high_5 + plt.axvspan(low_5, high_5, color = sns.xkcd_rgb['windows blue'], alpha = 0.25, label="Central 90\%") + lab = 'Median' + plt.axvline(median, c = sns.xkcd_rgb['windows blue'], alpha = 0.75, label = lab) + plt.xlim(min(min(self.low5), 1e1), max(max(self.high5), 1e8)) + plt.legend(loc=4, fontsize=18) + plt.savefig(self.analysispath + '/central_90_dNdE.png',bbox_inches='tight') + + print('90% Central Energy Range: {}, {} GeV'.format(round(low_5), round(high_5))) + self.save_items['energy_range'] = (self.low5, self.high5) + + def make_dNdE_multi(self): + r"""Make an E^-2 or E^-2.5 dNdE with the central 90% + for the most relevant declination band + (+/- 5 deg around source dec) + for multiple datasets + """ + low5 = [] + high5 = [] + fig, ax = plt.subplots(figsize = (8,5)) + fig.set_facecolor('white') + for enum in self.llh._samples: + llh = self.llh._samples[enum] + dataset = self.dataset[enum].replace('_', ' ') + + dec_mask_1 = llh.mc['dec'] > self.dec - (5. * np.pi / 180.) + dec_mask_2 = llh.mc['dec'] < self.dec + (5. * np.pi / 180.) + dec_mask_3, dec_mask_4 = None, None + dec_mask = dec_mask_1 * dec_mask_2 + + + lab = dataset + color = f'C{enum+1}' + delta_gamma = -1. * self.index + 1. + a = plt.hist(llh.mc['trueE'][dec_mask], bins = np.logspace(1., 8., 50), + weights = llh.mc['ow'][dec_mask] * np.power(llh.mc['trueE'][dec_mask], delta_gamma) / llh.mc['trueE'][dec_mask], + histtype = 'step', linewidth = 2., color = color, label = lab) + cdf = np.cumsum(a[0]) / np.sum(a[0]) + low_5 = np.interp(0.05, cdf, a[1][:-1]) + median = np.interp(0.5, cdf, a[1][:-1]) + high_5 = np.interp(0.95, cdf, a[1][:-1]) + + plt.axvspan(low_5, high_5, color = color, alpha = 0.25, label="Central 90\%") + lab = 'Median' + plt.axvline(median, c = color, alpha = 0.75, label = lab) + low5.append(low_5) + high5.append(high_5) + print('{} 90% Central Energy Range: {}, {} GeV'.format(dataset, round(low_5), round(high_5))) + + plt.yscale('log') + plt.xscale('log') + plt.grid(which = 'major', alpha = 0.25) + plt.xlabel('Energy (GeV)', fontsize = 24) + + plt.xlim(1e1, 1e8) + plt.legend(loc=4, fontsize=18) + plt.savefig(self.analysispath + '/central_90_dNdE.png',bbox_inches='tight') + + self.low5 = low5 + self.high5 = high5 + self.save_items['energy_range'] = tuple(zip(self.low5, self.high5)) + + def make_dNdE(self): + r"""Make an E^-2 or E^-2.5 dNdE with the central 90% + for the most relevant declination band + (+/- 5 deg around source dec) + """ + if self.multi: + self.make_dNdE_multi() + else: + self.make_dNdE_single() + + def write_circular(self): + pass diff --git a/fast_response/MultiAlertFollowup.py b/fast_response/MultiAlertFollowup.py new file mode 100644 index 0000000..ebc7851 --- /dev/null +++ b/fast_response/MultiAlertFollowup.py @@ -0,0 +1,161 @@ + +from abc import abstractmethod +import os, sys, time, subprocess +import pickle, dateutil.parser, logging, warnings +from argparse import Namespace +from copy import deepcopy + +import h5py +import healpy as hp +import numpy as np +import seaborn as sns +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy.lib.recfunctions as rf +from astropy.time import Time +from scipy.special import erfinv +from matplotlib.lines import Line2D +from scipy import sparse + +from skylab.datasets import Datasets +from skylab.llh_models import EnergyLLH +from skylab.priors import SpatialPrior +from skylab.ps_injector import PointSourceInjector +from skylab.ps_llh import PointSourceLLH, MultiPointSourceLLH +from skylab.ps_injector import PriorInjector +from skylab.spectral_models import PowerLaw +from skylab.temporal_models import BoxProfile, TemporalModel +# import meander + +from . import web_utils +from . import sensitivity_utils +from . import plotting_utils +from .reports import FastResponseReport +from .FastResponseAnalysis import FastResponseAnalysis, PriorFollowup, PointSourceFollowup +from .MultiFastResponseAnalysis import MultiPriorFollowup +from .GWFollowup import GWFollowup +from .AlertFollowup import AlertFollowup, CascadeFollowup, TrackFollowup + +from .reports import AlertReport + +mpl.use('agg') +current_palette = sns.color_palette('colorblind', 10) +logging.getLogger().setLevel(logging.ERROR) +warnings.simplefilter("ignore", UserWarning) +warnings.simplefilter("ignore", RuntimeWarning) + +# the real work comes with GWFollowup +# other than ExternalFollowup, it doesn't just add attributes +# but specify many methods, some which have to be adapted +# => the inheritance gets even deeper +# However Track and CascadeFollowup are relatively close + + +class MultiAlertFollowup(AlertFollowup, MultiPriorFollowup): + + # attributes that are identical for each analysis produced by one followup class + # will be broadcast to the constituent analyses + # this is clunky, but more explicit than some dir(...) shenanigans + # and excludes _season_names and _jitter which are not the same + static_attributes = ['_verbose', + '_index', + '_float_index', + '_fix_index', + '_index_range', + '_llh_seed', + '_nb_days', + '_ncpu', + '_pixel_scan_nsigma', + '_allow_neg', + '_containment', + '_nside', + '_smear', + ] + _base_dir = '/data/user/chraab/fast_response/multisample_alert_followup/' + _bg_trial_dir = os.path.join(_base_dir, 'alert_precomputed_trials') + _sens_dir = os.path.join(_base_dir, 'reference_sensitivity_curves') + + def run_background_trials(self, ntrials = 1000): + '''Not yet adapted to rate fluctuations as in PriorFollowup + ''' + # TODO include saving unpenalised trials + return PriorFollowup.run_background_trials(self, ntrials=ntrials) + + # TODO each class should have method + # - generate and save BG trials + # - load BG trials + # - calculate sensitivity curve + # - load sensitivity curve + # => same configuration can be used in script for + # - analysis + # - filling the cache + + def run_background_trials(self, ntrials = 1000): + # + # TODO other solutions for this + # - by month like in GW + # - + # - + # this is tricky - don't implement caching for now. + # as it deals with Multi-LLH trials (calculated where?) + # but uses self.llh.nbackground (which is defined for PSLLH only) + # to retrieve BG trials for the current rate - in multi + # and then it also requires a strategy for background rate in multiple samples. + # raise NotImplementedError('''Not decided how to handle background fluctuations + # in multiple samples, which may be partially correlated, + # but energy x selection adds another axis to seasonality.''') + duration_seconds = self.duration * 86400. + closest_rates = [] + for enum in self.llh._samples: + nbackground = self.llh._samples[enum].nbackground + # FIXME do it properly/correctly + #current_rate = nbackground / (self.duration * 86400.) * 1000. + current_rate = nbackground / (self.llh._samples[enum].on_livetime * 86400.) * 1000. + closest_rate = sensitivity_utils.find_nearest(np.linspace(4.2, 7.2, 6), current_rate) + closest_rates.append(closest_rate) + + rates_str = '_'.join([f'{_rate:.1f}' for _rate in closest_rates]) + trials_filename = '_'.join([ + f'precomputed_trials_delta_t_{self.duration * 86400.:.2e}', + f'index_{self._index}', + f'{rates_str}_mHz', + f'low_stats.npz', + ] + ) + + pre_ts_array = sparse.load_npz(os.path.join(trials_filename)) + ts_norm = np.log(np.amax(self.skymap)) + ts_prior = pre_ts_array.copy() + ts_prior.data += 2.*(np.log(self.skymap[pre_ts_array.indices]) - ts_norm) + ts_prior.data[~np.isfinite(ts_prior.data)] = 0. + ts_prior.data[ts_prior.data < 0] = 0. + tsd = ts_prior.max(axis=1).A + tsd = np.array(tsd) + self.tsd = tsd + return tsd + + #def sens_range_plot(self): + # TODO include method to produce sensitivity curve + + # so far have adapted report to be flexible + # should there be dedicated Multi...Report classes? TODO + def generate_report(self): + r""" + Generates report using ReportGenerator.MultiAlertReport attributes + and the ReportGenerator Class + + """ + super().generate_report() + + def write_circular(self, alert_results): + raise NotImplementedError('Analysis specific') + +class MultiTrackFollowup(MultiAlertFollowup, TrackFollowup): + # this inheritance structure should work as long as CascadeFollowup does not override methods + # but getting very spaghetti already and some refactoring might be in order + _smear = True + + +class MultiCascadeFollowup(MultiAlertFollowup, CascadeFollowup): + _smear = False + diff --git a/fast_response/MultiExternalFollowup.py b/fast_response/MultiExternalFollowup.py new file mode 100644 index 0000000..cc628c8 --- /dev/null +++ b/fast_response/MultiExternalFollowup.py @@ -0,0 +1,59 @@ +from .MultiFastResponseAnalysis import MultiPointSourceFollowup +from .FastResponseAnalysis import PointSourceFollowup + +import numpy as np + +# Consistent with existing structure: +# base class: methods +# specific class: data sample, class attributes, "analysis definition" +# instance: specific source/follow-up +class GFUFollowup(PointSourceFollowup): + _dataset = "GFUOnline_v001p02" + _season_names = [f"IC86, 201{y}" for y in range(7, 10)] + #_season_names = [f"IC86, 201{y}" for y in range(1, 10)] + _floor = np.radians(0.2) + +class GrecoFollowup(PointSourceFollowup): + _dataset = 'GrecoOnline_v002pFactor' + _season_names = [f"IC86, 20{y:02d}" for y in range(18, 20+1)] + #_season_names = [f"IC86, 20{y:02d}" for y in range(12, 22+1)] + _floor = np.radians(0.2) # can change this! + # extended/GRB-style LLH? + +class DNNFollowup(PointSourceFollowup): + _dataset = "DNNCascades_v001p01" # TODO switch to "online" version + _season_names = [f"IC86, 20{y:02d}" for y in range(18, 19+1)] + _floor = np.radians(1.5) # can change this! + _jitter = 3. # common default for DNN analyses + # Need to add analysis cuts? That's something SKATE analysers would know. + +#or any other number of definitions + +class DNNOnlineFollowup(PointSourceFollowup): + _dataset = "DNNCascadesOnline_v001p01" + _season_names = ["livestream"] + _floor = np.radians(1.5) # can change this! + _jitter = 3. # common default for DNN analyses + _background_days = 100. + # Need to add analysis cuts? That's something SKATE analysers would know. + +class DNNIceManFollowup(PointSourceFollowup): + _dataset = "DNNCascadesIceMan_v001p00" + # limited to one season ON PURPOSE for better comparison with DNNOnline + _season_names = [f"IC86, 20{y:02d}" for y in range(18, 18+1)] + _jitter = 3. # common default for DNN analyses + +class MultiFollowup(MultiPointSourceFollowup): + ''' + Class for external point-source or extended source followup. + By default, uses floating index of 2.5 in the LLH. Based on + the PointSourceFollowup class adapted to accept multiple samples + via the configuration of constituent follow-up analyses. + + ''' + _followups = [GFUFollowup, GrecoFollowup] # more consistent + _fix_index = True + _float_index = not _fix_index + _index = 2.5 + + \ No newline at end of file diff --git a/fast_response/MultiFastResponseAnalysis.py b/fast_response/MultiFastResponseAnalysis.py new file mode 100644 index 0000000..ad1e970 --- /dev/null +++ b/fast_response/MultiFastResponseAnalysis.py @@ -0,0 +1,431 @@ +r''' General Fast Response Analysis Class. + + Author: Alex Pizzuto + Date: 2021 + + One of two variants to include mutiple datasets. + Adding a new class that holds multiple FastResponseAnalysis as a container, + adds the appropriate LLH, injector, and adapted methods. + + ''' + +from abc import abstractmethod +import os, sys, time, subprocess +import pickle, dateutil.parser, logging, warnings +from argparse import Namespace +from copy import deepcopy + +import h5py +import healpy as hp +import numpy as np +import seaborn as sns +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy.lib.recfunctions as rf +from astropy.time import Time +from scipy.special import erfinv +from matplotlib.lines import Line2D + +from skylab.datasets import Datasets +from skylab.llh_models import EnergyLLH +from skylab.priors import SpatialPrior +from skylab.ps_injector import PointSourceInjector +from skylab.ps_llh import PointSourceLLH, MultiPointSourceLLH +from skylab.ps_injector import PriorInjector +from skylab.spectral_models import PowerLaw +from skylab.temporal_models import BoxProfile, TemporalModel +# import meander + +from . import web_utils +from . import sensitivity_utils +from . import plotting_utils +from .reports import FastResponseReport +from .FastResponseAnalysis import FastResponseAnalysis, PriorFollowup, PointSourceFollowup +from .GWFollowup import GWFollowup +from .AlertFollowup import AlertFollowup, CascadeFollowup, TrackFollowup + +mpl.use('agg') +current_palette = sns.color_palette('colorblind', 10) +logging.getLogger().setLevel(logging.ERROR) +warnings.simplefilter("ignore", UserWarning) +warnings.simplefilter("ignore", RuntimeWarning) + +class MultiFastResponseAnalysis(FastResponseAnalysis): + """ + Instead of sprinkling if-else statements around, could leave + the LLH contained within the original FastResponseAnalysis. + Analogous to MultiPointSourceLLH, this one then is a container + for one FRA per dataset, plus a MultiPointSourceLLH, and the injector. + Will involve some duplicated code, unless the methods are split up more, + or both FRA and MultiFRA get a new base class (analogous to BaseLLH). + But it separates code more cleanly, and can override inherited methods + instead of having two implementations with an if-else statement inbetween. + Requires some extra arguments to FastResponseAnalysis methods. + """ + + # attributes that will be set for a specific followup configuration inheriting from this base class + _followups = None + + # attributes that are identical for each analysis produced by one followup class + # will be broadcast to the constituent analyses + static_attributes = [] + + # inherits some default config attributes + def __init__(self, *args, **kwargs): + """The arguments dataset and index override the defaults + that may have been set by child classes. + + Same args and kwargs as FastResponseAnalysis + + """ + logging.debug('MultiFastResponseAnalysis') + + # basic config of this instance, and initialize_llh + super().__init__(*args, **kwargs) + + # TODO can we use just a bit of composition to clarify the inheritance maze? + + # save spectrum and time profile to help broadcast + self.spectrum = None + if self._fix_index: + self.spectrum = PowerLaw(A=1, gamma=self.index, E0=1000.) + self.time_profile = BoxProfile(self.start, self.stop) + + def initialize_analyses(self, *args, **kwargs): + # initialize individual dataset LLH + # needs to be here as they will have the same constructor signature + self.analyses = [] + for _followup in self._followups: + #if not issubclass(type(_followup), type(self)): + # raise TypeError(f'Trying to construct a {type(self)} with a {type(_followup)}') + _kwargs = deepcopy(kwargs) + _kwargs['save'] = False # this single FRA does not save output + # other class attributes one may want to broadcast + for attr in self.static_attributes: + setattr(_followup, attr, getattr(self, attr)) + # initialize with analysis specific args and kwargs + # such as source properties, or override settings + + # get tstart, tstop, ra, dec / skymap, and other config + _analysis = _followup(*args, **_kwargs) + self.analyses.append(_analysis) + + def initialize_llh(self, skipped=None, scramble=False): + if self._verbose: + print("Initializing MultiPointSourceLLH in Skylab") + + # kwargs for the BaseLLH instance + # TODO maybe this should be an attribute, + # so config can be changed in one place for both FRA snd MultiFRA? + base_kwargs = dict( + nsource=1., # seed for nsignal fit + nsource_bounds=(0., 1e3), # bounds on fitted ns + ncpu=self._ncpu, # use 10 CPUs when computing trials + seed=self.llh_seed, + ) + + multi_llh = MultiPointSourceLLH(**base_kwargs) + for enum, fra in enumerate(self.analyses): + fra.llh.do_trials_seed = 1 + multi_llh.add_sample(fra._dataset, fra.llh) + # TODO make sure sample_weights are still correct + # (as the temporal_model system wasn't used for this before) + return multi_llh + + def remove_event(self, exp, dset, skipped): + # should this ever be called from this instance? + # need to save event in self.save_items + raise NotImplemented('') + + # These properties are initialized for the base class + # could do something meaningful + @property + def skipped_event(self): + return self._skipped_event + # TODO could be a property + # that retrieves unique skipped_event from constituent analyses + @skipped_event.setter + def skipped_event(self, x): + self._skipped_event = x + # TODO make something meaningful of this + + + + @property + def exp(self): + '''Dictionary or flat array like llh_exp? used in remove_event''' + raise NotImplemented('') + + @exp.setter + def exp(self, x): + self._exp = x + + @property + def mc(self): # and livetime, grl, sinDec_bins, energy_bins... only used in LLH + '''Dictionary? i.e. self.llh.mc''' + raise NotImplemented('') + + @property + def livetime(self): + '''Dictionary?''' + raise NotImplemented('') + + @property + def dset(self): + '''Dictionary? used in remove_event''' + raise NotImplemented('') + + + @property + def datasets(self): + """Returns the datasets used""" + return [_a._dataset for _a in self.analyses] + + @property + def analyses(self): + '''Returns the constituent followup instances''' + return self._analyses + @analyses.setter + def analyses(self, x): + self._analyses = x + # Feature to come for easier comparison: + # re-initialize LLH from existing FRA's + # if isinstance(x, list): + # self._analyses = x + # elif isinstance(x, FastResponseAnalysis): + # self._analyses = [x] + # else: + # raise TypeError(f'trying to set analyses with {type(x)}') + + @property + def llh_exp(self): + """Returns a flat array of experimental data loaded across the used dataset(s) loaded into the LLH, + limited to fields used for plotting and amended with a field for the dataset.""" + if not hasattr(self, '_llh_exp'): + exp = {enum:_llh.exp for enum, _llh in self.llh._samples.items()} + for enum, _exp in exp.items(): + exp[enum] = self.unify_exp_array(_exp, enum=enum) + self._llh_exp = np.concatenate([exp[enum] for enum in exp]) + return self._llh_exp + + def plot_skymap(self, **kwargs): + # TODO maybe a "short name" description in the actual Skylab dataset? + labels = [] + for _ds in self.datasets: + _base = _ds.split('_')[0] # convention: version after underscore + _base = _base.replace('Online', '') # not necessary for legend + _base = _base.replace('Greco', 'GRECO') # some prefer this + labels.append(f'{_base} Event') + return super().plot_skymap(labels=labels, **kwargs) + + @property + def dataset_string(self): + string = '' + string += 'Datasets:\n' + string += '\n'.join(self.datasets) + string += '\n\n' + return string + +class MultiPriorFollowup(PriorFollowup, MultiFastResponseAnalysis): + + # attributes that are identical for each analysis produced by one followup class + # will be broadcast to the constituent analyses + static_attributes = ['_verbose', + '_index', + '_float_index', + '_fix_index', + '_index_range', + '_llh_seed', + '_nb_days', + '_ncpu', + '_pixel_scan_nsigma', + '_allow_neg', + '_containment', + '_nside', + ] + # TODO let this be defined in parent classes, extended by child classes + # could use a property and super()? + + def __init__(self, *args, **kwargs): + logging.debug('MultiPriorFollowup.__init__') + + # first, construct constituent analyses with same arguments + self.initialize_analyses(*args, **kwargs) + + # then prepare LLH and store prior-specific attributes + super().__init__(*args, **kwargs) + + def __str__(self): + string = super().__str__().rstrip() + string += '\n' + string += self.dataset_string + return string + + def initialize_injector(self, e_range=(0., np.inf)): + print("Initializing Prior Injector") + spatial_prior = SpatialPrior(self.skymap, containment = self._containment, allow_neg=self._allow_neg) + self.spatial_prior = spatial_prior + inj = PriorInjector( + spatial_prior, + gamma=self.index, + e_range = e_range, + E0=1000., + seed = self.llh_seed) + temporal_model = {enum:_llh.temporal_model for enum,_llh in self.llh._samples.items()} + inj.fill( + self.llh.exp, + self.llh.mc, + self.llh.livetime, + temporal_model=temporal_model) + self.inj = inj + self.save_items['E0'] = self.inj.E0 + + #def find_coincident_events(self): + # No implementation needed since replacing self.llh.exp with self.llh_exp + # TODO get feedback if that change is acceptable + + # TODO why does run_background_trials call initialize_llh to init a new LLH with scrambled data? + # is there no scrambling within the trials? + # FIXME my code breaks this: initialize_llh only constructs MultiPSLLH + + def make_dNdE(self): + r"""Make an E^-2 or E^-2.5 dNdE with the central 90% + for the minimum and maximum declinations on the skymap + for multiple datasets + """ + raise NotImplementedError("would be easier if we shared more plotting methods") + +class MultiPointSourceFollowup(PointSourceFollowup, MultiFastResponseAnalysis): + + # attributes that are identical for each analysis produced by one followup class + # will be broadcast to the constituent analyses + static_attributes = ['_verbose', + '_index', + '_float_index', + '_fix_index', + '_index_range', + '_llh_seed', + '_nb_days', + '_ncpu', + ] + + def __init__(self, *args, followups=None, **kwargs): + + logging.debug('MultiPointSourceFollowup.__init__') + if followups is not None: + self._followups = followups + + # first, construct constituent analyses with same arguments + self.initialize_analyses(*args, **kwargs) + + # then construct LLH + super().__init__(*args, **kwargs) + + def __str__(self): + string = super().__str__().rstrip() + string += self.dataset_string + return string + + def initialize_injector(self, e_range=(0., np.inf)): + inj = PointSourceInjector( + gamma = self.index, + E0 = 1000., + e_range=e_range) + temporal_model = {enum:_llh.temporal_model for enum,_llh in self.llh._samples.items()} + inj.fill( + self.dec, + self.llh.exp, + self.llh.mc, + self.llh.livetime, + temporal_model=temporal_model) + self.inj = inj + self.save_items['E0'] = self.inj.E0 + + def find_coincident_events(self): + '''Retrieve events with spatial x energy x temporal weight>10 + from all component analyses, store in self.coincident_events + as a list of dictionaries with columns for the report tables, + ['run', 'event', 'ra', 'dec', 'sigma', 'logE', 'time'] + added Delta Psi, spatial weight, energy weight, sample enum. + ''' + # TODO ideally samples should not overlap, however GFU and Greco do. + # decide whether to require de-duplication at analysis level, + # or in this table (and record tuples in enum column). + for enum, _ana in enumerate(self.analyses): + _ana.find_coincident_events(ns_params = self.ns_params) + for _event in _ana.coincident_events: + _event['enum'] = enum + _event['dataset'] = _ana.dataset + self.coincident_events = sum([_ana.coincident_events for _ana in self.analyses], []) + self.save_items['coincident_events'] = self.coincident_events + + def make_dNdE(self): + # easier if calculation and plotting were separate methods... + # but how to do it is already laid out in FRA_ifelse + r"""Make an E^-2 or E^-2.5 dNdE with the central 90% + for the most relevant declination band + (+/- 5 deg around source dec) + for multiple datasets + """ + low5 = [] + high5 = [] + fig, ax = plt.subplots(figsize = (8,5)) + fig.set_facecolor('white') + for enum in self.llh._samples: + llh = self.llh._samples[enum] + dataset = self.datasets[enum].replace('_', ' ') + style = plotting_utils.skymap_style[enum] + dec_mask_1 = llh.mc['dec'] > self.dec - (5. * np.pi / 180.) + dec_mask_2 = llh.mc['dec'] < self.dec + (5. * np.pi / 180.) + dec_mask_3, dec_mask_4 = None, None + dec_mask = dec_mask_1 * dec_mask_2 + + + lab = dataset + #color = f'C{enum+1}' + color = color = sns.xkcd_rgb['windows blue'] + delta_gamma = -1. * self.index + 1. + a = plt.hist(llh.mc['trueE'][dec_mask], bins = np.logspace(1., 8., 50), + weights = llh.mc['ow'][dec_mask] * np.power(llh.mc['trueE'][dec_mask], delta_gamma) / llh.mc['trueE'][dec_mask], + histtype = 'step', linewidth = 2., color = color, label = lab) + cdf = np.cumsum(a[0]) / np.sum(a[0]) + low_5 = np.interp(0.05, cdf, a[1][:-1]) + median = np.interp(0.5, cdf, a[1][:-1]) + high_5 = np.interp(0.95, cdf, a[1][:-1]) + + plt.axvspan(low_5, high_5, + color = color, linestyle = style['linestyle'], + linewidth = 2., + alpha = 0.25, label="Central 90%") + lab = 'Median' + plt.axvline(median, + c = color, linestyle = style['linestyle'], + linewidth = 2., + alpha = 0.75, label = lab) + low5.append(low_5) + high5.append(high_5) + print('{} 90% Central Energy Range: {}, {} GeV'.format(dataset, round(low_5), round(high_5))) + + plt.yscale('log') + plt.xscale('log') + plt.grid(which = 'major', alpha = 0.25) + plt.xlabel('Energy (GeV)', fontsize = 24) + + plt.xlim(1e1, 1e8) + plt.legend(loc=4, fontsize=18) + plt.savefig(self.analysispath + '/central_90_dNdE.png',bbox_inches='tight') + + self.low5 = low5 + self.high5 = high5 + self.save_items['energy_range'] = tuple(zip(self.low5, self.high5)) + return + + def write_circular(self): + raise NotImplemented('This method is not implemented for the parent class, either.') + + + + + + + \ No newline at end of file diff --git a/fast_response/MultiFollowup_ifelse.py b/fast_response/MultiFollowup_ifelse.py new file mode 100644 index 0000000..e386767 --- /dev/null +++ b/fast_response/MultiFollowup_ifelse.py @@ -0,0 +1,16 @@ +from .FastResponseAnalysis_ifelse import PointSourceFollowup + +class ExternalFollowup(PointSourceFollowup): + ''' + Class for external point-source or extended source followup. + By default, uses a fixed index of 2.5 in LLH. Based on + the PointSourceFollowup class adapted to accept multiple samples + + ''' + _dataset = ["GFUOnline_v001p02", + "GrecoOnline_v002pFactor", + ] + + _fix_index = True + _float_index = not _fix_index + _index = 2.5 \ No newline at end of file diff --git a/fast_response/MultiGWFollowup.py b/fast_response/MultiGWFollowup.py new file mode 100644 index 0000000..44844fe --- /dev/null +++ b/fast_response/MultiGWFollowup.py @@ -0,0 +1,48 @@ +from os.path import join +from .MultiFastResponseAnalysis import MultiPriorFollowup +from .GWFollowup import GWFollowup as GFUFollowup +from .FastResponseAnalysis import PriorFollowup + +import numpy as np + +class OnlyDNNOnlineFollowup(PriorFollowup): + _base_dir = "/data/user/chraab/fast_response/multisample/variable_jitter_new_environment" + _sens_dir = join(_base_dir, "precomputed_sensitivity") + _bg_dir = join(_base_dir, "precomputed_trials") + _bg_format = '_'.join([ + 'precomputed_trials_delta_t_{delta_t:.2e}', + 'nside_{nside}', + 'index_{index}', + '{lookup}', + 'low_stats.npz', + ]) + _dataset = "DNNCascadesOnline_v001p01" + _season_names = ["livestream"] + _floor = np.radians(1.5) # can change this! + _jitter = 3. # common default for DNN analyses + _background_days = 100. + # Need to add analysis cuts? That's something SKATE analysers would know. + +class DNNOnlineFollowup(MultiPriorFollowup): + _base_dir = "/data/user/chraab/fast_response/multisample/variable_jitter_new_environment" + _sens_dir = join(_base_dir, "precomputed_sensitivity") + _bg_dir = join(_base_dir, "precomputed_trials") + _bg_format = '_'.join([ + 'precomputed_trials_delta_t_{delta_t:.2e}', + 'nside_{nside}', + 'index_{index}', + '{lookup}', + 'low_stats.npz', + ]) + _followups = [OnlyDNNOnlineFollowup] + _fix_index = True # TODO check if this is what they do + _float_index = not _fix_index + _index = 2.0 + _nside = 64 + +class MultiGWFollowup(MultiPriorFollowup): + _followups = [GFUFollowup, OnlyDNNOnlineFollowup] + _fix_index = True # TODO check if this is what they do + _float_index = not _fix_index + _index = 2.0 + \ No newline at end of file diff --git a/fast_response/plotting_utils.py b/fast_response/plotting_utils.py index b024025..93760a3 100644 --- a/fast_response/plotting_utils.py +++ b/fast_response/plotting_utils.py @@ -4,6 +4,47 @@ import matplotlib.pyplot as plt import matplotlib as mpl import meander +from copy import copy + +skymap_style = [dict(linestyle='solid', marker='x', alpha=1.0), + dict(linestyle='dotted', marker='+', alpha=1.0), + dict(linestyle='dashed', marker='2', alpha=1.0), + dict(linestyle='dashdot', marker='d', alpha=1.0), + ] + +class TimeColormap: + """ + Class to colorize events in a skymap, optionally with a distinct sample per palette. + """ + def __init__(self, start, stop, n_maps=3): + self.norm = mpl.colors.Normalize(vmin=start, vmax=stop) + cmaps = [] + if n_maps == 1: + cmaps = [mpl.colors.ListedColormap(sns.color_palette('icefire', 512))] + elif n_maps<4: + dhue = 360/n_maps + hues = np.mod(136 + np.arange(0, 360, dhue), 360) + for i,hue in enumerate(hues): + cmaps.append(sns.diverging_palette(hue, (hue+dhue/1.5)%360, l=75, center="dark", as_cmap=True)) + else: + raise ValueError("Generating more than 5 color maps will be hard to distinguish") + self.cmaps = cmaps + self.n_maps = n_maps + + def get_cmap(self, i): + if self.n_maps == 1: + return self.cmaps[0] + elif i < self.n_maps: + return self.cmaps[i] + else: + raise ValueError("Have not configured enough color maps") + + def __call__(self, times, enum): + t = np.atleast_1d(times) + x = self.norm(t) + cmap = self.get_cmap(enum) + return cmap(x) + def plot_zoom(scan, ra, dec, title, reso=3, var="pVal", range=[0, 6],cmap=None): """ @@ -105,8 +146,9 @@ def plot_labels(src_dec, src_ra, reso): plt.text(np.radians(0), np.radians(-2.05*reso), r"right ascension", ha='center', va='center', fontsize=fontsize) +# FIXME reso here is not used def plot_events(dec, ra, sigmas, src_ra, src_dec, reso, sigma_scale=5., col = 'k', constant_sigma=False, - same_marker=False, energy_size=False, with_mark=True, with_dash=False, + same_marker=False, energy_size=False, with_mark=True, with_dash=False, kw_style={}, label=''): """ Adds events to a healpy zoom plot. Events are expected to be from self.llh.exp @@ -133,10 +175,12 @@ def plot_events(dec, ra, sigmas, src_ra, src_dec, reso, sigma_scale=5., col = 'k constant_sigma: bool Ignores sigma parameter and plots all markers with a size of 20. with_mark: bool - Uses an x marker instead of o + Include marker at event location in addition to error circle with_dash: bool Plot the angular error as a dashed contour. Usually used to indicated a removed event (e.g. alert event that triggered the analysis) + kw_style: dict + dictionary of style to use: marker and line style. Overridden by with_dash. same_marker, energy_size: bool Currently unused options. """ @@ -144,27 +188,44 @@ def plot_events(dec, ra, sigmas, src_ra, src_dec, reso, sigma_scale=5., col = 'k tmp = np.cos(src_ra - ra) * np.cos(src_dec) * cos_ev + np.sin(src_dec) * np.sin(dec) dist = np.arccos(tmp) + # with_dash overrides the given line style + kw_style = copy(kw_style) # because a dict is mutable + if with_dash: + kw_style['linestyle'] = ':' + # else, set default style + else: + kw_style.setdefault('linestyle', 'solid') + kw_style.setdefault('marker', 'x') + marker = kw_style.pop('marker') # have to pop it out as it's used in a different place + + if sigma_scale is not None: sigma = np.degrees(sigmas)/sigma_scale sizes = 5200*sigma**2 if constant_sigma: sizes = 20*np.ones_like(sizes) - if with_dash: - hp.projscatter(np.pi/2-dec, ra, marker='o', linewidth=2, - edgecolor=col, linestyle=':', facecolor="None", s=sizes, - alpha=1.0) - else: - hp.projscatter(np.pi/2-dec, ra, marker='o', linewidth=2, - edgecolor=col, facecolor="None", s=sizes, alpha=1.0) + + hp.projscatter(np.pi/2-dec, ra, marker='o', linewidth=2, + edgecolor=col, facecolor="None", s=sizes, + **kw_style, + ) if with_mark: - hp.projscatter(np.pi/2-dec, ra, marker='x', linewidth=2, + hp.projscatter(np.pi/2-dec, ra, marker=marker, linewidth=2, edgecolor=col, facecolor=col, s=60, alpha=1.0) +def auto_reso(events): + raise NotImplementedError('plot is ill defined') + #return 3. + + def load_plotting_settings(): """ Load settings to be used as default plot settings. Includes Times New Roman font and size 12 font """ + # undo eventual matplotlibrc in user config + mpl.rcdefaults() + mpl.use('agg') mpl.rcParams['text.usetex'] = True try: @@ -182,6 +243,9 @@ def load_plotting_settings(): mpl.rcParams['xtick.major.size'] = 5 mpl.rcParams['ytick.major.size'] = 5 + # increase figure resolution from default + mpl.rcParams['savefig.dpi'] = 200 + def contour(ra, dec, sigma, nside): r""" Function for plotting contours on skymaps diff --git a/fast_response/precomputed_background/glob_precomputed_trials_multi.py b/fast_response/precomputed_background/glob_precomputed_trials_multi.py new file mode 100644 index 0000000..40a8a5d --- /dev/null +++ b/fast_response/precomputed_background/glob_precomputed_trials_multi.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python + +from glob import glob +import healpy as hp +from scipy import sparse +import time +import argparse +import os + +parser = argparse.ArgumentParser(description='Glob precomp trials') +parser.add_argument('--dir',type=str, default='./', + help='directory for where trials are, will save globbed npz inside same ') +parser.add_argument('--nside',type=int, default=256, + help='nside used when running trials (default 256)') + +def get_glob_file(filename): + dirname = os.path.dirname(filename) + basename = os.path.basename(filename) + particles = basename.replace('seed_','seed').split('_') + basename_glob = '_'.join(_p for _p in particles if not 'seed' in _p) + return os.path.join(dirname, basename_glob) + +def sort_by_glob_file(files): + sorted = {} + for _file in files: + outfile = get_glob_file(_file) + sorted.setdefault(outfile, []) + sorted[outfile].append(_file) + return sorted + + + +def glob_allsky_scans(files, out, nside): + print(f'Creating {os.path.basename(out)}') + files = [fn for fn in files if f"nside_{nside}" in fn] + print('Found {} files to load'.format(len(files))) + if len(files)==0: return None + print('Nside: {}'.format(nside)) + npix = hp.nside2npix(nside) + print('Starting to load at {}'.format(time.ctime())) + maps = sparse.csr_matrix((0, npix), dtype=float) + for f in files: + scan = sparse.load_npz(f) + maps = sparse.vstack((maps, scan)) + print('.', end=' ') + print('') + + # Change format of sparse array + print("Starting to change from COO to CSR at {}".format(time.ctime())) + scans = maps.tocsr() + print("Finished at {}".format(time.ctime())) + + # Save the sparse array + sparse.save_npz(out, scans) + return maps + +def main(args): + """ + Glob the all-sky scans together + """ + # separate alert and GW trials by directory + files = sorted(glob(os.path.join(args.dir, '*seed_*.npz'))) + for outfile, filegroup in sort_by_glob_file(files).items(): + maps = glob_allsky_scans(filegroup, outfile, args.nside) + del maps + print ('done') + +args = parser.parse_args() +main(args) \ No newline at end of file diff --git a/fast_response/precomputed_background/precompute_ts_multi.py b/fast_response/precomputed_background/precompute_ts_multi.py new file mode 100644 index 0000000..c23de4f --- /dev/null +++ b/fast_response/precomputed_background/precompute_ts_multi.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python + +r""" +Run trials for background only, all-sky scans. Record TS and +number of true and fitted events around best fit location + +""" +import numpy as np +import healpy as hp +import os, sys, argparse +from astropy.time import Time +from numpy.lib.recfunctions import append_fields +from scipy import sparse +from glob import glob + +from fast_response.MultiExternalFollowup import MultiFollowup, GFUFollowup, GrecoFollowup, DNNFollowup, DNNOnlineFollowup + +parser = argparse.ArgumentParser(description='Precompute MultiAlertFollowup BG trials') +parser.add_argument('--deltaT', type=float, default=1000., + help='Time Window in seconds') +parser.add_argument('--ntrials', type=int, default = 1000, + help='Trials') +parser.add_argument('--start', type=str, required=False, + default='2025-06-01', + help="Start time of the analysis in ISO format") +parser.add_argument("--bkg", default=[0.206], type=float, nargs='+', + help="Expected background rates in mHz (default 6.4, 4.6)") +parser.add_argument('--seed', default=1, type=int, + help='Unique seed for running on the cluster') +parser.add_argument('--nside', default=64, type=int, + help='Skymap nside to scan') +parser.add_argument('--outdir',type=str, default=os.environ.get('FAST_RESPONSE_OUTPUT', './'), + help='Output directory to save npz (default = FAST_RESPONSE_OUTPUT env variable or cwd)') +parser.add_argument('--dataset', default="DNN", type=str, + help='Dataset(s) to include, joined by + signs') +parser.add_argument('--fix_index', action='store_true', help='Fix the spectral index during fitting') +parser.add_argument('--index', type=float, default=None, + help="Spectral index to assume in injected hypothesis") +args = parser.parse_args() + +followups = [] +if 'GFU' in args.dataset: + followups.append(GFUFollowup) +if 'Greco' in args.dataset: + followups.append(GrecoFollowup) +if 'DNN' in args.dataset: + followups.append(DNNOnlineFollowup) +MultiFollowup._followups = followups + +for attr in ['index', 'fix_index']: + setattr(MultiFollowup, f'_{attr}', getattr(args, attr)) +MultiFollowup._float_index = not MultiFollowup._fix_index + + + +outdir = os.path.join(args.outdir, 'alert_precomputed_trials/') +if not os.path.exists(outdir): + os.makedirs(outdir, exist_ok=True) + +# use the signal window to select events, as much as the desired time window +# TODO could this be better to include more events? + +start_mjd = Time(args.start).mjd +stop_mjd = start_mjd + (args.deltaT / 86400.) +start_iso = args.start +stop_iso = Time(stop_mjd, format='mjd').iso +deltaT = args.deltaT / 86400. + +#skymap required for initialization, but not used here +f = MultiFollowup('Precompute_trials_test', 0,0, + start_iso, stop_iso, save=False) +assert f.scramble +f._ncpu = 2 + +for i, enum in enumerate(f.llh._samples): + _llh = f.llh._samples[enum] + print(i, _llh.nbackground, _llh.on_livetime*86400) + print(_llh.nbackground/_llh.on_livetime/86400*1000) + print(_llh.nbackground/args.deltaT*1000) + # set background to required rate to simulate + #_llh.nbackground = args.bkg[i]*args.deltaT/1000. + _llh.nbackground = args.bkg[i]*_llh.on_livetime*86400/1000. + print(_llh.nbackground) +#inj = f.initialize_injector(gamma=2.5) #just put this here to initialize f.spatial_prior +#print f.llh.nbackground +#results_array = [] + +# per original seed of the job, carve out a block of unique seeds, one per trial +# this permits reproducing these trials later +seeds = range(args.seed, args.seed + args.ntrials) + +npix = hp.nside2npix(args.nside) +shape = (args.ntrials, npix) +maps = sparse.lil_matrix(shape, dtype=float) +for jj, seed in enumerate(seeds): + val = f.llh.scan(0.0, 0.0, scramble=True, seed = seed, + #spatial_prior = f.spatial_prior, + time_mask = [deltaT / 2., (start_mjd + stop_mjd) / 2.], + pixel_scan = [args.nside, 3.0], inject = None) + if val['TS'] is not None: + dtype = [('ts',float),('pixel',float)] + results = np.empty((val['TS'].size,), dtype=dtype) + pixels = hp.ang2pix(args.nside, np.pi/2. - val['dec'], val['ra']) + maps[jj, pixels] = val['TS'] +print("DONE") +hp_sparse = maps.tocsr() + +# TODO define this format centrally so it doesn't need to be copied +rates_str = '_'.join([f'{_rate:.2f}' for _rate in args.bkg]) +outfilename = '_'.join([ + f'precomputed_trials_delta_t_{args.deltaT:.2e}', + f'nside_{args.nside}', + f'index_{f._index}', + f'{rates_str}_mHz', + f'seed_{args.seed}', + f'low_stats.npz', + ] +) +outfilepath = os.path.join(outdir, outfilename) +sparse.save_npz(outfilepath, hp_sparse) +print("Saved to {}".format(outfilepath)) \ No newline at end of file diff --git a/fast_response/reports/ReportGenerator.py b/fast_response/reports/ReportGenerator.py index 9b34397..ad59a2c 100644 --- a/fast_response/reports/ReportGenerator.py +++ b/fast_response/reports/ReportGenerator.py @@ -355,6 +355,9 @@ def generate_report(self): precision=0).iso self.ontime['time_query']=now.strftime('%Y-%m-%d %H:%M:%S') + + # TODO add column per realtime data stream, right now it's only GFU. + # TODO for a start, add feature to realtime_tools to get times of online Greco filter (no reconstruction yet) self.query_events=icecube.realtime_tools.live.get_events( self.ontime['stream'], self.ontime['time_start'], @@ -375,7 +378,6 @@ def generate_report(self): s = self.source - dataset = Datasets[self.analysis._dataset] # Make rate plots and put them in analysis directory make_rate_plots( self.time_window, @@ -466,15 +468,24 @@ def generate_report(self): ("Time Window",r"{:1.1f}s".format(s["realtime"])), ] ) + + # retrieve metadata of used dataset(s) for the table + datasets = [Datasets[_ds] for _ds in self.analysis.datasets] - self.write_table(f,"skylabtable",[],[ + # + skylabtable = [ ("Skylab Version", skylab.__version__), ("IceTray Path", str(icetray.__path__).replace('_', '\_')), ("Created by", expanduser('~')[6:]), + ] + for dataset in datasets: + skylabtable.extend([ ("Dataset Used", str(dataset.subdir).replace('_',' ')), ("Dataset details", str(dataset.name)[:80]), - ("", str(dataset.name)[80:]) - ]) + ("", str(dataset.name)[80:]), + ]) + + self.write_table(f,"skylabtable",[],skylabtable) r1=[] r2=[] @@ -527,17 +538,29 @@ def generate_report(self): if self.analysis._float_index: if self.analysis.ts > 0.: - self.write_table( - f, - "results", - [], - [("$n_s$", "{:1.3f}".format(self.analysis.ns)), - ("$TS$", "{:1.3f}".format(self.analysis.ts)), - ("$\gamma$", f"{self.analysis.gamma:.2f}"), - ("$p-value$", "{:1.4f}".format(self.analysis.p)), - ("best-fit RA", "{:3.2f}\degree".format(np.rad2deg(self.analysis.skymap_fit_ra))), - ("best-fit dec", "{:3.2f}\degree".format(np.rad2deg(self.analysis.skymap_fit_dec)))] - ) + if self.source_type == 'skymap': + self.write_table( + f, + "results", + [], + [("$n_s$", "{:1.3f}".format(self.analysis.ns)), + ("$TS$", "{:1.3f}".format(self.analysis.ts)), + ("$\gamma$", f"{self.analysis.gamma:.2f}"), + ("$p-value$", "{:1.4f}".format(self.analysis.p)), + ("best-fit RA", "{:3.2f}\degree".format(np.rad2deg(self.analysis.skymap_fit_ra))), + ("best-fit dec", "{:3.2f}\degree".format(np.rad2deg(self.analysis.skymap_fit_dec)))] + ) + elif self.source_type == 'PS': + self.write_table( + f, + "results", + [], + [("$n_s$", "{:1.3f}".format(self.analysis.ns)), + ("$TS$", "{:1.3f}".format(self.analysis.ts)), + ("$\gamma$", f"{self.analysis.gamma:.2f}"), + ("$p-value$", "{:1.4f}".format(self.analysis.p)), + ] + ) else: self.write_table( f, diff --git a/fast_response/scripts/dnnc_prior_sensitivity.py b/fast_response/scripts/dnnc_prior_sensitivity.py new file mode 100755 index 0000000..1cf8dc4 --- /dev/null +++ b/fast_response/scripts/dnnc_prior_sensitivity.py @@ -0,0 +1,260 @@ +'''Script to calculate sensitivities of analyses defined +in the fast_response framework. + + Author: Christoph Raab + Date: 2026 + ''' + +from fast_response.MultiGWFollowup import DNNOnlineFollowup as MultiFollowup +import argparse +import subprocess +import warnings +import fast_response.web_utils as web_utils +import logging as log +import pyfiglet +import os +import numpy as np +import scipy + +from skylab import utils + +from astropy.time import Time, TimeDelta +import datetime + +# FIXME put this into Skylab instead! Ghastly to have it here +def fit_source(followup, seed, **kwargs): + llh = followup.llh + spatial_prior = followup.spatial_prior + injector = followup.inj + # fix random seed of injector + injector.set_rng_seed(seed) + + # this injects _fewer_ events than requested on purpose + # we want to pretend like we injected the full number + mean_signal = kwargs.pop("mean_signal", 0) + if kwargs.pop("poisson", True): + ni = np.random.poisson(mean_signal) + else: + ni = mean_signal + _, inject = injector.sample(ni, poisson=False) + + # fix random seed of likelihood + llh.set_rng_seed(seed) + + val = llh.scan( + 0.0,0.0, scramble=True, spatial_prior=spatial_prior, + time_mask = [followup.duration/2., followup.centertime], + pixel_scan=[followup.nside, followup._pixel_scan_nsigma], + inject=inject, + ) + if val.size == 0: + ts = 0 + else: + ts = val['TS_spatial_prior_0'].max() + return ni, ts + + +def do_trials(followup, n_iter=10, **kwargs): + print(".") + dtype = [("n_inj", int), ("TS", float)] + + rng = np.random.RandomState(followup.llh.do_trials_seed) + + results = [fit_source(followup, rng.randint(2**32), **kwargs) + for i in range(n_iter)] + + trials = np.empty((n_iter, ), dtype=dtype) + + for i in range(n_iter): + trials["n_inj"][i] = results[i][0] + trials["TS"][i] = results[i][1] + return trials + +def weighted_sensitivity(followup, bg_ts, alpha, beta, n_iter=100, trials=[]): + ts = np.percentile(bg_ts, 100*(1-alpha)) + for _mu in range(1, 5): + _trials = do_trials(followup, n_iter=n_iter, mean_signal=_mu, poisson=True) + trials.append(_trials) + # bg_trials = np.empty((bg_ts.size,), dtype=trials[0].dtype) + # for i, _ts in enumerate(bg_ts): + # bg_trials["n_inj"][i] = 0 + # bg_trials["TS"][i] = _ts + # trials.append(bg_trials) + trials = np.concatenate(trials) + + bounds = np.percentile( + trials["n_inj"][trials["n_inj"] > 0], + q=[followup.llh._ub_perc, 100. - followup.llh._ub_perc]) + + if bounds[0] == 1: + bounds[0] = np.count_nonzero(trials["n_inj"] == 1) /\ + np.sum(trials["n_inj"] < 2) + + def residual(n): + return np.log10((utils.poisson_percentile( + n, trials["n_inj"], trials["TS"], ts)[0] - beta)**2) + + seed = np.argmin([residual(n) for n in np.arange(0., bounds[-1])]) + + xmin, fmin, success = scipy.optimize.fmin_l_bfgs_b( + residual, [seed], bounds=[bounds], approx_grad=True) + + mu = xmin.item() + b, b_err = utils.poisson_percentile( + mu, trials["n_inj"], trials["TS"], ts) + + print( + "Best estimate: mu = {0:.2f} ({1:.2%} +/- {2:.2%})".format( + mu, b, b_err)) + + flux = followup.inj.mu2flux(mu) + print("mu = {0:.2f}, flux = {1:.2e}".format(mu, flux)) + + values = [(alpha, ts, beta, mu, flux)] + result = np.empty( + (1,), + dtype=[(k, float) for k in ["alpha", "TS", "beta", "mu", "flux"]]) + + result.flat = values + return result, trials + + +def run_trials(args): + results = ['prior_sensitivity', 'prior_trials', 'prior_weights'] + outdir = {} + for _r in results: + _outdir = os.path.join(args.out, _r) + if _r == 'prior_sensitivity': + _outdir = os.path.join(_outdir, f"{args.alpha}_{args.beta}") + if not os.path.exists(_outdir): + os.makedirs(_outdir) + outdir[_r] = _outdir + + # followups = [] + # if 'GFU' in args.dataset: + # followups.append(GFUFollowup) + # if 'DNN' in args.dataset: + # followups.append(DNNOnlineFollowup) + # MultiFollowup._followups = followups + + for attr in ['index', 'fix_index']: + setattr(MultiFollowup, f'_{attr}', getattr(args, attr)) + MultiFollowup._float_index = not MultiFollowup._fix_index + + + f = MultiFollowup( + args.name, args.skymap, args.start, + args.stop, extension=args.extension, + skipped=args.skip_events, save=False, + ) + assert f.scramble + + print(str(f)) + print(f.dataset) + print(args.out) + print(f.outdir) + print(f.save_output, + f.analysisid, + f.analysispath, + ) + if args.eband is None: + f.initialize_injector() + else: + halfwidth = np.sqrt(10**args.ewidth) + f.initialize_injector(e_range=(args.eband/halfwidth, args.eband*halfwidth)) + + dataset_string = '+'.join(sorted(f.datasets)) + skymap_string = os.path.split(args.skymap)[1].split('.')[0] + if args.eband is None: + label = f'{skymap_string}_{args.duration}_{f._index}_{dataset_string}' + else: + label = f'{skymap_string}_{args.duration}_{args.eband:.0f}_{dataset_string}' + outpath = {_r:os.path.join(outdir[_r], f'{label}.npy') for _r in results} + + if os.path.exists(outpath['prior_trials']): + trials = [np.load(outpath['prior_trials'])] + print(f'loading {trials[0].size} previous trial') + else: + trials = [] + + month = Time(args.start).datetime.month + # FIXME Horrible! do better than hardcoding this + rate = [ + 0.22, 0.22, 0.22, 0.21, + 0.21, 0.21, 0.21, 0.2 , + 0.22, 0.23, 0.22, 0.23, + ][month - 1] + bg_ts = f.load_background_trials(rate=rate, ntrials=args.ntrials) + + sensitivity, trials = weighted_sensitivity(f, + bg_ts, args.alpha, args.beta, + n_iter=args.n_iter, + trials=trials, + ) + + + # FIXME this won't work as it will fit a source at a fixed position + # sensitivity, trials, weights = f.llh.weighted_sensitivity(args.alpha, args.beta, f.inj, + # eps=args.eps, + # n_iter=args.n_iter, + # n_bckg=args.ntrials, + # trials = trials, + # ) + # print(sensitivity) + + + np.save(outpath['prior_sensitivity'], sensitivity) + np.save(outpath['prior_trials'], trials) + + +if __name__ == "__main__": + warnings.filterwarnings("ignore") + log.basicConfig(level=log.ERROR) + + parser = argparse.ArgumentParser(description='Fast Response Analysis sensitivity') + parser.add_argument('--name', type=str, default="test inherit sensitivity", + help='Name of the source (do not use underscores or LaTeX will be mad!)') + parser.add_argument("--skymap") + parser.add_argument("--alpha", default=0.5, type=float) + parser.add_argument("--beta", default=0.9, type=float) + parser.add_argument('--fix_index', action='store_true', help='Fix the spectral index during fitting') + parser.add_argument('--duration', default=10, type=float, + help='Duration of followup [days]') + parser.add_argument('--start', type=str, required=False, + default=Time(58933.0 - 100, format='mjd').iso, + help="Start time of the analysis in ISO format") + parser.add_argument('--extension', type=float, default=None, + help="Source extension in degrees") + parser.add_argument('--index', type=float, default=None, + help="Spectral index to assume in injected hypothesis") + parser.add_argument('--eband', default=None, type=float, + help='Determine differential sensitivity within an energy band', + ) + parser.add_argument('--ewidth', default=0.5, type=float, + help='Energy band width in decades', + ) + parser.add_argument('--skip-events', default=None, + type= lambda z:[ tuple(int(y) for y in x.split(':')) for x in z.split(',')], + help="Event to exclude from the analyses, eg." + "Example --skip-events=127853:67093193") + parser.add_argument('--ntrials', default=10000, type=int, + help="Number of background trials to perform") + parser.add_argument('--n_iter', default=10, type=int, + help="Number of signal trials per per sensitivity iteration") + parser.add_argument('--eps', default=0.05, type=float, + help="Cutoff uncertainty for the sensitivity calculation") + parser.add_argument('--out', default='/data/user/chraab/Output/fast_response/multisample/inherit/') + parser.add_argument('--dataset', default=None, type=str, + help='Dataset(s) to include, joined by + signs') + + args = parser.parse_args() + + if '_' in args.name: + print('Warning: underscores in source name cause LaTeX to fail!') + tic = datetime.datetime.now() + + args.stop = (Time(args.start, format='iso') + TimeDelta(args.duration, format='jd')).iso + + run_trials(args) + toc = datetime.datetime.now() + print((toc - tic).total_seconds()) diff --git a/fast_response/scripts/multisample_sensitivity.py b/fast_response/scripts/multisample_sensitivity.py new file mode 100755 index 0000000..197be0d --- /dev/null +++ b/fast_response/scripts/multisample_sensitivity.py @@ -0,0 +1,148 @@ +'''Script to calculate sensitivities of analyses defined +in the fast_response framework. + + Author: Christoph Raab + Date: 2025 + ''' + +from fast_response.MultiExternalFollowup import MultiFollowup, GFUFollowup, GrecoFollowup, DNNFollowup, DNNOnlineFollowup, DNNIceManFollowup +import argparse +import subprocess +import warnings +import fast_response.web_utils as web_utils +import logging as log +import pyfiglet +import os +import numpy as np + +from astropy.time import Time, TimeDelta +import datetime + +def calculate_sensitivity(args): + + results = ['sensitivity', 'trials', 'weights'] + outdir = {} + for _r in results: + _outdir = os.path.join(args.out, _r) + if _r == 'sensitivity': + _outdir = os.path.join(_outdir, f"{args.alpha}_{args.beta}") + if not os.path.exists(_outdir): + os.makedirs(_outdir) + outdir[_r] = _outdir + + followups = [] + if 'GFU' in args.dataset: + followups.append(GFUFollowup) + if 'Greco' in args.dataset: + followups.append(GrecoFollowup) + if 'DNN' in args.dataset: + followups.append(DNNOnlineFollowup) + if "IceMan" in args.dataset: + followups.append(DNNIceManFollowup) + MultiFollowup._followups = followups + + for attr in ['index', 'fix_index']: + setattr(MultiFollowup, f'_{attr}', getattr(args, attr)) + MultiFollowup._float_index = not MultiFollowup._fix_index + + + f = MultiFollowup( + args.name, args.ra, args.dec, args.start, + args.stop, extension=args.extension, + skipped=args.skip_events, save=False, + ) + assert f.scramble + + print(str(f)) + print(f.dataset) + print(args.out) + print(f.outdir) + print(f.save_output, + f.analysisid, + f.analysispath, + ) + if args.eband is None: + f.initialize_injector() + else: + halfwidth = np.sqrt(10**args.ewidth) + f.initialize_injector(e_range=(args.eband/halfwidth, args.eband*halfwidth)) + + + dataset_string = '+'.join(sorted(f.datasets)) + if args.eband is None: + label = f'{f.dec:.03f}_{args.duration}_{f._index}_{dataset_string}' + else: + label = f'{f.dec:.03f}_{args.duration}_{args.eband:.0f}_{dataset_string}' + outpath = {_r:os.path.join(outdir[_r], f'{label}.npy') for _r in results} + if os.path.exists(outpath['trials']): + trials = np.load(outpath['trials']) + print(f'loading {trials.size} previous trial') + else: + trials = None + + sensitivity, trials, weights = f.llh.weighted_sensitivity(args.alpha, args.beta, f.inj, + src_ra = f.ra, src_dec = f.dec, + eps=args.eps, + n_iter=args.n_iter, + n_bckg=args.ntrials, + trials = trials, + ) + np.save(outpath['trials'], trials) + np.save(outpath['sensitivity'], sensitivity) + print(sensitivity) + + +if __name__ == "__main__": + warnings.filterwarnings("ignore") + log.basicConfig(level=log.ERROR) + + parser = argparse.ArgumentParser(description='Fast Response Analysis sensitivity') + parser.add_argument('--name', type=str, default="test inherit sensitivity", + help='Name of the source (do not use underscores or LaTeX will be mad!)') + parser.add_argument('--ra', default=None, type=float, + help='Right ascension (in degrees)') + parser.add_argument('--dec', default=None, type=float, + help='Declination (in degrees)') + parser.add_argument("--alpha", default=0.5, type=float) + parser.add_argument("--beta", default=0.9, type=float) + parser.add_argument('--fix_index', action='store_true', help='Fix the spectral index during fitting') + parser.add_argument('--duration', default=10, type=float, + help='Duration of followup [days]') + parser.add_argument('--start', type=str, required=False, + default=Time(58933.0 - 100, format='mjd').iso, + help="Start time of the analysis in ISO format") + parser.add_argument('--extension', type=float, default=None, + help="Source extension in degrees") + parser.add_argument('--index', type=float, default=None, + help="Spectral index to assume in injected hypothesis") + parser.add_argument('--eband', default=None, type=float, + help='Determine differential sensitivity within an energy band', + ) + parser.add_argument('--ewidth', default=0.5, type=float, + help='Energy band width in decades', + ) + parser.add_argument('--skip-events', default=None, + type= lambda z:[ tuple(int(y) for y in x.split(':')) for x in z.split(',')], + help="Event to exclude from the analyses, eg." + "Example --skip-events=127853:67093193") + parser.add_argument('--ntrials', default=1000, type=int, + help="Number of background trials to perform") + parser.add_argument('--n_iter', default=10, type=int, + help="Number of signal trials per per sensitivity iteration") + parser.add_argument('--eps', default=0.05, type=float, + help="Cutoff uncertainty for the sensitivity calculation") + parser.add_argument('--out', default='/data/user/chraab/Output/fast_response/multisample/inherit/') + parser.add_argument('--dataset', default=None, type=str, + help='Dataset(s) to include, joined by + signs') + + args = parser.parse_args() + + if '_' in args.name: + print('Warning: underscores in source name cause LaTeX to fail!') + tic = datetime.datetime.now() + + args.stop = (Time(args.start, format='iso') + TimeDelta(args.duration, format='jd')).iso + + calculate_sensitivity(args) + toc = datetime.datetime.now() + print((toc - tic).total_seconds()) diff --git a/setup.py b/setup.py index 9b81fbc..ab4f37a 100644 --- a/setup.py +++ b/setup.py @@ -19,18 +19,29 @@ ], python_requires='>=3.1', install_requires=[ - 'astropy==2.0.16', - 'healpy==1.13.0', - 'matplotlib==2.2.5', - 'numpy==1.21.5', - 'pandas==1.3.5', - 'pyfiglet==0.8.post1', - 'python-dateutil==2.8.1', - 'pyzmq==19.0.1', - 'scipy==1.2.3', - 'seaborn==0.9.1', - 'zmq==0.0.0', - 'py27hash==1.0.2', - 'pygcn==1.1.2' +# 'astropy==2.0.16', + 'astropy', + #'healpy==1.13.0', + 'healpy', # already from CVMFS + 'matplotlib', + #'matplotlib==2.2.5', +# 'numpy==1.21.5', + 'numpy', + #'pandas==1.3.5', + #'pyfiglet==0.8.post1', + #'python-dateutil==2.8.1', + #'pyzmq==19.0.1', + #'scipy==1.2.3', + #'seaborn==0.9.1', + 'seaborn', + #'zmq==0.0.0', + #'py27hash==1.0.2', + #'pygcn==1.1.2' + 'pyfiglet', + #'pyzmq', + 'pygcn', + #'py27hash', + 'python-dateutil', + 'meander', ] )