lalinference_pipe_utils.py 125 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=20.0,gracedb="gracedb",basepath="./",downloadpsd=True,roq=False):
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
      if roq==False:
129 130 131 132 133 134 135 136
        try:
          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) ) ) )
          srate.append(pow(2.0, ceil( log(float(strlen.splitlines()[1].split()[5]), 2) ) ) * 2 )
        except:
          print "WARNING: lalapps_chirplen --flow",str(flow),"-m1",str(e.mass1),"-m2",str(e.mass2)," failed."
          print "WARNING: make sure you have set manually seglen and srate in the .ini file."
137 138 139 140 141 142 143
      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)
144 145 146 147 148 149 150 151 152 153 154
    if srate:
      if max(srate)<srate_psdfile:
        srate = max(srate)
      else:
        srate = srate_psdfile
        if downloadpsd:
          fhigh = srate_psdfile/2.0 * 0.95 # Because of the drop-off near Nyquist of the PSD from gstlal
    else:
      srate = None
    if dur:
      duration = max(dur)
155
    else:
156
      duration = None
157
    horizon_distance = max(horizon_distance) if len(horizon_distance) > 0 else None
158
    ev=Event(CoincInspiral=coinc, GID=gid, ifos = ifos, duration = duration, srate = srate,
159 160
             trigSNR = trigSNR, fhigh = fhigh, horizon_distance=horizon_distance)
    output.append(ev)
161

162
  print "Found %d coinc events in table." % len(coinc_events)
163
  os.chdir(cwd)
164 165
  return output

166 167 168 169 170 171 172
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
173
    import sqlite3
174 175 176 177
    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)
178
    #dbtables.DBTable_set_connection(connection)
179
    return (connection,working_filename)
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 212 213 214 215 216 217 218 219 220 221 222 223
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()

224
def get_zerolag_pipedown(database_connection, dumpfile=None, gpsstart=None, gpsend=None, max_cfar=-1, min_cfar=-1):
225 226 227 228 229 230 231 232 233 234 235 236 237 238
	"""
	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) \
239
		WHERE coinc_event.time_slide_id=='time_slide:time_slide_id:10049'\
240 241
		"
	if gpsstart is not None:
242
		get_coincs=get_coincs+' and coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1.0e-9 > %f'%(gpsstart)
243
	if gpsend is not None:
244
		get_coincs=get_coincs+' and coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1.0e-9 < %f'%(gpsend)
245 246
	if max_cfar !=-1:
		get_coincs=get_coincs+' and coinc_inspiral.combined_far < %f'%(max_cfar)
247 248
	if min_cfar != -1:
		get_coincs=get_coincs+' and coinc_inspiral.combined_far > %f'%(min_cfar)
249
	db_out=database_connection.cursor().execute(get_coincs)
John Douglas Veitch's avatar
John Douglas Veitch committed
250
    	extra={}
251
	for (sngl_time, ifo, coinc_id, snr, chisq, cfar) in db_out:
John Douglas Veitch's avatar
John Douglas Veitch committed
252 253
      		coinc_id=int(coinc_id.split(":")[-1])
	  	if not coinc_id in output.keys():
254
			output[coinc_id]=Event(trig_time=sngl_time,timeslide_dict={},event_id=int(coinc_id))
John Douglas Veitch's avatar
John Douglas Veitch committed
255 256 257 258
			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}
259 260 261 262 263 264
	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()
265
	return output.values()
266

267
def get_timeslides_pipedown(database_connection, dumpfile=None, gpsstart=None, gpsend=None, max_cfar=-1):
268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283
	"""
	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
284
	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 \
285
		    FROM sngl_inspiral join coinc_event_map on (coinc_event_map.table_name == 'sngl_inspiral' and coinc_event_map.event_id \
286 287
		    == 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)\
288
		    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
289
	joinstr = ' and '
290
	if gpsstart is not None:
291
		get_coincs=get_coincs+ joinstr + ' coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1e-9 > %f'%(gpsstart)
