Online new extinction

How the marginalize_likelihoods_online program worked before is:

1. For each registry, get the rankingstat from the inspiral job and run calc_rank_pdfs of it
2. Add it to the variable data
3. retry the failed bins
4. 3. data now contains the newly-created noise PDFS from (almost) every bin
5. get the marginalized zerolag from the trigger counter job and add it to the data variable. This should be the zerolag from the start of the analysis till now
6. If a previously saved marg dist_stat_pdf exists, load it and add it to data
7. data now contains not just the newly created noise PDFS, but also the noise PDFs from the start of the analysis till now
8. Save data to disk as the new version of the marg dist_stat_pdf file

With these changes, the marginalize_likelihoods_online implements the new extinction model. This involves first-round-extinction of the bin-specific PDFs with the bin-specific zerolag, adding these singly-extincted PDFs together, and adding the clustered zerolag so that it's ready for second-round-extinction. Now, the marginalize_likelihoods_online job works as: 0. Now the marg job also saves the bin-specific dist_stat_pdf to disk since they're required for first-round extinction. The path is determined by DataCache.generate, and not decided by the user. The filename of the marg dist_stat_pdf is still provided by the user

1. For each registry, load the old dist_stat_pdf file from disk if it exists
2. Get the rankingstat from the inspiral job and run calc_rank_pdfs of it, and add it to the old dist_stat_pdf
3. calc_rank_pdfs is retried 3 times in my version. This is because I found it difficult to coordinate the retries with both 1 and 2. But I think it's still possible to do the retries at the end like before, so I'll make that change soon
4. Save the new + old dist_stat_pdf to disk
5. Get the bin-specific zerolag from the corresponding inspiral job and add it to the new + old dist_stat_pdf (needed for dirst-round extinction)
6. Perform first round extinction on this pdf, and add it to the data variable
7. Since in each iteration of the loop we are adding the old and new dist_stat_pdfs to data, it contains the noise PDF from the start of the analysis to now. No need to do step 6 from the previous version of the code
8. Remove the bin-specific zerolags and add the marginalized zerolag from the trigger counter job. This should be the zerolag from the start of the analysis till now
9. Save data to disk as the new version of the marg dist_stat_pdf file

gstlal-inspiral/python/far.py

@@ -59,6 +59,7 @@ import math
 import multiprocessing
 import numpy
 import random
+import scipy
 from scipy import interpolate
 import sys
 import time
