lalinference_pipe_utils.py 96.1 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):
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
John Douglas Veitch's avatar
John Douglas Veitch committed
50 51 52
    if event_id is not None:
        self.event_id=event_id
    else:
John Douglas Veitch's avatar
John Douglas Veitch committed
53 54 55
        self.event_id=Event.new_id()
    if self.injection is not None:
        self.trig_time=self.injection.get_end()
56
        if event_id is None: self.event_id=int(str(self.injection.simulation_id).split(':')[2])
57 58 59
    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
60
    if self.sngltrigger is not None:
61
        self.trig_time=self.sngltrigger.get_end()
John Douglas Veitch's avatar
John Douglas Veitch committed
62
        self.event_id=int(str(self.sngltrigger.event_id).split(':')[2])
63 64 65 66
    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()))
67 68 69 70 71 72 73
    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
74 75

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

77
def readLValert(threshold_snr=None,gid=None,flow=40.0,gracedb="gracedb",basepath="./",downloadpsd=True):
78
  """
79
  Parse LV alert file, containing coinc, sngl, coinc_event_map.
80 81
  and create a list of Events as input for pipeline
  Based on Chris Pankow's script
82
  """
83
  output=[]
84 85
  from glue.ligolw import utils
  from glue.ligolw import lsctables
86
  from glue.ligolw import ligolw
Rory Smith's avatar
Rory Smith committed
87
  class PSDContentHandler(ligolw.LIGOLWContentHandler):
88
    pass
Rory Smith's avatar
Rory Smith committed
89
  lsctables.use_in(PSDContentHandler)
90 91
  from glue.ligolw import param
  from glue.ligolw import array
92
  import subprocess
Rory Smith's avatar
Rory Smith committed
93
  from lal import series as lalseries
94
  from subprocess import Popen, PIPE
95 96
  cwd=os.getcwd()
  os.chdir(basepath)
97
  print "%s download %s coinc.xml"%(gracedb,gid)
Vivien Raymond's avatar
Vivien Raymond committed
98
  subprocess.call([gracedb,"download", gid ,"coinc.xml"])
Rory Smith's avatar
Rory Smith committed
99
  xmldoc=utils.load_filename("coinc.xml",contenthandler = PSDContentHandler)
100
  coinctable = lsctables.CoincInspiralTable.get_table(xmldoc)
101
  coinc_events = [event for event in coinctable]
102
  sngltable = lsctables.SnglInspiralTable.get_table(xmldoc)
103
  sngl_events = [event for event in sngltable]
104 105 106
  #Issues to identify IFO with good data that did not produce a trigger
  #search_summary = lsctables.getTablesByType(xmldoc, lsctables.SearchSummaryTable)[0]
  #ifos = search_summary[0].ifos.split(",")
107 108
  #coinc_table = lsctables.getTablesByType(xmldoc, lsctables.CoincTable)[0]
  #ifos = coinc_table[0].instruments.split(",")
109
  trigSNR = 2.0*coinctable[0].snr #The factor of 2.0 is because detection pipelines recover SNR lower than PE can recover.
110
  # Parse PSD
111
  srate_psdfile=16384
112 113 114 115
  ifos=None
  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
116
    xmlpsd = lalseries.read_psd_xmldoc(utils.load_filename('psd.xml.gz',contenthandler = lalseries.PSDContentHandler))
117 118
    # Note: This finds the active IFOs by looking for available PSDs
    # Is there another way of getting this info?
119
    ifos = xmlpsd.keys()
120
  psdasciidic=None
121
  fhigh=None
122
  if os.path.exists("psd.xml.gz"):
123 124 125
    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
126
  else:
127
    print "Failed to gracedb download %s psd.xml.gz. lalinference will estimate the psd itself." % gid
128
  # Logic for template duration and sample rate disabled
129
  coinc_map = lsctables.CoincMapTable.get_table(xmldoc)
130 131
  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] ]
132 133 134
    dur=[]
    srate=[]
    for e in these_sngls:
135
      # Review: Replace this with a call to LALSimulation function at some point
136 137 138
      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) ) ) )
139
      srate.append(pow(2.0, ceil( log(float(strlen.splitlines()[1].split()[5]), 2) ) ) * 2 )
140 141 142 143 144 145 146
      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)
147 148 149 150
    if max(srate)<srate_psdfile:
      srate = max(srate)
    else:
      srate = srate_psdfile
Vivien Raymond's avatar
Vivien Raymond committed
151
      fhigh = srate_psdfile/2.0 * 0.95 # Because of the drop-off near Nyquist of the PSD from gstlal
