gstlal_inspiral/plots: add horizon.py module for plotting horizon distance.

gstlal-inspiral/python/plots/Makefile.am

@@ -4,6 +4,7 @@ plotsdir = $(pkgpythondir)/plots
 plots_PYTHON = \
 	bank.py \
 	far.py \
+	horizon.py \
 	lloid.py \
 	segments.py \
 	sensitivity.py \

gstlal-inspiral/python/plots/horizon.py

+# 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 gstlal import far
+from gstlal.plots import util as plotutil
+from gstlal.psd import read_psd, HorizonDistance
+
+
+class UnknownExtensionError(Exception):
+	pass
+
+
+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.verbose = verbose
+		self.horizon_history_dict = horizon_history_dict
+
+	@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.
+		"""
+		loader = cls.datasource_loader(sources)
+		urls = loader(sources)
+		rankingstat = far.marginalize_pdf_urls(list(urls), "RankingStat", verbose = verbose)
+		horizon_history_dict = rankingstat.numerator.horizon_history
+		return cls(horizon_history_dict)
+
+	@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): HorizonDistance(10., 2048., psd.deltaF, 1.4, 1.4)(psd, 8.)[0]})
+		return cls(horizon_history_dict)
+
+	@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').")
+
+	def savefig(self, output, figsize=(12,4)):
+		"""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.
+
+		Returns:
+			None
+		"""
+		fig, ax = pyplot.subplots(1, 2, figsize = figsize)
+		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(horizon_history.keys())
+			horizon_dist = 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 self.verbose:
+			sys.stderr.write("plotting " + output)
+
+		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_cachces(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)
+