Repository: RadioAstronomySoftwareGroup/pyuvdata
Pull request #805 (2020-04-15 19:52:32)

added fast H5PY slice indexing for UVH5.read_uvh5()
<!--- Provide a general summary of your changes in the Title above -->
Added smarter h5py dataset slicing upon UVH5 read call.

## Description
<!--- Describe your changes in detail -->
 In general, if one is loading a subset of the blts axis it defaults to a fancy index, which is very slow for HDF5 (see [this stack overflow post](https://stackoverflow.com/questions/21766145/h5py-correct-way-to-slice-array-datasets)). This can be sped up if we are able to slice the datasets when possible. 

Furthermore, default behavior was to slice only one data axis at a time. This PR lets us slice as many axes simultaneously as possible, so we don't need to load in more data than what is requested.

This also exposes the `chunks` kwarg to `UVH5.write_uvh5` and `UVData.write_uvh5` for writing HDF5 files with custom chunking, which can significantly accelerate or degrade I/O performance depending on how the data are typically being read.

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
Very slow read times when working with Nblts > 20k elements files but trying to select only a subset of the antennas.

<!--- If it fixes an open issue, please link to the issue here. If this PR closes an issue, put the word 'closes' before the issue link to auto-close the issue when the PR is merged. -->

## Types of changes
<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->
- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation change (documentation changes only)
- [ ] Version change
- [ ] Build or continuous integration change


## Checklist:
<!--- You may remove the checklists that don't apply to your change type(s) or just leave them empty -->
<!--- Go over all the following points, and replace the space with an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->
- [x] I have read the [contribution guide](https://github.com/RadioAstronomySoftwareGroup/pyuvdata/blob/master/.github/CONTRIBUTING.md).
- [ ] My code follows the code style of this project.

Bug fix checklist:
- [ ] My fix includes a new test that breaks as a result of the bug (if possible).
- [ ] All new and existing tests pass.
- [ ] I have updated the [CHANGELOG](https://github.com/RadioAstronomySoftwareGroup/pyuvdata/blob/master/CHANGELOG.md).

New feature checklist:
- [ ] I have added or updated the docstrings associated with my feature using the [numpy docstring format](https://numpydoc.readthedocs.io/en/latest/format.html).
- [ ] I have updated the tutorial to highlight my new feature (if appropriate).
- [x] I have added tests to cover my new feature.
- [ ] All new and existing tests pass.
- [x] I have updated the [CHANGELOG](https://github.com/RadioAstronomySoftwareGroup/pyuvdata/blob/master/CHANGELOG.md).

Breaking change checklist:
- [ ] I have updated the docstrings associated with my change using the [numpy docstring format](https://numpydoc.readthedocs.io/en/latest/format.html).
- [ ] I have updated the tutorial to reflect my changes (if appropriate).
- [ ] My change includes backwards compatibility and deprecation warnings (if possible).
- [ ] I have added tests to cover my changes.
- [ ] All new and existing tests pass.
- [ ] I have updated the [CHANGELOG](https://github.com/RadioAstronomySoftwareGroup/pyuvdata/blob/master/CHANGELOG.md).

Documentation change checklist:
- [ ] Any updated docstrings use the [numpy docstring format](https://numpydoc.readthedocs.io/en/latest/format.html).
- [ ] If this is a significant change to the readme or other docs, I have checked that they are rendered properly on ReadTheDocs. (you may need help to get this branch to build on RTD, just ask!)

Version change checklist:
- [ ] I have updated the [CHANGELOG](https://github.com/RadioAstronomySoftwareGroup/pyuvdata/blob/master/CHANGELOG.md) to put all the unreleased changes under the new version (leaving the unreleased section empty).
- [ ] I have noted any dependency changes since the last version and will update the conda package build accordingly.

Build or continuous integration change checklist:
- [ ] If required or optional dependencies have changed (including version numbers), I have updated the readme to reflect this.
- [ ] If this is a new CI setup, I have added the associated badge to the readme and to references/make_index.py (if appropriate).


nkern: <img width="825" alt="Screen Shot 2020-04-15 at 6 17 25 AM" src="https://user-images.githubusercontent.com/4307358/79382241-1535a000-7f18-11ea-8183-97460c580f53.png">

Example of a UVH5 read on 3 HERA IDR2.2 files calling 20 baselines a 1 pol with master (top) and with this branch (bottom). Roughly 1 OOM speed up
nkern: This turns out to be a more subtle issue than I originally thought. I/O benchmarking is a little tricky because of disk cache, which it seems like can only be completely reset by re-booting the computer. (on a mac, `sudo purge` comes close but doesn't seem to completely clear my cache).

Anyways, here is some benchmarking trying to account for this
<img width="750" alt="Screen Shot 2020-04-15 at 2 58 20 PM" src="https://user-images.githubusercontent.com/4307358/79420218-54daa700-7f6d-11ea-9b50-d165784f8a81.png">
What you are seeing is the new HDF5 slicing yielding a ~25% increase on a cold read (no cache). Interestingly, the new slicing is considerably faster upon re-read, which to me signals that it is using the disk cache more effectively than the master branch currently is.

Anyways, there are even more subtleties: the simultaneously slicing seems to yield speed ups only if the slices are well-matched to the dataset chunking shape (if it is chunked at all). If it is mismatched then it **actually slows it down considerably**. Given this, I think it makes sense to keep the simultaneous slicing off by default, but allow an experienced user to turn it on given an understanding of the data's chunking shape.

@plaplant let me know if that makes sense.
nkern: 
<img width="649" alt="Screen Shot 2020-04-15 at 11 20 40 PM" src="https://user-images.githubusercontent.com/4307358/79421707-3924d000-7f70-11ea-8e17-3ca22c7b92d8.png">

Here is an example on H1C IDR2 data how mismatched simultaneous slicing can slow the read in considerably (over 1 OOM!). In this final case, even though we are leading in only 1 pol vs. 4 pol, its near 30 times slower. In this case its better the load everything in first and then down select, which is currently the default.
codecov[bot]: # [Codecov](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/805?src=pr&el=h1) Report
> Merging [#805](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/805?src=pr&el=desc) into [master](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/commit/070b74372fec7ce5bdd21411039f7eb9e18fa1b6&el=desc) will **decrease** coverage by `0.19%`.
> The diff coverage is `86.66%`.

[![Impacted file tree graph](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/805/graphs/tree.svg?width=650&height=150&src=pr&token=4QykcCeB1H)](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/805?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #805      +/-   ##
==========================================
- Coverage   98.95%   98.76%   -0.20%     
==========================================
  Files          27       27              
  Lines        9566     9687     +121     
==========================================
+ Hits         9466     9567     +101     
- Misses        100      120      +20     
```


| [Impacted Files](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/805?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pyuvdata/uvdata/uvdata.py](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/805/diff?src=pr&el=tree#diff-cHl1dmRhdGEvdXZkYXRhL3V2ZGF0YS5weQ==) | `100.00% <ø> (ø)` | |
| [pyuvdata/uvdata/uvh5.py](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/805/diff?src=pr&el=tree#diff-cHl1dmRhdGEvdXZkYXRhL3V2aDUucHk=) | `96.67% <86.66%> (-3.33%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/805?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/805?src=pr&el=footer). Last update [070b743...c88df2c](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/805?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

plaplant: @nkern thanks very much for the detailed profiling! I think this type of study is very valuable, and certainly something we'd like to do more of in the future for specific access patterns for analysis. It also makes sense that caching help you a lot, both at the OS and HDF5 levels. HDF5 in particular has some tricks for caching chunked datasets that help speed it up.

I think it makes sense for the profiling to be quite sensitive to the chunk size. Part of the tradeoff of chunking is that when you need to read data from a part of a chunk, HDF5 must read _all_ of the chunk. If you only need a relatively small fraction of that chunk, you'll waste a lot of I/O. On the other hand, we need to have chunking enabled for the compression filters that we use. So it makes sense that there isn't a one-size-fits-all chunk size, and potentially let the user specify what is best for what they are trying to do.

I'm happy to expose the chunk size to the user as an option when writing datasets. I had initially used the automatic values for simplicity's sake, but it seems like for our use cases HDF5 is choosing sub-optimal defaults. I also think that doing multi-dimensional slicing into datasets is interesting, and something that's worth including as an option. I'll do a more detailed review in the next couple of days, but I think the high-level ideas in this PR are excellent ones. Thanks!
nkern: @plaplant great sounds good. Another thing I'm thinking about is that every time we re-write these files in our analysis pipe (say delay filtered or lst binned or xtalk subtracted) if we don't expose the chunking kwarg then we are writing the data w/ HDF5's sub-optimal auto-chunking, even if the raw data from the correlator are chunked well!

For posterity and others tuning in, here is the HDF5 group support page on chunking.
https://support.hdfgroup.org/HDF5/doc/Advanced/Chunking/
See "Pitfalls" for some ways in which chunking can *drastically* reduce performance (i.e. chunks that are too small or too large).
plaplant: @nkern I forgot to mention in my review, but we should also update the changelog. We're adding the `chunks` keyword to `UVData.write_uvh5` and `multidim_index` to `UVData.read`. We also probably want to mention that reading uvh5 files is optimized for multidimensional slicing (provided it's well-matched to the chunk size).
plaplant: Thanks very much @nkern! Looks like we've lost some test coverage, so I can add a few tests to get back to full coverage. I think that's the last remaining issue, and then we'll be good to go.
bhazelton: I'm happy to merge. Thank you both!

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index be4cca65..5d829901 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,8 @@ All notable changes to this project will be documented in this file.
 - Handling of strings in UVFlag files has been made more widely compatible.
 
 ### Added
+- `chunks` to `UVH5.write_uvh5` and `UVData.write_uvh5` for HDF5 dataset chunking.
+- `multidim_index` to `UVH5.read_uvh5` and `UVData.read` for multidimensional slicing into HDF5 datasets
 - Option to provide parameters for RectBivariatespline through interp
 - `weights_square_array` (optional) parameter on UVFlag - stores sum of squares of weights when converting to waterfall
 - `frequency_average` method on UVData to average data along the frequency axis.
diff --git a/pyuvdata/uvdata/uvdata.py b/pyuvdata/uvdata/uvdata.py
index 4f39444c..f974b117 100644
--- a/pyuvdata/uvdata/uvdata.py
+++ b/pyuvdata/uvdata/uvdata.py
@@ -6071,6 +6071,7 @@ class UVData(UVBase):
         run_check=True,
         check_extra=True,
         run_check_acceptability=True,
+        multidim_index=False,
     ):
         """
         Read a UVH5 file.
@@ -6163,6 +6164,11 @@ class UVData(UVBase):
             Option to check acceptable range of the values of parameters after
             reading in the file (the default is True, meaning the acceptable
             range check will be done). Ignored if read_data is False.
+        multidim_index : bool
+            [Only for HDF5] If True, attempt to index the HDF5 dataset
+            simultaneously along all data axes. Otherwise index one axis at-a-time.
+            This only works if data selection is sliceable along all but one axis.
+            If indices are not well-matched to data chunks, this can be slow.
 
         Raises
         ------
@@ -6203,6 +6209,7 @@ class UVData(UVBase):
             run_check_acceptability=run_check_acceptability,
             data_array_dtype=data_array_dtype,
             keep_all_metadata=keep_all_metadata,
+            multidim_index=multidim_index,
         )
         self._convert_from_filetype(uvh5_obj)
         del uvh5_obj
@@ -6248,6 +6255,7 @@ class UVData(UVBase):
         check_extra=True,
         run_check_acceptability=True,
         skip_bad_files=False,
+        multidim_index=False,
     ):
         """
         Read a generic file into a UVData object.
@@ -6427,6 +6435,11 @@ class UVData(UVBase):
             the read continues even if one or more files are corrupted. Files
             that produce errors will be printed. Default is False (files will
             not be skipped).
+        multidim_index : bool
+            [Only for HDF5] If True, attempt to index the HDF5 dataset
+            simultaneously along all data axes. Otherwise index one axis at-a-time.
+            This only works if data selection is sliceable along all but one axis.
+            If indices are not well-matched to data chunks, this can be slow.
 
         Raises
         ------
@@ -6800,6 +6813,7 @@ class UVData(UVBase):
                     run_check_acceptability=run_check_acceptability,
                     data_array_dtype=data_array_dtype,
                     keep_all_metadata=keep_all_metadata,
+                    multidim_index=multidim_index,
                 )
                 select = False
 
@@ -6985,6 +6999,7 @@ class UVData(UVBase):
         check_extra=True,
         run_check_acceptability=True,
         clobber=False,
+        chunks=True,
         data_compression=None,
         flags_compression="lzf",
         nsample_compression="lzf",
@@ -6999,15 +7014,18 @@ class UVData(UVBase):
              The UVH5 file to write to.
         clobber : bool
             Option to overwrite the file if it already exists.
+        chunks : tuple or bool
+            h5py.create_dataset chunks keyword. Tuple for chunk shape,
+            True for auto-chunking, None for no chunking. Default is True.
         data_compression : str
             HDF5 filter to apply when writing the data_array. Default is
-            None meaning no filter or compression.
+            None meaning no filter or compression. Dataset must be chunked.
         flags_compression : str
             HDF5 filter to apply when writing the flags_array. Default is "lzf"
-            for the LZF filter.
+            for the LZF filter. Dataset must be chunked.
         nsample_compression : str
             HDF5 filter to apply when writing the nsample_array. Default is "lzf"
-            for the LZF filter.
+            for the LZF filter. Dataset must be chunked.
         data_write_dtype : numpy dtype
             datatype of output visibility data. If 'None', then the same datatype
             as data_array will be used. Otherwise, a numpy dtype object must be
@@ -7032,6 +7050,7 @@ class UVData(UVBase):
             check_extra=check_extra,
             run_check_acceptability=run_check_acceptability,
             clobber=clobber,
+            chunks=chunks,
             data_compression=data_compression,
             flags_compression=flags_compression,
             nsample_compression=nsample_compression,
@@ -7043,6 +7062,7 @@ class UVData(UVBase):
         self,
         filename,
         clobber=False,
+        chunks=True,
         data_compression=None,
         flags_compression="lzf",
         nsample_compression="lzf",
@@ -7057,15 +7077,18 @@ class UVData(UVBase):
              The UVH5 file to write to.
         clobber : bool
             Option to overwrite the file if it already exists.
+        chunks : tuple or bool
+            h5py.create_dataset chunks keyword. Tuple for chunk shape,
+            True for auto-chunking, None for no chunking. Default is True.
         data_compression : str
             HDF5 filter to apply when writing the data_array. Default is
-            None meaning no filter or compression.
+            None meaning no filter or compression. Dataset must be chunked.
         flags_compression : str
             HDF5 filter to apply when writing the flags_array. Default is "lzf"
-            for the LZF filter.
+            for the LZF filter. Dataset must be chunked.
         nsample_compression : str
             HDF5 filter to apply when writing the nsample_array. Default is "lzf"
-            for the LZF filter.
+            for the LZF filter. Dataset must be chunked.
         data_write_dtype : numpy dtype
             datatype of output visibility data. If 'None', then the same datatype
             as data_array will be used. Otherwise, a numpy dtype object must be
@@ -7084,6 +7107,7 @@ class UVData(UVBase):
         uvh5_obj.initialize_uvh5_file(
             filename,
             clobber=clobber,
+            chunks=chunks,
             data_compression=data_compression,
             flags_compression=flags_compression,
             nsample_compression=nsample_compression,
diff --git a/pyuvdata/uvdata/uvh5.py b/pyuvdata/uvdata/uvh5.py
index 3f1033c9..5ccc18f2 100644
--- a/pyuvdata/uvdata/uvh5.py
+++ b/pyuvdata/uvdata/uvh5.py
@@ -91,17 +91,13 @@ def _read_complex_astype(dset, indices, dtype_out=np.complex64):
         raise ValueError(
             "output datatype must be one of (complex, np.complex64, np.complex128)"
         )
-    dset_shape = [0, 0, 0, 0]
-    for i in range(len(dset_shape)):
-        if indices[i] == np.s_[:]:
-            dset_shape[i] = dset.shape[i]
-        else:
-            dset_shape[i] = len(indices[i])
+    dset_shape = _get_dset_shape(dset, indices)
     output_array = np.empty(dset_shape, dtype=dtype_out)
     dtype_in = dset.dtype
     with dset.astype(dtype_in):
-        output_array.real = dset["r"][indices]
-        output_array.imag = dset["i"][indices]
+        # dset is indexed in native dtype, but is upcast upon assignment
+        output_array.real = _index_dset(dset["r"], indices)
+        output_array.imag = _index_dset(dset["i"], indices)
 
     return output_array
 
@@ -136,6 +132,235 @@ def _write_complex_astype(data, dset, indices):
     return
 
 
+def _convert_to_slices(indices, max_nslice_frac=0.1):
+    """
+    Convert list of indices to a list of slices.
+
+    Parameters
+    ----------
+    indices : list
+        A 1D list of integers for array indexing.
+    max_nslice_frac : float
+        A float from 0 -- 1. If the number of slices
+        needed to represent input 'indices' divided by len(indices)
+        exceeds this fraction, then we determine that we cannot
+        easily represent 'indices' with a list of slices.
+
+    Returns
+    -------
+    list
+        list of slice objects used to represent indices
+    bool
+        If True, indices is easily represented by slices
+        (max_nslice_frac condition met), otherwise False
+
+    Notes
+    -----
+    Example:
+        if: indices = [1, 2, 3, 4, 10, 11, 12, 13, 14]
+        then: slices = [slice(1, 5, 1), slice(11, 15, 1)]
+    """
+    # check for integer index
+    if isinstance(indices, (int, np.integer)):
+        indices = [indices]
+    # check for already a slice
+    if isinstance(indices, slice):
+        return [indices], True
+    # assert indices is longer than 2, or return trivial solutions
+    if len(indices) == 0:
+        return [slice(0, 0, 0)], False
+    elif len(indices) == 1:
+        return [slice(indices[0], indices[0] + 1, 1)], True
+    elif len(indices) == 2:
+        return [slice(indices[0], indices[1] + 1, indices[1] - indices[0])], True
+
+    # setup empty slices list
+    Ninds = len(indices)
+    slices = []
+
+    # iterate over indices
+    for i, ind in enumerate(indices):
+        if i == 0:
+            # start the first slice object
+            start = ind
+            last_step = indices[i + 1] - ind
+            continue
+
+        # calculate step from previous index
+        step = ind - indices[i - 1]
+
+        # if step != last_step, this ends the slice
+        if step != last_step:
+            # append to list
+            slices.append(slice(start, indices[i - 1] + 1, last_step))
+
+            # check if this is the last element
+            if i == Ninds - 1:
+                # append last element
+                slices.append(slice(ind, ind + 1, 1))
+                continue
+
+            # setup next step
+            start = ind
+            last_step = indices[i + 1] - ind
+
+        # check if this is the last element
+        elif i == Ninds - 1:
+            # end slice and append
+            slices.append(slice(start, ind + 1, step))
+
+    # determine whether slices are a reasonable representation
+    Nslices = len(slices)
+    passed = (float(Nslices) / len(indices)) < max_nslice_frac
+
+    return slices, passed
+
+
+def _get_slice_len(s, axlen):
+    """
+    Get length of a slice s into array of len axlen.
+
+    Parameters
+    ----------
+    s : slice object
+        Slice object to index with
+    axlen : int
+        Length of axis s slices into
+
+    Returns
+    -------
+    int
+        Length of slice object
+    """
+    if s.start is None:
+        start = 0
+    else:
+        start = s.start
+    if s.stop is None:
+        stop = axlen
+    else:
+        stop = np.min([s.stop, axlen])
+    if s.step is None:
+        step = 1
+    else:
+        step = s.step
+
+    return ((stop - 1 - start) // step) + 1
+
+
+def _get_dset_shape(dset, indices):
+    """
+    Given a 4-tuple of indices, determine the indexed array shape.
+
+    Parameters
+    ----------
+    dset : h5py dataset
+        A reference to an HDF5 dataset on disk.
+    indices : tuple
+        A 4-tuple with the indices to extract along each dimension of dset.
+        Each element should contain a list of indices, a slice element,
+        or a list of slice elements that will be concatenated after slicing.
+
+    Returns
+    -------
+    tuple
+        4-tuple with the shape of the indexed array
+    """
+    dset_shape = list(dset.shape)
+    for i, inds in enumerate(indices):
+        # check for integer
+        if isinstance(inds, (int, np.integer)):
+            dset_shape[i] = 1
+        # check for slice object
+        if isinstance(inds, slice):
+            dset_shape[i] = _get_slice_len(inds, dset_shape[i])
+        # check for list
+        if isinstance(inds, list):
+            # check for list of integers
+            if isinstance(inds[0], (int, np.integer)):
+                dset_shape[i] = len(inds)
+            elif isinstance(inds[0], slice):
+                dset_shape[i] = sum((_get_slice_len(s, dset_shape[i]) for s in inds))
+
+    return dset_shape
+
+
+def _index_dset(dset, indices):
+    """
+    Index a UVH5 data, flags or nsamples h5py dataset.
+
+    Parameters
+    ----------
+    dset : h5py dataset
+        A reference to an HDF5 dataset on disk.
+    indices : tuple
+        A 4-tuple with the indices to extract along each dimension of dset.
+        Each element should contain a list of indices, a slice element,
+        or a list of slice elements that will be concatenated after slicing.
+        Indices must be provided such that all dimensions can be indexed
+        simultaneously.
+
+    Returns
+    -------
+    ndarray
+        The indexed dset
+
+    Notes
+    -----
+    This makes and fills an empty array with dset indices.
+    For trivial indexing, (e.g. a trivial slice), constructing
+    a new array and filling it is suboptimal over direct
+    indexing, e.g. dset[indices].
+    This function specializes in repeated slices over the same axis,
+    e.g. if indices is [[slice(0, 5), slice(10, 15), ...], ..., ]
+    """
+    # get dset and arr shape
+    dset_shape = dset.shape
+    arr_shape = _get_dset_shape(dset, indices)
+
+    # create empty array of dset dtype
+    arr = np.empty(arr_shape, dtype=dset.dtype)
+
+    # get arr and dset indices for each dimension in indices
+    dset_indices = []
+    arr_indices = []
+    for i, dset_inds in enumerate(indices):
+        if isinstance(dset_inds, (int, np.integer)):
+            # this dimension is len 1, so slice is fine
+            arr_indices.append([slice(None)])
+            dset_indices.append([[dset_inds]])
+
+        elif isinstance(dset_inds, slice):
+            # this dimension is just a slice, so slice is fine
+            arr_indices.append([slice(None)])
+            dset_indices.append([dset_inds])
+
+        elif isinstance(dset_inds, list):
+            if isinstance(dset_inds[0], (int, np.integer)):
+                # this is a list of integers, append slice
+                arr_indices.append([slice(None)])
+                dset_indices.append([dset_inds])
+            elif isinstance(dset_inds[0], slice):
+                # this is a list of slices, need list of slice lens
+                slens = [_get_slice_len(s, dset_shape[i]) for s in dset_inds]
+                ssums = [sum(slens[:j]) for j in range(len(slens))]
+                arr_inds = [slice(s, s + l) for s, l in zip(ssums, slens)]
+                arr_indices.append(arr_inds)
+                dset_indices.append(dset_inds)
+
+    # iterate over each of the 4 axes and fill the array
+    for blt_arr, blt_dset in zip(arr_indices[0], dset_indices[0]):
+        for spw_arr, spw_dset in zip(arr_indices[1], dset_indices[1]):
+            for freq_arr, freq_dset in zip(arr_indices[2], dset_indices[2]):
+                for pol_arr, pol_dset in zip(arr_indices[3], dset_indices[3]):
+                    # index dset and assign to arr
+                    arr[blt_arr, spw_arr, freq_arr, pol_arr] = dset[
+                        blt_dset, spw_dset, freq_dset, pol_dset
+                    ]
+
+    return arr
+
+
 class UVH5(UVData):
     """
     A class for UVH5 file objects.
@@ -343,6 +568,7 @@ class UVH5(UVData):
         run_check_acceptability,
         data_array_dtype,
         keep_all_metadata,
+        multidim_index,
     ):
         """
         Read the data-size arrays (data, flags, nsamples) from a file.
@@ -384,6 +610,7 @@ class UVH5(UVData):
             blt_inds,
         )
 
+        # figure out which axis is the most selective
         if blt_inds is not None:
             blt_frac = len(blt_inds) / float(self.Nblts)
         else:
@@ -416,15 +643,15 @@ class UVH5(UVData):
 
         if min_frac == 1:
             # no select, read in all the data
+            inds = (np.s_[:], np.s_[:], np.s_[:], np.s_[:])
             if custom_dtype:
-                inds = (np.s_[:], np.s_[:], np.s_[:], np.s_[:])
                 self.data_array = _read_complex_astype(
                     dgrp["visdata"], inds, data_array_dtype
                 )
             else:
-                self.data_array = dgrp["visdata"][:, :, :, :]
-            self.flag_array = dgrp["flags"][:, :, :, :]
-            self.nsample_array = dgrp["nsamples"][:, :, :, :]
+                self.data_array = _index_dset(dgrp["visdata"], inds)
+            self.flag_array = _index_dset(dgrp["flags"], inds)
+            self.nsample_array = _index_dset(dgrp["nsamples"], inds)
         else:
             # do select operations on everything except data_array, flag_array
             # and nsample_array
@@ -432,62 +659,154 @@ class UVH5(UVData):
                 blt_inds, freq_inds, pol_inds, history_update_string, keep_all_metadata
             )
 
+            # determine which axes can be sliced, rather than fancy indexed
+            # max_nslice_frac of 0.1 yields slice speedup over fancy index for HERA data
+            # See pyuvdata PR #805
+            if blt_inds is not None:
+                blt_slices, blt_sliceable = _convert_to_slices(
+                    blt_inds, max_nslice_frac=0.1
+                )
+            else:
+                blt_inds, blt_slices = np.s_[:], np.s_[:]
+                blt_sliceable = True
+
+            if freq_inds is not None:
+                freq_slices, freq_sliceable = _convert_to_slices(
+                    freq_inds, max_nslice_frac=0.1
+                )
+            else:
+                freq_inds, freq_slices = np.s_[:], np.s_[:]
+                freq_sliceable = True
+
+            if pol_inds is not None:
+                pol_slices, pol_sliceable = _convert_to_slices(
+                    pol_inds, max_nslice_frac=0.5
+                )
+            else:
+                pol_inds, pol_slices = np.s_[:], np.s_[:]
+                pol_sliceable = True
+
             # open references to datasets
             visdata_dset = dgrp["visdata"]
             flags_dset = dgrp["flags"]
             nsamples_dset = dgrp["nsamples"]
 
+            # check that multidim_index is appropriate
+            if multidim_index:
+                # if more than one dim is not sliceable, then not appropriate
+                if sum([blt_sliceable, freq_sliceable, pol_sliceable]) < 2:
+                    multidim_index = False
+
             # just read in the right portions of the data and flag arrays
             if blt_frac == min_frac:
+                # construct inds list given simultaneous sliceability
+                inds = [blt_inds, np.s_[:], np.s_[:], np.s_[:]]
+                if blt_sliceable:
+                    inds[0] = blt_slices
+                if multidim_index:
+                    if freq_sliceable:
+                        inds[2] = freq_slices
+                    else:
+                        inds[2] = freq_inds
+                if multidim_index:
+                    if pol_sliceable:
+                        inds[3] = pol_slices
+                    else:
+                        inds[3] = pol_inds
+
+                inds = tuple(inds)
+
+                # index datasets
                 if custom_dtype:
-                    inds = (blt_inds, np.s_[:], np.s_[:], np.s_[:])
                     visdata = _read_complex_astype(visdata_dset, inds, data_array_dtype)
                 else:
-                    visdata = visdata_dset[blt_inds, :, :, :]
-                flags = flags_dset[blt_inds, :, :, :]
-                nsamples = nsamples_dset[blt_inds, :, :, :]
+                    visdata = _index_dset(visdata_dset, inds)
+                flags = _index_dset(flags_dset, inds)
+                nsamples = _index_dset(nsamples_dset, inds)
 
                 assert self.Nspws == visdata.shape[1]
 
-                if freq_frac < 1:
+                # down select on other dimensions if necessary
+                # use indices not slices here: generally not the bottleneck
+                if not multidim_index and freq_frac < 1:
                     visdata = visdata[:, :, freq_inds, :]
                     flags = flags[:, :, freq_inds, :]
                     nsamples = nsamples[:, :, freq_inds, :]
-                if pol_frac < 1:
+                if not multidim_index and pol_frac < 1:
                     visdata = visdata[:, :, :, pol_inds]
                     flags = flags[:, :, :, pol_inds]
                     nsamples = nsamples[:, :, :, pol_inds]
+
             elif freq_frac == min_frac:
+                # construct inds list given simultaneous sliceability
+                inds = [np.s_[:], np.s_[:], freq_inds, np.s_[:]]
+                if freq_sliceable:
+                    inds[2] = freq_slices
+                if multidim_index:
+                    if blt_sliceable:
+                        inds[0] = blt_slices
+                    else:
+                        inds[0] = blt_inds
+                if multidim_index:
+                    if pol_sliceable:
+                        inds[3] = pol_slices
+                    else:
+                        inds[3] = pol_inds
+
+                inds = tuple(inds)
+
+                # index datasets
                 if custom_dtype:
-                    inds = (np.s_[:], np.s_[:], freq_inds, np.s_[:])
                     visdata = _read_complex_astype(visdata_dset, inds, data_array_dtype)
                 else:
-                    visdata = visdata_dset[:, :, freq_inds, :]
-                flags = flags_dset[:, :, freq_inds, :]
-                nsamples = nsamples_dset[:, :, freq_inds, :]
+                    visdata = _index_dset(visdata_dset, inds)
+                flags = _index_dset(flags_dset, inds)
+                nsamples = _index_dset(nsamples_dset, inds)
 
-                if blt_frac < 1:
+                # down select on other dimensions if necessary
+                # use indices not slices here: generally not the bottleneck
+                if not multidim_index and blt_frac < 1:
                     visdata = visdata[blt_inds, :, :, :]
                     flags = flags[blt_inds, :, :, :]
                     nsamples = nsamples[blt_inds, :, :, :]
-                if pol_frac < 1:
+                if not multidim_index and pol_frac < 1:
                     visdata = visdata[:, :, :, pol_inds]
                     flags = flags[:, :, :, pol_inds]
                     nsamples = nsamples[:, :, :, pol_inds]
+
             else:
+                # construct inds list given simultaneous sliceability
+                inds = [np.s_[:], np.s_[:], np.s_[:], pol_inds]
+                if pol_sliceable:
+                    inds[3] = pol_slices
+                if multidim_index:
+                    if blt_sliceable:
+                        inds[0] = blt_slices
+                    else:
+                        inds[0] = blt_inds
+                if multidim_index:
+                    if freq_sliceable:
+                        inds[2] = freq_slices
+                    else:
+                        inds[2] = freq_inds
+
+                inds = tuple(inds)
+
+                # index datasets
                 if custom_dtype:
-                    inds = (np.s_[:], np.s_[:], np.s_[:], pol_inds)
                     visdata = _read_complex_astype(visdata_dset, inds, data_array_dtype)
                 else:
-                    visdata = visdata_dset[:, :, :, pol_inds]
-                flags = flags_dset[:, :, :, pol_inds]
-                nsamples = nsamples_dset[:, :, :, pol_inds]
+                    visdata = _index_dset(visdata_dset, inds)
+                flags = _index_dset(flags_dset, inds)
+                nsamples = _index_dset(nsamples_dset, inds)
 
-                if blt_frac < 1:
+                # down select on other dimensions if necessary
+                # use indices not slices here: generally not the bottleneck
+                if not multidim_index and blt_frac < 1:
                     visdata = visdata[blt_inds, :, :, :]
                     flags = flags[blt_inds, :, :, :]
                     nsamples = nsamples[blt_inds, :, :, :]
-                if freq_frac < 1:
+                if not multidim_index and freq_frac < 1:
                     visdata = visdata[:, :, freq_inds, :]
                     flags = flags[:, :, freq_inds, :]
                     nsamples = nsamples[:, :, freq_inds, :]
@@ -524,6 +843,7 @@ class UVH5(UVData):
         run_check_acceptability=True,
         data_array_dtype=np.complex128,
         keep_all_metadata=True,
+        multidim_index=False,
     ):
         """
         Read in data from a UVH5 file.
@@ -608,6 +928,11 @@ class UVH5(UVData):
             Option to check acceptable range of the values of parameters after
             reading in the file (the default is True, meaning the acceptable
             range check will be done). Ignored if read_data is False.
+        multidim_index : bool
+            If True, attempt to index the HDF5 dataset
+            simultaneously along all data axes. Otherwise index one axis at-a-time.
+            This only works if data selection is sliceable along all but one axis.
+            If indices are not well-matched to data chunks, this can be slow.
 
         Returns
         -------
@@ -659,6 +984,7 @@ class UVH5(UVData):
                 run_check_acceptability,
                 data_array_dtype,
                 keep_all_metadata,
+                multidim_index,
             )
 
         return
@@ -767,6 +1093,7 @@ class UVH5(UVData):
         check_extra=True,
         run_check_acceptability=True,
         clobber=False,
+        chunks=True,
         data_compression=None,
         flags_compression="lzf",
         nsample_compression="lzf",
@@ -790,15 +1117,18 @@ class UVH5(UVData):
             writing the file.
         clobber : bool
             Option to overwrite the file if it already exists.
+        chunks : tuple or bool
+            h5py.create_dataset chunks keyword. Tuple for chunk shape,
+            True for auto-chunking, None for no chunking. Default is True.
         data_compression : str
             HDF5 filter to apply when writing the data_array. Default is None
-            (no filter/compression).
+            (no filter/compression). Dataset must be chunked.
         flags_compression : str
             HDF5 filter to apply when writing the flags_array. Default is the
-            LZF filter.
+            LZF filter. Dataset must be chunked.
         nsample_compression : str
             HDF5 filter to apply when writing the nsample_array. Default is the
-            LZF filter.
+            LZF filter. Dataset must be chunked.
         data_write_dtype : numpy dtype
             The datatype of output visibility data. If 'None', then the same
             datatype as data_array will be used. The user may specify 'c8' for
@@ -862,7 +1192,7 @@ class UVH5(UVData):
                 visdata = dgrp.create_dataset(
                     "visdata",
                     self.data_array.shape,
-                    chunks=True,
+                    chunks=chunks,
                     compression=data_compression,
                     dtype=data_write_dtype,
                 )
@@ -871,20 +1201,20 @@ class UVH5(UVData):
             else:
                 visdata = dgrp.create_dataset(
                     "visdata",
-                    chunks=True,
+                    chunks=chunks,
                     data=self.data_array,
                     compression=data_compression,
                     dtype=data_write_dtype,
                 )
             dgrp.create_dataset(
                 "flags",
-                chunks=True,
+                chunks=chunks,
                 data=self.flag_array,
                 compression=flags_compression,
             )
             dgrp.create_dataset(
                 "nsamples",
-                chunks=True,
+                chunks=chunks,
                 data=self.nsample_array.astype(np.float32),
                 compression=nsample_compression,
             )
@@ -895,6 +1225,7 @@ class UVH5(UVData):
         self,
         filename,
         clobber=False,
+        chunks=True,
         data_compression=None,
         flags_compression="lzf",
         nsample_compression="lzf",
@@ -909,15 +1240,18 @@ class UVH5(UVData):
             The UVH5 file to write to.
         clobber : bool
             Option to overwrite the file if it already exists.
+        chunks : tuple or bool
+            h5py.create_dataset chunks keyword. Tuple for chunk shape,
+            True for auto-chunking, None for no chunking. Default is True.
         data_compression : str
             HDF5 filter to apply when writing the data_array. Default is None
-            (no filter/compression).
+            (no filter/compression). Dataset must be chunked.
         flags_compression : str
             HDF5 filter to apply when writing the flags_array. Default is the
-            LZF filter.
+            LZF filter. Dataset must be chunked.
         nsample_compression : str
             HDF5 filter to apply when writing the nsample_array. Default is the
-            LZF filter.
+            LZF filter. Dataset must be chunked.
         data_write_dtype : str or numpy dtype
             The datatype of output visibility data. If 'None', then double-
             precision floats will be used. The user may specify 'c8' for
@@ -981,21 +1315,21 @@ class UVH5(UVData):
             dgrp.create_dataset(
                 "visdata",
                 data_size,
-                chunks=True,
+                chunks=chunks,
                 dtype=data_write_dtype,
                 compression=data_compression,
             )
             dgrp.create_dataset(
                 "flags",
                 data_size,
-                chunks=True,
+                chunks=chunks,
                 dtype="b1",
                 compression=flags_compression,
             )
             dgrp.create_dataset(
                 "nsamples",
                 data_size,
-                chunks=True,
+                chunks=chunks,
                 dtype="f4",
                 compression=nsample_compression,
             )


Test patch:
diff --git a/pyuvdata/uvdata/tests/test_uvh5.py b/pyuvdata/uvdata/tests/test_uvh5.py
index effa01ab..73745579 100644
--- a/pyuvdata/uvdata/tests/test_uvh5.py
+++ b/pyuvdata/uvdata/tests/test_uvh5.py
@@ -267,6 +267,8 @@ def test_uvh5_compression_options(uv_uvfits):
     """
     Test writing data with compression filters.
     """
+    import h5py
+
     uv_in = uv_uvfits
     uv_out = UVData()
     testfile = os.path.join(DATA_PATH, "test", "outtest_uvfits_compression.uvh5")
@@ -275,6 +277,20 @@ def test_uvh5_compression_options(uv_uvfits):
     uv_in.write_uvh5(
         testfile,
         clobber=True,
+        chunks=True,
+        data_compression="lzf",
+        flags_compression=None,
+        nsample_compression=None,
+    )
+    uv_out.read(testfile)
+    assert uv_in == uv_out
+
+    # test with non-auto chunking
+    chunks = (680, 1, 16, 1)
+    uv_in.write_uvh5(
+        testfile,
+        clobber=True,
+        chunks=chunks,
         data_compression="lzf",
         flags_compression=None,
         nsample_compression=None,
@@ -282,6 +298,10 @@ def test_uvh5_compression_options(uv_uvfits):
     uv_out.read(testfile)
     assert uv_in == uv_out
 
+    # check that chunks match
+    with h5py.File(testfile, "r") as f:
+        assert f["Data"]["visdata"].chunks == chunks
+
     # clean up
     os.remove(testfile)
 
@@ -467,6 +487,14 @@ def test_uvh5_partial_read_pols(uv_uvfits):
     uvh5_uv2.select(polarizations=pols_to_keep)
     assert uvh5_uv == uvh5_uv2
 
+    # select on pols (non contiguous in file)
+    # and check consistent results with and without multidim_index
+    pols_to_keep = [-1, -2, -4]
+    uvh5_uv.read(testfile, polarizations=pols_to_keep, multidim_index=True)
+    uvh5_uv2.read(testfile, multidim_index=False)
+    uvh5_uv2.select(polarizations=pols_to_keep)
+    assert uvh5_uv == uvh5_uv2
+
     return
 
 
@@ -526,6 +554,30 @@ def test_uvh5_partial_read_multi1(uv_uvfits):
     )
     assert uvh5_uv == uvh5_uv2
 
+    # now check with multidim_index
+    # use different versions of blt_inds and pol_inds
+    # ones that are and are not sliceable by UVH5 standards
+    chans_to_keep = np.arange(20, 35)
+    uvh5_uv3 = UVData()
+    uvh5_uv4 = UVData()
+    random_blts = np.random.choice(np.arange(uv_in.Nblts), size=1000, replace=False)
+    for blts_to_keep in [random_blts, np.arange(1000)]:
+        for pols_to_keep in [[-1, -2], [-1, -2, -4]]:
+            uvh5_uv3.read(
+                testfile,
+                blt_inds=blts_to_keep,
+                polarizations=pols_to_keep,
+                freq_chans=chans_to_keep,
+                multidim_index=True,
+            )
+            uvh5_uv4.read(testfile)
+            uvh5_uv4.select(
+                blt_inds=blts_to_keep,
+                polarizations=pols_to_keep,
+                freq_chans=chans_to_keep,
+            )
+            assert uvh5_uv3 == uvh5_uv4
+
     # clean up
     os.remove(testfile)
 
@@ -562,6 +614,30 @@ def test_uvh5_partial_read_multi2(uv_uvfits):
     )
     assert uvh5_uv == uvh5_uv2
 
+    # now check with multidim_index
+    # use different versions of freq_inds and pol_inds
+    # ones that are and are not sliceable by UVH5 standards
+    blts_to_keep = np.arange(100)
+    uvh5_uv3 = UVData()
+    uvh5_uv4 = UVData()
+    random_freqs = np.random.choice(np.arange(uv_in.Nfreqs), size=50, replace=False)
+    for chans_to_keep in [random_freqs, np.arange(50)]:
+        for pols_to_keep in [[-1, -2], [-1, -2, -4]]:
+            uvh5_uv3.read(
+                testfile,
+                blt_inds=blts_to_keep,
+                freq_chans=chans_to_keep,
+                polarizations=pols_to_keep,
+                multidim_index=True,
+            )
+            uvh5_uv4.read(testfile)
+            uvh5_uv4.select(
+                blt_inds=blts_to_keep,
+                freq_chans=chans_to_keep,
+                polarizations=pols_to_keep,
+            )
+            assert uvh5_uv3 == uvh5_uv4
+
     # clean up
     os.remove(testfile)
 
@@ -598,6 +674,58 @@ def test_uvh5_partial_read_multi3(uv_uvfits):
     )
     assert uvh5_uv == uvh5_uv2
 
+    # now check with multidim_index
+    # use different versions of blt_inds and freq_inds
+    # ones that are and are not sliceable by UVH5 standards
+    pols_to_keep = [-1, -2]
+    uvh5_uv3 = UVData()
+    uvh5_uv4 = UVData()
+    random_blts = np.random.choice(np.arange(uv_in.Nblts), size=1000, replace=False)
+    random_freqs = np.random.choice(np.arange(uv_in.Nfreqs), size=50, replace=False)
+    for blts_to_keep in [random_blts, np.arange(1000)]:
+        for chans_to_keep in [random_freqs, np.arange(50)]:
+            uvh5_uv3.read(
+                testfile,
+                blt_inds=blts_to_keep,
+                freq_chans=chans_to_keep,
+                polarizations=pols_to_keep,
+                multidim_index=True,
+            )
+            uvh5_uv4.read(testfile)
+            uvh5_uv4.select(
+                blt_inds=blts_to_keep,
+                freq_chans=chans_to_keep,
+                polarizations=pols_to_keep,
+            )
+            assert uvh5_uv3 == uvh5_uv4
+
+    # clean up
+    os.remove(testfile)
+
+    return
+
+
+def test_uvh5_read_multdim_index(uv_uvfits):
+    """
+    Test some odd cases for UVH5 multdim indexing
+    """
+    uv_in = uv_uvfits
+    testfile = os.path.join(DATA_PATH, "test", "outtest.uvh5")
+    # change telescope name to avoid errors
+    uv_in.telescope_name = "PAPER"
+    uv_in.write_uvh5(testfile, clobber=True)
+    uvh5_uv = UVData()
+    uvh5_uv.read(testfile)
+
+    # check that non sliceable multidim index is caught
+    # and does not fail
+    ants_to_keep = np.array([0, 19, 11, 24, 3, 23, 1, 20, 21])
+    chans_to_keep = [15, 17, 20]
+    uvh5_uv = UVData()
+    uvh5_uv.read(
+        testfile, antenna_nums=ants_to_keep, freq_chans=chans_to_keep,
+    )
+
     # clean up
     os.remove(testfile)
 
@@ -2413,3 +2541,26 @@ def test_eq_coeffs_roundtrip(uv_uvfits):
     os.remove(testfile)
 
     return
+
+
+def test_read_slicing(uv_uvfits):
+    """Test HDF5 slicing helper functions"""
+    # check trivial slice representations
+    slices, _ = uvh5._convert_to_slices([])
+    assert slices == [slice(0, 0, 0)]
+    slices, _ = uvh5._convert_to_slices(10)
+    assert slices == [slice(10, 11, 1)]
+
+    # dataset shape checking
+    # check various kinds of indexing give the right answer
+    indices = [slice(0, 10), 0, [0, 1, 2], [0]]
+    dset = np.empty((100, 1, 1024, 2), dtype=np.float)
+    shape = uvh5._get_dset_shape(dset, indices)
+    assert tuple(shape) == (10, 1, 3, 1)
+
+    # dataset indexing
+    # check various kinds of indexing give the right answer
+    slices = [uvh5._convert_to_slices(ind)[0] for ind in indices]
+    slices[1] = 0
+    data = uvh5._index_dset(dset, slices)
+    assert data.shape == tuple(shape)