Updates to megaplot and adding BW_Flags

1. Adding support for cleaning model
2. Making megaplot python3 instead of 2
3. Log scaling on spectrograms
4. Plasma colomap on spectrograms
5. Support for CBC, full, clean_only, CBC+Glitch, CBC+Signal
6. Cornerplots for signal and glitch models
7. CBC Cornerplots for cbc models

BW_Flags, can now read in BW information into python Call via

    from bayeswave_plot import BW_Flags as bwf

    runFlags = bwf.Flags(TRIGTIME_DIRECTORY)
    bwf.readbwb(runFlags)

Megaplot Examples:

cleanOnly example: https://ldas-jobs.ligo.caltech.edu/~sophie.hourihane/forDerek/ScatteringInjs/H1_MULTIPLE_10/trigtime_8.000000000_0.0_0.0_0/

CBCGlitch example: https://ldas-jobs.ligo.caltech.edu/~sophie.hourihane/glitch_events/BlipLowFreq/BlipLowFreq1/GW150914/August/bliplowfreq1_GW150914/trigtime_1256909978.276999950_0.0_0.0_0/

fullOnly (ie signal + glitch) : https://ldas-jobs.ligo.caltech.edu/~sophie.hourihane/glitch_events/sudarshan/trigtime_1165069535.000000000_0.0_0.0_4/

GlitchOnly: https://ldas-jobs.ligo.caltech.edu/~sophie.hourihane/glitch_events/Whistle/Whistle1/noInj_Dmax300_Niter4e6_window2_srate8192

signal glitch example (separate) : https://ldas-jobs.ligo.caltech.edu/~sophie.hourihane/forNayyer/006_sophie_megaplot/trigtime_1240704078.500000000_0.0_0.0_0/

BayesWaveUtils/bayeswave_plot/BW_Flags.py

