Skip to content

Commit

Permalink
Add test of partial reading
Browse files Browse the repository at this point in the history
  • Loading branch information
@plaplant
plaplant committed Jun 30, 2018
1 parent a77c6b6 commit 1ea69ad
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 5 deletions.
69 changes: 69 additions & 0 deletions pyuvdata/tests/test_uvh5.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,3 +137,72 @@ def test_UVH5CompressionOptions():
os.remove(testfile)

return


def test_UVH5PartialRead():
"""
Test reading in only part of a dataset from disk
"""
uvh5_uv = UVData()
uvh5_uv2 = UVData()
uvfits_file = os.path.join(DATA_PATH, 'day2_TDEM0003_10s_norx_1src_1spw.uvfits')
uvtest.checkWarnings(uvh5_uv.read_uvfits, [uvfits_file], message='Telescope EVLA is not')
testfile = os.path.join(DATA_PATH, 'test', 'outtest.h5')
uvh5_uv.write_uvh5(testfile, clobber=True)

# select on antennas
ants_to_keep = np.array([0, 19, 11, 24, 3, 23, 1, 20, 21])
uvh5_uv.read_uvh5(testfile, antenna_nums=ants_to_keep)
uvh5_uv2.read_uvh5(testfile)
uvh5_uv2.select(antenna_nums=ants_to_keep)
nt.assert_equal(uvh5_uv, uvh5_uv2)

# select on frequency channels
chans_to_keep = np.arange(12, 22)
uvh5_uv.read_uvh5(testfile, freq_chans=chans_to_keep)
uvh5_uv2.read_uvh5(testfile)
uvh5_uv2.select(freq_chans=chans_to_keep)
nt.assert_equal(uvh5_uv, uvh5_uv2)

# select on pols
pols_to_keep = [-1, -2]
uvh5_uv.read_uvh5(testfile, polarizations=pols_to_keep)
uvh5_uv2.read_uvh5(testfile)
uvh5_uv2.select(polarizations=pols_to_keep)
nt.assert_equal(uvh5_uv, uvh5_uv2)

# now test selecting on multiple axes
# frequencies first
uvh5_uv.read_uvh5(testfile, antenna_nums=ants_to_keep,
freq_chans=chans_to_keep,
polarizations=pols_to_keep)
uvh5_uv2.read_uvh5(testfile)
uvh5_uv2.select(antenna_nums=ants_to_keep, freq_chans=chans_to_keep,
polarizations=pols_to_keep)
nt.assert_equal(uvh5_uv, uvh5_uv2)

# baselines first
ants_to_keep = np.array([0, 1])
uvh5_uv.read_uvh5(testfile, antenna_nums=ants_to_keep,
freq_chans=chans_to_keep,
polarizations=pols_to_keep)
uvh5_uv2.read_uvh5(testfile)
uvh5_uv2.select(antenna_nums=ants_to_keep, freq_chans=chans_to_keep,
polarizations=pols_to_keep)
nt.assert_equal(uvh5_uv, uvh5_uv2)

# polarizations first
ants_to_keep = np.array([0, 1, 2, 3, 6, 7, 8, 11, 14, 18, 19, 20, 21, 22])
chans_to_keep = np.arange(12, 64)
uvh5_uv.read_uvh5(testfile, antenna_nums=ants_to_keep,
freq_chans=chans_to_keep,
polarizations=pols_to_keep)
uvh5_uv2.read_uvh5(testfile)
uvh5_uv2.select(antenna_nums=ants_to_keep, freq_chans=chans_to_keep,
polarizations=pols_to_keep)
nt.assert_equal(uvh5_uv, uvh5_uv2)

# clean up
os.remove(testfile)

return
69 changes: 64 additions & 5 deletions pyuvdata/uvdata.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1767,13 +1767,55 @@ def write_miriad(self, filepath, run_check=True, check_extra=True,
clobber=clobber, no_antnums=no_antnums)
del(miriad_obj)

