def test_combine_callable(self):
def doublestack(data):
# concatenate frequencies, stack polarizations.
data = [np.stack(data[i:i + 2], axis=1) for i in range(0, 8, 2)]
return np.stack(data, axis=1)
fhs = [GetItem(self.fh, i) for i in range(8)]
ch = CombineStreams(fhs, doublestack)
expected = Reshape(self.fh, (4, 2))
assert ch.shape == expected.shape
assert ch.start_time == expected.start_time
assert ch.sample_rate == expected.sample_rate
assert ch.dtype == expected.dtype
assert_array_equal(ch.frequency, expected.frequency)
assert_array_equal(ch.sideband, expected.sideband)
assert_array_equal(ch.polarization, expected.polarization)
data = ch.read()
expected_data = expected.read()
assert_array_equal(data, expected_data)
r = repr(ch)
assert r.startswith('CombineStreams(ihs')
assert 'ihs: 8 streams' in r
ch.close()
expected.close()
def test_reshape(self, sample_shape):
fh = self.fh
ref_data = fh.read().reshape((-1, ) + sample_shape)
rt = Reshape(fh, sample_shape=sample_shape)
assert fh.sample_shape == (8, )
assert rt.sample_shape == sample_shape
assert rt.start_time == fh.start_time
assert rt.sample_rate == fh.sample_rate
data = rt.read()
assert_array_equal(data, ref_data)
r = repr(rt)
assert r.startswith(f'Reshape(ih, sample_shape={sample_shape})')
def test_specific3(self):
"""Selecting from both axes of two-dimensional samples."""
rh = Reshape(self.fh, (4, 2)) # freq, pol
gih = GetItem(rh, (slice(None, 2), 0))
assert gih.frequency.shape == (2, )
assert_array_equal(gih.frequency, rh.frequency[:2].squeeze())
assert gih.polarization.shape == ()
assert_array_equal(gih.polarization, rh.polarization[0])
def test_specific2(self):
"""Selecting one item in second axis of two-dimensional samples."""
rh = Reshape(self.fh, (4, 2)) # freq, pol
gih = GetItem(rh, (slice(None), 1))
assert gih.frequency.shape == (4, )
assert_array_equal(gih.frequency, rh.frequency.squeeze())
assert gih.polarization.shape == ()
assert_array_equal(gih.polarization, rh.polarization[1])
def test_specific1(self):
"""Selecting one item in first axis of two-dimensional samples."""
rh = Reshape(self.fh, (4, 2)) # freq, pol
gih = GetItem(rh, 0)
assert gih.frequency.shape == ()
assert_array_equal(gih.frequency, rh.frequency[0, 0])
assert gih.polarization.shape == (2, )
assert_array_equal(gih.polarization, rh.polarization)
def test_frequency_sideband_polarization_propagation1(self):
fh = self.fh
rt = Reshape(fh, (4, 2))
assert rt.frequency.shape == (4, 1)
assert np.all(rt.frequency == fh.frequency[::2].reshape(4, 1))
assert rt.sideband.shape == ()
assert np.all(rt.sideband == 1)
assert rt.polarization.shape == (2, )
assert np.all(rt.polarization == fh.polarization[:2])
def test_example(self):
"""Selecting time and from both axes of two-dimensional samples."""
rh = Reshape(self.fh, (4, 2)) # freq, pol
gsh = GetSlice(rh, (slice(10, -10), slice(None, 2), 0))
assert gsh.shape == (rh.shape[0] - 20, 2)
assert abs(gsh.start_time -
(rh.start_time + 10 / rh.sample_rate)) < 1. * u.ns
assert gsh.frequency.shape == (2, )
assert_array_equal(gsh.frequency, rh.frequency[:2].squeeze())
assert gsh.polarization.shape == ()
assert_array_equal(gsh.polarization, rh.polarization[0])
def test_frequency_sideband_polarization_propagation2(self):
# Add different frequency, sideband, and polarization information.
# (Note: these are incorrect; just for testing purposes.)
fh = SetAttribute(self.fh,
frequency=311.25 * u.MHz +
(np.arange(8.) // 4) * 16. * u.MHz,
sideband=np.tile([-1, 1], 4),
polarization=np.tile(['L', 'L', 'R', 'R'], 2))
rt = Reshape(fh, (2, 2, 2))
assert rt.frequency.shape == (2, 1, 1)
assert np.all(rt.frequency == fh.frequency[::4].reshape(2, 1, 1))
assert rt.sideband.shape == (2, )
assert np.all(rt.sideband == fh.sideband[:2])
assert rt.polarization.shape == (2, 1)
assert np.all(rt.polarization == fh.polarization[:4:2].reshape(2, 1))
def get_reshape_and_transpose(fh, sample_shape=(4, 2), sample_axes=(2, 1)):
rt = Reshape(fh, sample_shape=sample_shape)
return Transpose(rt, sample_axes=sample_axes)
def test_wrong_shape(self):
with pytest.raises(ValueError):
Reshape(self.fh, (4, 4))
with pytest.raises(ValueError):
Reshape(self.fh, ())
fname = sys.argv[1]
output_name = '/mnt/scratch-lustre/fsyed/B1133+16/Analysis2020/gk049e/numpy_arrays/yy/' + fname[:-5]
print("Output File Name: {}".format(output_name))
# Load Data
frequency = np.array([[316.00], [332.00]]) * u.MHz
sideband = np.array([[1, 1], [1, 1]])
polarization = ['R', 'L'] # Right circular polarization & left circular polarization
dispersion_measure = 4.84066 * u.pc / u.cm**3
polyco_file = '/mnt/scratch-lustre/fsyed/B1133+16/Analysis2020/gk049e/polycos/yy/polyco_new.dat'
fullpol = False
print("Parameters set")
# Creating stream reader. For other formats, such as vdif, use vdif.open(...)
fh = vdif.open(fdir + fname, 'rs')
rh = Reshape(fh, (2, 2))
dt = TimeDelta(10, format='sec')
# Rounding Time
start = Time(rh.time)
start_time_str = start.iso.__str__()
new_time = Time(start_time_str, precision = -1)
new_time_str = new_time.iso.__str__()
start_time = Time(new_time_str) + dt
print("Opened stream reader with sample shape:", rh.sample_shape)
print("Starting at time:", start_time)
# Initial waterfall interpretor
WF = sr.Fold(rh, dispersion_measure, frequency, sideband, polyco_file, polarization, fullpol,start=start_time, nthreads=1)
print("Initialized waterfall interpretor with shape:", WF.integrator.shape)