B
    d/                 @   s  d dl Z d dlZd dlZejdddddddddd	 d dlmZ d dlmZ d dl	Z	d dl
Z
d d	lmZ d d
lmZ d dlmZ d dlmZ ddlmZ dZddlmZ ddlmZ G dd deZd#ddZdd Zdd Zd$ddZe	 dZ!dd Z"d%d!d"Z#dS )&    Ng       @g      $@iX  T)	z	font.sizezaxes.titlesizezaxes.labelsizezxtick.labelsizezytick.labelsizezlegend.fontsizez
figure.dpizsavefig.dpiztext.usetex)figure)FigureCanvasAgg)	lsctables)dbtables)search_summary)segments   )offsetvectorz"Kipp Cannon <kipp.cannon@ligo.org>)date)versionc               @   s&   e Zd Zd	ddZdd Zdd ZdS )
CoincDatabaseexcesspowerNc             C   sP  || _ t|| _ytj| j| _W n tk
r@   d| _Y nX ytj	| j| _
W n tk
rp   d| _
Y nX y4tj| j| _tj| j| _tj| j| _W n& tk
r   d| _d| _d| _Y nX ytj| j| _W n tk
r   d| _Y nX t| j| | _t| j | _|dk	rBt| j| | _ntj | _| jdk	r4y| jj|ddd| _ W n t!k
r   d| _ Y nX y| jj|ddd| _"W n t!k
r   d| _"Y nX y| jj|ddd| _#W n t!k
r   d| _#Y nX y| jj|ddd| _$W n t!k
r0   d| _$Y nX nd| _ d| _"d| _#d| _$dS )z
		Compute and record some summary information about the
		database.  Call this after all the data has been inserted,
		and before you want any of this information.
		Nr   F)Z
create_newr         )%
connectionr   Zget_xmlZxmldocr   ZSnglBurstTableZ	get_tablesngl_burst_table
ValueErrorZSimBurstTablesim_burst_tableZCoincDefTablecoinc_def_tableZ
CoincTableZcoinc_tableTimeSlideTabletime_slide_tableZMultiBurstTableZmulti_burst_tableligolw_search_summaryZ!segmentlistdict_fromsearchsummaryZcoalesceseglistssetkeysZinstrumentsligolw_segmentsZsegmenttable_get_by_namevetoseglistsr   ZsegmentlistdictZget_coinc_def_idbb_definer_idKeyErrorZsb_definer_idsce_definer_idZscn_definer_id)selfr   Zlive_time_programsearchZveto_segments_name r"   d/work/yifan.wang/ringdown/master-ringdown-env/lib/python3.7/site-packages/lalburst/SnglBurstUtils.py__init__D   s`    
zCoincDatabase.__init__c             #   s    j  }xp j  d jfD ]V\}} fdd|d||fD }tdd |D }t| dd |D |fV  q"W |  dS )	z
		Generator function to return

			is_background, event_list, offsetvector

		tuples by querying the coinc_event and sngl_burst tables in
		the database.  Only coincs corresponding to
		sngl_burst<-->sngl_burst coincs will be retrieved.
		z\
	SELECT
		coinc_event_id,
		time_slide_id
	FROM
		coinc_event
	WHERE
		coinc_def_id == ?
		c                s    g | ]} j ||d  fqS ))r   row_from_cols).0row)r    r"   r#   
<listcomp>   s    z3CoincDatabase.get_noninjections.<locals>.<listcomp>a_  
	SELECT
		sngl_burst.*,
		time_slide.offset
	FROM
		coinc_event_map
		JOIN sngl_burst ON (
			coinc_event_map.table_name == 'sngl_burst'
			AND sngl_burst.event_id == coinc_event_map.event_id
		)
		JOIN time_slide ON (
			time_slide.instrument == sngl_burst.ifo
		)
	WHERE
		coinc_event_map.coinc_event_id == ?
		AND time_slide.time_slide_id == ?
			c             s   s   | ]\}}|j |fV  qd S )N)ifo)r'   eventoffsetr"   r"   r#   	<genexpr>   s    z2CoincDatabase.get_noninjections.<locals>.<genexpr>c             S   s   g | ]\}}|qS r"   r"   )r'   r+   r,   r"   r"   r#   r)      s    N)r   cursorexecuter   r	   anyvaluesclose)r    r.   coinc_event_idtime_slide_idrowsoffsetsr"   )r    r#   get_noninjections   s    

"zCoincDatabase.get_noninjectionsc             #   s    j  }xr j  d jfD ]X} j|}|d } fdd|d|fD }|dd |D tdd |D fV  q"W |  d	S )
z
		Generator function to return

			sim, event_list, offsetvector

		tuples by querying the sim_burst, coinc_event and
		sngl_burst tables in the database.  Only coincs
		corresponding to "exact" sim_burst<-->coinc_event coincs
		will be retrieved.
		aV  
	SELECT
		sim_burst.*,
		burst_coinc_event_map.event_id
	FROM
		sim_burst
		JOIN coinc_event_map AS sim_coinc_event_map ON (
			sim_coinc_event_map.table_name == 'sim_burst'
			AND sim_coinc_event_map.event_id == sim_burst.simulation_id
		)
		JOIN coinc_event AS sim_coinc_event ON (
			sim_coinc_event.coinc_event_id == sim_coinc_event_map.coinc_event_id
		)
		JOIN coinc_event_map AS burst_coinc_event_map ON (
			burst_coinc_event_map.coinc_event_id == sim_coinc_event_map.coinc_event_id
			AND burst_coinc_event_map.table_name == 'coinc_event'
		)
	WHERE
		sim_coinc_event.coinc_def_id == ?
		r%   c                s    g | ]} j ||d  fqS )r%   )r   r&   )r'   r(   )r    r"   r#   r)      s    z0CoincDatabase.get_injections.<locals>.<listcomp>a  
	SELECT
		sngl_burst.*,
		time_slide.offset
	FROM
		sngl_burst
		JOIN coinc_event_map ON (
			coinc_event_map.table_name == 'sngl_burst'
			AND coinc_event_map.event_id == sngl_burst.event_id
		)
		JOIN coinc_event ON (
			coinc_event.coinc_event_id == coinc_event_map.coinc_event_id
		)
		JOIN time_slide ON (
			coinc_event.time_slide_id == time_slide.time_slide_id
			AND time_slide.instrument == sngl_burst.ifo
		)
	WHERE
		coinc_event.coinc_event_id == ?
			c             S   s   g | ]\}}|qS r"   r"   )r'   r+   r,   r"   r"   r#   r)      s    c             s   s   | ]\}}|j |fV  qd S )N)r*   )r'   r+   r,   r"   r"   r#   r-      s    z/CoincDatabase.get_injections.<locals>.<genexpr>N)r   r.   r/   r   r   r&   r	   r2   )r    r.   r1   simr3   r5   r"   )r    r#   get_injections   s    
