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/bin/gstlal_inspiral_marginalize_likelihoods_online

@@ -152,23 +152,8 @@ def main():
 	# set up the output paths
 	#
 
-	marg_pdf_exists = os.path.isfile(options.output)
-	pdfs = DataCache.find(DataType.DIST_STAT_PDFS, svd_bins = "*")
-	if marg_pdf_exists and len(pdfs) == len(registries):
-		files_exist = True
-	elif not marg_pdf_exists and len(pdfs) == 0:
-		files_exist = False
-	elif marg_pdf_exists and len(pdfs) != len(registries): 
-		raise ValueError(f"Number of registry files provided ({len(registries)}) does not match number of DIST_STAT_PDF files found ({len(pdfs)})")
-	else:
-		raise ValueError("Could not find marg DIST_STAT_PDF file")
-
 	svd_bins = [reg[:4] for reg in registries]
-	if files_exist:
-		assert set(pdfs.groupby('svd_bin').keys()) == set(svd_bins), "svd bins of registry files are not the same as svd bins of found PDFs"
-	else:
-		pdfs = DataCache.generate(DataType.DIST_STAT_PDFS, CacheEntry.from_T050017(options.output).observatory, svd_bins = svd_bins)
-
+	pdfs = DataCache.generate(DataType.DIST_STAT_PDFS, CacheEntry.from_T050017(options.output).observatory, svd_bins = svd_bins)
 	pdfs = pdfs.groupby('svd_bin')
 
 	#
@@ -220,8 +205,12 @@ def main():
 			url = url_from_registry(reg, likelihood_path)
 
 			svd_bin = reg[:4]
-			# load the old ranking stat pdf for this bin:
-			old_pdf = far.parse_likelihood_control_doc(ligolw_utils.load_url(pdfs[svd_bin][0], verbose = options.verbose, contenthandler = far.RankingStat.LIGOLWContentHandler)) if files_exist else None
+			if os.path.isfile(pdfs[svd_bin].files[0]):
+				# load the old ranking stat pdf for this bin:
+				old_pdf = far.parse_likelihood_control_doc(ligolw_utils.load_url(pdfs[svd_bin].files[0], verbose = options.verbose, contenthandler = far.RankingStat.LIGOLWContentHandler))
+			else:
+				logging.warning(f"Couldn't find {pdfs[svd_bin].files[0]}, starting from scratch")
+				old_pdf = None
 
 			# create the new ranking stat pdf and marginalize as we go
 			new_pdf_status, pdf = calc_rank_pdfs(url, ranking_stat_samples, options.num_cores, verbose = options.verbose)
@@ -256,16 +245,23 @@ def main():
 				# get the zerolag pdf for this bin and use it to perform bin-specific extinction
 				zerolag_counts_url = url_from_registry(reg, zerolag_counts_path)
 				pdf += far.RankingStatPDF.from_xml(ligolw_utils.load_url(zerolag_counts_url, verbose = options.verbose, contenthandler = far.RankingStat.LIGOLWContentHandler), u"gstlal_inspiral_likelihood")
-				if data:
-					data += pdf.new_with_extinction()
+				if len(numpy.nonzero(pdf.zero_lag_lr_lnpdf.array)[0]) > 0:
+					# LR calculation has started and we are ready to perform first-round extinction
+					if data:
+						data += pdf.new_with_extinction()
+					else:
+						data = pdf.new_with_extinction()
 				else:
-					data = pdf.new_with_extinction()
+					logging.warning(f'Skipping first-round extinction for {pdfs[svd_bin].files[0]}')
 
 			# while looping through registries
 			# send heartbeat messages
 			if kafka_processor:
 				kafka_processor.heartbeat()
-
+		
+		# zero out the zerolag after the first round of extinction is finished
+		if data:
+			data.zero_lag_lr_lnpdf.count.array[:] = 0
 
 		# if we fail to complete more than 1% of the bins,
 		# this is a serious problem and we should just quit
@@ -291,7 +287,10 @@ def main():
 		zerolag_counts_url = url_from_registry("gstlal_ll_inspiral_trigger_counter_registry.txt", zerolag_counts_path)
 
 		# add zerolag counts url to marginalized data
-		data += far.RankingStatPDF.from_xml(ligolw_utils.load_url(zerolag_counts_url, verbose = options.verbose, contenthandler = far.RankingStat.LIGOLWContentHandler), u"gstlal_inspiral_likelihood")
+		if data:
+			data += far.RankingStatPDF.from_xml(ligolw_utils.load_url(zerolag_counts_url, verbose = options.verbose, contenthandler = far.RankingStat.LIGOLWContentHandler), u"gstlal_inspiral_likelihood")
+		else:
+			data = far.RankingStatPDF.from_xml(ligolw_utils.load_url(zerolag_counts_url, verbose = options.verbose, contenthandler = far.RankingStat.LIGOLWContentHandler), u"gstlal_inspiral_likelihood")
 
 		if kafka_processor:
 			kafka_processor.heartbeat()
@@ -309,10 +308,6 @@ def main():
 		ligolw_utils.write_filename(xmldoc, options.output, verbose = options.verbose)
 		logging.info(f"Done marginalizing likelihoods.")
 
-		# we just created the bin-specific and marg DIST_STAT_PDFs,
-		# so the files definitely exist for the next iteration of the loop
-		files_exist = True
-
 		if kafka_processor:
 			kafka_processor.heartbeat()