lalinference_pipe_utils.py 108 KB
Newer Older
1

2
#flow DAG Class definitions for LALInference Pipeline
3
# (C) 2012 John Veitch, Vivien Raymond, Kiersten Ruisard, Kan Wang
4

John Douglas Veitch's avatar
John Douglas Veitch committed
5
import itertools
6
import glue
7
from glue import pipeline,segmentsUtils,segments
8
import os
9
import socket
10
from lalapps import inspiralutils
11 12
import uuid
import ast
13
import pdb
John Douglas Veitch's avatar
John Douglas Veitch committed
14
import string
15
from math import floor,ceil,log,pow
16
import sys
17
import random
18
from itertools import permutations
19
import shutil
20
import numpy as np
21 22 23 24 25

# We use the GLUE pipeline utilities to construct classes for each
# type of job. Each class has inputs and outputs, which are used to
# join together types of jobs into a DAG.

John Douglas Veitch's avatar
John Douglas Veitch committed
26 27 28 29
class Event():
  """
  Represents a unique event to run on
  """
John Douglas Veitch's avatar
John Douglas Veitch committed
30
  new_id=itertools.count().next
31
  def __init__(self,trig_time=None,SimInspiral=None,SimBurst=None,SnglInspiral=None,CoincInspiral=None,event_id=None,timeslide_dict=None,GID=None,ifos=None, duration=None,srate=None,trigSNR=None,fhigh=None,horizon_distance=None):
John Douglas Veitch's avatar
John Douglas Veitch committed
32
    self.trig_time=trig_time
John Douglas Veitch's avatar
John Douglas Veitch committed
33
    self.injection=SimInspiral
34
    self.burstinjection=SimBurst
John Douglas Veitch's avatar
John Douglas Veitch committed
35
    self.sngltrigger=SnglInspiral
36 37 38 39
    if timeslide_dict is None:
      self.timeslides={}
    else:
      self.timeslides=timeslide_dict
40 41
    self.GID=GID
    self.coinctrigger=CoincInspiral
42 43 44 45
    if ifos is None:
      self.ifos = []
    else:
      self.ifos = ifos
46 47
    self.duration = duration
    self.srate = srate
48
    self.trigSNR = trigSNR
49
    self.fhigh = fhigh
50
    self.horizon_distance = horizon_distance
John Douglas Veitch's avatar
John Douglas Veitch committed
51 52 53
    if event_id is not None:
        self.event_id=event_id
    else:
John Douglas Veitch's avatar
John Douglas Veitch committed
54 55 56
        self.event_id=Event.new_id()
    if self.injection is not None:
        self.trig_time=self.injection.get_end()
57
        if event_id is None: self.event_id=int(str(self.injection.simulation_id).split(':')[2])
58 59 60
    if self.burstinjection is not None:
        self.trig_time=self.burstinjection.get_end()
        if event_id is None: self.event_id=int(str(self.burstinjection.simulation_id).split(':')[2])
John Douglas Veitch's avatar
John Douglas Veitch committed
61
    if self.sngltrigger is not None:
62
        self.trig_time=self.sngltrigger.get_end()
John Douglas Veitch's avatar
John Douglas Veitch committed
63
        self.event_id=int(str(self.sngltrigger.event_id).split(':')[2])
64 65 66 67
    if self.coinctrigger is not None:
        self.trig_time=self.coinctrigger.end_time + 1.0e-9 * self.coinctrigger.end_time_ns
    if self.GID is not None:
        self.event_id=int(''.join(i for i in self.GID if i.isdigit()))
68 69 70 71 72 73 74
    self.engine_opts={}
  def set_engine_option(self,opt,val):
    """
    Can set event-specific options for the engine nodes
    using this option, e.g. ev.set_engine_option('time-min','1083759273')
    """
    self.engine_opts[opt]=val
John Douglas Veitch's avatar
John Douglas Veitch committed
75 76

dummyCacheNames=['LALLIGO','LALVirgo','LALAdLIGO','LALAdVirgo']
John Douglas Veitch's avatar
John Douglas Veitch committed
77

78
def readLValert(threshold_snr=None,gid=None,flow=40.0,gracedb="gracedb",basepath="./",downloadpsd=True):
79
  """
80
  Parse LV alert file, containing coinc, sngl, coinc_event_map.
81 82
  and create a list of Events as input for pipeline
  Based on Chris Pankow's script
83
  """
84
  output=[]
85 86
  from glue.ligolw import utils
  from glue.ligolw import lsctables
87
  from glue.ligolw import ligolw
Rory Smith's avatar
Rory Smith committed
88
  class PSDContentHandler(ligolw.LIGOLWContentHandler):
89
    pass
Rory Smith's avatar
Rory Smith committed
90
  lsctables.use_in(PSDContentHandler)
91 92
  from glue.ligolw import param
  from glue.ligolw import array
93
  import subprocess
Rory Smith's avatar
Rory Smith committed
94
  from lal import series as lalseries
95
  from subprocess import Popen, PIPE
96 97
  cwd=os.getcwd()
  os.chdir(basepath)
98
  print "%s download %s coinc.xml"%(gracedb,gid)
Vivien Raymond's avatar
Vivien Raymond committed
99
  subprocess.call([gracedb,"download", gid ,"coinc.xml"])
Rory Smith's avatar
Rory Smith committed
100
  xmldoc=utils.load_filename("coinc.xml",contenthandler = PSDContentHandler)
101
  coinctable = lsctables.CoincInspiralTable.get_table(xmldoc)
102
  coinc_events = [event for event in coinctable]
103
  sngltable = lsctables.SnglInspiralTable.get_table(xmldoc)
104
  sngl_events = [event for event in sngltable]
105 106
  ifos = coinctable[0].ifos.split(",")
  trigSNR = coinctable[0].snr
107
  # Parse PSD
108
  srate_psdfile=16384
109
  fhigh=None