(zCoincDatabase.get_injections)r   N)__name__
__module____qualname__r$   r7   r9   r"   r"   r"   r#   r   C   s   
B*r   c             C   sb  |d krd}nd| }| j  }td| tjd xht| j D ]V\}}td||t|dt	t|| j
krv|| j
| @ pxd t	t| f tjd qBW | jd k	rtd|t| jf tjd | jd k	rtd	|t| jf tjd | jd k	rtd
||d d f tjd | jd k	rVx0|dD ]"\}}td|||f tjd q0W |  d S )N z%s: z%sdatabase stats:)filez"	%s%s livetime: %g s (%g%% vetoed)g      Y@g        z	%sburst events: %dz	%sburst injections: %dz	%stime slides: %dz5SELECT COUNT(DISTINCT(time_slide_id)) FROM time_slider   zsSELECT description, COUNT(*) FROM coinc_definer NATURAL JOIN coinc_event GROUP BY coinc_def_id ORDER BY descriptionz		%s%s: %d)r   r.   printsysstderrsortedr   itemsabsfloatr   r   lenr   r   r/   Zfetchoner   r2   )contentsfilenamer.   Z
instrumentZseglistdescriptionnr"   r"   r#   summarize_coinc_database   s"    
R

$rK   c             c   s0   x*| j  d|fD ]}| j|V  qW d S )Nz
SELECT sngl_burst.* FROM
	sngl_burst
	JOIN coinc_event_map ON (
		sngl_burst.event_id == coinc_event_map.event_id
		AND coinc_event_map.table_name == 'sngl_burst'
	)
WHERE
	coinc_event_map.coinc_event_id == ?
	)r   r.   r/   r   r&   )rG   r3   r1   r"   r"   r#   coinc_sngl_bursts  s    	rL   c             C   sD   t j| d }tdd | D }tdd | D }||fS )z
	Query the database for the IDs and offsets of all time slides, and
	return two dictionaries one containing the all-zero time slides and
	the other containing the not-all-zero time slides.
	)r   c             s   s&   | ]\}}t | s||fV  qd S )N)r0   r1   )r'   r4   r	   r"   r"   r#   r-   #  s    z"get_time_slides.<locals>.<genexpr>c             s   s&   | ]\}}t | r||fV  qd S )N)r0   r1   )r'   r4   r	   r"   r"   r#   r-   $  s    )r   r   as_dictdictrC   )r   time_slidesZzero_lag_time_slidesZbackground_time_slidesr"   r"   r#   get_time_slides  s    rP   Fc             C   s   d}| j  }t|}|r,td| tjd xZt|D ]N\}}|r^tdd| |  dtjd | j | |tt	| 
| 7 }q6W | j | |rtdtjd |S )	z
	Given a sequence of time slides (each of which is an instrument -->
	offset dictionary), use the segmentlistdict dictionary of segment
	lists to compute the live time in each time slide.  Return the sum
	of the live times from all slides.
	g        z+computing the live time for %d time slides:)r>   z	%.1f%%g      Y@ )endr>   z	100.0%)r6   copyrF   r?   r@   rA   	enumerateupdaterE   rD   intersectionr   )r   rO   verboseZlivetimeZold_offsetsNrJ   Z
time_slider"   r"   r#   time_slides_livetime)  s    
rY   z([+-]?[.0-9]+)[Ee]([+-]?[0-9]+)c             C   s"   t |  \}}d|t|f S )z
	Convert a string of the form "d.dddde-dd" to "d.dddd 	imes
	10^{-dd}"

	Example:

	>>> import math
	>>> print latexnumber("%.16e" % math.pi)
	3.1415926535897931 \times 10^{0}
	z%s \times 10^{%d})floatpatternmatchgroupsint)smer"   r"   r#   latexnumberK  s    ra        d@c             C   sd   t  }t| ||d |d dtd d   | }|d ||  |	| ||fS )z
	width is in mm
	gffffff9@r      r   T)
r   ZFigureFigureCanvasZset_size_inchesmathsqrtZgcagridZ
set_xlabelZ
set_ylabel)Zx_labelZy_labelwidthZfigZaxesr"   r"   r#   make_burst_plotc  s    &


ri   )N)F)rb   )$	itertoolsre   Z
matplotlibZrcParamsrU   r   Zmatplotlib.backends.backend_aggr   rd   rer@   Zligo.lwr   r   Zligo.lw.utilsr   r   r   r   r	   
__author__Zgit_versionr
   __date__r   __version__objectr   rK   rL   rP   rY   compilerZ   ra   ri   r"   r"   r"   r#   <module>   sB   
 -

 
