manifold_cbc_bank_plot_mass_model: add mass-model-investigation feature

bin/manifold_cbc_bank_plot_mass_model

@@ -33,19 +33,24 @@ def parse_args(args=None):
 	p = OptionParser(usage="Usage: %prog [options]")
 	p.add_option("", "--prior", help="Specify a mass model h5 file that is the base for the comparison")
 	p.add_option("", "--post", help="Specify a mass model h5 file that is compared against the base version.")
-	p.add_option("", "--title", default="PostOverPrior", help="Specify title of the plot. Default is PostOverPrior.")
+	p.add_option("", "--title", help="Specify title of the plot. Required if --compare or --mass-model-investigation.")
 	p.add_option("", "--compare", action="store_true", help="Plot comparison between two mass models assuming that the two mass models are generated from the same bank. --prior, --post, --title are required.")
 	p.add_option("", "--snr", type="float", default = 10, help="Specify SNR. Default is 10.")
+	p.add_option("", "--mass-model-investigation", action="store_true", help="Find the population model values for a given template id. --template-id, --title are required.")
+	p.add_option("", "--template-id", type="int", default=None, help="Required for --mass-model-investigation. Provide Gamma0 from the coinc file.")
 	(opt, args) = p.parse_args(args)
 	if opt.compare and None in [opt.title, opt.prior, opt.post]:
-		print("Terminating: --compare needs --title, --post, --prior")
+		print("Terminating: --compare requires --title, --post, --prior")
+		sys.exit()
+	if opt.mass_model_investigation and None in [opt.template_id, opt.title]: 
+		print("Terminating: --mass-model-investigation requires --template-id, --title")
 		sys.exit()
 	return opt, args
 
 def compute_mchirp(m1, m2):
 	return (m1*m2)**0.6/(m1 + m2)**0.2
 
-def get_mass_model_data(mass_model, snr):	
+def get_mass_model_data(mass_model, snr, template_id=None):	
 	model = io.h5read(mass_model)
 	ids = numpy.array(model['template_id'])
 	m1s = numpy.array(model['m1s'])
@@ -56,13 +61,25 @@ def get_mass_model_data(mass_model, snr):
 
 	coeffs = numpy.array([popmodel.lnP_template_signal(tid, snr) for tid in ids])
 	ix = coeffs.argsort()
+	ids = ids[ix]
 	m1s = m1s[ix]
 	m2s = m2s[ix]
 	chis = chis[ix]
 	mcs = mcs[ix]
 	coeffs = coeffs[ix]
 
-	return coeffs, m1s, m2s, chis, mcs
+	if template_id is not None:
+		index = numpy.where(ids == template_id)[0][0]
+		m1 = m1s[index]
+		m2 = m2s[index]
+		chi = chis[index]
+		coeff = coeffs[index]
+		mc = compute_mchirp(m1, m2)
+		
+		return coeffs, m1s, m2s, chis, mcs, coeff, m1, m2, chi, mc
+
+	else:
+		return coeffs, m1s, m2s, chis, mcs
 
 def plot_mass_model(coeffs, m1s, m2s, chis, mcs, snr, mass_model=None, title=None):
 	fig = pyplot.figure()
