diff --git a/.travis.yml b/.travis.yml index 118fea77..08bcd625 100644 --- a/.travis.yml +++ b/.travis.yml @@ -2,17 +2,18 @@ language: python python: # We don't actually use the Travis Python, but this keeps it organized. - "2.7" + - "3.6" # Cache pip-installed dependencies cache: pip: true env: - - PYUVDATA_VERSION="@d6e1c2f80d7b5b59d7794c8f28d867a7030a7c75" - PYUVDATA_VERSION="" + #- PYUVDATA_VERSION="@d6e1c2f80d7b5b59d7794c8f28d867a7030a7c75" -allow_failures: - - env: PYUVDATA_VERSION="" +#allow_failures: +# - env: PYUVDATA_VERSION="" install: # ensure that we have the full tag information available for version.py @@ -36,7 +37,7 @@ install: # create environment and install dependencies - conda create -q -n test-environment python=$TRAVIS_PYTHON_VERSION numpy scipy nose pip matplotlib coverage - source activate test-environment - - conda install -c conda-forge healpy aipy + - conda install -c conda-forge healpy aipy scikit-learn - pip install coveralls - pip install h5py - pip install git+https://github.com/HERA-Team/pyuvdata.git$PYUVDATA_VERSION diff --git a/hera_pspec/container.py b/hera_pspec/container.py index b097828e..80d4336a 100644 --- a/hera_pspec/container.py +++ b/hera_pspec/container.py @@ -48,7 +48,7 @@ def _open(self): self._update_header() # Denote as Container - if 'pspec_type' not in self.data.attrs.keys(): + if 'pspec_type' not in list(self.data.attrs.keys()): self.data.attrs['pspec_type'] = self.__class__.__name__ def _store_pspec(self, pspec_group, uvp): @@ -88,7 +88,7 @@ def _load_pspec(self, pspec_group): Returns a UVPSpec object constructed from the input HDF5 group. """ # Check that group is tagged as containing UVPSpec (pspec_type attribute) - if 'pspec_type' in pspec_group.attrs.keys(): + if 'pspec_type' in list(pspec_group.attrs.keys()): if pspec_group.attrs['pspec_type'] != uvpspec.UVPSpec.__name__: raise TypeError("HDF5 group is not tagged as a UVPSpec object.") else: @@ -104,13 +104,13 @@ def _update_header(self): Update the header in the HDF5 file with useful metadata, including the git version of hera_pspec. """ - if 'header' not in self.data.keys(): + if 'header' not in list(self.data.keys()): hdr = self.data.create_group('header') else: hdr = self.data['header'] # Check if versions of hera_pspec are the same - if 'hera_pspec.git_hash' in hdr.attrs.keys(): + if 'hera_pspec.git_hash' in list(hdr.attrs.keys()): if hdr.attrs['hera_pspec.git_hash'] != version.git_hash: print("WARNING: HDF5 file was created by a different version " "of hera_pspec.") @@ -138,13 +138,13 @@ def set_pspec(self, group, psname, pspec, overwrite=False): """ if self.mode == 'r': raise IOError("HDF5 file was opened read-only; cannot write to file.") - - if getattr(group, '__iter__', False): + + if isinstance(group, (tuple, list, dict)): raise ValueError("Only one group can be specified at a time.") # Handle input arguments that are iterable (i.e. sequences, but not str) - if getattr(psname, '__iter__', False): - if getattr(pspec, '__iter__', False) and len(pspec) == len(psname): + if isinstance(psname, list): + if isinstance(pspec, list) and len(pspec) == len(psname): # Recursively call set_pspec() on each item of the list for _psname, _pspec in zip(psname, pspec): if not isinstance(_pspec, uvpspec.UVPSpec): @@ -156,26 +156,26 @@ def set_pspec(self, group, psname, pspec, overwrite=False): # Raise exception if psname is a list, but pspec is not raise ValueError("If psname is a list, pspec must be a list of " "the same length.") - if getattr(pspec, '__iter__', False) \ - and not getattr(psname, '__iter__', False): + if isinstance(pspec, list) and not isinstance(psname, list): raise ValueError("If pspec is a list, psname must also be a list.") # No lists should pass beyond this point # Check that input is of the correct type if not isinstance(pspec, uvpspec.UVPSpec): + print("pspec:", type(pspec), pspec) raise TypeError("pspec must be a UVPSpec object.") key1 = "%s" % group key2 = "%s" % psname # Check that the group exists - if key1 not in self.data.keys(): + if key1 not in list(self.data.keys()): grp = self.data.create_group(key1) else: grp = self.data[key1] # Check that the psname exists - if key2 not in grp.keys(): + if key2 not in list(grp.keys()): # Create group if it doesn't exist psgrp = grp.create_group(key2) else: @@ -215,7 +215,7 @@ def get_pspec(self, group, psname=None): """ # Check that group is in keys and extract it if so key1 = "%s" % group - if key1 in self.data.keys(): + if key1 in list(self.data.keys()): grp = self.data[key1] else: raise KeyError("No group named '%s'" % key1) @@ -225,7 +225,7 @@ def get_pspec(self, group, psname=None): key2 = "%s" % psname # Load power spectrum if it exists - if key2 in grp.keys(): + if key2 in list(grp.keys()): return self._load_pspec(grp[key2]) else: raise KeyError("No pspec named '%s' in group '%s'" % (key2, key1)) @@ -234,7 +234,7 @@ def get_pspec(self, group, psname=None): # Otherwise, extract all available power spectra uvp = [] def pspec_filter(n, obj): - if u'pspec_type' in obj.attrs.keys(): + if u'pspec_type' in list(obj.attrs.keys()): uvp.append(self._load_pspec(obj)) # Traverse the entire set of groups/datasets looking for pspecs @@ -258,7 +258,7 @@ def spectra(self, group): """ # Check that group is in keys and extract it if so key1 = "%s" % group - if key1 in self.data.keys(): + if key1 in list(self.data.keys()): grp = self.data[key1] else: raise KeyError("No group named '%s'" % key1) @@ -266,7 +266,7 @@ def spectra(self, group): # Filter to look for pspec objects ps_list = [] def pspec_filter(n, obj): - if u'pspec_type' in obj.attrs.keys(): + if u'pspec_type' in list(obj.attrs.keys()): ps_list.append(n) # Traverse the entire set of groups/datasets looking for pspecs @@ -282,7 +282,7 @@ def groups(self): group_list : list of str List of group names. """ - groups = self.data.keys() + groups = list(self.data.keys()) if u'header' in groups: groups.remove(u'header') return groups @@ -373,7 +373,7 @@ def combine_psc_spectra(psc, groups=None, dset_split_str='_x_', ext_split_str='_ # Iterate over groups for grp in groups: # Get spectra in this group - spectra = psc.data[grp].keys() + spectra = list(psc.data[grp].keys()) # Get unique spectra by splitting and then re-joining unique_spectra = [] diff --git a/hera_pspec/conversions.py b/hera_pspec/conversions.py index fd7e5330..6d2eeabd 100644 --- a/hera_pspec/conversions.py +++ b/hera_pspec/conversions.py @@ -12,14 +12,14 @@ from astropy.cosmology import LambdaCDM _astropy = True except: - print("could not import astropy") + print("Could not import astropy") _astropy = False class units: """ - fundamental constants and conversion constants - in ** SI ** units + Fundamental constants and conversion constants + in ** SI ** units. c : speed of light m s-1 ckm : speed of light in km s-1 @@ -118,24 +118,26 @@ def __init__(self, Om_L=0.68440, Om_b=0.04911, Om_c=0.26442, H0=67.27, def get_params(self): """ - Return a dictionary with cosmological parameters + Return a dictionary with cosmological parameters. """ - return dict(map(lambda p: (p, getattr(self, p)), self.params)) + return dict([(p, getattr(self, p)) for p in self.params]) def f2z(self, freq, ghz=False): """ - convert frequency to redshift for 21cm line + Convert frequency to redshift for 21cm line. Parameters: ----------- - freq : frequency in Hz, type=float + freq : float + Frequency in Hz (or GHz if ghz=True). - ghz : boolean, if True: assume freq is GHz + ghz : bool, optional + If True: assume freq is GHz. Default: False. Output: ------- z : float - redshift + Redshift """ if ghz: freq = freq * 1e9 @@ -145,66 +147,70 @@ def f2z(self, freq, ghz=False): @staticmethod def z2f(z, ghz=False): """ - convert redshift to frequency in Hz for 21cm line + Convert redshift to frequency in Hz for 21cm line. Parameters: ----------- - z : redshift, type=float + z : float + Redshift. - ghz: boolean, if True: convert to GHz + ghz : bool, optional + If True: convert to GHz. Default: False. Output: ------- freq : float - frequency in Hz + Frequency in Hz (or GHz if ghz=True) """ freq = units.f21 / (z + 1) - if ghz: - freq /= 1e9 - + if ghz: freq /= 1e9 return freq def E(self, z): """ - ratio of hubble parameters: H(z) / H(z=0) - Hogg99 Eqn. 14 + Ratio of hubble parameters: H(z) / H(z=0) + (Hogg99 Eqn. 14) Parameters: ----------- - z : redshift, type=float + z : float + Redshift. """ return np.sqrt(self.Om_M*(1+z)**3 + self.Om_k*(1+z)**2 + self.Om_L) def DC(self, z, little_h=True): """ - line-of-sight comoving distance in Mpc - Hogg99 Eqn. 15 + Line-of-sight comoving distance in Mpc. + (Hogg99 Eqn. 15) Parameters: ----------- - z : redshift, type=float + z : float + Redshift. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc or Mpc - Default: h^-1 Mpc + Whether to have cosmological length units be h^-1 Mpc (True) or + Mpc (False). Default: True (h^-1 Mpc) """ + d = integrate.quad(lambda z: 1/self.E(z), 0, z)[0] if little_h: - return integrate.quad(lambda z: 1/self.E(z), 0, z)[0] * units.ckm / 100. + return d * units.ckm / 100. else: - return integrate.quad(lambda z: 1/self.E(z), 0, z)[0] * units.ckm / self.H0 + return d * units.ckm / self.H0 def DM(self, z, little_h=True): """ - transverse comoving distance in Mpc - Hogg99 Eqn. 16 + Transverse comoving distance in Mpc. + (Hogg99 Eqn. 16) Parameters: ----------- - z : redshift, type=float + z : float + Redshift. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc or Mpc - Default: h^-1 Mpc + Whether to have cosmological length units be h^-1 Mpc (True) or + Mpc (False). Default: True (h^-1 Mpc) """ if little_h: DH = units.ckm / 100. @@ -212,9 +218,13 @@ def DM(self, z, little_h=True): DH = units.ckm / self.H0 if self.Om_k > 0: - DM = DH * np.sinh(np.sqrt(self.Om_k) * self.DC(z, little_h=little_h) / DH) / np.sqrt(self.Om_k) + DM = DH * np.sinh(np.sqrt(self.Om_k) \ + * self.DC(z, little_h=little_h) / DH) \ + / np.sqrt(self.Om_k) elif self.Om_k < 0: - DM = DH * np.sin(np.sqrt(np.abs(self.Om_k)) * self.DC(z, little_h=little_h) / DH) / np.sqrt(np.abs(self.Om_k)) + DM = DH * np.sin(np.sqrt(np.abs(self.Om_k)) \ + * self.DC(z, little_h=little_h) / DH) \ + / np.sqrt(np.abs(self.Om_k)) else: DM = self.DC(z, little_h=little_h) @@ -222,47 +232,52 @@ def DM(self, z, little_h=True): def DA(self, z, little_h=True): """ - angular diameter (proper) distance in Mpc - Hogg99 Eqn. 18 + Angular diameter (proper) distance in Mpc. + (Hogg99 Eqn. 18) Parameters: ----------- - z : redshift, type=float + z : float + Redshift. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc or Mpc - Default: h^-1 Mpc + Whether to have cosmological length units be h^-1 Mpc (True) or + Mpc (False). Default: True (h^-1 Mpc) """ return self.DM(z, little_h=little_h) / (1 + z) def dRperp_dtheta(self, z, little_h=True): """ - conversion factor from angular size (radian) to transverse + Conversion factor from angular size (radian) to transverse comoving distance (Mpc) at a specific redshift: [Mpc / radians] Parameters: ----------- - z : float, redshift + z : float + Redshift. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc or Mpc - Default: h^-1 Mpc + Whether to have cosmological length units be h^-1 Mpc (True) or + Mpc (False). Default: True (h^-1 Mpc) """ return self.DM(z, little_h=little_h) def dRpara_df(self, z, ghz=False, little_h=True): """ - conversion from frequency bandwidth to radial - comoving distance at a specific redshift: [Mpc / Hz] + Conversion from frequency bandwidth to radial comoving distance at a + specific redshift: [Mpc / Hz] Parameters: ----------- - z : float, redshift - ghz : convert output to [Mpc / GHz] + z : float + Redshift. + + ghz : bool, optional + Whether to convert output to [Mpc / GHz] (if True). Default: False. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc or Mpc - Default: h^-1 Mpc + Whether to have cosmological length units be h^-1 Mpc (True) or + Mpc (False). Default: True (h^-1 Mpc) """ if little_h: y = (1 + z)**2.0 / self.E(z) * units.ckm / 100. / units.f21 @@ -275,34 +290,36 @@ def dRpara_df(self, z, ghz=False, little_h=True): def X2Y(self, z, little_h=True): """ - Conversion from radians^2 Hz -> Mpc^3 - at a specific redshift. + Conversion from radians^2 Hz -> Mpc^3 at a specific redshift. Parameters: ----------- - z : float, redshift + z : float + Redshift. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc or Mpc - Default: h^-1 Mpc + Whether to have cosmological length units be h^-1 Mpc (True) or + Mpc (False). Default: True (h^-1 Mpc) Notes: ------ Calls Cosmo_Conversions.dRperp_dtheta() and Cosmo_Conversions.dRpara_df(). """ - return self.dRperp_dtheta(z, little_h=little_h)**2 * self.dRpara_df(z, little_h=little_h) + return self.dRperp_dtheta(z, little_h=little_h)**2 \ + * self.dRpara_df(z, little_h=little_h) def bl_to_kperp(self, z, little_h=True): """ - Produce the conversion factor from baseline length [meters] to k_perpendicular mode [h Mpc-1] at a - specified redshift. + Produce the conversion factor from baseline length [meters] to + k_perpendicular mode [h Mpc-1] at a specified redshift. - Multiply this conversion factor by a baseline-separation length in [meters] - to get its corresponding k_perp mode in [h Mpc-1]. + Multiply this conversion factor by a baseline-separation length in + [meters] to get its corresponding k_perp mode in [h Mpc-1]. Parameters ---------- - z : float, redshift at which to perform calculation + z : float + Redshift. little_h : boolean, optional Whether to have cosmological length units be h^-1 Mpc or Mpc @@ -310,23 +327,26 @@ def bl_to_kperp(self, z, little_h=True): Return ------ - bl2kperp : float, conversion factor in units [h Mpc-1 / meters] + bl2kperp : float + Conversion factor in units [h Mpc-1 / meters] """ # Parsons 2012, Pober 2014, Kohn 2018 - bl2kpara = 2*np.pi / (self.dRperp_dtheta(z, little_h=little_h) * (units.c / self.z2f(z))) - + bl2kpara = 2*np.pi / (self.dRperp_dtheta(z, little_h=little_h) \ + * (units.c / self.z2f(z))) return bl2kpara def tau_to_kpara(self, z, little_h=True): """ - Produce the conversion factor from delay [seconds] to k_parallel mode [h Mpc-1] at a specified redshift. + Produce the conversion factor from delay [seconds] to k_parallel mode + [h Mpc-1] at a specified redshift. Multiply this conversion factor by a delay mode in [seconds] to get its corresponding k_para mode in [h Mpc-1]. Parameters ---------- - z : float, redshift at which to perform calculation + z : float + Redshift. little_h : boolean, optional Whether to have cosmological length units be h^-1 Mpc or Mpc @@ -334,7 +354,8 @@ def tau_to_kpara(self, z, little_h=True): Return ------ - tau2kpara : float, conversion factor in units [h Mpc-1 / seconds] + tau2kpara : float + Conversion factor in units [h Mpc-1 / seconds] """ # Parsons 2012, Pober 2014, Kohn 2018 tau2kpara = 2*np.pi / self.dRpara_df(z, little_h=little_h, ghz=False) @@ -343,10 +364,13 @@ def tau_to_kpara(self, z, little_h=True): def __str__(self): message = "Cosmo_Conversions object at <{}>\n".format(hex(id(self))) - message += "; ".join(map(lambda p: "{:s} : {:0.4f}".format(p, getattr(self, p)), self.params)) + message += "; ".join( ["{:s} : {:0.4f}".format(p, getattr(self, p)) + for p in self.params] ) return message def __eq__(self, other): - """ check two Cosmo_Conversion objects are equivalent """ + """ + Check two Cosmo_Conversion objects are equivalent + """ return self.get_params() == other.get_params() diff --git a/hera_pspec/grouping.py b/hera_pspec/grouping.py index 4d34f7e7..2350296b 100644 --- a/hera_pspec/grouping.py +++ b/hera_pspec/grouping.py @@ -47,7 +47,7 @@ def group_baselines(bls, Ngroups, keep_remainder=False, randomize=False, List of grouped baselines. """ Nbls = len(bls) # Total baselines - n = Nbls / Ngroups # Baselines per group + n = Nbls // Ngroups # Baselines per group rem = Nbls - n*Ngroups # Sanity check on number of groups @@ -192,15 +192,15 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, # Convert blpair_groups to list of blpair group integers if isinstance(blpair_groups[0][0], tuple): - new_blpair_grps = [map(lambda blp: uvp.antnums_to_blpair(blp), blpg) + new_blpair_grps = [[uvp.antnums_to_blpair(blp) for blp in blpg] for blpg in blpair_groups] blpair_groups = new_blpair_grps else: # If not, each baseline pair is its own group - blpair_groups = map(lambda blp: [blp], np.unique(uvp.blpair_array)) + blpair_groups = [[blp] for blp in np.unique(uvp.blpair_array)] assert blpair_weights is None, "Cannot specify blpair_weights if "\ "blpair_groups is None." - blpair_weights = map(lambda blp: [1.,], np.unique(uvp.blpair_array)) + blpair_weights = [[1.,] for blp in np.unique(uvp.blpair_array)] # Print warning if a blpair appears more than once in all of blpair_groups all_blpairs = [item for sublist in blpair_groups for item in sublist] @@ -236,8 +236,8 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, # For baseline pairs not in blpair_groups, add them as their own group extra_blpairs = set(uvp.blpair_array) - set(all_blpairs) - blpair_groups += map(lambda blp: [blp], extra_blpairs) - blpair_weights += map(lambda blp: [1.,], extra_blpairs) + blpair_groups += [[blp] for blp in extra_blpairs] + blpair_weights += [[1.,] for blp in extra_blpairs] # Create new data arrays data_array, wgts_array = odict(), odict() @@ -255,9 +255,9 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, spw_stats = odict([[stat, []] for stat in stat_l]) if store_cov: spw_cov = [] - + # Iterate over polarizations - for i, p in enumerate(uvp.pol_array): + for i, p in enumerate(uvp.polpair_array): pol_data, pol_wgts, pol_ints, pol_nsmp = [], [], [], [] pol_stats = odict([[stat, []] for stat in stat_l]) if store_cov: @@ -413,7 +413,8 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, # Update arrays bl_arr = np.array(sorted(set(bl_arr))) - bl_vecs = np.array(map(lambda bl: uvp.bl_vecs[uvp.bl_array.tolist().index(bl)], bl_arr)) + bl_vecs = np.array([uvp.bl_vecs[uvp.bl_array.tolist().index(bl)] + for bl in bl_arr]) # Assign arrays and metadata to UVPSpec object uvp.Ntimes = len(np.unique(time_avg_arr)) @@ -506,15 +507,15 @@ def fold_spectra(uvp): +leftright\ +rightleft\ +rightright) - uvp.cov_array[spw][:, :Ndlys/2, :, :] = 0.0 - uvp.cov_array[spw][:, :, :Ndlys/2, : :] = 0.0 + uvp.cov_array[spw][:, :Ndlys//2, :, :] = 0.0 + uvp.cov_array[spw][:, :, :Ndlys//2, : :] = 0.0 # fold stats array if it exists: sum in inverse quadrature if hasattr(uvp, 'stats_array'): for stat in uvp.stats_array.keys(): left = uvp.stats_array[stat][spw][:, 1:Ndlys//2, :][:, ::-1, :] right = uvp.stats_array[stat][spw][:, Ndlys//2+1:, :] - uvp.stats_array[stat][spw][:, Ndlys//2+1:, :] = (np.sum([1/left**2.0, 1/right**2.0], axis=0))**(-0.5) + uvp.stats_array[stat][spw][:, Ndlys//2+1:, :] = (np.sum([1./left**2.0, 1./right**2.0], axis=0))**(-0.5) uvp.data_array[spw][:, :Ndlys//2, :] = np.nan else: @@ -535,15 +536,15 @@ def fold_spectra(uvp): +leftright\ +rightleft\ +rightright) - uvp.cov_array[spw][:, :Ndlys/2, :, :] = 0.0 - uvp.cov_array[spw][:, :, :Ndlys/2, : :] = 0.0 + uvp.cov_array[spw][:, :Ndlys//2, :, :] = 0.0 + uvp.cov_array[spw][:, :, :Ndlys//2, : :] = 0.0 # fold stats array if it exists: sum in inverse quadrature if hasattr(uvp, 'stats_array'): for stat in uvp.stats_array.keys(): left = uvp.stats_array[stat][spw][:, :Ndlys//2, :][:, ::-1, :] right = uvp.stats_array[stat][spw][:, Ndlys//2+1:, :] - uvp.stats_array[stat][spw][:, Ndlys//2+1:, :] = (np.sum([1/left**2.0, 1/right**2.0], axis=0))**(-0.5) + uvp.stats_array[stat][spw][:, Ndlys//2+1:, :] = (np.sum([1./left**2.0, 1./right**2.0], axis=0))**(-0.5) uvp.data_array[spw][:, :Ndlys//2, :] = np.nan uvp.folded = True @@ -631,14 +632,15 @@ def bootstrap_average_blpairs(uvp_list, blpair_groups, time_avg=False, # Convert blpair tuples into blpair integers if necessary if isinstance(blpair_groups[0][0], tuple): - new_blp_grps = [map(lambda blp: uvputils._antnums_to_blpair(blp), blpg) + new_blp_grps = [[uvputils._antnums_to_blpair(blp) for blp in blpg] for blpg in blpair_groups] blpair_groups = new_blp_grps # Homogenise input UVPSpec objects in terms of available polarizations # and spectral windows if len(uvp_list) > 1: - uvp_list = uvpspec_utils.select_common(uvp_list, spws=True, pols=True, inplace=False) + uvp_list = uvpspec_utils.select_common(uvp_list, spws=True, + pols=True, inplace=False) # Loop over UVPSpec objects, looking for available blpairs in each avail_blpairs = [np.unique(uvp.blpair_array) for uvp in uvp_list] @@ -785,12 +787,14 @@ def bootstrap_resampled_error(uvp, blpair_groups=None, time_avg=False, Nsamples= # get all keys in uvp_avg and get data from each uvp_boot keys = uvp_avg.get_all_keys() - uvp_boot_data = odict([(k, np.array(map(lambda u: u.get_data(k), uvp_boots))) for k in keys]) + uvp_boot_data = odict([(k, np.array([u.get_data(k) for u in uvp_boots])) + for k in keys]) # calculate various error estimates if normal_std: for k in keys: - nstd = np.std(uvp_boot_data[k].real, axis=0) + 1j*np.std(uvp_boot_data[k].imag, axis=0) + nstd = np.std(uvp_boot_data[k].real, axis=0) \ + + 1j*np.std(uvp_boot_data[k].imag, axis=0) uvp_avg.set_stats("bs_std", k, nstd) if robust_std: @@ -895,7 +899,9 @@ def bootstrap_run(filename, spectra=None, blpair_groups=None, time_avg=False, Ns assert len(groups) > 0, "No groups exist in PSpecContainer" # get spectra if not fed - all_spectra = utils.flatten([map(lambda s: os.path.join(grp, s), psc.spectra(grp)) for grp in groups]) + all_spectra = utils.flatten([ [os.path.join(grp, s) + for s in psc.spectra(grp)] + for grp in groups ]) if spectra is None: spectra = all_spectra else: @@ -910,11 +916,15 @@ def bootstrap_run(filename, spectra=None, blpair_groups=None, time_avg=False, Ns # run boostrap_resampled_error uvp = psc.get_pspec(grp, spc) (uvp_avg, uvp_boots, - uvp_wgts) = bootstrap_resampled_error(uvp, blpair_groups=blpair_groups, time_avg=time_avg, - Nsamples=Nsamples, seed=seed, normal_std=normal_std, - robust_std=robust_std, cintervals=cintervals, - bl_error_tol=bl_error_tol, add_to_history=add_to_history, - verbose=verbose) + uvp_wgts) = bootstrap_resampled_error(uvp, blpair_groups=blpair_groups, + time_avg=time_avg, + Nsamples=Nsamples, seed=seed, + normal_std=normal_std, + robust_std=robust_std, + cintervals=cintervals, + bl_error_tol=bl_error_tol, + add_to_history=add_to_history, + verbose=verbose) # set averaged uvp psc.set_pspec(grp, spc+"_avg", uvp_avg, overwrite=overwrite) @@ -930,7 +940,7 @@ def get_bootstrap_run_argparser(): description="argument parser for grouping.bootstrap_run()") def list_of_lists_of_tuples(s): - s = map(lambda x: map(int, x.split()), s.split(',')) + s = [[int(_x) for _x in x.split()] for x in s.split(',')] return s # Add list of arguments diff --git a/hera_pspec/plot.py b/hera_pspec/plot.py index d00f751c..130526a9 100644 --- a/hera_pspec/plot.py +++ b/hera_pspec/plot.py @@ -294,10 +294,12 @@ def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False, return fig -def delay_waterfall(uvp, blpairs, spw, pol, component='real', average_blpairs=False, - fold=False, delay=True, deltasq=False, log=True, lst_in_hrs=True, - vmin=None, vmax=None, cmap='YlGnBu', axes=None, figsize=(14, 6), - force_plot=False, times=None, title_type='blpair'): +def delay_waterfall(uvp, blpairs, spw, pol, component='real', + average_blpairs=False, fold=False, delay=True, + deltasq=False, log=True, lst_in_hrs=True, + vmin=None, vmax=None, cmap='YlGnBu', axes=None, + figsize=(14, 6), force_plot=False, times=None, + title_type='blpair'): """ Plot a 1D delay spectrum waterfall (or spectra) for a group of baselines. @@ -484,7 +486,8 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='real', average_blpairs=Fa Nvert, Nhorz = axes.shape # Get LST range: setting y-ticks is tricky due to LST wrapping... - y = uvp_plt.lst_avg_array[uvp_plt.key_to_indices(waterfall.keys()[0])[1]] + y = uvp_plt.lst_avg_array[ + uvp_plt.key_to_indices(list(waterfall.keys())[0])[1] ] if lst_in_hrs: lst_units = "Hr" y = np.around(y * 24 / (2*np.pi), 2) @@ -516,7 +519,7 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='real', average_blpairs=Fa r"{\rm unnormed}") # Iterate over waterfall keys - keys = waterfall.keys() + keys = list(waterfall.keys()) k = 0 for i in range(Nvert): for j in range(Nhorz): @@ -533,8 +536,9 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='real', average_blpairs=Fa blp = uvp_plt.blpair_to_antnums(key[1]) # plot waterfall - cax = ax.matshow(waterfall[key], cmap=cmap, aspect='auto', vmin=vmin, vmax=vmax, - extent=[np.min(x), np.max(x), Ny, 0]) + cax = ax.matshow(waterfall[key], cmap=cmap, aspect='auto', + vmin=vmin, vmax=vmax, + extent=[np.min(x), np.max(x), Ny, 0]) # ax config ax.xaxis.set_ticks_position('bottom') @@ -593,13 +597,14 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True, red_tol=1.0, center_line=False, horizon_lines=False, title=None, ax=None, cmap='viridis', figsize=(8, 6), deltasq=False, colorbar=False, cbax=None, vmin=None, vmax=None, - edgecolor='none', flip_xax=False, flip_yax=False, lw=2, set_bl_tick_major=False, - set_bl_tick_minor=False, xtick_size=10, xtick_rot=0, ytick_size=10, ytick_rot=0, + edgecolor='none', flip_xax=False, flip_yax=False, lw=2, + set_bl_tick_major=False, set_bl_tick_minor=False, + xtick_size=10, xtick_rot=0, ytick_size=10, ytick_rot=0, **kwargs): """ Plot a 2D delay spectrum (or spectra) from a UVPSpec object. Note that - all integrations and redundant baselines are averaged (unless specifying times) - before plotting. + all integrations and redundant baselines are averaged (unless specifying + times) before plotting. Parameters ---------- @@ -609,8 +614,8 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True, spw : integer Which spectral window to plot. - pol : int or str - Polarization integer or string + pol : int or tuple + Polarization-pair integer or tuple, e.g. ('pI', 'pI') blpairs : list of tuples, optional List of baseline-pair tuples to use in plotting. @@ -696,8 +701,8 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True, Line-width of horizon and center lines if plotted. Default: 2. set_bl_tick_major : bool, optional - If True, use the baseline lengths as major ticks, rather than default uniform - grid. + If True, use the baseline lengths as major ticks, rather than default + uniform grid. set_bl_tick_minor : bool, optional If True, use the baseline lengths as minor ticks, which have no labels. @@ -714,7 +719,7 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True, uvp = copy.deepcopy(uvp) assert isinstance(uvp, uvpspec.UVPSpec), "input uvp must be a UVPSpec object" assert isinstance(spw, (int, np.integer)) - assert isinstance(pol, (int, str, np.integer, np.str)) + assert isinstance(pol, (int, np.integer, tuple)) # check pspec units for little h little_h = 'h^-3' in uvp.norm_units @@ -733,7 +738,8 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True, # Average across redundant groups and time # this also ensures blpairs are ordered from short_bl --> long_bl - blp_grps, lens, angs, tags = utils.get_blvec_reds(uvp, bl_error_tol=red_tol, match_bl_lens=True) + blp_grps, lens, angs, tags = utils.get_blvec_reds(uvp, bl_error_tol=red_tol, + match_bl_lens=True) uvp.average_spectra(blpair_groups=blp_grps, time_avg=True, inplace=True) # Convert to DeltaSq @@ -814,14 +820,16 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True, # Configure ticks if set_bl_tick_major: if rotate: - ax.set_xticks(map(lambda x: np.around(x, _get_sigfig(x)+2), x_axis)) + ax.set_xticks([np.around(x, _get_sigfig(x)+2) for x in x_axis]) else: - ax.set_yticks(map(lambda x: np.around(x, _get_sigfig(x)+2), y_axis)) + ax.set_yticks([np.around(x, _get_sigfig(x)+2) for x in y_axis]) if set_bl_tick_minor: if rotate: - ax.set_xticks(map(lambda x: np.around(x, _get_sigfig(x)+2), x_axis), minor=True) + ax.set_xticks([np.around(x, _get_sigfig(x)+2) for x in x_axis], + minor=True) else: - ax.set_yticks(map(lambda x: np.around(x, _get_sigfig(x)+2), y_axis), minor=True) + ax.set_yticks([np.around(x, _get_sigfig(x)+2) for x in y_axis], + minor=True) # Add colorbar if colorbar: diff --git a/hera_pspec/pspecbeam.py b/hera_pspec/pspecbeam.py index 33be3095..4e31b33e 100644 --- a/hera_pspec/pspecbeam.py +++ b/hera_pspec/pspecbeam.py @@ -1,6 +1,7 @@ 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 @@ -71,7 +72,7 @@ def _compute_pspec_scalar(cosmo, beam_freqs, omega_ratio, pspec_freqs, # Get redshifts and cosmological functions redshifts = cosmo.f2z(integration_freqs).flatten() - X2Y = np.array(map(lambda z: cosmo.X2Y(z, little_h=little_h), redshifts)) + X2Y = np.array([cosmo.X2Y(z, little_h=little_h) for z in redshifts]) # Use linear interpolation to interpolate the frequency-dependent # quantities derived from the beam model to the same frequency grid as the @@ -238,14 +239,15 @@ def Jy_to_mK(self, freqs, pol='pI'): return 1e-20 * conversions.cgs_units.c**2 \ / (2 * conversions.cgs_units.kb * freqs**2 * Op) - def get_Omegas(self, pols): + def get_Omegas(self, polpairs): """ - Get OmegaP and OmegaPP across beam_freqs for requested polarizations. + Get OmegaP and OmegaPP across beam_freqs for requested polarization + pairs. Parameters ---------- - pols : list - List of polarization strings or integers. + polpairs : list + List of polarization-pair tuples or integers. Returns ------- @@ -255,24 +257,38 @@ def get_Omegas(self, pols): OmegaPP : array_like Array containing power_sq_beam_int, shape: (Nbeam_freqs, Npols). """ - # type check - if isinstance(pols, (int, np.int, np.int32)): - pols = [pols] - elif isinstance(pols, (str, np.str)): - pols = [pols] - if isinstance(pols, (list, np.ndarray, tuple)): - if isinstance(pols[0], (int, np.int, np.int32)): - pols = map(lambda p: uvutils.polnum2str(p), pols) - - # initialize blank lists + # Unpack polpairs into tuples + if not isinstance(polpairs, (list, np.ndarray)): + if isinstance(polpairs, (tuple, int, np.integer)): + polpairs = [polpairs,] + else: + raise TypeError("polpairs is not a list of integers or tuples") + + # Convert integers to tuples + polpairs = [uvputils.polpair_int2tuple(p) + if isinstance(p, (int, np.integer, np.int32)) else p + for p in polpairs] + + # Calculate Omegas for each pol pair OmegaP, OmegaPP = [], [] - for p in pols: - OmegaP.append(self.power_beam_int(pol=p)) - OmegaPP.append(self.power_beam_sq_int(pol=p)) + for pol1, pol2 in polpairs: + if isinstance(pol1, (int, np.integer)): + pol1 = uvutils.polnum2str(pol1) + if isinstance(pol2, (int, np.integer)): + pol2 = uvutils.polnum2str(pol2) + + # Check for cross-pol; only auto-pol calculation currently supported + if pol1 != pol2: + raise NotImplementedError( + "get_Omegas does not support cross-polarized beams yet. " + "Could not calculate Omegas for (%s, %s)" % (pol1, pol2)) + + # Calculate Omegas + OmegaP.append(self.power_beam_int(pol=pol1)) + OmegaPP.append(self.power_beam_sq_int(pol=pol1)) OmegaP = np.array(OmegaP).T OmegaPP = np.array(OmegaPP).T - return OmegaP, OmegaPP diff --git a/hera_pspec/pspecdata.py b/hera_pspec/pspecdata.py index 8800013a..d738bfca 100644 --- a/hera_pspec/pspecdata.py +++ b/hera_pspec/pspecdata.py @@ -119,7 +119,7 @@ def add(self, dsets, wgts, labels=None, dsets_std=None): dsets_std = _dsets_std # Unpack dsets and wgts dicts - labels = dsets.keys() + labels = list(dsets.keys()) _dsets = [dsets[key] for key in labels] _wgts = [wgts[key] for key in labels] dsets = _dsets @@ -238,31 +238,44 @@ def validate_datasets(self, verbose=True): raise ValueError("all dsets must have the same Ntimes") # raise warnings if times don't match - lst_diffs = np.array(map(lambda dset: np.unique(self.dsets[0].lst_array) - np.unique(dset.lst_array), self.dsets[1:])) + lst_diffs = np.array( [ np.unique(self.dsets[0].lst_array) + - np.unique(dset.lst_array) + for dset in self.dsets[1:]] ) if np.max(np.abs(lst_diffs)) > 0.001: - raise_warning("Warning: taking power spectra between LST bins misaligned by more than 15 seconds", - verbose=verbose) + raise_warning("Warning: taking power spectra between LST bins " + "misaligned by more than 15 seconds", + verbose=verbose) # raise warning if frequencies don't match - freq_diffs = np.array(map(lambda dset: np.unique(self.dsets[0].freq_array) - np.unique(dset.freq_array), self.dsets[1:])) + freq_diffs = np.array( [ np.unique(self.dsets[0].freq_array) + - np.unique(dset.freq_array) + for dset in self.dsets[1:]] ) if np.max(np.abs(freq_diffs)) > 0.001e6: - raise_warning("Warning: taking power spectra between frequency bins misaligned by more than 0.001 MHz", + raise_warning("Warning: taking power spectra between frequency " + "bins misaligned by more than 0.001 MHz", verbose=verbose) # Check phase type phase_types = [] for d in self.dsets: phase_types.append(d.phase_type) if np.unique(phase_types).size > 1: - raise ValueError("all datasets must have the same phase type (i.e. 'drift', 'phased', ...)\ncurrent phase types are {}".format(phase_types)) + raise ValueError("all datasets must have the same phase type " + "(i.e. 'drift', 'phased', ...)\ncurrent phase " + "types are {}".format(phase_types)) # Check phase centers if phase type is phased if 'phased' in set(phase_types): - phase_ra = map(lambda d: d.phase_center_ra_degrees, self.dsets) - phase_dec = map(lambda d: d.phase_center_dec_degrees, self.dsets) - max_diff_ra = np.max(map(lambda d: np.diff(d), itertools.combinations(phase_ra, 2))) - max_diff_dec = np.max(map(lambda d: np.diff(d), itertools.combinations(phase_dec, 2))) + phase_ra = [d.phase_center_ra_degrees for d in self.dsets] + phase_dec = [d.phase_center_dec_degrees for d in self.dsets] + max_diff_ra = np.max( [np.diff(d) + for d in itertools.combinations(phase_ra, 2)]) + max_diff_dec = np.max([np.diff(d) + for d in itertools.combinations(phase_dec, 2)]) max_diff = np.sqrt(max_diff_ra**2 + max_diff_dec**2) - if max_diff > 0.15: raise_warning("Warning: maximum phase-center difference between datasets is > 10 arcmin", verbose=verbose) + if max_diff > 0.15: + raise_warning("Warning: maximum phase-center difference " + "between datasets is > 10 arcmin", + verbose=verbose) def check_key_in_dset(self, key, dset_ind): """ @@ -377,13 +390,13 @@ def parse_blkey(self, key): # get baseline bl = key[0] - if isinstance(bl, (int, np.int, np.int32)): + if isinstance(bl, (int, np.integer)): assert len(key) > 1, "baseline must be fed as a tuple" bl = tuple(key[:2]) key = key[2:] else: key = key[1:] - assert isinstance(bl, tuple), "baseline must be fed as a tuple" + assert isinstance(bl, tuple), "baseline must be fed as a tuple, %s" % bl # put pol into bl key if it exists if len(key) > 0: @@ -512,14 +525,15 @@ def C_model(self, key, model='empirical'): Ckey = key + (model,) # check cache - if not self._C.has_key(Ckey): + if Ckey not in self._C: # calculate covariance model if model == 'empirical': self.set_C({Ckey: utils.cov(self.x(key), self.w(key))}) return self._C[Ckey] - def cross_covar_model(self, key1, key2, model='empirical', conj_1=False, conj_2=True): + def cross_covar_model(self, key1, key2, model='empirical', + conj_1=False, conj_2=True): """ Return a covariance model having specified a key and model type. @@ -531,13 +545,16 @@ def cross_covar_model(self, key1, key2, model='empirical', conj_1=False, conj_2= subsequent indices specify the baseline index, in _key2inds format. model : string, optional - Type of covariance model to calculate, if not cached. options=['empirical'] + Type of covariance model to calculate, if not cached. + options=['empirical'] conj_1 : boolean, optional - Whether to conjugate first copy of data in covar or not. Default: False + Whether to conjugate first copy of data in covar or not. + Default: False conj_2 : boolean, optional - Whether to conjugate second copy of data in covar or not. Default: True + Whether to conjugate second copy of data in covar or not. + Default: True Returns ------- @@ -583,7 +600,7 @@ def I(self, key): dset, bl = self.parse_blkey(key) key = (dset,) + (bl,) - if not self._I.has_key(key): + if key not in self._I: self._I[key] = np.identity(self.spw_Nfreqs) return self._I[key] @@ -599,7 +616,8 @@ def iC(self, key, model='empirical'): subsequent indices specify the baseline index, in _key2inds format. model : string - Type of covariance model to calculate, if not cached. options=['empirical'] + Type of covariance model to calculate, if not cached. + options=['empirical'] Returns ------- @@ -614,7 +632,7 @@ def iC(self, key, model='empirical'): Ckey = key + (model,) # Calculate inverse covariance if not in cache - if not self._iC.has_key(Ckey): + if Ckey not in self._iC: C = self.C_model(key, model=model) U,S,V = np.linalg.svd(C.conj()) # conj in advance of next step @@ -659,9 +677,10 @@ def Y(self, key): dset, bl = self.parse_blkey(key) key = (dset,) + (bl,) - if not self._Y.has_key(key): + if key not in self._Y: self._Y[key] = np.diag(np.max(self.w(key), axis=1)) - if not np.all(np.isclose(self._Y[key], 0.0) + np.isclose(self._Y[key], 1.0)): + if not np.all(np.isclose(self._Y[key], 0.0) \ + + np.isclose(self._Y[key], 1.0)): raise NotImplementedError("Non-binary weights not currently implmented") return self._Y[key] @@ -720,17 +739,19 @@ def R(self, key): key = (dset,) + (bl,) Rkey = key + (self.data_weighting,) + (self.taper,) - if not self._R.has_key(Rkey): + if Rkey not in self._R: # form sqrt(taper) matrix if self.taper == 'none': sqrtT = np.ones(self.spw_Nfreqs).reshape(1, -1) else: sqrtT = np.sqrt(aipy.dsp.gen_window(self.spw_Nfreqs, self.taper)).reshape(1, -1) - # get flag weight vector: straight multiplication of vectors mimics matrix multiplication + # get flag weight vector: straight multiplication of vectors + # mimics matrix multiplication sqrtY = np.sqrt(self.Y(key).diagonal().reshape(1, -1)) - # replace possible nans with zero (when something dips negative in sqrt for some reason) + # replace possible nans with zero (when something dips negative + # in sqrt for some reason) sqrtT[np.isnan(sqrtT)] = 0.0 sqrtY[np.isnan(sqrtY)] = 0.0 @@ -946,7 +967,7 @@ def q_hat(self, key1, key2, allow_fft=False): else: q = [] - for i in xrange(self.spw_Ndlys): + for i in range(self.spw_Ndlys): Q = self.get_Q_alt(i) QRx2 = np.dot(Q, Rx2) qi = np.einsum('i...,i...->...', Rx1.conj(), QRx2) @@ -987,13 +1008,13 @@ def get_G(self, key1, key2): R2 = self.R(key2) iR1Q, iR2Q = {}, {} - for ch in xrange(self.spw_Ndlys): + for ch in range(self.spw_Ndlys): Q = self.get_Q_alt(ch) iR1Q[ch] = np.dot(R1, Q) # R_1 Q iR2Q[ch] = np.dot(R2, Q) # R_2 Q - for i in xrange(self.spw_Ndlys): - for j in xrange(self.spw_Ndlys): + for i in range(self.spw_Ndlys): + for j in range(self.spw_Ndlys): # tr(R_2 Q_i R_1 Q_j) G[i,j] += np.einsum('ab,ba', iR1Q[i], iR2Q[j]) @@ -1063,8 +1084,8 @@ def get_H(self, key1, key2, sampling=False): """ if self.spw_Ndlys == None: raise ValueError("Number of delay bins should have been set" - "by now! Cannot be equal to None") - + "by now! Cannot be equal to None.") + H = np.zeros((self.spw_Ndlys, self.spw_Ndlys), dtype=np.complex) R1 = self.R(key1) R2 = self.R(key2) @@ -1077,7 +1098,7 @@ def get_H(self, key1, key2, sampling=False): sinc_matrix = np.sinc(sinc_matrix / np.float(self.spw_Ndlys)) iR1Q_alt, iR2Q = {}, {} - for ch in xrange(self.spw_Ndlys): + for ch in range(self.spw_Ndlys): Q_alt = self.get_Q_alt(ch) iR1Q_alt[ch] = np.dot(R1, Q_alt) # R_1 Q_alt Q = Q_alt @@ -1087,8 +1108,8 @@ def get_H(self, key1, key2, sampling=False): iR2Q[ch] = np.dot(R2, Q) # R_2 Q - for i in xrange(self.spw_Ndlys): # this loop goes as nchan^4 - for j in xrange(self.spw_Ndlys): + for i in range(self.spw_Ndlys): # this loop goes as nchan^4 + for j in range(self.spw_Ndlys): # tr(R_2 Q_i R_1 Q_j) H[i,j] += np.einsum('ab,ba', iR1Q_alt[i], iR2Q[j]) @@ -1586,7 +1607,8 @@ def delays(self): return utils.get_delays(self.freqs[self.spw_range[0]:self.spw_range[1]], n_dlys=self.spw_Ndlys) * 1e9 # convert to ns - def scalar(self, pol, little_h=True, num_steps=2000, beam=None, taper_override='no_override'): + def scalar(self, polpair, little_h=True, num_steps=2000, beam=None, + taper_override='no_override'): """ Computes the scalar function to convert a power spectrum estimate in "telescope units" to cosmological units, using self.spw_range to set @@ -1599,9 +1621,9 @@ def scalar(self, pol, little_h=True, num_steps=2000, beam=None, taper_override=' Parameters ---------- - pol: str - Which polarization to compute the scalar for. - e.g. 'I', 'Q', 'U', 'V', 'XX', 'YY'... + polpair: tuple or int + Which pair of polarizations to compute the beam scalar for. + e.g. ('pI', 'pI') or ('XX', 'YY'). little_h : boolean, optional Whether to have cosmological length units be h^-1 Mpc or Mpc @@ -1627,6 +1649,16 @@ def scalar(self, pol, little_h=True, num_steps=2000, beam=None, taper_override=' [\int dnu (\Omega_PP / \Omega_P^2) ( B_PP / B_P^2 ) / (X^2 Y)]^-1 in h^-3 Mpc^3 or Mpc^3. """ + # make sure polarizations are the same + if isinstance(polpair, int): + polpair = uvutils.polpair_int2tuple(polpair) + if polpair[0] != polpair[1]: + raise NotImplementedError( + "Polarizations don't match. Beam scalar can only be " + "calculated for auto-polarization pairs at the moment.") + else: + pol = polpair[0] + # set spw_range and get freqs freqs = self.freqs[self.spw_range[0]:self.spw_range[1]] start = freqs[0] @@ -1751,8 +1783,8 @@ def validate_pol(self, dsets, pol_pair): assert isinstance(pol_pair[0], (int, np.integer)), err_msg assert isinstance(pol_pair[1], (int, np.integer)), err_msg - if pol_pair[0] != pol_pair[1]: - raise NotImplementedError("Only auto/equal polarizations are implement at the moment.") + #if pol_pair[0] != pol_pair[1]: + # raise NotImplementedError("Only auto/equal polarizations are implement at the moment.") dset_ind1 = self.dset_idx(dsets[0]) dset_ind2 = self.dset_idx(dsets[1]) @@ -1770,9 +1802,10 @@ def validate_pol(self, dsets, pol_pair): return valid - def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identity', - norm='I', taper='none', sampling=False, little_h=True, spw_ranges=None, - verbose=True, history='', store_cov=False): + def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, + input_data_weight='identity', norm='I', taper='none', + sampling=False, little_h=True, spw_ranges=None, + store_cov=False, verbose=True, history=''): """ Estimate the delay power spectrum from a pair of datasets contained in this object, using the optimal quadratic estimator of arXiv:1502.06016. @@ -1805,9 +1838,11 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit where the first index is for the Left-Hand dataset and second index is used for the Right-Hand dataset (see above). - pols : length-2 tuple of strings or integers or list of length-2 tuples of strings or integers + pols : length-2 tuple of strings or integers or list of length-2 + tuples of strings or integers Contains polarization pairs to use in forming power spectra - e.g. ('XX','XX') or [('XX','XX'),('XY','YX')] or list of polarization pairs. + e.g. ('XX','XX') or [('XX','XX'),('pI','pI')] or list of + polarization pairs. Only auto/equal polarization pairs are implemented at the moment. It uses the polarizations of the UVData onjects (specified in dsets) by default only if the UVData object consists of equal polarizations. @@ -1850,7 +1885,7 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit If True, calculate an analytic covariance between bandpowers given an input visibility noise model, and store the output in the UVPSpec object. - + verbose : bool, optional If True, print progress, warnings and debugging info to stdout. @@ -1917,22 +1952,30 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit verbose=verbose) # get datasets - assert isinstance(dsets, (list, tuple)), "dsets must be fed as length-2 tuple of integers" + assert isinstance(dsets, (list, tuple)), \ + "dsets must be fed as length-2 tuple of integers" assert len(dsets) == 2, "len(dsets) must be 2" - assert isinstance(dsets[0], (int, np.int)) and isinstance(dsets[1], (int, np.int)), "dsets must contain integer indices" + assert isinstance(dsets[0], (int, np.int)) \ + and isinstance(dsets[1], (int, np.int)), \ + "dsets must contain integer indices" dset1 = self.dsets[self.dset_idx(dsets[0])] dset2 = self.dsets[self.dset_idx(dsets[1])] # assert form of bls1 and bls2 - assert isinstance(bls1, list), "bls1 and bls2 must be fed as a list of antpair tuples" - assert isinstance(bls2, list), "bls1 and bls2 must be fed as a list of antpair tuples" - assert len(bls1) == len(bls2) and len(bls1) > 0, "length of bls1 must equal length of bls2 and be > 0" + assert isinstance(bls1, list), \ + "bls1 and bls2 must be fed as a list of antpair tuples" + assert isinstance(bls2, list), \ + "bls1 and bls2 must be fed as a list of antpair tuples" + 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) + assert isinstance(bls2[i], tuple), \ + "bls1[{}] type must match bls2[{}] type".format(i, i) else: - assert len(bls1[i]) == len(bls2[i]), "len(bls1[{}]) must match len(bls2[{}])".format(i, i) + assert len(bls1[i]) == len(bls2[i]), \ + "len(bls1[{}]) must match len(bls2[{}])".format(i, i) # construct list of baseline pairs bl_pairs = [] @@ -1942,7 +1985,8 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit elif isinstance(bls1[i], list) and len(bls1[i]) == 1: bl_pairs.append( (bls1[i][0], bls2[i][0]) ) else: - bl_pairs.append(map(lambda j: (bls1[i][j] , bls2[i][j]), range(len(bls1[i])))) + bl_pairs.append( + [ (bls1[i][j], bls2[i][j]) for j in range(len(bls1[i])) ] ) # validate bl-pair redundancy validate_blpairs(bl_pairs, dset1, dset2, baseline_tol=1.0) @@ -1951,24 +1995,29 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit if spw_ranges is None: spw_ranges = [(0, self.Nfreqs)] else: - assert np.isclose(map(lambda t: len(t), spw_ranges), 2).all(), "spw_ranges must be fed as a list of length-2 tuples" + assert np.isclose([len(t) for t in spw_ranges], 2).all(), \ + "spw_ranges must be fed as a list of length-2 tuples" - # if using default setting of number of delay bins equal to number of frequency channels + # if using default setting of number of delay bins equal to number + # of frequency channels if n_dlys is None: n_dlys = [None for i in range(len(spw_ranges))] elif isinstance(n_dlys, (int, np.integer)): n_dlys = [n_dlys] - # if using the whole band in the dataset, then there should just be one n_dly parameter specified + # if using the whole band in the dataset, then there should just be + # one n_dly parameter specified if spw_ranges is None and n_dlys != None: - assert len(n_dlys) == 1, "Only one spw, so cannot specify more than one n_dly value" + assert len(n_dlys) == 1, \ + "Only one spw, so cannot specify more than one n_dly value" - # assert that the same number of ndlys has been specified as the number of spws - assert len(spw_ranges) == len(n_dlys), "Need to specify number of delay bins for each spw" + # assert that the same number of ndlys has been specified as the + # number of spws + assert len(spw_ranges) == len(n_dlys), \ + "Need to specify number of delay bins for each spw" # setup polarization selection - if isinstance(pols, tuple): - pols = [pols] + if isinstance(pols, tuple): pols = [pols] # convert all polarizations to integers if fed as strings _pols = [] @@ -2010,12 +2059,12 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit # clear covariance cache self.clear_cache() - # setup emtpy data arrays + # setup empty data arrays spw_data = [] spw_wgts = [] spw_ints = [] spw_scalar = [] - spw_pol = [] + spw_polpair = [] spw_cov = [] d = self.delays() * 1e-9 @@ -2027,19 +2076,18 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit # Loop over polarizations for j, p in enumerate(pols): - p_str = tuple(map(lambda _p: uvutils.polnum2str(_p), p)) + p_str = tuple([uvutils.polnum2str(_p) for _p in p]) if verbose: print( "\nUsing polarization pair: {}".format(p_str)) # validating polarization pair on UVData objects valid = self.validate_pol(dsets, tuple(p)) if not valid: - # storing only one polarization as only equal polarization are allowed at the - # moment and UVPSpec object also understands one polarization - print ("Polarization pair: {} failed the validation test, continuing...".format(p_str)) + # Polarization pair is invalid; skip + print("Polarization pair: {} failed the validation test, " + "continuing...".format(p_str)) continue - - # UVPSpec only takes a single pol currently - spw_pol.append(p[0]) + + spw_polpair.append( uvputils.polpair_tuple2int(p) ) pol_data = [] pol_wgts = [] pol_ints = [] @@ -2047,13 +2095,22 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit # Compute scalar to convert "telescope units" to "cosmo units" if self.primary_beam is not None: - # using zero'th indexed poalrization as cross polarized beam are not yet implemented + + # Raise error if cross-pol is requested + if (p[0] != p[1]): + raise NotImplementedError("Cross-polarized beams are " + "not yet implemented") + + # using zero'th indexed polarization, as cross-polarized + # beams are not yet implemented if norm == 'H^-1': - # If using decorrelation, the H^-1 normalization already deals with the taper, - # so we need to override the taper when computing the scalar - scalar = self.scalar(p[0], little_h=True, taper_override='none') + # If using decorrelation, the H^-1 normalization + # already deals with the taper, so we need to override + # the taper when computing the scalar + scalar = self.scalar(p, little_h=True, + taper_override='none') else: - scalar = self.scalar(p[0], little_h=True) + scalar = self.scalar(p, little_h=True) else: raise_warning("Warning: self.primary_beam is not defined, " "so pspectra are not properly normalized", @@ -2087,21 +2144,24 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit key2_dof = np.sum(~np.isclose(self.Y(key2).diagonal(), 0.0)) if key1_dof < self.spw_Ndlys or key2_dof < self.spw_Ndlys: if verbose: - print("WARNING: Number of unflagged chans for key1 and/or key2 < n_dlys\n" \ - "which may lead to normalization instabilities.") + print("WARNING: Number of unflagged chans for key1 " + "and/or key2 < n_dlys\n which may lead to " + "normalization instabilities.") # Build Fisher matrix if input_data_weight == 'identity': # in this case, all Gv and Hv differ only by flagging pattern # so check if we've already computed this # First: get flag weighting matrices given key1 & key2 - Y = np.vstack([self.Y(key1).diagonal(), self.Y(key2).diagonal()]) + Y = np.vstack([self.Y(key1).diagonal(), + self.Y(key2).diagonal()]) # Second: check cache for Y - matches = [np.isclose(Y, y).all() for y in self._identity_Y.values()] + matches = [np.isclose(Y, y).all() + for y in self._identity_Y.values()] if True in matches: # This Y exists, so pick appropriate G and H and continue - match = self._identity_Y.keys()[matches.index(True)] + match = list(self._identity_Y.keys())[matches.index(True)] Gv = self._identity_G[match] Hv = self._identity_H[match] else: @@ -2138,8 +2198,8 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit if verbose: print(" Computing and multiplying scalar...") pv *= scalar - # Wide bin adjustment of scalar, which is only needed for the diagonal norm - # matrix mode (i.e., norm = 'I') + # Wide bin adjustment of scalar, which is only needed for + # the diagonal norm matrix mode (i.e., norm = 'I') if norm == 'I': pv *= self.scalar_delay_adjustment(Gv=Gv, Hv=Hv) @@ -2149,9 +2209,9 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit cov_qv = self.cov_q_hat(key1, key2) cov_pv = self.cov_p_hat(Mv, cov_qv) if self.primary_beam != None: - cov_pv *= \ - (scalar * self.scalar_delay_adjustment(key1, key2, - sampling=sampling))**2. + cov_pv *= (scalar * \ + self.scalar_delay_adjustment(key1, key2, + sampling=sampling))**2. pol_cov.extend(cov_pv) # Get baseline keys @@ -2172,7 +2232,7 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit wgts1 = self.w(key1).T wgts2 = self.w(key2).T - # get average of nsample across frequency axis, weighted by wgts + # get avg of nsample across frequency axis, weighted by wgts nsamp1 = np.sum(dset1.get_nsamples(bl1 + (p[0],))[:, self.spw_range[0]:self.spw_range[1]] * wgts1, axis=1) \ / np.sum(wgts1, axis=1).clip(1, np.inf) nsamp2 = np.sum(dset2.get_nsamples(bl2 + (p[1],))[:, self.spw_range[0]:self.spw_range[1]] * wgts2, axis=1) \ @@ -2186,13 +2246,14 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit blts2 = dset2.antpair2ind(bl2, ordered=False) integ2 = dset2.integration_time[blts2] * nsamp2 - # take inverse average of integ1 and integ2 to get total integration + # take inverse avg of integ1 and integ2 to get total integ # inverse avg is done b/c integ ~ 1/noise_var # and due to non-linear operation of V_1 * V_2 pol_ints.extend(1./np.mean([1./integ1, 1./integ2], axis=0)) # combined weight is geometric mean - pol_wgts.extend(np.concatenate([wgts1[:, :, None], wgts2[:, :, None]], axis=2)) + pol_wgts.extend(np.concatenate([wgts1[:, :, None], + wgts2[:, :, None]], axis=2)) # insert time and blpair info only once per blpair if i < 1 and j < 1: @@ -2205,7 +2266,8 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit lst2.extend(dset2.lst_array[inds2]) # insert blpair info - blp_arr.extend(np.ones_like(inds1, np.int) * uvputils._antnums_to_blpair(blp)) + blp_arr.extend(np.ones_like(inds1, np.int) \ + * uvputils._antnums_to_blpair(blp)) # insert into data and wgts integrations dictionaries spw_data.append(pol_data) @@ -2224,9 +2286,10 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit integration_array[i] = spw_ints sclr_arr.append(spw_scalar) - # raise error if none of pols are consistent witht the UVData objects - if len(spw_pol)==0: - raise ValueError("None of the specified polarization pair match that of the UVData objects") + # raise error if none of pols are consistent with the UVData objects + if len(spw_polpair) == 0: + raise ValueError("None of the specified polarization pairs " + "match that of the UVData objects") # fill uvp object uvp = uvpspec.UVPSpec() @@ -2237,15 +2300,17 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit uvp.time_avg_array = np.mean([uvp.time_1_array, uvp.time_2_array], axis=0) uvp.lst_1_array = np.array(lst1) uvp.lst_2_array = np.array(lst2) - uvp.lst_avg_array = np.mean([np.unwrap(uvp.lst_1_array), np.unwrap(uvp.lst_2_array)], axis=0) % (2*np.pi) + uvp.lst_avg_array = np.mean([np.unwrap(uvp.lst_1_array), + np.unwrap(uvp.lst_2_array)], axis=0) \ + % (2*np.pi) uvp.blpair_array = np.array(blp_arr) uvp.Nblpairs = len(np.unique(blp_arr)) uvp.Ntimes = len(np.unique(time1)) uvp.Nblpairts = len(time1) bls_arr = sorted(set(bls_arr)) - uvp.bl_array = np.array(map(lambda bl: uvp.antnums_to_bl(bl), bls_arr)) + uvp.bl_array = np.array([uvp.antnums_to_bl(bl) for bl in bls_arr]) antpos = dict(zip(dset1.antenna_numbers, dset1.antenna_positions)) - uvp.bl_vecs = np.array(map(lambda bl: antpos[bl[0]] - antpos[bl[1]], bls_arr)) + uvp.bl_vecs = np.array([antpos[bl[0]] - antpos[bl[1]] for bl in bls_arr]) uvp.Nbls = len(uvp.bl_array) uvp.spw_dly_array = np.array(dly_spws) uvp.spw_freq_array = np.array(freq_spws) @@ -2257,10 +2322,10 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit uvp.Nspwdlys = len(uvp.spw_dly_array) uvp.Nspwfreqs = len(uvp.spw_freq_array) uvp.Nfreqs = len(np.unique(freqs)) - uvp.pol_array = np.array(spw_pol, np.int) - uvp.Npols = len(spw_pol) + uvp.polpair_array = np.array(spw_polpair, np.int) + uvp.Npols = len(spw_polpair) uvp.scalar_array = np.array(sclr_arr) - uvp.channel_width = dset1.channel_width # all dsets are validated to agree + uvp.channel_width = dset1.channel_width # all dsets 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 @@ -2288,7 +2353,8 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit uvp.cosmo = self.primary_beam.cosmo # attach beam info uvp.beam_freqs = self.primary_beam.beam_freqs - uvp.OmegaP, uvp.OmegaPP = self.primary_beam.get_Omegas(uvp.pol_array) + uvp.OmegaP, uvp.OmegaPP = \ + self.primary_beam.get_Omegas(uvp.polpair_array) if hasattr(self.primary_beam, 'filename'): uvp.beamfile = self.primary_beam.filename @@ -2298,24 +2364,30 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, input_data_weight='identit uvp.cov_array = cov_array uvp.integration_array = integration_array uvp.wgt_array = wgt_array - uvp.nsample_array = dict(map(lambda k: (k, np.ones_like(uvp.integration_array[k], np.float)), uvp.integration_array.keys())) + uvp.nsample_array = dict( + [ (k, np.ones_like(uvp.integration_array[k], np.float)) + for k in uvp.integration_array.keys() ] ) # run check uvp.check() return uvp + def rephase_to_dset(self, dset_index=0, inplace=True): """ - Rephase visibility data in self.dsets to the LST grid of dset[dset_index] - using hera_cal.utils.lst_rephase. Each integration in all other dsets is - phased to the center of the corresponding LST bin (by index) in dset[dset_index]. + Rephase visibility data in self.dsets to the LST grid of + dset[dset_index] using hera_cal.utils.lst_rephase. + + Each integration in all other dsets is phased to the center of the + corresponding LST bin (by index) in dset[dset_index]. - Will only phase if the dataset's phase type is 'drift'. This is because the rephasing - algorithm assumes the data is drift-phased when applying phasor term. + Will only phase if the dataset's phase type is 'drift'. This is because + the rephasing algorithm assumes the data is drift-phased when applying + the phasor term. - Note that PSpecData.Jy_to_mK() must be run after rephase_to_dset(), if one intends - to use the former capability at any point. + Note that PSpecData.Jy_to_mK() must be run *after* rephase_to_dset(), + if one intends to use the former capability at any point. Parameters ---------- @@ -2375,7 +2447,7 @@ def rephase_to_dset(self, dset_index=0, inplace=True): pols) = hc.io.load_vis(dset, return_meta=True) # make bls dictionary - bls = dict(map(lambda k: (k, antpos[k[0]] - antpos[k[1]]), data.keys())) + bls = dict([(k, antpos[k[0]] - antpos[k[1]]) for k in data.keys()]) # Get dlst array dlst = lst_grid - lsts @@ -2390,8 +2462,10 @@ def rephase_to_dset(self, dset_index=0, inplace=True): for j, k in enumerate(data.keys()): # get blts indices of basline indices = dset.antpair2ind(k[:2], ordered=False) + # get index in polarization_array for this polarization polind = pol_list.index(uvutils.polstr2num(k[-1])) + # insert into dset dset.data_array[indices, 0, :, polind] = data[k] @@ -2404,13 +2478,14 @@ def rephase_to_dset(self, dset_index=0, inplace=True): def Jy_to_mK(self, beam=None): """ - Convert internal datasets from a Jy-scale to mK scale using a primary beam - model if available. Note that if you intend to rephase_to_dset(), Jy to mK conversion - must be done after that step. + Convert internal datasets from a Jy-scale to mK scale using a primary + beam model if available. Note that if you intend to rephase_to_dset(), + Jy to mK conversion must be done *after* that step. Parameters ---------- beam : PSpecBeam object + Beam object. """ # get all unique polarizations of all the datasets pols = set(np.ravel([dset.polarization_array for dset in self.dsets])) @@ -2420,13 +2495,16 @@ def Jy_to_mK(self, beam=None): beam = self.primary_beam else: if self.primary_beam is not None: - print("Warning: feeding a beam model when self.primary_beam already exists...") + print("Warning: feeding a beam model when self.primary_beam " + "already exists...") # Check beam is not None - assert beam is not None, "Cannot convert Jy --> mK b/c beam object is not defined..." + assert beam is not None, \ + "Cannot convert Jy --> mK b/c beam object is not defined..." # assert type of beam - assert isinstance(beam, pspecbeam.PSpecBeamBase), "beam model must be a subclass of pspecbeam.PSpecBeamBase" + assert isinstance(beam, pspecbeam.PSpecBeamBase), \ + "beam model must be a subclass of pspecbeam.PSpecBeamBase" # iterate over all pols and get conversion factors factors = {} @@ -2445,9 +2523,10 @@ def Jy_to_mK(self, beam=None): def trim_dset_lsts(self, lst_tol=6): """ - Assuming all datasets in self.dsets are locked to the same LST grid (but - each may have a constant offset), trim LSTs from each dset that aren't found - in all other dsets (within some decimal tolerance specified by lst_tol). + Assuming all datasets in self.dsets are locked to the same LST grid + (but each may have a constant offset), trim LSTs from each dset that + aren't found in all other dsets (within some decimal tolerance + specified by lst_tol). Warning: this edits the data in dsets in-place, and is not reversible. @@ -2461,34 +2540,41 @@ def trim_dset_lsts(self, lst_tol=6): for dset in self.dsets: _dlst = np.median(np.diff(np.unique(dset.lst_array))) if not np.isclose(dlst, _dlst, atol=10**(-lst_tol) / dset.Ntimes): - print("not all datasets in self.dsets are on the same LST grid, cannot LST trim.") + print("not all datasets in self.dsets are on the same LST " + "grid, cannot LST trim.") return - # get lst array of each dataset and turn into string and add to common_lsts + # get lst array of each dataset, turn into string and add to common_lsts lst_arrs = [] common_lsts = set() for i, dset in enumerate(self.dsets): - lsts = ["{lst:0.{tol}f}".format(lst=l, tol=lst_tol) for l in dset.lst_array] + lsts = ["{lst:0.{tol}f}".format(lst=l, tol=lst_tol) + for l in dset.lst_array] lst_arrs.append(lsts) if i == 0: common_lsts = common_lsts.union(set(lsts)) else: common_lsts = common_lsts.intersection(set(lsts)) - # iterate through dsets and trim off integrations whose lst isn't in common_lsts + # iterate through dsets and trim off integrations whose lst isn't + # in common_lsts for i, dset in enumerate(self.dsets): trim_inds = np.array([l not in common_lsts for l in lst_arrs[i]]) if np.any(trim_inds): self.dsets[i].select(times=dset.time_array[~trim_inds]) -def pspec_run(dsets, filename, dsets_std=None, groupname=None, dset_labels=None, dset_pairs=None, - psname_ext=None, spw_ranges=None, n_dlys=None, pol_pairs=None, blpairs=None, +def pspec_run(dsets, filename, dsets_std=None, 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, - Nblps_per_group=None, bl_len_range=(0, 1e10), bl_deg_range=(0, 180), bl_error_tol=1.0, - beam=None, cosmo=None, rephase_to_dset=None, trim_dset_lsts=False, broadcast_dset_flags=True, - time_thresh=0.2, Jy2mK=False, overwrite=True, verbose=True, store_cov=False, history=''): + Nblps_per_group=None, bl_len_range=(0, 1e10), + bl_deg_range=(0, 180), bl_error_tol=1.0, + beam=None, cosmo=None, rephase_to_dset=None, + trim_dset_lsts=False, broadcast_dset_flags=True, + time_thresh=0.2, Jy2mK=False, overwrite=True, + verbose=True, store_cov=False, history=''): """ Create a PSpecData object, run OQE delay spectrum estimation and write results to a PSpecContainer object. @@ -2582,7 +2668,8 @@ def pspec_run(dsets, filename, dsets_std=None, groupname=None, dset_labels=None, bl_deg_range : len-2 float tuple A tuple containing the min and max baseline angle (ENU frame in degrees) - to use in utils.calc_blpair_reds. Total range is between 0 and 180 degrees. + to use in utils.calc_blpair_reds. Total range is between 0 and 180 + degrees. bl_error_tol : float Baseline vector error tolerance when constructing redundant groups. @@ -2635,12 +2722,12 @@ def pspec_run(dsets, filename, dsets_std=None, groupname=None, dset_labels=None, Returns ------- psc : PSpecContainer object - A container for the output UVPSpec objects, which themselves contain the - power spectra and their metadata. + A container for the output UVPSpec objects, which themselves contain + the power spectra and their metadata. ds : PSpecData object - The PSpecData object used for OQE of power spectrum, with cached weighting - matrices. + The PSpecData object used for OQE of power spectrum, with cached + weighting matrices. """ # type check err_msg = "dsets must be fed as a list of dataset string paths or UVData objects." @@ -2678,7 +2765,8 @@ def pspec_run(dsets, filename, dsets_std=None, groupname=None, dset_labels=None, if dset_labels is None: dset_labels = ["dset{}".format(i) for i in range(Ndsets)] else: - assert not np.any(['_' in dl for dl in dset_labels]), "cannot accept underscores in input dset_labels: {}".format(dset_labels) + assert not np.any(['_' in dl for dl in dset_labels]), \ + "cannot accept underscores in input dset_labels: {}".format(dset_labels) # load data if fed as filepaths if isinstance(dsets[0], (str, np.str)): @@ -2698,7 +2786,8 @@ def pspec_run(dsets, filename, dsets_std=None, groupname=None, dset_labels=None, # check dsets_std input if dsets_std is not None: - err_msg = "input dsets_std must be a list of UVData objects or filepaths to miriad files" + err_msg = "input dsets_std must be a list of UVData objects or " \ + "filepaths to miriad files" assert isinstance(dsets_std,(list, tuple, np.ndarray)), err_msg assert len(dsets_std) == Ndsets, "len(dsets_std) must equal len(dsets)" @@ -2707,11 +2796,13 @@ def pspec_run(dsets, filename, dsets_std=None, groupname=None, dset_labels=None, try: # load data into UVData objects if fed as list of strings t0 = time.time() - dsets_std = _load_dsets(dsets_std, bls=bls, pols=pols, verbose=verbose) + dsets_std = _load_dsets(dsets_std, bls=bls, pols=pols, + verbose=verbose) utils.log("Loaded data in %1.1f sec." % (time.time() - t0), lvl=1, verbose=verbose) except ValueError: - # at least one of the dsets_std loads failed due to no data being present + # at least one of the dsets_std loads failed due to no data + # being present utils.log("One of the dsets_std loads failed due to no data overlap given the bls and pols selection", verbose=verbose) return None, None @@ -2719,7 +2810,7 @@ def pspec_run(dsets, filename, dsets_std=None, groupname=None, dset_labels=None, # configure polarization if pol_pairs is None: - unique_pols = reduce(operator.and_, [set(d.polarization_array) for d in dsets]) + unique_pols = np.unique(np.hstack([d.polarization_array for d in dsets])) pol_pairs = [(up, up) for up in unique_pols] assert len(pol_pairs) > 0, "no pol_pairs specified" @@ -2734,7 +2825,8 @@ def pspec_run(dsets, filename, dsets_std=None, groupname=None, dset_labels=None, beam.cosmo = cosmo # package into PSpecData - ds = PSpecData(dsets=dsets, wgts=[None for d in dsets], labels=dset_labels, dsets_std=dsets_std, beam=beam) + ds = PSpecData(dsets=dsets, wgts=[None for d in dsets], labels=dset_labels, + dsets_std=dsets_std, beam=beam) # Rephase if desired if rephase_to_dset is not None: @@ -2776,13 +2868,14 @@ def pspec_run(dsets, filename, dsets_std=None, groupname=None, dset_labels=None, # get bls if blpairs not fed if blpairs is None: (bls1, bls2, blps, xants1, - xants2) = utils.calc_blpair_reds(dsets[dsetp[0]], dsets[dsetp[1]], - filter_blpairs=True, - exclude_auto_bls=exclude_auto_bls, - exclude_permutations=exclude_permutations, - Nblps_per_group=Nblps_per_group, - bl_len_range=bl_len_range, - bl_deg_range=bl_deg_range) + xants2) = utils.calc_blpair_reds( + dsets[dsetp[0]], dsets[dsetp[1]], + filter_blpairs=True, + exclude_auto_bls=exclude_auto_bls, + exclude_permutations=exclude_permutations, + Nblps_per_group=Nblps_per_group, + bl_len_range=bl_len_range, + bl_deg_range=bl_deg_range ) bls1_list.append(bls1) bls2_list.append(bls2) @@ -2816,13 +2909,14 @@ def pspec_run(dsets, filename, dsets_std=None, groupname=None, dset_labels=None, # Run OQE uvp = ds.pspec(bls1_list[i], bls2_list[i], dset_idxs, pol_pairs, - spw_ranges=spw_ranges, n_dlys=n_dlys, store_cov=store_cov, - input_data_weight=input_data_weight, norm=norm, taper=taper, - history=history, verbose=verbose) + spw_ranges=spw_ranges, n_dlys=n_dlys, + store_cov=store_cov, input_data_weight=input_data_weight, + norm=norm, taper=taper, history=history, verbose=verbose) # Store output psname = '{}_x_{}{}'.format(dset_labels[dset_idxs[0]], dset_labels[dset_idxs[1]], psname_ext) + psc.set_pspec(group=groupname, psname=psname, pspec=uvp, overwrite=overwrite) @@ -2833,16 +2927,16 @@ def get_pspec_run_argparser(): a = argparse.ArgumentParser(description="argument parser for pspecdata.pspec_run()") def list_of_int_tuples(v): - v = map(lambda x: tuple(map(int, x.split())), v.split(",")) + v = [tuple([int(_x) for _x in x.split()]) for x in v.split(",")] return v def list_of_str_tuples(v): - v = map(lambda x: tuple(map(str, x.split())), v.split(",")) + v = [tuple([str(_x) for _x in x.split()]) for x in v.split(",")] return v def list_of_tuple_tuples(v): - v = map(lambda x: tuple(map(int, x.split())), v.split(",")) - v = map(lambda x: (x[:2], x[2:]), v) + v = [tuple([int(_x) for _x in x.split()]) for x in v.split(",")] + v = [(x[:2], x[2:]) for x in v] return v a.add_argument("dsets", nargs='*', help="List of UVData objects or miriad filepaths.") @@ -2885,16 +2979,20 @@ def validate_blpairs(blpairs, uvd1, uvd2, baseline_tol=1.0, verbose=True): Parameters ---------- - blpairs : list of baseline-pair tuples, Ex. [((1,2),(1,2)), ((2,3),(2,3))] + blpairs : list of baseline-pair tuples + Ex. [((1,2),(1,2)), ((2,3),(2,3))] See docstring of PSpecData.pspec() for details on format. - uvd1 : UVData instance containing visibility data that first bl in blpair will draw from - - uvd2 : UVData instance containing visibility data that second bl in blpair will draw from - - baseline_tol : float, distance tolerance for notion of baseline "redundancy" in meters + uvd1, uvd2 : UVData + UVData instances containing visibility data that first/second bl in + blpair will draw from + + baseline_tol : float, optional + Distance tolerance for notion of baseline "redundancy" in meters. + Default: 1.0. - verbose : bool, if True report feedback to stdout + verbose : bool, optional + If True report feedback to stdout. Default: True. """ # ensure uvd1 and uvd2 are UVData objects if isinstance(uvd1, UVData) == False: @@ -2911,11 +3009,13 @@ def validate_blpairs(blpairs, uvd1, uvd2, baseline_tol=1.0, verbose=True): # ensure shared antenna keys match within tolerance shared = sorted(set(ap1.keys()) & set(ap2.keys())) for k in shared: - assert np.linalg.norm(ap1[k] - ap2[k]) <= baseline_tol, "uvd1 and uvd2 don't agree on antenna positions within tolerance of {} m".format(baseline_tol) + assert np.linalg.norm(ap1[k] - ap2[k]) <= baseline_tol, \ + "uvd1 and uvd2 don't agree on antenna positions within tolerance of {} m".format(baseline_tol) ap = ap1 ap.update(ap2) - # iterate through baselines and check baselines crossed with each other are within tolerance + # iterate through baselines and check baselines crossed with each other + # are within tolerance for i, blg in enumerate(blpairs): if isinstance(blg, tuple): blg = [blg] @@ -2927,17 +3027,23 @@ def validate_blpairs(blpairs, uvd1, uvd2, baseline_tol=1.0, verbose=True): def raise_warning(warning, verbose=True): - '''warning function''' + """ + Warning function. + """ if verbose: print(warning) def _load_dsets(fnames, bls=None, pols=None, logf=None, verbose=True): - """ helper function for loading Miriad datasets in pspec_run """ + """ + Helper function for loading Miriad datasets in pspec_run. + """ dsets = [] Ndsets = len(fnames) for i, dset in enumerate(fnames): - utils.log("Reading {} / {} datasets...".format(i+1, Ndsets), f=logf, lvl=1, verbose=verbose) + utils.log("Reading {} / {} datasets...".format(i+1, Ndsets), + f=logf, lvl=1, verbose=verbose) + # read data uvd = UVData() uvd.read_miriad(glob.glob(dset), bls=bls, polarizations=pols) diff --git a/hera_pspec/pstokes.py b/hera_pspec/pstokes.py index 9e5903b2..bc32297b 100644 --- a/hera_pspec/pstokes.py +++ b/hera_pspec/pstokes.py @@ -8,8 +8,8 @@ from collections import OrderedDict as odict -# weights used in forming Stokes visibilities. -# see pyuvdata.utils.polstr2num for conversion between polarization string +# Weights used in forming Stokes visibilities. +# See pyuvdata.utils.polstr2num for conversion between polarization string # and polarization integer. Ex. {'XX': -5, ...} pol_weights = { 1: odict([(-5, 1.), (-6, 1.)]), @@ -26,7 +26,8 @@ def miriad2pyuvdata(dset, antenna_nums=None, bls=None, polarizations=None, Parameters ---------- dset : str - Miriad file to convert to UVData object containing visibilities and corresponding metadata + Miriad file to convert to UVData object containing visibilities and + corresponding metadata antenna_nums: integer list The antennas numbers to read into the object. @@ -56,7 +57,8 @@ def miriad2pyuvdata(dset, antenna_nums=None, bls=None, polarizations=None, """ uvd = pyuvdata.UVData() uvd.read_miriad(dset, antenna_nums=antenna_nums, bls=bls, - polarizations=polarizations, ant_str=ant_str, time_range=time_range) + polarizations=polarizations, ant_str=ant_str, + time_range=time_range) return uvd @@ -91,8 +93,10 @@ def _combine_pol(uvd1, uvd2, pol1, pol2, pstokes='pI'): ------- uvdS : UVData object """ - assert isinstance(uvd1, pyuvdata.UVData), "uvd1 must be a pyuvdata.UVData instance" - assert isinstance(uvd2, pyuvdata.UVData), "uvd2 must be a pyuvdata.UVData instance" + assert isinstance(uvd1, pyuvdata.UVData), \ + "uvd1 must be a pyuvdata.UVData instance" + assert isinstance(uvd2, pyuvdata.UVData), \ + "uvd2 must be a pyuvdata.UVData instance" # convert pol1 and/or pol2 to integer if fed as a string if isinstance(pol1, (str, np.str)): @@ -121,9 +125,12 @@ def _combine_pol(uvd1, uvd2, pol1, pol2, pstokes='pI'): pstokes_str = pyuvdata.utils.polnum2str(pstokes) # assert pstokes in pol_weights, and pol1 and pol2 in pol_weights[pstokes] - assert pstokes in pol_weights, "unrecognized pstokes parameter {}".format(pstokes_str) - assert pol1 in pol_weights[pstokes], "pol1 {} not used in constructing pstokes {}".format(pol1_str, pstokes_str) - assert pol2 in pol_weights[pstokes], "pol2 {} not used in constructing pstokes {}".format(pol2_str, pstokes_str) + assert pstokes in pol_weights, \ + "unrecognized pstokes parameter {}".format(pstokes_str) + assert pol1 in pol_weights[pstokes], \ + "pol1 {} not used in constructing pstokes {}".format(pol1_str, pstokes_str) + assert pol2 in pol_weights[pstokes], \ + "pol2 {} not used in constructing pstokes {}".format(pol2_str, pstokes_str) # constructing Stokes visibilities stdata = 0.5 * (pol_weights[pstokes][pol1]*data1 + pol_weights[pstokes][pol2]*data2) @@ -132,8 +139,8 @@ def _combine_pol(uvd1, uvd2, pol1, pol2, pstokes='pI'): uvdS = copy.deepcopy(uvd1) uvdS.data_array = stdata # pseudo-stokes data uvdS.flag_array = flag # flag array - uvdS.polarization_array = np.array([pstokes], dtype=np.int) # polarization number - uvdS.nsample_array = uvd1.nsample_array + uvd2.nsample_array # nsamples + uvdS.polarization_array = np.array([pstokes], dtype=np.int) # polarization number + uvdS.nsample_array = uvd1.nsample_array + uvd2.nsample_array # nsamples uvdS.history = "Merged into pseudo-stokes vis with hera_pspec version {} Git hash {}\n{}" \ "{}{}{}{}\n".format(version.version, version.git_hash, "-"*20+'\n', 'dset1 history:\n', uvd1.history, '\n'+'-'*20+'\ndset2 history:\n', @@ -195,8 +202,8 @@ def construct_pstokes(dset1, dset2, pstokes='pI', run_check=True, antenna_nums=N Ex: ['xx', 'yy', ...] time_range: float list - len-2 list containing min and max range of times (Julian Date) to read-in. - Ex: [2458115.20, 2458115.40] + len-2 list containing min and max range of times (Julian Date) to + read-in. Ex: [2458115.20, 2458115.40] history : str Extra history string to add to concatenated pseudo-Stokes visibility. @@ -207,15 +214,19 @@ def construct_pstokes(dset1, dset2, pstokes='pI', run_check=True, antenna_nums=N """ # convert dset1 and dset2 to UVData objects if they are miriad files if isinstance(dset1, pyuvdata.UVData) == False: - assert isinstance(dset1, (str, np.str)), "dset1 must be fed as a string or UVData object" + assert isinstance(dset1, (str, np.str)), \ + "dset1 must be fed as a string or UVData object" uvd1 = miriad2pyuvdata(dset1, antenna_nums=antenna_nums, bls=bls, - polarizations=polarizations, ant_str=ant_str, time_range=time_range) + polarizations=polarizations, ant_str=ant_str, + time_range=time_range) else: uvd1 = dset1 if isinstance(dset2, pyuvdata.UVData) == False: - assert isinstance(dset2, (str, np.str)), "dset2 must be fed as a string or UVData object" + assert isinstance(dset2, (str, np.str)), \ + "dset2 must be fed as a string or UVData object" uvd2 = miriad2pyuvdata(dset2, antenna_nums=antenna_nums, bls=bls, - polarizations=polarizations, ant_str=ant_str, time_range=time_range) + polarizations=polarizations, ant_str=ant_str, + time_range=time_range) else: uvd2 = dset2 @@ -246,23 +257,28 @@ def construct_pstokes(dset1, dset2, pstokes='pI', run_check=True, antenna_nums=N if np.array_equal(bls1, bls2) == False: raise ValueError("dset1 and dset2 must have the same baselines") - # makes the Npol length==1 so that the UVData carries data for the required polarization only - st_keys = pol_weights[pstokes].keys() + # makes the Npol length==1 so that the UVData carries data for the + # required polarization only + st_keys = list(pol_weights[pstokes].keys()) req_pol1 = st_keys[0] req_pol2 = st_keys[1] - # check polarizations of UVData objects are consistent with the required polarization - # to form the desired pseudo Stokes visibilities. If multiple exist, downselect on polarization. - assert req_pol1 in uvd1.polarization_array, "Polarization {} not found in dset1 object".format(req_pol1) + # check polarizations of UVData objects are consistent with the required + # polarization to form the desired pseudo Stokes visibilities. If multiple + # exist, downselect on polarization. + assert req_pol1 in uvd1.polarization_array, \ + "Polarization {} not found in dset1 object".format(req_pol1) if uvd1.Npols > 1: uvd1 = uvd1.select(polarizations=req_pol1, inplace=False) - assert req_pol2 in uvd2.polarization_array, "Polarization {} not found in dset2 object".format(req_pol2) + assert req_pol2 in uvd2.polarization_array, \ + "Polarization {} not found in dset2 object".format(req_pol2) if uvd2.Npols > 1: uvd2 = uvd2.select(polarizations=req_pol1, inplace=False) # combining visibilities to form the desired Stokes visibilties - uvdS = _combine_pol(uvd1=uvd1, uvd2=uvd2, pol1=req_pol1, pol2=req_pol2, pstokes=pstokes) + uvdS = _combine_pol(uvd1=uvd1, uvd2=uvd2, pol1=req_pol1, pol2=req_pol2, + pstokes=pstokes) uvdS.history += history if run_check: @@ -295,8 +311,10 @@ def filter_dset_on_stokes_pol(dsets, pstokes): to construct_pstokes to make the desired pseudo-Stokes visibility. """ # type check - assert isinstance(dsets, list), "dsets must be fed as a list of UVData objects" - assert np.all(isinstance(d, UVData) for d in dsets), "dsets must be fed as a list of UVData objects" + assert isinstance(dsets, list), \ + "dsets must be fed as a list of UVData objects" + assert np.all(isinstance(d, UVData) for d in dsets), \ + "dsets must be fed as a list of UVData objects" # get polarization of each dset pols = [d.polarization_array[0] for d in dsets] @@ -304,12 +322,15 @@ def filter_dset_on_stokes_pol(dsets, pstokes): # convert pstokes to integer if a string if isinstance(pstokes, (str, np.str)): pstokes = pyuvdata.utils.polstr2num(pstokes) - assert pstokes in [1, 2, 3, 4], "pstokes must be fed as a pseudo-Stokes parameter" + assert pstokes in [1, 2, 3, 4], \ + "pstokes must be fed as a pseudo-Stokes parameter" # get two necessary dipole pols given pstokes - desired_pols = pol_weights[pstokes].keys() - assert desired_pols[0] in pols and desired_pols[1] in pols, "necessary input pols {} and {} not found in dsets".format(*desired_pols) + desired_pols = list(pol_weights[pstokes].keys()) + assert desired_pols[0] in pols and desired_pols[1] in pols, \ + "necessary input pols {} and {} not found in dsets".format(*desired_pols) - inp_dsets = [dsets[pols.index(desired_pols[0])], dsets[pols.index(desired_pols[1])]] + inp_dsets = [dsets[pols.index(desired_pols[0])], + dsets[pols.index(desired_pols[1])]] return inp_dsets diff --git a/hera_pspec/testing.py b/hera_pspec/testing.py index 5224e73e..49b1d546 100644 --- a/hera_pspec/testing.py +++ b/hera_pspec/testing.py @@ -59,8 +59,8 @@ def build_vanilla_uvpspec(beam=None): freq_array = np.repeat(np.linspace(100e6, 105e6, Nfreqs, endpoint=False), Nspws) dly_array = np.repeat(utils.get_delays(freq_array, n_dlys=Ndlys), Nspws) - pol_array = np.array([-5]) - Npols = len(pol_array) + polpair_array = np.array([1515,]) # corresponds to ('xx','xx') + Npols = len(polpair_array) vis_units = 'unknown' norm_units = 'Hz str' weighting = 'identity' @@ -76,7 +76,8 @@ def build_vanilla_uvpspec(beam=None): 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]) + pol = beam.primary_beam.polarization_array[0] + OmegaP, OmegaPP = beam.get_Omegas((pol,pol)) beam_freqs = beam.beam_freqs # HERA coordinates in Karoo Desert, SA @@ -105,7 +106,8 @@ def build_vanilla_uvpspec(beam=None): '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', + 'dly_array', 'freq_array', 'polpair_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', @@ -220,6 +222,7 @@ def uvpspec_from_data(data, bl_grps, data_std=None, spw_ranges=None, uvp = ds.pspec(bls1, bls2, (0, 1), (pol, pol), input_data_weight='identity', spw_ranges=spw_ranges, taper=taper, verbose=verbose, store_cov=store_cov, n_dlys=n_dlys) + return uvp diff --git a/hera_pspec/tests/test_container.py b/hera_pspec/tests/test_container.py index dda9974f..44df8f9f 100644 --- a/hera_pspec/tests/test_container.py +++ b/hera_pspec/tests/test_container.py @@ -53,7 +53,7 @@ def test_PSpecContainer(self): psname=psname, pspec=self.uvp, overwrite=False) # Check that wrong pspec types are rejected by the set() method - assert_raises(ValueError, ps_store.set_pspec, group=group_names[2], + assert_raises(TypeError, ps_store.set_pspec, group=group_names[2], psname=psname, pspec=np.arange(11), overwrite=True) assert_raises(TypeError, ps_store.set_pspec, group=group_names[2], psname=psname, pspec=1, overwrite=True) diff --git a/hera_pspec/tests/test_conversions.py b/hera_pspec/tests/test_conversions.py index 47db136a..58ac345a 100644 --- a/hera_pspec/tests/test_conversions.py +++ b/hera_pspec/tests/test_conversions.py @@ -10,7 +10,8 @@ class Test_Cosmo(unittest.TestCase): def setUp(self): - self.C = conversions.Cosmo_Conversions(Om_L=0.68440, Om_b=0.04911, Om_c=0.26442, H0=100.0, + self.C = conversions.Cosmo_Conversions(Om_L=0.68440, Om_b=0.04911, + Om_c=0.26442, H0=100.0, Om_M=None, Om_k=None) def tearDown(self): pass @@ -40,9 +41,12 @@ def test_units(self): def test_distances(self): self.assertAlmostEqual(self.C.f2z(100e6), 13.20405751) + self.assertAlmostEqual(self.C.f2z(0.1, ghz=True), 13.20405751) self.assertAlmostEqual(self.C.z2f(10.0), 129127795.54545455) + self.assertAlmostEqual(self.C.z2f(10.0, ghz=True), 0.12912779554545455) self.assertAlmostEqual(self.C.E(10.0), 20.450997530682947) self.assertAlmostEqual(self.C.DC(10.0), 6499.708111027144) + self.assertAlmostEqual(self.C.DC(10.0, little_h=False), 6499.708111027144) self.assertAlmostEqual(self.C.DM(10.0), 6510.2536925709637) self.assertAlmostEqual(self.C.DA(10.0), 591.84124477917851) self.assertAlmostEqual(self.C.dRperp_dtheta(10.0), 6510.2536925709637) @@ -54,18 +58,25 @@ def test_distances(self): def test_little_h(self): # Test that putting in a really low value of H0 into the conversions object # has no effect on the result if little_h=True units are used - self.C = conversions.Cosmo_Conversions(Om_L=0.68440, Om_b=0.04911, Om_c=0.26442, H0=25.12, + self.C = conversions.Cosmo_Conversions(Om_L=0.68440, Om_b=0.04911, + Om_c=0.26442, H0=25.12, Om_M=None, Om_k=None) self.assertAlmostEqual(self.C.f2z(100e6), 13.20405751) self.assertAlmostEqual(self.C.z2f(10.0), 129127795.54545455) self.assertAlmostEqual(self.C.E(10.0), 20.450997530682947) self.assertAlmostEqual(self.C.DC(10.0, little_h=True), 6499.708111027144) + self.assertAlmostEqual(self.C.DC(10.0, little_h=False), + 6499.708111027144*100./25.12) self.assertAlmostEqual(self.C.DM(10.0, little_h=True), 6510.2536925709637) self.assertAlmostEqual(self.C.DA(10.0, little_h=True), 591.84124477917851) - self.assertAlmostEqual(self.C.dRperp_dtheta(10.0, little_h=True), 6510.2536925709637) - self.assertAlmostEqual(self.C.dRpara_df(10.0, little_h=True), 1.2487605057418872e-05) - self.assertAlmostEqual(self.C.dRpara_df(10.0, ghz=True, little_h=True), 12487.605057418872) - self.assertAlmostEqual(self.C.X2Y(10.0, little_h=True), 529.26719942209002) + self.assertAlmostEqual(self.C.dRperp_dtheta(10.0, little_h=True), + 6510.2536925709637) + self.assertAlmostEqual(self.C.dRpara_df(10.0, little_h=True), + 1.2487605057418872e-05) + self.assertAlmostEqual(self.C.dRpara_df(10.0, ghz=True, little_h=True), + 12487.605057418872) + self.assertAlmostEqual(self.C.X2Y(10.0, little_h=True), + 529.26719942209002) def test_params(self): params = self.C.get_params() diff --git a/hera_pspec/tests/test_grouping.py b/hera_pspec/tests/test_grouping.py index cd91c556..72f4b46d 100644 --- a/hera_pspec/tests/test_grouping.py +++ b/hera_pspec/tests/test_grouping.py @@ -149,8 +149,12 @@ def test_bootstrap_average_blpairs(self): [uvp3,], blpair_groups=[_blpairs,], time_avg=True) - ps_avg = uvp_avg.get_data((0, blpair, 'xx')) - ps_boot = uvp4[0].get_data((0, blpair, 'xx')) + try: + ps_avg = uvp_avg.get_data((0, blpair, ('xx','xx'))) + except: + print(uvp_avg.polpair_array) + raise + ps_boot = uvp4[0].get_data((0, blpair, ('xx','xx'))) np.testing.assert_array_almost_equal(ps_avg, ps_boot) def test_bootstrap_resampled_error(): @@ -258,11 +262,11 @@ def test_bootstrap_run(): # assert original uvp is unchanged nt.assert_true(uvp == psc.get_pspec("grp1", 'uvp')) # check stats array - np.testing.assert_array_equal([u'bs_cinterval_16.00', u'bs_cinterval_84.00', u'bs_robust_std', u'bs_std'], uvp_avg.stats_array.keys()) + np.testing.assert_array_equal([u'bs_cinterval_16.00', u'bs_cinterval_84.00', u'bs_robust_std', u'bs_std'], list(uvp_avg.stats_array.keys())) for stat in [u'bs_cinterval_16.00', u'bs_cinterval_84.00', u'bs_robust_std', u'bs_std']: - nt.assert_equal(uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), 'XX')).shape, (1, 50)) - nt.assert_false(np.any(np.isnan(uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), 'XX'))))) - nt.assert_equal(uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), 'XX')).dtype, np.complex128) + nt.assert_equal(uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), ('xx','xx'))).shape, (1, 50)) + nt.assert_false(np.any(np.isnan(uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), ('xx','xx')))))) + nt.assert_equal(uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), ('xx','xx'))).dtype, np.complex128) # test exceptions del psc diff --git a/hera_pspec/tests/test_noise.py b/hera_pspec/tests/test_noise.py index 9408f5d5..c78edbf1 100644 --- a/hera_pspec/tests/test_noise.py +++ b/hera_pspec/tests/test_noise.py @@ -12,11 +12,14 @@ class Test_Sensitivity(unittest.TestCase): - """ Test noise.Sensitivity object """ + """ + Test noise.Sensitivity object + """ def setUp(self): self.cosmo = conversions.Cosmo_Conversions() - self.beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, 'HERA_NF_pstokes_power.beamfits')) + self.beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, + 'HERA_NF_pstokes_power.beamfits')) self.sense = noise.Sensitivity(beam=self.beam, cosmo=self.cosmo) def tearDown(self): @@ -40,6 +43,10 @@ def test_set(self): self.beam.cosmo = C sense.set_beam(self.beam) nt.assert_equal(sense.cosmo.get_params(), sense.beam.cosmo.get_params()) + + bm = copy.deepcopy(self.beam) + delattr(bm, 'cosmo') + sense.set_beam(bm) def test_scalar(self): freqs = np.linspace(150e6, 160e6, 100, endpoint=False) @@ -48,18 +55,22 @@ def test_scalar(self): nt.assert_true(self.sense.pol, 'pI') def test_calc_P_N(self): + # calculate scalar freqs = np.linspace(150e6, 160e6, 100, endpoint=False) self.sense.calc_scalar(freqs, 'pI', num_steps=5000, little_h=True) + # basic execution k = np.linspace(0, 3, 10) Tsys = 500.0 t_int = 10.7 - P_N = self.sense.calc_P_N(Tsys, t_int, Ncoherent=1, Nincoherent=1, form='Pk') + P_N = self.sense.calc_P_N(Tsys, t_int, Ncoherent=1, Nincoherent=1, + form='Pk') nt.assert_true(isinstance(P_N, (float, np.float))) nt.assert_true(np.isclose(P_N, 908472312787.53491)) # calculate DelSq - Dsq = self.sense.calc_P_N(Tsys, t_int, k=k, Ncoherent=1, Nincoherent=1, form='DelSq') + Dsq = self.sense.calc_P_N(Tsys, t_int, k=k, Ncoherent=1, + Nincoherent=1, form='DelSq') nt.assert_equal(Dsq.shape, (10,)) nt.assert_true(Dsq[1] < P_N) @@ -78,15 +89,19 @@ def test_noise_validation(): # generate noise seed = 0 - uvd = testing.noise_sim(uvfile, Tsys, beam, seed=seed, whiten=True, inplace=False, Nextend=9) + uvd = testing.noise_sim(uvfile, Tsys, beam, seed=seed, whiten=True, + inplace=False, Nextend=9) # get redundant baseline group - reds, lens, angs = utils.get_reds(uvd, pick_data_ants=True, bl_len_range=(10, 20), + reds, lens, angs = utils.get_reds(uvd, pick_data_ants=True, + bl_len_range=(10, 20), bl_deg_range=(0, 1)) - bls1, bls2, blps = utils.construct_blpairs(reds[0], exclude_auto_bls=True, exclude_permutations=True) + bls1, bls2, blps = utils.construct_blpairs(reds[0], exclude_auto_bls=True, + exclude_permutations=True) # setup PSpecData - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None], beam=beam) + ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], + wgts=[None, None], beam=beam) ds.Jy_to_mK() # get pspec @@ -94,10 +109,13 @@ def test_noise_validation(): taper='none', sampling=False, little_h=True, spw_ranges=[(0, 50)], verbose=False) # get noise spectra from one of the blpairs - P_N = uvp.generate_noise_spectra(0, 'xx', Tsys, blpairs=uvp.get_blpairs()[:1], num_steps=2000).values()[0][0, 0] + P_N = list(uvp.generate_noise_spectra(0, ('xx','xx'), Tsys, + blpairs=uvp.get_blpairs()[:1], + num_steps=2000).values())[0][0, 0] # get P_std of real spectra for each baseline across time axis - P_stds = np.array([np.std(uvp.get_data((0, bl, 'xx')).real, axis=1) for bl in uvp.get_blpairs()]) + P_stds = np.array([np.std(uvp.get_data((0, bl, ('xx','xx'))).real, axis=1) + for bl in uvp.get_blpairs()]) # get average P_std_avg and its standard error P_std_avg = np.mean(P_stds) diff --git a/hera_pspec/tests/test_plot.py b/hera_pspec/tests/test_plot.py index 9d869be7..528cfbe9 100644 --- a/hera_pspec/tests/test_plot.py +++ b/hera_pspec/tests/test_plot.py @@ -99,21 +99,21 @@ def test_plot_average(self): blps = [blp for blp in blpairs] # Plot the spectra averaged over baseline-pairs and times - f1 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol='xx', + f1 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), average_blpairs=True, average_times=True) elements = [(matplotlib.lines.Line2D, 1),] self.assertTrue( axes_contains(f1.axes[0], elements) ) plt.close(f1) # Average over baseline-pairs but keep the time bins intact - f2 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol='xx', + f2 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), average_blpairs=True, average_times=False) elements = [(matplotlib.lines.Line2D, self.uvp.Ntimes),] self.assertTrue( axes_contains(f2.axes[0], elements) ) plt.close(f2) # Average over times, but keep the baseline-pairs separate - f3 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol='xx', + f3 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), average_blpairs=False, average_times=True) elements = [(matplotlib.lines.Line2D, self.uvp.Nblpairs),] self.assertTrue( axes_contains(f3.axes[0], elements) ) @@ -121,7 +121,7 @@ def test_plot_average(self): # Plot the spectra averaged over baseline-pairs and times, but also # fold the delay axis - f4 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol='xx', + f4 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), average_blpairs=True, average_times=True, fold=True) elements = [(matplotlib.lines.Line2D, 1),] @@ -129,7 +129,7 @@ def test_plot_average(self): plt.close(f4) # Plot imaginary part - f4 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol='xx', + f4 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), average_blpairs=False, average_times=True, component='imag') elements = [(matplotlib.lines.Line2D, self.uvp.Nblpairs),] @@ -137,7 +137,7 @@ def test_plot_average(self): plt.close(f4) # Plot abs - f5 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol='xx', + f5 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), average_blpairs=False, average_times=True, component='abs') elements = [(matplotlib.lines.Line2D, self.uvp.Nblpairs),] @@ -153,14 +153,16 @@ def test_plot_average(self): robust_std=False, verbose=False) f6 = plot.delay_spectrum(uvp_avg, uvp_avg.get_blpairs(), spw=0, - pol='xx', average_blpairs=False, average_times=False, + pol=('xx','xx'), average_blpairs=False, + average_times=False, component='real', error='bs_std', lines=False, markers=True) plt.close(f6) # plot errorbar instead of pspec f7 = plot.delay_spectrum(uvp_avg, uvp_avg.get_blpairs(), spw=0, - pol='xx', average_blpairs=False, average_times=False, + pol=('xx','xx'), average_blpairs=False, + average_times=False, component='real', lines=False, markers=True, plot_stats='bs_std') plt.close(f7) @@ -176,7 +178,7 @@ def test_plot_cosmo(self): # Set cosmology and plot in non-delay (i.e. cosmological) units self.uvp.set_cosmology(conversions.Cosmo_Conversions()) - f1 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol='xx', + f1 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), average_blpairs=True, average_times=True, delay=False) elements = [(matplotlib.lines.Line2D, 1), (matplotlib.legend.Legend, 0)] @@ -184,9 +186,10 @@ def test_plot_cosmo(self): plt.close(f1) # Plot in Delta^2 units - f2 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol='xx', + f2 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), average_blpairs=True, average_times=True, - delay=False, deltasq=True, legend=True, label_type='blpair') + delay=False, deltasq=True, legend=True, + label_type='blpair') # Should contain 1 line and 1 legend elements = [(matplotlib.lines.Line2D, 1), (matplotlib.legend.Legend, 1)] self.assertTrue( axes_contains(f2.axes[0], elements) ) @@ -201,7 +204,7 @@ def test_delay_spectrum_misc(self): blps = [blp for blp in blpairs] # times selection, label_type - f1 = plot.delay_spectrum(self.uvp, blpairs[:1], spw=0, pol='xx', + f1 = plot.delay_spectrum(self.uvp, blpairs[:1], spw=0, pol=('xx','xx'), times=self.uvp.time_avg_array[:1], lines=False, markers=True, logscale=False, label_type='key', force_plot=False) @@ -216,12 +219,14 @@ def test_delay_spectrum_misc(self): uvp = uvp + _uvp nt.assert_raises(ValueError, plot.delay_spectrum, uvp, uvp.get_blpairs(), 0, 'xx') - f2 = plot.delay_spectrum(uvp, uvp.get_blpairs(), 0, 'xx', force_plot=True, - label_type='blpairt', logscale=False, lines=True, markers=True) + f2 = plot.delay_spectrum(uvp, uvp.get_blpairs(), 0, ('xx','xx'), + force_plot=True, label_type='blpairt', + logscale=False, lines=True, markers=True) plt.close(f2) # exceptions - nt.assert_raises(ValueError, plot.delay_spectrum, uvp, uvp.get_blpairs()[:3], 0, 'xx', + nt.assert_raises(ValueError, plot.delay_spectrum, uvp, + uvp.get_blpairs()[:3], 0, ('xx','xx'), label_type='foo') @@ -235,25 +240,25 @@ def test_plot_waterfall(self): # Set cosmology and plot in non-delay (i.e. cosmological) units self.uvp.set_cosmology(conversions.Cosmo_Conversions(), overwrite=True) - f1 = plot.delay_waterfall(self.uvp, [blps,], spw=0, pol='xx', + f1 = plot.delay_waterfall(self.uvp, [blps,], spw=0, pol=('xx','xx'), average_blpairs=True, delay=False) plt.close(f1) # Plot in Delta^2 units - f2 = plot.delay_waterfall(self.uvp, [blps,], spw=0, pol='xx', + f2 = plot.delay_waterfall(self.uvp, [blps,], spw=0, pol=('xx','xx'), average_blpairs=True, delay=False, deltasq=True) plt.close(f2) # Try some other arguments - f3 = plot.delay_waterfall(self.uvp, [blpairs,], spw=0, pol='xx', + f3 = plot.delay_waterfall(self.uvp, [blpairs,], spw=0, pol=('xx','xx'), average_blpairs=False, delay=True, log=False, vmin=-1., vmax=3., cmap='RdBu', fold=True, component='abs') plt.close(f3) # Try with imaginary component - f4 = plot.delay_waterfall(self.uvp, [blpairs,], spw=0, pol='xx', + f4 = plot.delay_waterfall(self.uvp, [blpairs,], spw=0, pol=('xx','xx'), average_blpairs=False, delay=True, log=False, vmin=-1., vmax=3., cmap='RdBu', fold=True, component='imag') @@ -261,7 +266,8 @@ def test_plot_waterfall(self): # Try some more arguments fig, axes = plt.subplots(1, len(blps)) - plot.delay_waterfall(self.uvp, [blps,], spw=0, pol='xx', lst_in_hrs=False, + plot.delay_waterfall(self.uvp, [blps,], spw=0, pol=('xx','xx'), + lst_in_hrs=False, times=np.unique(self.uvp.time_avg_array)[:10], axes=axes, component='abs', title_type='blvec') plt.close() @@ -272,19 +278,24 @@ def test_plot_waterfall(self): _uvp = copy.deepcopy(self.uvp) _uvp.blpair_array += i * 20 uvp += _uvp - nt.assert_raises(ValueError, plot.delay_waterfall, uvp, uvp.get_blpairs(), 0, 'xx') - fig = plot.delay_waterfall(uvp, uvp.get_blpairs(), 0, 'xx', force_plot=True) + nt.assert_raises(ValueError, plot.delay_waterfall, uvp, + uvp.get_blpairs(), 0, ('xx','xx')) + fig = plot.delay_waterfall(uvp, uvp.get_blpairs(), 0, ('xx','xx'), + force_plot=True) plt.close() def test_uvdata_waterfalls(self): - """ Test waterfall plotter """ + """ + Test waterfall plotter + """ uvd = copy.deepcopy(self.uvd) basename = "test_waterfall_plots_3423523923_{bl}_{pol}" for d in ['data', 'flags', 'nsamples']: print("running on {}".format(d)) - plot.plot_uvdata_waterfalls(uvd, basename, vmin=0, vmax=100, data=d, plot_mode='real') + plot.plot_uvdata_waterfalls(uvd, basename, vmin=0, vmax=100, + data=d, plot_mode='real') figfiles = glob.glob("test_waterfall_plots_3423523923_*_*.png") nt.assert_equal(len(figfiles), 15) @@ -295,59 +306,81 @@ def test_delay_wedge(self): """ Tests for plot.delay_wedge """ # construct new uvp reds, lens, angs = utils.get_reds(self.ds.dsets[0], pick_data_ants=True) - bls1, bls2, blps, _, _ = utils.calc_blpair_reds(self.ds.dsets[0], self.ds.dsets[1], - exclude_auto_bls=False, exclude_permutations=True) + bls1, bls2, blps, _, _ = utils.calc_blpair_reds(self.ds.dsets[0], + self.ds.dsets[1], + exclude_auto_bls=False, + exclude_permutations=True) uvp = self.ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), spw_ranges=[(300, 350)], input_data_weight='identity', norm='I', taper='blackman-harris', verbose=False) # test basic delay_wedge call - f1 = plot.delay_wedge(uvp, 0, 'xx', blpairs=None, times=None, fold=False, delay=True, - rotate=False, component='real', log10=False, loglog=False, red_tol=1.0, - center_line=False, horizon_lines=False, title=None, ax=None, cmap='viridis', - figsize=(8, 6), deltasq=False, colorbar=False, cbax=None, vmin=None, vmax=None, - edgecolor='none', flip_xax=False, flip_yax=False, lw=2) + f1 = plot.delay_wedge(uvp, 0, ('xx','xx'), blpairs=None, times=None, + fold=False, delay=True, rotate=False, + component='real', log10=False, loglog=False, + red_tol=1.0, center_line=False, + horizon_lines=False, title=None, ax=None, + cmap='viridis', figsize=(8, 6), deltasq=False, + colorbar=False, cbax=None, vmin=None, vmax=None, + edgecolor='none', flip_xax=False, flip_yax=False, + lw=2) plt.close() # specify blpairs and times - f2 = plot.delay_wedge(uvp, 0, 'xx', blpairs=uvp.get_blpairs()[-5:], times=uvp.time_avg_array[:1], + f2 = plot.delay_wedge(uvp, 0, ('xx','xx'), + blpairs=uvp.get_blpairs()[-5:], + times=uvp.time_avg_array[:1], fold=False, delay=True, component='imag', rotate=False, log10=False, loglog=False, red_tol=1.0, - center_line=False, horizon_lines=False, title=None, ax=None, cmap='viridis', - figsize=(8, 6), deltasq=False, colorbar=False, cbax=None, vmin=None, vmax=None, - edgecolor='none', flip_xax=False, flip_yax=False, lw=2) + center_line=False, horizon_lines=False, title=None, + ax=None, cmap='viridis', + figsize=(8, 6), deltasq=False, colorbar=False, + cbax=None, vmin=None, vmax=None, + edgecolor='none', flip_xax=False, flip_yax=False, + lw=2) plt.close() # fold, deltasq, cosmo and log10, loglog - f3 = plot.delay_wedge(uvp, 0, 'xx', blpairs=None, times=None, fold=True, delay=False, component='abs', + f3 = plot.delay_wedge(uvp, 0, ('xx','xx'), blpairs=None, times=None, + fold=True, delay=False, component='abs', rotate=False, log10=True, loglog=True, red_tol=1.0, - center_line=False, horizon_lines=False, title='hello', ax=None, cmap='viridis', - figsize=(8, 6), deltasq=True, colorbar=False, cbax=None, vmin=None, vmax=None, - edgecolor='none', flip_xax=False, flip_yax=False, lw=2) + center_line=False, horizon_lines=False, + title='hello', ax=None, cmap='viridis', + figsize=(8, 6), deltasq=True, colorbar=False, + cbax=None, vmin=None, vmax=None, + edgecolor='none', flip_xax=False, flip_yax=False, + lw=2) plt.close() # colorbar, vranges, flip_axes, edgecolors, lines - f4 = plot.delay_wedge(uvp, 0, 'xx', blpairs=None, times=None, fold=False, delay=False, component='abs', + f4 = plot.delay_wedge(uvp, 0, ('xx','xx'), blpairs=None, times=None, + fold=False, delay=False, component='abs', rotate=False, log10=True, loglog=False, red_tol=1.0, - center_line=True, horizon_lines=True, title='hello', ax=None, cmap='viridis', - figsize=(8, 6), deltasq=True, colorbar=True, cbax=None, vmin=6, vmax=15, - edgecolor='grey', flip_xax=True, flip_yax=True, lw=2, set_bl_tick_minor=True) + center_line=True, horizon_lines=True, title='hello', + ax=None, cmap='viridis', figsize=(8, 6), + deltasq=True, colorbar=True, cbax=None, vmin=6, + vmax=15, edgecolor='grey', flip_xax=True, + flip_yax=True, lw=2, set_bl_tick_minor=True) plt.close() # feed axes, red_tol fig, ax = plt.subplots() cbax = fig.add_axes([0.85, 0.1, 0.05, 0.9]) cbax.axis('off') - plot.delay_wedge(uvp, 0, 'xx', blpairs=None, times=None, fold=False, delay=True, component='abs', - rotate=True, log10=True, loglog=False, red_tol=10.0, - center_line=False, horizon_lines=False, ax=ax, cmap='viridis', - figsize=(8, 6), deltasq=False, colorbar=True, cbax=cbax, vmin=None, vmax=None, - edgecolor='none', flip_xax=False, flip_yax=False, lw=2, set_bl_tick_major=True) + plot.delay_wedge(uvp, 0, ('xx','xx'), blpairs=None, times=None, + fold=False, delay=True, component='abs', + rotate=True, log10=True, loglog=False, red_tol=10.0, + center_line=False, horizon_lines=False, ax=ax, + cmap='viridis', figsize=(8, 6), deltasq=False, + colorbar=True, cbax=cbax, vmin=None, vmax=None, + edgecolor='none', flip_xax=False, flip_yax=False, + lw=2, set_bl_tick_major=True) plt.close() # test exceptions - nt.assert_raises(ValueError, plot.delay_wedge, uvp, 0, 'xx', component='foo') + nt.assert_raises(ValueError, plot.delay_wedge, uvp, 0, ('xx','xx'), + component='foo') plt.close() diff --git a/hera_pspec/tests/test_pspecbeam.py b/hera_pspec/tests/test_pspecbeam.py index 0862408d..971ea524 100644 --- a/hera_pspec/tests/test_pspecbeam.py +++ b/hera_pspec/tests/test_pspecbeam.py @@ -248,11 +248,14 @@ def test_PSpecBeamBase(self): def test_get_Omegas(self): beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') beam = pspecbeam.PSpecBeamUV(beamfile) - OP, OPP = beam.get_Omegas('xx') + OP, OPP = beam.get_Omegas(('xx','xx')) nt.assert_equal(OP.shape, (26, 1)) nt.assert_equal(OPP.shape, (26, 1)) - OP, OPP = beam.get_Omegas([-5, -6]) + OP, OPP = beam.get_Omegas([(-5,-5), (-6,-6)]) nt.assert_equal(OP.shape, (26, 2)) nt.assert_equal(OPP.shape, (26, 2)) + + nt.assert_raises(TypeError, beam.get_Omegas, 'xx') + nt.assert_raises(NotImplementedError, beam.get_Omegas, [('pI','pQ'),]) diff --git a/hera_pspec/tests/test_pspecdata.py b/hera_pspec/tests/test_pspecdata.py index eb9f3963..73403fd5 100644 --- a/hera_pspec/tests/test_pspecdata.py +++ b/hera_pspec/tests/test_pspecdata.py @@ -267,7 +267,7 @@ def test_get_Q_alt(self): self.assertAlmostEqual(np.imag(xQx), 0.) x_vect = np.ones(vect_length) - Q_matrix = self.ds.get_Q_alt(vect_length/2) + Q_matrix = self.ds.get_Q_alt(vect_length//2) xQx = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, x_vect)) self.assertAlmostEqual(xQx, np.abs(vect_length**2.)) # Sending in sinusoids for x and y should give delta functions @@ -291,7 +291,7 @@ def test_get_Q_alt(self): self.assertAlmostEqual(np.imag(xQx), 0.) x_vect = np.ones(vect_length) - Q_matrix = self.ds.get_Q_alt((vect_length-2)/2-1) + Q_matrix = self.ds.get_Q_alt((vect_length-2)//2-1) xQx = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, x_vect)) self.assertAlmostEqual(xQx, np.abs(vect_length**2.)) # Sending in sinusoids for x and y should give delta functions @@ -310,6 +310,9 @@ def test_get_Q_alt(self): Q_matrix = self.ds.get_Q_alt(alpha, allow_fft=False) Q_diff_norm = np.linalg.norm(Q_matrix - Q_matrix_fft) self.assertLessEqual(Q_diff_norm, multiplicative_tolerance) + + # Check for error handling + nt.assert_raises(ValueError, self.ds.set_Ndlys, vect_length+100) def test_get_unnormed_E(self): """ @@ -638,9 +641,9 @@ def test_get_H(self): for taper in taper_selection: self.ds.set_taper(taper) - self.ds.set_Ndlys(Nfreq/3) + self.ds.set_Ndlys(Nfreq//3) H = self.ds.get_H(key1, key2) - self.assertEqual(H.shape, (Nfreq/3, Nfreq/3)) # Test shape + self.assertEqual(H.shape, (Nfreq//3, Nfreq//3)) # Test shape self.ds.set_Ndlys() H = self.ds.get_H(key1, key2) @@ -801,7 +804,12 @@ def test_scalar(self): gauss = pspecbeam.PSpecBeamGauss(0.8, np.linspace(115e6, 130e6, 50, endpoint=False)) ds2 = pspecdata.PSpecData(dsets=self.d, wgts=self.w, beam=gauss) - + + # Check normal execution + scalar = self.ds.scalar(('xx','xx')) + scalar = self.ds.scalar(('xx','xx'), taper_override='none') + nt.assert_raises(NotImplementedError, self.ds.scalar, ('xx','yy')) + # Precomputed results in the following test were done "by hand" # using iPython notebook "Scalar_dev2.ipynb" in the tests/ directory # FIXME: Uncomment when pyuvdata support for this is ready @@ -814,26 +822,34 @@ def test_scalar(self): def test_validate_datasets(self): # test freq exception uvd = copy.deepcopy(self.d[0]) - uvd2 = uvd.select(frequencies=np.unique(uvd.freq_array)[:10], inplace=False) + uvd2 = uvd.select(frequencies=np.unique(uvd.freq_array)[:10], + inplace=False) ds = pspecdata.PSpecData(dsets=[uvd, uvd2], wgts=[None, None]) nt.assert_raises(ValueError, ds.validate_datasets) + # test time exception uvd2 = uvd.select(times=np.unique(uvd.time_array)[:10], inplace=False) ds = pspecdata.PSpecData(dsets=[uvd, uvd2], wgts=[None, None]) nt.assert_raises(ValueError, ds.validate_datasets) + # test std exception ds.dsets_std=ds.dsets_std[:1] nt.assert_raises(ValueError, ds.validate_datasets) + # test wgt exception ds.wgts = ds.wgts[:1] nt.assert_raises(ValueError, ds.validate_datasets) + # test warnings uvd = copy.deepcopy(self.d[0]) uvd2 = copy.deepcopy(self.d[0]) - uvd.select(frequencies=np.unique(uvd.freq_array)[:10], times=np.unique(uvd.time_array)[:10]) - uvd2.select(frequencies=np.unique(uvd2.freq_array)[10:20], times=np.unique(uvd2.time_array)[10:20]) + uvd.select(frequencies=np.unique(uvd.freq_array)[:10], + times=np.unique(uvd.time_array)[:10]) + uvd2.select(frequencies=np.unique(uvd2.freq_array)[10:20], + times=np.unique(uvd2.time_array)[10:20]) ds = pspecdata.PSpecData(dsets=[uvd, uvd2], wgts=[None, None]) ds.validate_datasets() + # test phasing uvd = copy.deepcopy(self.d[0]) uvd2 = copy.deepcopy(self.d[0]) @@ -842,6 +858,9 @@ def test_validate_datasets(self): nt.assert_raises(ValueError, ds.validate_datasets) uvd2.phase_to_time(Time(2458042.5, format='jd')) ds.validate_datasets() + + # test polarization + ds.validate_pol((0,1), ('xx', 'xx')) def test_rephase_to_dset(self): # generate two uvd objects w/ different LST grids @@ -857,7 +876,7 @@ def test_rephase_to_dset(self): # rephase and get pspec ds.rephase_to_dset(0) uvp2 = ds.pspec(bls, bls, (0, 1), pols=('xx','xx'), verbose=False) - blp = (0, ((37,39),(37,39)), 'xx') + blp = (0, ((37,39),(37,39)), ('xx','xx')) nt.assert_true(np.isclose(np.abs(uvp2.get_data(blp)/uvp1.get_data(blp)), 1.0).min()) def test_Jy_to_mK(self): @@ -869,7 +888,8 @@ def test_Jy_to_mK(self): ds.Jy_to_mK() nt.assert_true(ds.dsets[0].vis_units, 'mK') nt.assert_true(ds.dsets[1].vis_units, 'mK') - nt.assert_true(uvd.get_data(24, 25, 'xx')[30, 30] / ds.dsets[0].get_data(24, 25, 'xx')[30, 30] < 1.0) + nt.assert_true(uvd.get_data(24, 25, 'xx')[30, 30] \ + / ds.dsets[0].get_data(24, 25, 'xx')[30, 30] < 1.0) # test feeding beam ds2 = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], @@ -914,6 +934,7 @@ def test_units(self): nt.assert_raises(IndexError, ds.units) ds.add(self.uvd, None) # test basic execution + vis_u, norm_u = ds.units(little_h=False) vis_u, norm_u = ds.units() nt.assert_equal(vis_u, "UNCALIB") nt.assert_equal(norm_u, "Hz str [beam normalization not specified]") @@ -958,7 +979,7 @@ def test_pspec(self): # check with redundant baseline group list antpos, ants = uvd.get_ENU_antpos(pick_data_ants=True) antpos = dict(zip(ants, antpos)) - red_bls = map(lambda blg: sorted(blg), redcal.get_pos_reds(antpos))[2] + red_bls = [sorted(blg) for blg in redcal.get_pos_reds(antpos)][2] bls1, bls2, blps = utils.construct_blpairs(red_bls, exclude_permutations=True) uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', little_h=True, verbose=False) @@ -997,7 +1018,7 @@ def test_pspec(self): nt.assert_equal(uvp.Nspws, 3) nt.assert_equal(uvp.Nspwdlys, 43) nt.assert_equal(uvp.data_array[0].shape, (240, 14, 1)) - nt.assert_equal(uvp.get_data((0, 124125124125, 'xx')).shape, (60, 14)) + nt.assert_equal(uvp.get_data((0, 124125124125, ('xx','xx'))).shape, (60, 14)) uvp.select(spws=[1]) nt.assert_equal(uvp.Nspws, 1) @@ -1008,7 +1029,7 @@ def test_pspec(self): uvd = copy.deepcopy(self.uvd) ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm) uvp = ds.pspec(bls, bls, (0, 1), ('xx','xx'), spw_ranges=[(10, 24)], verbose=False) - nt.assert_raises(NotImplementedError, ds.pspec, bls, bls, (0, 1), pols=[('xx','yy')]) + #nt.assert_raises(NotImplementedError, ds.pspec, bls, bls, (0, 1), pols=[('xx','yy')]) uvd = copy.deepcopy(self.uvd) ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm) uvp = ds.pspec(bls, bls, (0, 1), [('xx','xx'), ('yy','yy')], spw_ranges=[(10, 24)], verbose=False) @@ -1063,7 +1084,7 @@ def test_pspec(self): spw_ranges=[(20, 30)]) nt.assert_equal(len(ds._identity_Y), len(ds._identity_G), len(ds._identity_H)) nt.assert_equal(len(ds._identity_Y), 1) - nt.assert_equal(ds._identity_Y.keys()[0], ((0, 24, 25, 'xx'), (1, 24, 25, 'xx'))) + nt.assert_equal(list(ds._identity_Y.keys())[0], ((0, 24, 25, 'xx'), (1, 24, 25, 'xx'))) # assert caching is not used when inappropriate ds.dsets[0].flag_array[ds.dsets[0].antpair2ind(37, 38, ordered=False), :, 25, :] = True uvp = ds.pspec([(24, 25), (37, 38)], [(24, 25), (37, 38)], (0, 1), ('xx', 'xx'), @@ -1108,7 +1129,7 @@ def test_normalization(self): # hera_pspec OQE ds = pspecdata.PSpecData(dsets=[d1, d2], wgts=[None, None], beam=beam) uvp = ds.pspec(bls1, bls2, (0, 1), pols=('xx','xx'), taper='none', input_data_weight='identity', norm='I', sampling=True) - oqe = uvp.get_data((0, ((24, 25), (37, 38)), 'xx'))[0] + oqe = uvp.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] # assert answers are same to within 3% nt.assert_true(np.isclose(np.real(oqe)/np.real(legacy), 1, atol=0.03, rtol=0.03).all()) # taper @@ -1121,7 +1142,7 @@ def test_normalization(self): # hera_pspec OQE ds = pspecdata.PSpecData(dsets=[d1, d2], wgts=[None, None], beam=beam) uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), taper='blackman-harris', input_data_weight='identity', norm='I') - oqe = uvp.get_data((0, ((24, 25), (37, 38)), 'xx'))[0] + oqe = uvp.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] # assert answers are same to within 3% nt.assert_true(np.isclose(np.real(oqe)/np.real(legacy), 1, atol=0.03, rtol=0.03).all()) @@ -1161,9 +1182,9 @@ def test_broadcast_dset_flags(self): uvp = ds.pspec([(24, 25), (37, 38), (38, 39)], [(24, 25), (37, 38), (38, 39)], (0, 1), ('xx', 'xx'), spw_ranges=[(400, 450)], verbose=False) # assert flag broadcast above hits weight arrays in uvp - nt.assert_true(np.all(np.isclose(uvp.get_wgts((0, ((24, 25), (24, 25)), 'xx'))[3], 0.0))) + nt.assert_true(np.all(np.isclose(uvp.get_wgts((0, ((24, 25), (24, 25)), ('xx','xx')))[3], 0.0))) # assert flag broadcast above hits integration arrays - nt.assert_true(np.isclose(uvp.get_integrations((0, ((24, 25), (24, 25)), 'xx'))[3], 0.0)) + nt.assert_true(np.isclose(uvp.get_integrations((0, ((24, 25), (24, 25)), ('xx','xx')))[3], 0.0)) # average spectra avg_uvp = uvp.average_spectra(blpair_groups=[sorted(np.unique(uvp.blpair_array))], time_avg=True, inplace=False) # repeat but change data in flagged portion @@ -1195,8 +1216,8 @@ def test_RFI_flag_propagation(self): uvp_unflagged = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', little_h=True, verbose=False) - qe_unflagged = uvp_unflagged.get_data((0, ((24, 25), (37, 38)), 'xx'))[0] - qe_flagged = uvp_flagged.get_data((0, ((24, 25), (37, 38)), 'xx'))[0] + qe_unflagged = uvp_unflagged.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] + qe_flagged = uvp_flagged.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] # assert answers are same to within 0.1% nt.assert_true(np.isclose(np.real(qe_unflagged)/np.real(qe_flagged), 1, atol=0.001, rtol=0.001).all()) @@ -1213,8 +1234,8 @@ def test_RFI_flag_propagation(self): uvp_flagged_mod = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', little_h=True, verbose=False) - qe_flagged_mod = uvp_flagged_mod.get_data((0, ((24, 25), (37, 38)), 'xx'))[0] - qe_flagged = uvp_flagged.get_data((0, ((24, 25), (37, 38)), 'xx'))[0] + qe_flagged_mod = uvp_flagged_mod.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] + qe_flagged = uvp_flagged.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] # assert answers are same to within 0.1% nt.assert_true(np.isclose(np.real(qe_flagged_mod), np.real(qe_flagged), atol=0.001, rtol=0.001).all()) @@ -1271,8 +1292,10 @@ def test_pspec_run(): # test basic execution if os.path.exists("./out.hdf5"): os.remove("./out.hdf5") - psc, ds = pspecdata.pspec_run(fnames, "./out.hdf5", Jy2mK=False, verbose=False, overwrite=True, - bl_len_range=(14, 15), bl_deg_range=(50, 70), psname_ext='_0') + psc, ds = pspecdata.pspec_run(fnames, "./out.hdf5", Jy2mK=False, + verbose=False, overwrite=True, + bl_len_range=(14, 15), bl_deg_range=(50, 70), + psname_ext='_0') nt.assert_true(isinstance(psc, container.PSpecContainer)) nt.assert_equal(psc.groups(), ['dset0_dset1']) nt.assert_equal(psc.spectra(psc.groups()[0]), ['dset0_x_dset1_0']) @@ -1297,7 +1320,7 @@ def test_pspec_run(): nt.assert_true(uvp.vis_units, "mK") # assert only blpairs that were fed are present nt.assert_equal(uvp.bl_array.tolist(), [137138, 152153]) - nt.assert_equal(uvp.pol_array.tolist(), [-5, -5]) + nt.assert_equal(uvp.polpair_array.tolist(), [1515, 1515]) # assert weird cosmology was passed nt.assert_equal(uvp.cosmo, cosmo) # assert cov_array was calculated b/c std files were passed and store_cov diff --git a/hera_pspec/tests/test_testing.py b/hera_pspec/tests/test_testing.py index c1ec847a..752e8c64 100644 --- a/hera_pspec/tests/test_testing.py +++ b/hera_pspec/tests/test_testing.py @@ -16,7 +16,7 @@ def test_build_vanilla_uvpspec(): beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits')) uvp, cosmo = testing.build_vanilla_uvpspec(beam=beam) - beam_OP = beam.get_Omegas(uvp.pol_array[0])[0] + beam_OP = beam.get_Omegas(uvp.polpair_array[0])[0] nt.assert_equal(beam_OP.tolist(), uvp.OmegaP.tolist()) def test_uvpspec_from_data(): diff --git a/hera_pspec/tests/test_utils.py b/hera_pspec/tests/test_utils.py index 7d19d8ab..98ef9e4a 100644 --- a/hera_pspec/tests/test_utils.py +++ b/hera_pspec/tests/test_utils.py @@ -157,7 +157,7 @@ def test_calc_blpair_reds(self): (bls1, bls2, blps, xants1, xants2) = utils.calc_blpair_reds(uvd, uvd, filter_blpairs=True, exclude_auto_bls=False, exclude_permutations=True) nt.assert_equal(len(bls1), len(bls2), 15) - nt.assert_equal(blps, zip(bls1, bls2)) + nt.assert_equal(blps, list(zip(bls1, bls2))) nt.assert_equal(xants1, xants2) nt.assert_equal(len(xants1), 42) @@ -198,13 +198,16 @@ def test_calc_blpair_reds(self): uvd2 = copy.deepcopy(uvd) uvd2.antenna_positions[0] += 2 nt.assert_raises(AssertionError, utils.calc_blpair_reds, uvd, uvd2) - + + def test_get_delays(self): + utils.get_delays(np.linspace(100., 200., 50)*1e6) + def test_get_reds(self): fname = os.path.join(DATA_PATH, 'zen.all.xx.LST.1.06964.uvA') uvd = UVData() uvd.read_miriad(fname, read_data=False) antpos, ants = uvd.get_ENU_antpos() - antpos_d = dict(zip(ants, antpos)) + antpos_d = dict(list(zip(ants, antpos))) # test basic execution xants = [0, 1, 2] @@ -237,17 +240,17 @@ def test_config_pspec_blpairs(self): uv_template = os.path.join(DATA_PATH, "zen.{group}.{pol}.LST.1.28828.uvOCRSA") groupings = utils.config_pspec_blpairs(uv_template, [('xx', 'xx')], [('even', 'odd')], verbose=False, exclude_auto_bls=True) nt.assert_equal(len(groupings), 1) - nt.assert_equal(groupings.keys()[0], (('even', 'odd'), ('xx', 'xx'))) - nt.assert_equal(len(groupings.values()[0]), 11833) + nt.assert_equal(list(groupings.keys())[0], (('even', 'odd'), ('xx', 'xx'))) + nt.assert_equal(len(list(groupings.values())[0]), 11833) # test multiple, some non-existant pairs groupings = utils.config_pspec_blpairs(uv_template, [('xx', 'xx'), ('yy', 'yy')], [('even', 'odd'), ('even', 'odd')], verbose=False, exclude_auto_bls=True) nt.assert_equal(len(groupings), 1) - nt.assert_equal(groupings.keys()[0], (('even', 'odd'), ('xx', 'xx'))) + nt.assert_equal(list(groupings.keys())[0], (('even', 'odd'), ('xx', 'xx'))) # test xants groupings = utils.config_pspec_blpairs(uv_template, [('xx', 'xx')], [('even', 'odd')], xants=[0, 1, 2], verbose=False, exclude_auto_bls=True) - nt.assert_equal(len(groupings.values()[0]), 9735) + nt.assert_equal(len(list(groupings.values())[0]), 9735) # test exceptions nt.assert_raises(AssertionError, utils.config_pspec_blpairs, uv_template, [('xx', 'xx'), ('xx', 'xx')], [('even', 'odd')], verbose=False) @@ -281,23 +284,16 @@ def test_log(): os.remove("logf.log") -def test_hash(): - """ - Check that MD5 hashing works. - """ - hsh = utils.hash(np.ones((8,16))) - - def test_get_blvec_reds(): fname = os.path.join(DATA_PATH, "zen.2458042.17772.xx.HH.uvXA") uvd = UVData() uvd.read_miriad(fname) antpos, ants = uvd.get_ENU_antpos(pick_data_ants=True) - reds = redcal.get_pos_reds(dict(zip(ants, antpos))) + reds = redcal.get_pos_reds(dict(list(zip(ants, antpos)))) uvp = testing.uvpspec_from_data(fname, reds[:2], spw_ranges=[(10, 40)]) # test execution w/ dictionary - blvecs = dict(zip(uvp.bl_array, uvp.get_ENU_bl_vecs())) + blvecs = dict(list(zip(uvp.bl_array, uvp.get_ENU_bl_vecs()))) (red_bl_grp, red_bl_len, red_bl_ang, red_bl_tag) = utils.get_blvec_reds(blvecs, bl_error_tol=1.0) nt.assert_equal(len(red_bl_grp), 2) diff --git a/hera_pspec/tests/test_uvpspec.py b/hera_pspec/tests/test_uvpspec.py index 812fbf82..9a8004a7 100644 --- a/hera_pspec/tests/test_uvpspec.py +++ b/hera_pspec/tests/test_uvpspec.py @@ -35,26 +35,26 @@ def test_eq(self): def test_get_funcs(self): # get_data - d = self.uvp.get_data((0, ((1, 2), (1, 2)), 'xx')) + d = self.uvp.get_data((0, ((1, 2), (1, 2)), ('xx','xx'))) nt.assert_equal(d.shape, (10, 30)) nt.assert_true(d.dtype == np.complex) nt.assert_almost_equal(d[0,0], (101.1021011020000001+0j)) - d = self.uvp.get_data((0, ((1, 2), (1, 2)), -5)) + d = self.uvp.get_data((0, ((1, 2), (1, 2)), 1515)) nt.assert_almost_equal(d[0,0], (101.1021011020000001+0j)) - d = self.uvp.get_data((0, 101102101102, -5)) + d = self.uvp.get_data((0, 101102101102, 1515)) nt.assert_almost_equal(d[0,0], (101.1021011020000001+0j)) # get_wgts - w = self.uvp.get_wgts((0, ((1, 2), (1, 2)), 'xx')) + w = self.uvp.get_wgts((0, ((1, 2), (1, 2)), ('xx','xx'))) nt.assert_equal(w.shape, (10, 50, 2)) # should have Nfreq dim, not Ndlys nt.assert_true(w.dtype == np.float) nt.assert_equal(w[0,0,0], 1.0) # get_integrations - i = self.uvp.get_integrations((0, ((1, 2), (1, 2)), 'xx')) + i = self.uvp.get_integrations((0, ((1, 2), (1, 2)), ('xx','xx'))) nt.assert_equal(i.shape, (10,)) nt.assert_true(i.dtype == np.float) nt.assert_almost_equal(i[0], 1.0) # get nsample - n = self.uvp.get_nsamples((0, ((1, 2), (1, 2)), 'xx')) + n = self.uvp.get_nsamples((0, ((1, 2), (1, 2)), ('xx', 'xx'))) nt.assert_equal(n.shape, (10,)) nt.assert_true(n.dtype == np.float) nt.assert_almost_equal(n[0], 1.0) @@ -70,11 +70,11 @@ def test_get_funcs(self): nt.assert_equal(len(k_perp), 30) nt.assert_equal(len(k_para), 30) # test key expansion - key = (0, ((1, 2), (1, 2)), 'xx') + key = (0, ((1, 2), (1, 2)), ('xx','xx')) d = self.uvp.get_data(key) nt.assert_equal(d.shape, (10, 30)) # test key as dictionary - key = {'spw':0, 'blpair':((1, 2), (1, 2)), 'pol': 'xx'} + key = {'spw':0, 'blpair':((1, 2), (1, 2)), 'polpair': ('xx','xx')} d = self.uvp.get_data(key) nt.assert_equal(d.shape, (10, 30)) # test get_blpairs @@ -82,11 +82,12 @@ def test_get_funcs(self): nt.assert_equal(blps, [((1, 2), (1, 2)), ((1, 3), (1, 3)), ((2, 3), (2, 3))]) # test get all keys keys = self.uvp.get_all_keys() - nt.assert_equal(keys, [(0, ((1, 2), (1, 2)), 'xx'), (0, ((1, 3), (1, 3)), 'xx'), - (0, ((2, 3), (2, 3)), 'xx')]) + nt.assert_equal(keys, [(0, ((1, 2), (1, 2)), ('xx','xx')), + (0, ((1, 3), (1, 3)), ('xx','xx')), + (0, ((2, 3), (2, 3)), ('xx','xx'))]) # test omit_flags self.uvp.integration_array[0][self.uvp.blpair_to_indices(((1, 2), (1, 2)))[:2]] = 0.0 - nt.assert_equal(self.uvp.get_integrations((0, ((1, 2), (1, 2)), 'xx'), omit_flags=True).shape, (8,)) + nt.assert_equal(self.uvp.get_integrations((0, ((1, 2), (1, 2)), ('xx','xx')), omit_flags=True).shape, (8,)) def test_stats_array(self): # test get_data and set_data @@ -151,11 +152,11 @@ def test_blpair_conversions(self): def test_indices_funcs(self): # key to indices - spw, blpairts, pol = self.uvp.key_to_indices( (0, ((1,2),(1,2)), -5) ) + spw, blpairts, pol = self.uvp.key_to_indices( (0, ((1,2),(1,2)), 1515) ) nt.assert_equal(spw, 0) nt.assert_equal(pol, 0) nt.assert_true(np.isclose(blpairts, np.array([0,3,6,9,12,15,18,21,24,27])).min()) - spw, blpairts, pol = self.uvp.key_to_indices( (0, 101102101102, 'xx') ) + spw, blpairts, pol = self.uvp.key_to_indices( (0, 101102101102, ('xx','xx')) ) nt.assert_equal(spw, 0) nt.assert_equal(pol, 0) nt.assert_true(np.isclose(blpairts, np.array([0,3,6,9,12,15,18,21,24,27])).min()) @@ -177,11 +178,11 @@ def test_indices_funcs(self): np.testing.assert_array_equal(spw3, spw3b) # pol to indices - pol = self.uvp.pol_to_indices('xx') + pol = self.uvp.polpair_to_indices(('xx','xx')) nt.assert_equal(len(pol), 1) - pol = self.uvp.pol_to_indices(-5) + pol = self.uvp.polpair_to_indices(1515) nt.assert_equal(len(pol), 1) - pol = self.uvp.pol_to_indices(['xx', 'xx']) + pol = self.uvp.polpair_to_indices([('xx','xx'), ('xx','xx')]) nt.assert_equal(len(pol), 1) # test blpair to indices @@ -227,8 +228,8 @@ def test_select(self): nt.assert_equal(uvp2.Nblpairs, 2) # pol select - uvp2 = uvp.select(pols=[-5], inplace=False) - nt.assert_equal(uvp2.pol_array[0], -5) + uvp2 = uvp.select(polpairs=[1515,], inplace=False) + nt.assert_equal(uvp2.polpair_array[0], 1515) # time select uvp2 = uvp.select(times=np.unique(uvp.time_avg_array)[:1], inplace=False) @@ -237,15 +238,19 @@ def test_select(self): # test pol and blpair select, and check dimensionality of output uvp = copy.deepcopy(self.uvp) uvp.set_stats('hi', uvp.get_all_keys()[0], np.ones(300).reshape(10, 30)) - uvp2 = uvp.select(blpairs=uvp.get_blpairs(), pols=uvp.pol_array, inplace=False) + uvp2 = uvp.select(blpairs=uvp.get_blpairs(), polpairs=uvp.polpair_array, + inplace=False) nt.assert_equal(uvp2.data_array[0].shape, (30, 30, 1)) nt.assert_equal(uvp2.stats_array['hi'][0].shape, (30, 30, 1)) # test when both blp and pol array are non-sliceable uvp2, uvp3, uvp4 = copy.deepcopy(uvp), copy.deepcopy(uvp), copy.deepcopy(uvp) - uvp2.pol_array[0], uvp3.pol_array[0], uvp4.pol_array[0] = -6, -7, -8 + uvp2.polpair_array[0] = 1414 + uvp3.polpair_array[0] = 1313 + uvp4.polpair_array[0] = 1212 uvp = uvp + uvp2 + uvp3 + uvp4 - uvp5 = uvp.select(blpairs=[101102101102], pols=[-5, -6, -7], inplace=False) + uvp5 = uvp.select(blpairs=[101102101102], polpairs=[1515, 1414, 1313], + inplace=False) nt.assert_equal(uvp5.data_array[0].shape, (10, 30, 3)) # select only on lst @@ -255,14 +260,14 @@ def test_select(self): # check non-sliceable select: both pol and blpairs are non-sliceable uvp = copy.deepcopy(self.uvp) - # extend pol_array axis - for i in [-6, -7, -8]: + # extend polpair_array axis + for i in [1414, 1313, 1212]: _uvp = copy.deepcopy(self.uvp) - _uvp.pol_array[0] = i + _uvp.polpair_array[0] = i uvp += _uvp # create a purposely non-sliceable select across *both* pol and blpair - uvp.select(pols=[-5, -7, -8], blpairs=[101102101102, 102103102103]) + uvp.select(polpairs=[1414, 1313, 1212], blpairs=[101102101102, 102103102103]) nt.assert_equal(uvp.Npols, 3) nt.assert_equal(uvp.Nblpairs, 2) @@ -277,7 +282,7 @@ def test_check(self): del uvp.Ntimes nt.assert_raises(AssertionError, uvp.check) uvp.Ntimes = self.uvp.Ntimes - uvp.data_array = uvp.data_array.values()[0] + uvp.data_array = list(uvp.data_array.values())[0] nt.assert_raises(AssertionError, uvp.check) uvp.data_array = copy.deepcopy(self.uvp.data_array) @@ -318,33 +323,35 @@ def test_sense(self): uvp = copy.deepcopy(self.uvp) # test generate noise spectra - P_N = uvp.generate_noise_spectra(0, -5, 500, form='Pk', real=True) + P_N = uvp.generate_noise_spectra(0, 1515, 500, form='Pk', real=True) nt.assert_equal(P_N[101102101102].shape, (10, 30)) # test smaller system temp - P_N2 = uvp.generate_noise_spectra(0, -5, 400, form='Pk', real=True) + P_N2 = uvp.generate_noise_spectra(0, 1515, 400, form='Pk', real=True) nt.assert_true((P_N[101102101102] > P_N2[101102101102]).all()) # test complex - P_N2 = uvp.generate_noise_spectra(0, -5, 500, form='Pk', real=False) + P_N2 = uvp.generate_noise_spectra(0, 1515, 500, form='Pk', real=False) nt.assert_true((P_N[101102101102] < P_N2[101102101102]).all()) # test Dsq - Dsq = uvp.generate_noise_spectra(0, -5, 500, form='DelSq', real=True) + Dsq = uvp.generate_noise_spectra(0, 1515, 500, form='DelSq', real=True) nt.assert_equal(Dsq[101102101102].shape, (10, 30)) nt.assert_true(Dsq[101102101102][0, 1] < P_N[101102101102][0, 1]) # test a blpair selection - P_N = uvp.generate_noise_spectra(0, -5, 500, form='Pk', real=True) + P_N = uvp.generate_noise_spectra(0, 1515, 500, form='Pk', real=True) def test_average_spectra(self): uvp = copy.deepcopy(self.uvp) # test blpair averaging - blpairs = uvp.get_blpair_groups_from_bl_groups([[101102, 102103, 101103]], only_pairs_in_bls=False) - uvp2 = uvp.average_spectra(blpair_groups=blpairs, time_avg=False, inplace=False) + blpairs = uvp.get_blpair_groups_from_bl_groups([[101102, 102103, 101103]], + only_pairs_in_bls=False) + uvp2 = uvp.average_spectra(blpair_groups=blpairs, time_avg=False, + inplace=False) nt.assert_equal(uvp2.Nblpairs, 1) - nt.assert_true(np.isclose(uvp2.get_nsamples((0, 101102101102, 'xx')), 3.0).all()) - nt.assert_equal(uvp2.get_data((0, 101102101102, 'xx')).shape, (10, 30)) + nt.assert_true(np.isclose(uvp2.get_nsamples((0, 101102101102, ('xx','xx'))), 3.0).all()) + nt.assert_equal(uvp2.get_data((0, 101102101102, ('xx','xx'))).shape, (10, 30)) # Test blpair averaging (with baseline-pair weights) # Results should be identical with different weights here, as the data @@ -358,16 +365,18 @@ def test_average_spectra(self): blpair_weights=blpair_wgts, inplace=False) #nt.assert_equal(uvp2.Nblpairs, 1) - nt.assert_true(np.isclose(uvp3a.get_data((0, 101102101102, 'xx')), - uvp3b.get_data((0, 101102101102, 'xx'))).all()) + nt.assert_true(np.isclose( + uvp3a.get_data((0, 101102101102, ('xx','xx'))), + uvp3b.get_data((0, 101102101102, ('xx','xx')))).all()) #nt.assert_equal(uvp2.get_data((0, 101102101102, 'xx')).shape, (10, 30)) - - + # test time averaging uvp2 = uvp.average_spectra(time_avg=True, inplace=False) nt.assert_true(uvp2.Ntimes, 1) - nt.assert_true(np.isclose(uvp2.get_nsamples((0, 101102101102, 'xx')), 10.0).all()) - nt.assert_true(uvp2.get_data((0, 101102101102, 'xx')).shape, (1, 30)) + nt.assert_true(np.isclose( + uvp2.get_nsamples((0, 101102101102, ('xx','xx'))), 10.0).all()) + nt.assert_true(uvp2.get_data((0, 101102101102, ('xx','xx'))).shape, + (1, 30)) # ensure averaging works when multiple repeated baselines are present, but only # if time_avg = True uvp.blpair_array[uvp.blpair_to_indices(102103102103)] = 101102101102 @@ -388,12 +397,14 @@ def test_fold_spectra(self): uvd_std = UVData() uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) uvd_std.read_miriad(os.path.join(DATA_PATH,'zen.even.xx.LST.1.28828.uvOCRSA')) - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits")) + beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, + "HERA_NF_dipole_power.beamfits")) bls = [(37, 38), (38, 39), (52, 53)] - uvp1 = testing.uvpspec_from_data(uvd, bls, data_std=uvd_std, spw_ranges=[(0,17)], beam=beam) + uvp1 = testing.uvpspec_from_data(uvd, bls, data_std=uvd_std, + spw_ranges=[(0,17)], beam=beam) uvp1.fold_spectra() - cov_folded = uvp1.get_cov((0, ((37, 38), (38, 39)), 'xx')) - data_folded = uvp1.get_data((0, ((37,38), (38, 39)), 'xx')) + cov_folded = uvp1.get_cov((0, ((37, 38), (38, 39)), ('xx','xx'))) + data_folded = uvp1.get_data((0, ((37,38), (38, 39)), ('xx','xx'))) def test_str(self): @@ -403,7 +414,7 @@ def test_str(self): def test_compute_scalar(self): uvp = copy.deepcopy(self.uvp) # test basic execution - s = uvp.compute_scalar(0, 'xx', num_steps=1000, noise_scalar=False) + s = uvp.compute_scalar(0, ('xx','xx'), num_steps=1000, noise_scalar=False) nt.assert_almost_equal(s/553995277.90425551, 1.0, places=5) # test execptions del uvp.OmegaP @@ -435,46 +446,63 @@ def test_combine_uvpspec(self): # setup uvp build uvd = UVData() uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits")) + beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, + "HERA_NF_dipole_power.beamfits")) bls = [(37, 38), (38, 39), (52, 53)] - uvp1 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(20, 30), (60, 90)], beam=beam) + uvp1 = testing.uvpspec_from_data(uvd, bls, + spw_ranges=[(20, 30), (60, 90)], + beam=beam) + + print("uvp1 unique blps:", np.unique(uvp1.blpair_array)) # test failure due to overlapping data uvp2 = copy.deepcopy(uvp1) + + print("uvp2 unique blps:", np.unique(uvp2.blpair_array)) + nt.assert_raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) # test success across pol - uvp2.pol_array[0] = -6 + uvp2.polpair_array[0] = 1414 out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) nt.assert_equal(out.Npols, 2) - nt.assert_true(len(set(out.pol_array) ^ set([-5, -6])) == 0) - key = (0, ((37, 38), (38, 39)), 'xx') - nt.assert_true(np.all(np.isclose(out.get_nsamples(key), np.ones(10, dtype=np.float64)))) - nt.assert_true(np.all(np.isclose(out.get_integrations(key), 190 * np.ones(10, dtype=np.float64), atol=5, rtol=2))) + nt.assert_true(len(set(out.polpair_array) ^ set([1515, 1414])) == 0) + key = (0, ((37, 38), (38, 39)), ('xx','xx')) + nt.assert_true(np.all(np.isclose(out.get_nsamples(key), + np.ones(10, dtype=np.float64)))) + nt.assert_true(np.all(np.isclose(out.get_integrations(key), + 190 * np.ones(10, dtype=np.float64), atol=5, rtol=2))) # test multiple non-overlapping data axes uvp2.freq_array[0] = 0.0 nt.assert_raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) # test partial data overlap failure - uvp2 = testing.uvpspec_from_data(uvd, [(37, 38), (38, 39), (53, 54)], spw_ranges=[(20, 30), (60, 90)], beam=beam) + uvp2 = testing.uvpspec_from_data(uvd, [(37, 38), (38, 39), (53, 54)], + spw_ranges=[(20, 30), (60, 90)], + beam=beam) nt.assert_raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) - uvp2 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(20, 30), (60, 105)], beam=beam) + uvp2 = testing.uvpspec_from_data(uvd, bls, + spw_ranges=[(20, 30), (60, 105)], + beam=beam) nt.assert_raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) uvp2 = copy.deepcopy(uvp1) - uvp2.pol_array[0] = -6 + uvp2.polpair_array[0] = 1414 uvp2 = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) nt.assert_raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) # test concat across spw - uvp2 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(85, 101)], beam=beam) + uvp2 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(85, 101)], + beam=beam) out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) nt.assert_equal(out.Nspws, 3) nt.assert_equal(out.Nfreqs, 51) nt.assert_equal(out.Nspwdlys, 56) # test concat across blpairts - uvp2 = testing.uvpspec_from_data(uvd, [(53, 54), (67, 68)], spw_ranges=[(20, 30), (60, 90)], beam=beam) + uvp2 = testing.uvpspec_from_data(uvd, [(53, 54), (67, 68)], + spw_ranges=[(20, 30), (60, 90)], + beam=beam) out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) nt.assert_equal(out.Nblpairs, 4) nt.assert_equal(out.Nbls, 5) @@ -490,7 +518,7 @@ def test_combine_uvpspec(self): # test feed as strings uvp1 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(20, 30)], beam=beam) uvp2 = copy.deepcopy(uvp1) - uvp2.pol_array[0] = -6 + uvp2.polpair_array[0] = 1414 uvp1.write_hdf5('uvp1.hdf5', overwrite=True) uvp2.write_hdf5('uvp2.hdf5', overwrite=True) out = uvpspec.combine_uvpspec(['uvp1.hdf5', 'uvp2.hdf5'], verbose=False) @@ -501,7 +529,7 @@ def test_combine_uvpspec(self): # test UVPSpec __add__ uvp3 = copy.deepcopy(uvp1) - uvp3.pol_array[0] = -7 + uvp3.polpair_array[0] = 1313 out = uvp1 + uvp2 + uvp3 nt.assert_equal(out.Npols, 3) @@ -510,7 +538,7 @@ def test_combine_uvpspec(self): spw_ranges=[(20, 30), (60, 90)], n_dlys=[5, 15]) uvp4b = copy.deepcopy(uvp4) - uvp4b.pol_array[0] = -6 + uvp4b.polpair_array[0] = 1414 out = uvpspec.combine_uvpspec([uvp4, uvp4b], verbose=False) @@ -531,10 +559,10 @@ def test_combine_uvpspec_std(self): uvp2 = copy.deepcopy(uvp1) nt.assert_raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) # test success across pol - uvp2.pol_array[0] = -6 + uvp2.polpair_array[0] = 1414 out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) nt.assert_equal(out.Npols, 2) - nt.assert_true(len(set(out.pol_array) ^ set([-5, -6])) == 0) + nt.assert_true(len(set(out.polpair_array) ^ set([1515, 1414])) == 0) # test multiple non-overlapping data axes uvp2.freq_array[0] = 0.0 nt.assert_raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) @@ -550,7 +578,7 @@ def test_combine_uvpspec_std(self): data_std=uvd_std, beam=beam) nt.assert_raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) uvp2 = copy.deepcopy(uvp1) - uvp2.pol_array[0] = -6 + uvp2.polpair_array[0] = 1414 uvp2 = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) nt.assert_raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) @@ -582,7 +610,7 @@ def test_combine_uvpspec_std(self): uvp1 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(20, 30)], data_std=uvd_std, beam=beam) uvp2 = copy.deepcopy(uvp1) - uvp2.pol_array[0] = -6 + uvp2.polpair_array[0] = 1414 uvp1.write_hdf5('uvp1.hdf5', overwrite=True) uvp2.write_hdf5('uvp2.hdf5', overwrite=True) out = uvpspec.combine_uvpspec(['uvp1.hdf5', 'uvp2.hdf5'], verbose=False) @@ -592,7 +620,7 @@ def test_combine_uvpspec_std(self): # test UVPSpec __add__ uvp3 = copy.deepcopy(uvp1) - uvp3.pol_array[0] = -7 + uvp3.polpair_array[0] = 1313 out = uvp1 + uvp2 + uvp3 nt.assert_equal(out.Npols, 3) diff --git a/hera_pspec/tests/test_uvpspec_utils.py b/hera_pspec/tests/test_uvpspec_utils.py index 670b0fbc..9a769c4e 100644 --- a/hera_pspec/tests/test_uvpspec_utils.py +++ b/hera_pspec/tests/test_uvpspec_utils.py @@ -35,7 +35,7 @@ def test_select_common(): # Check that selecting on common times works uvp_list = [uvp1, uvp2] uvp_new = uvputils.select_common(uvp_list, spws=True, blpairs=True, - times=True, pols=True, inplace=False) + times=True, polpairs=True, inplace=False) nt.assert_equal(uvp_new[0], uvp_new[1]) np.testing.assert_array_equal(uvp_new[0].time_avg_array, uvp_new[1].time_avg_array) @@ -43,7 +43,7 @@ def test_select_common(): # Check that selecting on common baseline-pairs works uvp_list_2 = [uvp1, uvp2, uvp3] uvp_new_2 = uvputils.select_common(uvp_list_2, spws=True, blpairs=True, - times=True, pols=True, inplace=False) + times=True, polpairs=True, inplace=False) nt.assert_equal(uvp_new_2[0], uvp_new_2[1]) nt.assert_equal(uvp_new_2[0], uvp_new_2[2]) np.testing.assert_array_equal(uvp_new_2[0].time_avg_array, @@ -52,29 +52,29 @@ def test_select_common(): # Check that zero overlap in times raises a ValueError nt.assert_raises(ValueError, uvputils.select_common, [uvp2, uvp6], spws=True, blpairs=True, times=True, - pols=True, inplace=False) + polpairs=True, inplace=False) # Check that zero overlap in times does *not* raise a ValueError if # not selecting on times uvp_new_3 = uvputils.select_common([uvp2, uvp6], spws=True, blpairs=True, times=False, - pols=True, inplace=False) + polpairs=True, inplace=False) # Check that zero overlap in baselines raises a ValueError nt.assert_raises(ValueError, uvputils.select_common, [uvp3, uvp5], spws=True, blpairs=True, times=True, - pols=True, inplace=False) + polpairs=True, inplace=False) # Check that matching times are ignored when set to False uvp_new = uvputils.select_common(uvp_list, spws=True, blpairs=True, - times=False, pols=True, inplace=False) + times=False, polpairs=True, inplace=False) nt.assert_not_equal( np.sum(uvp_new[0].time_avg_array - uvp_new[1].time_avg_array), 0.) nt.assert_equal(len(uvp_new), len(uvp_list)) # Check that in-place selection works uvputils.select_common(uvp_list, spws=True, blpairs=True, - times=True, pols=True, inplace=True) + times=True, polpairs=True, inplace=True) nt.assert_equal(uvp1, uvp2) # check uvplist > 2 @@ -92,11 +92,58 @@ def test_select_common(): # check pol overlap uvp7 = copy.deepcopy(uvp1) - uvp7.pol_array[0] = -8 - nt.assert_raises(ValueError, uvputils.select_common, [uvp1, uvp7], pols=True) + uvp7.polpair_array[0] = 1212 # = (-8,-8) + nt.assert_raises(ValueError, uvputils.select_common, [uvp1, uvp7], + polpairs=True) +def test_get_blpairs_from_bls(): + """ + Test conversion of + """ + # setup uvp + beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beam = pspecbeam.PSpecBeamUV(beamfile) + uvp, cosmo = testing.build_vanilla_uvpspec(beam=beam) + + # Check that bls can be specified in several different ways + blps = uvputils._get_blpairs_from_bls(uvp, bls=101102) + blps = uvputils._get_blpairs_from_bls(uvp, bls=(101,102)) + blps = uvputils._get_blpairs_from_bls(uvp, bls=[101102, 101103]) + + +def test_polpair_int2tuple(): + """ + Test conversion of polpair ints to tuples. + """ + # List of polpairs to test + polpairs = [('xx','xx'), ('xx','yy'), ('xy', 'yx'), + ('pI','pI'), ('pI','pQ'), ('pQ','pQ'), ('pU','pU'), + ('pV','pV') ] + + # Check that lists and single items work + pol_ints = uvputils.polpair_tuple2int(polpairs) + uvputils.polpair_tuple2int(polpairs[0]) + uvputils.polpair_int2tuple(1515) + uvputils.polpair_int2tuple([1515,1414]) + uvputils.polpair_int2tuple(np.array([1515,1414])) + + # Test converting to int and then back again + pol_pairs_returned = uvputils.polpair_int2tuple(pol_ints, pol_strings=True) + for i in range(len(polpairs)): + nt.assert_equal(polpairs[i], pol_pairs_returned[i]) + + # Check that errors are raised appropriately + nt.assert_raises(AssertionError, uvputils.polpair_int2tuple, ('xx','xx')) + nt.assert_raises(AssertionError, uvputils.polpair_int2tuple, 'xx') + nt.assert_raises(AssertionError, uvputils.polpair_int2tuple, 'pI') + nt.assert_raises(ValueError, uvputils.polpair_int2tuple, 999) + nt.assert_raises(ValueError, uvputils.polpair_int2tuple, [999,]) + + def test_subtract_uvp(): - """ Test subtraction of two UVPSpec objects """ + """ + Test subtraction of two UVPSpec objects + """ # setup uvp beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') beam = pspecbeam.PSpecBeamUV(beamfile) @@ -146,7 +193,7 @@ def test_fast_is_in(): times = [ 0.1, 0.15, 0.2, 0.25, 0.1, 0.15, 0.2, 0.25, 0.1, 0.15, 0.3, 0.3, ] - src_blpts = np.array(zip(blps, times)) + src_blpts = np.array(list(zip(blps, times))) nt.assert_true(uvputils._fast_is_in(src_blpts, [(101102104103, 0.2)])[0]) @@ -159,7 +206,7 @@ def test_fast_lookup_blpairts(): times = [ 0.1, 0.15, 0.2, 0.25, 0.1, 0.15, 0.2, 0.25, 0.1, 0.15, 0.3, 0.3, ] - src_blpts = np.array(zip(blps, times)) + src_blpts = np.array(list(zip(blps, times))) # List of blpair-times to look up query_blpts = [(102101103104, 0.1), (101102104103, 0.1), (101102104103, 0.25)] diff --git a/hera_pspec/tests/test_version.py b/hera_pspec/tests/test_version.py index 445a7225..1e78856f 100644 --- a/hera_pspec/tests/test_version.py +++ b/hera_pspec/tests/test_version.py @@ -2,14 +2,19 @@ import nose.tools as nt import sys import os -from StringIO import StringIO +try: + # Python 2 + from cStringIO import StringIO +except: + # Python 3 + from io import StringIO import subprocess import hera_pspec def test_main(): version_info = hera_pspec.version.construct_version_info() - + saved_stdout = sys.stdout try: out = StringIO() @@ -22,7 +27,13 @@ def test_main(): o=version_info['git_origin'], b=version_info['git_branch'], d=version_info['git_description'])) + + git_info = hera_pspec.version._get_gitinfo_file() + finally: sys.stdout = saved_stdout + + # Test history string function + history = hera_pspec.version.history_string() diff --git a/hera_pspec/utils.py b/hera_pspec/utils.py index e27ed664..f6d8b4a0 100644 --- a/hera_pspec/utils.py +++ b/hera_pspec/utils.py @@ -1,10 +1,10 @@ import numpy as np -import os, time, md5, yaml +import os, time, yaml import itertools, argparse, glob import traceback, operator import aipy, uvtools import pylab as plt -from conversions import Cosmo_Conversions +from hera_pspec.conversions import Cosmo_Conversions from hera_cal import redcal from collections import OrderedDict as odict from pyuvdata import utils as uvutils @@ -12,14 +12,6 @@ from datetime import datetime -def hash(w): - """ - Return an MD5 hash of a set of weights. - """ - DeprecationWarning("utils.hash is deprecated.") - return md5.md5(w.copy(order='C')).digest() - - def cov(d1, w1, d2=None, w2=None, conj_1=False, conj_2=True): """ Computes an empirical covariance matrix from data vectors. If d1 is of size @@ -124,7 +116,7 @@ def construct_blpairs(bls, exclude_auto_bls=False, exclude_permutations=False, g blpairs = list(itertools.permutations(bls, 2)) # explicitly add in auto baseline pairs - blpairs.extend(zip(bls, bls)) + blpairs.extend(list(zip(bls, bls))) # iterate through and eliminate all autos if desired if exclude_auto_bls: @@ -135,8 +127,8 @@ def construct_blpairs(bls, exclude_auto_bls=False, exclude_permutations=False, g blpairs = new_blpairs # create bls1 and bls2 list - bls1 = map(lambda blp: blp[0], blpairs) - bls2 = map(lambda blp: blp[1], blpairs) + bls1 = [blp[0] for blp in blpairs] + bls2 = [blp[1] for blp in blpairs] # group baseline pairs if desired if group: @@ -157,11 +149,14 @@ def construct_blpairs(bls, exclude_auto_bls=False, exclude_permutations=False, g return bls1, bls2, blpairs -def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, xant_flag_thresh=0.95, exclude_auto_bls=False, - exclude_permutations=True, Nblps_per_group=None, bl_len_range=(0, 1e10), bl_deg_range=(0, 180)): +def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, + xant_flag_thresh=0.95, exclude_auto_bls=False, + exclude_permutations=True, Nblps_per_group=None, + bl_len_range=(0, 1e10), bl_deg_range=(0, 180)): """ - Use hera_cal.redcal to get matching, redundant baseline-pair groups from uvd1 and uvd2 - within the specified baseline tolerance, not including flagged ants. + Use hera_cal.redcal to get matching, redundant baseline-pair groups from + uvd1 and uvd2 within the specified baseline tolerance, not including + flagged ants. Parameters ---------- @@ -173,8 +168,8 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, xant_flag_thre Baseline-vector redundancy tolerance in meters filter_blpairs : bool, optional - if True, calculate xants and filters-out baseline pairs based on xant lists - and actual baselines in the data. + if True, calculate xants and filters-out baseline pairs based on + xant lists and actual baselines in the data. xant_flag_thresh : float, optional Fraction of 2D visibility (per-waterfall) needed to be flagged to @@ -196,36 +191,31 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, xant_flag_thre Default: None bl_len_range : tuple, optional - len-2 tuple containing minimum baseline length and maximum baseline length [meters] - to keep in baseline type selection + len-2 tuple containing minimum baseline length and maximum baseline + length [meters] to keep in baseline type selection bl_deg_range : tuple, optional - len-2 tuple containing (minimum, maximum) baseline angle in degrees to keep in - baseline selection + len-2 tuple containing (minimum, maximum) baseline angle in degrees + to keep in baseline selection Returns ------- - baselines1 : list of baseline tuples - Contains list of baseline tuples that should be fed as first argument - to PSpecData.pspec(), corresponding to uvd1 - - baselines2 : list of baseline tuples - Contains list of baseline tuples that should be fed as second argument - to PSpecData.pspec(), corresponding to uvd2 + baselines1, baselines2 : lists of baseline tuples + Lists of baseline tuples that should be fed as first/second argument + to PSpecData.pspec(), corresponding to uvd1/uvd2 blpairs : list of baseline-pair tuples Contains the baseline-pair tuples. i.e. zip(baselines1, baselines2) - xants1 : list of bad antenna integers for uvd1 - - xants2 : list of bad antenna integers for uvd2 + xants1, xants2 : lists + List of bad antenna integers for uvd1 and uvd2 """ # get antenna positions antpos1, ants1 = uvd1.get_ENU_antpos(pick_data_ants=False) - antpos1 = dict(zip(ants1, antpos1)) + antpos1 = dict(list(zip(ants1, antpos1))) antpos2, ants2 = uvd2.get_ENU_antpos(pick_data_ants=False) - antpos2 = dict(zip(ants2, antpos2)) - antpos = dict(antpos1.items() + antpos2.items()) + antpos2 = dict(list(zip(ants2, antpos2))) + antpos = dict(list(antpos1.items()) + list(antpos2.items())) # assert antenna positions match for a in set(antpos1).union(set(antpos2)): @@ -251,7 +241,7 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, xant_flag_thre # get flags f1 = uvd1.get_flags(bl) # remove from bad list if unflagged data exists - if np.sum(f1) < reduce(operator.mul, f1.shape) * xant_flag_thresh: + if np.sum(f1) < np.prod(f1.shape) * xant_flag_thresh: if antnums[0] in xants1: xants1.remove(antnums[0]) if antnums[1] in xants2: @@ -262,7 +252,7 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, xant_flag_thre # get flags f2 = uvd2.get_flags(bl) # remove from bad list if unflagged data exists - if np.sum(f2) < reduce(operator.mul, f2.shape) * xant_flag_thresh: + if np.sum(f2) < np.prod(f2.shape) * xant_flag_thresh: if antnums[0] in xants2: xants2.remove(antnums[0]) if antnums[1] in xants2: @@ -279,8 +269,9 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, xant_flag_thre baselines1, baselines2, blpairs = [], [], [] for r in reds: (bls1, bls2, - blps) = construct_blpairs(r, exclude_auto_bls=exclude_auto_bls, group=False, - exclude_permutations=exclude_permutations) + blps) = construct_blpairs(r, exclude_auto_bls=exclude_auto_bls, + group=False, + exclude_permutations=exclude_permutations) if len(bls1) < 1: continue @@ -297,14 +288,17 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, xant_flag_thre _bls1.append(bl1) _bls2.append(bl2) bls1, bls2 = _bls1, _bls2 - blps = zip(bls1, bls2) + blps = list(zip(bls1, bls2)) # group if desired if Nblps_per_group is not None: Ngrps = int(np.ceil(float(len(blps)) / Nblps_per_group)) - bls1 = [bls1[Nblps_per_group*i:Nblps_per_group*(i+1)] for i in range(Ngrps)] - bls2 = [bls2[Nblps_per_group*i:Nblps_per_group*(i+1)] for i in range(Ngrps)] - blps = [blps[Nblps_per_group*i:Nblps_per_group*(i+1)] for i in range(Ngrps)] + bls1 = [bls1[Nblps_per_group*i:Nblps_per_group*(i+1)] + for i in range(Ngrps)] + bls2 = [bls2[Nblps_per_group*i:Nblps_per_group*(i+1)] + for i in range(Ngrps)] + blps = [blps[Nblps_per_group*i:Nblps_per_group*(i+1)] + for i in range(Ngrps)] baselines1.extend(bls1) baselines2.extend(bls2) @@ -661,12 +655,15 @@ def config_pspec_blpairs(uv_templates, pol_pairs, group_pairs, exclude_auto_bls= if xants is not None: bls1, bls2 = [], [] for bl1, bl2 in zip(_bls1, _bls2): - if bl1[0] not in xants and bl1[1] not in xants and bl2[0] not in xants and bl2[1] not in xants: + if bl1[0] not in xants \ + and bl1[1] not in xants \ + and bl2[0] not in xants \ + and bl2[1] not in xants: bls1.append(bl1) bls2.append(bl2) else: bls1, bls2 = _bls1, _bls2 - blps = zip(bls1, bls2) + blps = list(zip(bls1, bls2)) # iterate over pol-group pairs that exist groupings = odict() @@ -715,13 +712,15 @@ def get_blvec_reds(blvecs, bl_error_tol=1.0, match_bl_lens=False): """ from hera_pspec import UVPSpec # type check - assert isinstance(blvecs, (dict, odict, UVPSpec)), "blpairs must be fed as a dict or UVPSpec" + assert isinstance(blvecs, (dict, odict, UVPSpec)), \ + "blpairs must be fed as a dict or UVPSpec" if isinstance(blvecs, UVPSpec): # get baseline vectors uvp = blvecs bls = uvp.bl_array bl_vecs = uvp.get_ENU_bl_vecs()[:, :2] - blvecs = dict(zip(map(uvp.bl_to_antnums, bls), bl_vecs)) + blvecs = dict(list(zip( [uvp.bl_to_antnums(_bls) for _bls in bls], + bl_vecs ))) # get baseline-pairs blpairs = uvp.get_blpairs() # form dictionary @@ -820,7 +819,7 @@ def job_monitor(run_func, iterator, action_name, M=map, lf=None, maxiter=1, t_start = time.time() # run function over jobs - exit_codes = np.array(M(run_func, iterator)) + exit_codes = np.array(list(M(run_func, iterator))) tnow = datetime.utcnow() # check for len-0 @@ -844,7 +843,7 @@ def job_monitor(run_func, iterator, action_name, M=map, lf=None, maxiter=1, # break after certain number of tries break # re-run function over jobs that failed - exit_codes = np.array(M(run_func, failures)) + exit_codes = np.array(list(M(run_func, failures))) # update counter counter += 1 # update failures @@ -956,7 +955,7 @@ def get_reds(uvd, bl_error_tol=1.0, pick_data_ants=False, bl_len_range=(0, 1e4), uvd = _uvd # get antenna position dictionary antpos, ants = uvd.get_ENU_antpos(pick_data_ants=pick_data_ants) - antpos_dict = dict(zip(ants, antpos)) + antpos_dict = dict(list(zip(ants, antpos))) # use antenna position dictionary elif isinstance(uvd, (dict, odict)): @@ -972,7 +971,7 @@ def get_reds(uvd, bl_error_tol=1.0, pick_data_ants=False, bl_len_range=(0, 1e4), # put in autocorrs if add_autos: ants = antpos_dict.keys() - reds = [zip(ants, ants)] + reds + reds = [list(zip(ants, ants))] + reds lens = np.insert(lens, 0, 0) angs = np.insert(angs, 0, 0) diff --git a/hera_pspec/uvpspec.py b/hera_pspec/uvpspec.py index da6e5081..6d2c86f4 100644 --- a/hera_pspec/uvpspec.py +++ b/hera_pspec/uvpspec.py @@ -53,8 +53,8 @@ def __init__(self): self._spw_freq_array = PSpecParam("spw_freq_array", description="Spw integer array for the freq_array.", form="(Nspwfreqs,)", expected_type=np.uint16) self._freq_array = PSpecParam("freq_array", description="Frequency array of the original data in Hz.", form="(Nfreqs,)", expected_type=np.float64) self._dly_array = PSpecParam("dly_array", description="Delay array in seconds.", form="(Nspwdlys,)", expected_type=np.float64) - desc = "Polarization integers of power spectra. Stokes 1:4 (I,Q,U,V); circular -1:-4 (RR,LL,RL,LR); linear -5:-8 (XX,YY,XY,YX)" - self._pol_array = PSpecParam("pol_array", description=desc, form="(Npols,)", expected_type=np.int32) + desc = "Polarization pair integer, made up of two polarization integers concatenated in a standardized way." + self._polpair_array = PSpecParam("polpair_array", description=desc, form="(Npols,)", expected_type=np.int32) self._lst_1_array = PSpecParam("lst_1_array", description="LST array of the first bl in the bl-pair [radians].", form="(Nblpairts,)", expected_type=np.float64) self._lst_2_array = PSpecParam("lst_2_array", description="LST array of the second bl in the bl-pair [radians].", form="(Nblpairts,)", expected_type=np.float64) self._lst_avg_array = PSpecParam("lst_avg_array", description="Average of the lst_1_array and lst_2_array [radians].", form="(Nblpairts,)", expected_type=np.float64) @@ -88,8 +88,8 @@ def __init__(self): self._cosmo = PSpecParam("cosmo", description="Instance of conversion.Cosmo_Conversions class.", expected_type=conversions.Cosmo_Conversions) # Collect all parameters: required and non-required - self._all_params = sorted(map(lambda p: p[1:], - fnmatch.filter(self.__dict__.keys(), '_*'))) + self._all_params = sorted( [ p[1:] for p in + fnmatch.filter(self.__dict__.keys(), '_*')]) # Specify required params: these are required for read / write and # self.check() @@ -97,7 +97,8 @@ def __init__(self): "Nspws", "Ndlys", "Npols", "Nfreqs", "history", "Nspwdlys", "Nspwfreqs", "data_array", "wgt_array", "integration_array", - "spw_array", "freq_array", "dly_array", "pol_array", + "spw_array", "freq_array", "dly_array", + "polpair_array", "lst_1_array", "lst_2_array", "time_1_array", "time_2_array", "blpair_array", "Nbls", "bl_vecs", "bl_array", "channel_width", @@ -110,12 +111,14 @@ def __init__(self): # All parameters must fall into one and only one of the following # groups, which are used in __eq__ self._immutables = ["Ntimes", "Nblpairts", "Nblpairs", "Nspwdlys", - "Nspwfreqs", "Nspws", "Ndlys", "Npols", "Nfreqs", "history", + "Nspwfreqs", "Nspws", "Ndlys", "Npols", "Nfreqs", + "history", "Nbls", "channel_width", "weighting", "vis_units", "norm", "norm_units", "taper", "cosmo", "beamfile", 'folded'] - self._ndarrays = ["spw_array", "freq_array", "dly_array", "pol_array", - "lst_1_array", "lst_avg_array", "time_avg_array", + self._ndarrays = ["spw_array", "freq_array", "dly_array", + "polpair_array", "lst_1_array", "lst_avg_array", + "time_avg_array", "lst_2_array", "time_1_array", "time_2_array", "blpair_array", "OmegaP", "OmegaPP", "beam_freqs", "bl_vecs", "bl_array", "telescope_location", @@ -125,17 +128,22 @@ def __init__(self): "nsample_array", "cov_array"] self._dicts_of_dicts = ["stats_array"] - # define which attributes are considred meta data. Large attrs should be constructed as datasets + # define which attributes are considred meta data. Large attrs should + # be constructed as datasets self._meta_dsets = ["lst_1_array", "lst_2_array", "time_1_array", "time_2_array", "blpair_array", "bl_vecs", "bl_array", "lst_avg_array", "time_avg_array", "OmegaP", "OmegaPP", "label_1_array", "label_2_array", "spw_dly_array", "spw_freq_array"] - self._meta_attrs = sorted(set(self._all_params) - set(self._dicts) - set(self._meta_dsets) - set(self._dicts_of_dicts)) + self._meta_attrs = sorted( + set(self._all_params) - set(self._dicts) \ + - set(self._meta_dsets) - set(self._dicts_of_dicts)) self._meta = sorted(set(self._meta_dsets).union(set(self._meta_attrs))) # check all params are covered - assert len(set(self._all_params) - set(self._dicts) - set(self._immutables) - set(self._ndarrays) - set(self._dicts_of_dicts)) == 0 + assert len( set(self._all_params) - set(self._dicts) \ + - set(self._immutables) - set(self._ndarrays) \ + - set(self._dicts_of_dicts) ) == 0 # Default parameter values self.folded = False @@ -143,19 +151,20 @@ def __init__(self): def get_cov(self, key, omit_flags=False): """ Slice into covariance array with a specified data key in the format - (spw, ((ant1, ant2),(ant3, ant4)), pol) + (spw, ((ant1, ant2),(ant3, ant4)), (pol12, pol34)) or - (spw, blpair-integer, pol) + (spw, blpair-integer, pol-integer) - where spw is the spectral window integer, ant1 etc. arge integers, - and pol is either a polarization string (ex. 'XX') or integer (ex. -5). + where spw is the spectral window integer, ant1 etc. are integers, + and pol is either a tuple of polarization strings (ex. ('XX', 'YY')) + or integers (ex. (-5,-5)). Parameters ---------- key: tuple - Baseline-pair key + Contains the baseline-pair, spw, polpair keys omit_flags : bool, optional If True, remove time integrations (or spectra) that @@ -167,15 +176,16 @@ def get_cov(self, key, omit_flags=False): data : complex ndarray Shape (Ntimes, Ndlys, Ndlys) """ - spw, blpairts, pol = self.key_to_indices(key, omit_flags=omit_flags) + spw, blpairts, polpair = self.key_to_indices(key, omit_flags=omit_flags) + # Need to deal with folded data! # if data has been folded, return only positive delays if hasattr(self,'cov_array'): if self.folded: Ndlys = self.data_array[spw].shape[1] - return self.cov_array[spw][blpairts, Ndlys//2+1:, Ndlys//2+1:, pol] + return self.cov_array[spw][blpairts, Ndlys//2+1:, Ndlys//2+1:, polpair] else: - return self.cov_array[spw][blpairts, :, :, pol] + return self.cov_array[spw][blpairts, :, :, polpair] else: raise AttributeError("No covariance array has been calculated.") @@ -183,19 +193,20 @@ def get_data(self, key, omit_flags=False): """ Slice into data_array with a specified data key in the format - (spw, ((ant1, ant2), (ant3, ant4)), pol) + (spw, ((ant1, ant2), (ant3, ant4)), (pol12, pol34)) or - (spw, blpair-integer, pol) + (spw, blpair-integer, polpair) where spw is the spectral window integer, ant1 etc. are integers, - and pol is either a polarization string (ex. 'XX') or integer (ex. -5). + and polpair is either a tuple of polarization strings/ints or a + polarizatio-pair integer. Parameters ---------- key : tuple - Baseline-pair key + Contains the baseline-pair, spw, polpair keys omit_flags : bool, optional If True, remove time integrations (or spectra) that @@ -207,34 +218,34 @@ def get_data(self, key, omit_flags=False): data : complex ndarray Shape (Ntimes, Ndlys) """ - spw, blpairts, pol = self.key_to_indices(key, omit_flags=omit_flags) + spw, blpairts, polpair = self.key_to_indices(key, omit_flags=omit_flags) # if data has been folded, return only positive delays if self.folded: Ndlys = self.data_array[spw].shape[1] - return self.data_array[spw][blpairts, Ndlys//2+1:, pol] + return self.data_array[spw][blpairts, Ndlys//2+1:, polpair] # else return all delays else: - return self.data_array[spw][blpairts, :, pol] + return self.data_array[spw][blpairts, :, polpair] def get_wgts(self, key, omit_flags=False): """ Slice into wgt_array with a specified data key in the format - (spw, ((ant1, ant2), (ant3, ant4)), pol) + (spw, ((ant1, ant2), (ant3, ant4)), (pol12, pol34)) or - (spw, blpair-integer, pol) + (spw, blpair-integer, polpair-integer) where spw is the spectral window integer, ant1 etc. are integers, - and pol is either a polarization string (ex. 'XX') or integer (ex. -5). + and polpair is either a polarization pair tuple or integer. Parameters ---------- key : tuple - Contains the baseline-pair key + Contains the baseline-pair, spw, polpair keys omit_flags : bool, optional If True, remove time integrations (or spectra) that @@ -247,27 +258,27 @@ def get_wgts(self, key, omit_flags=False): Has shape (2, Ntimes, Ndlys), where the zeroth axis holds [wgt_1, wgt_2] in that order. """ - spw, blpairts, pol = self.key_to_indices(key, omit_flags=omit_flags) + spw, blpairts, polpair = self.key_to_indices(key, omit_flags=omit_flags) - return self.wgt_array[spw][blpairts, :, :, pol] + return self.wgt_array[spw][blpairts, :, :, polpair] def get_integrations(self, key, omit_flags=False): """ Slice into integration_array with a specified data key in the format:: - (spw, ((ant1, ant2), (ant3, ant4)), pol) + (spw, ((ant1, ant2), (ant3, ant4)), (pol12, pol34)) - or:: + or - (spw, blpair-integer, pol) + (spw, blpair-integer, polpair-integer) where spw is the spectral window integer, ant1 etc. are integers, - and pol is either a polarization string (ex. 'XX') or integer (ex. -5). + and polpair is either a polarization pair tuple or integer. Parameters ---------- key : tuple - Baseline-pair key + Contains the baseline-pair, spw, polpair keys omit_flags : bool, optional If True, remove time integrations (or spectra) that @@ -279,26 +290,27 @@ def get_integrations(self, key, omit_flags=False): data : float ndarray Has shape (Ntimes,) """ - spw, blpairts, pol = self.key_to_indices(key, omit_flags=omit_flags) + spw, blpairts, polpair = self.key_to_indices(key, omit_flags=omit_flags) - return self.integration_array[spw][blpairts, pol] + return self.integration_array[spw][blpairts, polpair] def get_nsamples(self, key, omit_flags=False): """ Slice into nsample_array with a specified data key in the format - (spw, ((ant1, ant2), (ant3, ant4)), pol) + (spw, ((ant1, ant2), (ant3, ant4)), (pol12, pol34)) or - (spw, blpair-integer, pol) + (spw, blpair-integer, polpair-integer) where spw is the spectral window integer, ant1 etc. are integers, - and pol is either a polarization string (ex. 'XX') or integer (ex. -5). + and polpair is either a polarization pair tuple or integer. Parameters ---------- - key : tuple, baseline-pair key + key : tuple + Contains the baseline-pair, spw, polpair keys omit_flags : bool, optional If True, remove time integrations (or spectra) that @@ -309,9 +321,9 @@ def get_nsamples(self, key, omit_flags=False): ------- data : float ndarray with shape (Ntimes,) """ - spw, blpairts, pol = self.key_to_indices(key, omit_flags=omit_flags) + spw, blpairts, polpair = self.key_to_indices(key, omit_flags=omit_flags) - return self.nsample_array[spw][blpairts, pol] + return self.nsample_array[spw][blpairts, polpair] def get_dlys(self, key): """ @@ -324,7 +336,7 @@ def get_dlys(self, key): Returns ------- - dlys : float ndarray, contains delays in nanosec of pspectra given spw + dlys : float ndarray, contains delays of pspectra in nanosec, given spw """ indices = self.spw_to_dly_indices(key) return self.dly_array[indices] @@ -342,7 +354,8 @@ def get_blpair_seps(self): # get list of bl separations bl_vecs = self.get_ENU_bl_vecs() bls = self.bl_array.tolist() - blseps = np.array(map(lambda bl: np.linalg.norm(bl_vecs[bls.index(bl)]), self.bl_array)) + blseps = np.array([np.linalg.norm(bl_vecs[bls.index(bl)]) + for bl in self.bl_array]) # construct empty blp_avg_sep array blp_avg_sep = np.empty(self.Nblpairts, np.float) @@ -378,11 +391,13 @@ def get_kperps(self, spw, little_h=True): k_perp : float ndarray, transverse wave-vectors, shape=(Nblpairs,) """ # assert cosmo - assert hasattr(self, 'cosmo'), "self.cosmo must exist to form cosmological " \ + assert hasattr(self, 'cosmo'), \ + "self.cosmo must exist to form cosmological " \ "wave-vectors. See self.set_cosmology()" # calculate mean redshift of band - avg_z = self.cosmo.f2z(np.mean(self.freq_array[self.spw_to_freq_indices(spw)])) + avg_z = self.cosmo.f2z( + np.mean(self.freq_array[self.spw_to_freq_indices(spw)]) ) # get kperps blpair_seps = self.get_blpair_seps() @@ -397,7 +412,8 @@ def get_kparas(self, spw, little_h=True): Parameters ---------- - spw : int, choice of spectral window + spw : int + Choice of spectral window little_h : boolean, optional Whether to have cosmological length units be h^-1 Mpc or Mpc @@ -405,31 +421,36 @@ def get_kparas(self, spw, little_h=True): Returns (k_perp, k_para) ------- - k_para : float ndarray, radial wave-vectors, shape=(Ndlys given spw) + k_para : float ndarray + Radial wave-vectors, shape=(Ndlys given spw) """ # assert cosmo - assert hasattr(self, 'cosmo'), "self.cosmo must exist to form cosmological " \ + assert hasattr(self, 'cosmo'), \ + "self.cosmo must exist to form cosmological " \ "wave-vectors. See self.set_cosmology()" # calculate mean redshift of band - avg_z = self.cosmo.f2z(np.mean(self.freq_array[self.spw_to_freq_indices(spw)])) + avg_z = self.cosmo.f2z( + np.mean(self.freq_array[self.spw_to_freq_indices(spw)]) ) # get kparas - k_para = self.get_dlys(spw) * self.cosmo.tau_to_kpara(avg_z, little_h=little_h) - + k_para = self.get_dlys(spw) \ + * self.cosmo.tau_to_kpara(avg_z, little_h=little_h) return k_para def get_blpairs(self): """ Returns a list of all blpair tuples in the data_array. """ - return [self.blpair_to_antnums(blp) for blp in np.unique(self.blpair_array)] + return [self.blpair_to_antnums(blp) + for blp in np.unique(self.blpair_array)] def get_all_keys(self): """ - Returns a list of all possible tuple keys in the data_array, in the format: + Returns a list of all possible tuple keys in the data_array, in the + format: - (spectral window, baseline-pair, polarization-string) + (spectral window, baseline-pair, polarization-string) """ # get unique blpair tuples blps = self.get_blpairs() @@ -438,8 +459,9 @@ def get_all_keys(self): # loop over spw and pols and add to keys for spw in range(self.Nspws): - for p in self.pol_array: - pstr = uvutils.polnum2str(p) + for p in self.polpair_array: + pol1, pol2 = uvputils.polpair_int2tuple(p) + pstr = (uvutils.polnum2str(pol1), uvutils.polnum2str(pol2)) all_keys.extend([(spw, blp, pstr) for blp in blps]) return all_keys @@ -467,6 +489,7 @@ def get_spw_ranges(self, spws=None): if spws is None: spws = np.arange(self.Nspws) spw_ranges = [] + # iterate over spectral windows for spw in spws: spw_freqs = self.freq_array[self.spw_to_freq_indices(spw)] @@ -483,11 +506,11 @@ def get_stats(self, stat, key, omit_flags=False): Parameters ---------- - stat : string + stat : str The statistic to return. key : tuple - A spw-blpair-pol key parseable by UVPSpec.key_to_indices. + A spw-blpair-polpair key parseable by UVPSpec.key_to_indices. omit_flags : bool, optional If True, remove time integrations (or spectra) that @@ -497,27 +520,28 @@ def get_stats(self, stat, key, omit_flags=False): Returns ------- statistic : array_like - A 2D (Ntimes, Ndlys) complex array holding desired bandpower statistic. + A 2D (Ntimes, Ndlys) complex array holding desired bandpower + statistic. """ if not hasattr(self, "stats_array"): raise AttributeError("No stats have been entered to this UVPSpec object") assert stat in self.stats_array.keys(), "Statistic name {} not found in stat keys.".format(stat) - spw, blpairts, pol = self.key_to_indices(key, omit_flags=omit_flags) + spw, blpairts, polpair = self.key_to_indices(key, omit_flags=omit_flags) statistic = self.stats_array[stat] # if bandpowers have been folded, return only positive delays if self.folded: Ndlys = statistic[spw].shape[1] - return statistic[spw][blpairts, Ndlys//2+1:, pol] + return statistic[spw][blpairts, Ndlys//2+1:, polpair] else: - return statistic[spw][blpairts, :, pol] + return statistic[spw][blpairts, :, polpair] def set_stats(self, stat, key, statistic): """ Set a statistic waterfall (Ntimes, Ndlys) in the stats_array - given the selection of spw, blpairts and pol in key. + given the selection of spw, blpairts and polpair in key. Parameters ---------- @@ -525,19 +549,20 @@ def set_stats(self, stat, key, statistic): Name of the statistic. key : tuple - A spw-blpair-pol key parseable by UVPSpec.key_to_indices. + A spw-blpair-polpair key parseable by UVPSpec.key_to_indices. statistic : ndarray 2D ndarray with statistics to set. Must conform to shape of the slice of data_array produced by the specified key. """ - spw, blpairts, pol = self.key_to_indices(key) + spw, blpairts, polpair = self.key_to_indices(key) statistic = np.asarray(statistic) - data_shape = self.data_array[spw][blpairts, :, pol].shape - + data_shape = self.data_array[spw][blpairts, :, polpair].shape + if data_shape != statistic.shape: - errmsg = "Input array shape {:s} must match " \ - "data_array shape {:s}.".format(statistic.shape, data_shape) + print(data_shape, statistic.shape) + errmsg = "Input array shape {} must match " \ + "data_array shape {}.".format(statistic.shape, data_shape) raise ValueError(errmsg) if not hasattr(self, "stats_array"): @@ -545,23 +570,29 @@ def set_stats(self, stat, key, statistic): dtype = statistic.dtype if stat not in self.stats_array.keys(): - self.stats_array[stat] = odict([[i, np.nan * np.ones(self.data_array[i].shape, dtype=dtype)] + self.stats_array[stat] = \ + odict([[i, + np.nan * np.ones(self.data_array[i].shape, dtype=dtype)] for i in range(self.Nspws)]) - self.stats_array[stat][spw][blpairts, :, pol] = statistic - + self.stats_array[stat][spw][blpairts, :, polpair] = statistic + + def convert_to_deltasq(self, little_h=True, inplace=True): """ Convert from P(k) to Delta^2(k) by multiplying by k^3 / (2pi^2). - The units of the output is therefore the current units (self.units) times h^3 Mpc^-3, - where the h^3 is only included if little_h == True. + The units of the output is therefore the current units (self.units) + times h^3 Mpc^-3, where the h^3 is only included if little_h == True. Parameters ---------- - inplace : boolean + little_h : bool, optional + Whether to use h^-1 units. Default: True. + + inplace : bool, optional If True edit and overwrite arrays in self, else make a copy of - self and return. + self and return. Default: True. """ # copy object if inplace: @@ -584,7 +615,8 @@ def convert_to_deltasq(self, little_h=True, inplace=True): if inplace == False: return uvp - + + def blpair_to_antnums(self, blpair): """ Convert baseline-pair integer to nested tuple of antenna numbers. @@ -601,7 +633,8 @@ def blpair_to_antnums(self, blpair): Ex. ((ant1, ant2), (ant3, ant4)) """ return uvputils._blpair_to_antnums(blpair) - + + def antnums_to_blpair(self, antnums): """ Convert nested tuple of antenna numbers to baseline-pair integer. @@ -618,10 +651,12 @@ def antnums_to_blpair(self, antnums): baseline-pair integer """ return uvputils._antnums_to_blpair(antnums) - + + def bl_to_antnums(self, bl): """ - Convert baseline (antenna-pair) integer to nested tuple of antenna numbers. + Convert baseline (antenna-pair) integer to nested tuple of antenna + numbers. Parameters ---------- @@ -634,7 +669,8 @@ def bl_to_antnums(self, bl): tuple containing baseline antenna numbers. Ex. (ant1, ant2) """ return uvputils._bl_to_antnums(bl) - + + def antnums_to_bl(self, antnums): """ Convert tuple of antenna numbers to baseline integer. @@ -651,7 +687,8 @@ def antnums_to_bl(self, antnums): Baseline integer. """ return uvputils._antnums_to_bl(antnums) - + + def blpair_to_indices(self, blpair): """ Convert a baseline-pair nested tuple ((ant1, ant2), (ant3, ant4)) or @@ -671,11 +708,14 @@ def blpair_to_indices(self, blpair): blpair = [blpair] elif isinstance(blpair, list): if isinstance(blpair[0], tuple): - blpair = map(lambda blp: self.antnums_to_blpair(blp), blpair) + blpair = [self.antnums_to_blpair(blp) for blp in blpair] # assert exists in data - assert np.array(map(lambda b: b in self.blpair_array, blpair)).all(), "blpairs {} not all found in data".format(blpair) - return np.arange(self.Nblpairts)[reduce(operator.add, map(lambda b: self.blpair_array == b, blpair))] - + assert np.array([b in self.blpair_array for b in blpair]).all(), \ + "blpairs {} not all found in data".format(blpair) + return np.arange(self.Nblpairts)[ + np.logical_or.reduce([self.blpair_array == b for b in blpair])] + + def spw_to_freq_indices(self, spw): """ Convert a spectral window integer into a list of indices to index @@ -684,7 +724,8 @@ def spw_to_freq_indices(self, spw): Parameters ---------- spw : int or tuple - Spectral window index, or spw tuple from get_spw_ranges(), or list of either. + Spectral window index, or spw tuple from get_spw_ranges(), or list + of either. Returns ------- @@ -701,16 +742,17 @@ def spw_to_freq_indices(self, spw): spw = [spw_ranges.index(s) for s in spw] # assert exists in data - assert np.array(map(lambda s: s in self.spw_freq_array, spw)).all(), "spws {} not all found in data".format(spw) + assert np.array([s in self.spw_freq_array for s in spw]).all(), \ + "spws {} not all found in data".format(spw) # get select boolean array - select = reduce(operator.add, map(lambda s: self.spw_freq_array == s, spw)) - + select = np.logical_or.reduce([self.spw_freq_array == s for s in spw]) + # get indices freq_indices = np.arange(self.Nspwfreqs)[select] - return freq_indices - + + def spw_to_dly_indices(self, spw): """ Convert a spectral window integer into a list of indices to index @@ -719,7 +761,8 @@ def spw_to_dly_indices(self, spw): Parameters ---------- spw : int or tuple - Spectral window index, or spw tuple from get_spw_ranges(), or list of either. + Spectral window index, or spw tuple from get_spw_ranges(), or list + of either. Returns ------- @@ -736,11 +779,12 @@ def spw_to_dly_indices(self, spw): spw = [spw_ranges.index(s) for s in spw] # assert exists in data - assert np.array(map(lambda s: s in self.spw_dly_array, spw)).all(), "spws {} not all found in data".format(spw) + assert np.array([s in self.spw_dly_array for s in spw]).all(), \ + "spws {} not all found in data".format(spw) # get select boolean array - select = reduce(operator.add, map(lambda s: self.spw_dly_array == s, spw)) - + select = np.logical_or.reduce([self.spw_dly_array == s for s in spw]) + if self.folded: select[self.dly_array < 1e-10] = False @@ -757,7 +801,8 @@ def spw_indices(self, spw): Parameters ---------- spw : int or tuple - Spectral window index, or spw tuple from get_spw_ranges(), or list of either. + Spectral window index, or spw tuple from get_spw_ranges(), or list + of either. Returns ------- @@ -774,47 +819,52 @@ def spw_indices(self, spw): spw = [spw_ranges.index(s) for s in spw] # assert exists in data - assert np.array(map(lambda s: s in self.spw_array, spw)).all(), "spws {} not all found in data".format(spw) + assert np.array([s in self.spw_array for s in spw]).all(), \ + "spws {} not all found in data".format(spw) # get select boolean array - select = reduce(operator.add, map(lambda s: self.spw_array == s, spw)) - + #select = reduce(operator.add, [self.spw_array == s for s in spw]) + select = np.logical_or.reduce([self.spw_array == s for s in spw]) + # get array spw_indices = np.arange(self.Nspws)[select] return spw_indices - - def pol_to_indices(self, pol): + + + def polpair_to_indices(self, polpair): """ - Map a polarization integer or str to its index in pol_array. + Map a polarization-pair integer or tuple to its index in polpair_array. Parameters ---------- - pol : str or int - Polarization string (ex. 'XX') or integer (ex. -5), or a list of - strs or ints. + polpair : (list of) int or tuple + Polarization-pair integer or tuple of polarization strings/ints. Returns ------- indices : int - Index of pol in pol_array. - """ - # convert pol to int if str - if isinstance(pol, (str, np.str)): - pol = [uvutils.polstr2num(pol)] - elif isinstance(pol, (int, np.integer)): - pol = [pol] - elif isinstance(pol, (list, tuple)): - for i in range(len(pol)): - if isinstance(pol[i], (np.str, str)): - pol[i] = uvutils.polstr2num(pol[i]) - - # ensure all pols exist in data - assert np.array(map(lambda p: p in self.pol_array, pol)).all(), "pols {} not all found in data".format(pol) - - indices = np.arange(self.Npols)[reduce(operator.add, map(lambda p: self.pol_array == p, pol))] + Index of polpair in polpair_array. + """ + # Convert to list if not already a list + if isinstance(polpair, (tuple, int, np.integer)): + polpair = [polpair,] + elif not isinstance(polpair, list): + raise TypeError("polpair must be list of tuple or int") + + # Convert list items to standard format (polpair integers) + polpair = [uvputils.polpair_tuple2int(p) if isinstance(p, tuple) else p + for p in polpair] + + # Ensure all pols exist in data + assert np.array([p in self.polpair_array for p in polpair]).all(), \ + "pols {} not all found in data".format(polpair) + + idxs = np.logical_or.reduce([self.polpair_array == pp for pp in polpair]) + indices = np.arange(self.polpair_array.size)[idxs] return indices - + + def time_to_indices(self, time, blpairs=None): """ Convert a time [Julian Date] from self.time_avg_array to an array that @@ -853,21 +903,21 @@ def key_to_indices(self, key, omit_flags=False): """ Convert a data key into relevant slice arrays. A data key takes the form - (spw_integer, ((ant1, ant2), (ant3, ant4)), pol_string) + (spw_integer, ((ant1, ant2), (ant3, ant4)), (pol12, pol34)) or - (spw, blpair-integer, pol) + (spw, blpair-integer, pol-integer) where spw is the spectral window integer, ant1 etc. are integers, - and pol is either a polarization string (ex. 'XX') or integer (ex. -5). + and polpair is either a polarization-pair integer or tuple. The key can also be a dictionary in the form:: key = { 'spw' : spw_integer, 'blpair' : ((ant1, ant2), (ant3, ant4)) - 'pol' : pol_string + 'polpair' : (pol12, pol34) } and it will parse the dictionary for you. @@ -875,7 +925,7 @@ def key_to_indices(self, key, omit_flags=False): Parameters ---------- key : tuple - Baseline-pair key. + Baseline-pair, spw, and pol-pair key. omit_flags : bool, optional If True, remove time integrations (or spectra) that @@ -890,56 +940,64 @@ def key_to_indices(self, key, omit_flags=False): blpairts : list List of integers to apply along blpairts axis. - pol : int - Polarization index. + polpair : int + Polarization pair index. """ # assert key length - assert len(key) == 3, "length of key must be 3: (spw, blpair, pol)" + assert len(key) == 3, "length of key must be 3: (spw, blpair, polpair)" if isinstance(key, (odict, dict)): - key = (key['spw'], key['blpair'], key['pol']) + key = (key['spw'], key['blpair'], key['polpair']) # assign key elements spw_ind = key[0] blpair = key[1] - pol = key[2] + polpair = key[2] # assert types assert isinstance(spw_ind, (int, np.integer)), "spw must be an integer" - assert isinstance(blpair, (int, np.integer, tuple)), "blpair must be an integer or nested tuple" - assert isinstance(pol, (np.str, str, np.integer, int)), "pol must be a string or integer" + assert isinstance(blpair, (int, np.integer, tuple)), \ + "blpair must be an integer or nested tuple" + assert isinstance(polpair, (tuple, np.integer, int)), \ + "polpair must be a tuple or integer: %s / %s" % (polpair, key) # convert blpair to int if not int if type(blpair) == tuple: blpair = self.antnums_to_blpair(blpair) - # convert pol to int if str - if type(pol) in (str, np.str): - pol = uvutils.polstr2num(pol) + # convert pol to int if not int + if type(polpair) == tuple: + polpair = uvputils.polpair_tuple2int(polpair) # check attributes exist in data - assert spw_ind in self.spw_freq_array and spw_ind in self.spw_dly_array, "spw {} not found in data".format(spw_ind) - assert blpair in self.blpair_array, "blpair {} not found in data".format(blpair) - assert pol in self.pol_array, "pol {} not found in data".format(pol) + assert spw_ind in self.spw_freq_array and spw_ind in self.spw_dly_array, \ + "spw {} not found in data".format(spw_ind) + assert blpair in self.blpair_array, \ + "blpair {} not found in data, blpairs are: {}".format(blpair, self.blpair_array) + assert polpair in self.polpair_array, \ + "polpair {} not found in data".format(polpair) # index polarization array - pol_ind = self.pol_to_indices(pol) - if isinstance(pol_ind, (list, np.ndarray)): - assert len(pol_ind) == 1, "only one polarization can be specified in key_to_indices" - pol_ind = pol_ind[0] + polpair_ind = self.polpair_to_indices(polpair) + if isinstance(polpair_ind, (list, np.ndarray)): + assert len(polpair_ind) == 1, \ + "only one polpair can be specified in key_to_indices" + polpair_ind = polpair_ind[0] # index blpairts blpairts_inds = self.blpair_to_indices(blpair) # omit flagged spectra: i.e. when integration_array == 0.0 if omit_flags: - integs = self.integration_array[spw_ind][blpairts_inds, pol_ind] + integs = self.integration_array[spw_ind][blpairts_inds, polpair_ind] keep = ~np.isclose(integs, 0.0) blpairts_inds = blpairts_inds[keep] - return spw_ind, blpairts_inds, pol_ind + return spw_ind, blpairts_inds, polpair_ind + def select(self, spws=None, bls=None, only_pairs_in_bls=True, blpairs=None, - times=None, lsts=None, pols=None, inplace=True, run_check=True): + times=None, lsts=None, polpairs=None, inplace=True, + run_check=True): """ Select function for selecting out certain slices of the data. @@ -973,10 +1031,9 @@ def select(self, spws=None, bls=None, only_pairs_in_bls=True, blpairs=None, Float ndarray of lsts from the lst_avg_array to select. If None, all lsts are kept. Default: None. - pols : list of str or int, optional - List of polarizations to keep. See pyuvdata.utils.polstr2num for - acceptable options. If None, all polarizations are kept. Default: - None. + polpairs : list of tuple or int, optional + List of polarization-pairs to keep. If None, all polarizations + are kept. Default: None. inplace : bool, optional If True, edit and overwrite arrays in self, else make a copy of @@ -998,14 +1055,13 @@ def select(self, spws=None, bls=None, only_pairs_in_bls=True, blpairs=None, uvputils._select(uvp, spws=spws, bls=bls, only_pairs_in_bls=only_pairs_in_bls, - blpairs=blpairs, times=times, lsts=lsts, pols=pols) - - if run_check: - uvp.check() - - if inplace == False: - return uvp + blpairs=blpairs, times=times, lsts=lsts, + polpairs=polpairs) + if run_check: uvp.check() + if inplace == False: return uvp + + def get_ENU_bl_vecs(self): """ Return baseline vector array in TOPO (ENU) frame in meters, with @@ -1018,10 +1074,11 @@ def get_ENU_bl_vecs(self): """ return uvutils.ENU_from_ECEF( (self.bl_vecs + self.telescope_location), \ - *uvutils.LatLonAlt_from_XYZ(self.telescope_location[None])) - + *uvutils.LatLonAlt_from_XYZ(self.telescope_location[None])) + + def read_from_group(self, grp, just_meta=False, spws=None, bls=None, - blpairs=None, times=None, lsts=None, pols=None, + blpairs=None, times=None, lsts=None, polpairs=None, only_pairs_in_bls=False): """ Clear current UVPSpec object and load in data from specified HDF5 group. @@ -1055,9 +1112,8 @@ def read_from_group(self, grp, just_meta=False, spws=None, bls=None, lsts : float ndarray lsts from the lst_avg_array to keep. - pols : list of str or int - List of polarization strings or integers to keep. See - pyuvdata.utils.polstr2num for acceptable options. + polpairs : list of tuple or int + List of polarization-pair integers or tuples to keep. only_pairs_in_bls : bool, optional If True, keep only baseline-pairs whose first _and_ second baseline @@ -1065,7 +1121,8 @@ def read_from_group(self, grp, just_meta=False, spws=None, bls=None, """ # Make sure the group is a UVPSpec object assert 'pspec_type' in grp.attrs, "This object is not a UVPSpec object" - assert grp.attrs['pspec_type'] == 'UVPSpec', "This object is not a UVPSpec object" + assert grp.attrs['pspec_type'] == 'UVPSpec', \ + "This object is not a UVPSpec object" # Clear all data in the current object self._clear() @@ -1082,21 +1139,22 @@ def read_from_group(self, grp, just_meta=False, spws=None, bls=None, if just_meta: uvputils._select(self, spws=spws, bls=bls, lsts=lsts, only_pairs_in_bls=only_pairs_in_bls, - blpairs=blpairs, times=times, pols=pols) + blpairs=blpairs, times=times, polpairs=polpairs) else: uvputils._select(self, spws=spws, bls=bls, lsts=lsts, only_pairs_in_bls=only_pairs_in_bls, - blpairs=blpairs, times=times, pols=pols, + blpairs=blpairs, times=times, polpairs=polpairs, h5file=grp) # handle cosmo if hasattr(self, 'cosmo'): - self.cosmo = conversions.Cosmo_Conversions(**ast.literal_eval(self.cosmo)) + self.cosmo = conversions.Cosmo_Conversions( + **ast.literal_eval(self.cosmo) ) self.check(just_meta=just_meta) def read_hdf5(self, filepath, just_meta=False, spws=None, bls=None, - blpairs=None, times=None, lsts=None, pols=None, + blpairs=None, times=None, lsts=None, polpairs=None, only_pairs_in_bls=False): """ Clear current UVPSpec object and load in data from an HDF5 file. @@ -1130,9 +1188,8 @@ def read_hdf5(self, filepath, just_meta=False, spws=None, bls=None, lsts : float ndarray lsts from the lst_avg_array to keep. - pols : list of str or int - List of polarization strings or integers to keep. See - pyuvdata.utils.polstr2num for acceptable options. + polpairs : list of polpair ints or tuples + List of polarization-pair integers or tuples to keep. only_pairs_in_bls : bool, optional If True, keep only baseline-pairs whose first _and_ second baseline @@ -1140,9 +1197,11 @@ def read_hdf5(self, filepath, just_meta=False, spws=None, bls=None, """ # Open file descriptor and read data with h5py.File(filepath, 'r') as f: - self.read_from_group(f, just_meta=just_meta, spws=spws, bls=bls, times=times, - lsts=lsts, only_pairs_in_bls=only_pairs_in_bls) - + self.read_from_group(f, just_meta=just_meta, spws=spws, bls=bls, + times=times, lsts=lsts, polpairs=polpairs, + only_pairs_in_bls=only_pairs_in_bls) + + def write_to_group(self, group, run_check=True): """ Write UVPSpec data into an HDF5 group. @@ -1171,7 +1230,19 @@ def write_to_group(self, group, run_check=True): if k == 'cosmo': group.attrs['cosmo'] = str(getattr(self, k).get_params()) continue - group.attrs[k] = getattr(self, k) + this_attr = getattr(self, k) + + # Do Unicode/string conversions, as HDF5 struggles with them + if k == 'labels': + this_attr = [np.string_(lbl) for lbl in this_attr] + + # Store attribute in group + try: + group.attrs[k] = this_attr + except(TypeError): + print("TypeError:", k, this_attr) + continue + for k in self._meta_dsets: if hasattr(self, k): group.create_dataset(k, data=getattr(self, k)) @@ -1206,7 +1277,8 @@ def write_to_group(self, group, run_check=True): # denote as a uvpspec object group.attrs['pspec_type'] = self.__class__.__name__ - + + def write_hdf5(self, filepath, overwrite=False, run_check=True): """ Write a UVPSpec object to HDF5 file. @@ -1232,7 +1304,8 @@ def write_hdf5(self, filepath, overwrite=False, run_check=True): # Write file with h5py.File(filepath, 'w') as f: self.write_to_group(f, run_check=run_check) - + + def set_cosmology(self, new_cosmo, overwrite=False, new_beam=None, verbose=True): """ @@ -1293,7 +1366,7 @@ def set_cosmology(self, new_cosmo, overwrite=False, new_beam=None, # PSpecBeamUV will adopt a default cosmology upon instantiation, # but this doesn't matterfor what we need from it new_beam = pspecbeam.PSpecBeamUV(new_beam) - self.OmegaP, self.OmegaPP = new_beam.get_Omegas(self.pol_array) + self.OmegaP, self.OmegaPP = new_beam.get_Omegas(self.polpair_array) self.beam_freqs = new_beam.beam_freqs # Update cosmo @@ -1306,12 +1379,12 @@ def set_cosmology(self, new_cosmo, overwrite=False, new_beam=None, if verbose: print("Updating scalar array and re-normalizing power spectra") for spw in range(self.Nspws): - for j, pol in enumerate(self.pol_array): - scalar = self.compute_scalar(spw, pol, num_steps=1000, + for j, polpair in enumerate(self.polpair_array): + scalar = self.compute_scalar(spw, polpair, num_steps=1000, little_h=True, noise_scalar=False) # renormalize power spectra with new scalar - self.data_array[spw][:, :, j] *= scalar / self.scalar_array[spw, j] + self.data_array[spw][:, :, j] *= scalar/self.scalar_array[spw, j] # update self.scalar_array element self.scalar_array[spw, j] = scalar @@ -1319,7 +1392,8 @@ def set_cosmology(self, new_cosmo, overwrite=False, new_beam=None, # Update self.units if pspectra were not originally in cosmological units if "Mpc" not in self.norm_units: self.norm_units = "h^-3 Mpc^3" - + + def check(self, just_meta=False): """ Run checks to make sure metadata and data arrays are properly defined @@ -1332,33 +1406,41 @@ def check(self, just_meta=False): """ # iterate over all possible parameters for p in self._all_params: + # only enforce existance if not just_meta if not just_meta: if p in self._req_params: - assert hasattr(self, p), "required parameter {} doesn't exist".format(p) + assert hasattr(self, p), \ + "required parameter {} doesn't exist".format(p) # if attribute exists, check its type if hasattr(self, p): - a = getattr(self, '_'+p) + a = getattr(self, '_' + p) if hasattr(a, 'expected_type'): err_msg = "attribute {} does not have expected data type {}".format(p, a.expected_type) + # ndarrays if p in self._ndarrays: - assert isinstance(getattr(self, p), np.ndarray), "attribute {} needs to be an ndarray".format(p) - if issubclass(getattr(self, p).dtype.type, a.expected_type): + assert isinstance(getattr(self, p), np.ndarray), \ + "attribute {} needs to be an ndarray, {}".format(p, getattr(self, p).tostring()) + if issubclass(getattr(self, p).dtype.type, + a.expected_type): pass else: # try to cast into its dtype try: - setattr(self, p, getattr(self, p).astype(a.expected_type)) + setattr(self, p, + getattr(self, p).astype(a.expected_type)) except: raise ValueError(err_msg) # dicts elif p in self._dicts: - assert isinstance(getattr(self, p), (dict, odict)), "attribute {} needs to be a dictionary".format(p) + assert isinstance(getattr(self, p), (dict, odict)), \ + "attribute {} needs to be a dictionary".format(p) # iterate over keys for k in getattr(self, p).keys(): - assert isinstance(getattr(self, p)[k], np.ndarray), "values of attribute {} need to be ndarrays".format(p) + assert isinstance(getattr(self, p)[k], np.ndarray), \ + "values of attribute {} need to be ndarrays".format(p) assert issubclass(getattr(self, p)[k].dtype.type, a.expected_type), err_msg # immutables elif p in self._immutables: @@ -1381,7 +1463,8 @@ def check(self, just_meta=False): raise AssertionError(err_msg) # check spw convention assert set(self.spw_array) == set(np.arange(self.Nspws)), "spw_array must be np.arange(Nspws)" - + + def _clear(self): """ Clear UVPSpec of all parameters. Warning: this cannot be undone. @@ -1441,16 +1524,21 @@ def __eq__(self, other, params=None, verbose=False): return True def __add__(self, other, verbose=False): - """ Combine the data of two UVPSpec objects together along a single axis """ + """ + Combine the data of two UVPSpec objects together along a single axis + """ return combine_uvpspec([self, other], verbose=verbose) @property def units(self): - """Return power spectrum units. See self.vis_units and self.norm_units.""" + """ + Return power spectrum units. See self.vis_units and self.norm_units. + """ return "({})^2 {}".format(self.vis_units, self.norm_units) - def generate_noise_spectra(self, spw, pol, Tsys, blpairs=None, little_h=True, - form='Pk', num_steps=2000, real=True): + def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None, + little_h=True, form='Pk', num_steps=2000, + real=True): """ Generate the expected 1-sigma noise power spectrum given a selection of spectral window, system temp., and polarization. This estimate is @@ -1465,7 +1553,7 @@ def generate_noise_spectra(self, spw, pol, Tsys, blpairs=None, little_h=True, number of incoherent averaging samples and comes from self.nsample_array. - If the polarization specified is a pseudo Stokes pol (I, Q, U or V) + If the polarizations specified are pseudo Stokes pol (I, Q, U or V) then an extra factor of 2 is divided. If form == 'DelSq' then a factor of k^3 / (2pi^2) is multiplied. If real is True, a factor of sqrt(2) is divided to account for @@ -1482,9 +1570,9 @@ def generate_noise_spectra(self, spw, pol, Tsys, blpairs=None, little_h=True, spw : int Spectral window index to generate noise curve for. - pol : str or int - Polarization selection in form of str (e.g. 'I' or 'Q' or 'xx') or - int (e.g. -5 or -6). + polpair : tuple or int + Polarization-pair selection in form of tuple (e.g. ('I', 'Q')) or + polpair int. Tsys : float System temperature in Kelvin. @@ -1517,15 +1605,15 @@ def generate_noise_spectra(self, spw, pol, Tsys, blpairs=None, little_h=True, of noise power spectra as values, with ndarrays having shape (Ntimes, Ndlys). """ - # Assert polarization type - if isinstance(pol, (np.integer, int)): - pol = uvutils.polnum2str(pol) + # Assert polarization-pair type + if isinstance(polpair, tuple): + polpair = uvputils.polpair_tuple2int(polpair) # Get polarization index - pol_ind = self.pol_to_indices(pol) + polpair_ind = self.polpair_to_indices(polpair) # Calculate scalar - scalar = self.compute_scalar(spw, pol, num_steps=num_steps, + scalar = self.compute_scalar(spw, polpair, num_steps=num_steps, little_h=little_h, noise_scalar=True) # Get k vectors @@ -1538,7 +1626,7 @@ def generate_noise_spectra(self, spw, pol, Tsys, blpairs=None, little_h=True, if blpairs is None: blpairs = np.unique(self.blpair_array) elif isinstance(blpairs[0], tuple): - blpairs = map(lambda blp: self.antnums_to_blpair(blp), blpairs) + blpairs = [self.antnums_to_blpair(blp) for blp in blpairs] # Get delays dlys = self.get_dlys(spw) @@ -1553,8 +1641,8 @@ def generate_noise_spectra(self, spw, pol, Tsys, blpairs=None, little_h=True, # iterate over time axis for j, ind in enumerate(inds): # get integration time and n_samp - t_int = self.integration_array[spw][ind, pol_ind] - n_samp = self.nsample_array[spw][ind, pol_ind] + t_int = self.integration_array[spw][ind, polpair_ind] + n_samp = self.nsample_array[spw][ind, polpair_ind] # get kvecs if form == 'DelSq': @@ -1674,7 +1762,8 @@ def fold_spectra(self): WARNING: This operation cannot be undone. """ grouping.fold_spectra(self) - + + def get_blpair_groups_from_bl_groups(self, blgroups, only_pairs_in_bls=False): """ Get baseline pair matches from a list of baseline groups. @@ -1707,7 +1796,7 @@ def get_blpair_groups_from_bl_groups(self, blgroups, only_pairs_in_bls=False): return blpair_groups - def compute_scalar(self, spw, pol, num_steps=1000, little_h=True, + def compute_scalar(self, spw, polpair, num_steps=1000, little_h=True, noise_scalar=False): """ Compute power spectrum normalization scalar given an adopted cosmology @@ -1719,8 +1808,8 @@ def compute_scalar(self, spw, pol, num_steps=1000, little_h=True, spw : integer Spectral window selection. - pol : str or int - Which polarization to calculate the scalar for. + polpair : tuple or int + Which polarization pair to calculate the scalar for. num_steps : int, optional Number of integration bins along frequency in computing scalar. @@ -1749,8 +1838,8 @@ def compute_scalar(self, spw, pol, num_steps=1000, little_h=True, spw_freqs = self.freq_array[self.spw_to_freq_indices(spw)] # compute scalar - OP = self.OmegaP[:, self.pol_to_indices(pol)].squeeze() - OPP = self.OmegaPP[:, self.pol_to_indices(pol)].squeeze() + OP = self.OmegaP[:, self.polpair_to_indices(polpair)].squeeze() + OPP = self.OmegaPP[:, self.polpair_to_indices(polpair)].squeeze() scalar = pspecbeam._compute_pspec_scalar(self.cosmo, self.beam_freqs, OPP / OP**2, spw_freqs, num_steps=num_steps, @@ -1781,7 +1870,7 @@ def combine_uvpspec(uvps, verbose=True): A UVPSpec object with the data of all the inputs combined. """ # Perform type checks and get concatenation axis - (uvps, concat_ax, new_spws, new_blpts, new_pols, + (uvps, concat_ax, new_spws, new_blpts, new_polpairs, static_meta) = get_uvp_overlap(uvps, just_meta=False, verbose=verbose) Nuvps = len(uvps) @@ -1791,13 +1880,13 @@ def combine_uvpspec(uvps, verbose=True): # Sort new data axes new_spws = sorted(new_spws) new_blpts = sorted(new_blpts) - new_pols = [new_pols[i] for i in np.argsort(np.abs(new_pols))] + new_polpairs = [new_polpairs[i] for i in np.argsort(np.abs(new_polpairs))] Nspws = len(new_spws) Nblpairts = len(new_blpts) - Npols = len(new_pols) + Npols = len(new_polpairs) # Store covariance only if all uvps have stored covariance. - store_cov = np.all([hasattr(uvp,'cov_array') for uvp in uvps]) + store_cov = np.all([hasattr(uvp, 'cov_array') for uvp in uvps]) # Create new empty data arrays and fill spw arrays u.data_array = odict() @@ -1842,7 +1931,7 @@ def combine_uvpspec(uvps, verbose=True): u.spw_dly_array = np.array(u.spw_dly_array) u.freq_array = np.array(u.freq_array) u.dly_array = np.array(u.dly_array) - u.pol_array = np.array(new_pols) + u.polpair_array = np.array(new_polpairs) # Number of spectral windows, delays etc. u.Nspws = Nspws @@ -1867,8 +1956,9 @@ def combine_uvpspec(uvps, verbose=True): # get each uvp's data axes uvp_spws = [_uvp.get_spw_ranges() for _uvp in uvps] - uvp_blpts = [zip(_uvp.blpair_array, _uvp.time_avg_array) for _uvp in uvps] - uvp_pols = [_uvp.pol_array.tolist() for _uvp in uvps] + uvp_blpts = [list(zip(_uvp.blpair_array, _uvp.time_avg_array)) + for _uvp in uvps] + uvp_polpairs = [_uvp.polpair_array.tolist() for _uvp in uvps] # Construct dict of label indices, to be used for re-ordering later u_lbls = {lbl: ll for ll, lbl in enumerate(u.labels)} @@ -1886,8 +1976,8 @@ def combine_uvpspec(uvps, verbose=True): if i == 0: blpts_idxs0 = blpts_idxs.copy() # Loop over polarizations - for k, p in enumerate(new_pols): - q = uvp_pols[l].index(p) + for k, p in enumerate(new_polpairs): + q = uvp_polpairs[l].index(p) u.scalar_array[i, k] = uvps[l].scalar_array[m, q] # Loop over blpair-times @@ -1935,8 +2025,8 @@ def combine_uvpspec(uvps, verbose=True): m = [spw == _spw for _spw in uvp_spws[l]].index(True) # Loop over polarizations - for k, p in enumerate(new_pols): - q = uvp_pols[l].index(p) + for k, p in enumerate(new_polpairs): + q = uvp_polpairs[l].index(p) u.data_array[i][j, :, k] = uvps[l].data_array[m][n, :, q] if store_cov: @@ -1964,12 +2054,12 @@ def combine_uvpspec(uvps, verbose=True): u.lst_avg_array[j] = uvps[l].lst_avg_array[n] u.blpair_array[j] = uvps[l].blpair_array[n] - elif concat_ax == 'pol': + elif concat_ax == 'polpairs': # Concatenate polarizations - for k, p in enumerate(new_pols): - l = [p in _pols for _pols in uvp_pols].index(True) - q = uvp_pols[l].index(p) + for k, p in enumerate(new_polpairs): + l = [p in _polpairs for _polpairs in uvp_polpairs].index(True) + q = uvp_polpairs[l].index(p) # Get mapping of blpair-time indices between old UVPSpec objects # and the new one @@ -2006,11 +2096,15 @@ def combine_uvpspec(uvps, verbose=True): u.lst_2_array[j] = uvps[0].lst_2_array[n] u.lst_avg_array[j] = uvps[0].lst_avg_array[n] u.blpair_array[j] = uvps[0].blpair_array[n] - + + else: + # Make sure we have properly identified the concat_ax + raise ValueError("concat_ax {} not recognized.".format(concat_ax)) + # Set baselines u.Nblpairs = len(np.unique(u.blpair_array)) uvp_bls = [uvp.bl_array for uvp in uvps] - new_bls = sorted(reduce(operator.or_, [set(bls) for bls in uvp_bls])) + new_bls = np.unique(np.hstack(uvp_bls)) u.bl_array = np.array(new_bls) u.Nbls = len(u.bl_array) u.bl_vecs = [] @@ -2020,7 +2114,7 @@ def combine_uvpspec(uvps, verbose=True): u.bl_vecs.append(uvps[l].bl_vecs[h]) u.bl_vecs = np.array(u.bl_vecs) u.Ntimes = len(np.unique(u.time_avg_array)) - u.history = reduce(operator.add, [uvp.history for uvp in uvps]) + u.history = "".join([uvp.history for uvp in uvps]) u.labels = np.array(u.labels, np.str) for k in static_meta.keys(): @@ -2077,8 +2171,8 @@ def get_uvp_overlap(uvps, just_meta=True, verbose=True): List of unique baseline-pair-time tuples (blpair_integer, time_avg_array JD float) across all input uvps - unique_pols : list - List of unique polarization integers across all input uvps + unique_polpairs : list + List of unique polarization-pair integers across all input uvps """ # type check assert isinstance(uvps, (list, tuple, np.ndarray)), \ @@ -2116,7 +2210,7 @@ def get_uvp_overlap(uvps, just_meta=True, verbose=True): # create unique data axis lists unique_spws = [] unique_blpts = [] - unique_pols = [] + unique_polpairs = [] data_concat_axes = odict() blpts_comb = [] # Combined blpair + time @@ -2124,8 +2218,8 @@ def get_uvp_overlap(uvps, just_meta=True, verbose=True): for uvp1 in uvps: for s in uvp1.get_spw_ranges(): if s not in unique_spws: unique_spws.append(s) - for p in uvp1.pol_array: - if p not in unique_pols: unique_pols.append(p) + for p in uvp1.polpair_array: + if p not in unique_polpairs: unique_polpairs.append(p) uvp1_blpts = zip(uvp1.blpair_array, uvp1.time_avg_array) uvp1_blpts_comb = [bl[0] + 1.j*bl[1] for bl in uvp1_blpts] @@ -2138,21 +2232,21 @@ def get_uvp_overlap(uvps, just_meta=True, verbose=True): for i, uvp1 in enumerate(uvps): # get uvp1 sets and append to unique lists uvp1_spws = uvp1.get_spw_ranges() - uvp1_blpts = zip(uvp1.blpair_array, uvp1.time_avg_array) - uvp1_pols = uvp1.pol_array + uvp1_blpts = list(zip(uvp1.blpair_array, uvp1.time_avg_array)) + uvp1_polpairs = uvp1.polpair_array # iterate over uvps for j, uvp2 in enumerate(uvps): if j <= i: continue # get uvp2 sets uvp2_spws = uvp2.get_spw_ranges() - uvp2_blpts = zip(uvp2.blpair_array, uvp2.time_avg_array) - uvp2_pols = uvp2.pol_array + uvp2_blpts = list(zip(uvp2.blpair_array, uvp2.time_avg_array)) + uvp2_polpairs = uvp2.polpair_array # determine if uvp1 and uvp2 are an identical match spw_match = len(set(uvp1_spws) ^ set(uvp2_spws)) == 0 blpts_match = len(set(uvp1_blpts) ^ set(uvp2_blpts)) == 0 - pol_match = len(set(uvp1_pols) ^ set(uvp2_pols)) == 0 + polpair_match = len(set(uvp1_polpairs) ^ set(uvp2_polpairs)) == 0 # ensure no partial-overlaps if not spw_match: @@ -2161,30 +2255,34 @@ def get_uvp_overlap(uvps, just_meta=True, verbose=True): if not blpts_match: assert len(set(uvp1_blpts) & set(uvp2_blpts)) == 0, \ "uvp {} and {} have partial overlap across blpairts, cannot combine".format(i, j) - if not pol_match: - assert len(set(uvp1_pols) & set(uvp2_pols)) == 0, \ - "uvp {} and {} have partial overlap across pol, cannot combine".format(i, j) + if not polpair_match: + assert len(set(uvp1_polpairs) & set(uvp2_polpairs)) == 0, \ + "uvp {} and {} have partial overlap across pol-pair, cannot combine".format(i, j) # assert all except 1 axis overlaps - matches = [spw_match, blpts_match, pol_match] + matches = [spw_match, blpts_match, polpair_match] assert sum(matches) != 3, \ "uvp {} and {} have completely overlapping data, cannot combine".format(i, j) assert sum(matches) > 1, \ "uvp {} and {} are non-overlapping across multiple data axes, cannot combine".format(i, j) - concat_ax = ['spw', 'blpairts', 'pol'][matches.index(False)] + concat_ax = ['spw', 'blpairts', 'polpairs'][matches.index(False)] data_concat_axes[(i, j)] = concat_ax if verbose: print("uvp {} and {} are concatable across {} axis".format(i, j, concat_ax)) # assert all uvp pairs have the same (single) non-overlap (concat) axis err_msg = "Non-overlapping data in uvps span multiple data axes:\n{}" \ - "".format("\n".join(map(lambda i: "{} & {}: {}".format(i[0][0], i[0][1], i[1]), data_concat_axes.items()))) + "".format("\n".join( ["{} & {}: {}".format(i[0][0],i[0][1],i[1]) + for i in data_concat_axes.items()] )) assert len(set(data_concat_axes.values())) == 1, err_msg # perform additional checks given concat ax if concat_ax == 'blpairts': # check scalar_array - assert np.all([np.isclose(uvps[0].scalar_array, u.scalar_array) for u in uvps[1:]]), "" \ - "scalar_array must be the same for all uvps given concatenation along blpairts." + assert np.all([np.isclose(uvps[0].scalar_array, u.scalar_array) + for u in uvps[1:]]), \ + "scalar_array must be the same for all uvps given " \ + "concatenation along blpairts." - return uvps, concat_ax, unique_spws, unique_blpts, unique_pols, static_meta + return uvps, concat_ax, unique_spws, unique_blpts, \ + unique_polpairs, static_meta diff --git a/hera_pspec/uvpspec_utils.py b/hera_pspec/uvpspec_utils.py index 8a2e33b9..70aad8af 100644 --- a/hera_pspec/uvpspec_utils.py +++ b/hera_pspec/uvpspec_utils.py @@ -1,20 +1,19 @@ import numpy as np import copy, operator from collections import OrderedDict as odict -from pyuvdata.utils import polstr2num - +from pyuvdata.utils import polstr2num, polnum2str def subtract_uvp(uvp1, uvp2, run_check=True, verbose=False): """ Subtract uvp2.data_array from uvp1.data_array. Subtract matching - spw, blpair-lst, pol keys. For non-overlapping keys, remove from output - uvp. + spw, blpair-lst, polpair keys. For non-overlapping keys, remove from + output uvp. Note: Entries in stats_array are added in quadrature. nsample_arrays, integration_arrays and wgt_arrays are inversely added in quadrature. If these arrays are identical in the two objects, this is equivalent - to multiplying (dividing) the stats (nsamp, int & wgt) arrays(s) - by sqrt(2). + to multiplying (dividing) the stats (nsamp, int & wgt) arrays(s) by + sqrt(2). Parameters ---------- @@ -34,11 +33,12 @@ def subtract_uvp(uvp1, uvp2, run_check=True, verbose=False): """ # select out common parts uvp1, uvp2 = select_common([uvp1, uvp2], spws=True, blpairs=True, lsts=True, - pols=True, times=False, inplace=False, verbose=verbose) + polpairs=True, times=False, inplace=False, + verbose=verbose) # get metadata spws1 = [spw for spw in uvp1.get_spw_ranges()] - pols1 = uvp1.pol_array.tolist() + polpairs1 = uvp1.polpair_array.tolist() blps1 = sorted(set(uvp1.blpair_array)) spws2 = [spw for spw in uvp2.get_spw_ranges()] @@ -48,55 +48,64 @@ def subtract_uvp(uvp1, uvp2, run_check=True, verbose=False): i2 = spws2.index(spw) # iterate over pol - for j, pol in enumerate(pols1): + for j, polpair in enumerate(polpairs1): # iterate over blp for k, blp in enumerate(blps1): # form keys - key1 = (i, blp, pol) - key2 = (i2, blp, pol) + key1 = (i, blp, polpair) + key2 = (i2, blp, polpair) # subtract data blp1_inds = uvp1.blpair_to_indices(blp) uvp1.data_array[i][blp1_inds, :, j] -= uvp2.get_data(key2) # add nsample inversely in quadrature - uvp1.nsample_array[i][blp1_inds, j] = np.sqrt(1. / (1./uvp1.get_nsamples(key1)**2 + 1./uvp2.get_nsamples(key2)**2)) + uvp1.nsample_array[i][blp1_inds, j] \ + = np.sqrt( 1. / (1./uvp1.get_nsamples(key1)**2 + + 1. / uvp2.get_nsamples(key2)**2) ) # add integration inversely in quadrature - uvp1.integration_array[i][blp1_inds, j] = np.sqrt(1. / (1./uvp1.get_integrations(key1)**2 + 1./uvp2.get_integrations(key2)**2)) + uvp1.integration_array[i][blp1_inds, j] \ + = np.sqrt(1. / (1./uvp1.get_integrations(key1)**2 + + 1. / uvp2.get_integrations(key2)**2)) # add wgts inversely in quadrature - uvp1.wgt_array[i][blp1_inds, :, :, j] = np.sqrt(1. / (1./uvp1.get_wgts(key1)**2 + 1./uvp2.get_wgts(key2)**2)) - uvp1.wgt_array[i][blp1_inds, :, :, j] /= uvp1.wgt_array[i][blp1_inds, :, :, j].max() + uvp1.wgt_array[i][blp1_inds, :, :, j] \ + = np.sqrt(1. / (1./uvp1.get_wgts(key1)**2 + + 1. / uvp2.get_wgts(key2)**2)) + uvp1.wgt_array[i][blp1_inds, :, :, j] /= \ + uvp1.wgt_array[i][blp1_inds, :, :, j].max() # add stats in quadrature: real imag separately if hasattr(uvp1, "stats_array") and hasattr(uvp2, "stats_array"): for s in uvp1.stats_array.keys(): stat1 = uvp1.get_stats(s, key1) stat2 = uvp2.get_stats(s, key2) - uvp1.stats_array[s][i][blp1_inds, :, j] = np.sqrt(stat1.real**2 + stat2.real**2) + 1j*np.sqrt(stat1.imag**2 + stat2.imag**2) + uvp1.stats_array[s][i][blp1_inds, :, j] \ + = np.sqrt(stat1.real**2 + stat2.real**2) \ + + 1j*np.sqrt(stat1.imag**2 + stat2.imag**2) # add cov in quadrature: real and imag separately if hasattr(uvp1, "cov_array") and hasattr(uvp2, "cov_array"): cov1 = uvp1.get_cov(key1) cov2 = uvp2.get_cov(key2) - uvp1.cov_array[i][blp1_inds, :, :, j] = np.sqrt(cov1.real**2 + cov2.real**2) + 1j*np.sqrt(cov1.imag**2 + cov2.imag**2) + uvp1.cov_array[i][blp1_inds, :, :, j] \ + = np.sqrt(cov1.real**2 + cov2.real**2) \ + + 1j*np.sqrt(cov1.imag**2 + cov2.imag**2) # run check - if run_check: - uvp1.check() - + if run_check: uvp1.check() return uvp1 -def select_common(uvp_list, spws=True, blpairs=True, times=True, pols=True, +def select_common(uvp_list, spws=True, blpairs=True, times=True, polpairs=True, lsts=False, inplace=False, verbose=False): """ - Find spectral windows, baseline-pairs, times, and/or polarizations that a - set of UVPSpec objects have in common and return new UVPSpec objects that - contain only those data. + Find spectral windows, baseline-pairs, times, and/or polarization-pairs + that a set of UVPSpec objects have in common and return new UVPSpec objects + that contain only those data. If there is no overlap, an error will be raised. @@ -129,9 +138,9 @@ def select_common(uvp_list, spws=True, blpairs=True, times=True, pols=True, Whether to retain only the (average) lsts that all UVPSpec objects have in common. Similar algorithm to times. Default: False. - pols : bool, optional - Whether to retain only the polarizations that all UVPSpec objects have - in common. Default: True. + polpairs : bool, optional + Whether to retain only the polarization pairs that all UVPSpec objects + have in common. Default: True. inplace : bool, optional Whether the selection should be applied to the UVPSpec objects @@ -170,18 +179,20 @@ def select_common(uvp_list, spws=True, blpairs=True, times=True, pols=True, common_blpairs = common_blpairs[np.all(has_blpairs, axis=0)] if verbose: print("common_blpairs:", common_blpairs) - # Get polarizations that are common to all - if pols: - common_pols = np.unique(uvp_list[0].pol_array) - has_pols = [np.isin(common_pols, uvp.pol_array) for uvp in uvp_list] - common_pols = common_pols[np.all(has_pols, axis=0)] - if verbose: print("common_pols:", common_pols) + # Get polarization-pairs that are common to all + if polpairs: + common_polpairs = np.unique(uvp_list[0].polpair_array) + has_polpairs = [np.isin(common_polpairs, uvp.polpair_array) + for uvp in uvp_list] + common_polpairs = common_polpairs[np.all(has_polpairs, axis=0)] + if verbose: print("common_polpairs:", common_polpairs) # Get common spectral windows (the entire window must match) # Each row of common_spws is a list of that spw's index in each UVPSpec if spws: common_spws = uvp_list[0].get_spw_ranges() - has_spws = [map(lambda x: x in uvp.get_spw_ranges(), common_spws) for uvp in uvp_list] + has_spws = [[x in uvp.get_spw_ranges() for x in common_spws] + for uvp in uvp_list] common_spws = [common_spws[i] for i, f in enumerate(np.all(has_spws, axis=0)) if f] if verbose: print("common_spws:", common_spws) @@ -199,48 +210,135 @@ def select_common(uvp_list, spws=True, blpairs=True, times=True, pols=True, if lsts and len(common_lsts) == 0: raise ValueError("No lsts were found that exist in all spectra.") - if pols and len(common_pols) == 0: - raise ValueError("No polarizations were found that exist in all spectra.") + if polpairs and len(common_polpairs) == 0: + raise ValueError("No polarization-pairs were found that exist in all spectra.") # Apply selections out_list = [] for i, uvp in enumerate(uvp_list): - _spws, _blpairs, _times, _lsts, _pols = None, None, None, None, None + _spws, _blpairs, _times, _lsts, _polpairs = None, None, None, None, None # Set indices of blpairs, times, and pols to keep if blpairs: _blpairs = common_blpairs if times: _times = common_times if lsts: _lsts = common_lsts - if pols: _pols = common_pols + if polpairs: _pols = common_polpairs if spws: _spws = [uvp.get_spw_ranges().index(j) for j in common_spws] _uvp = uvp.select(spws=_spws, blpairs=_blpairs, times=_times, - pols=_pols, lsts=_lsts, inplace=inplace) + polpairs=_polpairs, lsts=_lsts, inplace=inplace) if not inplace: out_list.append(_uvp) # Return if not inplace if not inplace: return out_list +def polpair_int2tuple(polpair, pol_strings=False): + """ + Convert a pol-pair integer into a tuple pair of polarization + integers. See polpair_tuple2int for more details. + + Parameters + ---------- + polpair : int or list of int + Integer representation of polarization pair. + + pol_strings : bool, optional + If True, return polarization pair tuples with polarization strings. + Otherwise, use polarization integers. Default: True. + + Returns + ------- + polpair : tuple, length 2 + A length-2 tuple containing a pair of polarization + integers, e.g. (-5, -5). + """ + # Recursive evaluation + if isinstance(polpair, (list, np.ndarray)): + return [polpair_int2tuple(p, pol_strings=pol_strings) for p in polpair] + + # Check for integer type + assert isinstance(polpair, (int, np.integer)), \ + "polpair must be integer: %s" % type(polpair) + + # Split into pol1 and pol2 integers + pol1 = int(str(polpair)[:-2]) - 20 + pol2 = int(str(polpair)[-2:]) - 20 + + # Check that pol1 and pol2 are in the allowed range (-8, 4) + if (pol1 < -8 or pol1 > 4) or (pol2 < -8 or pol2 > 4): + raise ValueError("polpair integer evaluates to an invalid " + "polarization pair: (%d, %d)" + % (pol1, pol2)) + # Convert to strings if requested + if pol_strings: + return (polnum2str(pol1), polnum2str(pol2)) + else: + return (pol1, pol2) + + +def polpair_tuple2int(polpair): + """ + Convert a tuple pair of polarization strings/integers into + an pol-pair integer. + + The polpair integer is formed by adding 20 to each standardized + polarization integer (see polstr2num and AIPS memo 117) and + then concatenating them. For example, polarization pair + ('pI', 'pQ') == (1, 2) == 2122. + + Parameters + ---------- + polpair : tuple, length 2 + A length-2 tuple containing a pair of polarization strings + or integers, e.g. ('XX', 'YY') or (-5, -5). + + Returns + ------- + polpair : int + Integer representation of polarization pair. + """ + # Recursive evaluation + if isinstance(polpair, (list, np.ndarray)): + return [polpair_tuple2int(p) for p in polpair] + + # Check types + assert type(polpair) in (tuple,), "pol must be a tuple" + assert len(polpair) == 2, "polpair tuple must have 2 elements" + + # Convert strings to ints if necessary + pol1, pol2 = polpair + if type(pol1) in (str, np.str): pol1 = polstr2num(pol1) + if type(pol2) in (str, np.str): pol2 = polstr2num(pol2) + + # Convert to polpair integer + ppint = (20 + pol1)*100 + (20 + pol2) + return ppint + + def _get_blpairs_from_bls(uvp, bls, only_pairs_in_bls=False): """ - Get baseline pair matches from a list of baseline antenna-pairs in a UVPSpec object. + Get baseline pair matches from a list of baseline antenna-pairs in a + UVPSpec object. Parameters ---------- - uvp : UVPSpec object with at least meta-data in required params loaded in. - If only meta-data is loaded in then h5file must be specified. + uvp : UVPSpec object + Must at least have meta-data in required params loaded in. If only + meta-data is loaded in then h5file must be specified. - bls : list of i6 baseline integers or baseline tuples, Ex. (2, 3) - Select all baseline-pairs whose first _or_ second baseline are in bls list. - This changes if only_pairs_in_bls == True. + bls : list of i6 baseline integers or baseline tuples + Select all baseline-pairs whose first _or_ second baseline are in bls + list. This changes if only_pairs_in_bls == True. - only_pairs_in_bls : bool, if True, keep only baseline-pairs whose first _and_ second baseline - are both found in bls list. + only_pairs_in_bls : bool + If True, keep only baseline-pairs whose first _and_ second baseline are + both found in bls list. - Returns blp_select + Returns ------- - blp_select : boolean ndarray used to index into uvp.blpair_array to get relevant baseline-pairs + blp_select : bool ndarray + Used to index into uvp.blpair_array to get relevant baseline-pairs """ # get blpair baselines in integer form bl1 = np.floor(uvp.blpair_array / 1e6) @@ -248,24 +346,28 @@ def _get_blpairs_from_bls(uvp, bls, only_pairs_in_bls=False): # ensure bls is in integer form if isinstance(bls, tuple): - assert isinstance(bls[0], (int, np.integer)), "bls must be fed as a list of baseline tuples Ex: [(1, 2), ...]" + assert isinstance(bls[0], (int, np.integer)), \ + "bls must be fed as a list of baseline tuples Ex: [(1, 2), ...]" bls = [uvp.antnums_to_bl(bls)] elif isinstance(bls, list): if isinstance(bls[0], tuple): - bls = map(lambda b: uvp.antnums_to_bl(b), bls) + bls = [uvp.antnums_to_bl(b) for b in bls] elif isinstance(bls, (int, np.integer)): bls = [bls] + # get indices if only_pairs_in_bls: - blp_select = np.array(map(lambda blp: np.bool((blp[0] in bls) * (blp[1] in bls)), blpair_bls)) + blp_select = np.array( [np.bool((blp[0] in bls) * (blp[1] in bls)) + for blp in blpair_bls] ) else: - blp_select = np.array(map(lambda blp: np.bool((blp[0] in bls) + (blp[1] in bls)), blpair_bls)) + blp_select = np.array( [np.bool((blp[0] in bls) + (blp[1] in bls)) + for blp in blpair_bls] ) return blp_select def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, - times=None, lsts=None, pols=None, h5file=None): + times=None, lsts=None, polpairs=None, h5file=None): """ Select function for selecting out certain slices of the data, as well as loading in data from HDF5 file. @@ -301,9 +403,8 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, List of lsts from the lst_avg_array to keep. Cannot be fed if times is fed. - pols : list - A list of polarization strings or integers to keep. - See pyuvdata.utils.polstr2num for acceptable options. + polpairs : list + A list of polarization-pair tuples or integers to keep. h5file : h5py file descriptor Used for loading in selection of data from HDF5 file. @@ -352,32 +453,41 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, if blpairs is not None: if bls is None: blp_select = np.zeros(uvp.Nblpairts, np.bool) + # assert form - assert isinstance(blpairs[0], (tuple, int, np.integer)), "blpairs must be fed as a list of baseline-pair tuples or baseline-pair integers" + assert isinstance(blpairs[0], (tuple, int, np.integer)), \ + "blpairs must be fed as a list of baseline-pair tuples or baseline-pair integers" + # if fed as list of tuples, convert to integers if isinstance(blpairs[0], tuple): - blpairs = map(lambda blp: uvp.antnums_to_blpair(blp), blpairs) - blpair_select = np.array(reduce(operator.add, map(lambda blp: uvp.blpair_array == blp, blpairs))) + blpairs = [uvp.antnums_to_blpair(blp) for blp in blpairs] + blpair_select = np.logical_or.reduce( + [uvp.blpair_array == blp for blp in blpairs]) blp_select += blpair_select if times is not None: if bls is None and blpairs is None: blp_select = np.ones(uvp.Nblpairts, np.bool) - time_select = np.array(reduce(operator.add, map(lambda t: np.isclose(uvp.time_avg_array, t, rtol=1e-16), times))) + time_select = np.logical_or.reduce( + [np.isclose(uvp.time_avg_array, t, rtol=1e-16) + for t in times]) blp_select *= time_select if lsts is not None: assert times is None, "Cannot select on lsts and times simultaneously." if bls is None and blpairs is None: blp_select = np.ones(uvp.Nblpairts, np.bool) - lst_select = np.array(reduce(operator.add, map(lambda t: np.isclose(uvp.lst_avg_array, t, rtol=1e-16), lsts))) + lst_select = np.logical_or.reduce( + [ np.isclose(uvp.lst_avg_array, t, rtol=1e-16) + for t in lsts] ) blp_select *= lst_select if bls is None and blpairs is None and times is None and lsts is None: blp_select = slice(None) else: # assert something was selected - assert blp_select.any(), "no selections provided matched any of the data... " + assert blp_select.any(), \ + "no selections provided matched any of the data... " # turn blp_select into slice if possible blp_select = np.where(blp_select)[0] @@ -411,36 +521,40 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, uvp.bl_vecs = uvp.bl_vecs[bl_select] uvp.Nbls = len(uvp.bl_array) - if pols is not None: + if polpairs is not None: # assert form - assert isinstance(pols[0], (str, np.str, int, np.integer)), "pols must be fed as a list of pol strings or pol integers" - - # if fed as strings convert to integers - if isinstance(pols[0], (np.str, str)): - pols = map(lambda p: polstr2num(p), pols) - + assert isinstance(polpairs[0], (tuple, int, np.integer)), \ + "polpairs must be fed as a list of tuples or pol integers" + + # convert to polpair integers if needed + polpairs = [polpair_tuple2int(p) if isinstance(p, tuple) + else p for p in polpairs] + # create selection - pol_select = np.array(reduce(operator.add, map(lambda p: uvp.pol_array == p, pols))) + polpair_select = np.logical_or.reduce( [uvp.polpair_array == p + for p in polpairs] ) # turn into slice object if possible - pol_select = np.where(pol_select)[0] - if len(set(np.diff(pol_select))) == 0: + polpair_select = np.where(polpair_select)[0] + if len(set(np.diff(polpair_select))) == 0: # sliceable - pol_select = slice(pol_select[0], pol_select[-1] + 1) - elif len(set(np.diff(pol_select))) == 1: + polpair_select = slice(polpair_select[0], polpair_select[-1] + 1) + elif len(set(np.diff(polpair_select))) == 1: # sliceable - pol_select = slice(pol_select[0], pol_select[-1] + 1, np.diff(pol_select)[0]) + polpair_select = slice(polpair_select[0], + polpair_select[-1] + 1, + np.diff(polpair_select)[0]) # edit metadata - uvp.pol_array = uvp.pol_array[pol_select] - uvp.Npols = len(uvp.pol_array) + uvp.polpair_array = uvp.polpair_array[polpair_select] + uvp.Npols = len(uvp.polpair_array) if hasattr(uvp, 'scalar_array'): - uvp.scalar_array = uvp.scalar_array[:, pol_select] + uvp.scalar_array = uvp.scalar_array[:, polpair_select] else: - pol_select = slice(None) + polpair_select = slice(None) # determine if data arrays are sliceable - if isinstance(pol_select, slice) or isinstance(blp_select, slice): + if isinstance(polpair_select, slice) or isinstance(blp_select, slice): sliceable = True else: sliceable = False @@ -463,7 +577,8 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, # get stats_array keys if h5file if h5file is not None: - statnames = np.unique([f[f.find("_")+1: f.rfind("_")] for f in h5file.keys() + statnames = np.unique([f[f.find("_")+1: f.rfind("_")] + for f in h5file.keys() if f.startswith("stats")]) else: if hasattr(uvp, "stats_array"): @@ -511,24 +626,24 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, # slice data arrays and assign to dictionaries if sliceable: # can slice in 1 step - data[s] = _data[blp_select, :, pol_select] - wgts[s] = _wgts[blp_select, :, :, pol_select] - ints[s] = _ints[blp_select, pol_select] - nsmp[s] = _nsmp[blp_select, pol_select] + data[s] = _data[blp_select, :, polpair_select] + wgts[s] = _wgts[blp_select, :, :, polpair_select] + ints[s] = _ints[blp_select, polpair_select] + nsmp[s] = _nsmp[blp_select, polpair_select] if store_cov: - cov[s] = _covs[blp_select, :, :, pol_select] + cov[s] = _covs[blp_select, :, :, polpair_select] for statname in statnames: - stats[statname][s] = _stat[statname][blp_select, :, pol_select] + stats[statname][s] = _stat[statname][blp_select, :, polpair_select] else: # need to slice in 2 steps - data[s] = _data[blp_select, :, :][:, :, pol_select] - wgts[s] = _wgts[blp_select, :, :, :][:, :, :, pol_select] - ints[s] = _ints[blp_select, :][:, pol_select] - nsmp[s] = _nsmp[blp_select, :][:, pol_select] + data[s] = _data[blp_select, :, :][:, :, polpair_select] + wgts[s] = _wgts[blp_select, :, :, :][:, :, :, polpair_select] + ints[s] = _ints[blp_select, :][:, polpair_select] + nsmp[s] = _nsmp[blp_select, :][:, polpair_select] if store_cov: - cov[s] = _covs[blp_select, :, :, :][:, :, :, pol_select] + cov[s] = _covs[blp_select, :, :, :][:, :, :, polpair_select] for statname in statnames: - stats[statname][s] = _stat[statname][blp_select, :, :][:, :, pol_select] + stats[statname][s] = _stat[statname][blp_select, :, :][:, :, polpair_select] # assign arrays to uvp uvp.data_array = data @@ -553,7 +668,8 @@ def _blpair_to_antnums(blpair): Returns ------- antnums : tuple - nested tuple containing baseline-pair antenna numbers. Ex. ((ant1, ant2), (ant3, ant4)) + nested tuple containing baseline-pair antenna numbers. + Ex. ((ant1, ant2), (ant3, ant4)) """ # get antennas ant1 = int(np.floor(blpair / 1e9)) @@ -599,6 +715,7 @@ def _antnums_to_blpair(antnums): return blpair + def _bl_to_antnums(bl): """ Convert baseline integer to tuple of antenna numbers. @@ -624,6 +741,7 @@ def _bl_to_antnums(bl): return antnums + def _antnums_to_bl(antnums): """ Convert tuple of antenna numbers to baseline integer. @@ -651,6 +769,7 @@ def _antnums_to_bl(antnums): return bl + def _blpair_to_bls(blpair): """ Convert a blpair integer or nested tuple of antenna pairs @@ -670,6 +789,7 @@ def _blpair_to_bls(blpair): return bl1, bl2 + def _conj_blpair_int(blpair): """ Conjugate a baseline-pair integer @@ -809,6 +929,7 @@ def _fast_lookup_blpairts(src_blpts, query_blpts, time_prec=8): # array lookup functions to be used src_blpts = np.asarray(src_blpts) query_blpts = np.asarray(query_blpts) + src_blpts = src_blpts[:,0] + 1.j*np.around(src_blpts[:,1], time_prec) query_blpts = query_blpts[:,0] + 1.j*np.around(query_blpts[:,1], time_prec) # Applies rounding to time values to ensure reliable float comparisons diff --git a/hera_pspec/version.py b/hera_pspec/version.py index b70cd15a..cc76d255 100644 --- a/hera_pspec/version.py +++ b/hera_pspec/version.py @@ -1,51 +1,127 @@ +# -*- coding: utf-8 -*- +# Copyright 2019 the HERA Project +# Licensed under the MIT License + +from __future__ import print_function, division, absolute_import + import os +import six import subprocess import json +import inspect hera_pspec_dir = os.path.dirname(os.path.realpath(__file__)) -version_file = os.path.join(hera_pspec_dir, 'VERSION') -version = open(version_file).read().strip() + + +def _get_git_output(args, capture_stderr=False): + """ + Get output from git, ensuring that it is of the ``str`` type, not bytes. + """ + + argv = ['git', '-C', hera_pspec_dir] + args + + if capture_stderr: + data = subprocess.check_output(argv, stderr=subprocess.STDOUT) + else: + data = subprocess.check_output(argv) + + data = data.strip() + + if six.PY2: + return data + return data.decode('utf8') + + +def _get_gitinfo_file(git_file=None): + """ + Get saved info from GIT_INFO file that was created when installing package + """ + if git_file is None: + git_file = os.path.join(hera_pspec_dir, 'GIT_INFO') + + with open(git_file) as data_file: + data = [_unicode_to_str(x) for x in json.loads(data_file.read().strip())] + git_origin = data[0] + git_hash = data[1] + git_description = data[2] + git_branch = data[3] + + return {'git_origin': git_origin, 'git_hash': git_hash, + 'git_description': git_description, 'git_branch': git_branch} + + +def _unicode_to_str(u): + if six.PY2: + return u.encode('utf8') + return u def construct_version_info(): - hera_pspec_dir = os.path.dirname(os.path.realpath(__file__)) version_file = os.path.join(hera_pspec_dir, 'VERSION') - version = open(version_file).read().strip() + with open(version_file) as f: + version = f.read().strip() + + git_origin = '' + git_hash = '' + git_description = '' + git_branch = '' + + version_info = {'version': version, 'git_origin': '', 'git_hash': '', + 'git_description': '', 'git_branch': ''} try: - git_origin = subprocess.check_output(['git', '-C', hera_pspec_dir, 'config', - '--get', 'remote.origin.url'], - stderr=subprocess.STDOUT).strip() - git_hash = subprocess.check_output(['git', '-C', hera_pspec_dir, 'rev-parse', 'HEAD'], - stderr=subprocess.STDOUT).strip() - git_description = subprocess.check_output(['git', '-C', hera_pspec_dir, - 'describe', '--dirty', '--tag', '--always']).strip() - git_branch = subprocess.check_output(['git', '-C', hera_pspec_dir, 'rev-parse', - '--abbrev-ref', 'HEAD'], - stderr=subprocess.STDOUT).strip() - git_version = subprocess.check_output(['git', '-C', hera_pspec_dir, 'describe', '--always', - '--tags', '--abbrev=0']).strip() - except: # pragma: no cover - can't figure out how to test exception. + version_info['git_origin'] = _get_git_output( + ['config', '--get', 'remote.origin.url'], + capture_stderr=True) + version_info['git_hash'] = _get_git_output( + ['rev-parse', 'HEAD'], + capture_stderr=True) + version_info['git_description'] = _get_git_output( + ['describe', '--dirty', '--tag', '--always']) + version_info['git_branch'] = _get_git_output( + ['rev-parse', '--abbrev-ref', 'HEAD'], + capture_stderr=True) + except subprocess.CalledProcessError: # pragma: no cover try: # Check if a GIT_INFO file was created when installing package - git_file = os.path.join(hera_pspec_dir, 'GIT_INFO') - with open(git_file) as data_file: - data = [x.encode('UTF8') for x in json.loads(data_file.read().strip())] - git_origin = data[0] - git_hash = data[1] - git_description = data[2] - git_branch = data[3] - except: - git_origin = '' - git_hash = '' - git_description = '' - git_branch = '' - - version_info = {'version': version, 'git_origin': git_origin, - 'git_hash': git_hash, 'git_description': git_description, - 'git_branch': git_branch} + version_info.update(_get_gitinfo_file()) + except (IOError, OSError): + pass + return version_info + +def history_string(notes=''): + """ + Creates a standardized history string that all functions that write to + disk can use. Optionally add notes. + """ + history = '\n------------\nThis file was produced by the function ' \ + + str(inspect.stack()[1][3]) + '()' + # inspect.stack()[1][3] is the name of the function that called this fn + + history += ' in ' + os.path.basename(inspect.stack()[1][1]) + ' using: ' + # inspect.stack()[1][1] is path to the file that contains the function + # that called this function + version_info = construct_version_info() + + for v in sorted(version_info.keys()): + history += '\n ' + v + ': ' + version_info[v] + + if (notes is not None) and (notes != ''): + history += '\n\nNotes:\n' + history += notes + return history + '\n------------\n' + +def print_version_info(): + """ + Print git/version info. + """ + print('Version = {0}'.format(version)) + print('git origin = {0}'.format(git_origin)) + print('git branch = {0}'.format(git_branch)) + print('git description = {0}'.format(git_description)) + version_info = construct_version_info() version = version_info['version'] git_origin = version_info['git_origin'] @@ -53,20 +129,9 @@ def construct_version_info(): git_description = version_info['git_description'] git_branch = version_info['git_branch'] -# String to add to history of any files written with this version of pyuvdata -hera_pspec_version_str = ('hera_pspec version: ' + version + '.') -if git_hash is not '': - hera_pspec_version_str += (' Git origin: ' + str(git_origin) + - '. Git hash: ' + str(git_hash) + - '. Git branch: ' + str(git_branch) + - '. Git description: ' + str(git_description) + '.') - def main(): - print('Version = {0}'.format(version)) - print('git origin = {0}'.format(git_origin)) - print('git branch = {0}'.format(git_branch)) - print('git description = {0}'.format(git_description)) + print_version_info() if __name__ == '__main__': main() diff --git a/setup.py b/setup.py index 6af90132..93eb6f29 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ data = [version.git_origin, version.git_hash, version.git_description, version.git_branch] with open(os.path.join('hera_pspec', 'GIT_INFO'), 'w') as outfile: - json.dump(data, outfile) + json.dump(data, outfile, default=str) def package_files(package_dir, subdirectory): # walk the input package_dir/subdirectory @@ -38,4 +38,4 @@ def package_files(package_dir, subdirectory): } if __name__ == '__main__': - apply(setup, (), setup_args) + setup(**setup_args)