152 153 154 155
    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)
156

157
  print "Found %d coinc events in table." % len(coinc_events)
158
  os.chdir(cwd)
159 160
  return output

161 162 163 164 165 166 167
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
168
    import sqlite3
169 170 171 172
    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)
173
    #dbtables.DBTable_set_connection(connection)
174
    return (connection,working_filename)
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 212 213 214 215 216 217 218
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()

219

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

263 264

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

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

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


349
def scan_timefile(timefile):
350 351 352 353 354 355 356 357 358 359 360 361
    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))
362 363
    timefilehandle.close()
    return times
364

365 366 367 368 369 370 371 372 373 374 375 376 377 378 379
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
380
  try:
Rory Smith's avatar
Rory Smith committed
381 382
    #from pylal import series
    from lal import series as series
383
    lal=0
384 385 386
  except ImportError:
    print "ERROR, cannot import pylal.series in bppu/get_xml_psds()\n"
    exit(1)
387

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

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

443
    combine=[]
444

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

453 454 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 480 481 482 483 484 485 486
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)
  # check that trigger masses are identical in each IFO    
  assert len(set(mass1)) == 1
  assert len(set(mass2)) == 1

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

  return mchirp

def get_roq_mchirp_priors(path, roq_paths, roq_params, key, gid):
487 488 489 490 491 492 493 494

  mc_priors = {}

  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]
495 496
  # below is to construct non-overlapping mc priors for multiple roq mass-bin runs
  '''i=0
497 498 499 500 501 502 503
  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.
504
    i+=1'''
505 506 507 508
  if gid is not None:
  	trigger_mchirp = get_trigger_chirpmass(gid)
  else:
	trigger_mchirp = None
509 510 511

  return mc_priors, trigger_mchirp

512 513 514 515 516
def get_roq_mass_freq_scale_factor(mc_priors, trigger_mchirp):
  mc_max = mc_priors['4s'][1]
  mc_min = mc_priors['128s'][0]
  scale_factor = 1
  if trigger_mchirp >= mc_max: 
517
  	scale_factor = 2**(floor(trigger_mchirp/mc_max))
518
  if trigger_mchirp <= mc_min:
519 520
	scale_factor = 1./2**(ceil(trigger_mchirp/mc_min))
  return scale_factor
521

522 523
def create_pfn_tuple(filename,protocol='file://',site='local'):
    return( (os.path.basename(filename),protocol+os.path.abspath(filename),site) )
524

525
class LALInferencePipelineDAG(pipeline.CondorDAG):
526
  def __init__(self,cp,dax=False,first_dag=True,previous_dag=None,site='local'):
527 528
    self.subfiles=[]
    self.config=cp
529
    self.engine=get_engine_name(cp)
John Douglas Veitch's avatar
John Douglas Veitch committed
530
    self.EngineNode=chooseEngineNode(self.engine)
531
    self.site=site
532
    if cp.has_option('paths','basedir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
533
      self.basepath=cp.get('paths','basedir')
534 535 536
    else:
      self.basepath=os.getcwd()
      print 'No basepath specified, using current directory: %s'%(self.basepath)
537
    mkdirs(self.basepath)
538 539
    if dax:
        os.chdir(self.basepath)
540
    self.posteriorpath=os.path.join(self.basepath,'posterior_samples')
John Douglas Veitch's avatar
John Douglas Veitch committed
541
    mkdirs(self.posteriorpath)
542 543 544 545 546 547 548 549 550 551 552 553
    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)