110 111 112
  if downloadpsd:
    print "gracedb download %s psd.xml.gz" % gid
    subprocess.call([gracedb,"download", gid ,"psd.xml.gz"])
Rory Smith's avatar
Rory Smith committed
113
    xmlpsd = lalseries.read_psd_xmldoc(utils.load_filename('psd.xml.gz',contenthandler = lalseries.PSDContentHandler))
114 115 116 117 118 119
    if os.path.exists("psd.xml.gz"):
      psdasciidic=get_xml_psds(os.path.realpath("./psd.xml.gz"),ifos,os.path.realpath('./PSDs'),end_time=None)
      combine=np.loadtxt(psdasciidic[psdasciidic.keys()[0]])
      srate_psdfile = pow(2.0, ceil( log(float(combine[-1][0]), 2) ) ) * 2
    else:
      print "Failed to gracedb download %s psd.xml.gz. lalinference will estimate the psd itself." % gid
120
  coinc_map = lsctables.CoincMapTable.get_table(xmldoc)
121 122
  for coinc in coinc_events:
    these_sngls = [e for e in sngl_events if e.event_id in [c.event_id for c in coinc_map if c.coinc_event_id == coinc.coinc_event_id] ]
123 124
    dur=[]
    srate=[]
125
    horizon_distance=[]
126
    for e in these_sngls:
127
      # Review: Replace this with a call to LALSimulation function at some point
128 129 130
      p=Popen(["lalapps_chirplen","--flow",str(flow),"-m1",str(e.mass1),"-m2",str(e.mass2)],stdout=PIPE, stderr=PIPE, stdin=PIPE)
      strlen = p.stdout.read()
      dur.append(pow(2.0, ceil( log(max(8.0,float(strlen.splitlines()[2].split()[5]) + 2.0), 2) ) ) )
131
      srate.append(pow(2.0, ceil( log(float(strlen.splitlines()[1].split()[5]), 2) ) ) * 2 )
132 133 134 135 136 137 138
      snr = e.snr
      eff_dist = e.eff_distance
      if threshold_snr is not None:
          if snr > threshold_snr:
              horizon_distance.append(eff_dist * snr/threshold_snr)
          else:
              horizon_distance.append(2 * eff_dist)
139 140 141 142
    if max(srate)<srate_psdfile:
      srate = max(srate)
    else:
      srate = srate_psdfile
143 144
      if downloadpsd:
        fhigh = srate_psdfile/2.0 * 0.95 # Because of the drop-off near Nyquist of the PSD from gstlal
145 146 147 148
    horizon_distance = max(horizon_distance) if len(horizon_distance) > 0 else None
    ev=Event(CoincInspiral=coinc, GID=gid, ifos = ifos, duration = max(dur), srate = srate,
             trigSNR = trigSNR, fhigh = fhigh, horizon_distance=horizon_distance)
    output.append(ev)
149

150
  print "Found %d coinc events in table." % len(coinc_events)
151
  os.chdir(cwd)
152 153
  return output

154 155 156 157 158 159 160
def open_pipedown_database(database_filename,tmp_space):
    """
    Open the connection to the pipedown database
    """
    if not os.access(database_filename,os.R_OK):
	raise Exception('Unable to open input file: %s'%(database_filename))
    from glue.ligolw import dbtables
161
    import sqlite3
162 163 164 165
    working_filename=dbtables.get_connection_filename(database_filename,tmp_path=tmp_space)
    connection = sqlite3.connect(working_filename)
    if tmp_space:
	dbtables.set_temp_store_directory(connection,tmp_space)
166
    #dbtables.DBTable_set_connection(connection)
167
    return (connection,working_filename)
168

169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211
def get_zerolag_lloid(database_connection, dumpfile=None, gpsstart=None, gpsend=None, max_cfar=-1, min_cfar=-1):
	"""
	Returns a list of Event objects
	from pipedown data base. Can dump some stats to dumpfile if given,
	and filter by gpsstart and gpsend to reduce the nunmber or specify
	max_cfar to select by combined FAR
	"""
	output={}
	if gpsstart is not None: gpsstart=float(gpsstart)
	if gpsend is not None: gpsend=float(gpsend)
	# Get coincs
	get_coincs = "SELECT sngl_inspiral.end_time+sngl_inspiral.end_time_ns*1e-9,sngl_inspiral.ifo,coinc_event.coinc_event_id,sngl_inspiral.snr,sngl_inspiral.chisq,coinc_inspiral.combined_far \
		FROM sngl_inspiral join coinc_event_map on (coinc_event_map.table_name=='sngl_inspiral' and coinc_event_map.event_id ==\
		sngl_inspiral.event_id) join coinc_event on (coinc_event.coinc_event_id==coinc_event_map.coinc_event_id) \
		join coinc_inspiral on (coinc_event.coinc_event_id==coinc_inspiral.coinc_event_id) \
        WHERE coinc_event.time_slide_id=='time_slide:time_slide_id:1'\
		"
	if gpsstart is not None:
		get_coincs=get_coincs+' and coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1.0e-9 > %f'%(gpsstart)
	if gpsend is not None:
		get_coincs=get_coincs+' and coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1.0e-9 < %f'%(gpsend)
	if max_cfar !=-1:
		get_coincs=get_coincs+' and coinc_inspiral.combined_far < %f'%(max_cfar)
	if min_cfar != -1:
		get_coincs=get_coincs+' and coinc_inspiral.combined_far > %f'%(min_cfar)
	db_out=database_connection.cursor().execute(get_coincs)
    	extra={}
	for (sngl_time, ifo, coinc_id, snr, chisq, cfar) in db_out:
      		coinc_id=int(coinc_id.split(":")[-1])
	  	if not coinc_id in output.keys():
			output[coinc_id]=Event(trig_time=sngl_time,timeslide_dict={},event_id=int(coinc_id))
			extra[coinc_id]={}
		output[coinc_id].timeslides[ifo]=0
		output[coinc_id].ifos.append(ifo)
		extra[coinc_id][ifo]={'snr':snr,'chisq':chisq,'cfar':cfar}
	if dumpfile is not None:
		fh=open(dumpfile,'w')
		for co in output.keys():
			for ifo in output[co].ifos:
				fh.write('%s %s %s %s %s %s %s\n'%(str(co),ifo,str(output[co].trig_time),str(output[co].timeslides[ifo]),str(extra[co][ifo]['snr']),str(extra[co][ifo]['chisq']),str(extra[co][ifo]['cfar'])))
		fh.close()
	return output.values()

