Commit 78904f2b authored by Ben Farr's avatar Ben Farr
Browse files

Work in the FD by default

parent e3d3bdab
......@@ -25,11 +25,11 @@ def specgrams(data, NFFT=1024, noverlap=None, ymin=None, ymax=None, axs=None):
if noverlap is None:
noverlap = NFFT/2
for ax, td in zip(axs, data):
time_offset = td.times[0]
Pxx, fs, bins, im = ax.specgram(td, NFFT=NFFT,
xextent=[td.times[0], td.times[-1]],
Fs=td.sample_rate, noverlap=noverlap,
for ax, fd in zip(axs, data):
time_offset = fd.times[0]
Pxx, fs, bins, im = ax.specgram(fd.tseries, NFFT=NFFT,
xextent=[fd.times[0], fd.times[-1]],
Fs=fd.sample_rate, noverlap=noverlap,
cmap=option_d)
ax.set_ylim(ymin=ymin, ymax=ymax)
......
......@@ -13,74 +13,95 @@ from pycbc.filter import resample_to_delta_t as _resample_to_delta_t
from .utils import tukey_window
class Strain(TimeSeries):
class Strain(FrequencySeries):
"""
A class for smoothly moving between time and frequency domain
representations of a strain.
:param series:
A pycbc :class:`TimeSeries` or :class:`FrequencySeries`.
A pycbc :class:`TimeSeries` or :class:`FrequencySeries`. If an
array is given, it's assumed to be a frequency series and `**kwargs`
are passed to :class:`FrequencySeries`.
:param white:
If ``True``, the input series is assumed to be whitened, and amplitudes
will be preserved when moving between domains.
"""
def __init__(self, series, white=False, **kwargs):
self._bandpass_correction = 1.
self._white = white
if isinstance(series, TimeSeries):
tseries = series
elif isinstance(series, FrequencySeries):
tseries = series.to_timeseries()
if isinstance(series, FrequencySeries):
fseries = series
elif isinstance(series, TimeSeries):
fseries = series.to_frequencyseries()
if self._white:
fseries.data /= series.delta_t * np.sqrt(len(series)/2)
if white:
sel = series.data != 0.
self._bandpass_correction = np.sqrt(len(sel)/np.count_nonzero(sel))
tseries.data *= tseries.delta_t * np.sqrt(len(tseries)/2)
tseries.data *= self._bandpass_correction
else:
tseries = TimeSeries(series, **kwargs)
fseries = FrequencySeries(series, **kwargs)
sel = fseries.data != 0.
self._bandpass_correction = np.sqrt(len(sel)/np.count_nonzero(sel))
super(Strain, self).__init__(tseries.data,
tseries.delta_t,
tseries.start_time)
super(Strain, self).__init__(fseries.data,
fseries.delta_f,
fseries.epoch)
def window(self, window=tukey_window):
"""Window the (time-domain) strain using `window`."""
self.data *= window(len(self))
return self
def whiten(self, psd):
self._white = True
self.data = whiten(self, psd(self.freqs))
sel = self.data != 0.
self._bandpass_correction = np.sqrt(len(sel)/np.count_nonzero(sel))
def unwhiten(self, psd):
self._white = False
self.data = unwhiten(self, psd(self.freqs))
self._bandpass_correction = 1.
@property
def tilde(self):
"""The frequency-domain representation of the strain.
def tseries(self):
"""The time-domain representation of the strain.
:returns: A pycbc :class:`FrequencySeries`
:returns: A pycbc :class:`TimeSeries`
"""
fseries = self.to_frequencyseries()
tseries = self.to_timeseries()
if self._white:
#fseries.data /= 4 * np.sqrt(self.duration)
fseries.data /= self._bandpass_correction
fseries.data /= self.delta_t * np.sqrt(len(self)/2)
return fseries
tseries.data *= self._bandpass_correction
tseries.data *= tseries.delta_t * np.sqrt(len(tseries)/2)
return tseries
@property
def times(self):
"""Sample times of the strain"""
return np.array(self.sample_times)
return float(self.epoch) + self.delta_t*np.arange(2*(len(self) - 1))
@property
def duration(self):
"""Length of time (in seconds)"""
return 1/self.delta_f
@property
def sample_rate(self):
"""Sampling rate"""
return 1/self.delta_t
@property
def freqs(self):
"""Sample frequencies of the strain"""
return np.fft.rfftfreq(len(self), self.delta_t)
return np.array(self.sample_frequencies)
@property
def delta_f(self):
"""Frequency spacing of samples"""
return self.sample_rate/len(self)
def delta_t(self):
"""Time spacing of samples"""
ntimes = 2*(len(self) - 1)
return self.duration/ntimes
def copy(self):
"""Make a copy """
"""Make a copy"""
return Strain(self, white=self._white)
def __add__(self, strain):
......@@ -95,22 +116,28 @@ class Strain(TimeSeries):
result.data -= strain
return result
def window(strain, window=tukey_window):
"""Window the (time-domain) strain using `window`."""
return strain.copy().window(window)
def whiten_strain(series, psd):
"""Whiten `series` using the provided `psd`, resulting in
noise which is unit-normal distributed"""
strain = Strain(series)
strain.whiten(psd)
return strain
def whiten(series, psd, window=tukey_window):
"""Whiten `series` using the provided `psd` after applying `window` to the
time domain data"""
N = len(series)
dt = series.delta_t
def whiten(strain, psd):
"""Return a numpy array containing whitened frequency domain data"""
return np.sqrt(8.) * np.array(strain.data) / (np.sqrt(psd) * np.sqrt(strain.duration))
def unwhiten_strain(series, psd):
"""Unwhiten `series` using the provided `psd`"""
strain = Strain(series)
if window:
strain.window(window)
wf = strain.tilde
wf.data /= .5 * np.sqrt(psd(strain.freqs)) * np.sqrt(strain.duration)
return Strain(wf, white=True)
strain.unwhiten(psd)
return strain
def unwhiten(strain, psd):
"""Return a numpy array containing unwhitened frequency domain data"""
data = np.array(strain.data) * np.sqrt(psd) * np.sqrt(strain.duration)/np.sqrt(8.)
data[np.isnan(data) | np.isinf(data)] = 0.
return data
def resample(strain, srate):
"""Resample `strain` at `srate`"""
......
......@@ -277,8 +277,7 @@ def project_on_network(hplus, hcross, time, ra, dec, psi=0., sample_rate=16384.,
fplus, fcross = lal.ComputeDetAMResponse(lal.CachedDetectors[diff].response, ra, dec, psi, gmst)
hf = combine_and_timeshift(fplus, fcross, hplus, hcross, timeshift, sample_rate, duration)
hfseries = FrequencySeries(hf, 1./duration, epoch)
waveforms.append(Strain(hfseries))
waveforms.append(Strain(hf, delta_f=1./duration, epoch=epoch))
return waveforms
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment