lalinference_pipe_utils.py 126 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.

26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81
def guess_url(fslocation):
    """
    Try to work out the web address of a given path
    """
    SERVER="localhost"
    USER=os.environ['USER']
    HOST=socket.getfqdn()
    if 'public_html' in fslocation:
        k='public_html/'
    elif 'WWW' in fslocation:
        k='WWW/'
    elif 'www_html' in fslocation:
        k='www_html/'
    else:
        k=None
    if k is not None:
        (a,b)=fslocation.split(k)
        webpath=os.path.join('~%s'%(USER),b)
        onweb=True
    else:
        (c,d)=fslocation.split(USER)
        for k in ['public_html','WWW','www_html']:
            trypath=c+os.environ['USER']+'/'+k+d
            #Follow symlinks
            if os.path.realpath(trypath)==os.path.normpath(fslocation):
                (a,b)=trypath.split(k)
                webpath=os.path.join('~%s'%(USER),b,webpath)
                onweb=True
                break
            else:
                webpath=fslocation
                onweb=False
    if 'atlas' in HOST:
        url="https://atlas1.atlas.aei.uni-hannover.de/"
    elif 'cit' in HOST or 'caltech' in HOST:
        url="https://ldas-jobs.ligo.caltech.edu/"
    elif 'ligo-wa' in HOST:
        url="https://ldas-jobs.ligo-wa.caltech.edu/"
    elif 'ligo-la' in HOST:
        url="https://ldas-jobs.ligo-la.caltech.edu/"
    elif 'uwm' in HOST or 'nemo' in HOST:
        url="https://ldas-jobs.phys.uwm.edu/"
    elif 'phy.syr.edu' in HOST:
        url="https://sugar-jobs.phy.syr.edu/"
    elif 'arcca.cf.ac.uk' in HOST:
        url="https://geo2.arcca.cf.ac.uk/"
    elif 'vulcan' in HOST:
        url="https://galahad.aei.mpg.de/"
    else:
        if onweb:
          url="http://%s/"%(HOST)
        else:
          url=HOST+':'
    url=url+webpath
    return(url)

John Douglas Veitch's avatar
John Douglas Veitch committed
82 83 84 85
class Event():
  """
  Represents a unique event to run on
  """
John Douglas Veitch's avatar
John Douglas Veitch committed
86
  new_id=itertools.count().next
87
  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
88
    self.trig_time=trig_time
John Douglas Veitch's avatar
John Douglas Veitch committed
89
    self.injection=SimInspiral
90
    self.burstinjection=SimBurst
John Douglas Veitch's avatar
John Douglas Veitch committed
91
    self.sngltrigger=SnglInspiral
92 93 94 95
    if timeslide_dict is None:
      self.timeslides={}
    else:
      self.timeslides=timeslide_dict
96 97
    self.GID=GID
    self.coinctrigger=CoincInspiral
98 99 100 101
    if ifos is None:
      self.ifos = []
    else:
      self.ifos = ifos
102 103
    self.duration = duration
    self.srate = srate
104
    self.trigSNR = trigSNR
105
    self.fhigh = fhigh
106
    self.horizon_distance = horizon_distance
John Douglas Veitch's avatar
John Douglas Veitch committed
107 108 109
    if event_id is not None:
        self.event_id=event_id
    else:
John Douglas Veitch's avatar
John Douglas Veitch committed
110 111 112
        self.event_id=Event.new_id()
    if self.injection is not None:
        self.trig_time=self.injection.get_end()
113
        if event_id is None: self.event_id=int(str(self.injection.simulation_id).split(':')[2])
114 115 116
    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
117
    if self.sngltrigger is not None:
118
        self.trig_time=self.sngltrigger.get_end()
John Douglas Veitch's avatar
John Douglas Veitch committed
119
        self.event_id=int(str(self.sngltrigger.event_id).split(':')[2])
120 121 122 123
    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()))
124 125 126 127 128 129 130
    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
131 132

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

134
def readLValert(threshold_snr=None,gid=None,flow=20.0,gracedb="gracedb",basepath="./",downloadpsd=True,roq=False):
135
  """
136
  Parse LV alert file, containing coinc, sngl, coinc_event_map.
137 138
  and create a list of Events as input for pipeline
  Based on Chris Pankow's script
139
  """
140
  output=[]
141 142
  from glue.ligolw import utils
  from glue.ligolw import lsctables
143
  from glue.ligolw import ligolw
Rory Smith's avatar
Rory Smith committed
144
  class PSDContentHandler(ligolw.LIGOLWContentHandler):
145
    pass