212
def get_zerolag_pipedown(database_connection, dumpfile=None, gpsstart=None, gpsend=None, max_cfar=-1, min_cfar=-1):
213 214 215 216 217 218 219 220 221 222 223 224 225 226
	"""
	Returns a list of Event objects
	from pipedown data base. Can dump some stats to dumpfile if given,
	and filter by gpsstart and gpsend to reduce the nunmber or specify
	max_cfar to select by combined FAR
	"""
	output={}
	if gpsstart is not None: gpsstart=float(gpsstart)
	if gpsend is not None: gpsend=float(gpsend)
	# Get coincs
	get_coincs = "SELECT sngl_inspiral.end_time+sngl_inspiral.end_time_ns*1e-9,sngl_inspiral.ifo,coinc_event.coinc_event_id,sngl_inspiral.snr,sngl_inspiral.chisq,coinc_inspiral.combined_far \
		FROM sngl_inspiral join coinc_event_map on (coinc_event_map.table_name=='sngl_inspiral' and coinc_event_map.event_id ==\
		sngl_inspiral.event_id) join coinc_event on (coinc_event.coinc_event_id==coinc_event_map.coinc_event_id) \
		join coinc_inspiral on (coinc_event.coinc_event_id==coinc_inspiral.coinc_event_id) \
227
		WHERE coinc_event.time_slide_id=='time_slide:time_slide_id:10049'\
228 229
		"
	if gpsstart is not None:
230
		get_coincs=get_coincs+' and coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1.0e-9 > %f'%(gpsstart)
231
	if gpsend is not None:
232
		get_coincs=get_coincs+' and coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1.0e-9 < %f'%(gpsend)
233 234
	if max_cfar !=-1:
		get_coincs=get_coincs+' and coinc_inspiral.combined_far < %f'%(max_cfar)
235 236
	if min_cfar != -1:
		get_coincs=get_coincs+' and coinc_inspiral.combined_far > %f'%(min_cfar)
237
	db_out=database_connection.cursor().execute(get_coincs)
John Douglas Veitch's avatar
John Douglas Veitch committed
238
    	extra={}
239
	for (sngl_time, ifo, coinc_id, snr, chisq, cfar) in db_out:
John Douglas Veitch's avatar
John Douglas Veitch committed
240 241
      		coinc_id=int(coinc_id.split(":")[-1])
	  	if not coinc_id in output.keys():
242
			output[coinc_id]=Event(trig_time=sngl_time,timeslide_dict={},event_id=int(coinc_id))
John Douglas Veitch's avatar
John Douglas Veitch committed
243 244 245 246
			extra[coinc_id]={}
		output[coinc_id].timeslides[ifo]=0
		output[coinc_id].ifos.append(ifo)
		extra[coinc_id][ifo]={'snr':snr,'chisq':chisq,'cfar':cfar}
247 248 249 250 251 252
	if dumpfile is not None:
		fh=open(dumpfile,'w')
		for co in output.keys():
			for ifo in output[co].ifos:
				fh.write('%s %s %s %s %s %s %s\n'%(str(co),ifo,str(output[co].trig_time),str(output[co].timeslides[ifo]),str(extra[co][ifo]['snr']),str(extra[co][ifo]['chisq']),str(extra[co][ifo]['cfar'])))
		fh.close()
253
	return output.values()
254

255
def get_timeslides_pipedown(database_connection, dumpfile=None, gpsstart=None, gpsend=None, max_cfar=-1):
256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271
	"""
	Returns a list of Event objects
	with times and timeslide offsets
	"""
	output={}
	if gpsstart is not None: gpsstart=float(gpsstart)
	if gpsend is not None: gpsend=float(gpsend)
	db_segments=[]
	sql_seg_query="SELECT search_summary.out_start_time, search_summary.out_end_time from search_summary join process on process.process_id==search_summary.process_id where process.program=='thinca'"
	db_out = database_connection.cursor().execute(sql_seg_query)
	for d in db_out:
		if d not in db_segments:
			db_segments.append(d)
	seglist=segments.segmentlist([segments.segment(d[0],d[1]) for d in db_segments])
	db_out_saved=[]
	# Get coincidences
272
	get_coincs="SELECT sngl_inspiral.end_time+sngl_inspiral.end_time_ns*1e-9,time_slide.offset,sngl_inspiral.ifo,coinc_event.coinc_event_id,sngl_inspiral.snr,sngl_inspiral.chisq,coinc_inspiral.combined_far \
273
		    FROM sngl_inspiral join coinc_event_map on (coinc_event_map.table_name == 'sngl_inspiral' and coinc_event_map.event_id \
274 275
		    == sngl_inspiral.event_id) join coinc_event on (coinc_event.coinc_event_id==coinc_event_map.coinc_event_id) join time_slide\
		    on (time_slide.time_slide_id == coinc_event.time_slide_id and time_slide.instrument==sngl_inspiral.ifo)\
276
		    join coinc_inspiral on (coinc_inspiral.coinc_event_id==coinc_event.coinc_event_id) where coinc_event.time_slide_id!='time_slide:time_slide_id:10049'"
John Douglas Veitch's avatar
John Douglas Veitch committed
277
	joinstr = ' and '
278
	if gpsstart is not None:
279
		get_coincs=get_coincs+ joinstr + ' coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1e-9 > %f'%(gpsstart)