292
	if gpsend is not None:
293
		get_coincs=get_coincs+ joinstr+' coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1e-9 <%f'%(gpsend)
294 295
	if max_cfar!=-1:
		get_coincs=get_coincs+joinstr+' coinc_inspiral.combined_far < %f'%(max_cfar)
296
	db_out=database_connection.cursor().execute(get_coincs)
297 298
	from pylal import SnglInspiralUtils
	extra={}
299
	for (sngl_time, slide, ifo, coinc_id, snr, chisq, cfar) in db_out:
300 301 302 303 304 305 306 307 308
		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}
309 310 311 312 313 314
	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()
315
	return output.values()
316

317 318 319 320 321 322 323 324
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
325 326 327
def chooseEngineNode(name):
  if name=='lalinferencenest':
    return LALInferenceNestNode
328 329
  if name=='lalinferenceburst':
    return LALInferenceBurstNode
John Douglas Veitch's avatar
John Douglas Veitch committed
330 331
  if name=='lalinferencemcmc':
    return LALInferenceMCMCNode
332
  if name=='lalinferencebambi' or name=='lalinferencebambimpi':
333
    return LALInferenceBAMBINode
334 335
  if name=='lalinferencedatadump':
    return LALInferenceDataDumpNode
336 337
  if name=='bayeswavepsd':
    return BayesWavePSDNode
John Douglas Veitch's avatar
John Douglas Veitch committed
338 339
  return EngineNode

340 341 342
def get_engine_name(cp):
    name=cp.get('analysis','engine')
    if name=='random':
343
        engine_list=['lalinferencenest','lalinferencemcmc','lalinferencebambimpi']
344 345 346 347 348 349 350 351 352
        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

353
def scan_timefile(timefile):
354 355 356 357 358 359 360 361 362 363 364 365
    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))
366 367
    timefilehandle.close()
    return times
368

369 370 371 372
def get_xml_psds(psdxml,ifos,outpath,end_time=None):
  """
  Get a psd.xml.gz file and:
  1) Reads it
373 374
  2) Checks the psd file contains all the IFO we want to analyze
  3) Writes down the PSDs 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.
375 376 377
  Input:
    psdxml: psd.xml.gz file
    ifos: list of ifos used for the analysis
378 379
    outpath: path where the ascii PSD will be written to
    (end_time): trigtime for this event. Will be used a part of the PSD file name
380 381 382
  """
  lal=1
  from glue.ligolw import utils
383
  try:
Rory Smith's avatar
Rory Smith committed
384 385
    #from pylal import series
    from lal import series as series
386
    lal=0
387 388 389
  except ImportError:
    print "ERROR, cannot import pylal.series in bppu/get_xml_psds()\n"
    exit(1)
390

391 392 393 394 395 396 397 398 399 400 401
  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')
402
    # Check we don't already have that ascii (e.g. because we are running parallel runs of the save event
403 404 405 406 407 408 409 410 411
    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
412

413 414 415 416
  # 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
417
  xmlpsd =  series.read_psd_xmldoc(utils.load_filename(psdxml,contenthandler = series.PSDContentHandler))
418 419 420 421 422 423 424 425 426
  # 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')
427
    # Check we don't already have that ascii (e.g. because we are running parallel runs of the save event
428
    if os.path.isfile(path_to_ascii_psd):
429
      continue
430 431 432 433 434 435 436 437 438 439 440 441 442 443 444
    # 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
445

446
    combine=[]
447

Rory Smith's avatar
Rory Smith committed
448
    for i in np.arange(len(data.data.data)) :
449
      combine.append([f0+i*deltaF,data.data.data[i]])
450
    np.savetxt(path_to_ascii_psd,combine)
451
    ifo=instrument.encode('ascii')
452 453 454
    # set node.psds dictionary with the path to the ascii files
    out[ifo]=os.path.join(outpath,ifo+'_psd_'+time+'.txt')
  return out
455

456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479
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)
480
  # check that trigger masses are identical in each IFO