Rory Smith's avatar
Rory Smith committed
146
  lsctables.use_in(PSDContentHandler)
147 148
  from glue.ligolw import param
  from glue.ligolw import array
149
  import subprocess
Rory Smith's avatar
Rory Smith committed
150
  from lal import series as lalseries
151
  from subprocess import Popen, PIPE
152 153
  cwd=os.getcwd()
  os.chdir(basepath)
154
  print "%s download %s coinc.xml"%(gracedb,gid)
Vivien Raymond's avatar
Vivien Raymond committed
155
  subprocess.call([gracedb,"download", gid ,"coinc.xml"])
Rory Smith's avatar
Rory Smith committed
156
  xmldoc=utils.load_filename("coinc.xml",contenthandler = PSDContentHandler)
157
  coinctable = lsctables.CoincInspiralTable.get_table(xmldoc)
158
  coinc_events = [event for event in coinctable]
159
  sngltable = lsctables.SnglInspiralTable.get_table(xmldoc)
160
  sngl_events = [event for event in sngltable]
161 162
  ifos = coinctable[0].ifos.split(",")
  trigSNR = coinctable[0].snr
163
  # Parse PSD
164
  srate_psdfile=16384
165
  fhigh=None
166 167 168
  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
169
    xmlpsd = lalseries.read_psd_xmldoc(utils.load_filename('psd.xml.gz',contenthandler = lalseries.PSDContentHandler))
170 171 172 173 174 175
    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
176
  coinc_map = lsctables.CoincMapTable.get_table(xmldoc)
177 178
  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] ]
179 180
    dur=[]
    srate=[]
181
    horizon_distance=[]
182
    for e in these_sngls:
183
      # Review: Replace this with a call to LALSimulation function at some point
184
      if roq==False:
185 186 187 188 189 190 191 192
        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."
193 194 195 196 197 198 199
      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)
200 201 202 203 204 205 206 207 208 209 210
    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)
211
    else:
212
      duration = None
213
    horizon_distance = max(horizon_distance) if len(horizon_distance) > 0 else None
214
    ev=Event(CoincInspiral=coinc, GID=gid, ifos = ifos, duration = duration, srate = srate,
215 216
             trigSNR = trigSNR, fhigh = fhigh, horizon_distance=horizon_distance)
    output.append(ev)
217

218
  print "Found %d coinc events in table." % len(coinc_events)
219
  os.chdir(cwd)
220 221
  return output

222 223 224 225 226 227 228
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
229
    import sqlite3
230 231 232 233
    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)
234
    #dbtables.DBTable_set_connection(connection)
235
    return (connection,working_filename)
236

237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279
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()

280
def get_zerolag_pipedown(database_connection, dumpfile=None, gpsstart=None, gpsend=None, max_cfar=-1, min_cfar=-1):
281 282 283 284 285 286 287 288 289 290 291 292 293 294
	"""
	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) \
295
		WHERE coinc_event.time_slide_id=='time_slide:time_slide_id:10049'\
296 297
		"
	if gpsstart is not None:
298
		get_coincs=get_coincs+' and coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1.0e-9 > %f'%(gpsstart)
299
	if gpsend is not None:
300
		get_coincs=get_coincs+' and coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1.0e-9 < %f'%(gpsend)
301 302
	if max_cfar !=-1:
		get_coincs=get_coincs+' and coinc_inspiral.combined_far < %f'%(max_cfar)
303 304
	if min_cfar != -1:
		get_coincs=get_coincs+' and coinc_inspiral.combined_far > %f'%(min_cfar)
305
	db_out=database_connection.cursor().execute(get_coincs)
John Douglas Veitch's avatar
John Douglas Veitch committed
306
    	extra={}
307
	for (sngl_time, ifo, coinc_id, snr, chisq, cfar) in db_out:
John Douglas Veitch's avatar
John Douglas Veitch committed
308 309
      		coinc_id=int(coinc_id.split(":")[-1])
	  	if not coinc_id in output.keys():
310
			output[coinc_id]=Event(trig_time=sngl_time,timeslide_dict={},event_id=int(coinc_id))
John Douglas Veitch's avatar
John Douglas Veitch committed
311 312 313 314
			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}
315 316 317 318 319 320
	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()
321
	return output.values()
322

323
def get_timeslides_pipedown(database_connection, dumpfile=None, gpsstart=None, gpsend=None, max_cfar=-1):
324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339
	"""
	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
340
	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 \
341
		    FROM sngl_inspiral join coinc_event_map on (coinc_event_map.table_name == 'sngl_inspiral' and coinc_event_map.event_id \
342 343
		    == 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)\