@@ -107,6 +124,57 @@ def plot_mass_model(coeffs, m1s, m2s, chis, mcs, snr, mass_model=None, title=Non
 		fig.tight_layout()
 		pyplot.savefig('snr-'+str(snr)+'_mc_chi.png')
 
+def plot_mass_model_annotated(coeffs, m1s, m2s, chis, mcs, snr, coeff, m1, m2, chi, mc, mass_model=None, title=None):
+	fig = pyplot.figure()
+	pyplot.scatter(m1s, m2s, c=coeffs, vmax=max(coeffs),vmin=min(coeffs), s=0.5)
+	ax = pyplot.gca()
+	pyplot.xlabel('m1')
+	pyplot.ylabel('m2')
+	ax.set_yscale('log')
+	ax.set_xscale('log')
+	pyplot.xlim(m1-0.5, m1+0.5)
+	pyplot.ylim(m2-0.5, m2+0.5)
+	pyplot.plot(m1, m2, marker='o', markersize=2, color='red')
+	pyplot.colorbar()
+	if title is not None:
+		ax.set_title("SNR="+str(snr)+"  "+title)
+		fig.tight_layout()
+		pyplot.savefig(title+'_snr-'+str(snr)+'_m1_m2.png')
+	elif mass_model is not None:
+		ax.set_title("SNR="+str(snr))
+		fig.tight_layout()
+		pyplot.savefig(mass_model.replace('.h5','snr-'+str(snr)+'_m1_m2.png'))
+	else:
+		ax.set_title("SNR="+str(snr))
+		fig.tight_layout()
+		pyplot.savefig('snr-'+str(snr)+'_m1_m2.png')
+
+	fig = pyplot.figure()
+	pyplot.scatter(mcs, chis, c=coeffs, vmax=max(coeffs),vmin=min(coeffs), s=0.5)
+	ax = pyplot.gca()
+	pyplot.xlabel(r'$\mathcal{M}$')
+	pyplot.ylabel(r'$\chi_{\rm{eff}}$')
+	ax.set_yscale('linear')
+	ax.set_xscale('log')
+	pyplot.xlim(mc-2.0, mc+2.0)
+	pyplot.ylim(chi-0.05, chi+0.05)
+	pyplot.plot(mc, chi, marker='o', markersize=2, color='red')
+	pyplot.colorbar()
+	if title is not None:
+		ax.set_title("SNR="+str(snr)+"  "+title)
+		fig.tight_layout()
+		pyplot.savefig(title+'_snr-'+str(snr)+'_mc_chi.png')
+	elif mass_model is not None:
+		ax.set_title("SNR="+str(snr))
+		fig.tight_layout()
+		pyplot.savefig(mass_model.replace('.h5','snr-'+str(snr)+'_mc_chi.png'))
+	else:
+		ax.set_title("SNR="+str(snr))
+		fig.tight_layout()
+		pyplot.savefig('snr-'+str(snr)+'_mc_chi.png')
+
+	print('coeff: %f, m1: %f, m2: %f, chi: %f, mc: %f' % (coeff, m1, m2, chi, mc))
+
 def main(args=None, show_plot: bool = True, verbose: bool = True):
 	opt, args = parse_args(args)
 	try:
@@ -115,8 +183,12 @@ def main(args=None, show_plot: bool = True, verbose: bool = True):
 		snr=10   
 
 	if not opt.compare:
-		coeffs, m1s, m2s, chis, mcs = get_mass_model_data(args[0], snr)
-		plot_mass_model(coeffs, m1s, m2s, chis, mcs, snr, args[0])
+		if opt.mass_model_investigation:
+			coeffs, m1s, m2s, chis, mcs, coeff, m1, m2, chi, mc = get_mass_model_data(args[0], snr, opt.template_id)
+			plot_mass_model_annotated(coeffs, m1s, m2s, chis, mcs, snr, coeff, m1, m2, chi, mc, args[0], title=opt.title)
+		else:
+			coeffs, m1s, m2s, chis, mcs = get_mass_model_data(args[0], snr)
+			plot_mass_model(coeffs, m1s, m2s, chis, mcs, snr, args[0], title=opt.title)
 
 	if opt.compare:
 		prior_coeffs, prior_m1s, prior_m2s, prior_chis, prior_mcs = get_mass_model_data(opt.prior, snr)
@@ -124,6 +196,6 @@ def main(args=None, show_plot: bool = True, verbose: bool = True):
 		diff_coeffs = post_coeffs - prior_coeffs 
 		plot_mass_model(diff_coeffs, prior_m1s, prior_m2s, prior_chis, prior_mcs, snr, title=opt.title)
 
-
 if __name__ == '__main__':
 	main()
+