Merge pull request #218 from HERA-Team/partial_load_from_container

Partial load from container.get_pspec

hera_pspec/__init__.py

@@ -1,13 +1,13 @@
 """
 __init__.py file for hera_pspec
 """
-from hera_pspec import version, conversions, grouping, pspecbeam, plot, pstokes, testing
-from hera_pspec import uvpspec_utils as uvputils
+from . import version, conversions, grouping, pspecbeam, plot, pstokes, testing
+from . import uvpspec_utils as uvputils
 
-from hera_pspec.uvpspec import UVPSpec
-from hera_pspec.pspecdata import PSpecData
-from hera_pspec.container import PSpecContainer
-from hera_pspec.parameter import PSpecParam
-from hera_pspec.pspecbeam import PSpecBeamUV, PSpecBeamGauss, PSpecBeamFromArray
+from .uvpspec import UVPSpec
+from .pspecdata import PSpecData
+from .container import PSpecContainer
+from .parameter import PSpecParam
+from .pspecbeam import PSpecBeamUV, PSpecBeamGauss, PSpecBeamFromArray
 
 __version__ = version.version

hera_pspec/container.py

@@ -1,10 +1,11 @@
 import numpy as np
 import h5py
-from hera_pspec import uvpspec, version, utils
 import argparse
 import time
 from functools import wraps
 
+from . import uvpspec, version, utils
+
 
 def transactional(fn):
     """
@@ -53,7 +54,7 @@ class PSpecContainer(object):
     Container class for managing multiple UVPSpec objects.
     """
     def __init__(self, filename, mode='r', keep_open=True, swmr=False,
-                 tsleep=0.5, maxiter=2):
+                 tsleep=0.1, maxiter=1):
         """
         Manage a collection of UVPSpec objects that are stored in a structured
         HDF5 file.
@@ -208,7 +209,7 @@ def _store_pspec(self, pspec_group, uvp):
         # Write UVPSpec to group
         uvp.write_to_group(pspec_group, run_check=True)
 
-    def _load_pspec(self, pspec_group):
+    def _load_pspec(self, pspec_group, **kwargs):
         """
         Load a new UVPSpec object from a HDF5 group.
 
@@ -218,6 +219,9 @@ def _load_pspec(self, pspec_group):
             Group containing datasets that contain power spectrum and
             supporting information, in a standard format expected by UVPSpec.
         
+        kwargs : dict, optional
+            UVPSpec.read_from_group partial IO keyword arguments
+
         Returns
         -------
         uvp : UVPSpec
@@ -236,7 +240,7 @@ def _load_pspec(self, pspec_group):
 
         # Create new UVPSpec object and fill with data from this group
         uvp = uvpspec.UVPSpec()
-        uvp.read_from_group(pspec_group)
+        uvp.read_from_group(pspec_group, **kwargs)
         return uvp
 
     def _update_header(self):
@@ -345,7 +349,7 @@ def set_pspec(self, group, psname, pspec, overwrite=False):
         psgrp.attrs['pspec_type'] = pspec.__class__.__name__
 
     @transactional
-    def get_pspec(self, group, psname=None):
+    def get_pspec(self, group, psname=None, **kwargs):
         """
         Get a UVPSpec power spectrum object from a given group.
 
@@ -357,6 +361,9 @@ def get_pspec(self, group, psname=None):
         psname : str, optional
             The name of the power spectrum to return.
         
+        kwargs : dict
+            UVPSpec.read_from_group partial IO keyword arguments
+
         Returns
         -------
         uvp : UVPSpec or list of UVPSpec
@@ -376,15 +383,15 @@ def get_pspec(self, group, psname=None):
 
             # Load power spectrum if it exists
             if key2 in list(grp.keys()):
-                return self._load_pspec(grp[key2])
+                return self._load_pspec(grp[key2], **kwargs)
             else:
                 raise KeyError("No pspec named '%s' in group '%s'" % (key2, key1))
 
         # Otherwise, extract all available power spectra
         uvp = []
         def pspec_filter(n, obj):
             if u'pspec_type' in list(obj.attrs.keys()):
-                uvp.append(self._load_pspec(obj))
+                uvp.append(self._load_pspec(obj, **kwargs))
 
         # Traverse the entire set of groups/datasets looking for pspecs
         grp.visititems(pspec_filter) # This adds power spectra to the uvp list

hera_pspec/grouping.py

@@ -1,13 +1,12 @@
 import numpy as np
 from collections import OrderedDict as odict
-from hera_pspec import uvpspec_utils as uvputils
-from hera_pspec import utils, version
 import random
 import copy
 import argparse
 from astropy import stats as astats
 import os
 
+from . import utils, version, uvpspec_utils as uvputils
 
 def group_baselines(bls, Ngroups, keep_remainder=False, randomize=False,
                     seed=None):

hera_pspec/noise.py

@@ -1,10 +1,11 @@
 import numpy as np
 import os
-from hera_pspec import conversions, pspecbeam
 import copy
 import ast
 from collections import OrderedDict as odict
 
+from . import conversions, pspecbeam
+
 
 def calc_P_N(scalar, Tsys, t_int, Ncoherent=1, Nincoherent=None, form='Pk', k=None, component='real'):
     """

