https://github.com/gwastro/pycbc
Tip revision: 2f773a01d4d91ec521b458914e8e711a5e3e2e1f authored by Ian Harry on 06 January 2022, 15:40:15 UTC
Update setup.py (#3908)
Update setup.py (#3908)
Tip revision: 2f773a0
matchedfilter.py
# Copyright (C) 2012 Alex Nitz
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# =============================================================================
#
# Preamble
#
# =============================================================================
#
"""
This modules provides functions for matched filtering along with associated
utilities.
"""
import logging
from math import sqrt
from pycbc.types import TimeSeries, FrequencySeries, zeros, Array
from pycbc.types import complex_same_precision_as, real_same_precision_as
from pycbc.fft import fft, ifft, IFFT
import pycbc.scheme
from pycbc import events
from pycbc.events import ranking
import pycbc
import numpy
BACKEND_PREFIX="pycbc.filter.matchedfilter_"
@pycbc.scheme.schemed(BACKEND_PREFIX)
def correlate(x, y, z):
err_msg = "This function is a stub that should be overridden using the "
err_msg += "scheme. You shouldn't be seeing this error!"
raise ValueError(err_msg)
class BatchCorrelator(object):
""" Create a batch correlation engine
"""
def __init__(self, xs, zs, size):
""" Correlate x and y, store in z. Arrays need not be equal length, but
must be at least size long and of the same dtype. No error checking
will be performed, so be careful. All dtypes must be complex64.
Note, must be created within the processing context that it will be used in.
"""
self.size = int(size)
self.dtype = xs[0].dtype
self.num_vectors = len(xs)
# keep reference to arrays
self.xs = xs
self.zs = zs
# Store each pointer as in integer array
self.x = Array([v.ptr for v in xs], dtype=int)
self.z = Array([v.ptr for v in zs], dtype=int)
@pycbc.scheme.schemed(BACKEND_PREFIX)
def batch_correlate_execute(self, y):
pass
execute = batch_correlate_execute
@pycbc.scheme.schemed(BACKEND_PREFIX)
def _correlate_factory(x, y, z):
err_msg = "This class is a stub that should be overridden using the "
err_msg += "scheme. You shouldn't be seeing this error!"
raise ValueError(err_msg)
class Correlator(object):
""" Create a correlator engine
Parameters
---------
x : complex64
Input pycbc.types.Array (or subclass); it will be conjugated
y : complex64
Input pycbc.types.Array (or subclass); it will not be conjugated
z : complex64
Output pycbc.types.Array (or subclass).
It will contain conj(x) * y, element by element
The addresses in memory of the data of all three parameter vectors
must be the same modulo pycbc.PYCBC_ALIGNMENT
"""
def __new__(cls, *args, **kwargs):
real_cls = _correlate_factory(*args, **kwargs)
return real_cls(*args, **kwargs) # pylint:disable=not-callable
# The class below should serve as the parent for all schemed classes.
# The intention is that this class serves simply as the location for
# all documentation of the class and its methods, though that is not
# yet implemented. Perhaps something along the lines of:
#
# http://stackoverflow.com/questions/2025562/inherit-docstrings-in-python-class-inheritance
#
# will work? Is there a better way?
class _BaseCorrelator(object):
def correlate(self):
"""
Compute the correlation of the vectors specified at object
instantiation, writing into the output vector given when the
object was instantiated. The intention is that this method
should be called many times, with the contents of those vectors
changing between invocations, but not their locations in memory
or length.
"""
pass
class MatchedFilterControl(object):
def __init__(self, low_frequency_cutoff, high_frequency_cutoff, snr_threshold, tlen,
delta_f, dtype, segment_list, template_output, use_cluster,
downsample_factor=1, upsample_threshold=1, upsample_method='pruned_fft',
gpu_callback_method='none', cluster_function='symmetric'):
""" Create a matched filter engine.
Parameters
----------
low_frequency_cutoff : {None, float}, optional
The frequency to begin the filter calculation. If None, begin at the
first frequency after DC.
high_frequency_cutoff : {None, float}, optional
The frequency to stop the filter calculation. If None, continue to the
the nyquist frequency.
snr_threshold : float
The minimum snr to return when filtering
segment_list : list
List of FrequencySeries that are the Fourier-transformed data segments
template_output : complex64
Array of memory given as the 'out' parameter to waveform.FilterBank
use_cluster : boolean
If true, cluster triggers above threshold using a window; otherwise,
only apply a threshold.
downsample_factor : {1, int}, optional
The factor by which to reduce the sample rate when doing a heirarchical
matched filter
upsample_threshold : {1, float}, optional
The fraction of the snr_threshold to trigger on the subsampled filter.
upsample_method : {pruned_fft, str}
The method to upsample or interpolate the reduced rate filter.
cluster_function : {symmetric, str}, optional
Which method is used to cluster triggers over time. If 'findchirp', a
sliding forward window; if 'symmetric', each window's peak is compared
to the windows before and after it, and only kept as a trigger if larger
than both.
"""
# Assuming analysis time is constant across templates and segments, also
# delta_f is constant across segments.
self.tlen = tlen
self.flen = self.tlen / 2 + 1
self.delta_f = delta_f
self.delta_t = 1.0/(self.delta_f * self.tlen)
self.dtype = dtype
self.snr_threshold = snr_threshold
self.flow = low_frequency_cutoff
self.fhigh = high_frequency_cutoff
self.gpu_callback_method = gpu_callback_method
if cluster_function not in ['symmetric', 'findchirp']:
raise ValueError("MatchedFilter: 'cluster_function' must be either 'symmetric' or 'findchirp'")
self.cluster_function = cluster_function
self.segments = segment_list
self.htilde = template_output
if downsample_factor == 1:
self.snr_mem = zeros(self.tlen, dtype=self.dtype)
self.corr_mem = zeros(self.tlen, dtype=self.dtype)
if use_cluster and (cluster_function == 'symmetric'):
self.matched_filter_and_cluster = self.full_matched_filter_and_cluster_symm
# setup the threasholding/clustering operations for each segment
self.threshold_and_clusterers = []
for seg in self.segments:
thresh = events.ThresholdCluster(self.snr_mem[seg.analyze])
self.threshold_and_clusterers.append(thresh)
elif use_cluster and (cluster_function == 'findchirp'):
self.matched_filter_and_cluster = self.full_matched_filter_and_cluster_fc
else:
self.matched_filter_and_cluster = self.full_matched_filter_thresh_only
# Assuming analysis time is constant across templates and segments, also
# delta_f is constant across segments.
self.kmin, self.kmax = get_cutoff_indices(self.flow, self.fhigh,
self.delta_f, self.tlen)
# Set up the correlation operations for each analysis segment
corr_slice = slice(self.kmin, self.kmax)
self.correlators = []
for seg in self.segments:
corr = Correlator(self.htilde[corr_slice],
seg[corr_slice],
self.corr_mem[corr_slice])
self.correlators.append(corr)
# setup up the ifft we will do
self.ifft = IFFT(self.corr_mem, self.snr_mem)
elif downsample_factor >= 1:
self.matched_filter_and_cluster = self.heirarchical_matched_filter_and_cluster
self.downsample_factor = downsample_factor
self.upsample_method = upsample_method
self.upsample_threshold = upsample_threshold
N_full = self.tlen
N_red = N_full / downsample_factor
self.kmin_full, self.kmax_full = get_cutoff_indices(self.flow,
self.fhigh, self.delta_f, N_full)
self.kmin_red, _ = get_cutoff_indices(self.flow,
self.fhigh, self.delta_f, N_red)
if self.kmax_full < N_red:
self.kmax_red = self.kmax_full
else:
self.kmax_red = N_red - 1
self.snr_mem = zeros(N_red, dtype=self.dtype)
self.corr_mem_full = FrequencySeries(zeros(N_full, dtype=self.dtype), delta_f=self.delta_f)
self.corr_mem = Array(self.corr_mem_full[0:N_red], copy=False)
self.inter_vec = zeros(N_full, dtype=self.dtype)
else:
raise ValueError("Invalid downsample factor")
def full_matched_filter_and_cluster_symm(self, segnum, template_norm, window, epoch=None):
""" Returns the complex snr timeseries, normalization of the complex snr,
the correlation vector frequency series, the list of indices of the
triggers, and the snr values at the trigger locations. Returns empty
lists for these for points that are not above the threshold.
Calculated the matched filter, threshold, and cluster.
Parameters
----------
segnum : int
Index into the list of segments at MatchedFilterControl construction
against which to filter.
template_norm : float
The htilde, template normalization factor.
window : int
Size of the window over which to cluster triggers, in samples
Returns
-------
snr : TimeSeries
A time series containing the complex snr.
norm : float
The normalization of the complex snr.
corrrelation: FrequencySeries
A frequency series containing the correlation vector.
idx : Array
List of indices of the triggers.
snrv : Array
The snr values at the trigger locations.
"""
norm = (4.0 * self.delta_f) / sqrt(template_norm)
self.correlators[segnum].correlate()
self.ifft.execute()
snrv, idx = self.threshold_and_clusterers[segnum].threshold_and_cluster(self.snr_threshold / norm, window)
if len(idx) == 0:
return [], [], [], [], []
logging.info("%s points above threshold" % str(len(idx)))
snr = TimeSeries(self.snr_mem, epoch=epoch, delta_t=self.delta_t, copy=False)
corr = FrequencySeries(self.corr_mem, delta_f=self.delta_f, copy=False)
return snr, norm, corr, idx, snrv
def full_matched_filter_and_cluster_fc(self, segnum, template_norm, window, epoch=None):
""" Returns the complex snr timeseries, normalization of the complex snr,
the correlation vector frequency series, the list of indices of the
triggers, and the snr values at the trigger locations. Returns empty
lists for these for points that are not above the threshold.
Calculated the matched filter, threshold, and cluster.
Parameters
----------
segnum : int
Index into the list of segments at MatchedFilterControl construction
against which to filter.
template_norm : float
The htilde, template normalization factor.
window : int
Size of the window over which to cluster triggers, in samples
Returns
-------
snr : TimeSeries
A time series containing the complex snr.
norm : float
The normalization of the complex snr.
corrrelation: FrequencySeries
A frequency series containing the correlation vector.
idx : Array
List of indices of the triggers.
snrv : Array
The snr values at the trigger locations.
"""
norm = (4.0 * self.delta_f) / sqrt(template_norm)
self.correlators[segnum].correlate()
self.ifft.execute()
idx, snrv = events.threshold(self.snr_mem[self.segments[segnum].analyze],
self.snr_threshold / norm)
idx, snrv = events.cluster_reduce(idx, snrv, window)
if len(idx) == 0:
return [], [], [], [], []
logging.info("%s points above threshold" % str(len(idx)))
snr = TimeSeries(self.snr_mem, epoch=epoch, delta_t=self.delta_t, copy=False)
corr = FrequencySeries(self.corr_mem, delta_f=self.delta_f, copy=False)
return snr, norm, corr, idx, snrv
def full_matched_filter_thresh_only(self, segnum, template_norm, window=None, epoch=None):
""" Returns the complex snr timeseries, normalization of the complex snr,
the correlation vector frequency series, the list of indices of the
triggers, and the snr values at the trigger locations. Returns empty
lists for these for points that are not above the threshold.
Calculated the matched filter, threshold, and cluster.
Parameters
----------
segnum : int
Index into the list of segments at MatchedFilterControl construction
against which to filter.
template_norm : float
The htilde, template normalization factor.
window : int
Size of the window over which to cluster triggers, in samples.
This is IGNORED by this function, and provided only for API compatibility.
Returns
-------
snr : TimeSeries
A time series containing the complex snr.
norm : float
The normalization of the complex snr.
corrrelation: FrequencySeries
A frequency series containing the correlation vector.
idx : Array
List of indices of the triggers.
snrv : Array
The snr values at the trigger locations.
"""
norm = (4.0 * self.delta_f) / sqrt(template_norm)
self.correlators[segnum].correlate()
self.ifft.execute()
idx, snrv = events.threshold_only(self.snr_mem[self.segments[segnum].analyze],
self.snr_threshold / norm)
logging.info("%s points above threshold" % str(len(idx)))
snr = TimeSeries(self.snr_mem, epoch=epoch, delta_t=self.delta_t, copy=False)
corr = FrequencySeries(self.corr_mem, delta_f=self.delta_f, copy=False)
return snr, norm, corr, idx, snrv
def heirarchical_matched_filter_and_cluster(self, segnum, template_norm, window):
""" Returns the complex snr timeseries, normalization of the complex snr,
the correlation vector frequency series, the list of indices of the
triggers, and the snr values at the trigger locations. Returns empty
lists for these for points that are not above the threshold.
Calculated the matched filter, threshold, and cluster.
Parameters
----------
segnum : int
Index into the list of segments at MatchedFilterControl construction
template_norm : float
The htilde, template normalization factor.
window : int
Size of the window over which to cluster triggers, in samples
Returns
-------
snr : TimeSeries
A time series containing the complex snr at the reduced sample rate.
norm : float
The normalization of the complex snr.
corrrelation: FrequencySeries
A frequency series containing the correlation vector.
idx : Array
List of indices of the triggers.
snrv : Array
The snr values at the trigger locations.
"""
from pycbc.fft.fftw_pruned import pruned_c2cifft, fft_transpose
htilde = self.htilde
stilde = self.segments[segnum]
norm = (4.0 * stilde.delta_f) / sqrt(template_norm)
correlate(htilde[self.kmin_red:self.kmax_red],
stilde[self.kmin_red:self.kmax_red],
self.corr_mem[self.kmin_red:self.kmax_red])
ifft(self.corr_mem, self.snr_mem)
if not hasattr(stilde, 'red_analyze'):
stilde.red_analyze = \
slice(stilde.analyze.start/self.downsample_factor,
stilde.analyze.stop/self.downsample_factor)
idx_red, snrv_red = events.threshold(self.snr_mem[stilde.red_analyze],
self.snr_threshold / norm * self.upsample_threshold)
if len(idx_red) == 0:
return [], None, [], [], []
idx_red, _ = events.cluster_reduce(idx_red, snrv_red, window / self.downsample_factor)
logging.info("%s points above threshold at reduced resolution"\
%(str(len(idx_red)),))
# The fancy upsampling is here
if self.upsample_method=='pruned_fft':
idx = (idx_red + stilde.analyze.start/self.downsample_factor)\
* self.downsample_factor
idx = smear(idx, self.downsample_factor)
# cache transposed versions of htilde and stilde
if not hasattr(self.corr_mem_full, 'transposed'):
self.corr_mem_full.transposed = zeros(len(self.corr_mem_full), dtype=self.dtype)
if not hasattr(htilde, 'transposed'):
htilde.transposed = zeros(len(self.corr_mem_full), dtype=self.dtype)
htilde.transposed[self.kmin_full:self.kmax_full] = htilde[self.kmin_full:self.kmax_full]
htilde.transposed = fft_transpose(htilde.transposed)
if not hasattr(stilde, 'transposed'):
stilde.transposed = zeros(len(self.corr_mem_full), dtype=self.dtype)
stilde.transposed[self.kmin_full:self.kmax_full] = stilde[self.kmin_full:self.kmax_full]
stilde.transposed = fft_transpose(stilde.transposed)
correlate(htilde.transposed, stilde.transposed, self.corr_mem_full.transposed)
snrv = pruned_c2cifft(self.corr_mem_full.transposed, self.inter_vec, idx, pretransposed=True)
idx = idx - stilde.analyze.start
idx2, snrv = events.threshold(Array(snrv, copy=False), self.snr_threshold / norm)
if len(idx2) > 0:
correlate(htilde[self.kmax_red:self.kmax_full],
stilde[self.kmax_red:self.kmax_full],
self.corr_mem_full[self.kmax_red:self.kmax_full])
idx, snrv = events.cluster_reduce(idx[idx2], snrv, window)
else:
idx, snrv = [], []
logging.info("%s points at full rate and clustering" % len(idx))
return self.snr_mem, norm, self.corr_mem_full, idx, snrv
else:
raise ValueError("Invalid upsample method")
def compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr,
hpnorm=None, hcnorm=None,
out=None, thresh=0,
analyse_slice=None):
"""
Matched filter maximised over polarization and orbital phase.
This implements the statistic derived in 1603.02444. It is encouraged
to read that work to understand the limitations and assumptions implicit
in this statistic before using it.
Parameters
-----------
hplus : TimeSeries
This is the IFFTed complex SNR time series of (h+, data). If not
normalized, supply the normalization factor so this can be done!
It is recommended to normalize this before sending through this
function
hcross : TimeSeries
This is the IFFTed complex SNR time series of (hx, data). If not
normalized, supply the normalization factor so this can be done!
hphccorr : float
The real component of the overlap between the two polarizations
Re[(h+, hx)]. Note that the imaginary component does not enter the
detection statistic. This must be normalized and is sign-sensitive.
thresh : float
Used for optimization. If we do not care about the value of SNR
values below thresh we can calculate a quick statistic that will
always overestimate SNR and then only calculate the proper, more
expensive, statistic at points where the quick SNR is above thresh.
hpsigmasq : float
The normalization factor (h+, h+). Default = None (=1, already
normalized)
hcsigmasq : float
The normalization factor (hx, hx). Default = None (=1, already
normalized)
out : TimeSeries (optional, default=None)
If given, use this array to store the output.
Returns
--------
det_stat : TimeSeries
The SNR maximized over sky location
"""
# NOTE: Not much optimization has been done here! This may need to be
# Cythonized.
if out is None:
out = zeros(len(hplus))
out.non_zero_locs = numpy.array([], dtype=out.dtype)
else:
if not hasattr(out, 'non_zero_locs'):
# Doing this every time is not a zero-cost operation
out.data[:] = 0
out.non_zero_locs = numpy.array([], dtype=out.dtype)
else:
# Only set non zero locations to zero
out.data[out.non_zero_locs] = 0
# If threshold is given we can limit the points at which to compute the
# full statistic
if thresh:
# This is the statistic that always overestimates the SNR...
# It allows some unphysical freedom that the full statistic does not
idx_p, _ = events.threshold_only(hplus[analyse_slice],
thresh / (2**0.5 * hpnorm))
idx_c, _ = events.threshold_only(hcross[analyse_slice],
thresh / (2**0.5 * hcnorm))
idx_p = idx_p + analyse_slice.start
idx_c = idx_c + analyse_slice.start
hp_red = hplus[idx_p] * hpnorm
hc_red = hcross[idx_p] * hcnorm
stat_p = hp_red.real**2 + hp_red.imag**2 + \
hc_red.real**2 + hc_red.imag**2
locs_p = idx_p[stat_p > (thresh*thresh)]
hp_red = hplus[idx_c] * hpnorm
hc_red = hcross[idx_c] * hcnorm
stat_c = hp_red.real**2 + hp_red.imag**2 + \
hc_red.real**2 + hc_red.imag**2
locs_c = idx_c[stat_c > (thresh*thresh)]
locs = numpy.unique(numpy.concatenate((locs_p, locs_c)))
hplus = hplus[locs]
hcross = hcross[locs]
hplus = hplus * hpnorm
hcross = hcross * hcnorm
# Calculate and sanity check the denominator
denom = 1 - hphccorr*hphccorr
if denom < 0:
if hphccorr > 1:
err_msg = "Overlap between hp and hc is given as %f. " %(hphccorr)
err_msg += "How can an overlap be bigger than 1?"
raise ValueError(err_msg)
else:
err_msg = "There really is no way to raise this error!?! "
err_msg += "If you're seeing this, it is bad."
raise ValueError(err_msg)
if denom == 0:
# This case, of hphccorr==1, makes the statistic degenerate
# This case should not physically be possible luckily.
err_msg = "You have supplied a real overlap between hp and hc of 1. "
err_msg += "Ian is reasonably certain this is physically impossible "
err_msg += "so why are you seeing this?"
raise ValueError(err_msg)
assert(len(hplus) == len(hcross))
# Now the stuff where comp. cost may be a problem
hplus_magsq = numpy.real(hplus) * numpy.real(hplus) + \
numpy.imag(hplus) * numpy.imag(hplus)
hcross_magsq = numpy.real(hcross) * numpy.real(hcross) + \
numpy.imag(hcross) * numpy.imag(hcross)
rho_pluscross = numpy.real(hplus) * numpy.real(hcross) + numpy.imag(hplus)*numpy.imag(hcross)
sqroot = (hplus_magsq - hcross_magsq)**2
sqroot += 4 * (hphccorr * hplus_magsq - rho_pluscross) * \
(hphccorr * hcross_magsq - rho_pluscross)
# Sometimes this can be less than 0 due to numeric imprecision, catch this.
if (sqroot < 0).any():
indices = numpy.arange(len(sqroot))[sqroot < 0]
# This should not be *much* smaller than 0 due to numeric imprecision
if (sqroot[indices] < -0.0001).any():
err_msg = "Square root has become negative. Something wrong here!"
raise ValueError(err_msg)
sqroot[indices] = 0
sqroot = numpy.sqrt(sqroot)
det_stat_sq = 0.5 * (hplus_magsq + hcross_magsq - \
2 * rho_pluscross*hphccorr + sqroot) / denom
det_stat = numpy.sqrt(det_stat_sq)
if thresh:
out.data[locs] = det_stat
out.non_zero_locs = locs
return out
else:
return Array(det_stat, copy=False)
def compute_u_val_for_sky_loc_stat(hplus, hcross, hphccorr,
hpnorm=None, hcnorm=None, indices=None):
"""The max-over-sky location detection statistic maximizes over a phase,
an amplitude and the ratio of F+ and Fx, encoded in a variable called u.
Here we return the value of u for the given indices.
"""
if indices is not None:
hplus = hplus[indices]
hcross = hcross[indices]
if hpnorm is not None:
hplus = hplus * hpnorm
if hcnorm is not None:
hcross = hcross * hcnorm
# Sanity checking in func. above should already have identified any points
# which are bad, and should be used to construct indices for input here
hplus_magsq = numpy.real(hplus) * numpy.real(hplus) + \
numpy.imag(hplus) * numpy.imag(hplus)
hcross_magsq = numpy.real(hcross) * numpy.real(hcross) + \
numpy.imag(hcross) * numpy.imag(hcross)
rho_pluscross = numpy.real(hplus) * numpy.real(hcross) + \
numpy.imag(hplus)*numpy.imag(hcross)
a = hphccorr * hplus_magsq - rho_pluscross
b = hplus_magsq - hcross_magsq
c = rho_pluscross - hphccorr * hcross_magsq
sq_root = b*b - 4*a*c
sq_root = sq_root**0.5
sq_root = -sq_root
# Catch the a->0 case
bad_lgc = (a == 0)
dbl_bad_lgc = numpy.logical_and(c == 0, b == 0)
dbl_bad_lgc = numpy.logical_and(bad_lgc, dbl_bad_lgc)
# Initialize u
u = sq_root * 0.
# In this case u is completely degenerate, so set it to 1
u[dbl_bad_lgc] = 1.
# If a->0 avoid overflow by just setting to a large value
u[bad_lgc & ~dbl_bad_lgc] = 1E17
# Otherwise normal statistic
u[~bad_lgc] = (-b[~bad_lgc] + sq_root[~bad_lgc]) / (2*a[~bad_lgc])
snr_cplx = hplus * u + hcross
coa_phase = numpy.angle(snr_cplx)
return u, coa_phase
def compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross, hphccorr,
hpnorm=None, hcnorm=None,
out=None, thresh=0,
analyse_slice=None):
"""
Matched filter maximised over polarization phase.
This implements the statistic derived in 1709.09181. It is encouraged
to read that work to understand the limitations and assumptions implicit
in this statistic before using it.
In contrast to compute_max_snr_over_sky_loc_stat this function
performs no maximization over orbital phase, treating that as an intrinsic
parameter. In the case of aligned-spin 2,2-mode only waveforms, this
collapses to the normal statistic (at twice the computational cost!)
Parameters
-----------
hplus : TimeSeries
This is the IFFTed complex SNR time series of (h+, data). If not
normalized, supply the normalization factor so this can be done!
It is recommended to normalize this before sending through this
function
hcross : TimeSeries
This is the IFFTed complex SNR time series of (hx, data). If not
normalized, supply the normalization factor so this can be done!
hphccorr : float
The real component of the overlap between the two polarizations
Re[(h+, hx)]. Note that the imaginary component does not enter the
detection statistic. This must be normalized and is sign-sensitive.
thresh : float
Used for optimization. If we do not care about the value of SNR
values below thresh we can calculate a quick statistic that will
always overestimate SNR and then only calculate the proper, more
expensive, statistic at points where the quick SNR is above thresh.
hpsigmasq : float
The normalization factor (h+, h+). Default = None (=1, already
normalized)
hcsigmasq : float
The normalization factor (hx, hx). Default = None (=1, already
normalized)
out : TimeSeries (optional, default=None)
If given, use this array to store the output.
Returns
--------
det_stat : TimeSeries
The SNR maximized over sky location
"""
# NOTE: Not much optimization has been done here! This may need to be
# Cythonized.
if out is None:
out = zeros(len(hplus))
out.non_zero_locs = numpy.array([], dtype=out.dtype)
else:
if not hasattr(out, 'non_zero_locs'):
# Doing this every time is not a zero-cost operation
out.data[:] = 0
out.non_zero_locs = numpy.array([], dtype=out.dtype)
else:
# Only set non zero locations to zero
out.data[out.non_zero_locs] = 0
# If threshold is given we can limit the points at which to compute the
# full statistic
if thresh:
# This is the statistic that always overestimates the SNR...
# It allows some unphysical freedom that the full statistic does not
#
# For now this is copied from the max-over-phase statistic. One could
# probably make this faster by removing the imaginary components of
# the matched filter, as these are not used here.
idx_p, _ = events.threshold_only(hplus[analyse_slice],
thresh / (2**0.5 * hpnorm))
idx_c, _ = events.threshold_only(hcross[analyse_slice],
thresh / (2**0.5 * hcnorm))
idx_p = idx_p + analyse_slice.start
idx_c = idx_c + analyse_slice.start
hp_red = hplus[idx_p] * hpnorm
hc_red = hcross[idx_p] * hcnorm
stat_p = hp_red.real**2 + hp_red.imag**2 + \
hc_red.real**2 + hc_red.imag**2
locs_p = idx_p[stat_p > (thresh*thresh)]
hp_red = hplus[idx_c] * hpnorm
hc_red = hcross[idx_c] * hcnorm
stat_c = hp_red.real**2 + hp_red.imag**2 + \
hc_red.real**2 + hc_red.imag**2
locs_c = idx_c[stat_c > (thresh*thresh)]
locs = numpy.unique(numpy.concatenate((locs_p, locs_c)))
hplus = hplus[locs]
hcross = hcross[locs]
hplus = hplus * hpnorm
hcross = hcross * hcnorm
# Calculate and sanity check the denominator
denom = 1 - hphccorr*hphccorr
if denom < 0:
if hphccorr > 1:
err_msg = "Overlap between hp and hc is given as %f. " %(hphccorr)
err_msg += "How can an overlap be bigger than 1?"
raise ValueError(err_msg)
else:
err_msg = "There really is no way to raise this error!?! "
err_msg += "If you're seeing this, it is bad."
raise ValueError(err_msg)
if denom == 0:
# This case, of hphccorr==1, makes the statistic degenerate
# This case should not physically be possible luckily.
err_msg = "You have supplied a real overlap between hp and hc of 1. "
err_msg += "Ian is reasonably certain this is physically impossible "
err_msg += "so why are you seeing this?"
raise ValueError(err_msg)
assert(len(hplus) == len(hcross))
# Now the stuff where comp. cost may be a problem
hplus_magsq = numpy.real(hplus) * numpy.real(hplus)
hcross_magsq = numpy.real(hcross) * numpy.real(hcross)
rho_pluscross = numpy.real(hplus) * numpy.real(hcross)
det_stat_sq = (hplus_magsq + hcross_magsq - 2 * rho_pluscross*hphccorr)
det_stat = numpy.sqrt(det_stat_sq / denom)
if thresh:
out.data[locs] = det_stat
out.non_zero_locs = locs
return out
else:
return Array(det_stat, copy=False)
def compute_u_val_for_sky_loc_stat_no_phase(hplus, hcross, hphccorr,
hpnorm=None , hcnorm=None, indices=None):
"""The max-over-sky location (no phase) detection statistic maximizes over
an amplitude and the ratio of F+ and Fx, encoded in a variable called u.
Here we return the value of u for the given indices.
"""
if indices is not None:
hplus = hplus[indices]
hcross = hcross[indices]
if hpnorm is not None:
hplus = hplus * hpnorm
if hcnorm is not None:
hcross = hcross * hcnorm
rhoplusre=numpy.real(hplus)
rhocrossre=numpy.real(hcross)
overlap=numpy.real(hphccorr)
denom = (-rhocrossre+overlap*rhoplusre)
# Initialize tan_kappa array
u_val = denom * 0.
# Catch the denominator -> 0 case
numpy.putmask(u_val, denom == 0, 1E17)
# Otherwise do normal statistic
numpy.putmask(u_val, denom != 0, (-rhoplusre+overlap*rhocrossre)/(-rhocrossre+overlap*rhoplusre))
coa_phase = numpy.zeros(len(indices), dtype=numpy.float32)
return u_val, coa_phase
class MatchedFilterSkyMaxControl(object):
# FIXME: This seems much more simplistic than the aligned-spin class.
# E.g. no correlators. Is this worth updating?
def __init__(self, low_frequency_cutoff, high_frequency_cutoff,
snr_threshold, tlen, delta_f, dtype):
"""
Create a matched filter engine.
Parameters
----------
low_frequency_cutoff : {None, float}, optional
The frequency to begin the filter calculation. If None, begin
at the first frequency after DC.
high_frequency_cutoff : {None, float}, optional
The frequency to stop the filter calculation. If None, continue
to the nyquist frequency.
snr_threshold : float
The minimum snr to return when filtering
"""
self.tlen = tlen
self.delta_f = delta_f
self.dtype = dtype
self.snr_threshold = snr_threshold
self.flow = low_frequency_cutoff
self.fhigh = high_frequency_cutoff
self.matched_filter_and_cluster = \
self.full_matched_filter_and_cluster
self.snr_plus_mem = zeros(self.tlen, dtype=self.dtype)
self.corr_plus_mem = zeros(self.tlen, dtype=self.dtype)
self.snr_cross_mem = zeros(self.tlen, dtype=self.dtype)
self.corr_cross_mem = zeros(self.tlen, dtype=self.dtype)
self.snr_mem = zeros(self.tlen, dtype=self.dtype)
self.cached_hplus_hcross_correlation = None
self.cached_hplus_hcross_hplus = None
self.cached_hplus_hcross_hcross = None
self.cached_hplus_hcross_psd = None
def full_matched_filter_and_cluster(self, hplus, hcross, hplus_norm,
hcross_norm, psd, stilde, window):
"""
Return the complex snr and normalization.
Calculated the matched filter, threshold, and cluster.
Parameters
----------
h_quantities : Various
FILL ME IN
stilde : FrequencySeries
The strain data to be filtered.
window : int
The size of the cluster window in samples.
Returns
-------
snr : TimeSeries
A time series containing the complex snr.
norm : float
The normalization of the complex snr.
correlation: FrequencySeries
A frequency series containing the correlation vector.
idx : Array
List of indices of the triggers.
snrv : Array
The snr values at the trigger locations.
"""
I_plus, Iplus_corr, Iplus_norm = matched_filter_core(hplus, stilde,
h_norm=hplus_norm,
low_frequency_cutoff=self.flow,
high_frequency_cutoff=self.fhigh,
out=self.snr_plus_mem,
corr_out=self.corr_plus_mem)
I_cross, Icross_corr, Icross_norm = matched_filter_core(hcross,
stilde, h_norm=hcross_norm,
low_frequency_cutoff=self.flow,
high_frequency_cutoff=self.fhigh,
out=self.snr_cross_mem,
corr_out=self.corr_cross_mem)
# The information on the complex side of this overlap is important
# we may want to use this in the future.
if not id(hplus) == self.cached_hplus_hcross_hplus:
self.cached_hplus_hcross_correlation = None
if not id(hcross) == self.cached_hplus_hcross_hcross:
self.cached_hplus_hcross_correlation = None
if not id(psd) == self.cached_hplus_hcross_psd:
self.cached_hplus_hcross_correlation = None
if self.cached_hplus_hcross_correlation is None:
hplus_cross_corr = overlap_cplx(hplus, hcross, psd=psd,
low_frequency_cutoff=self.flow,
high_frequency_cutoff=self.fhigh,
normalized=False)
hplus_cross_corr = numpy.real(hplus_cross_corr)
hplus_cross_corr = hplus_cross_corr / (hcross_norm*hplus_norm)**0.5
self.cached_hplus_hcross_correlation = hplus_cross_corr
self.cached_hplus_hcross_hplus = id(hplus)
self.cached_hplus_hcross_hcross = id(hcross)
self.cached_hplus_hcross_psd = id(psd)
else:
hplus_cross_corr = self.cached_hplus_hcross_correlation
snr = self._maximized_snr(I_plus,I_cross,
hplus_cross_corr,
hpnorm=Iplus_norm,
hcnorm=Icross_norm,
out=self.snr_mem,
thresh=self.snr_threshold,
analyse_slice=stilde.analyze)
# FIXME: This should live further down
# Convert output to pycbc TimeSeries
delta_t = 1.0 / (self.tlen * stilde.delta_f)
snr = TimeSeries(snr, epoch=stilde.start_time, delta_t=delta_t,
copy=False)
idx, snrv = events.threshold_real_numpy(snr[stilde.analyze],
self.snr_threshold)
if len(idx) == 0:
return [], 0, 0, [], [], [], [], 0, 0, 0
logging.info("%s points above threshold", str(len(idx)))
idx, snrv = events.cluster_reduce(idx, snrv, window)
logging.info("%s clustered points", str(len(idx)))
# erased self.
u_vals, coa_phase = self._maximized_extrinsic_params\
(I_plus.data, I_cross.data, hplus_cross_corr,
indices=idx+stilde.analyze.start, hpnorm=Iplus_norm,
hcnorm=Icross_norm)
return snr, Iplus_corr, Icross_corr, idx, snrv, u_vals, coa_phase,\
hplus_cross_corr, Iplus_norm, Icross_norm
def _maximized_snr(self, hplus, hcross, hphccorr, **kwargs):
return compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr,
**kwargs)
def _maximized_extrinsic_params(self, hplus, hcross, hphccorr, **kwargs):
return compute_u_val_for_sky_loc_stat(hplus, hcross, hphccorr,
**kwargs)
class MatchedFilterSkyMaxControlNoPhase(MatchedFilterSkyMaxControl):
# Basically the same as normal SkyMaxControl, except we use a slight
# variation in the internal SNR functions.
def _maximized_snr(self, hplus, hcross, hphccorr, **kwargs):
return compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross,
hphccorr, **kwargs)
def _maximized_extrinsic_params(self, hplus, hcross, hphccorr, **kwargs):
return compute_u_val_for_sky_loc_stat_no_phase(hplus, hcross, hphccorr,
**kwargs)
def make_frequency_series(vec):
"""Return a frequency series of the input vector.
If the input is a frequency series it is returned, else if the input
vector is a real time series it is fourier transformed and returned as a
frequency series.
Parameters
----------
vector : TimeSeries or FrequencySeries
Returns
-------
Frequency Series: FrequencySeries
A frequency domain version of the input vector.
"""
if isinstance(vec, FrequencySeries):
return vec
if isinstance(vec, TimeSeries):
N = len(vec)
n = N // 2 + 1
delta_f = 1.0 / N / vec.delta_t
vectilde = FrequencySeries(zeros(n, dtype=complex_same_precision_as(vec)),
delta_f=delta_f, copy=False)
fft(vec, vectilde)
return vectilde
else:
raise TypeError("Can only convert a TimeSeries to a FrequencySeries")
def sigmasq_series(htilde, psd=None, low_frequency_cutoff=None,
high_frequency_cutoff=None):
"""Return a cumulative sigmasq frequency series.
Return a frequency series containing the accumulated power in the input
up to that frequency.
Parameters
----------
htilde : TimeSeries or FrequencySeries
The input vector
psd : {None, FrequencySeries}, optional
The psd used to weight the accumulated power.
low_frequency_cutoff : {None, float}, optional
The frequency to begin accumulating power. If None, start at the beginning
of the vector.
high_frequency_cutoff : {None, float}, optional
The frequency to stop considering accumulated power. If None, continue
until the end of the input vector.
Returns
-------
Frequency Series: FrequencySeries
A frequency series containing the cumulative sigmasq.
"""
htilde = make_frequency_series(htilde)
N = (len(htilde)-1) * 2
norm = 4.0 * htilde.delta_f
kmin, kmax = get_cutoff_indices(low_frequency_cutoff,
high_frequency_cutoff, htilde.delta_f, N)
sigma_vec = FrequencySeries(zeros(len(htilde), dtype=real_same_precision_as(htilde)),
delta_f = htilde.delta_f, copy=False)
mag = htilde.squared_norm()
if psd is not None:
mag /= psd
sigma_vec[kmin:kmax] = mag[kmin:kmax].cumsum()
return sigma_vec*norm
def sigmasq(htilde, psd = None, low_frequency_cutoff=None,
high_frequency_cutoff=None):
"""Return the loudness of the waveform. This is defined (see Duncan
Brown's thesis) as the unnormalized matched-filter of the input waveform,
htilde, with itself. This quantity is usually referred to as (sigma)^2
and is then used to normalize matched-filters with the data.
Parameters
----------
htilde : TimeSeries or FrequencySeries
The input vector containing a waveform.
psd : {None, FrequencySeries}, optional
The psd used to weight the accumulated power.
low_frequency_cutoff : {None, float}, optional
The frequency to begin considering waveform power.
high_frequency_cutoff : {None, float}, optional
The frequency to stop considering waveform power.
Returns
-------
sigmasq: float
"""
htilde = make_frequency_series(htilde)
N = (len(htilde)-1) * 2
norm = 4.0 * htilde.delta_f
kmin, kmax = get_cutoff_indices(low_frequency_cutoff,
high_frequency_cutoff, htilde.delta_f, N)
ht = htilde[kmin:kmax]
if psd:
try:
numpy.testing.assert_almost_equal(ht.delta_f, psd.delta_f)
except AssertionError:
raise ValueError('Waveform does not have same delta_f as psd')
if psd is None:
sq = ht.inner(ht)
else:
sq = ht.weighted_inner(ht, psd[kmin:kmax])
return sq.real * norm
def sigma(htilde, psd = None, low_frequency_cutoff=None,
high_frequency_cutoff=None):
""" Return the sigma of the waveform. See sigmasq for more details.
Parameters
----------
htilde : TimeSeries or FrequencySeries
The input vector containing a waveform.
psd : {None, FrequencySeries}, optional
The psd used to weight the accumulated power.
low_frequency_cutoff : {None, float}, optional
The frequency to begin considering waveform power.
high_frequency_cutoff : {None, float}, optional
The frequency to stop considering waveform power.
Returns
-------
sigmasq: float
"""
return sqrt(sigmasq(htilde, psd, low_frequency_cutoff, high_frequency_cutoff))
def get_cutoff_indices(flow, fhigh, df, N):
"""
Gets the indices of a frequency series at which to stop an overlap
calculation.
Parameters
----------
flow: float
The frequency (in Hz) of the lower index.
fhigh: float
The frequency (in Hz) of the upper index.
df: float
The frequency step (in Hz) of the frequency series.
N: int
The number of points in the **time** series. Can be odd
or even.
Returns
-------
kmin: int
kmax: int
"""
if flow:
kmin = int(flow / df)
if kmin < 0:
err_msg = "Start frequency cannot be negative. "
err_msg += "Supplied value and kmin {} and {}".format(flow, kmin)
raise ValueError(err_msg)
else:
kmin = 1
if fhigh:
kmax = int(fhigh / df )
if kmax > int((N + 1)/2.):
kmax = int((N + 1)/2.)
else:
# int() truncates towards 0, so this is
# equivalent to the floor of the float
kmax = int((N + 1)/2.)
if kmax <= kmin:
err_msg = "Kmax cannot be less than or equal to kmin. "
err_msg += "Provided values of freqencies (min,max) were "
err_msg += "{} and {} ".format(flow, fhigh)
err_msg += "corresponding to (kmin, kmax) of "
err_msg += "{} and {}.".format(kmin, kmax)
raise ValueError(err_msg)
return kmin,kmax
def matched_filter_core(template, data, psd=None, low_frequency_cutoff=None,
high_frequency_cutoff=None, h_norm=None, out=None, corr_out=None):
""" Return the complex snr and normalization.
Return the complex snr, along with its associated normalization of the template,
matched filtered against the data.
Parameters
----------
template : TimeSeries or FrequencySeries
The template waveform
data : TimeSeries or FrequencySeries
The strain data to be filtered.
psd : {FrequencySeries}, optional
The noise weighting of the filter.
low_frequency_cutoff : {None, float}, optional
The frequency to begin the filter calculation. If None, begin at the
first frequency after DC.
high_frequency_cutoff : {None, float}, optional
The frequency to stop the filter calculation. If None, continue to the
the nyquist frequency.
h_norm : {None, float}, optional
The template normalization. If none, this value is calculated internally.
out : {None, Array}, optional
An array to use as memory for snr storage. If None, memory is allocated
internally.
corr_out : {None, Array}, optional
An array to use as memory for correlation storage. If None, memory is allocated
internally. If provided, management of the vector is handled externally by the
caller. No zero'ing is done internally.
Returns
-------
snr : TimeSeries
A time series containing the complex snr.
corrrelation: FrequencySeries
A frequency series containing the correlation vector.
norm : float
The normalization of the complex snr.
"""
htilde = make_frequency_series(template)
stilde = make_frequency_series(data)
if len(htilde) != len(stilde):
raise ValueError("Length of template and data must match")
N = (len(stilde)-1) * 2
kmin, kmax = get_cutoff_indices(low_frequency_cutoff,
high_frequency_cutoff, stilde.delta_f, N)
if corr_out is not None:
qtilde = corr_out
else:
qtilde = zeros(N, dtype=complex_same_precision_as(data))
if out is None:
_q = zeros(N, dtype=complex_same_precision_as(data))
elif (len(out) == N) and type(out) is Array and out.kind =='complex':
_q = out
else:
raise TypeError('Invalid Output Vector: wrong length or dtype')
correlate(htilde[kmin:kmax], stilde[kmin:kmax], qtilde[kmin:kmax])
if psd is not None:
if isinstance(psd, FrequencySeries):
try:
numpy.testing.assert_almost_equal(stilde.delta_f, psd.delta_f)
except AssertionError:
raise ValueError("PSD delta_f does not match data")
qtilde[kmin:kmax] /= psd[kmin:kmax]
else:
raise TypeError("PSD must be a FrequencySeries")
ifft(qtilde, _q)
if h_norm is None:
h_norm = sigmasq(htilde, psd, low_frequency_cutoff, high_frequency_cutoff)
norm = (4.0 * stilde.delta_f) / sqrt( h_norm)
return (TimeSeries(_q, epoch=stilde._epoch, delta_t=stilde.delta_t, copy=False),
FrequencySeries(qtilde, epoch=stilde._epoch, delta_f=stilde.delta_f, copy=False),
norm)
def smear(idx, factor):
"""
This function will take as input an array of indexes and return every
unique index within the specified factor of the inputs.
E.g.: smear([5,7,100],2) = [3,4,5,6,7,8,9,98,99,100,101,102]
Parameters
-----------
idx : numpy.array of ints
The indexes to be smeared.
factor : idx
The factor by which to smear out the input array.
Returns
--------
new_idx : numpy.array of ints
The smeared array of indexes.
"""
s = [idx]
for i in range(factor+1):
a = i - factor/2
s += [idx + a]
return numpy.unique(numpy.concatenate(s))
def matched_filter(template, data, psd=None, low_frequency_cutoff=None,
high_frequency_cutoff=None, sigmasq=None):
""" Return the complex snr.
Return the complex snr, along with its associated normalization of the
template, matched filtered against the data.
Parameters
----------
template : TimeSeries or FrequencySeries
The template waveform
data : TimeSeries or FrequencySeries
The strain data to be filtered.
psd : FrequencySeries
The noise weighting of the filter.
low_frequency_cutoff : {None, float}, optional
The frequency to begin the filter calculation. If None, begin at the
first frequency after DC.
high_frequency_cutoff : {None, float}, optional
The frequency to stop the filter calculation. If None, continue to the
the nyquist frequency.
sigmasq : {None, float}, optional
The template normalization. If none, this value is calculated
internally.
Returns
-------
snr : TimeSeries
A time series containing the complex snr.
"""
snr, _, norm = matched_filter_core(template, data, psd=psd,
low_frequency_cutoff=low_frequency_cutoff,
high_frequency_cutoff=high_frequency_cutoff, h_norm=sigmasq)
return snr * norm
_snr = None
def match(vec1, vec2, psd=None, low_frequency_cutoff=None,
high_frequency_cutoff=None, v1_norm=None, v2_norm=None,
subsample_interpolation=False, return_phase=False):
""" Return the match between the two TimeSeries or FrequencySeries.
Return the match between two waveforms. This is equivalent to the overlap
maximized over time and phase.
Parameters
----------
vec1 : TimeSeries or FrequencySeries
The input vector containing a waveform.
vec2 : TimeSeries or FrequencySeries
The input vector containing a waveform.
psd : Frequency Series
A power spectral density to weight the overlap.
low_frequency_cutoff : {None, float}, optional
The frequency to begin the match.
high_frequency_cutoff : {None, float}, optional
The frequency to stop the match.
v1_norm : {None, float}, optional
The normalization of the first waveform. This is equivalent to its
sigmasq value. If None, it is internally calculated.
v2_norm : {None, float}, optional
The normalization of the second waveform. This is equivalent to its
sigmasq value. If None, it is internally calculated.
subsample_interpolation : {False, bool}, optional
If True the peak will be interpolated between samples using a simple
quadratic fit. This can be important if measuring matches very close to
1 and can cause discontinuities if you don't use it as matches move
between discrete samples. If True the index returned will be a float.
return_phase : {False, bool}, optional
If True, also return the phase shift that gives the match.
Returns
-------
match: float
index: int
The number of samples to shift to get the match.
phi: float
Phase to rotate complex waveform to get the match, if desired.
"""
htilde = make_frequency_series(vec1)
stilde = make_frequency_series(vec2)
N = (len(htilde)-1) * 2
global _snr
if _snr is None or _snr.dtype != htilde.dtype or len(_snr) != N:
_snr = zeros(N,dtype=complex_same_precision_as(vec1))
snr, _, snr_norm = matched_filter_core(htilde, stilde, psd,
low_frequency_cutoff,
high_frequency_cutoff,
v1_norm, out=_snr)
maxsnr, max_id = snr.abs_max_loc()
if v2_norm is None:
v2_norm = sigmasq(stilde, psd, low_frequency_cutoff,
high_frequency_cutoff)
if subsample_interpolation:
# This uses the implementation coded up in sbank. Thanks Nick!
# The maths for this is well summarized here:
# https://ccrma.stanford.edu/~jos/sasp/Quadratic_Interpolation_Spectral_Peaks.html
# We use adjacent points to interpolate, but wrap off the end if needed
left = abs(snr[-1]) if max_id == 0 else abs(snr[max_id - 1])
middle = maxsnr
right = abs(snr[0]) if max_id == (len(snr)-1) else abs(snr[max_id + 1])
# Get derivatives
id_shift, maxsnr = quadratic_interpolate_peak(left, middle, right)
max_id = max_id + id_shift
if return_phase:
rounded_max_id = int(round(max_id))
phi = numpy.angle(snr[rounded_max_id])
return maxsnr * snr_norm / sqrt(v2_norm), max_id, phi
else:
return maxsnr * snr_norm / sqrt(v2_norm), max_id
def overlap(vec1, vec2, psd=None, low_frequency_cutoff=None,
high_frequency_cutoff=None, normalized=True):
""" Return the overlap between the two TimeSeries or FrequencySeries.
Parameters
----------
vec1 : TimeSeries or FrequencySeries
The input vector containing a waveform.
vec2 : TimeSeries or FrequencySeries
The input vector containing a waveform.
psd : Frequency Series
A power spectral density to weight the overlap.
low_frequency_cutoff : {None, float}, optional
The frequency to begin the overlap.
high_frequency_cutoff : {None, float}, optional
The frequency to stop the overlap.
normalized : {True, boolean}, optional
Set if the overlap is normalized. If true, it will range from 0 to 1.
Returns
-------
overlap: float
"""
return overlap_cplx(vec1, vec2, psd=psd, \
low_frequency_cutoff=low_frequency_cutoff,\
high_frequency_cutoff=high_frequency_cutoff,\
normalized=normalized).real
def overlap_cplx(vec1, vec2, psd=None, low_frequency_cutoff=None,
high_frequency_cutoff=None, normalized=True):
"""Return the complex overlap between the two TimeSeries or FrequencySeries.
Parameters
----------
vec1 : TimeSeries or FrequencySeries
The input vector containing a waveform.
vec2 : TimeSeries or FrequencySeries
The input vector containing a waveform.
psd : Frequency Series
A power spectral density to weight the overlap.
low_frequency_cutoff : {None, float}, optional
The frequency to begin the overlap.
high_frequency_cutoff : {None, float}, optional
The frequency to stop the overlap.
normalized : {True, boolean}, optional
Set if the overlap is normalized. If true, it will range from 0 to 1.
Returns
-------
overlap: complex
"""
htilde = make_frequency_series(vec1)
stilde = make_frequency_series(vec2)
kmin, kmax = get_cutoff_indices(low_frequency_cutoff,
high_frequency_cutoff, stilde.delta_f, (len(stilde)-1) * 2)
if psd:
inner = (htilde[kmin:kmax]).weighted_inner(stilde[kmin:kmax], psd[kmin:kmax])
else:
inner = (htilde[kmin:kmax]).inner(stilde[kmin:kmax])
if normalized:
sig1 = sigma(vec1, psd=psd, low_frequency_cutoff=low_frequency_cutoff,
high_frequency_cutoff=high_frequency_cutoff)
sig2 = sigma(vec2, psd=psd, low_frequency_cutoff=low_frequency_cutoff,
high_frequency_cutoff=high_frequency_cutoff)
norm = 1 / sig1 / sig2
else:
norm = 1
return 4 * htilde.delta_f * inner * norm
def quadratic_interpolate_peak(left, middle, right):
""" Interpolate the peak and offset using a quadratic approximation
Parameters
----------
left : numpy array
Values at a relative bin value of [-1]
middle : numpy array
Values at a relative bin value of [0]
right : numpy array
Values at a relative bin value of [1]
Returns
-------
bin_offset : numpy array
Array of bins offsets, each in the range [-1/2, 1/2]
peak_values : numpy array
Array of the estimated peak values at the interpolated offset
"""
bin_offset = 1.0/2.0 * (left - right) / (left - 2 * middle + right)
peak_value = middle - 0.25 * (left - right) * bin_offset
return bin_offset, peak_value
class LiveBatchMatchedFilter(object):
"""Calculate SNR and signal consistency tests in a batched progression"""
def __init__(self, templates, snr_threshold, chisq_bins, sg_chisq,
maxelements=2**27,
snr_abort_threshold=None,
newsnr_threshold=None,
max_triggers_in_batch=None):
"""Create a batched matchedfilter instance
Parameters
----------
templates: list of `FrequencySeries`
List of templates from the FilterBank class.
snr_threshold: float
Minimum value to record peaks in the SNR time series.
chisq_bins: str
Str that determines how the number of chisq bins varies as a
function of the template bank parameters.
sg_chisq: pycbc.vetoes.SingleDetSGChisq
Instance of the sg_chisq class to calculate sg_chisq with.
maxelements: {int, 2**27}
Maximum size of a batched fourier transform.
snr_abort_threshold: {float, None}
If the SNR is above this threshold, do not record any triggers.
newsnr_threshold: {float, None}
Only record triggers that have a re-weighted NewSNR above this
threshold.
max_triggers_in_batch: {int, None}
Record X number of the loudest triggers by newsnr in each mpi
process group. Signal consistency values will also only be calculated
for these triggers.
"""
self.snr_threshold = snr_threshold
self.snr_abort_threshold = snr_abort_threshold
self.newsnr_threshold = newsnr_threshold
self.max_triggers_in_batch = max_triggers_in_batch
from pycbc import vetoes
self.power_chisq = vetoes.SingleDetPowerChisq(chisq_bins, None)
self.sg_chisq = sg_chisq
durations = numpy.array([1.0 / t.delta_f for t in templates])
lsort = durations.argsort()
durations = durations[lsort]
templates = [templates[li] for li in lsort]
# Figure out how to chunk together the templates into groups to process
_, counts = numpy.unique(durations, return_counts=True)
tsamples = [(len(t) - 1) * 2 for t in templates]
grabs = maxelements / numpy.unique(tsamples)
chunks = numpy.array([])
num = 0
for count, grab in zip(counts, grabs):
chunks = numpy.append(chunks, numpy.arange(num, count + num, grab))
chunks = numpy.append(chunks, [count + num])
num += count
chunks = numpy.unique(chunks).astype(numpy.uint32)
# We now have how many templates to grab at a time.
self.chunks = chunks[1:] - chunks[0:-1]
self.out_mem = {}
self.cout_mem = {}
self.ifts = {}
chunk_durations = [durations[i] for i in chunks[:-1]]
self.chunk_tsamples = [tsamples[int(i)] for i in chunks[:-1]]
samples = self.chunk_tsamples * self.chunks
# Create workspace memory for correlate and snr
mem_ids = [(a, b) for a, b in zip(chunk_durations, self.chunks)]
mem_types = set(zip(mem_ids, samples))
self.tgroups, self.mids = [], []
for i, size in mem_types:
dur, count = i
self.out_mem[i] = zeros(size, dtype=numpy.complex64)
self.cout_mem[i] = zeros(size, dtype=numpy.complex64)
self.ifts[i] = IFFT(self.cout_mem[i], self.out_mem[i],
nbatch=count,
size=len(self.cout_mem[i]) // count)
# Split the templates into their processing groups
for dur, count in mem_ids:
tgroup = templates[0:count]
self.tgroups.append(tgroup)
self.mids.append((dur, count))
templates = templates[count:]
# Associate the snr and corr memory block to each template
self.corr = []
for i, tgroup in enumerate(self.tgroups):
psize = self.chunk_tsamples[i]
s = 0
e = psize
mid = self.mids[i]
for htilde in tgroup:
htilde.out = self.out_mem[mid][s:e]
htilde.cout = self.cout_mem[mid][s:e]
s += psize
e += psize
self.corr.append(BatchCorrelator(tgroup, [t.cout for t in tgroup], len(tgroup[0])))
def set_data(self, data):
"""Set the data reader object to use"""
self.data = data
self.block_id = 0
def combine_results(self, results):
"""Combine results from different batches of filtering"""
result = {}
for key in results[0]:
result[key] = numpy.concatenate([r[key] for r in results])
return result
def process_data(self, data_reader):
"""Process the data for all of the templates"""
self.set_data(data_reader)
return self.process_all()
def process_all(self):
"""Process every batch group and return as single result"""
results = []
veto_info = []
while 1:
result, veto = self._process_batch()
if result is False: return False
if result is None: break
results.append(result)
veto_info += veto
result = self.combine_results(results)
if self.max_triggers_in_batch:
sort = result['snr'].argsort()[::-1][:self.max_triggers_in_batch]
for key in result:
result[key] = result[key][sort]
tmp = veto_info
veto_info = [tmp[i] for i in sort]
result = self._process_vetoes(result, veto_info)
return result
def _process_vetoes(self, results, veto_info):
"""Calculate signal based vetoes"""
chisq = numpy.array(numpy.zeros(len(veto_info)), numpy.float32, ndmin=1)
dof = numpy.array(numpy.zeros(len(veto_info)), numpy.uint32, ndmin=1)
sg_chisq = numpy.array(numpy.zeros(len(veto_info)), numpy.float32,
ndmin=1)
results['chisq'] = chisq
results['chisq_dof'] = dof
results['sg_chisq'] = sg_chisq
keep = []
for i, (snrv, norm, l, htilde, stilde) in enumerate(veto_info):
correlate(htilde, stilde, htilde.cout)
c, d = self.power_chisq.values(htilde.cout, snrv,
norm, stilde.psd, [l], htilde)
chisq[i] = c[0] / d[0]
dof[i] = d[0]
sgv = self.sg_chisq.values(stilde, htilde, stilde.psd,
snrv, norm, c, d, [l])
if sgv is not None:
sg_chisq[i] = sgv[0]
if self.newsnr_threshold:
newsnr = ranking.newsnr(results['snr'][i], chisq[i])
if newsnr >= self.newsnr_threshold:
keep.append(i)
if self.newsnr_threshold:
keep = numpy.array(keep, dtype=numpy.uint32)
for key in results:
results[key] = results[key][keep]
return results
def _process_batch(self):
"""Process only a single batch group of data"""
if self.block_id == len(self.tgroups):
return None, None
tgroup = self.tgroups[self.block_id]
psize = self.chunk_tsamples[self.block_id]
mid = self.mids[self.block_id]
stilde = self.data.overwhitened_data(tgroup[0].delta_f)
psd = stilde.psd
valid_end = int(psize - self.data.trim_padding)
valid_start = int(valid_end - self.data.blocksize * self.data.sample_rate)
seg = slice(valid_start, valid_end)
self.corr[self.block_id].execute(stilde)
self.ifts[mid].execute()
self.block_id += 1
snr = numpy.zeros(len(tgroup), dtype=numpy.complex64)
time = numpy.zeros(len(tgroup), dtype=numpy.float64)
templates = numpy.zeros(len(tgroup), dtype=numpy.uint64)
sigmasq = numpy.zeros(len(tgroup), dtype=numpy.float32)
time[:] = self.data.start_time
result = {}
tkeys = tgroup[0].params.dtype.names
for key in tkeys:
result[key] = []
veto_info = []
# Find the peaks in our SNR times series from the various templates
i = 0
for htilde in tgroup:
if hasattr(htilde, 'time_offset'):
if 'time_offset' not in result:
result['time_offset'] = []
l = htilde.out[seg].abs_arg_max()
sgm = htilde.sigmasq(psd)
norm = 4.0 * htilde.delta_f / (sgm ** 0.5)
l += valid_start
snrv = numpy.array([htilde.out[l]])
# If nothing is above threshold we can exit this template
s = abs(snrv[0]) * norm
if s < self.snr_threshold:
continue
time[i] += float(l - valid_start) / self.data.sample_rate
# We have an SNR so high that we will drop the entire analysis
# of this chunk of time!
if self.snr_abort_threshold is not None and s > self.snr_abort_threshold:
logging.info("We are seeing some *really* high SNRs, lets"
" assume they aren't signals and just give up")
return False, []
veto_info.append((snrv, norm, l, htilde, stilde))
snr[i] = snrv[0] * norm
sigmasq[i] = sgm
templates[i] = htilde.id
if not hasattr(htilde, 'dict_params'):
htilde.dict_params = {}
for key in tkeys:
htilde.dict_params[key] = htilde.params[key]
for key in tkeys:
result[key].append(htilde.dict_params[key])
if hasattr(htilde, 'time_offset'):
result['time_offset'].append(htilde.time_offset)
i += 1
result['snr'] = abs(snr[0:i])
result['coa_phase'] = numpy.angle(snr[0:i])
result['end_time'] = time[0:i]
result['template_id'] = templates[0:i]
result['sigmasq'] = sigmasq[0:i]
for key in tkeys:
result[key] = numpy.array(result[key])
if 'time_offset' in result:
result['time_offset'] = numpy.array(result['time_offset'])
return result, veto_info
def followup_event_significance(ifo, data_reader, bank,
template_id, coinc_times,
coinc_threshold=0.005,
lookback=150, duration=0.095):
""" Followup an event in another detector and determine its significance
"""
from pycbc.waveform import get_waveform_filter_length_in_time
tmplt = bank.table[template_id]
length_in_time = get_waveform_filter_length_in_time(tmplt['approximant'],
tmplt)
# calculate onsource time range
from pycbc.detector import Detector
onsource_start = -numpy.inf
onsource_end = numpy.inf
fdet = Detector(ifo)
for cifo in coinc_times:
time = coinc_times[cifo]
dtravel = Detector(cifo).light_travel_time_to_detector(fdet)
if time - dtravel > onsource_start:
onsource_start = time - dtravel
if time + dtravel < onsource_end:
onsource_end = time + dtravel
# Source must be within this time window to be considered a possible
# coincidence
onsource_start -= coinc_threshold
onsource_end += coinc_threshold
# Calculate how much time needed to calculate significance
trim_pad = (data_reader.trim_padding * data_reader.strain.delta_t)
bdur = int(lookback + 2.0 * trim_pad + length_in_time)
if bdur > data_reader.strain.duration * .75:
bdur = data_reader.strain.duration * .75
# Require all strain be valid within lookback time
if data_reader.state is not None:
state_start_time = data_reader.strain.end_time \
- data_reader.reduced_pad * data_reader.strain.delta_t - bdur
if not data_reader.state.is_extent_valid(state_start_time, bdur):
return None, None, None, None
# We won't require that all DQ checks be valid for now, except at
# onsource time.
if data_reader.dq is not None:
dq_start_time = onsource_start - duration / 2.0
dq_duration = onsource_end - onsource_start + duration
if not data_reader.dq.is_extent_valid(dq_start_time, dq_duration):
return None, None, None, None
# Calculate SNR time series for this duration
htilde = bank.get_template(template_id, min_buffer=bdur)
stilde = data_reader.overwhitened_data(htilde.delta_f)
sigma2 = htilde.sigmasq(stilde.psd)
snr, _, norm = matched_filter_core(htilde, stilde, h_norm=sigma2)
# Find peak in on-source and determine p-value
onsrc = snr.time_slice(onsource_start, onsource_end)
peak = onsrc.abs_arg_max()
peak_time = peak * snr.delta_t + onsrc.start_time
peak_value = abs(onsrc[peak])
bstart = float(snr.start_time) + length_in_time + trim_pad
bkg = abs(snr.time_slice(bstart, onsource_start)).numpy()
window = int((onsource_end - onsource_start) * snr.sample_rate)
nsamples = int(len(bkg) / window)
peaks = bkg[:nsamples*window].reshape(nsamples, window).max(axis=1)
pvalue = (1 + (peaks >= peak_value).sum()) / float(1 + nsamples)
# Return recentered source SNR for bayestar, along with p-value, and trig
peak_full = int((peak_time - snr.start_time) / snr.delta_t)
half_dur_samples = int(snr.sample_rate * duration / 2)
snr_slice = slice(peak_full - half_dur_samples,
peak_full + half_dur_samples + 1)
baysnr = snr[snr_slice]
logging.info('Adding %s to candidate, pvalue %s, %s samples', ifo,
pvalue, nsamples)
return baysnr * norm, peak_time, pvalue, sigma2
def compute_followup_snr_series(data_reader, htilde, trig_time,
duration=0.095, check_state=True,
coinc_window=0.05):
"""Given a StrainBuffer, a template frequency series and a trigger time,
compute a portion of the SNR time series centered on the trigger for its
rapid sky localization and followup.
If the trigger time is too close to the boundary of the valid data segment
the SNR series is calculated anyway and might be slightly contaminated by
filter and wrap-around effects. For reasonable durations this will only
affect a small fraction of the triggers and probably in a negligible way.
Parameters
----------
data_reader : StrainBuffer
The StrainBuffer object to read strain data from.
htilde : FrequencySeries
The frequency series containing the template waveform.
trig_time : {float, lal.LIGOTimeGPS}
The trigger time.
duration : float (optional)
Duration of the computed SNR series in seconds. If omitted, it defaults
to twice the Earth light travel time plus 10 ms of timing uncertainty.
check_state : boolean
If True, and the detector was offline or flagged for bad data quality
at any point during the inspiral, then return (None, None) instead.
coinc_window : float (optional)
Maximum possible time between coincident triggers at different
detectors. This is needed to properly determine data padding.
Returns
-------
snr : TimeSeries
The portion of SNR around the trigger. None if the detector is offline
or has bad data quality, and check_state is True.
"""
if check_state:
# was the detector observing for the full amount of involved data?
state_start_time = trig_time - duration / 2 - htilde.length_in_time
state_end_time = trig_time + duration / 2
state_duration = state_end_time - state_start_time
if data_reader.state is not None:
if not data_reader.state.is_extent_valid(state_start_time,
state_duration):
return None
# was the data quality ok for the full amount of involved data?
dq_start_time = state_start_time - data_reader.dq_padding
dq_duration = state_duration + 2 * data_reader.dq_padding
if data_reader.dq is not None:
if not data_reader.dq.is_extent_valid(dq_start_time, dq_duration):
return None
stilde = data_reader.overwhitened_data(htilde.delta_f)
snr, _, norm = matched_filter_core(htilde, stilde,
h_norm=htilde.sigmasq(stilde.psd))
valid_end = int(len(snr) - data_reader.trim_padding)
valid_start = int(valid_end - data_reader.blocksize * snr.sample_rate)
half_dur_samples = int(snr.sample_rate * duration / 2)
coinc_samples = int(snr.sample_rate * coinc_window)
valid_start -= half_dur_samples + coinc_samples
valid_end += half_dur_samples
if valid_start < 0 or valid_end > len(snr)-1:
raise ValueError(('Requested SNR duration ({0} s)'
' too long').format(duration))
# Onsource slice for Bayestar followup
onsource_idx = float(trig_time - snr.start_time) * snr.sample_rate
onsource_idx = int(round(onsource_idx))
onsource_slice = slice(onsource_idx - half_dur_samples,
onsource_idx + half_dur_samples + 1)
return snr[onsource_slice] * norm
__all__ = ['match', 'matched_filter', 'sigmasq', 'sigma', 'get_cutoff_indices',
'sigmasq_series', 'make_frequency_series', 'overlap',
'overlap_cplx', 'matched_filter_core', 'correlate',
'MatchedFilterControl', 'LiveBatchMatchedFilter',
'MatchedFilterSkyMaxControl', 'MatchedFilterSkyMaxControlNoPhase',
'compute_max_snr_over_sky_loc_stat_no_phase',
'compute_max_snr_over_sky_loc_stat',
'compute_followup_snr_series',
'compute_u_val_for_sky_loc_stat_no_phase',
'compute_u_val_for_sky_loc_stat',
'followup_event_significance']