344
		    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
345
	joinstr = ' and '
346
	if gpsstart is not None:
347
		get_coincs=get_coincs+ joinstr + ' coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1e-9 > %f'%(gpsstart)
348
	if gpsend is not None:
349
		get_coincs=get_coincs+ joinstr+' coinc_inspiral.end_time+coinc_inspiral.end_time_ns*1e-9 <%f'%(gpsend)
350 351
	if max_cfar!=-1:
		get_coincs=get_coincs+joinstr+' coinc_inspiral.combined_far < %f'%(max_cfar)
352
	db_out=database_connection.cursor().execute(get_coincs)
353 354
	from pylal import SnglInspiralUtils
	extra={}
355
	for (sngl_time, slide, ifo, coinc_id, snr, chisq, cfar) in db_out:
356 357 358 359 360 361 362 363 364
		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}
365 366 367 368 369 370
	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()
371
	return output.values()
372

373 374 375 376 377 378 379 380
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
381 382 383
def chooseEngineNode(name):
  if name=='lalinferencenest':
    return LALInferenceNestNode
384 385
  if name=='lalinferenceburst':
    return LALInferenceBurstNode
John Douglas Veitch's avatar
John Douglas Veitch committed
386 387
  if name=='lalinferencemcmc':
    return LALInferenceMCMCNode
388
  if name=='lalinferencebambi' or name=='lalinferencebambimpi':
389
    return LALInferenceBAMBINode
390 391
  if name=='lalinferencedatadump':
    return LALInferenceDataDumpNode
392 393
  if name=='bayeswavepsd':
    return BayesWavePSDNode
John Douglas Veitch's avatar
John Douglas Veitch committed
394 395
  return EngineNode

396 397 398
def get_engine_name(cp):
    name=cp.get('analysis','engine')
    if name=='random':
399
        engine_list=['lalinferencenest','lalinferencemcmc','lalinferencebambimpi']
400 401 402 403 404 405 406 407 408
        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

409
def scan_timefile(timefile):
410 411 412 413 414 415 416 417 418 419 420 421
    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))
422 423
    timefilehandle.close()
    return times
424

425 426 427 428
def get_xml_psds(psdxml,ifos,outpath,end_time=None):
  """
  Get a psd.xml.gz file and:
  1) Reads it
429 430
  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.
431 432 433
  Input:
    psdxml: psd.xml.gz file
    ifos: list of ifos used for the analysis
434 435
    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
436 437 438
  """
  lal=1
  from glue.ligolw import utils
439
  try:
Rory Smith's avatar
Rory Smith committed
440 441
    #from pylal import series
    from lal import series as series
442
    lal=0
443 444 445
  except ImportError:
    print "ERROR, cannot import pylal.series in bppu/get_xml_psds()\n"
    exit(1)
446

447 448 449 450 451 452 453 454 455 456 457
  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')
458
    # Check we don't already have that ascii (e.g. because we are running parallel runs of the save event
459 460 461 462 463 464 465 466 467
    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
468

469 470 471 472
  # 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
473
  xmlpsd =  series.read_psd_xmldoc(utils.load_filename(psdxml,contenthandler = series.PSDContentHandler))
474 475 476 477 478 479 480 481 482
  # 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')
483
    # Check we don't already have that ascii (e.g. because we are running parallel runs of the save event
484
    if os.path.isfile(path_to_ascii_psd):
485
      continue
486 487 488 489 490 491 492 493 494 495 496 497 498 499 500
    # 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
501

502
    combine=[]
503

Rory Smith's avatar
Rory Smith committed
504
    for i in np.arange(len(data.data.data)) :
505
      combine.append([f0+i*deltaF,data.data.data[i]])
506
    np.savetxt(path_to_ascii_psd,combine)
507
    ifo=instrument.encode('ascii')
508 509 510
    # set node.psds dictionary with the path to the ascii files
    out[ifo]=os.path.join(outpath,ifo+'_psd_'+time+'.txt')
  return out
511

512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535
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)
536
  # check that trigger masses are identical in each IFO
537 538 539
  assert len(set(mass1)) == 1
  assert len(set(mass2)) == 1

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

  return mchirp

545
def get_roq_mchirp_priors(path, roq_paths, roq_params, key, gid=None,sim_inspiral=None):
546

547 548
  ## XML and GID cannot be given at the same time
  ## sim_inspiral must already point at the right row
549
  mc_priors = {}
Vivien Raymond's avatar
Vivien Raymond committed
550

551 552 553
  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)
554 555 556 557 558 559

  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]
