Merge pull request #382 from HERA-Team/support-future-array-shapes

Implemented compatibility with UVData future_array_shapes

examples/DesignReview_WarmUp.ipynb

@@ -22,7 +22,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 1,
    "metadata": {
     "collapsed": true
    },
@@ -218,25 +218,39 @@
    },
    "outputs": [],
    "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 2",
+   "display_name": "hera_pspec",
    "language": "python",
-   "name": "python2"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
-    "version": 2
+    "version": 3
    },
    "file_extension": ".py",
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.13"
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,

hera_pspec/pspecbeam.py

@@ -415,12 +415,13 @@ def __init__(self, uvbeam, cosmo=None):
         # setup uvbeam object
         if isinstance(uvbeam, str):
             uvb = UVBeam()
-            uvb.read_beamfits(uvbeam)
+            uvb.read_beamfits(uvbeam, use_future_array_shapes=True)
         else:
             uvb = uvbeam
+            uvb.use_future_array_shapes()
 
         # get frequencies and set cosmology
-        self.beam_freqs = uvb.freq_array[0]
+        self.beam_freqs = uvb.freq_array
         if cosmo is not None:
             self.cosmo = cosmo
         else:
@@ -484,12 +485,11 @@ def beam_normalized_response(self, pol='pI', freq=None, x_orientation=None):
 
         if pol in pol_array:
             stokes_p_ind = np.where(np.isin(pol_array, pol))[0][0]
-            beam_res = beam_res[0, 0, stokes_p_ind] # extract the beam with the correct polarization, dim (nfreq X npix)
+            beam_res = beam_res[0, stokes_p_ind] # extract the beam with the correct polarization, dim (nfreq X npix)
         else:
             raise ValueError('Do not have the right polarization information')
 
         omega = np.sum(beam_res, axis=-1) * np.pi / (3. * nside**2) #compute beam solid angle as a function of frequency
-
         return beam_res, omega, nside
 
     def power_beam_int(self, pol='pI'):

hera_pspec/pspecdata.py