280
	if gpsend is not None:
281
		get_coincs=get_coincs+ joinstr+' coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1e-9 <%f'%(gpsend)
282 283
	if max_cfar!=-1:
		get_coincs=get_coincs+joinstr+' coinc_inspiral.combined_far < %f'%(max_cfar)
284
	db_out=database_connection.cursor().execute(get_coincs)
285 286
	from pylal import SnglInspiralUtils
	extra={}
287
	for (sngl_time, slide, ifo, coinc_id, snr, chisq, cfar) in db_out:
288 289 290 291 292 293 294 295 296
		coinc_id=int(coinc_id.split(":")[-1])
		seg=filter(lambda seg:sngl_time in seg,seglist)[0]
		slid_time = SnglInspiralUtils.slideTimeOnRing(sngl_time,slide,seg)
		if not coinc_id in output.keys():
			output[coinc_id]=Event(trig_time=slid_time,timeslide_dict={},event_id=int(coinc_id))
			extra[coinc_id]={}
		output[coinc_id].timeslides[ifo]=slid_time-sngl_time
		output[coinc_id].ifos.append(ifo)
		extra[coinc_id][ifo]={'snr':snr,'chisq':chisq,'cfar':cfar}
297 298 299 300 301 302
	if dumpfile is not None:
		fh=open(dumpfile,'w')
		for co in output.keys():
			for ifo in output[co].ifos:
				fh.write('%s %s %s %s %s %s %s\n'%(str(co),ifo,str(output[co].trig_time),str(output[co].timeslides[ifo]),str(extra[co][ifo]['snr']),str(extra[co][ifo]['chisq']),str(extra[co][ifo]['cfar'])))
		fh.close()
303
	return output.values()
304

305 306 307 308 309 310 311 312
def mkdirs(path):
  """
  Helper function. Make the given directory, creating intermediate
  dirs if necessary, and don't complain about it already existing.
  """
  if os.access(path,os.W_OK) and os.path.isdir(path): return
  else: os.makedirs(path)

John Douglas Veitch's avatar
John Douglas Veitch committed
313 314 315
def chooseEngineNode(name):
  if name=='lalinferencenest':
    return LALInferenceNestNode
316 317
  if name=='lalinferenceburst':
    return LALInferenceBurstNode
John Douglas Veitch's avatar
John Douglas Veitch committed
318 319
  if name=='lalinferencemcmc':
    return LALInferenceMCMCNode
320
  if name=='lalinferencebambi' or name=='lalinferencebambimpi':
321
    return LALInferenceBAMBINode
322 323
  if name=='lalinferencedatadump':
    return LALInferenceDataDumpNode
324 325
  if name=='bayeswavepsd':
    return BayesWavePSDNode
John Douglas Veitch's avatar
John Douglas Veitch committed
326 327
  return EngineNode

328 329 330
def get_engine_name(cp):
    name=cp.get('analysis','engine')
    if name=='random':
331
        engine_list=['lalinferencenest','lalinferencemcmc','lalinferencebambimpi']
332 333 334 335 336 337 338 339 340
        if cp.has_option('input','gid'):
            gid=cp.get('input','gid')
            engine_number=int(''.join(i for i in gid if i.isdigit())) % 2
        else:
            engine_number=random.randint(0,1)
        return engine_list[engine_number]
    else:
        return name

341
def scan_timefile(timefile):
342 343 344 345 346 347 348 349 350 351 352 353
    import re
    p=re.compile('[\d.]+')
    times=[]
    timefilehandle=open(timefile,'r')
    for time in timefilehandle:
      if not p.match(time):
	continue
      if float(time) in times:
	print 'Skipping duplicate time %s'%(time)
	continue
      print 'Read time %s'%(time)
      times.append(float(time))
354 355
    timefilehandle.close()
    return times
356

357 358 359 360 361 362 363 364 365 366 367 368 369 370 371
def get_xml_psds(psdxml,ifos,outpath,end_time=None):
  """
  Get a psd.xml.gz file and:
  1) Reads it
  2) Converts PSD (10e-44) -> ASD ( ~10e-22)
  3) Checks the psd file contains all the IFO we want to analyze
  4) Writes down the ASDs into an ascii file for each IFO in psd.xml.gz. The name of the file contains the trigtime (if given) and the ifo name.
  Input:
    psdxml: psd.xml.gz file
    ifos: list of ifos used for the analysis
    outpath: path where the ascii ASD will be written to
    (end_time): trigtime for this event. Will be used a part of the ASD file name
  """
  lal=1
  from glue.ligolw import utils
372
  try:
Rory Smith's avatar
Rory Smith committed
373 374
    #from pylal import series
    from lal import series as series
375
    lal=0
376 377 378
  except ImportError:
    print "ERROR, cannot import pylal.series in bppu/get_xml_psds()\n"
    exit(1)
379

380 381 382 383 384 385 386 387 388 389 390
  out={}
  if not os.path.isdir(outpath):
    os.makedirs(outpath)
  if end_time is not None:
    time=repr(float(end_time))
  else:
    time=''
  #check we don't already have ALL the psd files #
  got_all=1
  for ifo in ifos:
    path_to_ascii_psd=os.path.join(outpath,ifo+'_psd_'+time+'.txt')
391
    # Check we don't already have that ascii (e.g. because we are running parallel runs of the save event
392 393 394 395 396 397 398 399 400
    if os.path.isfile(path_to_ascii_psd):
      got_all*=1
    else:
      got_all*=0
  if got_all==1:
    #print "Already have PSD files. Nothing to do...\n"
    for ifo in ifos:
      out[ifo]=os.path.join(outpath,ifo+'_psd_'+time+'.txt')
    return out
401

402 403 404 405
  # We need to convert the PSD for one or more IFOS. Open the file
  if not os.path.isfile(psdxml):
    print "ERROR: impossible to open the psd file %s. Exiting...\n"%psdxml
    sys.exit(1)