481 482 483
  assert len(set(mass1)) == 1
  assert len(set(mass2)) == 1

484
  mchirp = (mass1[0]*mass2[0])**(3./5.) / ( (mass1[0] + mass2[0])**(1./5.) )
485 486 487 488
  os.remove("coinc.xml")

  return mchirp

489
def get_roq_mchirp_priors(path, roq_paths, roq_params, key, gid=None,sim_inspiral=None):
490

491 492
  ## XML and GID cannot be given at the same time
  ## sim_inspiral must already point at the right row
493
  mc_priors = {}
Vivien Raymond's avatar
Vivien Raymond committed
494

495 496 497
  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)
498 499 500 501 502 503

  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]
504 505
  # below is to construct non-overlapping mc priors for multiple roq mass-bin runs
  '''i=0
506 507 508 509 510 511 512
  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.
513
    i+=1'''
514 515
  if gid is not None:
  	trigger_mchirp = get_trigger_chirpmass(gid)
516 517
  elif sim_inspiral is not None:
        trigger_mchirp = sim_inspiral.mchirp
518 519
  else:
	trigger_mchirp = None
520 521 522

  return mc_priors, trigger_mchirp

Vivien Raymond's avatar
Vivien Raymond committed
523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548
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

549
def get_roq_mass_freq_scale_factor(mc_priors, trigger_mchirp, force_flow=None):
550 551 552 553 554
  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]
555
  scale_factor = 1.
556
  if force_flow == None and trigger_mchirp != None:
557 558 559
     if trigger_mchirp >= mc_max:
        scale_factor = 2.**(floor(trigger_mchirp/mc_max))
     if trigger_mchirp <= mc_min:
560
        scale_factor = (2./3.2)**(ceil(trigger_mchirp/mc_min))
561
  elif force_flow != None:
562
     scale_factor = 20./force_flow
563
  return scale_factor
564

565 566
def create_pfn_tuple(filename,protocol='file://',site='local'):
    return( (os.path.basename(filename),protocol+os.path.abspath(filename),site) )
567

Vivien Raymond's avatar
Vivien Raymond committed
568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588
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

589
class LALInferencePipelineDAG(pipeline.CondorDAG):
590
  def __init__(self,cp,dax=False,site='local'):
591 592
    self.subfiles=[]
    self.config=cp
593
    self.engine=get_engine_name(cp)
John Douglas Veitch's avatar
John Douglas Veitch committed
594
    self.EngineNode=chooseEngineNode(self.engine)
595
    self.site=site
596
    if cp.has_option('paths','basedir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
597
      self.basepath=cp.get('paths','basedir')
598 599 600
    else:
      self.basepath=os.getcwd()
      print 'No basepath specified, using current directory: %s'%(self.basepath)
601
    mkdirs(self.basepath)
602
    print("Generating LALInference DAG in {0}".format(self.basepath))
603 604
    if dax:
        os.chdir(self.basepath)
605
    self.posteriorpath=os.path.join(self.basepath,'posterior_samples')
John Douglas Veitch's avatar
John Douglas Veitch committed
606
    mkdirs(self.posteriorpath)
607 608 609 610
    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)