560 561
  # below is to construct non-overlapping mc priors for multiple roq mass-bin runs
  '''i=0
562 563 564 565 566 567 568
  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.
569
    i+=1'''
570 571
  if gid is not None:
  	trigger_mchirp = get_trigger_chirpmass(gid)
572 573
  elif sim_inspiral is not None:
        trigger_mchirp = sim_inspiral.mchirp
574 575
  else:
	trigger_mchirp = None
576 577 578

  return mc_priors, trigger_mchirp

Vivien Raymond's avatar
Vivien Raymond committed
579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604
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

605
def get_roq_mass_freq_scale_factor(mc_priors, trigger_mchirp, force_flow=None):
606 607 608 609 610
  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]
611
  scale_factor = 1.
612
  if force_flow == None and trigger_mchirp != None:
613 614 615
     if trigger_mchirp >= mc_max:
        scale_factor = 2.**(floor(trigger_mchirp/mc_max))
     if trigger_mchirp <= mc_min:
616
        scale_factor = (2./3.2)**(ceil(trigger_mchirp/mc_min))
617
  elif force_flow != None:
618
     scale_factor = 20./force_flow
619
  return scale_factor
620

621 622
def create_pfn_tuple(filename,protocol='file://',site='local'):
    return( (os.path.basename(filename),protocol+os.path.abspath(filename),site) )
623

Vivien Raymond's avatar
Vivien Raymond committed
624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644
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

645
class LALInferencePipelineDAG(pipeline.CondorDAG):
646
  def __init__(self,cp,dax=False,site='local'):
647 648
    self.subfiles=[]
    self.config=cp
649
    self.engine=get_engine_name(cp)
John Douglas Veitch's avatar
John Douglas Veitch committed
650
    self.EngineNode=chooseEngineNode(self.engine)
651
    self.site=site
652
    if cp.has_option('paths','basedir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
653
      self.basepath=cp.get('paths','basedir')
654 655 656
    else:
      self.basepath=os.getcwd()
      print 'No basepath specified, using current directory: %s'%(self.basepath)
657
    mkdirs(self.basepath)
658
    print("Generating LALInference DAG in {0}".format(self.basepath))
659 660
    if dax:
        os.chdir(self.basepath)
661
    self.posteriorpath=os.path.join(self.basepath,'posterior_samples')
John Douglas Veitch's avatar
John Douglas Veitch committed
662
    mkdirs(self.posteriorpath)
663 664 665 666
    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)
667
    if cp.has_option('paths','cachedir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
668
      self.cachepath=cp.get('paths','cachedir')
669 670
    else:
      self.cachepath=os.path.join(self.basepath,'caches')
671
    mkdirs(self.cachepath)
672
    if cp.has_option('paths','logdir'):
John Douglas Veitch's avatar
John Douglas Veitch committed
673
      self.logpath=cp.get('paths','logdir')
674 675
    else:
      self.logpath=os.path.join(self.basepath,'log')
676
    mkdirs(self.logpath)
677
    if cp.has_option('analysis','ifos'):
678
      self.ifos=ast.literal_eval(cp.get('analysis','ifos'))
679 680 681
    else:
      self.ifos=['H1','L1','V1']
    self.segments={}
682 683 684 685
    if cp.has_option('datafind','veto-categories'):
      self.veto_categories=cp.get('datafind','veto-categories')
    else: self.veto_categories=[]
    for ifo in self.ifos:
686
      self.segments[ifo]=[]
687 688 689
    self.computeroqweightsnode={}
    self.bayeslinenode={}
    self.bayeswavepsdnode={}
690
    self.dq={}
691
    self.frtypes=ast.literal_eval(cp.get('datafind','types'))
692
    self.channels=ast.literal_eval(cp.get('data','channels'))
John Douglas Veitch's avatar
John Douglas Veitch committed
693
    self.use_available_data=False
694
    self.webdir=cp.get('paths','webdir')
695 696 697 698
    if cp.has_option('analysis','dataseed'):
      self.dataseed=cp.getint('analysis','dataseed')
    else:
      self.dataseed=None
699
    # Set up necessary job files.
700
    self.prenodes={}
701
    self.datafind_job = pipeline.LSCDataFindJob(self.cachepath,self.logpath,self.config,dax=self.is_dax())
702
    self.datafind_job.add_opt('url-type','file')
703 704 705
    # 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')
706 707 708 709
    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'))
710
    self.datafind_job.set_sub_file(os.path.abspath(os.path.join(self.basepath,'datafind.sub')))
711
    self.preengine_job = EngineJob(self.config, os.path.join(self.basepath,'prelalinference.sub'),self.logpath,engine='lalinferencedatadump',ispreengine=True,dax=self.is_dax())
712 713
    self.preengine_job.set_grid_site('local')
    self.preengine_job.set_universe('vanilla')
714
    if self.config.getboolean('analysis','roq'):
715 716 717 718 719
      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')
720 721 722 723 724
    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')
725 726 727 728
    # Need to create a job file for each IFO combination
    self.engine_jobs={}
    ifocombos=[]
    for N in range(1,len(self.ifos)+1):
729
        for a in permutations(self.ifos,N):
730 731
            ifocombos.append(a)
    for ifos in ifocombos:
732
        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)