Rory Smith's avatar
Rory Smith committed
406
  xmlpsd =  series.read_psd_xmldoc(utils.load_filename(psdxml,contenthandler = series.PSDContentHandler))
407 408 409 410 411 412 413 414 415
  # Check the psd file contains all the IFOs we want to analize
  for ifo in ifos:
    if not ifo in [i.encode('ascii') for i in xmlpsd.keys()]:
      print "ERROR. The PSD for the ifo %s does not seem to be contained in %s\n"%(ifo,psdxml)
      sys.exit(1)
  #loop over ifos in psd xml file
  for instrument in xmlpsd.keys():
    #name of the ascii file we are going to write the PSD into
    path_to_ascii_psd=os.path.join(outpath,instrument.encode('ascii')+'_psd_'+time+'.txt')
416
    # Check we don't already have that ascii (e.g. because we are running parallel runs of the save event
417
    if os.path.isfile(path_to_ascii_psd):
418
      continue
419 420 421 422 423 424 425 426 427 428 429 430 431 432 433
    # get data for the IFO
    ifodata=xmlpsd[instrument]
    #check data is not empty
    if ifodata is None:
      continue
    # we have data. Get psd array
    if lal==0:
      #pylal stores the series in ifodata.data
      data=ifodata
    else:
      # lal stores it in ifodata.data.data
      data=ifodata.data
    # Fill a two columns array of (freq, psd) and save it in the ascii file
    f0=ifodata.f0
    deltaF=ifodata.deltaF
434

435
    combine=[]
436

Rory Smith's avatar
Rory Smith committed
437 438
    for i in np.arange(len(data.data.data)) :
      combine.append([f0+i*deltaF,np.sqrt(data.data.data[i])])
439
    np.savetxt(path_to_ascii_psd,combine)
440
    ifo=instrument.encode('ascii')
441 442 443
    # set node.psds dictionary with the path to the ascii files
    out[ifo]=os.path.join(outpath,ifo+'_psd_'+time+'.txt')
  return out
444

445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468
def get_trigger_chirpmass(gid=None,gracedb="gracedb"):
  from glue.ligolw import lsctables
  from glue.ligolw import ligolw
  from glue.ligolw import utils
  class PSDContentHandler(ligolw.LIGOLWContentHandler):
    pass
  lsctables.use_in(PSDContentHandler)
  import subprocess

  cwd=os.getcwd()
  subprocess.call([gracedb,"download", gid ,"coinc.xml"])
  xmldoc=utils.load_filename("coinc.xml",contenthandler = PSDContentHandler)
  coinctable = lsctables.CoincInspiralTable.get_table(xmldoc)
  coinc_events = [event for event in coinctable]
  sngltable = lsctables.SnglInspiralTable.get_table(xmldoc)
  sngl_events = [event for event in sngltable]
  coinc_map = lsctables.CoincMapTable.get_table(xmldoc)
  mass1 = []
  mass2 = []
  for coinc in coinc_events:
    these_sngls = [e for e in sngl_events if e.event_id in [c.event_id for c in coinc_map if c.coinc_event_id == coinc.coinc_event_id] ]
    for e in these_sngls:
      mass1.append(e.mass1)
      mass2.append(e.mass2)
469
  # check that trigger masses are identical in each IFO
470 471 472
  assert len(set(mass1)) == 1
  assert len(set(mass2)) == 1

473
  mchirp = (mass1[0]*mass2[0])**(3./5.) / ( (mass1[0] + mass2[0])**(1./5.) )
474 475 476 477
  os.remove("coinc.xml")

  return mchirp

478
def get_roq_mchirp_priors(path, roq_paths, roq_params, key, gid=None,sim_inspiral=None):
479

480 481
  ## XML and GID cannot be given at the same time
  ## sim_inspiral must already point at the right row
482
  mc_priors = {}
Vivien Raymond's avatar
Vivien Raymond committed
483

484 485 486
  if gid is not None and sim_inspiral is not None:
    print "Error in get_roq_mchirp_priors, cannot use both gid and sim_inspiral\n"
    sys.exit(1)
487 488 489 490 491 492

  for roq in roq_paths:
    params=os.path.join(path,roq,'params.dat')
    roq_params[roq]=np.genfromtxt(params,names=True)
    mc_priors[roq]=[float(roq_params[roq]['chirpmassmin']),float(roq_params[roq]['chirpmassmax'])]
  ordered_roq_paths=[item[0] for item in sorted(roq_params.items(), key=key)][::-1]
493 494
  # below is to construct non-overlapping mc priors for multiple roq mass-bin runs
  '''i=0
495 496 497 498 499 500 501
  for roq in ordered_roq_paths:
    if i>0:
      # change min, just set to the max of the previous one since we have already aligned it in the previous iteration of this loop
      #mc_priors[roq][0]+= (mc_priors[roq_lengths[i-1]][1]-mc_priors[roq][0])/2.
      mc_priors[roq][0]=mc_priors[ordered_roq_paths[i-1]][1]
    if i<len(roq_paths)-1:
      mc_priors[roq][1]-= (mc_priors[roq][1]- mc_priors[ordered_roq_paths[i+1]][0])/2.
502
    i+=1'''
503 504
  if gid is not None:
  	trigger_mchirp = get_trigger_chirpmass(gid)
505 506
  elif sim_inspiral is not None:
        trigger_mchirp = sim_inspiral.mchirp
507 508
  else:
	trigger_mchirp = None
509 510 511

  return mc_priors, trigger_mchirp