hera_pspec/plot.py

@@ -1,15 +1,14 @@
 import numpy as np
 import pyuvdata
-from hera_pspec import conversions, uvpspec, utils
-import matplotlib
-import matplotlib.pyplot as plt
 import copy
 from collections import OrderedDict as odict
 import astropy.units as u
 import astropy.constants as c
 from pyuvdata import UVData
 import uvtools
 
+from . import conversions, uvpspec, utils
+
 
 def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False, 
                    average_times=False, fold=False, plot_noise=False, 
@@ -107,6 +106,9 @@ def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False,
     fig : matplotlib.pyplot.Figure
         Matplotlib Figure instance.
     """
+    import matplotlib
+    import matplotlib.pyplot as plt
+
     # Create new Axes if none specified
     new_plot = False
     if ax is None:
@@ -373,6 +375,9 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='real',
     fig : matplotlib.pyplot.Figure
         Matplotlib Figure instance if input ax is None.
     """
+    import matplotlib
+    import matplotlib.pyplot as plt
+
     # assert component
     assert component in ['real', 'abs', 'imag'], "Can't parse specified component {}".format(component)
 
@@ -715,6 +720,9 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True,
     fig : matplotlib.pyplot.Figure
         Matplotlib Figure instance if ax is None.
     """
+    import matplotlib
+    import matplotlib.pyplot as plt
+
     # type checking
     uvp = copy.deepcopy(uvp)
     assert isinstance(uvp, uvpspec.UVPSpec), "input uvp must be a UVPSpec object"
@@ -964,6 +972,9 @@ def plot_uvdata_waterfalls(uvd, basename, data='data', plot_mode='log',
         Keyword arguments passed to uvtools.plot.waterfall, which passes them 
         on to matplotlib.imshow.
     """
+    import matplotlib
+    import matplotlib.pyplot as plt
+
     assert isinstance(uvd, UVData), "'uvd' must be a UVData object."
     assert data in ['data', 'flags', 'nsamples'], \
             "'%s' not a valid data array; use 'data', 'flags', or 'nsamples'" \

hera_pspec/pspecbeam.py

@@ -1,13 +1,13 @@
 import numpy as np
 import os
-import hera_pspec.conversions as conversions
-import hera_pspec.uvpspec_utils as uvputils
 import scipy.integrate as integrate
 from scipy.interpolate import interp1d
 from pyuvdata import UVBeam, utils as uvutils
 import aipy
 from collections import OrderedDict as odict
 
