def stamp_snr(channel1, channel2, stride):
"""
calculates stamp snr
Parameters
----------
channel1 : TimeSeries
channel1 timeseries
channel2 : TimeSeries
channel2 TimeSeries
Returns
-------
snr : Spectrogram
stamp snr spectrogram
"""
y = stamp_y(channel1, channel2, stride)
variance = stamp_variance(channel1, channel2, stride)
snr = Spectrogram(y.value / variance.value**0.5,
unit=None,
dt=y.dt,
f0=y.f0,
df=y.df,
epoch=y.epoch,
copy=True)
return snr
def plot_spectrogram_from_ts(ts):
plot = SpectrogramPlot()
ax = plot.gca()
ax.plot(Spectrogram(spec))
#pyplot.ylim(1e-9, 1e-2)
#pyplot.xlim(0.1, 500)
#pyplot.loglog()
pyplot.savefig("specgram.png", dpi=300)
pyplot.close()
def csdgram(channel1, channel2, stride, overlap=None, pad=False):
"""
calculates one-sided csd spectrogram between two timeseries
or fftgrams. Allows for flexibility for holding DARM
fftgram in memory while looping over others.
Parameters
----------
channel1 : TimeSeries or Spectrogram object
timeseries from channel 1
timeseries2 : TimeSeries or Spectrogram object
timeseries from channel 2
Returns
-------
csdgram : spectrogram object
csd spectrogram for two objects
"""
if isinstance(channel1, TimeSeries):
fftgram1 = fftgram(channel1, stride, pad=pad, overlap=overlap)
elif isinstance(channel1, Spectrogram):
fftgram1 = channel1
else:
raise TypeError('First arg is either TimeSeries or Spectrogram object')
if isinstance(channel2, TimeSeries):
fftgram2 = fftgram(channel2, stride, pad=pad, overlap=overlap)
elif isinstance(channel2, Spectrogram):
fftgram2 = channel2
else:
raise TypeError('First arg is either TimeSeries or Spectrogram object')
# clip off first 2 and last 2 segments to be consistent with psd
# calculation
out = np.conj(fftgram1.value) * fftgram2.value
if pad:
# if zero padded, take every other frequency
out = out[:, 0::2]
df = fftgram1.df * 2
f0 = fftgram1.f0 * 2
else:
df = fftgram1.df
f0 = fftgram1.f0
csdname = 'csd'
out = Spectrogram(out,
name=csdname,
epoch=fftgram1.epoch.value,
df=df,
dt=fftgram1.dt,
copy=True,
unit=fftgram1.unit * fftgram2.unit,
f0=f0)
return out
def plot_spectrogram_from_ts(ts, fname='specgram.png'):
'''
Plot spectrogram
'''
plot = SpectrogramPlot()
ax = plot.gca()
ax.plot(Spectrogram(spec))
#plt.ylim(1e-9, 1e-2)
#plt.xlim(0.1, 500)
#plt.loglog()
plt.savefig(fname, transparent=True)
plt.close()
def plot_spectrogram(spec, dt, df, ymax, t0, t1, fname="specgram.png"):
'''
Plot standard Fourier-based spectrogram
'''
plot = SpectrogramPlot()
ax = plot.gca()
ax.plot(Spectrogram(spec, dt=dt, df=df, epoch=float(t0)), cmap='viridis')
plot.add_colorbar(label='Amplitude')
plt.xlim(t0, t1)
plt.ylim(0, ymax)
plt.savefig(fname, transparent=True) #,dpi=300)
plt.close()
def create_matrix_from_file(coh_file, channels):
"""
Creates coherence matrix from data that's in a file.
Used typically as helper function for plotting
Parameters
----------
coh_file : str
File containing coherence data
channels : list (str)
channels to plot
Returns
-------
coh_matrix : Spectrogram object
coherence matrix in form of spectrogram object
returns automatically in terms of coherence SNRr -
coherence * N.
(not actually a spectrogram, though)
frequencies : numpy array
numpy array of frequencies associated with coherence matrix
labels : list (str)
labels for coherence matrix
N : int
Number of time segment used to create coherence spectra
"""
labels = []
counter = 0
f = h5py.File(coh_file, 'r')
# get number of averages
N = f['info'].value
channels = f['psd2s'].keys()
failed_channels = f['failed_channels'].value
print failed_channels
First = 1
for channel in channels:
if First:
# initialize matrix!
darm_psd = FrequencySeries.from_hdf5(
f['psd1'][f['psd1'].keys()[0]])
First = 0
coh_matrix = np.zeros((darm_psd.size, len(channels)))
if channel in failed_channels:
continue
data = FrequencySeries.from_hdf5(f['coherences'][channel])
labels.append(channel[3:-3].replace('_', '-'))
coh_matrix[:data.size, counter] = data
counter += 1
coh_matrix = Spectrogram(coh_matrix * N)
return coh_matrix, darm_psd.frequencies.value, labels, N
def spectrogram(self):
# check if times are continuous
# check if freqs are continuous
kwargs = {'name': 'magnetic data', 'channel': None, 'unit': units.nT}
if self.cont_times:
kwargs['t0'] = self.times[0]
kwargs['dt'] = self.times[2] - self.times[1]
else:
kwargs['times'] = self.times
if self.cont_freqs:
kwargs['f0'] = self.frequencies[0]
kwargs['df'] = self.frequencies[2] - self.frequencies[1]
else:
kwargs['frequencies'] = self.frequencies
return Spectrogram(self['data']['spectra'].value.squeeze(), **kwargs)
def plot_spectrogram(spec,
dt,
df,
sample_rate,
start_time,
end_time,
fname="specgram.png"):
plot = SpectrogramPlot()
ax = plot.gca()
ax.plot(Spectrogram(spec, dt=dt, df=df, epoch=start_time), cmap='viridis')
plot.add_colorbar(label='Amplitude')
pyplot.xlim(start_time, end_time)
pyplot.ylim(0, sample_rate / 2.)
pyplot.savefig(fname) #,dpi=300)
pyplot.close()
def _coherence_spectrogram(fftgram1,
fftgram2,
stride,
segmentDuration,
pad=False):
max_f = min(fftgram1.shape[1], fftgram2.shape[1])
csd12 = csdgram(fftgram1[:, 0:max_f],
fftgram2[:, 0:max_f],
stride,
pad=pad)
if pad:
psd1 = np.abs(fftgram1[:, 0::2])**2
psd2 = np.abs(fftgram2[:, 0::2])**2
else:
psd1 = np.abs(fftgram1)**2
psd2 = np.abs(fftgram2)**2
# average over strides if necessary
if not segmentDuration:
segmentDuration = stride
if segmentDuration < stride:
raise ValueError('segmentDuration must be longer than or equal\
to stride')
if segmentDuration % stride:
raise ValueError('stride must evenly divide segmentDuration')
navgs = segmentDuration / stride
if not fftgram1.shape[0] % navgs:
print 'WARNING: will cut off last segment because segmentDuration doesnt evenly divide total time window'
nsegs = int(fftgram1.shape[0] / navgs)
nfreqs = csd12.frequencies.size
coh_spec = np.zeros((nsegs, nfreqs))
for i in range(nsegs):
idx1 = i * navgs
idx2 = idx1 + navgs
coh_spec[i, :] = np.abs(np.mean(csd12[idx1:idx2, :], 0) ** 2)\
/ ((np.mean(psd1[idx1:idx2, :nfreqs], 0)) *
np.mean(psd2[idx1:idx2, :nfreqs], 0))
coh_spec = Spectrogram(coh_spec,
df=csd12.df,
dt=segmentDuration,
epoch=csd12.epoch,
unit=None,
name='coherence spectrogram')
return coh_spec
def fft_amp_spectrogram(self):
# vals = self['data']['fftamp'].value.squeeze()
# intermediate = vals['real'] + 1j * vals['imag']
intermediate = self['data']['fftamp'].value.squeeze()['real'] +\
1j * self['data']['fftamp'].value.squeeze()['imag']
# intermediate = np.zeros((self.times.size, self.frequencies.size))
# intermediate = np.zeros(vals.flatten().size)
kwargs = {'name': 'magnetic data', 'channel': None, 'unit': units.nT}
if self.cont_times:
kwargs['t0'] = self.times[0]
kwargs['dt'] = self.times[2] - self.times[1]
else:
kwargs['times'] = self.times
if self.cont_freqs:
kwargs['f0'] = self.frequencies[0]
kwargs['df'] = self.frequencies[2] - self.frequencies[1]
else:
kwargs['frequencies'] = self.frequencies
return Spectrogram(intermediate, **kwargs)
def cross_corr(self, other):
if npt.assert_almost_equal(self.frequencies, other.frequencies):
raise ValueError('frequencies must match')
newmat = np.zeros((self.spectrogram.shape))
this_specgram = self.fft_amp_spectrogram.copy()
other_specgram = other.fft_amp_spectrogram.copy()
final_times = []
if self.cont_times and other.cont_times:
et = int(
min(this_specgram.times[-1].value,
other_specgram.times[-1].value))
st = int(
max(this_specgram.times[0].value,
other_specgram.times[0].value))
this_specgram2 = this_specgram.crop(st, et)
other_specgram2 = other_specgram.crop(st, et)
newmat = this_specgram2.value * other_specgram2.value
final_times = this_specgram2.times.value
else:
for ii, time in enumerate(self.times):
# check if time appears in other specgram
if np.isin(time, other.times):
# find where it appears
other_idx = np.where(other.times == time)[0]
# do cross-corr
newmat[ii, :] = this_specgram[ii, :].value *\
other_specgram[other_idx, :].value
# add this to final times list
final_times.append(time)
# return
return Spectrogram(newmat.squeeze(),
dt=(final_times[2] - final_times[1]),
t0=final_times[0],
f0=this_specgram.f0,
df=this_specgram.df,
unit=this_specgram.unit,
name=this_specgram.name + other_specgram.name,
channel=None)
def path_to_contour(tf_path, times, frequencies, t_unc=0, f_unc=0):
"""
Convert a time-frequency path as a TimeSeries object to a Spectrogram.
Parameters
----------
tf_path : gwpy.TimeSeries object
Time-frequency path.
times : array
Time bins of output spectrogram
frequencies : array
Frequency bins of output spectrogram
t_unc : float, int, optional
Time uncertainty; widens path contour by this amount in time.
f_unc : float, int, optional
Frequency uncertainty; widens path contour by this amount in frequency.
Returns
-------
contour_path : gwpy.Spectrogram object
Spectrogram whose values are 1 in pixels overlapping the time-frequency path and 0 elsewhere.
"""
path_times = tf_path.times.value
path_frequencies = tf_path.value
values = np.zeros([times.size, frequencies.size])
for i in range(times.size - 1):
for j in range(frequencies.size - 1):
t_idx = np.where((path_times >= times[i] - t_unc)
& (path_times < times[i + 1] + t_unc))[0]
f_idx = np.where(
(path_frequencies[t_idx] >= frequencies[j] - f_unc)
& (path_frequencies[t_idx] < frequencies[j + 1] + f_unc))[0]
if len(f_idx) > 0:
values[i, j] = 1
contour_path = Spectrogram(values, times=times, frequencies=frequencies)
contour_path.name = tf_path.name
return contour_path
def stamp_variance(channel1, channel2, stride):
"""
calculates stamp-pem variance from two time-series.
Parameters
----------
channel1 : TimeSeries or Spectrogram object
timeseries or PSD for channel 1
channel2 : TimeSeries or Spectrogram object
timeseries or PSD for channel 2
stride : int
fft stride
Returns
-------
stamp variance : Spectrogram object
"""
# set units
if isinstance(channel1, TimeSeries):
psd1 = psdgram(channel1, stride)
else:
psd1 = channel1
if isinstance(channel2, TimeSeries):
psd2 = psdgram(channel2, stride)
else:
psd2 = channel2
# set units
if psd1.unit and psd2.unit:
var_unit = psd1.unit * psd2.unit
else:
var_unit = (u.Hz)**-2
variance = Spectrogram(0.5 * psd1.value * psd2.value,
epoch=psd1.epoch,
dt=psd1.dt,
copy=True,
unit=var_unit,
f0=psd1.f0,
df=psd1.df)
return variance
def get_channel_online_data(
channel,
st,
et,
format='spectrogram',
remove_nonlocked_times=False,
normalize_coherence=False,
config_file='/home/stochastic/config_files/ini_files/H1.ini'):
"""
Returns a list of PEMCoherenceSegment
objects.
Parameters
----------
channel : str
channel name you want to load
st : str or int
start time (in string format) or gps time
et : str or int
end time (in string format) or gps time
format : str, optional, default='spectrogram'
format to return. either spectrogram or seglist. Spectrogram returns a
`gwpy.spectrogram.Spectrogram` and seglist returns a list of
`stamp_pem.coherence_segment.PEMCoherenceSegment`.
remove_nonlocked_times: bool, optional, default=False
Removes non locked times from a spectrogram
normalize_coherence : bool, optional, default=False
Normalizes each column of spectrogram by the number of averages
Returns
-------
out : `gwpy.spectrogram.Spectrogram` or list of
`stamp_pem.coherence_segment.PEMCoherencSegment` objects
representation of data between start and end times for a given channel
"""
pipeline_dict = coh_io.read_pipeline_ini(config_file)
env_params, run_params = coh_io.check_ini_params(pipeline_dict)
channel_dict = ChannelDict.read(env_params['list'])
jobdur = int(env_params['job_duration'])
darm_channel = run_params['darm_channel']
basedir = env_params['base_directory']
if isinstance(st, str):
st = int(time.to_gps(st))
if isinstance(et, str):
et = int(time.to_gps(et))
starttimes = np.arange(st, et, jobdur)
subsys = get_channels_subsystem(channel, channel_dict)
seglist = []
for starttime in starttimes:
cohdir = coh_io.get_directory_structure(subsys, starttime, basedir)
cohfile = coh_io.create_coherence_data_filename(darm_channel,
subsys,
starttime,
starttime + jobdur,
directory=cohdir)
try:
subsystem_data = PEMSubsystem.read(subsys, cohfile)
except IOError:
print "No data found between %d and %d for %s" % (
starttime, starttime + jobdur, channel)
continue
if np.isnan(subsystem_data[channel].psd1.value[0]):
continue
seglist.append(subsystem_data[channel])
N = 1
if format == 'spectrogram':
if remove_nonlocked_times:
foundtimes = np.asarray(
[seglist[ii].starttime for ii in range(len(seglist))])
data = np.zeros((len(seglist), seglist[0].psd1.size))
for ii in range(len(seglist)):
if normalize_coherence:
N = seglist[ii].N
if seglist[ii].get_coh()[0] == np.nan:
continue
data[ii, :] = seglist[ii].get_coh() * N
specgram = Spectrogram(data,
epoch=foundtimes[0],
dt=jobdur,
df=seglist[0].psd1.df)
else:
foundtimes = np.asarray(
[seglist[ii].starttime for ii in range(len(seglist))])
count = 0
data = np.nan * np.zeros((starttimes.size, seglist[0].psd1.size))
for ii, starttime in enumerate(starttimes):
if np.any(foundtimes == starttime):
if normalize_coherence:
N = seglist[count].N
data[ii, :] = seglist[count].get_coh() * N
count += 1
specgram = Spectrogram(data,
dt=jobdur,
epoch=starttimes[0],
df=seglist[0].psd1.df)
return specgram
elif format == 'seglist':
return seglist
else:
raise ValueError('format needs to be "spectrogram" or "seglist"')
def generate_magnetometer_csd(ifoparams1, ifoparams2, generalparams,
stochparams):
from datetime import datetime
from astropy.time import Time
# get start time and duration from our paramfile
st = Time(datetime.strptime(generalparams['mag_day_start'], '%Y%m%d')).gps
dur = int(float(generalparams['ndays']) * 86400)
# general frame name
rawmag_fname = generalparams[
'output_prefix'] + '/correlated_mag_data/%s-RAW-MAG-DATA-%d-%d.gwf'
# specific frame names
rawmag_fname1 = rawmag_fname % (ifoparams1['name'], st, dur)
rawmag_fname2 = rawmag_fname % (ifoparams2['name'], st, dur)
# load them up
magts1 = MagTimeSeries.read(str(rawmag_fname1),
'%s:mag_data' % ifoparams1['name'])
magts2 = MagTimeSeries.read(str(rawmag_fname2),
'%s:mag_data' % ifoparams2['name'])
# take csd spectrogram between two magnetic timeseries
magcsd = magts1.csd_spectrogram(magts2,
float(stochparams['segdur']),
fftlength=float(stochparams['segdur']),
overlap=float(stochparams['segdur']) / 2.,
nproc=8)
# generate hdf5 file name for output
csd_fname = generalparams[
'output_prefix'] + '/correlated_mag_data/%s-%s-COMPLEX-CSD-%d-%d.hdf5'
# coarsegrain
# note there's a gwpy bug for slicing on rows right now
# it produces a FrequencySeries with the wrong metadata.
# we'll fix it on the fly
# TODO: open a GWpy issue about this problem
# coarsegrain once to get shape
ntimes = magcsd.shape[0]
spec1 = coarseGrain(FrequencySeries(magcsd[0, :].value,
df=magcsd.df,
f0=magcsd.f0),
flowy=float(stochparams['deltaf']),
deltaFy=float(stochparams['deltaf']))
# initialize array for speed
# make sure to tell it you want a complex array
# otherwise it'll automatically pick out
# real part of what you put into it
newarr = np.zeros((ntimes, spec1.size), dtype=complex)
print('COARSEGRAINING MAG DATA FOR STOCHASTIC USE')
# loop over each individual spectrum and coarse grain
for ii in range(ntimes):
newarr[ii, :] = coarseGrain(FrequencySeries(magcsd[ii, :].value,
df=magcsd.df,
f0=magcsd.f0),
flowy=float(stochparams['deltaf']),
deltaFy=float(stochparams['deltaf'])).value
# create final spectrogram object
final_csd = Spectrogram(newarr,
df=spec1.df,
f0=spec1.f0,
name=magcsd.name,
channel=magcsd.channel,
times=magcsd.times)
# write to an hdf5 file
final_csd.write(str(csd_fname %
(ifoparams1['name'], ifoparams2['name'], st, dur)),
overwrite=True)
# add final message so you know where to look after
print('WROTE CORRELATED SPECTROGRAM TO %s' %
str(csd_fname % (ifoparams1['name'], ifoparams2['name'], st, dur)))
def fftgram(timeseries, stride, pad=False):
"""
calculates fourier-gram with automatic
50% overlapping hann windowing.
Parameters
----------
timeseries : gwpy TimeSeries
time series to take fftgram of
stride : `int`
number of seconds in single PSD
Returns
-------
fftgram : gwpy spectrogram
a fourier-gram
"""
fftlength = stride
dt = stride
df = 1. / fftlength
# number of values in a step
stride *= timeseries.sample_rate.value
# number of steps
nsteps = 2 * int(timeseries.size // stride) - 1
# only get positive frequencies
if pad:
nfreqs = int(fftlength * timeseries.sample_rate.value)
df = df / 2
f0 = df
else:
nfreqs = int(fftlength * timeseries.sample_rate.value) / 2
dtype = np.complex
# initialize the spectrogram
out = Spectrogram(np.zeros((nsteps, nfreqs), dtype=dtype),
name=timeseries.name,
epoch=timeseries.epoch,
f0=df,
df=df,
dt=dt,
copy=True,
unit=1 / u.Hz**0.5,
dtype=dtype)
# stride through TimeSeries, recording FFTs as columns of Spectrogram
for step in range(nsteps):
# indexes for this step
idx = (stride / 2) * step
idx_end = idx + stride
stepseries = timeseries[idx:idx_end]
# zeropad, window, fft, shift zero to center, normalize
# window
stepseries = np.multiply(stepseries, np.hanning(stepseries.value.size))
# take fft
if pad:
stepseries = TimeSeries(np.hstack(
(stepseries, np.zeros(stepseries.size))),
name=stepseries.name,
x0=stepseries.x0,
dx=timeseries.dx)
tempfft = stepseries.fft(stepseries.size)
else:
tempfft = stepseries.fft(stepseries.size)
tempfft.override_unit(out.unit)
out[step] = tempfft[1:]
return out
def psdgram(timeseries, stride):
"""
calculates one-sided PSD from timeseries
properly using welch's method by averaging adjacent non-ovlping
segments. Since default fftgram overlaps segments
we have to be careful here...
Parameters
----------
fftgram : Spectrogram object
complex fftgram
adjacent : `int`
number of adjacent segments
to calculate PSD of middle segment
Returns
-------
psdgram : Spectrogram object
psd spectrogram calculated in
spirit of STAMP/stochastic
"""
fftlength = stride
dt = stride
df = 1. / fftlength
# number of values in a step
stride *= timeseries.sample_rate.value
# number of steps
nsteps = 2 * int(timeseries.size // stride) - 1
# only get positive frequencies
nfreqs = int(fftlength * timeseries.sample_rate.value) / 2.
# initialize the spectrogram
if timeseries.unit:
unit = timeseries.unit / u.Hz
else:
unit = 1 / u.Hz
out = Spectrogram(np.zeros((nsteps, int(nfreqs))),
name=timeseries.name,
epoch=timeseries.epoch,
f0=df / 2,
df=df,
dt=dt,
copy=True,
unit=unit)
# stride through TimeSeries, recording FFTs as columns of Spectrogram
for step in range(nsteps):
# indexes for this step
idx = (stride / 2) * step
idx_end = idx + stride
stepseries = timeseries[idx:idx_end]
steppsd = stepseries.psd()[1:]
out.value[step, :] = steppsd.value
psdleft = np.hstack((out.T, np.zeros((out.shape[1], 4))))
psdright = np.hstack((np.zeros((out.shape[1], 4)), out.T))
psd_temp = ((psdleft + psdright) / 2).T
# create spectrogram object. multiply by 2 for one-sided.
psd = Spectrogram(2 * psd_temp.value[4:-4],
name=timeseries.name,
epoch=timeseries.epoch.value + 2 * dt,
f0=df,
df=df,
dt=dt,
copy=True,
unit=unit)
return psd
def csdgram(channel1, channel2, stride):
"""
calculates one-sided csd spectrogram between two timeseries
or fftgrams. Allows for flexibility for holding DARM
fftgram in memory while looping over others.
Parameters
----------
channel1 : TimeSeries or Spectrogram object
timeseries from channel 1
timeseries2 : TimeSeries or Spectrogram object
timeseries from channel 2
Returns
-------
csdgram : spectrogram object
csd spectrogram for two objects
"""
if isinstance(channel1, TimeSeries):
fftgram1 = fftgram(channel1, stride, pad=True)
elif isinstance(channel1, Spectrogram):
fftgram1 = channel1
else:
raise TypeError('First arg is either TimeSeries or Spectrogram object')
if isinstance(channel2, TimeSeries):
fftgram2 = fftgram(channel2, stride, pad=True)
elif isinstance(channel2, Spectrogram):
fftgram2 = channel2
else:
raise TypeError('First arg is either TimeSeries or Spectrogram object')
# clip off first 2 and last 2 segments to be consistent with psd
# calculation
out = (fftgram1.value * np.conj(fftgram2.value))[2:-2]
csdname = 'csd spectrogram between %s and %s' % (fftgram1.name,
fftgram2.name)
out = Spectrogram(out,
name=csdname,
epoch=fftgram1.epoch.value + 2 * fftgram1.dt.value,
df=fftgram1.df,
dt=fftgram1.dt,
copy=True,
unit=fftgram1.unit * fftgram2.unit,
f0=fftgram1.f0)
df = fftgram1.df.value * 2
f0 = fftgram1.f0.value * 2
csdgram = Spectrogram(np.zeros((out.shape[0], out.shape[1] / 2),
dtype=np.complex),
df=df,
dt=fftgram1.dt,
copy=True,
unit=out.unit,
f0=f0,
epoch=out.epoch)
for ii in range(csdgram.shape[0]):
# multiply by 2 for one-sided spectrum
temp = Spectrum(2 * out.data[ii],
df=out.df,
f0=out.f0,
epoch=out.epoch,
unit=out.unit)
N = out.shape[1] / 2
csdgram[ii] = coarseGrain(temp, df, f0, N)
return csdgram
def fftgram(timeseries, stride, overlap=None, pad=False):
"""
calculates fourier-gram with auto 50% overlapping
segments and hann windowed
Parameters
----------
timeseries : :py:mod:`gwpy.timeseries.TimeSeries`
time series to take fftgram of
stride : `int`
number of seconds in single PSD
overalp : `int`
number of seconds of overlap. Defaults to half of stride.
Returns
-------
fftgram : :py:mod:`gwpy.spectrogram.Spectrogram`
a fourier-gram
"""
fftlength = stride
if overlap is None:
overlap = stride / 2.
dt = stride - overlap
df = 1. / fftlength
# number of values in a step
stride *= timeseries.sample_rate.value
overlap *= timeseries.sample_rate.value
step = stride - overlap
# number of steps
nsteps = int(timeseries.size // step) - 1
# only get positive frequencies
if pad:
nfreqs = int(fftlength * timeseries.sample_rate.value)
df = df / 2
f0 = df
else:
nfreqs = int(fftlength * timeseries.sample_rate.value) / 2
dtype = np.complex
# initialize the spectrogram
out = Spectrogram(np.zeros((nsteps, nfreqs), dtype=dtype),
name=timeseries.name,
epoch=timeseries.epoch,
f0=df,
df=df,
dt=dt,
copy=True,
unit=u.Hz**(-0.5),
dtype=dtype)
# stride through TimeSeries, recording FFTs as columns of Spectrogram
if (timeseries.sample_rate.value * stride) % 1:
raise ValueError('1/stride must evenly divide sample rate of channel')
for step in range(nsteps):
# indexes for this step
idx = int((stride / 2) * step)
idx_end = int(idx + stride)
stepseries = timeseries[idx:idx_end]
# zeropad, window, fft
stepseries = np.multiply(stepseries, np.hanning(stepseries.value.size))
# hann windowing normalization
norm = np.sum(np.hanning(stepseries.value.size))
if pad:
stepseries = TimeSeries(np.hstack(
(stepseries, np.zeros(stepseries.size))),
name=stepseries.name,
x0=stepseries.x0,
dx=timeseries.dx)
tempfft = stepseries.fft(stepseries.size)
else:
tempfft = stepseries.fft(stepseries.size)
# reset proper unit
tempfft.override_unit(out.unit)
# normalize
out[step] = tempfft[1:] / norm
return out
def coherence_spectrogram(channel1,
channel2,
stride,
segmentDuration,
overlap=None,
pad=False,
st=None,
et=None):
"""
Calculates coherence spectrogram for two channels.
Can be called in two ways: either channel 1 and
channel 2 are channel names or they are time series.
If they are strings, then each segment is loaded
separately as a way to save memory.
TODO: This function is way too complicated. Just give it a single
calling method.
Parameters
----------
channel 1 : `str` or TimeSeries object
channel 1 name or time series
channel 2 : `str` or TimeSeries object
channel 2 name or time series
stride : `int`
stride for individual ffts (in seconds)
segmentDuration : `int`
duration of segments in spectrogram (in seconds)
overlap : `int`, optional
amount of overlap between ffts (in seconds)
pad : `bool`, optional
decide whether or not to pad ffts for taking csd
frames : `bool`, optional
decide whether to use frames or nds2 to load data
st : `int`, optional
start time if we're loading channel 1 and channel 2
segments on the fly
et : `int`, optional
end time if we're loading channel 1 and channel 2 on the fly
Returns
-------
coherence_spectrogram : Spectrogram object
coherence spectrogram between channels 1 and 2
psd1_spectrogram : Spectrogram object
channel 1 psd spectrogram
psd2_spectrogram : Spectrogram object
channel 2 psd spectrogram
N : `int`
number of averages used for each pixel
"""
if isinstance(channel1, TimeSeries):
nsegs = (channel1.times.value[-1] - channel1.times.value[0]) \
/ segmentDuration
nfreqs = min(channel1.sample_rate.value,
channel2.sample_rate.value) * stride * 0.5
epoch = channel1.epoch
elif isinstance(channel1, str):
nsegs = int((et - st) / segmentDuration)
test1 = _read_data(channel1, st, st + 1)
test2 = _read_data(channel2, st, st + 1)
nfreqs = int(
min(test1.sample_rate.value, test2.sample_rate.value) * (stride) *
0.5)
epoch = test1.epoch
df = 1. / stride
coherence_spectrogram = Spectrogram(np.zeros((nsegs, nfreqs)),
epoch=epoch,
dt=segmentDuration,
df=df,
f0=df)
psd1_spectrogram = Spectrogram(np.zeros((nsegs, nfreqs)),
epoch=epoch,
dt=segmentDuration,
df=df,
f0=df,
unit=u.Hz**-1)
psd2_spectrogram = Spectrogram(np.zeros((nsegs, nfreqs)),
epoch=epoch,
dt=segmentDuration,
df=df,
f0=df,
unit=u.Hz**-1)
if isinstance(channel1, str) and isinstance(channel2, str):
for i in range(int(nsegs)):
startTime = st + i * segmentDuration
endTime = startTime + segmentDuration
stepseries1 = _read_data(
channel1,
startTime,
endTime,
)
stepseries2 = _read_data(
channel2,
startTime,
endTime,
)
test, csd12, psd1, psd2, N = coherence(stepseries1,
stepseries2,
stride,
overlap=None,
pad=pad)
coherence_spectrogram[i] = test
psd1_spectrogram[i] = psd1
psd2_spectrogram[i] = psd2
else:
samples1 = segmentDuration * channel1.sample_rate.value
samples2 = segmentDuration * channel2.sample_rate.value
for i in range(int(nsegs)):
idx_start1 = samples1 * i
idx_end1 = idx_start1 + samples1
idx_start2 = samples2 * i
idx_end2 = idx_start2 + samples2
stepseries1 = channel1[idx_start1:idx_end1]
stepseries2 = channel2[idx_start2:idx_end2]
test, csd12, psd1, psd2, N = coherence(stepseries1,
stepseries2,
stride,
overlap=None,
pad=pad)
coherence_spectrogram[i] = test
psd1_spectrogram[i] = psd1
psd2_spectrogram[i] = psd2
return coherence_spectrogram, psd1_spectrogram, psd2_spectrogram, N
