Merge pull request #357 from HERA-Team/wf_delay

Update on the UVWindow and FTBeam modules

Author: adeliegorce

hera_pspec/grouping.py

@@ -664,15 +664,20 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa
     if isinstance(bin_widths, (float, int)):
         bin_widths = np.ones_like(kbins) * bin_widths
 
+    # copy input
+    uvp = copy.deepcopy(uvp_in) 
+
+    # transform kgrid to little_h units
+    if not little_h:
+        kbins = kbins / uvp.cosmo.h
+        bin_widths = bin_widths / uvp.cosmo.h
+
     # ensure bins don't overlap
     assert len(kbins) == len(bin_widths)
     kbin_left = kbins - bin_widths / 2
     kbin_right = kbins + bin_widths / 2
     assert np.all(kbin_left[1:] >= kbin_right[:-1] - 1e-6), "kbins must not overlap"
 
-    # copy input
-    uvp = copy.deepcopy(uvp_in) 
-
     # perform time and cylindrical averaging upfront if requested
     if not uvp.exact_windows and (blpair_groups is not None or time_avg):
         uvp.average_spectra(blpair_groups=blpair_groups, time_avg=time_avg,
@@ -695,11 +700,6 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa
     if store_window:
         window_function_array = odict()
 
-    # transform kgrid to little_h units
-    if not little_h:
-        kbins = kbins / uvp.cosmo.h
-        bin_widths = bin_widths / uvp.cosmo.h
-
     # iterate over spectral windows
     spw_ranges = uvp.get_spw_ranges()
     for spw in uvp.spw_array:
@@ -887,7 +887,7 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa
                                                                time_avg=time_avg,
                                                                error_weights=error_weights,
                                                                spw_array=spw,
-                                                               little_h=little_h,
+                                                               little_h=True,
                                                                verbose=True)[spw]
 
     # handle data arrays
