lalsuite 7.26.2.dev20251210__cp312-cp312-macosx_12_0_arm64.whl

lalinspiral/inspinjfind.py

@@ -0,0 +1,485 @@
+# Copyright (C) 2006--2011,2013,2014,2016,2017  Kipp Cannon
+#
+# 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
+#
+# =============================================================================
+#
+
+
+"""
+Inspiral injection identification library.  Contains code providing the
+capacity to search a list of sngl_inspiral candidates for events
+matching entries in a sim_inspiral list of software injections, recording the
+matches as inspiral <--> injection coincidences using the standard coincidence
+infrastructure.  Also, any pre-recorded inspiral <--> inspiral coincidences are
+checked for cases where all the inspiral events in a coincidence match an
+injection, and these are recorded as coinc <--> injection coincidences,
+again using the standard coincidence infrastructure.
+"""
+
+
+import bisect
+import functools
+import sys
+import tqdm
+
+
+from lal import iterutils
+from igwn_ligolw import ligolw
+from igwn_ligolw import lsctables
+from igwn_ligolw.utils import coincs as ligolw_coincs
+from igwn_ligolw.utils import time_slide as ligolw_time_slide
+from . import thinca
+
+
+__author__ = "Kipp Cannon <kipp.cannon@ligo.org>"
+from .git_version import date as __date__
+from .git_version import version as __version__
+
+
+#
+# =============================================================================
+#
+#                               Content Handler
+#
+# =============================================================================
+#
+
+
+@functools.total_ordering
+class SnglInspiral(lsctables.SnglInspiral):
+	"""
+	Version of lsctables.SnglInspiral who's .__cmp__() method compares
+	this object's .end value directly to the value of other.  Allows a
+	list of instances of this class sorted by .end to be bisection
+	searched for a LIGOTimeGPS end time.
+	"""
+	__slots__ = ()
+
+	def __lt__(self, other):
+		return self.end < other
+
+	def __eq__(self, other):
+		return self.end == other
+
+
+#
+# =============================================================================
+#
+#                              Document Interface
+#
+# =============================================================================
+#
+
+
+InspiralSICoincDef = lsctables.CoincDef(search = u"inspiral", search_coinc_type = 1, description = u"sim_inspiral<-->sngl_inspiral coincidences")
+InspiralSCNearCoincDef = lsctables.CoincDef(search = u"inspiral", search_coinc_type = 2, description = u"sim_inspiral<-->coinc_event coincidences (nearby)")
+InspiralSCExactCoincDef = lsctables.CoincDef(search = u"inspiral", search_coinc_type = 3, description = u"sim_inspiral<-->coinc_event coincidences (exact)")
+InspiralSTCoincDef = lsctables.CoincDef(search = u"inspiral", search_coinc_type = 4, description = u"sim_inspiral<-->sngl_inspiral template coincidences")
+
+
+class DocContents(object):
+	"""
+	A wrapper interface to the XML document.
+	"""
+	def __init__(self, xmldoc, bbdef, sbdef, scedef, scndef, process, end_time_bisect_window):
+		#
+		# store the process row
+		#
+
+		self.process = process
+
+		#
+		# locate the sngl_inspiral and sim_inspiral tables
+		#
+
+		self.snglinspiraltable = lsctables.SnglInspiralTable.get_table(xmldoc)
+		self.siminspiraltable = lsctables.SimInspiralTable.get_table(xmldoc)
+
+		#
+		# get the offset vectors from the document
+		#
+
+		self.offsetvectors = lsctables.TimeSlideTable.get_table(xmldoc).as_dict()
+
+		#
+		# construct the zero-lag time slide needed to cover the
+		# instruments listed in all the triggers, then determine
+		# its ID (or create it if needed)
+		#
+		# FIXME:  in the future, the sim_inspiral table should
+		# indicate time slide at which the injection was done
+		#
+
+		self.tisi_id = ligolw_time_slide.get_time_slide_id(xmldoc, {}.fromkeys(self.snglinspiraltable.getColumnByName("ifo"), 0.0), create_new = process, superset_ok = True, nonunique_ok = True)
+
+		#
+		# get coinc_definer row for sim_inspiral <--> sngl_inspiral
+		# coincs; this creates a coinc_definer table if the
+		# document doesn't have one
+		#
+
+		self.sb_coinc_def_id = ligolw_coincs.get_coinc_def_id(xmldoc, sbdef.search, sbdef.search_coinc_type, create_new = True, description = sbdef.description)
+
+		#
+		# get coinc_def_id's for sngl_inspiral <--> sngl_inspiral, and
+		# both kinds of sim_inspiral <--> coinc_event coincs.  set all
+		# to None if this document does not contain any sngl_inspiral
+		# <--> sngl_inspiral coincs.
+		#
+
+		try:
+			ii_coinc_def_id = ligolw_coincs.get_coinc_def_id(xmldoc, bbdef.search, bbdef.search_coinc_type, create_new = False)
+		except KeyError:
+			ii_coinc_def_id = None
+			self.sce_coinc_def_id = None
+			self.scn_coinc_def_id = None
+		else:
+			self.sce_coinc_def_id = ligolw_coincs.get_coinc_def_id(xmldoc, scedef.search, scedef.search_coinc_type, create_new = True, description = scedef.description)
+			self.scn_coinc_def_id = ligolw_coincs.get_coinc_def_id(xmldoc, scndef.search, scndef.search_coinc_type, create_new = True, description = scndef.description)
+
+		#
+		# get coinc table, create one if needed
+		#
+
+		try:
+			self.coinctable = lsctables.CoincTable.get_table(xmldoc)
+		except ValueError:
+			self.coinctable = lsctables.CoincTable.new()
+			xmldoc.childNodes[0].appendChild(self.coinctable)
+		self.coinctable.sync_next_id()
+
+		#
+		# get coinc_map table, create one if needed
+		#
+
+		try:
+			self.coincmaptable = lsctables.CoincMapTable.get_table(xmldoc)
+		except ValueError:
+			self.coincmaptable = lsctables.CoincMapTable.new()
+			xmldoc.childNodes[0].appendChild(self.coincmaptable)
+
+		#
+		# index the document
+		#
+		# FIXME:  inspiral<-->inspiral coincs should be organized by time
+		# slide ID, but since injections are only done at zero lag
+		# for now this is ignored.
+		#
+
+		# index the sngl_inspiral table
+		index = dict((row.event_id, row) for row in self.snglinspiraltable)
+		# find IDs of inspiral<-->inspiral coincs
+		self.sngls = dict((row.coinc_event_id, []) for row in self.coinctable if row.coinc_def_id == ii_coinc_def_id)
+		# construct event list for each inspiral<-->inspiral coinc
+		for row in self.coincmaptable:
+			try:
+				self.sngls[row.coinc_event_id].append(index[row.event_id])
+			except KeyError:
+				pass
+		del index
+		# construct a sngl-->coincs look-up table
+		self.coincs = dict((event.event_id, set()) for events in self.sngls.values() for event in events)
+		for row in self.coincmaptable:
+			if row.event_id in self.coincs and row.coinc_event_id in self.sngls:
+				self.coincs[row.event_id].add(row.coinc_event_id)
+		# create a coinc_event_id to offset vector look-up table
+		self.coincoffsets = dict((row.coinc_event_id, self.offsetvectors[row.time_slide_id]) for row in self.coinctable if row.coinc_def_id == ii_coinc_def_id)
+
+		#
+		# sort sngl_inspiral table by end time
+		#
+
+		self.snglinspiraltable.sort(key = lambda row: row.end)
+
+		#
+		# set the window for inspirals_near_endtime().  this window
+		# is the amount of time such that if an injection's end
+		# time and a inspiral event's end time differ by more than
+		# this it is *impossible* for them to match one another.
+		#
+
+		self.end_time_bisect_window = lsctables.LIGOTimeGPS(end_time_bisect_window)
+
+
+	def inspirals_near_endtime(self, t):
+		"""
+		Return a list of the inspiral events whose end times are
+		within self.end_time_bisect_window of t.
+		"""
+		return self.snglinspiraltable[bisect.bisect_left(self.snglinspiraltable, t - self.end_time_bisect_window):bisect.bisect_right(self.snglinspiraltable, t + self.end_time_bisect_window)]
+
+	def coincs_near_endtime(self, t):
+		"""
+		Return a list of the (coinc_event_id, event list) tuples in
+		which at least one inspiral event's end time is within
+		self.end_time_bisect_window of t.
+		"""
+		# FIXME:  this test does not consider the time slide
+		# offsets that should be applied to the coinc, but for now
+		# injections are done at zero lag so this isn't a problem
+		# yet
+		coinc_event_ids = set()
+		for event in self.inspirals_near_endtime(t):
+			try:
+				coinc_event_ids |= self.coincs[event.event_id]
+			except KeyError:
+				# this single isn't in any coincs
+				pass
+		return [(coinc_event_id, self.sngls[coinc_event_id]) for coinc_event_id in coinc_event_ids]
+
+	def sort_triggers_by_id(self):
+		"""
+		Sort the sngl_inspiral table's rows by ID (tidy-up document
+		for output).
+		"""
+		self.snglinspiraltable.sort(key = lambda row: row.event_id)
+
+	def new_coinc(self, coinc_def_id):
+		"""
+		Construct a new coinc_event row attached to the given
+		process, and belonging to the set of coincidences defined
+		by the given coinc_def_id.
+		"""
+		coinc = lsctables.Coinc()
+		coinc.process_id = self.process.process_id
+		coinc.coinc_def_id = coinc_def_id
+		coinc.coinc_event_id = self.coinctable.get_next_id()
+		coinc.time_slide_id = self.tisi_id
+		coinc.insts = None
+		coinc.nevents = 0
+		coinc.likelihood = None
+		self.coinctable.append(coinc)
+		return coinc
+
+
+#
+# =============================================================================
+#
+#                 Build sim_inspiral <--> sngl_inspiral Coincidences
+#
+# =============================================================================
+#
+
+
+def find_sngl_inspiral_matches(contents, sim, comparefunc):
+	"""
+	Scan the inspiral table for triggers matching sim.
+	"""
+	return [inspiral for inspiral in contents.inspirals_near_endtime(sim.time_geocent) if not comparefunc(sim, inspiral)]
+
+
+def add_sim_inspiral_coinc(contents, sim, inspirals):
+	"""
+	Create a coinc_event in the coinc table, and add arcs in the
+	coinc_event_map table linking the sim_inspiral row and the list of
+	sngl_inspiral rows to the new coinc_event row.
+	"""
+	coinc = contents.new_coinc(contents.sb_coinc_def_id)
+	coinc.insts = (event.ifo for event in inspirals)
+	coinc.nevents = len(inspirals)
+
+	contents.coincmaptable.append(lsctables.CoincMap(
+		coinc_event_id = coinc.coinc_event_id,
+		table_name = u"sim_inspiral",
+		event_id = sim.simulation_id
+	))
+
+	for event in inspirals:
+		contents.coincmaptable.append(lsctables.CoincMap(
+			coinc_event_id = coinc.coinc_event_id,
+			table_name = u"sngl_inspiral",
+			event_id = event.event_id
+		))
+
+
+#
+# =============================================================================
+#
+#                   Build sim_inspiral <--> coinc Coincidences
+#
+# =============================================================================
+#
+
+
+def find_exact_coinc_matches(coincs, sim, comparefunc):
+	"""
+	Return a set of the coinc_event_ids of the inspiral<-->inspiral
+	coincs in which all inspiral events match sim.
+	"""
+	# comparefunc is True --> inspiral does not match sim
+	# any(...) --> at least one inspiral does not match sim
+	# not any(...) --> all inspirals match sim
+	#
+	# FIXME:  this test does not consider the time slide offsets that
+	# should be applied to the coinc, but for now injections are done
+	# at zero lag so this isn't a problem yet
+	return set(coinc_event_id for coinc_event_id, inspirals in coincs if not any(comparefunc(sim, inspiral) for inspiral in inspirals))
+
+
+def find_near_coinc_matches(coincs, sim, comparefunc):
+	"""
+	Return a set of the coinc_event_ids of the inspiral<-->inspiral
+	coincs in which at least one inspiral event matches sim.
+	"""
+	# comparefunc is True --> inspiral does not match sim
+	# all(...) --> no inspirals match sim
+	# not all(...) --> at least one inspiral matches sim
+	#
+	# FIXME:  this test does not consider the time slide offsets that
+	# should be applied to the coinc, but for now injections are done
+	# at zero lag so this isn't a problem yet
+	return set(coinc_event_id for coinc_event_id, inspirals in coincs if not all(comparefunc(sim, inspiral) for inspiral in inspirals))
+
+
+def add_sim_coinc_coinc(contents, sim, coinc_event_ids, coinc_def_id):
+	"""
+	Create a coinc_event in the coinc table, and add arcs in the
+	coinc_event_map table linking the sim_inspiral row and the list of
+	coinc_event rows to the new coinc_event row.
+	"""
+	coinc = contents.new_coinc(coinc_def_id)
+	coinc.nevents = len(coinc_event_ids)
+
+	contents.coincmaptable.append(lsctables.CoincMap(
+		coinc_event_id = coinc.coinc_event_id,
+		table_name = u"sim_inspiral",
+		event_id = sim.simulation_id
+	))
+
+	for coinc_event_id in coinc_event_ids:
+		contents.coincmaptable.append(lsctables.CoincMap(
+			coinc_event_id = coinc.coinc_event_id,
+			table_name = u"coinc_event",
+			event_id = coinc_event_id
+		))
+
+
+#
+# =============================================================================
+#
+#                                 Library API
+#
+# =============================================================================
+#
+
+
+def ligolw_inspinjfind(xmldoc, process, search, snglcomparefunc, nearcoinccomparefunc, end_time_bisect_window = 1.0, verbose = False):
+	#
+	# Analyze the document's contents.
+	#
+
+	if verbose:
+		print("indexing ...", file=sys.stderr)
+
+	bbdef = {"inspiral": thinca.InspiralCoincDef}[search]
+	sbdef = {"inspiral": InspiralSICoincDef}[search]
+	scedef = {"inspiral": InspiralSCExactCoincDef}[search]
+	scndef = {"inspiral": InspiralSCNearCoincDef}[search]
+
+	contents = DocContents(xmldoc = xmldoc, bbdef = bbdef, sbdef = sbdef, scedef = scedef, scndef = scndef, process = process, end_time_bisect_window = end_time_bisect_window)
+
+	#
+	# Find sim_inspiral <--> sngl_inspiral coincidences.
+	#
+
+	for sim in tqdm.tqdm(contents.siminspiraltable, desc = sbdef.description, disable = not verbose):
+		inspirals = find_sngl_inspiral_matches(contents, sim, snglcomparefunc)
+		if inspirals:
+			add_sim_inspiral_coinc(contents, sim, inspirals)
+
+	#
+	# Find sim_inspiral <--> coinc_event coincidences.
+	#
+
+	if contents.scn_coinc_def_id:
+		for sim in tqdm.tqdm(contents.siminspiraltable, desc = scndef.description, disable = not verbose):
+			coincs = contents.coincs_near_endtime(sim.time_geocent)
+			exact_coinc_event_ids = find_exact_coinc_matches(coincs, sim, snglcomparefunc)
+			near_coinc_event_ids = find_near_coinc_matches(coincs, sim, nearcoinccomparefunc)
+			assert exact_coinc_event_ids.issubset(near_coinc_event_ids)
+			if exact_coinc_event_ids:
+				add_sim_coinc_coinc(contents, sim, exact_coinc_event_ids, contents.sce_coinc_def_id)
+			if near_coinc_event_ids:
+				add_sim_coinc_coinc(contents, sim, near_coinc_event_ids, contents.scn_coinc_def_id)
+
+	#
+	# Restore the original event order.
+	#
+
+	if verbose:
+		print("finishing ...", file=sys.stderr)
+	contents.sort_triggers_by_id()
+
+	#
+	# Done.
+	#
+
+	return xmldoc
+
+
+#
+# =============================================================================
+#
+#                                    Revert
+#
+# =============================================================================
+#
+
+
+def revert(xmldoc, program, verbose = False):
+	#
+	# remove entries from process metadata tables
+	#
+
+	if verbose:
+		print("removing process metadata ...", file=sys.stderr)
+	process_table = lsctables.ProcessTable.get_table(xmldoc)
+	# IDs of things to delete
+	process_ids = process_table.get_ids_by_program(program)
+	iterutils.inplace_filter((lambda row: row.process_id not in process_ids), process_table)
+	iterutils.inplace_filter((lambda row: row.process_id not in process_ids), lsctables.ProcessParamsTable.get_table(xmldoc))
+
+	#
+	# remove coinc_event and coinc_event_map entries
+	#
+
+	if verbose:
+		print("removing coincs ...", file=sys.stderr)
+	coinc_event_table = lsctables.CoincTable.get_table(xmldoc)
+	# IDs of things to delete
+	coinc_ids = frozenset(row.coinc_event_id for row in coinc_event_table if row.process_id in process_ids)
+	iterutils.inplace_filter((lambda row: row.coinc_event_id not in coinc_ids), coinc_event_table)
+	iterutils.inplace_filter((lambda row: row.coinc_event_id not in coinc_ids), lsctables.CoincMapTable.get_table(xmldoc))
+	# IDs of things to keep
+	time_slide_ids = frozenset(row.time_slide_id for row in coinc_event_table)
+	coinc_def_ids = frozenset(row.coinc_def_id for row in coinc_event_table)
+
+	#
+	# remove time_slide and coinc_definer entries
+	#
+
+	if verbose:
+		print("removing coinc metadata ...", file=sys.stderr)
+	# coinc types to delete
+	coinc_defs = frozenset((row.search, row.search_coinc_type) for row in (InspiralSICoincDef, InspiralSCNearCoincDef, InspiralSCExactCoincDef))
+	iterutils.inplace_filter((lambda row: row.process_id not in process_ids or row.time_slide_id in time_slide_ids), lsctables.TimeSlideTable.get_table(xmldoc))
+	iterutils.inplace_filter((lambda row: (row.search, row.search_coinc_type) not in coinc_defs or row.coinc_def_id in coinc_def_ids), lsctables.CoincDefTable.get_table(xmldoc))