+#!/usr/bin/env python
+from bayeswave_pipe import bayeswave_pipe_utils as pipe_utils
+import numpy as np
+import os
+import sys
+
+"""
+This script contains a class to display the results of a BayesWave run.
+Written by Sophie Hourihane based on code written by Meg Millhouse and probably others 
+"""
+"""
+    Class containing run information for a BayesWave run
+    The parameters are kept inside of trigdir / bayeswave.run
+
+    Must use readbwb to actually initialize values 
+
+    Example use:
+        # Initialize runFlags
+        runFlags = Flags('/home/sophie.hourihane/public_html/glitch_events/noNoise/seeds_noClean/0noise_2048_5chains_heterodyne_distance600')
+        # Set runFlags
+        
+"""
+class Flags:
+    def __init__(self, trigdir):
+        # Parameters describing where things live
+        if trigdir[-1] == '/':
+            self.trigdir = trigdir[:-1]
+        else:
+            self.trigdir = trigdir
+        self.XMLfname = None
+        self.jobName=''
+
+        # Run settings 
+        self.trigtime = 0 # gps time of event 
+        self.flow = 32
+        self.srate = 2048
+        self.segment_start = 0
+        self.seglen = 0
+        self.Nchain = 20  # 20 is default number of chains 
+        self.Niter =  4000000
+        self.Ncycle = 100
+        self.chirplets = False
+
+        # Windows, cleaning, and bayesLine
+        self.bayesLine = False
+        self.noClean = 0
+        self.CBCwindow = -1 #window around cbc event
+        self.window    = 0
+
+        # Parameters describing what run type was used
+        self.signalGlitch_flag = False
+        self.fullOnly_flag = False
+        self.GlitchCBC_flag = False
+        self.SignalCBC_flag = False
+        self.glitchOnly_flag = False
+        self.cbcOnly_flag = False
+        self.CBCType = False # ie, literally if CBC is in the modellist 
+        self.CBCName = False # I think this is unused now
+        self.multi_type = False # true for full, cbcsignal, cbcglitch
+        self.noNoiseFlag = False
+        self.cleanOnly_flag = False
+
+        self.modelList = [] # Gets filled in readbwb, theres probably a clearer way to to this
+
+        # Injection settings 
+        self.injFlag = False
+        self.event = 0
+        self.mdc = False
+        self.inj_dat = None  #If there is an injection, this holds an injection table
+        self.inj_fref = 0
+        
+        # Likelihood calculation
+        self.heterodyneL = False
+
+        # Printed output 
+        self.verbose = 0
+        self.checkpoint = False
+
+        # Lists of IFOs  
+        self.ifoList = []
+        self.ifoNames = []
+        self.snrList = []
+
+    def run_dirname(self):
+        if self.fullOnly_flag:
+            return('full')
+        if self.SignalCBC_flag:
+            return('cbcsignal')
+        if self.GlitchCBC_flag:
+            return('cbcglitch')
+        if self.signalGlitch_flag:
+            # case where signal and glitch are run together
+            # but not on top of eachother 
+            return('signal')
+        # if there is only one thing in the modellist, return that one 
+        if len(self.modelList) == 1:
+            return(self.modelList[0])
+        if 'clean' in self.modelList:
+            # if there is clean in the modelList and only one other thing, return the other thing 
+            if len(self.modelList) - 1 == 1:
+                return(self.modelList[0])
+        else:
+            print("uh oh! what's the dirname?")
+            return('')
+
+    """
+    Given an xml parameter name, will return that value
+    Accepted names are: 
+    param_names = ['M1_s', 'M2_s', 'Mchirp_s', 'Mtot_s', 
+                 's1', 's2', 'chi_eff', 
+                 'phase', 't', 
+                 'distance', 'redshift',
+                 'RA', 'sindec', 'psi', 'cosiota',
+                 'M1_d', 'M2_d', 'Mchirp_d', 'Mtot_d'] 
+    """   
+    def get_xml_parameter(self, p):
+        if p[-2:] == '_d':
+            p = p[:-2]
+        # if we get a value in the source frame, divide detector frame by (1 + z)
+        if p[-2:] == '_s':
+            p = p[:-2]
+            z = redshift_from_distance(self.inj_dat.getColumnByName('distance')[self.event])
+            return(self.get_xml_parameter(p) / (1.0 + z))
+        if p == 'Mc' or p == 'Mchirp':
+            return(self.inj_dat.getColumnByName('mchirp')[self.event])
+        elif p == 'Z' or p == 'redshift':
+            return(redshift_from_distance(self.inj_dat.getColumnByName('distance')[self.event]))
+        elif p == 'spin1' or p == 'spin2' or p == 's1' or p == 's2':
+            # dimensionless spin parameter a1 and a2 to be more precise 
+            # TODO not sure what units the xml files are in
+            # but I will assume they are in units that match up with conventions here
+            # https://arxiv.org/abs/1409.7215 (pg 6, a_i = |s_i| / m_i^2, otherwise missing a factor of c / G)
+            if p[-1] == '1':
+                p = 'spin1'
+            else:
+                p = 'spin2'
+            spin = np.zeros(3)
+            mass = self.get_xml_parameter('mass{0}'.format(p[-1])) #detector frame
+            for i, k in enumerate(['x', 'y', 'z']):
+                spin[i] = self.inj_dat.getColumnByName(p + k)[self.event]
+            return(np.linalg.norm(spin) / mass ** 2)              
+        elif (p == 'chi_eff'):
+            a1 = self.get_xml_parameter('spin1')
+            m1 = self.get_xml_parameter('mass1')
+            a2 = self.get_xml_parameter('spin2')
+            m2 = self.get_xml_parameter('mass2')
+            
+            # Note! We assume all spins are aligned with the orbital angular momentum 
+            theta1 = 0
+            theta2 = 0
+            return((m1 * a1 * np.cos(theta1) + m2 * a2 * np.cos(theta2)) / (m1 + m2))
+        elif p == 'Q':
+            return(self.inj_dat.getColumnByName('mass2')[self.event] / self.inj_dat.getColumnByName('mass1')[self.event])
+        elif p == 'mtot' or p == 'Mt' or p == 'Mtot':
+            return(self.inj_dat.getColumnByName('mass2')[self.event] + self.inj_dat.getColumnByName('mass1')[self.event])
+        elif p == 'coa_t' or p == 't':
+            return(self.inj_dat.getColumnByName('geocent_end_time')[self.event] + self.inj_dat.getColumnByName('geocent_end_time_ns')[self.event] * 1e-9) #- self.segment_start)
+        elif p == 't_segment':
+            return(self.inj_dat.getColumnByName('geocent_end_time')[self.event] + self.inj_dat.getColumnByName('geocent_end_time_ns')[self.event] * 1e-9 - self.segment_start)
+        elif p == 'psi':
+            return(self.inj_dat.getColumnByName('polarization')[self.event])
+        elif p == 'ra' or p == 'RA':
+            return(self.inj_dat.getColumnByName('longitude')[self.event])
+        elif p == 'sin_dec' or p == 'sindec':
+            return(np.sin(self.inj_dat.getColumnByName('latitude')[self.event]))
+        elif p == 'dec':
+            return(self.inj_dat.getColumnByName('latitude')[self.event])
+        elif p == 'cos_iota' or p == 'cosiota':
+            return(np.cos(self.inj_dat.getColumnByName('inclination')[self.event]))
+        elif p == 'iota':
+            return(self.inj_dat.getColumnByName('inclination')[self.event])
+        elif p == 'phase' or p == 'phi' :
+            return(self.inj_dat.getColumnByName('coa_phase')[self.event])
+        else:
+            return(self.inj_dat.getColumnByName(p)[self.event])
+
+
+# ----------------------------------
+# Function for initializing runFlags 
+# ----------------------------------
+def readbwb(runFlags, verbose = True):
+    # -- Initialize variables with info about the job
+    info = ""
+    jobName = ''
+    restrictModel = ''
+    lalsimFlag = False
+    found_model = False # Flag to see if we found what model we should be using 
+
+    # -- Open *bayeswave.run
+    bayeswaverunfile = runFlags.trigdir + '/bayeswave.run'
+    if verbose:
+        print("BWrun file is {0}".format(bayeswaverunfile))
+    runFlags.bayeswaverunfile = bayeswaverunfile
+    bayeswave = open(bayeswaverunfile, 'r')
+    # -- Parse command line arguments
+    raw_lines = bayeswave.readlines()
+    lines = [l.lstrip() for l in raw_lines]
+    cmdline = ''
+    for l in lines:
+        if 'Command' in l:
+            cmdline = l.rstrip().split(' ')
+
+    # Go through the different commands in bayeswave.run
+    for index, arg in enumerate(cmdline):
+        # Parameters describing where things live
+        if arg=='--runName':
+            jobName = cmdline[index+1]
+            if verbose:
+                print("The job name is: {0}".format(jobName))
+            runFlags.jobName = jobName+'_'
+
+        # Run settings
+        elif arg=='--trigtime':
+            gps =  float(cmdline[index+1])
+            runFlags.gps = gps
+            runFlags.trigtime = gps
+            info = info + "The trigger time is GPS: {0}\n".format(gps)
+        elif arg=='--srate':
+            runFlags.srate = float(cmdline[index + 1])
+        elif arg=='--segment-start':
+            runFlags.segment_start = float(cmdline[index + 1])
+        elif arg=='--seglen':
+            runFlags.seglen = float(cmdline[index + 1])
+        elif arg == '--Nchain':
+            runFlags.Nchain = int(cmdline[index + 1])
+        elif arg == '--Niter':
+            runFlags.Niter = int(cmdline[index + 1])
+        elif arg == '--Ncycle':
+            runFlags.Ncycle = int(cmdline[index + 1])
+        elif arg == '--chirplets':
+            runFlags.chirplets = True
+
+        # IFOs
+        elif arg=='--ifo':
+            runFlags.ifoNames.append(cmdline[index+1])
+
+        # Windows, cleaning, and bayesLine
+        elif arg == '--bayesLine':
+            runFlags.bayesLine = True
+        elif arg == '--window':
+            runFlags.window = float(cmdline[index + 1])
+        elif arg == '--CBCwindow':
+            runFlags.CBCwindow = float(cmdline[index + 1])
+        elif arg == '--noClean':
+            if verbose:
+                print('\nThis run had no cleaning phase, --noClean\n')
+            runFlags.noClean = True
+        # Parameters describing what run type was
+        elif arg=='--glitchOnly':
+            if verbose:
+                print('\nThis run was executed with the --glitchOnly flag\n')
+            runFlags.modelList = ['glitch']
+            runFlags.glitchOnly_flag = True
+            runFlags.noNoiseFlag = True
+            found_model = True
+        elif arg=='--signalOnly':
+            if verbose:
+                print('\nThis run was executed with the --signalOnly flag\n')
+            runFlags.modelList = ['signal']
+            found_model = True
+        elif arg=='--cbcOnly':
+            if verbose:
+                print('\nThis run was executed with the --cbcOnly flag\n')
+            runFlags.modelList = ['cbc']
+            runFlags.CBCType = True
+            runFlags.cbcOnly_flag = True
+            found_model = True
+        elif arg=='--noiseOnly':
+            if verbose:
+                print('\nThis run was executed with the --noiseOnly flag\n')
+            runFlags.modelList = ['noise']
+            found_model = True
+        elif arg=='--skipNoise':
+            if verbose:
+                print('\nThis run was executed with the --skipNoise flag\n')
+            runFlags.noNoiseFlag = True
+        # Combined types
+        elif arg=='--fullOnly':
+            if verbose:
+                print('\nThis run was executed with the --fullOnly flag\n')
+            runFlags.fullOnly_flag = True
+            runFlags.modelList = ['signal', 'glitch']
+            found_model = True
+        elif arg=='--GlitchCBC':
+            if verbose:
+                print('\nThis run was executed with the --GlitchCBC flag\n')
+            runFlags.GlitchCBC_flag = True
+            runFlags.CBCType = True
+            runFlags.modelList = ['glitch', 'cbc']
+            runFlags.multi_type = True
+            found_model = True
+        elif arg=='--SignalCBC':
+            if verbose:
+                print('\nThis run was executed with the --SignalCBC flag\n')
+            runFlags.SignalCBC_flag = True
+            runFlags.CBCType = True
+            runFlags.multi_type = True
+            runFlags.modelList = ['signal', 'cbc']
+            found_model = True
+        elif arg=='--cleanOnly':
+            if verbose:
+                print('\nThis run was executed with the --cleanOnly flag\n')
+            runFlags.cleanOnly_flag = True
+            runFlags.modelList = ['clean']
+            found_model = True
+        # Injection settings
+        elif arg=='--inj':
+            runFlags.injFlag = True
+            runFlags.mdc = True
+            runFlags.XMLfname = cmdline[index+1]
+            xmldoc = pipe_utils.ligolw_utils.load_filename(runFlags.trigdir + '/' + runFlags.XMLfname, contenthandler = pipe_utils.LIGOLWContentHandler, verbose = False)
+            runFlags.inj_dat = pipe_utils.lsctables.SimInspiralTable.get_table(xmldoc)
+            if verbose:
+                print('xml filename is {0}'.format(runFlags.trigdir + '/' + runFlags.XMLfname))
+        elif arg=='--MDC-cache':
+            runFlags.mdc = True
+        elif arg == '--MDC-channel':
+            runFlags.mdc = True
+        elif arg =='--event':
+            # if injection then this indexes the event in the .xml file
+            runFlags.event = int(cmdline[index + 1])
+        elif arg=='--MDC-cache':
+            mdc = True
+        elif arg=='--lalinspiralinjection':
+            lalsimFlag = True
+        elif arg=='--inj-fref':
+            runFlags.inj_fref = float(cmdline[index + 1])
+        # Likelihood calculation
+        elif arg=='--heterodyneL':
+            runFlags.heterodyneL = True
+
+        # Printed output settings
+        elif arg == '--checkpoint':
+            runFlags.checkpoint = True
+        elif arg=='--verbose':
+            if verbose:
+                print('\nThis run was executed with the --verbose flag\n')
+            runFlags.verbose = True
+
+    if lalsimFlag:
+        runFlags.injFlag = False
+
+    if not found_model:
+        runFlags.signalGlitch_flag = True
+        runFlags.modelList = ['signal', 'glitch']
+    if not runFlags.noClean and not runFlags.cleanOnly_flag:
+        runFlags.modelList.append('clean')
+
+    # -- Determine number of IFOs
+    ifoNum = len(runFlags.ifoNames)
+    # -- Convert the IFO name list into a list of numbers useful for opening datafiles
+    runFlags.ifoList = [str(ii) for ii in range(0,len(runFlags.ifoNames))]
+
+    # -- Getting flow
+    for index, arg in enumerate(cmdline):
+        for ifo in runFlags.ifoNames:
+            if arg == '--{0}-flow'.format(ifo):
+                #TODO THIS IS WRONG IF DIFFERENT IFOS HAVE DIFFERENT FLOW
+                runFlags.flow = float(cmdline[index+1])
+
+    info = info + "Detectors(s) found: {0}\n".format(len(runFlags.ifoList))
+    info = info + "The detectors are: {0}\n".format(', '.join(runFlags.ifoNames))
+    info = info + "BayesLine is: {0}\n".format(runFlags.bayesLine )
+    info = info + "Models are: {0}\n".format(', '.join(runFlags.modelList))
+    info = info + "\tflow = {0}\n".format(runFlags.flow)
+    info = info + "\tseglen = {0}\n".format(runFlags.seglen)
+    info = info + "\tNchains = {0}\n".format(runFlags.Nchain)
+    info = info + "\tCleaning phase = {0}\n".format(not runFlags.noClean)
+
+    if 'cbc' in runFlags.modelList:
+        info = info + "\theterodyneL = {0}\n".format(runFlags.heterodyneL)
+
+    # -- If MDC, read SNR info
+    if runFlags.mdc:
+        for ifo in runFlags.ifoList:
+            snrFile = 'post/injection_whitened_moments_%s.dat'%(runFlags.ifoNames[int(ifo)])
+            if not os.path.exists(snrFile):
+                # post hasn't been run yet most likely
+                #sys.exit("\n {0} not found! \n".format(snrFile))
+                break
+            snrdata = open(snrFile, 'r')
+            snrdata.readline()
+            snr = float(snrdata.readline().split(' ')[0])
+            runFlags.snrList.append(snr)
+            snrdata.close()
+            info = info + 'Injected SNR in detector {0} = {1}\n'.format(runFlags.ifoNames[int(ifo)],runFlags.snrList[-1])
+    bayeswave.close()
+    # -- Report to user
+    if verbose:
+        print("{0}".format(info))
+
+""" 
+    Functions for calculating redshift (used in Flags.get_xml_paramter)
+"""
+def redshift_from_distance(DL): # distance (MPC)
+    #double zlow, zhigh;
+    #double zmid;
+    #double Dlow, Dhigh, Dmid;
+    #double u, v, w;
+    #int i;
+
+    zlow = 1.0;
+    zhigh = 1000.0;
+
+    if(DL < 1.0):
+        zlow = 1.0e-10
+        zhigh = 3.0e-4
+    elif (DL < 100.0):
+        zlow = 2.0e-4
+        zhigh = 5.0e-2
+    elif (DL < 1000.0):
+        zlow = 1.0e-2
+        zhigh = 3.0e-1
+    elif (DL < 10000.0):
+        zlow = 1.0e-1
+        zhigh = 2.0
+
+    zmid = (zhigh + zlow) / 2.0;
+
+    Dlow = DL_fit(zlow);
+    Dhigh = DL_fit(zhigh);
+    Dmid = DL_fit(zmid);
+    zmid = (zhigh + zlow) / 2.0;
+
+    Dlow = DL_fit(zlow);
+    Dhigh = DL_fit(zhigh);
+    Dmid = DL_fit(zmid);
+
+    for i in range(30):
+
+        u = Dlow-DL;
+        v = Dmid-DL;
+
+        if(u*v < 0.0):
+            zhigh = zmid;
+            Dhigh = Dmid;
+        else:
+            zlow = zmid;
+            Dlow = Dmid;
+
+        zmid = (zhigh+zlow)/2.0;
+        Dmid = DL_fit(zmid);
+
+    return(zmid);
+        
+# fits distance to redshift assuming flat cosmological constant
+def DL_fit(z):
+    # Planck values https://arxiv.org/abs/1807.06209
+    Om = 0.315;
+    H0 = 67.4;
+    #double RH, x1, x2, Om1, Om2;
+
+    CLIGHT = 2.99792458e8 # m / s
+    MPC = 3.085677581e+22 # meters
+
+    # See http://arxiv.org/pdf/1111.6396v1.pdf
+
+    RH = CLIGHT / H0 * (1.0e-3 * MPC);
+
+    x1 = (1.0-Om)/(Om);
+
+    x2 = (1.0-Om)/(Om*pow((1.0+z),3.0));
+
+    Om1 = (1.0+1.32*x1+0.4415*x1*x1+0.02656*x1*x1*x1)/(1.0+1.392*x1+0.5121*x1*x1+0.03944*x1*x1*x1);
+
+    Om2 = (1.0+1.32*x2+0.4415*x2*x2+0.02656*x2*x2*x2)/(1.0+1.392*x2+0.5121*x2*x2+0.03944*x2*x2*x2);
+
+    D = 2.0*RH*(1.0+z) / np.sqrt(Om)*(Om1-Om2 / np.sqrt(1.0+z));
+
+    return(D / MPC);