554
    if cp.has_option('paths','cachedir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
555
      self.cachepath=cp.get('paths','cachedir')
556 557
    else:
      self.cachepath=os.path.join(self.basepath,'caches')
558
    mkdirs(self.cachepath)
559
    if cp.has_option('paths','logdir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
560
      self.logpath=cp.get('paths','logdir')
561 562
    else:
      self.logpath=os.path.join(self.basepath,'log')
563
    mkdirs(self.logpath)
564
    if cp.has_option('analysis','ifos'):
565
      self.ifos=ast.literal_eval(cp.get('analysis','ifos'))
566 567 568
    else:
      self.ifos=['H1','L1','V1']
    self.segments={}
569 570 571 572
    if cp.has_option('datafind','veto-categories'):
      self.veto_categories=cp.get('datafind','veto-categories')
    else: self.veto_categories=[]
    for ifo in self.ifos:
573
      self.segments[ifo]=[]
574 575
    self.computeroqweightsnodes={}
    self.bayeslinenodes={}
576
    self.dq={}
577
    self.frtypes=ast.literal_eval(cp.get('datafind','types'))
578
    self.channels=ast.literal_eval(cp.get('data','channels'))
John Douglas Veitch's avatar
John Douglas Veitch committed
579
    self.use_available_data=False
580
    self.webdir=cp.get('paths','webdir')
581 582 583 584
    if cp.has_option('analysis','dataseed'):
      self.dataseed=cp.getint('analysis','dataseed')
    else:
      self.dataseed=None
585
    # Set up necessary job files.
586
    self.prenodes={}
587
    self.datafind_job = pipeline.LSCDataFindJob(self.cachepath,self.logpath,self.config,dax=self.is_dax())
588
    self.datafind_job.add_opt('url-type','file')
589 590
    if cp.has_option('analysis','accounting_group'):
      self.datafind_job.add_condor_cmd('accounting_group',cp.get('analysis','accounting_group'))
591 592
    if cp.has_option('analysis','accounting_group_user'):
      self.datafind_job.add_condor_cmd('accounting_group_user',cp.get('analysis','accounting_group_user'))
593
    self.datafind_job.set_sub_file(os.path.abspath(os.path.join(self.basepath,'datafind.sub')))
594
    self.preengine_job = EngineJob(self.config, os.path.join(self.basepath,'prelalinference.sub'),self.logpath,engine='lalinferencedatadump',ispreengine=True,dax=self.is_dax())
595 596
    self.preengine_job.set_grid_site('local')
    self.preengine_job.set_universe('vanilla')
597 598 599 600 601 602
    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')
603 604 605 606
    # Need to create a job file for each IFO combination
    self.engine_jobs={}
    ifocombos=[]
    for N in range(1,len(self.ifos)+1):
607
        for a in permutations(self.ifos,N):
608 609
            ifocombos.append(a)
    for ifos in ifocombos:
610
        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)
611
    self.results_page_job = ResultsPageJob(self.config,os.path.join(self.basepath,'resultspage.sub'),self.logpath,dax=self.is_dax())
612
    self.results_page_job.set_grid_site('local')
613
    self.cotest_results_page_job = ResultsPageJob(self.config,os.path.join(self.basepath,'resultspagecoherent.sub'),self.logpath,dax=self.is_dax())
614
    self.cotest_results_page_job.set_grid_site('local')
615
    self.merge_job = MergeNSJob(self.config,os.path.join(self.basepath,'merge_runs.sub'),self.logpath,dax=self.is_dax())
616
    self.merge_job.set_grid_site('local')
617
    self.coherence_test_job = CoherenceTestJob(self.config,os.path.join(self.basepath,'coherence_test.sub'),self.logpath,dax=self.is_dax())
618
    self.coherence_test_job.set_grid_site('local')
619
    self.gracedbjob = GraceDBJob(self.config,os.path.join(self.basepath,'gracedb.sub'),self.logpath,dax=self.is_dax())
620
    self.gracedbjob.set_grid_site('local')
621
    # Process the input to build list of analyses to do
John Douglas Veitch's avatar
Working  
John Douglas Veitch committed
622
    self.events=self.setup_from_inputs()
623

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

629
    # Set up the segments
630
    if not (self.config.has_option('input','gps-start-time') and self.config.has_option('input','gps-end-time')) and len(self.times)>0:
631
      (mintime,maxtime)=self.get_required_data(self.times)
632 633 634 635
      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))))
636
    self.add_science_segments()
637

638
    # Save the final configuration that is being used
639
    # first to the run dir
640 641
    conffilename=os.path.join(self.basepath,'config.ini')
    with open(conffilename,'wb') as conffile:
John Douglas Veitch's avatar
John Douglas Veitch committed
642
      self.config.write(conffile)
643 644 645 646
    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)
647

648
    # Generate the DAG according to the config given
649
    for event in self.events: self.add_full_analysis(event)
650
    self.add_skyarea_followup()
651 652 653
    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)
654 655 656 657 658 659 660 661 662
    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'):
663 664
      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())
665 666 667 668
      if self.engine=='lalinferenceburst':
          prefix='LIB_'
      else:
          prefix='LALInference_'
669
      for p in respagenodes:
670
          skyareanode=SkyAreaNode(self.skyareajob,prefix=prefix)
671
          skyareanode.add_resultspage_parent(p)
672
          skyareanode.set_ifos(p.ifos)
673
          self.add_node(skyareanode)
674

675
  def add_full_analysis(self,event):
676
    if self.engine=='lalinferencenest' or  self.engine=='lalinferenceburst':
