diff --git a/fast_response/AlertFollowup.py b/fast_response/AlertFollowup.py index f5747783..7fec44b9 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 4c8be36b..bc97eef0 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 @@ -32,6 +35,7 @@ from . import sensitivity_utils from . import plotting_utils from .reports import FastResponseReport +from .precomputed_background import glob_precomputed_trials_multi as pt mpl.use('agg') current_palette = sns.color_palette('colorblind', 10) @@ -55,15 +59,32 @@ class FastResponseAnalysis(object): _angScale = 2.145966 _llh_seed = 1 _season_names = [f"IC86, 201{y}" for y in range(1, 10)] - _nb_days = 10 + _nb_days = 10 # BG window around analysis window used to estimate BG rate. _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 +115,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 @@ -118,6 +140,8 @@ def __init__(self, name, tstart, tstop, skipped=None, seed=None, self.exp = None self.llh = self.initialize_llh(skipped=skipped, scramble=self.scramble) self.inj = None + # not yet unblinded or loaded BG trials: + self.ts, self.tsd = None, None @property def dataset(self): @@ -127,6 +151,11 @@ def dataset(self): def dataset(self, x): self._dataset = x + @property + def datasets(self): + """Returns the datasets used""" + return [self.dataset] + @property def index(self): """Returns the spectral index""" @@ -141,8 +170,44 @@ def llh_seed(self): return self._llh_seed @llh_seed.setter def llh_seed(self, x): + # FIXME the llh does not get initialized again + # so this does not repeat the initial scramble + # FIXME what would be a consistent way to set_rng_seed on the LLH's if they exist? 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: + # TODO this relies on the assumption the archival GFU is equal to the default + # 1a) legacy behaviour, these are added to GFU archival + 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 +265,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,9 +286,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 self.grl = grl @@ -273,6 +363,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 +598,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 +616,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 +642,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 +700,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 +726,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 +745,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 +803,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 +857,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 +869,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 +890,20 @@ 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}', + '*', # individual analyses will need to specify here whether they glob them on the fly or have a single file + ]) 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) @@ -767,6 +916,7 @@ def __init__(self, name, skymap_path, tstart, tstop, skipped=None, seed=None, probs = hdf_data['PROBDENSITY'][()] area = 4*np.pi/probs.size probs *= area + # FIXME the below should be redundant with self.format_skymap(skymap) skymap = hp.pixelfunc.ud_grade( probs, self._nside, power=-2, order_in='NESTED', order_out='RING' @@ -779,6 +929,15 @@ def __init__(self, name, skymap_path, tstart, tstop, skipped=None, seed=None, self.ra, self.dec, self.extension = None, None, extension self.save_items['skymap'] = skymap + # FIXME initialize_llh in GWFollowup mostly just waits (we do this different) + # BUT + # - does llh.set_temporal_model(box) instead of passing it to PSLLH constructor which does the same + # - does NOT set timescramble=True + # - (also not through the update_timescramble kwarg) + # ...so then does it even use the time scrambling? + # ...(should it? maybe I'm getting something wrong about the ontime) + # => but I guess it shouldn't matter because GWFollowup.llh is never used for trials. + def __str__(self): int_str = super().__str__() int_str += ' '*10 + 'Skymap file:' + self.skymap_path @@ -787,7 +946,7 @@ def __str__(self): def format_skymap(self, skymap): r"""Method to up or downgrade nside of a skymap to - the nside used in the analysis + the nside used in the analysis. Normalizes the sum to 1. Parameters ----------- @@ -799,6 +958,7 @@ def format_skymap(self, skymap): skymap: array Healpix skymap, with correct nside for use in FRA """ + # FIXME ud_grade vs. get_interp_val? if hp.pixelfunc.get_nside(skymap) != self._nside: skymap = hp.pixelfunc.ud_grade(skymap, self._nside, power=-2) skymap = skymap/skymap.sum() @@ -844,6 +1004,11 @@ def run_background_trials(self, ntrials=1000): ntrials: int number of trials to run (default 1000) """ + # If they are already loaded, no need to run new ones; maybe being called from calc_pvalue + if self.tsd is not None: + if self.tsd.size >= ntrials: + return self.tsd[:ntrials] + # Else assume we're calling this method because we WANT to run new ones. tsd = [] spatial_prior = SpatialPrior(self.skymap, containment = self._containment, allow_neg=self._allow_neg) @@ -872,27 +1037,99 @@ def run_background_trials(self, ntrials=1000): self.tsd = tsd self.save_items['tsd'] = tsd + def load_background_trials(self, ntrials=None, rate=None, month=None) -> np.ndarray: + """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) + Generalizing the method from GWFollowup.run_background_trials used for durations > 1 day, except + - relying on the precomputed scans to have been concatenated into one file already + - setting TS=0 in empty trials according to the new convention + + + 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 + + Returns + ------- + tsd : np.ndarray + Background TS values + """ + 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(map(str, Path(self._bg_dir).glob(filename))) + if not bg_files: + raise FileNotFoundError(f"Did not find precomputed bg {filename} in {self._bg_dir}") + glob_file = pt.get_glob_file(bg_files[0]) + + # Load sparse matrix of background scans + if os.path.exists(glob_file): + pre_ts_array = pt.load_maps(glob_file) + else: + pre_ts_array = pt.concatenate_maps(bg_files, self.nside) + pt.save_maps(pre_ts_array, glob_file) + 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 + # 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] = 0 # new convention: 0 for empty trial + 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): 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 = np.copy(self.llh_exp[t_mask]) + # Using the new llh_exp property exp_theta = 0.5*np.pi - events['dec'] exp_phi = events['ra'] exp_pix = hp.ang2pix(self.nside, exp_theta, exp_phi) overlap = np.isin(exp_pix, self.ipix_90) events = events[overlap] - # print nearby events, as a check (if needed) - # msk1 = (self.llh.exp[t_mask]['ra'] < (self.skymap_fit_ra+np.radians(5)))*(self.llh.exp[t_mask]['ra'] > (self.skymap_fit_ra-np.radians(5))) - # msk2 = (self.llh.exp[t_mask]['dec'] < (self.skymap_fit_dec+np.radians(5)))*((self.llh.exp[t_mask]['dec'] > self.skymap_fit_dec-np.radians(5))) - # msk3 = msk1*msk2 - # print('Nearby events:') - # print("[run, event, ra, dec, sigma, logE, time]") - # for e in self.llh.exp[t_mask][msk3]: print([e[k] for k in ['run', 'event', 'ra', 'dec', 'sigma', 'logE', 'time']]) - if len(events) == 0: coincident_events = [] else: @@ -908,14 +1145,14 @@ def find_coincident_events(self): self.coincident_events = coincident_events self.save_items['coincident_events'] = coincident_events - def unblind_TS(self, custom_events=None): + def unblind_TS(self, scramble: bool=False): r""" Unblind TS, either sky scan for spatial prior, or just at one location for a point source Parameters ----------- - custom_events: array - specific events for use in scan (UNUSED) + scramble: bool + Let the LLH scramble with its current RNG seed before unblinding Returns ----------- @@ -936,7 +1173,10 @@ def unblind_TS(self, custom_events=None): t1 = time.time() print("Starting scan") val = self.llh.scan( - 0.0,0.0, scramble = False, spatial_prior=spatial_prior, + 0.0,0.0, + # if scrambling, llh.scan() takes a seed from its own kwargs + scramble = scramble, seed=self.llh_seed, + spatial_prior=spatial_prior, time_mask = [self.duration/2., self.centertime], pixel_scan=[self.nside, self._pixel_scan_nsigma] ) @@ -963,7 +1203,7 @@ def unblind_TS(self, custom_events=None): self.scanned_pixels = hp.ang2pix( self.nside, np.pi/2. - val['dec'], val['ra'] ) - except Exception as e: + except Exception as e: # TODO be more specific, eg empty scan print(e) ts, ns = 0., 0. if self._float_index: @@ -993,6 +1233,40 @@ def unblind_TS(self, custom_events=None): else: return ts, ns + # Moved from GWFollowup: will work also for Multi*Followups due to llh.scan using the `enum` field in custom_events. + def per_event_scan(self, custom_events): + """Runs the all-sky scan for only one (or certain) events on the sky + + Parameters + ------------ + custom_events: masked array + Ontime event(s) loaded in Skylab to use when running the all sky scan. + Must have the field "enum" if used with a Multi*Followup. + + Returns + ----------- + ts: float + best-fit TS using only this event + p: float + p-value for the given event(s) + + """ + from skylab.priors import SpatialPrior + + spatial_prior = SpatialPrior(self.skymap, containment = self._containment, allow_neg=self._allow_neg) + val = self.llh.scan( + 0.0,0.0, scramble = False, spatial_prior=spatial_prior, + time_mask = [self.duration/2., self.centertime], + pixel_scan=[self.nside, self._pixel_scan_nsigma], + custom_events=custom_events + ) + if val['TS'].size == 0: + ts = -1.*np.inf # FIXME why is the convention different than during unblind_TS? + else: + ts = val['TS_spatial_prior_0'].max() + p = np.count_nonzero(self.tsd >= ts) / float(len(self.tsd)) + return ts, p + def upper_limit(self): """ UPPER LIMIT WITH SPATIAL PRIOR NOT YET IMPLEMENTED """ @@ -1085,7 +1359,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 +1397,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) @@ -1193,7 +1469,7 @@ def initialize_injector(self, e_range=(0., np.inf)): temporal_model=self.llh.temporal_model) self.inj = inj - def unblind_TS(self): + def unblind_TS(self, scramble: bool=False): r""" Unblind TS at one location for a point source Returns @@ -1205,28 +1481,57 @@ def unblind_TS(self): """ # 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) + ts, ns = self.llh.fit_source(src_ra=self.ra, src_dec=self.dec, scramble=scramble) params = ns.copy() params.pop('nsignal') self.ns_params = params 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 """ + # TODO spatial weight means that cascades are never "coincident" + # (different story in PriorFollowup - then resolution plays no role) + # can have an overfluctuation in 30 days without a coincidence! + 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 +1636,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 +1651,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 +1696,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/GWFollowup.py b/fast_response/GWFollowup.py index 3dbd6f0f..bfefdd9a 100644 --- a/fast_response/GWFollowup.py +++ b/fast_response/GWFollowup.py @@ -38,6 +38,9 @@ class GWFollowup(PriorFollowup): _season_names = ['IC86, 2017', 'IC86, 2018', 'IC86, 2019'] _nb_days = 5. _ncpu = 10 + _background_days = 6. # Background window, can adjust per dataset + _sens_dir = '/data/ana/analyses/NuSources/2023_realtime_gw_analysis/' \ + + 'fast_response/ps_sensitivities' def __init__(self, name, skymap_path, tstart, tstop, skipped=None, seed=None, outdir=None, save=True, extension=None): @@ -203,7 +206,7 @@ def initialize_llh(self, skipped=None, scramble=False): t0 = Time(datetime.datetime.utcnow()).mjd if self.stop > t0 + 60./86400.: - self.get_data(livestream_start=self.start-6., livestream_stop=self.start) + self.get_data(livestream_start=self.start-self._background_days, livestream_stop=self.start) print('Loading off-time data') elif self.exp is None: dset = Datasets[self.dataset] @@ -591,6 +594,23 @@ def upper_limit(self): self.save_items['sens_range'] = sens_range self.make_dec_pdf() + def load_ps_sensitivities(self): + sens_pickle = f'{self._sens_dir}/ps_sensitivities_deltaT_{self.duration*86400.:.2e}s.pkl' + from os.path import isfile + if isfile(sens_pickle): + with open(sens_pickle, 'rb') as f: + saved_sens=pickle.load(f) + dec_range=saved_sens['dec'] + sens=saved_sens['sens_flux'] + return dec_range, sens + + sens_npy = f"{self._sens_dir}/deltaT_{self.duration*86400:.2e}_index_{self.index:.1f}.npy" + if isfile(sens_npy): + saved_sens = np.load(sens_npy) + return saved_sens['dec'], saved_sens['flux'] + raise FileNotFoundError(f"Could find neither {sens_pickle} nor {sens_npy}") + + def ps_sens_range(self): r""" Compute minimum and maximum sensitivities within @@ -604,20 +624,7 @@ def ps_sens_range(self): highest sensitivity within dec range """ - sens_dir = '/data/ana/analyses/NuSources/2023_realtime_gw_analysis/' \ - + 'fast_response/ps_sensitivities' - - with open(f'{sens_dir}/ps_sensitivities_deltaT_{self.duration*86400.:.2e}s.pkl','rb') as f: - saved_sens=pickle.load(f) - dec_range=saved_sens['dec'] - sens=saved_sens['sens_flux'] - #dec_range = np.linspace(-85,85,35) - #sens = [1.15, 1.06, .997, .917, .867, .802, .745, .662, - # .629, .573, .481, .403, .332, .250, .183, .101, - # .035, .0286, .0311, .0341, .0361, .0394, .0418, - # .0439, .0459, .0499, .0520, .0553, .0567, .0632, - # .0679, .0732, .0788, .083, .0866] - + dec_range, sens = self.load_ps_sensitivities() src_theta, src_phi = hp.pix2ang(self.nside, self.ipix_90) src_dec = np.pi/2. - src_theta src_dec = np.unique(src_dec) @@ -708,21 +715,10 @@ def make_dec_pdf(self): sinDec_bins = np.linspace(-1,1,30) bin_centers = (sinDec_bins[:-1] + sinDec_bins[1:]) / 2 - - sens_dir = '/data/ana/analyses/NuSources/2023_realtime_gw_analysis/' \ - + 'fast_response/ps_sensitivities' - - with open(f'{sens_dir}/ps_sensitivities_deltaT_{self.duration*86400.:.2e}s.pkl','rb') as f: - saved_sens=pickle.load(f) - dec_range=np.sin(saved_sens['dec']*np.pi/180) - sens=saved_sens['sens_flux'] - - #dec_range = np.linspace(-1,1,35) - #sens = [1.15, 1.06, .997, .917, .867, .802, .745, .662, - # .629, .573, .481, .403, .332, .250, .183, .101, - # .035, .0286, .0311, .0341, .0361, .0394, .0418, - # .0439, .0459, .0499, .0520, .0553, .0567, .0632, - # .0679, .0732, .0788, .083, .0866] + + dec_range, sens = self.load_ps_sensitivities() + sindec_range = np.sin(np.deg2rad(dec_range)) + sens = np.array(sens) pixels = np.arange(len(self.skymap)) @@ -748,7 +744,7 @@ def make_dec_pdf(self): ax2 = ax1.twinx() # instantiate a second axes that shares the same x-axis ax2.set_ylabel('E$^2$F (GeVcm$^2$)') # we already handled the x-label with ax1 - ax2.plot(dec_range, sens, color='C1', label='PS Sensitivity') + ax2.plot(sindec_range, sens, color='C1', label='PS Sensitivity') ax2.set_yscale('log') ax2.set_xlim(-1,1) ax2.tick_params(axis='y') diff --git a/fast_response/MultiAlertFollowup.py b/fast_response/MultiAlertFollowup.py new file mode 100644 index 00000000..ebc78511 --- /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 00000000..94fc23d5 --- /dev/null +++ b/fast_response/MultiExternalFollowup.py @@ -0,0 +1,60 @@ +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 + + +#or any other number of definitions + +class DNNOnlineFollowup(PointSourceFollowup): + _dataset = "DNNCascadesOnline_v001p01" + _season_names = [f"IC86, 20{y:02d}" for y in range(24, 25+1)] + _floor = np.radians(1.5) # can change this! + _jitter = 3. # common default for DNN analyses + _background_days = 60. + _nb_days = 60. + + +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 = False + _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 00000000..e7a25bee --- /dev/null +++ b/fast_response/MultiFastResponseAnalysis.py @@ -0,0 +1,476 @@ +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 +from collections import defaultdict + +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)}') + + # We share an LLH seed + @property + def llh_seed(self): + return self._llh_seed + + # when setting a new one, broadcast it + @llh_seed.setter + def llh_seed(self, value): + self._llh_seed = value + #self.llh.set_rng_seed(value) # this won't work as seed gets value before we have llh + + + @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) + # precomputed sensitivity and trial can have a generic path + # but actually specified by the configured datasets + self._bg_dir = os.path.join(self._bg_dir, "+".join(self.datasets)) + self._sens_dir = os.path.join(self._sens_dir, "+".join(self.datasets)) + + 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 + + + # 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 + """ + min_dec, max_dec = self.dec_skymap_range() + low5 = [] + high5 = [] + fig, ax = plt.subplots(figsize = (8,5)) + fig.set_facecolor('white') + # iterate over dataset + + for enum in self.llh._samples: + llh = self.llh._samples[enum] + dataset = self.datasets[enum].replace('_', ' ') + style = plotting_utils.skymap_style[enum] + + energy_range = defaultdict(list) + + # then iterate over the min- and max- declination of the skymap + for (dec_label, dec, color) in [ + ("min. dec", min_dec, sns.xkcd_rgb['windows blue']), + ("max. dec", max_dec, sns.xkcd_rgb['dark navy blue']), + + ]: + label = "" + if len(self.analyses)==1: + label = dec_label + else: + label = " ".join((dec_label, dataset)) + + # obtain the quantities for this dataset + energy_band = plotting_utils.get_energy_band(llh.mc, self.index, dec, + half_width=5., coverage=0.9) + # then make the plot, differentiating by linestyle + plotting_utils.plot_energy_band(**energy_band, color=color, linestyle=style["linestyle"], + label_prefix=label) + for key in ["low", "high", "median"]: + energy_range[key].append(energy_band[key]) + + low5.append(np.min(energy_range["low"])) + high5.append(np.max(energy_range["high"])) + + + 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.energy_range = tuple(zip(self.low5, self.high5)) + self.save_items['energy_range'] = self.energy_range + +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. + ''' + # NOTE ideally samples should not overlap and this is prevented by configuring the dataset + # if not, duplicates will appear appear here twice + # (which is a warning that the combined LLH is iffy) + + 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): + 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] + color = sns.xkcd_rgb['windows blue'] + + # obtain the quantities for this dataset + energy_band = plotting_utils.get_energy_band(llh.mc, self.index, self.dec, + half_width=5., coverage=0.9) + # then make the plot, differentiating by linestyle + plotting_utils.plot_energy_band(**energy_band, color=color, linestyle=style["linestyle"], + label_prefix=dataset) + + low5.append(energy_band["low"]) + high5.append(energy_band["high"]) + + + 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 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/MultiGWFollowup.py b/fast_response/MultiGWFollowup.py new file mode 100644 index 00000000..a8e24443 --- /dev/null +++ b/fast_response/MultiGWFollowup.py @@ -0,0 +1,154 @@ +from os.path import join +from .MultiFastResponseAnalysis import MultiPriorFollowup +from .GWFollowup import GWFollowup as GFUFollowup +from .FastResponseAnalysis import PriorFollowup + +import numpy as np +import numpy.lib.recfunctions as rf +import healpy as hp + +# These methods are very similar to (Multi)PriorFollowup +# but GW analyses make some particular choices we want to reproduce here +class MultiGWFollowup(MultiPriorFollowup): + + def find_coincident_events(self): + r""" + Find "coincident events" for a skymap + based analysis. These are ALL ontime events, + with a bool to indicate if they are in the 90% contour + + This works with a MultiPointSourceLLH. + """ + if self.ts_scan is None: + raise ValueError("Need to unblind TS before finding events") + exp_theta = 0.5*np.pi - self.llh_exp['dec'] + exp_phi = self.llh_exp['ra'] + exp_pix = hp.ang2pix(self.nside, exp_theta, exp_phi) + + t_mask=(self.llh_exp['time'] <= self.stop) & (self.llh_exp['time'] >= self.start) + events = self.llh_exp[t_mask] + ontime_pix = hp.ang2pix(self.nside, 0.5*np.pi - events['dec'], events['ra']) + overlap = np.isin(ontime_pix, self.ipix_90) + + events = rf.append_fields( + events, names=['in_contour', 'ts', 'ns', 'gamma', 'B'], + data=np.empty((5, events['ra'].size)), + usemask=False) + + for i in range(events['ra'].size): + events['in_contour'][i]=overlap[i] + enum_i = events['enum'][i] + events['B'][i] = self.llh._samples[enum_i].llh_model.background(events[i]) + + val_pix = self.scanned_pixels + for i in range(events['ra'].size): + idx, = np.where(val_pix == exp_pix[t_mask][i]) + # scan was restricted to containment fraction of spatial prior + # if a given each event centroid overlaps, look up the scan values + if idx.size > 0: + events['ts'][i] = self.ts_scan['TS_spatial_prior_0'][idx[0]] + events['ns'][i] = self.ts_scan['nsignal'][idx[0]] + events['gamma'][i] = self.ts_scan['gamma'][idx[0]] + + self.events_rec_array = events + self.coincident_events = [dict(zip(events.dtype.names, x)) for x in events] + self.save_items['coincident_events'] = self.coincident_events + + def per_event_pvalue(self): + """ + Calculate per-event p-values. There are a few cases here: + + - overall p < 0.1: + Redoes the all-sky scan, using per_event_scan, with only that single event. + This is the same as asking the question: + If that single event is the only one on the sky, with this given skymap, + what TS/p-value would we get for that event? + + - 1.0 > overall p > 0.1: + Calculates the p-value at the reconstructed event direction. + Takes the TS at that location, and calculates the p-value at that location. + Does not re-run the scan, to save time in realtime + + - p=1.0: + Does not get p-values for the events (all are set to None) + + """ + self.events_rec_array = rf.append_fields( + self.events_rec_array, + names=['pvalue'], + data=np.empty((1, self.events_rec_array['ra'].size)), + usemask=False + ) + + if self.tsd is None: # can not determine p-values + for i in range(self.events_rec_array.size): + self.events_rec_array['pvalue'][i] = None + elif self.events_rec_array.size > 0: # always scan, most general + for i in range(self.events_rec_array.size): + ts, p = self.per_event_scan(self.events_rec_array[i]) + self.events_rec_array['pvalue'][i] = p + else: + pass # nothing to do + + self.coincident_events = [dict(zip(self.events_rec_array.dtype.names, x)) for x in self.events_rec_array] + self.save_items['coincident_events'] = self.coincident_events + + +class OnlyDNNOnlineFollowup(PriorFollowup): + _base_dir = "/data/user/chraab/fast_response/multisample/extended_archival_floating" + _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}', + 'seed_*npz', + ]) + _dataset = "DNNCascadesOnline_v001p01" + _season_names = [f"IC86, 20{y:02d}" for y in range(24, 25+1)] + _floor = np.radians(1.5) # can change this! + _jitter = 3. # common default for DNN analyses + _background_days = 60. + _nb_days = 60. + +class OnlyIceManFollowup(PriorFollowup): + _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 DNNOnlineFollowup(MultiGWFollowup): + _base_dir = "/data/user/chraab/fast_response/multisample/extended_archival_floating" + _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}', + 'seed_*npz', + ]) + _followups = [OnlyDNNOnlineFollowup] + _fix_index = False + _float_index = not _fix_index + _index = 2.0 + _nside = 128 + +class IceManFollowup(MultiGWFollowup): + _base_dir = "/data/user/chraab/fast_response/multisample/extended_archival_floating" + _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}', + 'seed_*npz', + ]) + _followups = [OnlyIceManFollowup] + _fix_index = False + _float_index = not _fix_index + _index = 2.0 + _nside = 128 diff --git a/fast_response/plotting_utils.py b/fast_response/plotting_utils.py index b024025c..2b2ae537 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 @@ -312,4 +376,54 @@ def make_public_zoom_skymap(skymap, events, ra, dec, with_contour=True, name='te labels=[f'{min(skymap):.1e}',f'{max(skymap):.1e}']) plt.savefig(f'./{name}_skymap_zoom_public.png', bbox_inches='tight', dpi=300) - plt.close() \ No newline at end of file + plt.close() + +def get_energy_band(mc, index, dec, half_width=5., coverage=0.9): + """Get the weighted true energy distribution and 90% band + for a given declination, coupled out from make_dNdE to be used + by multiple plotting methods""" + + dec_mask_1 = mc['dec'] > dec - np.deg2rad(half_width) + dec_mask_2 = mc['dec'] < dec + np.deg2rad(half_width) + dec_mask = dec_mask_1 & dec_mask_2 + + delta_gamma = -1. * index + 1. # TODO I don't understand why we add a power of E then divide it out again + + a = np.histogram(mc['trueE'][dec_mask], bins = np.logspace(1., 8., 50), + weights = mc['ow'][dec_mask] * np.power(mc['trueE'][dec_mask], delta_gamma) / mc['trueE'][dec_mask], + ) + + cdf = np.cumsum(a[0]) / np.sum(a[0]) + low = np.interp((1 - coverage)/2, cdf, a[1][:-1]) + median = np.interp(0.5, cdf, a[1][:-1]) + high = np.interp((1 + coverage)/2, cdf, a[1][:-1]) + return {'histogram':a, 'low': low, 'median': median, 'high': high} + + +def plot_energy_band(histogram=None, low=None, median=None, high=None, quantile=0.9, + color=sns.xkcd_rgb['windows blue'], + linestyle="solid", + label_prefix=""): + if histogram: + plt.stairs(histogram[0], histogram[1], fill=False, + linewidth = 2., color = color, label = label_prefix) + + plt.yscale('log') + plt.xscale('log') + plt.grid(which = 'major', alpha = 0.25) + plt.xlabel('Energy (GeV)', fontsize = 24) + + if low is not None and high is not None: + plt.axvspan(low, high, color=color, alpha = 0.25, + label=" ".join((label_prefix, f"Central {quantile:.0%}")), + linestyle = linestyle, + linewidth = 2., + ) + if median: + plt.axvline(median, c=color, alpha=0.75, + linestyle = linestyle, + linewidth = 2., + label=" ".join((label_prefix, "Median")), + ) + plt.xlim(1e1, 1e8) + \ No newline at end of file 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 00000000..321e6d88 --- /dev/null +++ b/fast_response/precomputed_background/glob_precomputed_trials_multi.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python +""" +This module allows combining precomputed background sky scans +from multiple files saves as SciPy sparse matrices. +This is based on glob_precomputed_trials.py but generalized +in the assumptions on the file name pattern, so it can be used +with different conventions, the only requirement that they contain +some seed index in the form `seed_*`. +Running as a script, it will save the combined maps again as a sparse matrix. +""" + +import logging +from glob import glob +import healpy as hp +from scipy import sparse +import time +import argparse +import os +import sys +try: + from tqdm import tqdm +except ImportError: + tqdm = lambda x:x + +logger = logging.getLogger(__name__) + +parser = argparse.ArgumentParser(description='Glob precomputed 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 load_maps(fn): + return sparse.load_npz(fn) + +def concatenate_maps(files, nside, max_ntrials=None) -> sparse.csr_matrix: + files = [fn for fn in files if f"nside_{nside}" in fn] + logger.info('Found {} files to load'.format(len(files))) + if len(files)==0: + return None + logger.info('Nside: {}'.format(nside)) + npix = hp.nside2npix(nside) + logger.info('Starting to load at {}'.format(time.ctime())) + maps = sparse.csr_matrix((0, npix), dtype=float) + for fn in tqdm(files): + scan = load_maps(fn) + maps = sparse.vstack((maps, scan)) + if max_ntrials is not None and maps.shape[0] > max_ntrials: + return maps + return maps + +def save_maps(maps, out): + logger.info(f'Creating {os.path.basename(out)}') + # Change format of sparse array + logger.info("Starting to change from COO to CSR at {}".format(time.ctime())) + # NOTE: I am not sure we need to keep this; maps are already CSR + scans = maps.tocsr() + logger.info("Finished at {}".format(time.ctime())) + # Save the sparse array + sparse.save_npz(out, scans) + + +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 = concatenate_maps(filegroup, args.nside) + save_maps(maps, outfile) + del maps + logger.info('done') + return 0 + +if __name__ == "__main__": + args = parser.parse_args() + logger.setLevel(logging.INFO) + sys.exit(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 00000000..12cb2c9b --- /dev/null +++ b/fast_response/precomputed_background/precompute_ts_multi.py @@ -0,0 +1,123 @@ +#!/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, DNNIceManFollowup + +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=128, 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) +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 + + + +# 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 + +outdir = os.path.join(args.outdir, 'precomputed_trials/', "+".join(f.datasets)) +if not os.path.exists(outdir): + os.makedirs(outdir, exist_ok=True) + +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)) diff --git a/fast_response/reports/ReportGenerator.py b/fast_response/reports/ReportGenerator.py index 9b34397d..ad59a2c2 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/convert_moc_to_healpix.py b/fast_response/scripts/convert_moc_to_healpix.py index 74537d6a..678939f2 100755 --- a/fast_response/scripts/convert_moc_to_healpix.py +++ b/fast_response/scripts/convert_moc_to_healpix.py @@ -12,7 +12,7 @@ parser = argparse.ArgumentParser(description='GW Followup') parser.add_argument('--skymap', type=str, default=None, help='path to skymap (should be the *.multiorder.fits downloaded from GraceDB)') -parser.add_argument('--nside', type=str, default=512, +parser.add_argument('--nside', type=int, default=512, help='nside to be used with the skymap (default=512)') args = parser.parse_args() @@ -28,4 +28,5 @@ new_map = m.rasterize(args.nside, 'NESTED') # save the new map -new_map.write_map(args.skymap.replace('multiorder','converted'), overwrite=True) \ No newline at end of file +outfile = args.skymap[:args.skymap.find('.fits')] + '.converted.fits' +new_map.write_map(outfile, overwrite=True) diff --git a/setup.py b/setup.py index 9b81fbc6..ab4f37aa 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', ] )