@@ -519,6 +520,9 @@ def binned_log_likelihood_ratio_rates_from_samples(signal_lr_lnpdf, noise_lr_lnp
 
 class RankingStatPDF(object):
 	ligo_lw_name_suffix = u"gstlal_inspiral_rankingstatpdf"
+	extinction_fitting_limits = (1/2., 1/100.)
+	# extinction curve fitting is done for the range of LRs
+	# corresponding to the 1/2 and 1/100 point of the zl CCDF
 
 	@staticmethod
 	def density_estimate(lnpdf, name, kernel = rate.gaussian_window(4.)):
@@ -712,89 +716,73 @@ WHERE
 		return self
 
 
-	def new_with_extinction(self):
-		self = self.copy()
+	def ready_for_extinction(self):
+		# ensure we have sufficient zerolag and noise samples before attempting the extinction curve fit
+		# none of the checks below should evaluate to True even with a small amount of zerolag and noise
+		# samples. They should only evaluate to True at the start of an online analysis
+		bg = self.noise_lr_lnpdf.copy().array
+		fg = self.zero_lag_lr_lnpdf.copy().array
+		bg[:10] = 0.
+		fg[:10] = 0.
+		if fg.sum() == 0 or bg.sum() == 0:
+			return False
+
+		fg_ccdf = numpy.cumsum(fg[::-1])[::-1]
+		ix_min = (fg_ccdf < fg_ccdf[0] * self.extinction_fitting_limits[0]).argmax()
+		ix_max = (fg_ccdf < fg_ccdf[0] * self.extinction_fitting_limits[1]).argmax()
+		if ix_min == ix_max:
+			return False
+
+		if fg[ix_min: ix_max + 1].sum() == 0 or bg[ix_min: ix_max + 1].sum() == 0:
+			return False
+
+		if (fg_ccdf[ix_min: ix_max + 1] == 0).any():
+			# log will evaluate to -inf, and the curve fitting will crash
+			return False
+
+		return True
 
-		# the probability that an event survives clustering is the
-		# probability that no event with a higher value of the
-		# ranking statistic is within +/- dt of the event in
-		# question.  .noise_lr_lnpdf.count contains an accurate
-		# model of the counts of pre-clustered events in each
-		# ranking statistic bin, but we know the events are not
-		# independent and so the total number of them cannot be
-		# used to form a Poisson rate for the purpose of computing
-		# the probability we desire.  it has been found,
-		# empirically, that if the CCDF of pre-clustered background
-		# event counts is raised to some power and the clustering
-		# survival probability computed from that as though it were
-		# the CCDF of a Poisson process with some effective mean
-		# event rate, the result is a very good model for the
-		# post-clustering distribution of ranking statistics.  we
-		# have two unknown parameters:  the power to which to raise
-		# the pre-clustered ranking statistic's CCDF, and the mean
-		# event rate to assume.  we solve for these parameters by
-		# fitting to the observed zero-lag ranking statistic
-		# histogram
+
+	def new_with_extinction(self, verbose = False):
+		self = self.copy()
 
 		x = self.noise_lr_lnpdf.bins[0].centres()
 		assert (x == self.zero_lag_lr_lnpdf.bins[0].centres()).all()
 
-		# some candidates are rejected by the ranking statistic,
-		# causing there to be a spike in the zero-lag density at ln
-		# L = -inf.  if enough candidates get rejected this spike
-		# becomes the mode of the PDF which messes up the mask
-		# constructed below for the fit.  we zero the first 40 bins
-		# here to prevent that from happening (assume density
-		# estimation kernel = 4 bins wide, with 10 sigma impulse
-		# length)
-		zl_counts = self.zero_lag_lr_lnpdf.array.copy()
-		zl_counts[:40] = 0.
-		if not zl_counts.any():
-			raise ValueError("zero-lag counts are all zero")
-
-		# get the noise counts
-		noise_counts = self.noise_lr_lnpdf.array.copy()
-
-		# get the zerolag counts.
-		# we model the tail of the distribution - top 0.1 - 1% - where
-		# clustering only effects the result at a < 1% level.
-		if zl_counts.sum() < 100 * 100:
-			tail_zl_counts = zl_counts.sum() * 0.99
-		else:
-			tail_zl_counts = zl_counts.sum() - 100
-		onepercent = zl_counts.cumsum().searchsorted(tail_zl_counts)
-
-		# normalize the counts
-		noise_counts /= noise_counts.sum()
-		zl_counts /= zl_counts.sum()
-
-		# compute survival probability
-		norm = zl_counts[onepercent] / noise_counts[onepercent]
-		zl_counts[onepercent:] = 0
-		noise_counts[onepercent:] = 0
-		survival_probability = zl_counts / noise_counts
-		survival_probability[onepercent:] = norm
-		survival_probability[numpy.isnan(survival_probability)] = 0.0
-
-		# apply to background counts and signal counts
-		self.noise_lr_lnpdf.array *= survival_probability
-		self.noise_lr_lnpdf.normalize()
-		self.signal_lr_lnpdf.array *= survival_probability
-		self.signal_lr_lnpdf.normalize()
+		#FIXME: Add comprehensive explanation about the extinction model
 
-		#
-		# never allow PDFs that have had the extinction model
-		# applied to be written to disk:  on-disk files must only
-		# ever provide the original data.  forbid PDFs that have
-		# been extincted from being re-extincted.
-		#
+		lrs = self.noise_lr_lnpdf.centres()[0]
+		bg = self.noise_lr_lnpdf.array
+		fg = self.zero_lag_lr_lnpdf.array
+
+		# zero out the beginning bins of each because they are notoriously bad and should just be ignored
+		bg[:10] = 0.
+		fg[:10] = 0.
+
+
+		# fitting is done between ix_min and ix_max
+		fg_ccdf = numpy.cumsum(fg[::-1])[::-1]
+		ix_min = (fg_ccdf < fg_ccdf[0] * self.extinction_fitting_limits[0]).argmax()
+		ix_max = (fg_ccdf < fg_ccdf[0] * self.extinction_fitting_limits[1]).argmax()
+
+
+		bgtotal = bg[ix_min: ix_max + 1].sum()
+		bg_ccdf = numpy.cumsum(bg[::-1])[::-1] / bgtotal
 
-		def new_with_extinction(*args, **kwargs):
-			raise NotImplementedError("re-extincting an extincted RankingStatPDF object is forbidden")
-		self.new_with_extinction = new_with_extinction
-		def to_xml(*args, **kwargs):
-			raise NotImplementedError("writing extincted RankingStatPDF object to disk is forbidden")
-		self.to_xml = to_xml
+		# define a function for the extincted bg for scipy.optimize.curve_fit to call
+		def bg_ccdf_extinct_func(idx, c, A):
+			return numpy.log(A * bgtotal / c) + numpy.log1p(-numpy.exp(-1 * bg_ccdf[idx] * c))
+
+		# find the best fit c for extinction
+		c = scipy.optimize.curve_fit(bg_ccdf_extinct_func, range(ix_min, ix_max + 1), numpy.log(fg_ccdf[ix_min: ix_max + 1]), bounds = [0, numpy.inf], sigma = numpy.sqrt(1 / (bg_ccdf[ix_min: ix_max + 1] * bgtotal)))#, maxfev = 5000)
+		if verbose:
+			print(f"Best value of c is {c[0][0]}, A is {c[0][1]} with covariance {c[1]}", file = sys.stderr)
+
+		# calculate the extincted PDF
+		bg_pdf_extinct = c[0][1] * bg * numpy.exp(-1 * bg_ccdf * c[0][0])
+
+		self.noise_lr_lnpdf.array = bg_pdf_extinct
+		self.noise_lr_lnpdf.normalize()
 
 		return self