Vivien Raymond's avatar
Vivien Raymond committed
512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537
def get_roq_component_mass_priors(path, roq_paths, roq_params, key, gid=None,sim_inspiral=None):

  ## XML and GID cannot be given at the same time
  ## sim_inspiral must already point at the right row
  m1_priors = {}
  m2_priors = {}

  if gid is not None and sim_inspiral is not None:
    print "Error in get_roq_mchirp_priors, cannot use both gid and sim_inspiral\n"
    sys.exit(1)

  for roq in roq_paths:
    params=os.path.join(path,roq,'params.dat')
    roq_params[roq]=np.genfromtxt(params,names=True)
    m1_priors[roq]=[float(roq_params[roq]['mass1min']),float(roq_params[roq]['mass1max'])]
    m2_priors[roq]=[float(roq_params[roq]['mass2min']),float(roq_params[roq]['mass2max'])]

  if gid is not None:
    trigger_mchirp = get_trigger_chirpmass(gid)
  elif sim_inspiral is not None:
    trigger_mchirp = sim_inspiral.mchirp
  else:
    trigger_mchirp = None

  return m1_priors, m2_priors, trigger_mchirp

538
def get_roq_mass_freq_scale_factor(mc_priors, trigger_mchirp, force_flow=None):
539 540 541 542 543
  mc_priors_keys_int = map(lambda k : int(k[:-1]), mc_priors.keys())
  roq_min = mc_priors.keys()[np.argmin(mc_priors_keys_int)]
  roq_max = mc_priors.keys()[np.argmax(mc_priors_keys_int)]
  mc_max = mc_priors[roq_min][1]
  mc_min = mc_priors[roq_max][0]
544
  scale_factor = 1.
545
  if force_flow == None and trigger_mchirp != None:
546 547 548 549
     if trigger_mchirp >= mc_max:
        scale_factor = 2.**(floor(trigger_mchirp/mc_max))
     if trigger_mchirp <= mc_min:
        scale_factor = (2./3.)**(ceil(trigger_mchirp/mc_min))
550
  elif force_flow != None:
551
     scale_factor = 20./force_flow
552
  return scale_factor
553

554 555
def create_pfn_tuple(filename,protocol='file://',site='local'):
    return( (os.path.basename(filename),protocol+os.path.abspath(filename),site) )
556

Vivien Raymond's avatar
Vivien Raymond committed
557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577
def mchirp_from_components(m1, m2):
  return (m1*m2)**(3.0/5.0) / (m1+m2)**(1.0/5.0)

def Query_ROQ_Bounds_Type(path, roq_paths):
  # Assume that parametrization of ROQ bounds is independent of seglen; just look at first one
  import numpy as np
  roq = roq_paths[0]
  params = os.path.join(path,roq,'params.dat')
  roq_params0 = np.genfromtxt(params,names=True)
  roq_names_set = set(roq_params0.dtype.names)
  component_mass_bounds_set = set(['mass1min', 'mass1max', 'mass2min', 'mass2max'])
  chirp_mass_q_bounds_set = set(['chirpmassmin', 'chirpmassmax', 'qmin', 'qmax'])
  if roq_names_set.issuperset(component_mass_bounds_set):
    roq_bounds = 'component_mass'
  elif roq_names_set.issuperset(chirp_mass_q_bounds_set):
    roq_bounds = 'chirp_mass_q'
  else:
    print 'Invalid bounds for ROQ. Ether (m1,m2) or (mc,q) bounds are supported.'
    sys.exit(1)
  return roq_bounds

578
class LALInferencePipelineDAG(pipeline.CondorDAG):
579
  def __init__(self,cp,dax=False,first_dag=True,previous_dag=None,site='local'):
580 581
    self.subfiles=[]
    self.config=cp
582
    self.engine=get_engine_name(cp)
John Douglas Veitch's avatar
John Douglas Veitch committed
583
    self.EngineNode=chooseEngineNode(self.engine)
584
    self.site=site