+from . import conversions as conversions, uvpspec_utils as uvputils
+
 
 def _compute_pspec_scalar(cosmo, beam_freqs, omega_ratio, pspec_freqs, 
                           num_steps=5000, taper='none', little_h=True, 

hera_pspec/pspecdata.py

@@ -4,8 +4,6 @@
 import copy, operator, itertools, sys
 from collections import OrderedDict as odict
 import hera_cal as hc
-from hera_pspec import uvpspec, utils, version, pspecbeam, container
-from hera_pspec import uvpspec_utils as uvputils
 from pyuvdata import utils as uvutils
 import datetime
 import time
@@ -16,6 +14,9 @@
 import json
 import uvtools.dspec as dspec
 
+from . import uvpspec, utils, version, pspecbeam, container, uvpspec_utils as uvputils
+
+
 class PSpecData(object):
 
     def __init__(self, dsets=[], wgts=None, dsets_std=None, labels=None,
@@ -2892,7 +2893,7 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True,
               groupname=None, dset_labels=None, dset_pairs=None, psname_ext=None,
               spw_ranges=None, n_dlys=None, pol_pairs=None, blpairs=None,
               input_data_weight='identity', norm='I', taper='none',
-              exclude_auto_bls=False, exclude_permutations=True,
+              exclude_auto_bls=False, exclude_cross_bls=False, exclude_permutations=True,
               Nblps_per_group=None, bl_len_range=(0, 1e10),
               bl_deg_range=(0, 180), bl_error_tol=1.0,
               beam=None, cosmo=None, interleave_times=False, rephase_to_dset=None,
@@ -2985,6 +2986,11 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True,
         exclude_auto_bls is True, eliminate all instances of a bl crossed
         with itself. Default: False
 
+    exclude_cross_bls : boolean
+        If True and if blpairs is None, exclude all bls crossed with a
+        different baseline. Note if this and exclude_auto_bls are True
+        then no blpairs will exist. Default: False
+
     exclude_permutations : boolean
         If blpairs is None, redundant baseline groups will be formed and
         all cross-multiplies will be constructed. In doing so, if
@@ -3276,6 +3282,7 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True,
                                       dsets[dsetp[0]], dsets[dsetp[1]],
                                       filter_blpairs=True,
                                       exclude_auto_bls=exclude_auto_bls,
+                                      exclude_cross_bls=exclude_cross_bls,
                                       exclude_permutations=exclude_permutations,
                                       Nblps_per_group=Nblps_per_group,
                                       bl_len_range=bl_len_range,
@@ -3367,6 +3374,7 @@ def list_of_tuple_tuples(v):
     a.add_argument("--time_thresh", default=0.2, type=float, help="Fractional flagging threshold across time to trigger flag broadcast if broadcast_dset_flags is True")
     a.add_argument("--Jy2mK", default=False, action='store_true', help="Convert datasets from Jy to mK if a beam model is provided.")
     a.add_argument("--exclude_auto_bls", default=False, action='store_true', help='If blpairs is not provided, exclude all baselines paired with itself.')
+    a.add_argument("--exclude_cross_bls", default=False, action='store_true', help='If blpairs is not provided, exclude all baselines paired with a different baseline.')
     a.add_argument("--exclude_permutations", default=False, action='store_true', help='If blpairs is not provided, exclude a basline-pair permutations. Ex: if (A, B) exists, exclude (B, A).')
     a.add_argument("--Nblps_per_group", default=None, type=int, help="If blpairs is not provided and group == True, set the number of blpairs in each group.")
     a.add_argument("--bl_len_range", default=(0, 1e10), nargs='+', type=float, help="If blpairs is not provided, limit the baselines used based on their minimum and maximum length in meters.")

hera_pspec/pstokes.py

@@ -4,9 +4,9 @@
 import numpy as np, os
 import pyuvdata
 import copy
-from hera_pspec import version
 from collections import OrderedDict as odict
 
+from . import version
 
 # Weights used in forming Stokes visibilities.
 # See pyuvdata.utils.polstr2num for conversion between polarization string

hera_pspec/testing.py

@@ -2,11 +2,11 @@
 import numpy as np
 import copy, operator, itertools
 from collections import OrderedDict as odict
-from hera_pspec import uvpspec, pspecdata, conversions, pspecbeam, utils
 from pyuvdata import UVData
 from hera_cal.utils import JD2LST
 from scipy import stats
-import hera_pspec.uvpspec_utils as uvputils
+
+from . import uvpspec, pspecdata, conversions, pspecbeam, utils, uvpspec_utils as uvputils
 
 
 def build_vanilla_uvpspec(beam=None):

hera_pspec/tests/test_container.py

@@ -118,6 +118,14 @@ def test_PSpecContainer(self, keep_open=True, swmr=False):
                 ps = ps_store.get_pspec(g, psname=psname)
                 assert(isinstance(ps, UVPSpec))
 
+        # check partial IO in get_pspec
+        ps = ps_store.get_pspec(group_names[0], pspec_names[0], just_meta=True)
+        assert not hasattr(ps, 'data_array')
+        assert hasattr(ps, 'time_avg_array')
+        ps = ps_store.get_pspec(group_names[0], pspec_names[0], blpairs=[((1, 2), (1, 2))])
+        assert hasattr(ps, 'data_array')
+        assert np.all(np.isclose(ps.blpair_array, 101102101102))
+
         # Check that invalid list arguments raise errors in set_pspec()
         assert_raises(ValueError, ps_store.set_pspec, group=group_names[:2], 
                       psname=pspec_names[0], pspec=self.uvp, overwrite=True)

hera_pspec/tests/test_utils.py

@@ -147,7 +147,6 @@ def test_spw_range_from_redshifts(self):
         nt.ok_( isinstance(spw5, tuple) )
         nt.ok_( spw5[0] is not None )
 
-
     def test_calc_blpair_reds(self):
         fname = os.path.join(DATA_PATH, 'zen.all.xx.LST.1.06964.uvA')
         uvd = UVData()
@@ -194,10 +193,17 @@ def test_calc_blpair_reds(self):
                                    bl_len_range=(10.0, 20.0))
         nt.assert_equal(blps, [((24, 25), (37, 38))])
 
+        # test exclude_cross_bls
+        (bls1, bls2, blps, xants1,
+         xants2) = utils.calc_blpair_reds(uvd, uvd, filter_blpairs=True, exclude_cross_bls=True)
+        for bl1, bl2 in blps:
+            assert bl1 == bl2
+
         # test exceptions
         uvd2 = copy.deepcopy(uvd)
         uvd2.antenna_positions[0] += 2
         nt.assert_raises(AssertionError, utils.calc_blpair_reds, uvd, uvd2)
+        nt.assert_raises(AssertionError, utils.calc_blpair_reds, uvd, uvd, exclude_auto_bls=True, exclude_cross_bls=True)
 
     def test_get_delays(self):
         utils.get_delays(np.linspace(100., 200., 50)*1e6)

hera_pspec/tests/test_uvpspec.py

@@ -80,6 +80,9 @@ def test_get_funcs(self):
         # test get_blpairs
         blps = self.uvp.get_blpairs()
         nt.assert_equal(blps, [((1, 2), (1, 2)), ((1, 3), (1, 3)), ((2, 3), (2, 3))])
+        # test get_polpairs
+        polpairs = self.uvp.get_polpairs()
+        nt.assert_equal(polpairs, [('xx', 'xx')])
         # test get all keys
         keys = self.uvp.get_all_keys()
         nt.assert_equal(keys, [(0, ((1, 2), (1, 2)), ('xx','xx')),