inspiral_lr.py: Add population model to ranking statistic. Population model...

inspiral_lr.py: Add population model to ranking statistic. Population model information in inspiral_intrinsics.py.

gstlal-inspiral/python/stats/Makefile.am

@@ -4,6 +4,7 @@ statsdir = $(pkgpythondir)/stats
 stats_PYTHON = \
 	horizonhistory.py \
 	inspiral_extrinsics.py \
+	inspiral_intrinsics.py \
 	inspiral_lr.py \
 	trigger_rate.py
 

gstlal-inspiral/python/stats/inspiral_intrinsics.py

+# Copyright (C) 2017,2018  Heather Fong
+#
+# 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.
+
+
+#
+# =============================================================================
+#
+#                                   Preamble
+#
+# =============================================================================
+#
+
+
+import h5py
+import math
+import numpy
+import os
+from scipy.interpolate import PPoly
+
+
+from gstlal import stats as gstlalstats
+
+
+# FIXME:  caution, this information might get organized differently later.
+# for now we just need to figure out where the gstlal-inspiral directory in
+# share/ is.  don't write anything that assumes that this module will
+# continue to define any of these symbols
+from gstlal import paths as gstlal_config_paths
+
+
+__all__ = [
+	"UniformInTemplatePopulationModel",
+	"SourcePopulationModel"
+]
+
+
+#
+# =============================================================================
+#
+#                               Population Model
+#
+# =============================================================================
+#
+
+
+class UniformInTemplatePopulationModel(object):
+	def __init__(self, template_ids):
+		"""
+		Assumes uniform in template population model, no
+                astrophysical prior.
+		"""
+		self.lnP = -math.log(len(template_ids))
+
+
+	@gstlalstats.assert_ln_probability
+	def lnP_template_signal(self, template_id, snr):
+		assert snr >= 0.
+		return self.lnP
+
+
+class SourcePopulationModel(object):
+	DEFAULT_FILENAME = os.path.join(gstlal_config_paths["pkgdatadir"], "lnP_template_signal_BNS_gaussian_lowspin_Ozel.hdf5")
+
+
+	def __init__(self, template_ids, filename = None):
+		"""
+		Sets the polynomial coefficients, given the template ID 
+                and SNR for a source population model, from which
+                lnP is then computed using PPoly.
+		"""
+		with h5py.File(filename if filename is not None else self.DEFAULT_FILENAME, 'r') as model:
+			coefficients = model['coefficients'].value
+			snr_bp = model['SNR'].value
+		# PPoly can construct an array of polynomials by just
+		# feeding it the coefficients array all in one go, but then
+		# it insists on evaluating all of them at once.  we don't
+		# want to suffer that cost, so we have to make an array of
+		# PPoly objects ourselves, and index into it to evaluate
+		# just one.  since we have to do this anyway, we use a
+		# dictionary to also solve the problem of mapping
+		# template_id to a specific polynomial
+		self.polys = dict((template_id, PPoly(coefficients[:,:,[template_id]], snr_bp)) for template_id in template_ids)
+		self.max_snr = snr_bp.max()
+
+
+	@gstlalstats.assert_ln_probability
+	def lnP_template_signal(self, template_id, snr):
+		assert snr >= 0.
+		try:
+			lnP_vs_snr = self.polys[template_id]
+		except KeyError:
+			raise KeyError("template ID %d is not in this model" % template_id)
+		# PPoly's .__call__() returns an array, so we need the
+		# final [0] to flatten it
+		return lnP_vs_snr(min(snr, self.max_snr))[0]

gstlal-inspiral/python/stats/inspiral_lr.py

@@ -44,6 +44,7 @@ from glue.ligolw import utils as ligolw_utils
 from glue import segments
 from gstlal.stats import horizonhistory
 from gstlal.stats import inspiral_extrinsics
+from gstlal.stats import inspiral_intrinsics
 from gstlal.stats import trigger_rate
 import lal
 from lal import rate
@@ -259,6 +260,17 @@ class LnSignalDensity(LnLRDensity):
 		# record of horizon distances for all instruments in the
 		# network
 		self.horizon_history = horizonhistory.HorizonHistories((instrument, horizonhistory.NearestLeafTree()) for instrument in self.instruments)
+
+		# source population model
+		if self.template_ids:
+			self.population_model = inspiral_intrinsics.UniformInTemplatePopulationModel(self.template_ids)
+			# FIXME:  switch to this when a model file becomes
+			# available
+			#self.population_model = inspiral_intrinsics.SourcePopulationModel(self.template_ids)
+		else:
+			# default lnP = 1/len(templates) = 0
+			self.population_model = inspiral_intrinsics.UniformInTemplatePopulationModel([0])
+
 		self.InspiralExtrinsics = inspiral_extrinsics.InspiralExtrinsics(self.min_instruments)
 
 	def __call__(self, segments, snrs, phase, dt, template_id, **kwargs):
@@ -310,6 +322,9 @@ class LnSignalDensity(LnLRDensity):
 		except ValueError:
 			return float("-inf")
 
+		# evaluate population model
+		lnP += self.population_model.lnP_template_signal(template_id, max(snrs.values()))
+
 		# evalute the (snr, \chi^2 | snr) PDFs (same for all
 		# instruments)
 		interp = self.interps["snr_chi"]
@@ -475,6 +490,19 @@ class LnSignalDensity(LnLRDensity):
 		xml = cls.get_xml_root(xml, name)
 		self = super(LnSignalDensity, cls).from_xml(xml, name)
 		self.horizon_history = horizonhistory.HorizonHistories.from_xml(xml, u"horizon_history")
+		# source population model
+		# FIXME:  this should probably be stored in the ranking
+		# statistic file somehow.  maybe the HDF5 filename could be
+		# stored.  whatever would allow the correct model to be
+		# re-initialized
+		if self.template_ids:
+			self.population_model = inspiral_intrinsics.UniformInTemplatePopulationModel(self.template_ids)
+			# FIXME:  switch to this when a model file becomes
+			# available
+			#self.population_model = inspiral_intrinsics.SourcePopulationModel(self.template_ids)
+		else:
+			# default lnP = 1/len(templates) = 0
+			self.population_model = inspiral_intrinsics.UniformInTemplatePopulationModel([0])
 		return self
 
 
@@ -519,6 +547,9 @@ class DatalessLnSignalDensity(LnSignalDensity):
 		except ValueError:
 			return float("-inf")
 
+		# evaluate population model
+		lnP += self.population_model.lnP_template_signal(template_id, max(snrs.values()))
+
 		# evalute the (snr, \chi^2 | snr) PDFs (same for all
 		# instruments)
 		interp = self.interps["snr_chi"]