def read_uvh5(self, filename, run_check=True, check_extra=True,
run_check_acceptability=True):
def read_uvh5(self, filename, antenna_nums=None, antenna_names=None,
ant_str=None, bls=None, frequencies=None, freq_chans=None,
times=None, polarizations=None, blt_inds=None, read_data=True,
run_check=True, check_extra=True, run_check_acceptability=True):
"""
Read a UVH5 file.
Args:
filename: The UVH5 file to read.
antenna_nums: The antennas numbers to include when reading data into
the object (antenna positions and names for the excluded antennas
will be retained). This cannot be provided if antenna_names is
also provided. Ignored if read_data is False.
antenna_names: The antennas names to include when reading data into
the object (antenna positions and names for the excluded antennas
will be retained). This cannot be provided if antenna_nums is
also provided. Ignored if read_data is False.
bls: A list of antenna number tuples (e.g. [(0,1), (3,2)]) or a list of
baseline 3-tuples (e.g. [(0,1,'xx'), (2,3,'yy')]) specifying baselines
to keep in the object. For length-2 tuples, the ordering of the numbers
within the tuple does not matter. For length-3 tuples, the polarization
string is in the order of the two antennas. If length-3 tuples are provided,
the polarizations argument below must be None. Ignored if read_data is False.
ant_str: A string containing information about what antenna numbers
and polarizations to include when reading data into the object.
Can be 'auto', 'cross', 'all', or combinations of antenna numbers
and polarizations (e.g. '1', '1_2', '1x_2y').
See tutorial for more examples of valid strings and
the behavior of different forms for ant_str.
If '1x_2y,2y_3y' is passed, both polarizations 'xy' and 'yy' will
be kept for both baselines (1,2) and (2,3) to return a valid
pyuvdata object.
An ant_str cannot be passed in addition to any of the above antenna
args or the polarizations arg.
Ignored if read_data is False.
frequencies: The frequencies to include when reading data into the
object. Ignored if read_data is False.
freq_chans: The frequency channel numbers to include when reading
data into the object. Ignored if read_data is False.
times: The times to include when reading data into the object.
Ignored if read_data is False.
polarizations: The polarizations to include when reading data into
the object. Ignored if read_data is False.
blt_inds: The baseline-time indices to include when reading data into
the object. This is not commonly used. Ignored if read_data is False.
read_data: Read in the visibility and flag data. If set to false,
only the basic header info and metadata (if read_metadata is True)
will be read in. Results in an incompletely defined object
(check will not pass). Default True.
run_check: Option to check for the existence and proper shapes of
parameters after reading in the file. Default is True.
check_extra: Option to check optional parameters as well as required
Expand All @@ -1786,18 +1828,35 @@ def read_uvh5(self, filename, run_check=True, check_extra=True,
"""
from . import uvh5
if isinstance(filename, (list, tuple)):
self.read_uvh5(filename[0], run_check=run_check, check_extra=check_extra,
if not read_data:
raise ValueError('read_data cannot be False for a list of uvfits files')

self.read_uvh5(filename[0], antenna_nums=antenna_nums,
antenna_names=antenna_names, ant_str=ant_str, bls=bls,
frequencies=frequencies, freq_chans=freq_chans, times=times,
polarizations=polarizations, blt_inds=blt_inds,
read_data=read_data, run_check=run_check,
check_extra=check_extra,
run_check_acceptability=run_check_acceptability)
if len(filename) > 1:
for f in filename[1:]:
uv2 = UVData()
uv2.read_uvh5(f, run_check=run_check, check_extra=check_extra,
uv2.read_uvh5(f, antenna_nums=antenna_nums,
antenna_names=antenna_names, ant_str=ant_str, bls=bls,
frequencies=frequencies, freq_chans=freq_chans,
times=times, polarizations=polarizations,
blt_inds=blt_inds, read_data=read_data,
run_check=run_check, check_extra=check_extra,
run_check_acceptability=run_check_acceptability)
self += uv2
del(uv2)
else:
uvh5_obj = uvh5.UVH5()
uvh5_obj.read_uvh5(filename, run_check=run_check, check_extra=check_extra,
uvh5_obj.read_uvh5(filename, antenna_nums=antenna_nums,
antenna_names=antenna_names, ant_str=ant_str, bls=bls,
frequencies=frequencies, freq_chans=freq_chans, times=times,
polarizations=polarizations, blt_inds=blt_inds,
read_data=read_data, run_check=run_check, check_extra=check_extra,
run_check_acceptability=run_check_acceptability)
self._convert_from_filetype(uvh5_obj)
del(uvh5_obj)
Expand Down

0 comments on commit 1ea69ad

Please sign in to comment.