# Copyright (C) 2021 ChiWai Chan
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
import sys
import numpy
import matplotlib
matplotlib.use("Agg")
from pathlib import Path
from collections import defaultdict
from matplotlib import pyplot
from lal.utils import CacheEntry
from gstlal import far
from gstlal.plots import util as plotutil
from gstlal.psd import read_psd, HorizonDistance as calc_horizon_distance
[docs]class UnknownExtensionError(Exception):
pass
[docs]class HorizonDistance:
def __init__(self, horizon_history_dict, verbose=False):
"""Construct a HorizonDistance instance from a list of ranking
statistic files.
Args:
horizon_history_dict:
dict, a dictionary containing keyed by ifo,
whose value is a dictionary of horizon history
(see notes).
verbose:
bool, default is False, toggle error message.
Notes:
You should use the two class methods
(`from_rankingstats` and `from_psds`) instead of
directly calling the init method. In case you need to
construct an instance from the init method, you prepare
a horizon history dict containing horizon history keyed
by ifo. A horizon history is a dictionary of GPS time
and horizon distance pairs.
Returns:
An instance of HorizionDistance.
"""
self._min = None
self._max = None
self.verbose = verbose
self.horizon_history_dict = horizon_history_dict
[docs] @classmethod
def from_rankingstats(cls, sources, verbose=False):
"""Construct a HorizonDistance instance from a list of ranking
statistic files.
Args:
sources:
list of str, the data sources to be loaded.
verbose:
bool, default is False, toggle error message.
Returns:
An instance of HorizionDistance.
Note:
The ranking statistic files are expectd to be in the same
template bank bin.
"""
loader = cls.datasource_loader(sources)
urls = loader(sources)
rankingstat = far.marginalize_pdf_urls(urls, "RankingStat", verbose = verbose)
horizon_history_dict = rankingstat.numerator.horizon_history
return cls(horizon_history_dict, verbose=verbose)
[docs] @classmethod
def from_psds(cls, sources, verbose=False):
"""Construct a HorizonDistance instance from a list of psds.
Args:
sources:
list of str, the data sources to be loaded.
verbose:
bool, default is False, toggle error message.
Returns:
An instance of HorizionDistance.
"""
loader = cls.datasource_loader(sources)
horizon_history_dict = defaultdict(dict)
for f in loader(sources):
psds = read_psd(f, verbose = verbose)
for ifo, psd in psds.items():
if psd is not None:
horizon_history_dict[ifo].update({int(psd.epoch): calc_horizon_distance(10., 2048., psd.deltaF, 1.4, 1.4)(psd, 8.)[0]})
return cls(horizon_history_dict, verbose=verbose)
[docs] @staticmethod
def datasource_loader(sources):
"""Determine how to load the data sources based on their extensions.
Args:
sources:
list of str, the data sources to be loaded.
Raises:
UnknownExtensionError:
If the sources do not ends with .gz or .cache.
Returns:
None
"""
assert type(sources) == list, "Please provides <class 'list'> as input data sources."
assert len(sources) >= 1, "Please provides at least one data source."
# default loader
extension = Path(sources[0]).suffix
if extension == ".gz":
return _load_files
if extension == ".cache":
return _load_caches
raise UnknownExtensionError("Cannot determine the extension of the input data sources, please prepare files with the following extension: ('.gz' or '.cache').")
@property
def min(self):
"""Return the global minimum of the horizon distance."""
if self._min:
return self._min
gminh = 1e32
for ifo, history in self.horizon_history_dict.items():
gminh = min(gminh, min(history.values()))
self._min = gminh
return gminh
@property
def max(self):
"""Return the global minimum of the horizon distance."""
if self._max:
return self._max
gmaxh = -1e32
for ifo, history in self.horizon_history_dict.items():
gmaxh = max(gmaxh, max(history.values()))
self._max = gmaxh
return gmaxh
[docs] def savefig(self, output, figsize=(12,4), limits=None, title=None):
"""Save the horizon distance plot to disk.
Args:
output:
str, the output file name.
figsize:
tuple of float, default (12,4), set the output
figure size in inch.
limits:
list of float, default to auto detection, set
the limits on horizon distance.
title:
str, the title of the plots.
Returns:
None
"""
fig, ax = pyplot.subplots(1, 2, figsize = figsize)
if title is not None:
fig.suptitle(title)
pyplot.tight_layout(pad = 4, w_pad = 4, h_pad = 4)
mint = int(min([min(horizon_history.keys()) for _, horizon_history in self.horizon_history_dict.items()]))
for ifo, horizon_history in self.horizon_history_dict.items():
GPSTime = numpy.array(list(horizon_history.keys()))
horizon_dist = list(horizon_history.values())
minh, maxh = (float("inf"), 0)
maxh = max(maxh, max(horizon_dist))
minh = min(minh, min(horizon_dist))
SinceGPSTime = (GPSTime - mint)/1000.
binvec = numpy.linspace(minh, maxh, 25)
ax[0].semilogy(SinceGPSTime, horizon_dist, "x", color = plotutil.colour_from_instruments([ifo]), label = ifo)
ax[1].hist(horizon_dist, binvec, alpha = 0.5, color = plotutil.colour_from_instruments([ifo]), label = ifo)
if limits is not None:
ax[0].set_ylim(limits)
ax[1].set_xlim(limits)
if self.verbose:
sys.stderr.write("plotting " + str(output) + "\n")
ax[0].set_xlabel("Time (ks) from GPS {:d}".format(mint))
ax[0].set_ylabel("Mpc")
ax[0].legend(loc = "best")
ax[0].grid()
ax[1].set_xlabel("Mpc")
ax[1].set_ylabel("Count")
ax[1].legend(loc = "best")
fig.savefig(output)
pyplot.close()
def _load_files(files):
return files
def _load_caches(caches):
urls = (CacheEntry(line).url for cache in caches for line in open(cache))
return sorted(urls, key = lambda x: CacheEntry.from_T050017(x).description)