733
    self.results_page_job = ResultsPageJob(self.config,os.path.join(self.basepath,'resultspage.sub'),self.logpath,dax=self.is_dax())
734
    self.results_page_job.set_grid_site('local')
735
    self.cotest_results_page_job = ResultsPageJob(self.config,os.path.join(self.basepath,'resultspagecoherent.sub'),self.logpath,dax=self.is_dax())
736
    self.cotest_results_page_job.set_grid_site('local')
737 738 739 740 741 742 743 744
    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')
745
    self.coherence_test_job = CoherenceTestJob(self.config,os.path.join(self.basepath,'coherence_test.sub'),self.logpath,dax=self.is_dax())
746
    self.coherence_test_job.set_grid_site('local')
747
    self.gracedbjob = GraceDBJob(self.config,os.path.join(self.basepath,'gracedb.sub'),self.logpath,dax=self.is_dax())
748
    self.gracedbjob.set_grid_site('local')
749
    # Process the input to build list of analyses to do
John Douglas Veitch's avatar
Working  
John Douglas Veitch committed
750
    self.events=self.setup_from_inputs()
751

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

757
    # Set up the segments
758
    if not (self.config.has_option('input','gps-start-time') and self.config.has_option('input','gps-end-time')) and len(self.times)>0:
759
      (mintime,maxtime)=self.get_required_data(self.times)
760 761 762 763
      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))))
764
    self.add_science_segments()
765

766
    # Save the final configuration that is being used
767
    # first to the run dir
768 769
    conffilename=os.path.join(self.basepath,'config.ini')
    with open(conffilename,'wb') as conffile:
John Douglas Veitch's avatar
John Douglas Veitch committed
770
      self.config.write(conffile)
771 772 773 774
    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)
775

776
    # Generate the DAG according to the config given
777
    for event in self.events: self.add_full_analysis(event)
778
    self.add_skyarea_followup()
779 780 781
    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)
782
    self.dagfilename="lalinference_%s-%s"%(self.config.get('input','gps-start-time'),self.config.get('input','gps-end-time'))
783
    self.set_dag_file(os.path.join(self.basepath,self.dagfilename))
784 785 786 787 788 789 790
    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'):
791 792
      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())
793
      if self.engine=='lalinferenceburst':
794
          prefix='LIB'
795
      else:
796
          prefix='LALInference'
797
      for p in respagenodes:
798
          skyareanode=SkyAreaNode(self.skyareajob,prefix=prefix)
799
          skyareanode.add_resultspage_parent(p)
800
          skyareanode.set_ifos(p.ifos)
801
          self.add_node(skyareanode)
802

803
  def add_full_analysis(self,event):
804
    if self.engine=='lalinferencenest' or  self.engine=='lalinferenceburst':
805
      result=self.add_full_analysis_lalinferencenest(event)
806
    elif self.engine=='lalinferencemcmc':
807
      result=self.add_full_analysis_lalinferencemcmc(event)
808
    elif self.engine=='lalinferencebambi' or self.engine=='lalinferencebambimpi':
809 810 811
      result=self.add_full_analysis_lalinferencebambi(event)
    return result

812 813 814 815 816 817 818 819 820
  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
821

John Douglas Veitch's avatar
John Douglas Veitch committed
822 823
  def get_required_data(self,times):
    """
John Douglas Veitch's avatar
John Douglas Veitch committed
824
    Calculate the data that will be needed to process all events
John Douglas Veitch's avatar
John Douglas Veitch committed
825
    """
826
    #psdlength = self.config.getint('input','max-psd-length')
827
    padding=self.config.getint('input','padding')
828 829
    if self.config.has_option('engine','seglen') or self.config.has_option('lalinference','seglen'):
      if self.config.has_option('engine','seglen'):
830
        seglen = int(np.ceil(self.config.getfloat('engine','seglen')))
831 832 833
      if self.config.has_option('lalinference','seglen'):
        seglen = self.config.getint('lalinference','seglen')

834
      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'):