611
    if cp.has_option('paths','cachedir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
612
      self.cachepath=cp.get('paths','cachedir')
613 614
    else:
      self.cachepath=os.path.join(self.basepath,'caches')
615
    mkdirs(self.cachepath)
616
    if cp.has_option('paths','logdir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
617
      self.logpath=cp.get('paths','logdir')
618 619
    else:
      self.logpath=os.path.join(self.basepath,'log')
620
    mkdirs(self.logpath)
621
    if cp.has_option('analysis','ifos'):
622
      self.ifos=ast.literal_eval(cp.get('analysis','ifos'))
623 624 625
    else:
      self.ifos=['H1','L1','V1']
    self.segments={}
626 627 628 629
    if cp.has_option('datafind','veto-categories'):
      self.veto_categories=cp.get('datafind','veto-categories')
    else: self.veto_categories=[]
    for ifo in self.ifos:
630
      self.segments[ifo]=[]
631 632 633
    self.computeroqweightsnode={}
    self.bayeslinenode={}
    self.bayeswavepsdnode={}
634
    self.dq={}
635
    self.frtypes=ast.literal_eval(cp.get('datafind','types'))
636
    self.channels=ast.literal_eval(cp.get('data','channels'))
John Douglas Veitch's avatar
John Douglas Veitch committed
637
    self.use_available_data=False
638
    self.webdir=cp.get('paths','webdir')
639 640 641 642
    if cp.has_option('analysis','dataseed'):
      self.dataseed=cp.getint('analysis','dataseed')
    else:
      self.dataseed=None
643
    # Set up necessary job files.
644
    self.prenodes={}
645
    self.datafind_job = pipeline.LSCDataFindJob(self.cachepath,self.logpath,self.config,dax=self.is_dax())
646
    self.datafind_job.add_opt('url-type','file')
647 648 649
    # If running on OSG use its datafind server
    if cp.has_option('analysis','osg') and cp.getboolean('analysis','osg'):
        self.datafind_job.add_opt('server','datafind.ligo.org')
650 651 652 653
    if cp.has_option('condor','accounting_group'):
      self.datafind_job.add_condor_cmd('accounting_group',cp.get('condor','accounting_group'))
    if cp.has_option('condor','accounting_group_user'):
      self.datafind_job.add_condor_cmd('accounting_group_user',cp.get('condor','accounting_group_user'))
654
    self.datafind_job.set_sub_file(os.path.abspath(os.path.join(self.basepath,'datafind.sub')))
655
    self.preengine_job = EngineJob(self.config, os.path.join(self.basepath,'prelalinference.sub'),self.logpath,engine='lalinferencedatadump',ispreengine=True,dax=self.is_dax())
656 657
    self.preengine_job.set_grid_site('local')
    self.preengine_job.set_universe('vanilla')
658 659 660 661 662 663
    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')
664 665 666 667 668
    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')
669 670 671 672
    # Need to create a job file for each IFO combination
    self.engine_jobs={}
    ifocombos=[]
    for N in range(1,len(self.ifos)+1):
673
        for a in permutations(self.ifos,N):
674 675
            ifocombos.append(a)
    for ifos in ifocombos:
676
        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)
677
    self.results_page_job = ResultsPageJob(self.config,os.path.join(self.basepath,'resultspage.sub'),self.logpath,dax=self.is_dax())
678
    self.results_page_job.set_grid_site('local')
679
    self.cotest_results_page_job = ResultsPageJob(self.config,os.path.join(self.basepath,'resultspagecoherent.sub'),self.logpath,dax=self.is_dax())
680
    self.cotest_results_page_job.set_grid_site('local')
681 682 683 684 685 686 687 688
    if self.engine=='lalinferencemcmc':
        self.combine_job = CombineMCMCJob(self.config,os.path.join(self.basepath,'combine_files.sub'),self.logpath,dax=self.is_dax())
        self.combine_job.set_grid_site('local')
        self.merge_job = MergeJob(self.config,os.path.join(self.basepath,'merge_runs.sub'),self.logpath,dax=self.is_dax(),engine='mcmc')
        self.merge_job.set_grid_site('local')
    else:
        self.merge_job = MergeJob(self.config,os.path.join(self.basepath,'merge_runs.sub'),self.logpath,dax=self.is_dax(),engine='nest')
        self.merge_job.set_grid_site('local')
689
    self.coherence_test_job = CoherenceTestJob(self.config,os.path.join(self.basepath,'coherence_test.sub'),self.logpath,dax=self.is_dax())
690
    self.coherence_test_job.set_grid_site('local')
691
    self.gracedbjob = GraceDBJob(self.config,os.path.join(self.basepath,'gracedb.sub'),self.logpath,dax=self.is_dax())
692
    self.gracedbjob.set_grid_site('local')
693
    # Process the input to build list of analyses to do
John Douglas Veitch's avatar
Working  
John Douglas Veitch committed
694
    self.events=self.setup_from_inputs()
695

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

701
    # Set up the segments
702
    if not (self.config.has_option('input','gps-start-time') and self.config.has_option('input','gps-end-time')) and len(self.times)>0:
703
      (mintime,maxtime)=self.get_required_data(self.times)
704 705 706 707
      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))))