@@ -943,9 +943,9 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa
     return uvp
 
 def spherical_wf_from_uvp(uvp_in, kbins, bin_widths,
-                        blpair_groups=None, blpair_lens=None, blpair_weights=None,
-                        error_weights=None, time_avg=False, spw_array=None,
-                        little_h=True, verbose=False):
+                          blpair_groups=None, blpair_lens=None, blpair_weights=None,
+                          error_weights=None, time_avg=False, spw_array=None,
+                          little_h=True, verbose=False):
 
     """
     Obtains exact spherical window functions from an UVPspec object,
@@ -984,8 +984,9 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths,
         Spectral window indices.
 
     little_h : bool, optional
-        If True, kgrid is in h Mpc^-1 units, otherwise just Mpc^-1 units.
-        If False, user must ensure adopted h is consistent with uvp_in.cosmo
+        If True, kbins is in h Mpc^-1 units, otherwise just Mpc^-1 units.
+        The code ensures adopted h is consistent with uvp_in.cosmo. If not,
+        it modifies the unit of kbins.
 
     verbose : bool, optional
         If True, print progress, warnings and debugging info to stdout.
@@ -1004,6 +1005,21 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths,
     if isinstance(bin_widths, (float, int)):
         bin_widths = np.ones_like(kbins) * bin_widths
 
+    # if window functions have been computed without little h
+    # it is not possible to re adjust so kbins need to be in Mpc-1
+    # and reciprocally
+    if little_h != ('h^-3' in uvp_in.norm_units):
+        warnings.warn('Changed little_h units to make kbins consistent ' \
+                      'with uvp.window_function_array. Might be inconsistent ' \
+                      'with the power spectrum units.')
+        if little_h:
+            kbins *= uvp_in.cosmo.h 
+            bin_widths *= uvp_in.cosmo.h 
+        else:
+            kbins /= uvp_in.cosmo.h
+            bin_widths /= uvp_in.cosmo.h
+        little_h = 'h^-3' in uvp_in.norm_units
+
     # ensure bins don't overlap
     assert len(kbins) == len(bin_widths)
     kbin_left = kbins - bin_widths / 2
@@ -1062,11 +1078,6 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths,
                         time_avg=time_avg,
                         inplace=True)
 
-    # transform kgrid to little_h units
-    if not little_h:
-        kbins = kbins / uvp.cosmo.h
-        bin_widths = bin_widths / uvp.cosmo.h
-
     # initialize blank arrays and dicts
     window_function_array = odict()
 
@@ -1078,7 +1089,7 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths,
         # construct array giving the k probed by each baseline-tau pair
         kperps = uvp.cosmo.bl_to_kperp(uvp.cosmo.f2z(avg_nu), little_h=little_h) * blpair_lens
         kparas = uvp.cosmo.tau_to_kpara(uvp.cosmo.f2z(avg_nu), little_h=little_h) * uvp.get_dlys(spw)
-        kmags = np.sqrt(kperps[:, None]**2+kparas**2)
+        kmags = np.sqrt(kperps[:, None]**2 + kparas**2)
 
         # setup arrays 
         window_function_array[spw] = np.zeros((uvp.Ntimes, Nk, Nk, uvp.Npols), dtype=np.float64)
@@ -1092,17 +1103,17 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths,
             kpara_bins = uvp.window_function_kpara[spw][:, ip]
             ktot = np.sqrt(kperp_bins[:, None]**2 + kpara_bins**2)
 
-            cyl_wf = uvp.window_function_array[spw][:, :, :, :, ip]
+            cyl_wf = uvp.window_function_array[spw][..., ip]
             # separate baseline-time axis to iterate over times
-            cyl_wf = cyl_wf.reshape((uvp.Ntimes, uvp.Nblpairs, *cyl_wf.shape[1:] ))
+            cyl_wf = cyl_wf.reshape((uvp.Ntimes, uvp.Nblpairs, *cyl_wf.shape[1:]))
 
             # take average for each time
             for it in range(uvp.Ntimes):
                 wf_spherical = np.zeros((Nk, Nk))
                 for m1 in range(Nk):
                     mask1 = (kbin_left[m1] <= kmags) & (kmags < kbin_right[m1])
                     if np.any(mask1):
-                        wf_temp = np.sum(cyl_wf[it, :, :, :, :]*mask1[:, :, None, None].astype(int), axis=(0, 1))/np.sum(mask1)
+                        wf_temp = np.sum(cyl_wf[it, ...]*mask1[:, :, None, None].astype(int), axis=(0, 1))/np.sum(mask1)
                         if np.sum(wf_temp) > 0.: 
                             for m2 in range(Nk):
                                 mask2 = (kbin_left[m2] <= ktot) & (ktot < kbin_right[m2])
@@ -1113,7 +1124,7 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths,
                                                            where = np.sum(wf_spherical[m1,:]) != 0)
                 spw_window_function.append(wf_spherical)
 
-            window_function_array[spw][:, :, :, ip] = np.copy(spw_window_function)
+            window_function_array[spw][..., ip] = np.copy(spw_window_function)
 
     return window_function_array
 

hera_pspec/tests/test_grouping.py

@@ -608,7 +608,7 @@ def test_spherical_average():
                    spw_ranges=(175,195), taper='bh',verbose=False)
     # get exact window functions
     uvp.get_exact_window_functions(ftbeam_file = os.path.join(DATA_PATH, 'FT_beam_HERA_dipole_test'),
-                                       spw_array=None, inplace=True, verbose=False)
+                                   spw_array=None, inplace=True, verbose=False)
     # spherical window functions for redundant groups
     sph = grouping.spherical_average(uvp, kbins, bin_widths)
     # spherical average for input blpair groups
@@ -619,9 +619,10 @@ def test_spherical_average():
 def test_spherical_wf_from_uvp():
 
     # parameters
-    kbins = np.arange(0, 2.9, 0.25)
+    dk = 0.25
+    kbins = np.arange(0.1, 2.9, dk)
     Nk = len(kbins)
-    bin_widths = 0.25
+    print(kbins[0])
 
     basename = 'FT_beam_HERA_dipole_test'
     # obtain uvp object
@@ -641,50 +642,68 @@ def test_spherical_wf_from_uvp():
     # obtain exact_windows (fiducial usage)
     uvp.get_exact_window_functions(ftbeam_file = os.path.join(DATA_PATH, basename),
                                    inplace=True)
-    wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths)
-                                   #   blpair_groups=blpair_groups,
-                                   # blpair_lens=blpair_lens,
-                                   # blpair_weights=blpair_weights,
-                                   # time_avg=time_avg,
-                                   # error_weights=error_weights,
-                                   # spw_array=spw,
-                                   # little_h=little_h,
-                                   # verbose=True)[spw]
+    wf_array = grouping.spherical_wf_from_uvp(uvp,
+                                              kbins=np.arange(0.1, 2.9, dk),
+                                              bin_widths=dk,
+                                              little_h='h^-3' in uvp.norm_units)
+
     # test different options
     # little_h
-    wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths, little_h=False)
+    wf_array = grouping.spherical_wf_from_uvp(uvp, 
+                                              kbins = np.arange(0.1, 2.9, dk)/uvp.cosmo.h,
+                                              bin_widths= dk/uvp.cosmo.h,
+                                              little_h=True)
     # spw_array
-    wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths, spw_array=0)
-    pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, bin_widths,
-                  spw_array=2)
+    wf_array = grouping.spherical_wf_from_uvp(uvp,
+                                              kbins=np.arange(0.1, 2.9, dk),
+                                              bin_widths=dk, spw_array=0,
+                                              little_h='h^-3' in uvp.norm_units)
+    pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp,
+                  kbins=np.arange(0.1, 2.9, dk), bin_widths=dk,
+                  spw_array=2, little_h='h^-3' in uvp.norm_units)
      # blpair_groups
     blpair_groups, blpair_lens, _ = uvp.get_red_blpairs()
-    wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths,
+    wf_array = grouping.spherical_wf_from_uvp(uvp,
+                                              kbins=np.arange(0.1, 2.9, dk),
+                                              bin_widths=dk,
                                               blpair_groups=blpair_groups,
-                                              blpair_lens=blpair_lens)
-    pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, bin_widths,
-                  blpair_groups=blpair_groups[0])
+                                              blpair_lens=blpair_lens,
+                                              little_h='h^-3' in uvp.norm_units)
+    pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp,
+                  kbins=np.arange(0.1, 2.9, dk), bin_widths=dk,
+                  blpair_groups=blpair_groups[0], little_h='h^-3' in uvp.norm_units)
 
     # raise warning or error if blpair_groups inconsistent with blpair_lens
-    wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths,
+    wf_array = grouping.spherical_wf_from_uvp(uvp, kbins=np.arange(0.1, 2.9, dk),
+                                              bin_widths=dk,
                                               blpair_groups=None,
-                                              blpair_lens=blpair_lens)
-    wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths,
+                                              blpair_lens=blpair_lens,
+                                              little_h='h^-3' in uvp.norm_units)
+    wf_array = grouping.spherical_wf_from_uvp(uvp, kbins=np.arange(0.1, 2.9, dk),
+                                              bin_widths=dk,
                                               blpair_groups=blpair_groups,
-                                              blpair_lens=None)
-    pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, bin_widths,
-                  blpair_groups=blpair_groups, blpair_lens=[blpair_lens[0],blpair_lens[0]])
+                                              blpair_lens=None,
+                                              little_h='h^-3' in uvp.norm_units)
+    pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp,
+                  kbins=np.arange(0.1, 2.9, dk), bin_widths=dk,
+                  blpair_groups=blpair_groups, blpair_lens=[blpair_lens[0],blpair_lens[0]],
+                  little_h='h^-3' in uvp.norm_units)
     # error if overlapping bins
-    pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, 1.0)
-    # blpair_weights
-    wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths,
+    pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp,
+                 kbins=np.arange(0.1, 2.9, dk), bin_widths=1.0)
+    # # blpair_weights
+    wf_array = grouping.spherical_wf_from_uvp(uvp, kbins=np.arange(0.1, 2.9, dk),
+                                              bin_widths=dk,
                                               blpair_groups=blpair_groups,
                                               blpair_lens=blpair_lens,
-                                              blpair_weights=[[1. for item in grp] for grp in blpair_groups])
+                                              little_h='h^-3' in uvp.norm_units)
+                                              # blpair_weights=[[1. for item in grp] for grp in blpair_groups],
 
     # raise error if uvp.exact_windows is False
     uvp.exact_windows = False
-    pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, bin_widths)
+    pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp,
+                  kbins=np.arange(0.1, 2.9, dk), bin_widths=dk,
+                  little_h='h^-3' in uvp.norm_units)
 
 
 if __name__ == "__main__":

hera_pspec/tests/test_uvwindow.py

@@ -125,6 +125,30 @@ def test_from_file(self):
         pytest.raises(AssertionError, uvwindow.FTBeam.from_file, spw_range=(1001, 1022),
                       ftfile=self.ft_file)
 
+    def test_gaussian(self):
+
+        # fiducial use
+        widths = -0.0343 * self.freq_array/1e6 + 11.30 
+        test = uvwindow.FTBeam.gaussian(freq_array=self.freq_array,
+                                        widths=widths,
+                                        pol=self.pol)
+        # if widths given as unique number, this value is used for all freqs
+        test2 = uvwindow.FTBeam.gaussian(freq_array=self.freq_array,
+                                        widths=np.mean(widths),
+                                        pol=self.pol)
+
+        # tests on freq_array consistency
+        pytest.raises(AssertionError, uvwindow.FTBeam.gaussian,
+                      freq_array=self.freq_array[:2], pol=self.pol,
+                      widths=np.mean(widths))
+        pytest.raises(AssertionError, uvwindow.FTBeam.gaussian,
+                      freq_array=self.freq_array, pol=self.pol,
+                      widths=widths[:10])
+
+        # make sure widths are given in degrees (raises warning)
+        test = uvwindow.FTBeam.gaussian(freq_array=self.freq_array,
+                                        pol=self.pol, widths=0.10)
+
     def test_get_bandwidth(self):
 
         test_bandwidth = uvwindow.FTBeam.get_bandwidth(self.ft_file)
@@ -431,12 +455,11 @@ def test_get_cylindrical_wf(self):
                                                      return_bins='unweighted')
         assert np.all(cyl_wf2 == cyl_wf)
         assert np.all(kperp2 == kperp)  # unweighted option to return_bins
-        # warning raised if kperp_bins not linearly spaced
+        # ValueError raised if kperp_bins not linearly spaced
         kperp_log = np.logspace(-2, 0, 100)
-        _, _, _ = test.get_cylindrical_wf(bl_len,
-                                          kperp_bins=kperp_log*test.kunits,
-                                          kpara_bins=None,
-                                          return_bins='unweighted')
+        pytest.raises(ValueError, test.get_cylindrical_wf, bl_len,
+                      kperp_bins=kperp_log*test.kunits, kpara_bins=None,
+                      return_bins='unweighted')
 
         # kpara bins
         _, kpara3, cyl_wf3 = test.get_cylindrical_wf(bl_len,
@@ -445,12 +468,11 @@ def test_get_cylindrical_wf(self):
                                                      return_bins='unweighted')
         assert np.all(cyl_wf3 == cyl_wf)
         assert np.all(kpara == kpara3)
-        # warning raised if kpara_bins not linearly spaced
+        # ValueError raised if kpara_bins not linearly spaced
         kpara_log = np.logspace(-1, 1, 100)
-        _, _, _ = test.get_cylindrical_wf(bl_len,
-                                          kperp_bins=None,
-                                          kpara_bins=kpara_log*test.kunits,
-                                          return_bins='unweighted')
+        pytest.raises(ValueError, test.get_cylindrical_wf, bl_len,
+                      kperp_bins=None, kpara_bins=kpara_log*test.kunits,
+                      return_bins='unweighted')
 
         # test filling array by delay symmetry for odd number of delays
         ft_beam_test = uvwindow.FTBeam.from_file(ftfile=self.ft_file, 
@@ -507,12 +529,11 @@ def test_cylindrical_to_spherical(self):
                                                           ktot=ktot, 
                                                           bl_lens=bl_len)        
 
-        # raise warning if kbins not linearly spaced
+        # raise ValueError if kbins not linearly spaced
         kbins_log = np.logspace(-2, 2, 20)
-        _, _ = test.cylindrical_to_spherical(cyl_wf=cyl_wf, 
-                                             kbins=kbins_log*test.kunits, 
-                                             ktot=ktot, 
-                                             bl_lens=bl_len)  
+        pytest.raises(ValueError, test.cylindrical_to_spherical, cyl_wf=cyl_wf, 
+                      kbins=kbins_log*test.kunits, ktot=ktot, 
+                      bl_lens=bl_len)  
 
     def test_get_spherical_wf(self):
 
@@ -530,6 +551,7 @@ def test_get_spherical_wf(self):
                                            verbose=True)
         kperp_bins = test.get_kperp_bins(self.lens[:1])
         kpara_bins = test.get_kpara_bins(test.freq_array)
+        print(np.diff(kpara_bins), np.diff(kperp_bins))
 
         WF = test.get_spherical_wf(kbins=self.kbins,
                                    kperp_bins=kperp_bins,
@@ -559,18 +581,16 @@ def test_get_spherical_wf(self):
                                                         step=kpara_centre)
                                    * test.kunits,
                                    bl_lens=self.lens[:1])
-        # warning raised if kbins not linearly spaced
+        # ValueError raised if kbins not linearly spaced
         kperp_log = np.logspace(-2, 0, 100)
-        _ = test.get_spherical_wf(kbins=self.kbins, 
-                                  kperp_bins=kperp_log*test.kunits,
-                                  bl_lens=self.lens[:1])
+        pytest.raises(ValueError, test.get_spherical_wf, kbins=self.kbins, 
+                      kperp_bins=kperp_log*test.kunits, bl_lens=self.lens[:1])
         kpara_log = np.logspace(-1, 1, 100)
-        _ = test.get_spherical_wf(kbins=self.kbins, 
-                                  kpara_bins=kpara_log*test.kunits,
-                                  bl_lens=self.lens[:1])
+        pytest.raises(ValueError, test.get_spherical_wf, kbins=self.kbins, 
+                      kpara_bins=kpara_log*test.kunits, bl_lens=self.lens[:1])
         kbins_log = np.logspace(-2, 2, 20)
-        _ = test.get_spherical_wf(kbins=kbins_log*test.kunits,
-                                  bl_lens=self.lens[:1])
+        pytest.raises(ValueError, test.get_spherical_wf, kbins=kbins_log*test.kunits,
+                      bl_lens=self.lens[:1])
 
 
     def test_check_kunits(self):