gstlal_inspiral_set_svdbin_option : add a standalone function to write...

gstlal_inspiral_set_svdbin_option : add a standalone function to write bin-specific option to svd manifest file

gstlal-inspiral/bin/Makefile.am

@@ -63,6 +63,7 @@ dist_bin_SCRIPTS = \
 	gstlal_inspiral_svd_bank_workflow \
 	gstlal_inspiral_workflow \
 	gstlal_inspiral_diststatpdf_collect_zerolag \
+	gstlal_inspiral_set_svdbin_option \
 	gstlal_ll_inspiral_calculate_range \
 	gstlal_ll_inspiral_event_plotter \
 	gstlal_ll_inspiral_event_uploader \

gstlal-inspiral/bin/gstlal_inspiral_set_svdbin_option

+#!/usr/bin/env python3
+#
+# Copyright (C) 2023  Leo Tsukada (leo.tsukada@ligo.org)
+#
+# 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 argparse
+import os
+import sys
+from typing import Iterable
+from collections.abc import Mapping
+
+from lal import rate
+
+from gstlal.config.inspiral import Config
+from gstlal.dags.inspiral import DAG
+from gstlal.datafind import DataCache, DataType
+from gstlal.workflows import write_makefile
+
+
+def calc_gate_threshold(config, svd_bin, aggregate="max"):
+	"""
+	Given a configuration, svd bin and aggregate, this calculates
+	the h(t) gate threshold used for a given svd bin.
+	"""
+	if isinstance(config.filter.ht_gate_threshold, str):
+		bank_mchirp = config.svd.stats["bins"][svd_bin][f"{aggregate}_mchirp"]
+		min_mchirp, min_threshold, max_mchirp, max_threshold = [
+			float(y) for x in config.filter.ht_gate_threshold.split("-") for y in x.split(":")
+		]
+		gate_mchirp_ratio = (max_threshold - min_threshold) / (max_mchirp - min_mchirp)
+		threshold = round(gate_mchirp_ratio * (bank_mchirp - min_mchirp) + min_threshold, 3)
+	else: # uniform threshold
+		threshold = config.filter.ht_gate_threshold
+	config.svd.stats.bins[svd_bin]["ht_gate_threshold"] = threshold
+	return threshold
+
+def autocorrelation_length_map(ac_length_range):
+	"""
+	Given autocorrelation length ranges (e.g. 0:15:701)
+	or a single autocorrelation value, returns a function that
+	maps a given chirp mass to an autocorrelation length.
+	"""
+	if isinstance(ac_length_range, str):
+		ac_length_range = [ac_length_range]
+
+	# handle case with AC length ranges
+	if isinstance(ac_length_range, Iterable):
+		ac_lengths = []
+		min_mchirps = []
+		max_mchirps = []
+		for this_range in ac_length_range:
+			min_mchirp, max_mchirp, ac_length = this_range.split(":")
+			min_mchirps.append(float(min_mchirp))
+			max_mchirps.append(float(max_mchirp))
+			ac_lengths.append(int(ac_length))
+
+		# sanity check inputs
+		for bound1, bound2 in zip(min_mchirps[1:], max_mchirps[:-1]):
+			assert bound1 == bound2, "gaps not allowed in autocorrelation length ranges"
+
+		# convert to binning
+		bins = rate.IrregularBins([min_mchirps[0]] + max_mchirps)
+
+	# handle single value case
+	else:
+		ac_lengths = [ac_length_range]
+		bins = rate.IrregularBins([0., numpy.inf])
+
+	# create mapping
+	def mchirp_to_ac_length(mchirp):
+		idx = bins[mchirp]
+		return ac_lengths[idx]
+
+	return mchirp_to_ac_length
+
+def svd_bin_to_dtdphi_file(config, svd_bin, stats_bin, aggregate="mean"):
+	category_condition = {
+			"IMBH" : lambda stats_bin: stats_bin[f"{aggregate}_mtotal"] > 100 and stats_bin[f"{aggregate}_mratio"] < 10,
+			"others" : lambda stats_bin: True
+			}
+
+	if isinstance(config.prior.dtdphi, Mapping):
+		if "bank_name" in stats_bin:
+			sub_bank = stats_bin["bank_name"]
+			dtdphi_file = config.prior.dtdphi[sub_bank]
+		else:
+			assert all(category in category_condition for category in config.prior.dtdphi if category != "others"), "At least one of the categories set in config, [%s], is not defined in those in the source code, [%s]." % (",".join(config.prior.dtdphi.keys()), ",".join(category_condition.keys()))
+			dtdphi_files = {category: filename for category, filename in config.prior.dtdphi.items() if category_condition[category](stats_bin)}
+			if not len(dtdphi_files):
+				raise ValueError("SVD bin %s does not meet a condition of any category given in config.prior.dtdphi option.\
+						Add 'others' category explicitly and point that to the default dtdphi file to catch such bins." % (svd_bin,))
+			elif "others" in dtdphi_files:
+				# pick the default dtdphi file from 'others' category
+				dtdphi_file = dtdphi_files.pop("others")
+			if len(dtdphi_files) >= 2:
+				raise ValueError("SVD bin id %s falls onto the multiple categories (%s). It needs to be assigned to only one or none."\
+					% (svd_bin, ",".join([tpl[0] for tpl in dtdphi_files])))
+			elif len(dtdphi_files):
+				# pick the dtdphi file specified in the config
+				dtdphi_file = [*dtdphi_files.values()][0]
+			assert "dtdphi_file" in locals(), "dtdphi_file is not defined even after the validation"
+	else:
+		dtdphi_file = config.prior.dtdphi
+
+	return dtdphi_file
+
+#
+# =============================================================================
+#
+#                                 Command Line
+#
+# =============================================================================
+#
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-c", "--config", help="Sets the path to read configuration from.")
+
+
+#
+# =============================================================================
+#
+#                                     Main
+#
+# =============================================================================
+#
+
+
+#
+# Parse command line
+#
+
+
+args = parser.parse_args()
+
+#
+# load config
+#
+
+config = Config.load(args.config)
+config.setup()
+
+# set up autocorrelation mapping
+mchirp_to_ac_length = autocorrelation_length_map(config.svd.autocorrelation_length)
+
+for svd_bin, stats_bin in config.svd.stats.bins.items():
+	stats_bin["ac_length"] = mchirp_to_ac_length(stats_bin["mean_mchirp"])
+	stats_bin["ht_gate_threshold"] = calc_gate_threshold(config, svd_bin)
+	stats_bin["mass_model_file"] = config.prior.mass_model
+	if config.prior.idq_timeseries:
+		stats_bin["idq_file"] = config.prior.idq_timeseries
+	if config.prior.dtdphi:
+		stats_bin["dtdphi_file"] = svd_bin_to_dtdphi_file(config, svd_bin, stats_bin)
+	else:
+		stats_bin["dtdphi_file"] = os.path.join(gstlal_config_paths["pkgdatadir"], "inspiral_dtdphi_pdf.h5")
+
+config.write_svd_manifest(config.svd.manifest)