708
    self.add_science_segments()
709

710
    # Save the final configuration that is being used
711
    # first to the run dir
712 713
    conffilename=os.path.join(self.basepath,'config.ini')
    with open(conffilename,'wb') as conffile:
John Douglas Veitch's avatar
John Douglas Veitch committed
714
      self.config.write(conffile)
715 716 717 718
    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)
719

720
    # Generate the DAG according to the config given
721
    for event in self.events: self.add_full_analysis(event)
722
    self.add_skyarea_followup()
723 724 725
    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)
726
    self.dagfilename="lalinference_%s-%s"%(self.config.get('input','gps-start-time'),self.config.get('input','gps-end-time'))
727
    self.set_dag_file(os.path.join(self.basepath,self.dagfilename))
728 729 730 731 732 733 734
    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'):
735 736
      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())
737
      if self.engine=='lalinferenceburst':
738
          prefix='LIB'
739
      else:
740
          prefix='LALInference'
741
      for p in respagenodes:
742
          skyareanode=SkyAreaNode(self.skyareajob,prefix=prefix)
743
          skyareanode.add_resultspage_parent(p)
744
          skyareanode.set_ifos(p.ifos)
745
          self.add_node(skyareanode)
746

747
  def add_full_analysis(self,event):
748
    if self.engine=='lalinferencenest' or  self.engine=='lalinferenceburst':
749
      result=self.add_full_analysis_lalinferencenest(event)
750
    elif self.engine=='lalinferencemcmc':
751
      result=self.add_full_analysis_lalinferencemcmc(event)
752
    elif self.engine=='lalinferencebambi' or self.engine=='lalinferencebambimpi':
753 754 755
      result=self.add_full_analysis_lalinferencebambi(event)
    return result

756 757 758 759 760 761 762 763 764
  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
765

John Douglas Veitch's avatar
John Douglas Veitch committed
766 767
  def get_required_data(self,times):
    """
John Douglas Veitch's avatar
John Douglas Veitch committed
768
    Calculate the data that will be needed to process all events
John Douglas Veitch's avatar
John Douglas Veitch committed
769
    """
770
    #psdlength = self.config.getint('input','max-psd-length')
771
    padding=self.config.getint('input','padding')
772 773
    if self.config.has_option('engine','seglen') or self.config.has_option('lalinference','seglen'):
      if self.config.has_option('engine','seglen'):
774
        seglen = int(np.ceil(self.config.getfloat('engine','seglen')))
775 776 777
      if self.config.has_option('lalinference','seglen'):
        seglen = self.config.getint('lalinference','seglen')

778
      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'):
779
        psdlength = 0
780 781
        padding = 0
        self.config.set('input','padding',str(padding))
782 783 784
        if self.config.has_option('condor','bayeswave'):
          if (np.log2(seglen)%1):
            seglen = np.power(2., np.ceil(np.log2(seglen)))
785

786 787
      else:
        psdlength = 32*seglen
788 789
    else:
      seglen = max(e.duration for e in self.events)
790
      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'):
791
        psdlength = 0
792 793
        padding = 0
        self.config.set('input','padding',str(padding))
794 795 796
        if self.config.has_option('condor','bayeswave'):
          if (np.log2(seglen)%1):
            seglen = np.power(2., np.ceil(np.log2(seglen)))
797 798 799
      else:
        psdlength = 32*seglen
    # Assume that the data interval is (end_time - seglen -padding , end_time + psdlength +padding )
800
    # -> change to (trig_time - seglen - padding - psdlength + 2 , trig_time + padding + 2) to estimate the psd before the trigger for online follow-up.