def setup_class(self):
# Chose signals that are not commensurate with quarter-sample offsets,
# or with the frames.
f_signal = (self.full_sample_rate * 2 / self.samples_per_full_frame *
np.array([31.092, 65.1234]))
cosine = PureTone(f_signal, self.start_time, np.pi * u.deg)
# Create a stream that just contains the two tones, and one sampled
# at 4 times lower rate.
self.full_fh = StreamGenerator(
cosine,
shape=((self.samples_per_full_frame * self.n_frames, ) +
self.sideband.shape),
sample_rate=self.full_sample_rate,
samples_per_frame=self.samples_per_full_frame,
frequency=self.frequency,
sideband=self.sideband,
start_time=self.start_time,
dtype=self.dtype)
self.part_fh = StreamGenerator(
cosine,
shape=((self.samples_per_full_frame // 4 * self.n_frames, ) +
self.sideband.shape),
sample_rate=self.full_sample_rate / 4,
samples_per_frame=self.samples_per_full_frame // 4,
frequency=self.frequency,
sideband=self.sideband,
start_time=self.start_time,
dtype=self.dtype)
def test_real_to_complex_delta():
"""Test converting a real delta function to complex."""
def real_delta(handle):
real_delta = np.zeros(handle.samples_per_frame, dtype=np.float64)
if handle.offset == 0:
real_delta[0] = 1.0
return real_delta
delta_fh = StreamGenerator(real_delta,
samples_per_frame=1024,
start_time=Time('2010-11-12T13:14:15'),
sample_rate=1. * u.kHz,
frequency=400 * u.kHz,
sideband=1,
shape=(2048, ),
dtype='f8')
real_data = delta_fh.read()
assert real_data[0] == 1.
assert np.all(real_data[1:] == 0.)
complex_delta = np.zeros(2048 // 2, dtype=np.complex128)
complex_delta[0] = 1.0
real2complex = Real2Complex(delta_fh)
complex_signal = real2complex.read()
assert complex_signal.shape == (1024, )
assert np.iscomplexobj(complex_signal)
assert np.isclose(complex_signal, complex_delta).all()
assert real2complex.frequency == 400.5 * u.kHz
assert real2complex.sideband == 1
r = repr(real2complex)
assert r.startswith('Real2Complex(ih)')
def setup_class(self):
self.shape = (1000, 5, 3)
self.ih = StreamGenerator(self.make_arange_data,
self.shape,
Time('2010-11-12'),
1. * u.Hz,
samples_per_frame=100,
dtype=float)
self.ih.intensity = 2
self.ih.non_uniform_axis = 1
def test_real_to_complex_sine(f_nyquist):
"""Test converting a real sine function to complex."""
def real_sine(handle):
real_sine = np.sin(f_nyquist * np.pi *
np.arange(handle.samples_per_frame))
return real_sine
sine_fh = StreamGenerator(real_sine,
samples_per_frame=1024,
start_time=Time('2010-11-12T13:14:15'),
sample_rate=1. * u.kHz,
frequency=400 * u.kHz,
sideband=-1,
shape=(2048, ),
dtype='f8')
f_complex = f_nyquist - 0.5
complex_dc = np.exp(2j * np.pi *
(-0.25 + np.arange(2048 // 2) * f_complex))
real2complex = Real2Complex(sine_fh)
complex_signal = real2complex.read()
assert complex_signal.shape == (1024, )
assert np.iscomplexobj(complex_signal)
assert_allclose(complex_signal, complex_dc, atol=1e-8)
assert real2complex.frequency == 399.5 * u.kHz
assert real2complex.sideband == -1
def test_sample_slice(self):
sh = StreamGenerator(self.my_source,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
samples_per_frame=1)
sliced = sh[-400:]
assert sliced.shape == (400, ) + sh.sample_shape
assert abs(sliced.stop_time - sh.stop_time) < 1. * u.ns
assert abs(sliced.start_time -
(sh.stop_time - 400 / sh.sample_rate)) < 1. * u.ns
sh.seek(-400, 2)
expected = sh.read()
data = sliced.read()
assert np.all(data == expected)
r = repr(sliced)
assert r.startswith('GetSlice(ih, item=')
assert '\nih: StreamGenerator(' in r
def test_exceptions(self):
with StreamGenerator(self.my_source,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
samples_per_frame=1) as sh:
with pytest.raises(EOFError):
sh.seek(-10, 2)
sh.read(20)
with pytest.raises(AttributeError):
sh.frequency
with pytest.raises(AttributeError):
sh.sideband
with pytest.raises(AttributeError):
sh.polarization
with pytest.raises(ValueError):
StreamGenerator(self.my_source,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
sideband=np.ones((2, 2), dtype='i1'))
def test_frequency_sideband_setting(self):
frequency = np.array([320., 350., 380., 410.])[:, np.newaxis] * u.MHz
sideband = np.array([-1, 1])
with StreamGenerator(self.my_source,
frequency=frequency,
sideband=sideband,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
samples_per_frame=1) as sh:
assert np.all(sh.frequency == frequency)
assert np.all(sh.sideband == sideband)
with pytest.raises(AttributeError):
sh.polarization
def test_frequency_sideband_polarization_setting(self):
frequency = np.array([320., 320., 350., 350.])[:, np.newaxis] * u.MHz
sideband = np.array([-1, 1, -1, 1])[:, np.newaxis]
polarization = np.array(['X', 'Y'])
with StreamGenerator(self.my_source,
polarization=polarization,
frequency=frequency,
sideband=sideband,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
samples_per_frame=1) as sh:
assert np.all(sh.frequency == frequency)
assert np.all(sh.sideband == sideband)
assert np.all(sh.polarization == polarization)
def setup_class(self):
self.full_shape = ((self.samples_per_full_frame * self.n_frames, ) +
self.sideband.shape)
self.downsample = (16 if self.dtype.kind == 'c' else 8)
self.sample_rate = self.full_sample_rate / self.downsample
# Create the IF (which can produce a complex tone for quadrature).
self.mixer = PureTone(self.lo, self.start_time, self.phi0_mixer)
# Create a real-valued stream with a test-specific signal.
self.raw = StreamGenerator(
self.signal,
shape=self.full_shape,
start_time=self.start_time,
sample_rate=self.full_sample_rate,
dtype=np.dtype('f4'),
samples_per_frame=self.samples_per_full_frame)
def setup_class(self):
self.start_time = Time('2010-11-12T13:14:15')
self.sample_rate = 128. * u.kHz
self.shape = (164000, 2)
self.gp_sample = 64000
# Complex timestream
self.gp = StreamGenerator(self.make_giant_pulse,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
samples_per_frame=1000,
dtype=np.complex64,
frequency=300 * u.MHz,
sideband=np.array((1, -1)))
# Time delay of 0.05 s over 128 kHz band.
self.dm = DispersionMeasure(1000. * 0.05 / 0.039342251)
def setup(self):
self.start_time = Time('2010-11-12T13:14:15')
self.sample_rate = 128. * u.kHz
self.shape = (164000, 2)
self.gp_sample = 64000
# Real timestream; mean frequecies of the two bands are the same.
self.gp = StreamGenerator(self.make_giant_pulse,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
samples_per_frame=1000,
dtype=np.float32,
frequency=300. * u.MHz,
sideband=np.array((1, -1)))
# Time delay of 0.05 s over 128 kHz band.
self.dm = DispersionMeasure(1000. * 0.05 / 0.039342251)
def test_repr(self):
frequency = np.array([320., 350., 380., 410.])[:, np.newaxis] * u.MHz
sideband = np.array([-1, 1])
with StreamGenerator(self.my_source,
frequency=frequency,
sideband=sideband,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
samples_per_frame=1) as sh:
r = repr(sh)
assert r.startswith('StreamGenerator(')
assert 'start_time=' in r
assert 'samples_per_frame' not in r # has default
assert 'frequency=' in r
assert 'polarization' not in r
def test_basics(self):
with StreamGenerator(self.my_source,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
samples_per_frame=1) as sh:
assert sh.size == np.prod(self.shape)
assert sh.shape == self.shape
assert isinstance(sh.dtype, np.dtype)
assert sh.dtype == np.dtype('c8')
assert sh.samples_per_frame == 1
assert abs(sh.stop_time - sh.start_time - 1. * u.s) < 1. * u.ns
sh.seek(980)
data1 = sh.read(1)
assert data1.dtype == sh.dtype
assert np.all(data1 == 980.)
data2 = sh.read()
assert data2.shape == (1000 - 981, 4, 2)
assert np.all(data2 == np.arange(981, 1000).reshape(19, 1, 1))
def test_zeros_generation(self):
def generate_zero(sh):
return np.zeros((sh.samples_per_frame, ) + sh.shape[1:], sh.dtype)
with StreamGenerator(generate_zero,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
samples_per_frame=10) as sh:
assert sh.size == np.prod(self.shape)
assert sh.shape == self.shape
assert sh.samples_per_frame == 10
assert abs(sh.stop_time - sh.start_time - 1. * u.s) < 1. * u.ns
sh.seek(10)
data1 = sh.read(2)
assert data1.dtype == sh.dtype == np.dtype('c8')
assert np.all(data1 == 0.)
data2 = sh.read()
assert data2.shape == (1000 - 12, 4, 2)
assert np.all(data2 == 0.)
def test_disperse(self):
gp = StreamGenerator(self.make_giant_pulse,
shape=self.shape,
start_time=self.start_time,
sample_rate=self.sample_rate,
samples_per_frame=1000,
dtype=np.float32,
frequency=300. * u.MHz,
sideband=np.array((1, -1)))
disperse = Disperse(gp, self.dm)
assert_quantity_allclose(disperse.reference_frequency, 300. * u.MHz)
disperse.seek(self.start_time + self.gp_sample / self.sample_rate)
disperse.seek(-self.gp_sample // 2, 1)
around_gp = disperse.read(self.gp_sample)
assert around_gp.dtype == np.float32
p = (around_gp**2).reshape(-1, self.gp_sample // 20, 2).sum(1)
# Note: FT leakage means that not everything outside of the dispersed
# pulse is zero. But the total power there is small.
assert np.all(p[:9] < 0.006)
assert np.all(p[11:] < 0.006)
# Lower sideband [1] has lower frequencies and thus is dispersed
# to later.
assert p[9, 0] > 0.99 and p[10, 0] < 0.006
assert p[10, 1] > 0.99 and p[9, 1] < 0.006
class TestResampleReal:
"""Tests for resampling a signal.
Here used for a real signal, but subclassed below for complex.
"""
dtype = np.dtype('f4')
full_sample_rate = 1 * u.kHz
samples_per_full_frame = 4096 # Per full frame.
start_time = Time('2010-11-12T13:14:15')
frequency = 400. * u.kHz
sideband = np.array([-1, 1])
n_frames = 3
pad = 32 # Size of response = 2*pad + 1.
atol = 7e-4 # Tolerance within which we expect to reproduce signal.
@classmethod
def setup_class(self):
# Chose signals that are not commensurate with quarter-sample offsets,
# or with the frames.
f_signal = (self.full_sample_rate * 2 / self.samples_per_full_frame *
np.array([31.092, 65.1234]))
cosine = PureTone(f_signal, self.start_time, np.pi * u.deg)
# Create a stream that just contains the two tones, and one sampled
# at 4 times lower rate.
self.full_fh = StreamGenerator(
cosine,
shape=((self.samples_per_full_frame * self.n_frames, ) +
self.sideband.shape),
sample_rate=self.full_sample_rate,
samples_per_frame=self.samples_per_full_frame,
frequency=self.frequency,
sideband=self.sideband,
start_time=self.start_time,
dtype=self.dtype)
self.part_fh = StreamGenerator(
cosine,
shape=((self.samples_per_full_frame // 4 * self.n_frames, ) +
self.sideband.shape),
sample_rate=self.full_sample_rate / 4,
samples_per_frame=self.samples_per_full_frame // 4,
frequency=self.frequency,
sideband=self.sideband,
start_time=self.start_time,
dtype=self.dtype)
def test_setup(self):
self.full_fh.seek(0)
self.part_fh.seek(0)
full = self.full_fh.read()
part = self.part_fh.read()
assert_allclose(part, full[::4], atol=1e-8, rtol=0)
@pytest.mark.parametrize('offset',
(34, 34.5, 35.75, 50. * u.ms, 0.065 * u.s,
Time('2010-11-12T13:14:15.073')))
def test_resample(self, offset):
# Offsets equal to quarter samples to allow check with full_fh.
ih = Resample(self.part_fh, offset, pad=self.pad)
# Always lose 2 * pad per frame.
assert ih.shape[0] == self.part_fh.shape[0] - 2 * self.pad
assert ih.sample_shape == self.part_fh.sample_shape
# Check we are at the given offset.
if isinstance(offset, Time):
expected_time = offset
elif isinstance(offset, u.Quantity):
expected_time = self.part_fh.start_time + offset
else:
expected_time = (self.part_fh.start_time +
offset / self.part_fh.sample_rate)
assert abs(ih.time - expected_time) < 1. * u.ns
part_fh_offset = seek_float(self.part_fh, offset)
ioffset = round(part_fh_offset)
fraction = part_fh_offset - ioffset
assert ih.offset + self.pad == ioffset
expected_start_time = (self.part_fh.start_time + (
(self.pad + fraction) / self.part_fh.sample_rate))
assert abs(ih.start_time - expected_start_time) < 1. * u.ns
# Check that data is correctly resampled.
ih.seek(0)
data = ih.read()
# Find corresponding fully sampled data.
self.full_fh.seek(ih.start_time)
# Check time: should be exact, since our offsets are quarter samples.
assert np.abs(self.full_fh.time - ih.start_time) < 1. * u.ns
expected = self.full_fh.read(data.shape[0] * 4)[::4]
assert_allclose(data, expected, atol=self.atol, rtol=0)
def test_repr(self):
ih = Resample(self.part_fh, 0.5, samples_per_frame=511)
r = repr(ih)
assert r.startswith('Resample(ih')
assert 'offset=0.5' in r
@pytest.mark.parametrize('shift',
(0., 0.25, -5.25, [1.75, 10.25], [-1., 13] * u.ms)
)
@pytest.mark.parametrize('offset', (None, 0, 0.25))
def test_shift_and_resample(self, shift, offset):
# Shifts and offsets at quarter samples to allow check with full_fh.
ih = ShiftAndResample(self.part_fh, shift, offset=offset, pad=self.pad)
# start_time should be at expected offset from old grid.
expected_offset = seek_float(
self.part_fh, offset if offset is not None else np.mean(shift))
d_off = ((ih.start_time - self.start_time) * ih.sample_rate -
expected_offset).to_value(u.one)
assert abs(d_off - np.around(d_off)) < u.ns * ih.sample_rate
expected_length = (self.part_fh.shape[0] - 2 * self.pad -
seek_float(self.part_fh, np.ptp(shift)))
assert abs(ih.shape[0] - expected_length) <= 0.5
# Data should be shifted by the expected amounts.
shift = np.atleast_1d(shift)
for i, s in enumerate(shift):
ih.seek(0)
time_shift = seek_float(self.part_fh, s) / ih.sample_rate
fh_pos = self.full_fh.seek(ih.time - time_shift)
assert fh_pos >= 0
assert abs(ih.time - time_shift - self.full_fh.time) < 1. * u.ns
data = ih.read()
expected = self.full_fh.read(len(data) * 4)[::4]
sel = i if shift.size > 1 else Ellipsis
assert_allclose(data[:, sel],
expected[:, sel],
atol=self.atol,
rtol=0)
class TestShiftSamples:
@classmethod
def make_arange_data(self, ih):
test_data = (np.arange(
ih.offset, ih.offset +
ih.samples_per_frame).reshape((ih.samples_per_frame, ) +
(1, ) * len(ih.shape[1:])))
new_shape = (ih.samples_per_frame, ) + ih.shape[1:]
return np.broadcast_to(test_data, new_shape)
@classmethod
def make_non_uniform_arange_data(self, ih):
axis = ih.non_uniform_axis
data = self.make_arange_data(ih)
multiplier = np.arange(1, data.shape[axis] + 1) * ih.intensity
return data * multiplier.reshape((data.shape[axis], ) + (1, ) *
(data.ndim - 1 - axis))
@classmethod
def setup_class(self):
self.shape = (1000, 5, 3)
self.ih = StreamGenerator(self.make_arange_data,
self.shape,
Time('2010-11-12'),
1. * u.Hz,
samples_per_frame=100,
dtype=float)
self.ih.intensity = 2
self.ih.non_uniform_axis = 1
@pytest.mark.parametrize('start, n', [(0, 5), (90, 20)])
def test_shift_back(self, start, n):
shift_axis = 1
shift = np.arange(-self.shape[shift_axis] + 1, 1)
assert shift.max() == 0
shifter = ShiftSamples(self.ih,
shift.reshape(-1, 1),
samples_per_frame=100)
assert shifter.start_time == self.ih.start_time
shifter.seek(start)
shifted = shifter.read(n)
self.ih.seek(start)
raw_data = self.ih.read(100)
for i, sf in enumerate(-shift):
assert np.all(shifted[:, i] == raw_data[sf:sf + n, i])
@pytest.mark.parametrize('start, n', [(0, 5), (100, 20)])
def test_shift_both(self, start, n):
shift = np.array([-2, 0, 3])
shifter = ShiftSamples(self.ih, shift, samples_per_frame=100)
assert abs(shifter.start_time - 3 / self.ih.sample_rate -
self.ih.start_time) < 1. * u.ns
shifter.seek(start)
shifted = shifter.read(n)
self.ih.seek(start)
raw_data = self.ih.read(100)
for i, sf in enumerate(3 - shift):
assert np.all(shifted[:, :, i] == raw_data[sf:sf + n, :, i])
def test_compare_with_shift_and_resample(self):
shift = np.array([-2, 1, 4])
shifter = ShiftSamples(self.ih, shift, samples_per_frame=100)
resampler = ShiftAndResample(self.ih,
shift,
offset=0,
pad=32,
samples_per_frame=200)
# Note: resampler has larger padding, so start time is later.
shifter.seek(90)
resampler.seek(shifter.time)
# integer shifts, so time should be same.
assert abs(resampler.time - shifter.time) < 1. * u.ns
shifted = shifter.read(20)
resampled = resampler.read(20)
assert_allclose(shifted, resampled)
@pytest.mark.parametrize(
'fshift, ishift', [(np.array([1., 2., 3.25]), [1, 2, 3]),
(np.array([[-1.9], [-5.], [5.25], [3.49], [-1.2]]),
np.reshape([-2, -5, 5, 3, -1], (-1, 1)))])
def test_non_integer_shift(self, fshift, ishift):
shifter1 = ShiftSamples(self.ih, fshift)
comparison = ShiftSamples(self.ih, ishift)
assert np.all(shifter1._shift == comparison._shift)
data1 = shifter1.read()
expected = comparison.read()
assert np.all(data1 == expected)
shifter2 = ShiftSamples(self.ih, fshift / self.ih.sample_rate)
assert np.all(shifter2._shift == comparison._shift)
data2 = shifter2.read()
assert np.all(data2 == expected)
def test_wrong_shape(self):
with pytest.raises(ValueError, match='broadcast to sample shape'):
ShiftSamples(self.ih, np.array([[1], [2]]))
