svd_bank_snr.py: modify output format

gstlal-inspiral/bin/gstlal_inspiral_calc_snr

@@ -26,7 +26,10 @@ Typical Usages:
 	--output-width (default = 32bits)
 	--instrument (require)
 	--outdir (require)
+	--start (default = None)
+	--end (default = None)
 	--verbose (default = False)
+	--complex (defaule = False)
 
 --mode 1 (calculate SNR using Finite Impulse Response):
 	1. GW options: also see datasource.GWDataSourceInfo()
@@ -54,7 +57,10 @@ Typical Usages:
 	--output-width (default = 32bits)
 	--instrument (require)
 	--outdir (require)
+	--start (default = None)
+	--end (default = None)
 	--verbose (default = False)
+	--complex (default = False)
 """
 import numpy
 from optparse import OptionParser, OptionGroup, IndentedHelpFormatter
@@ -129,6 +135,7 @@ def parse_command_line():
 	group = OptionGroup(parser, "Output Control Options", "Control SNR output")
 	group.add_option("--outdir", metavar = "directory", type = "str", help = "Output directory for SNR(s) (requires).")
 	group.add_option("--mode", metavar = "method", type = "int", default = 0, help = "The method (0 = LLOID / 1 = FIR) that is used to calculate SNR (default = 0).")
+	group.add_option("--complex", action = "store_true", help = "Choose whether to output the complex snr or not.")
 	group.add_option("--drop-first", metavar = "seconds", type = "int", default = 0, help = "Dropping the initital '--drop-first' seconds of SNR data (default = 0).")
 	group.add_option("--drop-last", metavar = "seconds", type = "int", default = 0, help = "Dropping the last '--drop-last' seconds of SNR data (default = 0).")
 	group.add_option("--output-width", metavar = "bits", type = "int", default = 32, help = "The size of the output data, can only be 32 or 64 bits (default = 32 bits).")
@@ -141,6 +148,8 @@ def parse_command_line():
 
 	# Setting up GW data
 	gw_data_source_info = datasource.GWDataSourceInfo(options)
+	if options.instrument not in gw_data_source_info.channel_dict.keys():
+		raise ValueError("No such instrument: %s in GWDataSourceInfo: (%s)"% (options.instrument, ", ".join(gw_data_source_info.channel_dict.keys())))
 
 	# Setting up PSD
 	if options.reference_psd:
@@ -174,14 +183,10 @@ def parse_command_line():
 			try:
 				bank = banks[options.sub_bank_id]
 			except IndexError:
-				raise IndexError("Invaild --bank-id %d. Possible id (0-%d)\n" % (options.sub_bank_id, len(banks) - 1))
+				raise IndexError("Invaild --sub-bank-id %d. Possible id (0-%d)\n" % (options.sub_bank_id, len(banks) - 1))
 		if options.row_number >=  bank.bank_fragments[0].mix_matrix.shape[1] / 2:
 			raise ValueError("No such template: Invaild --row-number %d. Possible range (0-%d)\n" % (options.row_number, bank.bank_fragments[0].mix_matrix.shape[1] / 2 - 1))
 
-		# Setting up output
-		if options.instrument not in set(gw_data_source_info.channel_dict):
-			raise ValueError("No such instrument: %s in detectos: (%s)"% (options.instrument, ", ".join(set(gw_data_source_info.channel_dict))))
-
 		return options, gw_data_source_info, bank, psd
 
 	# Use Finite Impulse Response
@@ -242,14 +247,15 @@ if options.mode == 0:
 		verbose = options.verbose
 	)
 
-	if options.row_number is None:
-		# FIXME: put it in one xmldoc
-		for row in range(num_of_row):
-			snrdict = {options.instrument : lloid_snr(row_number = row, drop_first = options.drop_first, drop_last = options.drop_last)}
-			svd_bank_snr.write_url(svd_bank_snr.make_xmldoc(snrdict), os.path.join(options.outdir, "snr_%d.xml.gz" % row), verbose = options.verbose)
-	else:
-		snrdict = {options.instrument : lloid_snr(row_number = options.row_number, drop_first = options.drop_first, drop_last = options.drop_last)}
-		svd_bank_snr.write_url(svd_bank_snr.make_xmldoc(snrdict), os.path.join(options.outdir, "snr.xml.gz"), verbose = options.verbose)
+	for index, snr in enumerate(lloid_snr(COMPLEX = options.complex, row_number = options.row_number, drop_first = options.drop_first, drop_last = options.drop_last)):
+		snrdict = {options.instrument: [snr]}
+		svd_bank_snr.write_url(svd_bank_snr.make_xmldoc(snrdict), os.path.join(options.outdir, "%s-SNR_%d-%d-%d.xml.gz" % (options.instrument, index, int(snr.epoch), int(snr.data.length * snr.deltaT))), verbose = options.verbose)
+	#
+	# uncomment to save all snrs in one single XML
+	#
+	#snrdict = {options.instrument : lloid_snr(COMPLEX = options.complex, row_number = options.row_number, drop_first = options.drop_first, drop_last = options.drop_last)}
+	#tempsnr = 
+	#svd_bank_snr.write_url(svd_bank_snr.make_xmldoc(snrdict), os.path.join(options.outdir, "%s-SNR-%d-%d.xml.gz" % (options.instrument, int(snrdict.values()[0].epoch), int(snrdict.values()[0].data.length * snrdict.values()[0].deltaT))), verbose = options.verbose)
 
 elif options.mode == 1:
 	#FIXME: proper handle for latency
@@ -266,7 +272,5 @@ elif options.mode == 1:
 		verbose = options.verbose
 		)
 
-	#FIXME: allow multiple instruments
-	snrdict = {options.instrument : fir_snr(drop_first = options.drop_first, drop_last = options.drop_last)}
-
-	svd_bank_snr.write_url(svd_bank_snr.make_xmldoc(snrdict),os.path.join(options.outdir, "snr.xml.gz"), verbose = options.verbose)
+	snrdict = {options.instrument : fir_snr(COMPLEX = options.complex, drop_first = options.drop_first, drop_last = options.drop_last)}
+	svd_bank_snr.write_url(svd_bank_snr.make_xmldoc(snrdict),os.path.join(options.outdir, "%s-SNR-%d-%d.xml.gz" % (options.instrument, int(snrdict.values()[0].epoch), int(snrdict.values()[0].data.length * snrdict.values()[0].deltaT))), verbose = options.verbose)

gstlal-inspiral/bin/gstlal_inspiral_plot_snr

@@ -2,32 +2,45 @@
 """
 Plotter for gstlal_inspiral_calc_snr
 """
+import matplotlib
+matplotlib.use("Agg")
+from matplotlib import pyplot
+import os
 import sys
 from optparse import OptionParser
 
 from gstlal import plotsnr
 
+from lal.utils import CacheEntry
+
 def parse_command_line():
 	parser = OptionParser(description = __doc__)
 
-	parser.add_option("-o", "--output", metavar = "filename", help = "The output filename for the SNR plot (require).")
-	parser.add_option("-i", "--input", metavar = "filename", help = "A LIGO light-weight XML file containing SNR time series data (require).")
+	parser.add_option("-o", "--outdir", metavar = "directory", help = "The output directory for the SNR plot (require).")
+	parser.add_option("-i", "--input", metavar = "cache", help = "The input cache containing SNR urls (require).")
+	parser.add_option("-f", "--format", metavar = "file_type", help = "The format of the output plot, can only be '.png' or '.svg' (default = .svg).")
 	parser.add_option("--center", metavar = "gpsSeconds", type = "float", help = "Center the plot to --center (optional).")
+	parser.add_option("--width", metavar = "inch", default = 8, type = "int", help = "The output width of the figure.")
 	parser.add_option("--span", metavar = "seconds", type = "float", help = "The time span around --center (optional).")
 	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
 
 	options, args = parser.parse_args()
 
-	missing_required_options = []
 	if options.input is None:
-		missing_required_options.append("--input")
-	if options.output is None:
-		missing_required_options.append("--output")
-	if missing_required_options:
-		raise ValueError("Missing required option(s) %s" % ", ".join(missing_required_options))
+		raise ValueError("Missing --input.")
+	if options.outdir is None:
+		raise ValueError("Missing --outdir.")
+	if options.format is not None:
+		if options.format != ".png" and options.format != ".svg":
+			raise ValueError("--format can only be '.png' or '.svg'.")
 
 	return options
 
 options = parse_command_line()
 
-plotsnr.plot_snr(options.input, output = options.output, center = options.center, span = options.span, verbose = options.verbose)
+with open(options.input) as f:
+	suffix = ".svg" if (options.format is None) else options.format
+	for snr_file in f:
+		figure = plotsnr.plot_snr(CacheEntry(snr_file).url, width = options.width, center = options.center, span = options.span, verbose = options.verbose)
+		figure.savefig(os.path.join(options.outdir, CacheEntry(snr_file).description+suffix))
+		pyplot.close(pyplot.gcf())

gstlal-inspiral/python/plotsnr.py

@@ -8,51 +8,118 @@ import matplotlib
 matplotlib.use("Agg")
 from matplotlib import pyplot
 
+from gstlal import plotutil
 from gstlal import svd_bank_snr
 
-def plot_snr(filename, output = None, center = None, span = None, verbose = False):
-	xmldoc = svd_bank_snr.read_url(filename, contenthandler = svd_bank_snr.SNRContentHandler, verbose = verbose)
-	SNR_dict = svd_bank_snr.read_xmldoc(xmldoc)
+def axes_add_snr(axes, snrdict, center = None, span = None):
+	"""
+
+	Add snr time series to the plot
+
+	Args:
+		axes (object): matplotlib axes
+		snrdict (dict): dictionary of snrs
+		center (float): the center gpstime of the plot
+		span (float): seconds to span around center
+
+	"""
+	col = 0
+	for instrument, SNRs in snrdict.items():
+		if len(SNRs) > 1:
+			raise ValueError("Too many snr time series in snrdict.")
+
+		data, center_gps_time, relative_gps_time = locate_center(SNRs[0], center, span)
+		if numpy.iscomplexobj(data):
+			axes[col].plot(relative_gps_time, data.real, color = plotutil.colour_from_instruments([instrument]), linestyle = "-", label = "%s_Real" % instrument, linewidth = 1)
+			axes[col].plot(relative_gps_time, data.imag, color = plotutil.colour_from_instruments([instrument]), linestyle = "--", label = "%s_Imaginary" % instrument, linewidth = 1)
+		else:
+			axes[col].plot(relative_gps_time, data, color = plotutil.colour_from_instruments([instrument]), label = "%s" %instrument, linewidth = 1)
+		axes[col].set_xlabel("GPS Time Since %f" % center_gps_time)
+		axes[col].set_ylabel("SNR")
+		axes[col].legend(loc = "upper left")
+		axes[col].tick_params(labelbottom = True)
+		col += 1
+
+def plot_snr(filename, width = 6, center = None, span = None, verbose = False):
+	"""
+
+	Plot snr time series from LIGO light-weighted XML file
+
+	Args:
+		filename: A LIGO light-weighted XML file
+		width (int): width of the output figure in inch
+		center (float): the center gpstime of the plot
+		span (float): seconds to span around center
+		verbose (bool): be verbose
+
+	Return:
+
+		fig (object): matplotlib figure
+	"""
+
+	xmldoc = svd_bank_snr.read_url(filename, svd_bank_snr.SNRContentHandler, verbose = verbose)
+	SNRs_dict = svd_bank_snr.read_xmldoc(xmldoc)
+
+	nrows = len(SNRs_dict.keys())
+	ncols = 1
+
+	fig, axes = pyplot.subplots(nrows = nrows, ncols = ncols, sharex = True)
+	if nrows == 1:
+		axes = [axes]
+	axes_add_snr(axes, SNRs_dict, center = center, span = span)
+	fig.set_size_inches(width, int(round(width/plotutil.golden_ratio)))
+	fig.tight_layout(pad = 0.8)
+
+	return fig
+
+def locate_center(snr_series, center, span):
+	"""
+
+	locate the snr_series at nearest gpstime with certain span
+
+	Args:
+		snr_series (lal.series): A (snr) lal.series
+		center (float): gps time
+		span (float): seconds to span around center
+
+	Return:
+		data (numpy.array <type 'float'>): snr_series.data.data [center-span, center+span]
+		gps_time (float): nearest gps_time at center
+		relative_gps_time (numpy.array <type 'float'>): gpstime relative to gps_time
+
+	"""
+
 	start = None
 	mid = 0
 	end = None
+	if center and span:
+		start = find_nearest(snr_series, center - span)[1]
+		mid = find_nearest(snr_series, center)[1]
+		end = find_nearest(snr_series, center + span)[1]
 
-	if verbose:
-		sys.stderr.write("Ploting SNR ...\n")
-
-	for instrument, SNR in SNR_dict.items():
-		if verbose and center and span:
-			sys.stderr.write("Locating SNR at GPSTime: %f spanning %f s\n" % (center, span))
-		if center and span:
-			start = find_nearest(SNR, center - span)[1]
-			mid = find_nearest(SNR, center)[1]
-			end = find_nearest(SNR, center + span)[1]
-
-		gps_time = (SNR.epoch.gpsSeconds + SNR.epoch.gpsNanoSeconds * 10.**-9 + (numpy.arange(SNR.data.length) * SNR.deltaT))
-		relative_gps_time = (gps_time - gps_time[mid])[start:end]
-		data = SNR.data.data[start:end]
-
-		fig, ax = pyplot.subplots(nrows = 1, ncols = 1, figsize = [15,6])
-		ax.plot(relative_gps_time, data)
-		ax.set_xlabel("GPS Time Since %f" % gps_time[mid])
-		ax.set_ylabel("SNR")
-
-		pyplot.tight_layout()
-		if output is None:
-			output = "SNR_%s_since_%d.svg" %(SNR.name, gps_time[mid])
-			if verbose:
-				sys.stderr.write("%s --> Done\n" % output)
-			fig.savefig(output)
-		else:
-			if verbose:
-				sys.stderr.write("%s --> Done\n" % output)
-			fig.savefig(output)
-		pyplot.close()
+	gps_time = (snr_series.epoch.gpsSeconds + snr_series.epoch.gpsNanoSeconds * 1.e-9 + (numpy.arange(snr_series.data.length) * snr_series.deltaT))
+	relative_gps_time = (gps_time - gps_time[mid])[start:end]
+	data = snr_series.data.data[start:end]
+
+	return data, gps_time[mid], relative_gps_time
 
 def find_nearest(snr_series, time):
+	"""
+
+	Find the nearest gps time available from the time series
+
+	Args:
+		snr_series (lal.series): A (snr) lal-series
+		time (float): requested gpstime
+
+	Return:
+		(snr[gpstime_index], gpstime_index)
+
+	"""
+
 	gps_start = snr_series.epoch.gpsSeconds + snr_series.epoch.gpsNanoSeconds * 10.**-9
 	gps = gps_start + numpy.arange(snr_series.data.length) * snr_series.deltaT
-	if time - gps[0] > 0 and time - gps[-1] <0:
+	if time - gps[0] >= 0 and time - gps[-1] <= 0:
 		index = abs(gps - time).argmin()
 	else:
 		raise ValueError("Invalid choice of center time %f." % time)

gstlal-inspiral/python/svd_bank_snr.py

@@ -72,18 +72,11 @@ class SNR_Pipeline(object):
 		if self.pipeline.set_state(Gst.State.NULL) != Gst.StateChangeReturn.SUCCESS:
 			raise RuntimeError("pipeline could not be set to NULL.")
 
-	def get_snr_series(self, row_number = 0, drop_first = 0, drop_last = 0):
-		assert drop_first >= 0, "must drop positive number of data"
-		assert drop_last >= 0, "must drop positive number of data"
-		bps = drop_first * int(round(1 / self.snr_info["deltaT"]))
-		bpe = -drop_last * int(round(1 / self.snr_info["deltaT"])) if drop_last != 0 else None
-
-		data = numpy.abs(self.snr_info["data"])[:,row_number][bps:bpe]
-
+	def make_series(self, data):
 		if data.dtype == numpy.float32:
 			tseries = lal.CreateREAL4TimeSeries(
 					name = self.snr_info["instrument"],
-					epoch = self.snr_info["epoch"] + drop_first,
+					epoch = self.snr_info["epoch"],
 					deltaT = self.snr_info["deltaT"],
 					f0 = 0,
 					sampleUnits = lal.DimensionlessUnit,
@@ -93,18 +86,66 @@ class SNR_Pipeline(object):
 		elif data.dtype == numpy.float64:
 			tseries = lal.CreateREAL8TimeSeries(
 					name = self.snr_info["instrument"],
-					epoch = self.snr_info["epoch"] + drop_first,
+					epoch = self.snr_info["epoch"],
+					deltaT = self.snr_info["deltaT"],
+					f0 = 0,
+					sampleUnits = lal.DimensionlessUnit,
+					length = len(data)
+					)
+			tseries.data.data = data
+		elif data.dtype == numpy.complex64:
+			tseries = lal.CreateCOMPLEX8TimeSeries(
+					name = self.snr_info["instrument"],
+					epoch = self.snr_info["epoch"],
+					deltaT = self.snr_info["deltaT"],
+					f0 = 0,
+					sampleUnits = lal.DimensionlessUnit,
+					length = len(data)
+					)
+			tseries.data.data = data
+		elif data.dtype == numpy.complex128:
+			tseries = lal.CreateCOMPLEX16TimeSeries(
+					name = self.snr_info["instrument"],
+					epoch = self.snr_info["epoch"],
 					deltaT = self.snr_info["deltaT"],
 					f0 = 0,
 					sampleUnits = lal.DimensionlessUnit,
 					length = len(data)
 					)
 			tseries.data.data = data
+
 		else:
 			raise ValueError("unsupported type : %s " % data.dtype)
 
 		return tseries
 
+	def get_snr_series(self, COMPLEX = False, row_number = None, drop_first = 0, drop_last = 0):
+		assert drop_first >= 0, "must drop positive number of data"
+		assert drop_last >= 0, "must drop positive number of data"
+		bps = drop_first * int(round(1 / self.snr_info["deltaT"]))
+		bpe = -drop_last * int(round(1 / self.snr_info["deltaT"])) if drop_last != 0 else None
+
+		self.snr_info["epoch"] += drop_first
+		self.snr_info["data"] = self.snr_info["data"][bps:bpe].T
+
+		if row_number is None:
+			temp = []
+			if COMPLEX:
+				for data in self.snr_info["data"]:
+					temp.append(self.make_series(data))
+				return temp
+			else:
+				for data in self.snr_info["data"]:
+					temp.append(numpy.abs(self.make_series(data)))
+				return temp
+		else:
+			self.snr_info["data"] = self.snr_info["data"][row_number]
+			if COMPLEX:
+				return [self.make_series(self.snr_info["data"])]
+			else:
+				return [self.make_series(numpy.abs(self.snr_info["data"]))]
+
+
         def new_preroll_handler(self, elem):
                 with self.lock:
                         # ignore preroll buffers
@@ -192,8 +233,8 @@ class LLOID_SNR(SNR_Pipeline):
 		self.run(gw_data_source_info.seg)
                 self.snr_info["data"] = numpy.concatenate(numpy.array(self.snr_info["data"]), axis = 0)
 
-	def __call__(self, row_number = 0, drop_first = 0, drop_last = 0):
-		return self.get_snr_series(row_number, drop_first, drop_last)
+	def __call__(self, COMPLEX = False, row_number = 0, drop_first = 0, drop_last = 0):
+		return self.get_snr_series(COMPLEX, row_number, drop_first, drop_last)
 
 class FIR_SNR(SNR_Pipeline):
 	def __init__(self, gw_data_source_info, template, instrument, rate, latency, psd = None, psd_fft_length = 32, ht_gate_threshold = float("inf"), veto_segments = None, width = 32, track_psd = False, verbose = False):
@@ -241,8 +282,8 @@ class FIR_SNR(SNR_Pipeline):
 		self.snr_info["data"] = numpy.vectorize(complex)(self.snr_info["data"][:,0], self.snr_info["data"][:,1])
 		self.snr_info["data"].shape = len(self.snr_info["data"]), 1
 
-	def __call__(self, row_number = 0 , drop_first = 0, drop_last = 0):
-		return self.get_snr_series(row_number, drop_first, drop_last)
+	def __call__(self, COMPLEX = False, row_number = 0 , drop_first = 0, drop_last = 0):
+		return self.get_snr_series(COMPLEX, row_number, drop_first, drop_last)
 
 #=============================================================================================
 #
@@ -256,32 +297,49 @@ def make_xmldoc(snrdict, xmldoc = None, root_name = u"gstlal_inspiral_snr"):
 
 	root = xmldoc.appendChild(ligolw.LIGO_LW())
 	root.Name = root_name
-	for instrument, snr in snrdict.items():
-		if snr.data.data.dtype == numpy.float32:
-			tseries = root.appendChild(lal.series.build_REAL4TimeSeries(snr))
-		elif snr.data.data.dtype == numpy.float64:
-			tseries = root.appendChild(lal.series.build_REAL8TimeSeries(snr))
-		else:
-			raise ValueError("unsupported type : %s" % snr.data.data.dtype)
-		if instrument is not None:
-			tseries.appendChild(ligolw_param.Param.from_pyvalue(u"instrument", instrument))
-
+	for instrument, snrs in snrdict.items():
+		for snr in snrs:
+			if snr.data.data.dtype == numpy.float32:
+				tseries = root.appendChild(lal.series.build_REAL4TimeSeries(snr))
+			elif snr.data.data.dtype == numpy.float64:
+				tseries = root.appendChild(lal.series.build_REAL8TimeSeries(snr))
+			elif snr.data.data.dtype == numpy.complex64:
+				tseries = root.appendChild(lal.series.build_COMPLEX8TimeSeries(snr))
+			elif snr.data.data.dtype == numpy.complex128:
+				tseries = root.appendChild(lal.series.build_COMPLEX16TimeSeries(snr))
+			else:
+				raise ValueError("unsupported type : %s" % snr.data.data.dtype)
 	return xmldoc
 
 def read_xmldoc(xmldoc, root_name = u"gstlal_inspiral_snr"):
 	if root_name is not None:
 		xmldoc, = (elem for elem in xmldoc.getElementsByTagName(ligolw.LIGO_LW.tagName) if elem.hasAttribute(u"Name") if elem.Name == root_name)
 
-	result = []
+	result = {}
+	temp = []
 	for elem in (elem for elem in xmldoc.getElementsByTagName(ligolw.LIGO_LW.tagName) if elem.hasAttribute(u"Name")):
 		if elem.Name == u"REAL4TimeSeries":
-			result.append((ligolw_param.get_pyvalue(elem, u"instrument"), lal.series.parse_REAL4TimeSeries(elem)))
+			tseries = lal.series.parse_REAL4TimeSeries(elem)
+			temp.append([tseries.name, tseries])
 		elif elem.Name == u"REAL8TimeSeries":
-			result.append((ligolw_param.get_pyvalue(elem, u"instrument"), lal.series.parse_REAL8TimeSeries(elem)))
+			tseries = lal.series.parse_REAL8TimeSeries(elem)
+			temp.append([tseries.name, tseries])
+		elif elem.Name == u"COMPLEX8TimeSeries":
+			tseries = lal.series.parse_COMPLEX8TimeSeries(elem)
+			temp.append([tseries.name, tseries])
+		elif elem.Name == u"COMPLEX16imeSeries":
+			tseries = lal.series.parse_COMPLEX16TimeSeries(elem)
+			temp.append([tseries.name, tseries])
+
+	for i in temp:
+		if i[0] in result.keys():
+			result[i[0]].append(i[1])
+		else:
+			result[i[0]] = [i[1]]
 
 	assert result is not None, "xmldoc contains no LAL Series or LAL Series is unsupported"
 
-	return dict(result)
+	return result
 
 # wrapper for writing snr series to URL
 def write_url(xmldoc, filename, verbose = False):