677
      result=self.add_full_analysis_lalinferencenest(event)
678
    elif self.engine=='lalinferencemcmc':
679
      result=self.add_full_analysis_lalinferencemcmc(event)
680
    elif self.engine=='lalinferencebambi' or self.engine=='lalinferencebambimpi':
681 682 683
      result=self.add_full_analysis_lalinferencebambi(event)
    return result

684

685 686 687 688 689 690 691 692 693
  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
694

John Douglas Veitch's avatar
John Douglas Veitch committed
695 696
  def get_required_data(self,times):
    """
John Douglas Veitch's avatar
John Douglas Veitch committed
697
    Calculate the data that will be needed to process all events
John Douglas Veitch's avatar
John Douglas Veitch committed
698
    """
699
    #psdlength = self.config.getint('input','max-psd-length')
700
    padding=self.config.getint('input','padding')
701 702 703 704 705 706
    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')

707
      if os.path.isfile(os.path.join(self.basepath,'psd.xml.gz')):
708 709 710
        psdlength = 0
      else:
        psdlength = 32*seglen
711 712
    else:
      seglen = max(e.duration for e in self.events)
713
      if os.path.isfile(os.path.join(self.basepath,'psd.xml.gz')):
714 715 716 717
        psdlength = 0
      else:
        psdlength = 32*seglen
    # Assume that the data interval is (end_time - seglen -padding , end_time + psdlength +padding )
718
    # -> change to (trig_time - seglen - padding - psdlength + 2 , trig_time + padding + 2) to estimate the psd before the trigger for online follow-up.
719
    # Also require padding before start time
720
    return (min(times)-padding-seglen-psdlength+2,max(times)+padding+2)
John Douglas Veitch's avatar
John Douglas Veitch committed
721

722 723 724 725 726
  def setup_from_times(self,times):
    """
    Generate a DAG from a list of times
    """
    for time in self.times:
727
      self.add_full_analysis(Event(trig_time=time))
728

729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752
  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
753 754 755 756 757 758 759
  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
    """
760
    events=[]
761 762
    gpsstart=None
    gpsend=None
763 764 765 766
    if self.config.has_option('input','gps-start-time'):
      gpsstart=self.config.getfloat('input','gps-start-time')
    if self.config.has_option('input','gps-end-time'):
      gpsend=self.config.getfloat('input','gps-end-time')
767
    inputnames=['gps-time-file','burst-injection-file','injection-file','sngl-inspiral-file','coinc-inspiral-file','pipedown-db','gid','gstlal-db']
768
    ReadInputFromList=sum([ 1 if self.config.has_option('input',name) else 0 for name in inputnames])
769
    # If no input events given, just return an empty list (e.g. for PP pipeline)
770
    if ReadInputFromList!=1 and (gpsstart is None or gpsend is None):
771
        return []
772
    # Review: Clean up this section
773
    if self.config.has_option('input','events'):
774
      selected_events=self.config.get('input','events')
775
      print 'Selected events %s'%(str(selected_events))
776

777 778
      if selected_events=='all':
          selected_events=None
779 780
      else:
          selected_events=self.select_events()
781 782
    else:
        selected_events=None
783

784
    # No input file given, analyse the entire time stretch between gpsstart and gpsend
785 786
    if self.config.has_option('input','analyse-all-time') and self.config.getboolean('input','analyse-all-time')==True:
        print 'Setting up for analysis of continuous time stretch %f - %f'%(gpsstart,gpsend)
787
        seglen=self.config.getfloat('engine','seglen')
788
        if(self.config.has_option('input','segment-overlap')):
789 790 791 792 793 794 795 796 797 798
          overlap=self.config.getfloat('input','segment-overlap')
        else:
          overlap=32.;
        if(overlap>seglen):
          print 'ERROR: segment-overlap is greater than seglen'
          sys.exit(1)
        # Now divide gpsstart - gpsend into jobs of seglen - overlap length
        t=gpsstart
        events=[]
        while(t<gpsend):
799
            ev=Event(trig_time=t+seglen-2)
800 801 802 803 804 805 806 807
            ev.set_engine_option('segment-start',str(t-overlap))
            ev.set_engine_option('time-min',str(t))
            tMax=t + seglen - overlap
            if tMax>=gpsend:
                tMax=gpsend
            ev.set_engine_option('time-max',str(tMax))
            events.append(ev)
            t=tMax
808
        return events
809

John Douglas Veitch's avatar
Working  
John Douglas Veitch committed
810 811
    # ASCII list of GPS times
    if self.config.has_option('input','gps-time-file'):
<