@@ -198,6 +198,9 @@ def add(self, dsets, wgts, labels=None, dsets_std=None, cals=None, cal_flag=True
         for d, w, s in zip(dsets, wgts, dsets_std):
             if not isinstance(d, UVData):
                 raise TypeError("Only UVData objects can be used as datasets.")
+            elif not d.future_array_shapes:
+                warnings.warn('Converting data to future_array_shapes...')
+                d.use_future_array_shapes()
             if not isinstance(w, UVData) and w is not None:
                 raise TypeError("Only UVData objects (or None) can be used as "
                                 "weights.")
@@ -248,7 +251,7 @@ def add(self, dsets, wgts, labels=None, dsets_std=None, cals=None, cal_flag=True
         self.Ntimes = self.dsets[0].Ntimes
 
         # Store the actual frequencies
-        self.freqs = self.dsets[0].freq_array[0]
+        self.freqs = self.dsets[0].freq_array
         self.spw_range = (0, self.Nfreqs)
         self.spw_Nfreqs = self.Nfreqs
         self.spw_Ndlys = self.spw_Nfreqs
@@ -292,7 +295,7 @@ def validate_datasets(self, verbose=True):
         channel_widths = [d.channel_width for d in self.dsets]
         if np.unique(Nfreqs).size > 1:
             raise ValueError("all dsets must have the same Nfreqs")
-        if np.unique(channel_widths).size > 1:
+        if not np.allclose(channel_widths[0], channel_widths):
             raise ValueError("all dsets must have the same channel_widths")
 
         # Check shape along time axis
@@ -2439,7 +2442,7 @@ def broadcast_dset_flags(self, spw_ranges=None, time_thresh=0.2,
                 # unflag
                 if unflag:
                     # unflag for all times
-                    dset.flag_array[:,:,self.spw_range[0]:self.spw_range[1],:] = False
+                    dset.flag_array[:, self.spw_range[0]:self.spw_range[1], :] = False
                     continue
                 # enact time threshold on flag waterfalls
                 # iterate over polarizations
@@ -2450,7 +2453,7 @@ def broadcast_dset_flags(self, spw_ranges=None, time_thresh=0.2,
                         # get baseline-times indices
                         bl_inds = np.where(np.in1d(dset.baseline_array, bl))[0]
                         # get flag waterfall
-                        flags = dset.flag_array[bl_inds, 0, :, i].copy()
+                        flags = dset.flag_array[bl_inds, :, i].copy()
                         Ntimes = float(flags.shape[0])
                         Nfreqs = float(flags.shape[1])
                         # get time- and freq-continguous flags
@@ -2459,12 +2462,12 @@ def broadcast_dset_flags(self, spw_ranges=None, time_thresh=0.2,
                         # get freq channels where non-contiguous flags exceed threshold
                         exceeds_thresh = np.sum(flags[~freq_contig_flgs], axis=0, dtype=float) / Ntimes_noncontig > time_thresh
                         # flag channels for all times that exceed time_thresh
-                        dset.flag_array[bl_inds, :, np.where(exceeds_thresh)[0][:, None], i] = True
+                        dset.flag_array[bl_inds, np.where(exceeds_thresh)[0][:, None], i] = True
                         # for pixels that have flags but didn't meet broadcasting limit
                         # flag the integration within the spw
                         flags[:, np.where(exceeds_thresh)[0]] = False
                         flag_ints = np.max(flags[:, self.spw_range[0]:self.spw_range[1]], axis=1)
-                        dset.flag_array[bl_inds[flag_ints], :, self.spw_range[0]:self.spw_range[1], i] = True
+                        dset.flag_array[bl_inds[flag_ints], self.spw_range[0]:self.spw_range[1], i] = True
 
     def units(self, little_h=True):
         """
@@ -3513,7 +3516,7 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None,
         uvp.polpair_array = np.array(spw_polpair, int)
         uvp.Npols = len(spw_polpair)
         uvp.scalar_array = np.array(sclr_arr)
-        uvp.channel_width = dset1.channel_width  # all dsets validated to agree
+        uvp.channel_width = np.array(dset1.channel_width)  # all dsets validated to agree
         uvp.exact_windows = False
         uvp.weighting = input_data_weight
         uvp.vis_units, uvp.norm_units = self.units(little_h=little_h)
@@ -3678,7 +3681,7 @@ def rephase_to_dset(self, dset_index=0, inplace=True):
                 polind = pol_list.index(uvutils.polstr2num(k[-1], x_orientation=self.dsets[0].x_orientation))
 
                 # insert into dset
-                dset.data_array[indices, 0, :, polind] = data[k]
+                dset.data_array[indices, :, polind] = data[k]
 
             # set phasing in UVData object to unknown b/c there isn't a single
             # consistent phasing for the entire data set.
@@ -3732,7 +3735,7 @@ def Jy_to_mK(self, beam=None):
                 )
                 continue
             for j, p in enumerate(dset.polarization_array):
-                dset.data_array[:, :, :, j] *= factors[p][None, None, :]
+                dset.data_array[:, :, j] *= factors[p][None, :]
             dset.vis_units = 'mK'
 
     def trim_dset_lsts(self, lst_tol=6):
@@ -4176,7 +4179,7 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True,
             assert 0 <= spw[0] < spw[1], f"spw_ranges must be a list of tuples of length 2, with both elements being non-negative integers of increasing value. Got {spw}"
             for dset in dsets:
                 assert spw[1] <= dset.Nfreqs, f"spw_range of {spw} out of range for dset {dset.filename[0]} with the second element being less than the number of frequencies in the data"
-            utils.log(f"Using spw_range: {dsets[0].freq_array[:, spw[0]] /1e6} - {dsets[0].freq_array[:, spw[1] - 1]/1e6} MHz", verbose=verbose)
+            utils.log(f"Using spw_range: {dsets[0].freq_array[spw[0]] /1e6} - {dsets[0].freq_array[spw[1] - 1]/1e6} MHz", verbose=verbose)
 
     # read calibration if provided (calfits partial IO not yet supported)
     if cals is not None:
@@ -4542,7 +4545,7 @@ def _load_dsets(fnames, bls=None, pols=None, logf=None, verbose=True,
         else:
             dfiles = dset
         uvd.read(dfiles, bls=bls, polarizations=pols,
-                 file_type=file_type)
+                 file_type=file_type, use_future_array_shapes=True)
         uvd.extra_keywords['filename'] = json.dumps(dfiles)
         dsets.append(uvd)
 
@@ -4576,9 +4579,9 @@ def _load_cals(cnames, logf=None, verbose=True):
         # read data
         uvc = UVCal()
         if isinstance(cfile, str):
-            uvc.read_calfits(glob.glob(cfile))
+            uvc.read_calfits(glob.glob(cfile), use_future_array_shapes=True)
         else:
-            uvc.read_calfits(cfile)
+            uvc.read_calfits(cfile, use_future_array_shapes=True)
         uvc.extra_keywords['filename'] = json.dumps(cfile)
         cals.append(uvc)
 

hera_pspec/testing.py

@@ -71,8 +71,8 @@ def build_vanilla_uvpspec(beam=None):
     spw_freq_array = np.tile(np.arange(Nspws), Nfreqs)
     spw_dly_array = np.tile(np.arange(Nspws), Ndlys)
     spw_array = np.arange(Nspws)
-    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)
+    freq_array = np.linspace(100e6, 105e6, Nfreqs, endpoint=False)
+    dly_array = utils.get_delays(freq_array, n_dlys=Ndlys)
     polpair_array = np.array(
         [
             1515,
@@ -82,7 +82,7 @@ def build_vanilla_uvpspec(beam=None):
     vis_units = "unknown"
     norm_units = "Hz str"
     weighting = "identity"
-    channel_width = np.median(np.diff(freq_array))
+    channel_width = np.ones(Nfreqs) * np.median(np.diff(freq_array))
     history = "example"
     taper = "none"
     norm = "I"
@@ -291,15 +291,17 @@ def uvpspec_from_data(
     # load data
     if isinstance(data, str):
         uvd = UVData()
-        uvd.read_miriad(data)
+        uvd.read_miriad(data, use_future_array_shapes=True)
     elif isinstance(data, UVData):
         uvd = data
+        uvd.use_future_array_shapes()
 
     if isinstance(data_std, str):
         uvd_std = UVData()
-        uvd_std.read_miriad(data_std)
+        uvd_std.read_miriad(data_std, use_future_array_shapes=True)
     elif isinstance(data_std, UVData):
         uvd_std = data_std
+        uvd_std.use_future_array_shapes()
     else:
         uvd_std = None
     if uvd_std is not None:
@@ -422,11 +424,12 @@ def noise_sim(
     # Read data files
     if isinstance(data, str):
         _data = UVData()
-        _data.read_miriad(data)
+        _data.read_miriad(data, use_future_array_shapes=True)
         data = _data
     elif isinstance(data, UVData):
         if not inplace:
             data = copy.deepcopy(data)
+        data.use_future_array_shapes()
     assert isinstance(data, UVData)
 
     # whiten input data
@@ -461,7 +464,7 @@ def noise_sim(
     if not isinstance(int_time, np.ndarray):
         int_time = np.array([int_time])
     Trms = Tsys / np.sqrt(
-        int_time[:, None, None, None] * data.nsample_array * data.channel_width
+        int_time[:, None, None] * data.nsample_array * data.channel_width[None, :, None]
     )
 
     # if a pol is pStokes pol, divide by extra sqrt(2)
@@ -472,12 +475,12 @@ def noise_sim(
 
     # Get Vrms in Jy using beam
     if beam is not None:
-        freqs = np.unique(data.freq_array)[None, None, :, None]
+        freqs = np.unique(data.freq_array)[None, :, None]
         K_to_Jy = [
             1e3 / (beam.Jy_to_mK(freqs.squeeze(), pol=p))
             for p in data.polarization_array
         ]
-        K_to_Jy = np.array(K_to_Jy).T[None, None, :, :]
+        K_to_Jy = np.array(K_to_Jy).T[None, :, :]
         K_to_Jy /= np.array(
             [np.sqrt(2) if p in [1, 2, 3, 4] else 1.0 for p in data.polarization_array]
         )
@@ -646,9 +649,10 @@ def sky_noise_sim(
 
     if isinstance(data, str):
         uvd = UVData()
-        uvd.read(data)
+        uvd.read(data, use_future_array_shapes=True)
     else:
         uvd = copy.deepcopy(data)
+        uvd.use_future_array_shapes()
     assert (
         -7 not in uvd.polarization_array and -8 not in uvd.polarization_array
     ), "Does not operate on cross-hand polarizations"
@@ -657,7 +661,7 @@ def sky_noise_sim(
         beam = pspecbeam.PSpecBeamUV(beam)
 
     # get metadata
-    freqs = uvd.freq_array[0]
+    freqs = uvd.freq_array
     Ntimes = uvd.Ntimes
     lsts = np.unique(uvd.lst_array)
     int_time = np.median(uvd.integration_time)
@@ -673,7 +677,7 @@ def sky_noise_sim(
             OmegaP[pol] = beam.power_beam_int(pol)
         # interpolate to freq_array of data
         OmegaP[pol] = interpolate.interp1d(
-            beam.primary_beam.freq_array[0],
+            beam.primary_beam.freq_array,
             OmegaP[pol],
             kind="linear",
             bounds_error=False,
@@ -739,6 +743,6 @@ def sky_noise_sim(
         # fill data
         blt_inds = uvd.antpair2ind(bl[:2])
         pol_ind = pols.index(bl[2])
-        uvd.data_array[blt_inds, 0, :, pol_ind] = n + sig
+        uvd.data_array[blt_inds, :, pol_ind] = n + sig
 
     return uvd