Merge 1a33ec5 into a111e56

hera_pspec/__init__.py

@@ -2,7 +2,7 @@
 __init__.py file for hera_pspec
 """
 
-from hera_pspec import version, conversions, grouping, pspecbeam, plot
+from hera_pspec import version, conversions, grouping, pspecbeam, plot, testing
 from hera_pspec import uvpspec_utils as uvputils
 
 from hera_pspec.uvpspec import UVPSpec

hera_pspec/pspecdata.py

@@ -7,6 +7,7 @@
 from hera_pspec import uvpspec, utils, version
 from hera_pspec import uvpspec_utils as uvputils
 from pyuvdata import utils as uvutils
+import datetime
 
 
 class PSpecData(object):
@@ -164,8 +165,11 @@ def validate_datasets(self, verbose=True):
 
         # Check if dsets are all the same shape along freq axis
         Nfreqs = [d.Nfreqs for d in self.dsets]
+        channel_widths = [d.channel_width for d in self.dsets]
         if np.unique(Nfreqs).size > 1:
             raise ValueError("all dsets must have the same Nfreqs")
+        if np.unique(channel_widths).size > 1:
+            raise ValueError("all dsets must have the same channel_widths")
 
         # Check shape along time axis
         Ntimes = [d.Ntimes for d in self.dsets]
@@ -1165,7 +1169,7 @@ def pspec(self, bls1, bls2, dsets, input_data_weight='identity', norm='I',
         pol_arr = map(lambda p: pyuvdata.utils.polnum2str(p), dset1.polarization_array)
 
         # assert form of bls1 and bls2
-        assert len(bls1) == len(bls2), "length of bls1 must equal length of bls2"
+        assert len(bls1) == len(bls2) and len(bls1) > 0, "length of bls1 must equal length of bls2 and be > 0"
         for i in range(len(bls1)):
             if isinstance(bls1[i], tuple):
                 assert isinstance(bls2[i], tuple), "bls1[{}] type must match bls2[{}] type".format(i, i)
@@ -1276,7 +1280,6 @@ def pspec(self, bls1, bls2, dsets, input_data_weight='identity', norm='I',
                         pass
                     else:
                         if verbose: print("  Building G...")
-                        print key1, "---", key2
                         Gv = self.get_G(key1, key2)
                         Hv = self.get_H(key1, key2, taper=taper)
                         built_GH = True
@@ -1379,14 +1382,28 @@ def pspec(self, bls1, bls2, dsets, input_data_weight='identity', norm='I',
         uvp.pol_array = np.array(map(lambda p: uvutils.polstr2num(p), pol_arr))
         uvp.Npols = len(pol_arr)
         uvp.scalar_array = np.array(sclr_arr)
-        uvp.channel_width = dset1.channel_width
+        uvp.channel_width = dset1.channel_width  # all dsets are validated to agree
         uvp.weighting = input_data_weight
         uvp.vis_units, uvp.norm_units = self.units(little_h=little_h)
         uvp.telescope_location = dset1.telescope_location
-        uvp.history = dset1.history + dset2.history + history
+        filename1 = getattr(dset1.extra_keywords, 'filename', None)
+        filename2 = getattr(dset2.extra_keywords, 'filename', None)
+        label1 = self.labels[self.dset_idx(dsets[0])]
+        label2 = self.labels[self.dset_idx(dsets[1])]
+        uvp.labels = sorted(set([label1, label2]))
+        uvp.label_1_array = np.ones((uvp.Nspws, uvp.Nblpairts, uvp.Npols), np.int) \
+                            * uvp.labels.index(label1)
+        uvp.label_2_array = np.ones((uvp.Nspws, uvp.Nblpairts, uvp.Npols), np.int) \
+                            * uvp.labels.index(label2)
+        uvp.labels = np.array(uvp.labels, np.str)
+        uvp.history = "UVPSpec written on {} with hera_pspec git hash {}\n{}\n" \
+                      "dataset1: filename: {}, label: {}, history:\n{}\n{}\n" \
+                      "dataset2: filename: {}, label: {}, history:\n{}\n{}\n" \
+                      "".format(datetime.datetime.now(), version.git_hash, '-'*20, 
+                                filename1, label1, dset1.history, '-'*20, 
+                                filename2, label2, dset2.history, '-'*20)
         uvp.taper = taper
         uvp.norm = norm
-        uvp.git_hash = version.git_hash
 
         if self.primary_beam is not None:
             # attach cosmology
@@ -1396,14 +1413,6 @@ def pspec(self, bls1, bls2, dsets, input_data_weight='identity', norm='I',
             uvp.OmegaP, uvp.OmegaPP = self.primary_beam.get_Omegas(uvp.pol_array)
             if hasattr(self.primary_beam, 'filename'):
                 uvp.beamfile = self.primary_beam.filename
-        if hasattr(dset1.extra_keywords, 'filename'):
-            uvp.filename1 = dset1.extra_keywords['filename']
-        if hasattr(dset2.extra_keywords, 'filename'):
-            uvp.filename2 = dset2.extra_keywords['filename']
-        lbl1 = self.labels[self.dset_idx(dsets[0])]
-        lbl2 = self.labels[self.dset_idx(dsets[1])]
-        if lbl1 is not None: uvp.label1 = lbl1
-        if lbl2 is not None: uvp.label2 = lbl2
 
         # fill data arrays
         uvp.data_array = data_array

hera_pspec/testing.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python2
+import numpy as np
+import copy, operator, itertools
+from collections import OrderedDict as odict
+from hera_pspec import uvpspec, pspecdata, conversions, pspecbeam
+from pyuvdata import UVData
+
+
+def build_vanilla_uvpspec(beam=None):
+    """
+    Build an example vanilla UVPSpec object from scratch, with all necessary metadata.
+
+    Parameters
+    ----------
+    beam : PSpecBeamUV object
+
+    Returns
+    -------
+    uvp : UVPSpec object
+    """
+    uvp = uvpspec.UVPSpec()
+
+    Ntimes = 10
+    Nfreqs = 50
+    Ndlys = Nfreqs
+    Nspws = 1
+    Nspwdlys = Nspws * Nfreqs
+
+    # [((1, 2), (1, 2)), ((2, 3), (2, 3)), ((1, 3), (1, 3))]
+    blpairs = [1002001002, 2003002003, 1003001003]
+    bls = [1002, 2003, 1003]
+    Nbls = len(bls)
+    Nblpairs = len(blpairs)
+    Nblpairts = Nblpairs * Ntimes
+
+    blpair_array = np.tile(blpairs, Ntimes)
+    bl_array = np.array(bls)
+    bl_vecs = np.array([[  5.33391548e+00,  -1.35907816e+01,  -7.91624188e-09],
+                        [ -8.67982998e+00,   4.43554478e+00,  -1.08695203e+01],
+                        [ -3.34591450e+00,  -9.15523687e+00,  -1.08695203e+01]])
+    time_array = np.repeat(np.linspace(2458042.1, 2458042.2, Ntimes), Nblpairs)
+    time_1_array = time_array
+    time_2_array = time_array
+    lst_array = np.repeat(np.ones(Ntimes, dtype=np.float), Nblpairs)
+    lst_1_array = lst_array
+    lst_2_array = lst_array
+    time_avg_array = time_array
+    lst_avg_array = lst_array
+    spws = np.arange(Nspws)
+    spw_array = np.tile(spws, Ndlys)
+    freq_array = np.repeat(np.linspace(100e6, 105e6, Nfreqs, endpoint=False), Nspws)
+    dly_array = np.fft.fftshift(np.repeat(np.fft.fftfreq(Nfreqs, np.median(np.diff(freq_array))), Nspws))
+    pol_array = np.array([-5])
+    Npols = len(pol_array)
+    vis_units = 'unknown'
+    norm_units = 'Hz str'
+    weighting = 'identity'
+    channel_width = np.median(np.diff(freq_array))
+    history = 'example'
+    taper = "none"
+    norm = "I"
+    git_hash = "random"
+    scalar_array = np.ones((Nspws, Npols), np.float)
+    label1 = 'red'
+    label2 = 'blue'
+    labels = np.array([label1, label2])
+    label_1_array = np.ones((Nspws, Nblpairts, Npols), np.int) * 0
+    label_2_array = np.ones((Nspws, Nblpairts, Npols), np.int) * 1
+    if beam is not None:
+        OmegaP, OmegaPP = beam.get_Omegas(beam.primary_beam.polarization_array[0])
+        beam_freqs = beam.beam_freqs
+
+    telescope_location = np.array([5109325.85521063, 
+                                   2005235.09142983, 
+                                  -3239928.42475397])
+
+    cosmo = conversions.Cosmo_Conversions()
+
+    data_array, wgt_array, integration_array, nsample_array = {}, {}, {}, {}
+    for s in spws:
+        data_array[s] = np.ones((Nblpairts, Ndlys, Npols), dtype=np.complex) \
+                      * blpair_array[:, None, None] / 1e9
+        wgt_array[s] = np.ones((Nblpairts, Ndlys, 2, Npols), dtype=np.float)
+        integration_array[s] = np.ones((Nblpairts, Npols), dtype=np.float)
+        nsample_array[s] = np.ones((Nblpairts, Npols), dtype=np.float)
+
+    params = ['Ntimes', 'Nfreqs', 'Nspws', 'Nspwdlys', 'Nblpairs', 'Nblpairts', 
+              'Npols', 'Ndlys', 'Nbls', 'blpair_array', 'time_1_array', 
+              'time_2_array', 'lst_1_array', 'lst_2_array', 'spw_array', 
+              'dly_array', 'freq_array', 'pol_array', 'data_array', 'wgt_array',
+              'integration_array', 'bl_array', 'bl_vecs', 'telescope_location', 
+              'vis_units', 'channel_width', 'weighting', 'history', 'taper', 'norm', 
+              'git_hash', 'nsample_array', 'time_avg_array', 'lst_avg_array', 
+              'cosmo', 'scalar_array', 'labels', 'norm_units', 'labels', 'label_1_array',
+              'label_2_array']
+
+    if beam is not None:
+        params += ['OmegaP', 'OmegaPP', 'beam_freqs']
+
+    # Set all parameters
+    for p in params:
+        setattr(uvp, p, locals()[p])
+
+    uvp.check()
+
+    return uvp, cosmo
+
+def build_uvpspec_from_data(data, bls, spw_ranges=None, beam=None, taper='none', cosmo=None, verbose=False):
+    """
+    Build an example UVPSpec object from a visibility file and PSpecData.
+
+    Parameters
+    ----------
+    data : UVData object or path to miriad file
+
+    bls : list of at least 2 baseline tuples
+        Ex: [(24, 25), (37, 38), ...]
+
+    spw_ranges : list of spw tuples
+        See PSpecData.pspec docstring
+
+    beam : PSpecBeamUV object or path to beamfits healpix beam
+
+    taper : string
+        Optional tapering applied to the data before OQE.
+
+    cosmo : Cosmo_Conversions object
+
+    verbose : bool
+        if True, report feedback to standard output
+
+    Returns
+    -------
+    uvp : UVPSpec object
+    """
+    # load data
+    if isinstance(data, str):
+        uvd = UVData()
+        uvd.read_miriad(dfile)
+    elif isinstance(data, UVData):
+        uvd = data
+
+    # load beam
+    if isinstance(beam, str):
+        beam = pspecbeam.PSpecBeamUV(beam, cosmo=cosmo)
+    if beam is not None and cosmo is not None:
+        beam.cosmo = cosmo
+
+    # instantiate pspecdata
+    ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], labels=['d1', 'd2'], beam=beam)
+
+    # get red bls
+    bls1, bls2, _ = pspecdata.construct_blpairs(bls, exclude_auto_bls=True)
+
+    # run pspec
+    uvp = ds.pspec(bls1, bls2, (0, 1), input_data_weight='identity', spw_ranges=spw_ranges, 
+                   taper=taper, verbose=verbose)
+
+    return uvp
+

hera_pspec/tests/test_container.py

@@ -3,14 +3,14 @@
 import numpy as np
 import os, sys, copy
 from hera_pspec.data import DATA_PATH
-from hera_pspec import PSpecContainer, UVPSpec
-from test_uvpspec import build_example_uvpspec
+from hera_pspec import PSpecContainer, UVPSpec, testing
+
 
 class Test_PSpecContainer(unittest.TestCase):
 
     def setUp(self):
         self.fname = os.path.join(DATA_PATH, '_test_container.hdf5')
-        self.uvp, self.cosmo = build_example_uvpspec()
+        self.uvp, self.cosmo = testing.build_vanilla_uvpspec()
 
     def tearDown(self):
         # Remove HDF5 file

hera_pspec/tests/test_grouping.py

@@ -3,8 +3,7 @@
 import numpy as np
 import os
 from hera_pspec.data import DATA_PATH
-from test_uvpspec import build_example_uvpspec
-from hera_pspec import uvpspec, conversions, parameter, pspecbeam, pspecdata
+from hera_pspec import uvpspec, conversions, parameter, pspecbeam, pspecdata, testing
 from hera_pspec import uvpspec_utils as uvputils
 from hera_pspec import grouping
 
@@ -13,7 +12,7 @@ class Test_grouping(unittest.TestCase):
     def setUp(self):
         beamfile = os.path.join(DATA_PATH, 'NF_HERA_Beams.beamfits')
         self.beam = pspecbeam.PSpecBeamUV(beamfile)
-        uvp, cosmo = build_example_uvpspec(beam=self.beam)
+        uvp, cosmo = testing.build_vanilla_uvpspec(beam=self.beam)
         uvp.check()
         self.uvp = uvp
 

hera_pspec/tests/test_utils.py

@@ -3,10 +3,10 @@
 import numpy as np
 import os, sys, copy
 from hera_pspec.data import DATA_PATH
-from hera_pspec import utils
+from hera_pspec import utils, testing
 from collections import OrderedDict as odict
 from pyuvdata import UVData
-from test_uvpspec import build_example_uvpspec
+
 
 def test_cov():
     # load another data file
@@ -60,7 +60,7 @@ def setUp(self):
                                           "zen.2458042.17772.xx.HH.uvXA"))
 
         # Create UVPSpec object
-        self.uvp, cosmo = build_example_uvpspec()
+        self.uvp, cosmo = testing.build_vanilla_uvpspec()
 
     def tearDown(self):
         pass