585
    if cp.has_option('paths','basedir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
586
      self.basepath=cp.get('paths','basedir')
587 588 589
    else:
      self.basepath=os.getcwd()
      print 'No basepath specified, using current directory: %s'%(self.basepath)
590
    mkdirs(self.basepath)
591 592
    if dax:
        os.chdir(self.basepath)
593
    self.posteriorpath=os.path.join(self.basepath,'posterior_samples')
John Douglas Veitch's avatar
John Douglas Veitch committed
594
    mkdirs(self.posteriorpath)
595 596 597 598 599 600 601 602 603 604 605 606
    if first_dag:
      daglogdir=cp.get('paths','daglogdir')
      mkdirs(daglogdir)
      self.daglogfile=os.path.join(daglogdir,'lalinference_pipeline-'+str(uuid.uuid1())+'.log')
      pipeline.CondorDAG.__init__(self,self.daglogfile,dax=dax)
    elif not first_dag and previous_dag is not None:
      daglogdir=cp.get('paths','daglogdir')
      mkdirs(daglogdir)
      self.daglogfile=os.path.join(daglogdir,'lalinference_pipeline-'+str(uuid.uuid1())+'.log')
      pipeline.CondorDAG.__init__(self,self.daglogfile,dax=dax)
      for node in previous_dag.get_nodes():
        self.add_node(node)
607
    if cp.has_option('paths','cachedir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
608
      self.cachepath=cp.get('paths','cachedir')
609 610
    else:
      self.cachepath=os.path.join(self.basepath,'caches')
611
    mkdirs(self.cachepath)
612
    if cp.has_option('paths','logdir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
613
      self.logpath=cp.get('paths','logdir')
614 615
    else:
      self.logpath=os.path.join(self.basepath,'log')
616
    mkdirs(self.logpath)
617
    if cp.has_option('analysis','ifos'):
618
      self.ifos=ast.literal_eval(cp.get('analysis','ifos'))
619 620 621
    else:
      self.ifos=['H1','L1','V1']
    self.segments={}
622 623 624 625
    if cp.has_option('datafind','veto-categories'):
      self.veto_categories=cp.get('datafind','veto-categories')
    else: self.veto_categories=[]
    for ifo in self.ifos:
626
      self.segments[ifo]=[]
627 628 629
    self.computeroqweightsnode={}
    self.bayeslinenode={}
    self.bayeswavepsdnode={}
630
    self.dq={}
631
    self.frtypes=ast.literal_eval(cp.get('datafind','types'))
632
    self.channels=ast.literal_eval(cp.get('data','channels'))
John Douglas Veitch's avatar
John Douglas Veitch committed
633
    self.use_available_data=False
634
    self.webdir=cp.get('paths','webdir')
635 636 637 638
    if cp.has_option('analysis','dataseed'):
      self.dataseed=cp.getint('analysis','dataseed')
    else:
      self.dataseed=None
639
    # Set up necessary job files.
640
    self.prenodes={}
641
    self.datafind_job = pipeline.LSCDataFindJob(self.cachepath,self.logpath,self.config,dax=self.is_dax())
642
    self.datafind_job.add_opt('url-type','file')
643 644
    if cp.has_option('analysis','accounting_group'):
      self.datafind_job.add_condor_cmd('accounting_group',cp.get('analysis','accounting_group'))
645 646
    if cp.has_option('analysis','accounting_group_user'):
      self.datafind_job.add_condor_cmd('accounting_group_user',cp.get('analysis','accounting_group_user'))
647
    self.datafind_job.set_sub_file(os.path.abspath(os.path.join(self.basepath,'datafind.sub')))
648
    self.preengine_job = EngineJob(self.config, os.path.join(self.basepath,'prelalinference.sub'),self.logpath,engine='lalinferencedatadump',ispreengine=True,dax=self.is_dax())
649 650
    self.preengine_job.set_grid_site('local')
    self.preengine_job.set_universe('vanilla')
651 652 653 654 655 656
    if self.config.has_option('condor','computeroqweights'):
      self.computeroqweights_job = ROMJob(self.config,os.path.join(self.basepath,'computeroqweights.sub'),self.logpath,dax=self.is_dax())
      self.computeroqweights_job.set_grid_site('local')
    if self.config.has_option('condor','bayesline'):
      self.bayesline_job = BayesLineJob(self.config,os.path.join(self.basepath,'bayesline.sub'),self.logpath,dax=self.is_dax())
      self.bayesline_job.set_grid_site('local')
657 658 659 660 661
    self.bayeswavepsd_job={}
    if self.config.has_option('condor','bayeswave'):
      for ifo in self.ifos:
         self.bayeswavepsd_job[ifo] = BayesWavePSDJob(self.config,os.path.join(self.basepath,'bayeswavepsd_%s.sub'%(ifo)),self.logpath,dax=self.is_dax())
         self.bayeswavepsd_job[ifo].set_grid_site('local')
662 663 664 665
    # Need to create a job file for each IFO combination
    self.engine_jobs={}
    ifocombos=[]
    for N in range(1,len(self.ifos)+1):
666
        for a in permutations(self.ifos,N):
667 668
            ifocombos.append(a)
    for ifos in ifocombos:
669
        self.engine_jobs[ifos] = EngineJob(self.config, os.path.join(self.basepath,'engine_%s.sub'%(reduce(lambda x,y:x+y, map(str,ifos)))),self.logpath,engine=self.engine,dax=self.is_dax(), site=site)
670
    self.results_page_job = ResultsPageJob(self.config,os.path.join(self.basepath,'resultspage.sub'),self.logpath,dax=self.is_dax())
671
    self.results_page_job.set_grid_site('local')
672
    self.cotest_results_page_job = ResultsPageJob(self.config,os.path.join(self.basepath,'resultspagecoherent.sub'),self.logpath,dax=self.is_dax())
673
    self.cotest_results_page_job.set_grid_site('local')
674
    self.merge_job = MergeNSJob(self.config,os.path.join(self.basepath,'merge_runs.sub'),self.logpath,dax=self.is_dax())
675
    self.merge_job.set_grid_site('local')
676
    self.coherence_test_job = CoherenceTestJob(self.config,os.path.join(self.basepath,'coherence_test.sub'),self.logpath,dax=self.is_dax())
677
    self.coherence_test_job.set_grid_site('local')
678
    self.gracedbjob = GraceDBJob(self.config,os.path.join(self.basepath,'gracedb.sub'),self.logpath,dax=self.is_dax())
679
    self.gracedbjob.set_grid_site('local')
680
    # Process the input to build list of analyses to do
John Douglas Veitch's avatar
Working  
John Douglas Veitch committed
681
    self.events=self.setup_from_inputs()
682

683
    # Sanity checking
John Douglas Veitch's avatar
Working  
John Douglas Veitch committed
684
    if len(self.events)==0:
685 686
      print 'No input events found, please check your config if you expect some events'
    self.times=[e.trig_time for e in self.events]
687

688
    # Set up the segments
689
    if not (self.config.has_option('input','gps-start-time') and self.config.has_option('input','gps-end-time')) and len(self.times)>0:
690
      (mintime,maxtime)=self.get_required_data(self.times)
691 692 693 694
      if not self.config.has_option('input','gps-start-time'):
        self.config.set('input','gps-start-time',str(int(floor(mintime))))
      if not self.config.has_option('input','gps-end-time'):
        self.config.set('input','gps-end-time',str(int(ceil(maxtime))))
695
    self.add_science_segments()
696

697
    # Save the final configuration that is being used
698
    # first to the run dir
699 700
    conffilename=os.path.join(self.basepath,'config.ini')
    with open(conffilename,'wb') as conffile:
John Douglas Veitch's avatar
John Douglas Veitch committed
701
      self.config.write(conffile)
702 703 704 705
    if self.config.has_option('paths','webdir'):
      mkdirs(self.config.get('paths','webdir'))
      with open(os.path.join(self.config.get('paths','webdir'),'config.ini'),'wb') as conffile:
        self.config.write(conffile)
706

707
    # Generate the DAG according to the config given
708
    for event in self.events: self.add_full_analysis(event)
709
    self.add_skyarea_followup()
710 711 712
    if self.config.has_option('analysis','upload-to-gracedb'):
      if self.config.getboolean('analysis','upload-to-gracedb'):
        self.add_gracedb_FITSskymap_upload(self.events[0],engine=self.engine)
713 714 715 716 717 718 719 720 721
    self.dagfilename="lalinference_%s-%s"%(self.config.get('input','gps-start-time'),self.config.get('input','gps-end-time'))
    self.set_dag_file(self.dagfilename)
    if self.is_dax():
      self.set_dax_file(self.dagfilename)

  def add_skyarea_followup(self):
    # Add skyarea jobs if the executable is given
    # Do one for each results page for now
    if self.config.has_option('condor','skyarea'):
722 723
      self.skyareajob=SkyAreaJob(self.config,os.path.join(self.basepath,'skyarea.sub'),self.logpath,dax=self.is_dax())
      respagenodes=filter(lambda x: isinstance(x,ResultsPageNode) ,self.get_nodes())
724 725 726 727
      if self.engine=='lalinferenceburst':
          prefix='LIB_'
      else:
          prefix='LALInference_'
728
      for p in respagenodes:
729
          skyareanode=SkyAreaNode(self.skyareajob,prefix=prefix)
730
          skyareanode.add_resultspage_parent(p)
731
          skyareanode.set_ifos(p.ifos)
732
          self.add_node(skyareanode)
733

734
  def add_full_analysis(self,event):
735
    if self.engine=='lalinferencenest' or  self.engine=='lalinferenceburst':
736
      result=self.add_full_analysis_lalinferencenest(event)
737
    elif self.engine=='lalinferencemcmc':
738
      result=self.add_full_analysis_lalinferencemcmc(event)
739
    elif self.engine=='lalinferencebambi' or self.engine=='lalinferencebambimpi':
740 741 742
      result=self.add_full_analysis_lalinferencebambi(event)
    return result

743 744 745 746 747 748 749 750 751
  def create_frame_pfn_file(self):
    """
    Create a pegasus cache file name, uses inspiralutils
    """
    import inspiralutils as iu
    gpsstart=self.config.get('input','gps-start-time')
    gpsend=self.config.get('input','gps-end-time')
    pfnfile=iu.create_frame_pfn_file(self.frtypes,gpsstart,gpsend)
    return pfnfile
752

John Douglas Veitch's avatar
John Douglas Veitch committed
753 754
  def get_required_data(self,times):
    """
John Douglas Veitch's avatar
John Douglas Veitch committed
755
    Calculate the data that will be needed to process all events
John Douglas Veitch's avatar
John Douglas Veitch committed
756
    """
757
    #psdlength = self.config.getint('input','max-psd-length')
758
    padding=self.config.getint('input','padding')
759 760 761 762 763 764
    if self.config.has_option('engine','seglen') or self.config.has_option('lalinference','seglen'):
      if self.config.has_option('engine','seglen'):
        seglen = self.config.getint('engine','seglen')
      if self.config.has_option('lalinference','seglen'):
        seglen = self.config.getint('lalinference','seglen')

765
      if os.path.isfile(os.path.join(self.basepath,'psd.xml.gz')) or self.config.has_option('condor','bayesline') or self.config.has_option('condor','bayeswave'):
766
        psdlength = 0
767 768 769 770 771
        padding = 0
        self.config.set('input','padding',str(padding))
        if (np.log2(seglen)%1):
          seglen = np.power(2., np.ceil(np.log2(seglen)))

772 773
      else:
        psdlength = 32*seglen
774 775
    else:
      seglen = max(e.duration for e in self.events)
776
      if os.path.isfile(os.path.join(self.basepath,'psd.xml.gz')) or self.config.has_option('condor','bayesline') or self.config.has_option('condor','bayeswave'):
777
        psdlength = 0
778 779 780 781
        padding = 0
        self.config.set('input','padding',str(padding))
        if (np.log2(seglen)%1):
          seglen = np.power(2., np.ceil(np.log2(seglen)))
782 783 784
      else:
        psdlength = 32*seglen
    # Assume that the data interval is (end_time - seglen -padding , end_time + psdlength +padding )
785
    # -> change to (trig_time - seglen - padding - psdlength + 2 , trig_time + padding + 2) to estimate the psd before the trigger for online follow-up.
786
    # Also require padding before start time
787
    return (min(times)-padding-seglen-psdlength+2,max(times)+padding+2)
John Douglas Veitch's avatar
John Douglas Veitch committed
788

789 790 791 792 793
  def setup_from_times(self,times):
    """
    Generate a DAG from a list of times
    """
    for time in self.times:
794
      self.add_full_analysis(Event(trig_time=time))
795

796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819
  def select_events(self):
    """
    Read events from the config parser. Understands both ranges and comma separated events, or combinations
    eg. events=[0,1,5:10,21] adds to the analysis the events: 0,1,5,6,7,8,9,10 and 21
    """
    events=[]
    times=[]
    raw_events=self.config.get('input','events').replace('[','').replace(']','').split(',')
    for raw_event in raw_events:
        if ':' in raw_event:
            limits=raw_event.split(':')
            if len(limits) != 2:
                print "Error: in event config option; ':' must separate two numbers."
                exit(0)
            low=int(limits[0])
            high=int(limits[1])
            if low>high:
                events.extend(range(int(high),int(low)+1))
            elif high>low:
                events.extend(range(int(low),int(high)+1))
        else:
            events.append(int(raw_event))
    return events

John Douglas Veitch's avatar
Working  
John Douglas Veitch committed
820 821 822 823 824 825 826
  def setup_from_inputs(self):
    """
    Scan the list of inputs, i.e.
    gps-time-file, injection-file, sngl-inspiral-file, coinc-inspiral-file, pipedown-database
    in the [input] section of the ini file.
    And process the events found therein
    """
827
    events=[]
828 829
    gpsstart=None
    gpsend=None
830 831 832 833
    if self.config.has_option('input','gps-start-time'):
      gpsstart=self.config.getfloat('i