added time-dependent PN in generate_noise

hera_pspec/grouping.py

@@ -334,7 +334,6 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False,
                         if store_cov:
                             bpg_cov.append(cov * w[:, None])
 
-
                 # Take integration-weighted averages, with clipping to deal
                 # with zeros
                 bpg_data = np.sum(bpg_data, axis=0) \

hera_pspec/noise.py

@@ -44,11 +44,8 @@ def calc_P_N(scalar, Tsys, t_int, Ncoherent=1, Nincoherent=None, form='Pk', k=No
     assert form in ('Pk', 'DelSq'), "form must be either 'Pk' or 'DelSq' for P(k) or Delta^2(k) respectively"
     assert component in ['abs', 'real', 'imag'], "component must be one of 'real', 'imag', 'abs'"
 
-    # convert to mK
-    Tsys *= 1e3
-
-    # construct prefactor
-    P_N = scalar * Tsys**2
+    # construct prefactor in mK^2
+    P_N = scalar * (Tsys * 1e3)**2
 
     # Multiply in effective integration time
     P_N /= (t_int * Ncoherent)

hera_pspec/tests/test_grouping.py

@@ -215,17 +215,16 @@ def test_validate_bootstrap_errorbar():
                    taper='none', sampling=False, little_h=False, spw_ranges=[(0, 50)], verbose=False)
 
     # bootstrap resample
+    Nsamples = 1000
     seed = 0
-    Nsamples = 200
     uvp_avg, uvp_boots, uvp_wgts = grouping.bootstrap_resampled_error(uvp, time_avg=False, Nsamples=Nsamples,
                                                                       seed=seed, normal_std=True,
                                                                       blpair_groups=[uvp.get_blpairs()])
-
     # assert z-score has std of ~1.0 along time ax to within 1/sqrt(Nsamples)
-    bs_std_zscr_real = np.std(uvp_avg.data_array[0].real) / np.mean(uvp_avg.stats_array['bs_std'][0].real)
-    nt.assert_true(np.abs(1.0 - bs_std_zscr_real) < 1/np.sqrt(Nsamples))
-    bs_std_zscr_imag = np.std(uvp_avg.data_array[0].imag) / np.mean(uvp_avg.stats_array['bs_std'][0].imag)
-    nt.assert_true(np.abs(1.0 - bs_std_zscr_imag) < 1/np.sqrt(Nsamples))
+    zscr_real = np.std(uvp_avg.data_array[0].real / uvp_avg.stats_array['bs_std'][0].real)
+    zscr_imag = np.std(uvp_avg.data_array[0].imag / uvp_avg.stats_array['bs_std'][0].imag)
+    nt.assert_true(np.abs(1.0 - zscr_real) < 1/np.sqrt(Nsamples))
+    nt.assert_true(np.abs(1.0 - zscr_imag) < 1/np.sqrt(Nsamples))
 
 
 def test_bootstrap_run():

hera_pspec/tests/test_uvpspec.py

@@ -335,24 +335,33 @@ def test_sense(self):
         uvp = copy.deepcopy(self.uvp)
 
         # test generate noise spectra
-        P_N = uvp.generate_noise_spectra(0, 1515, 500, form='Pk', component='real')
+        polpair = ('xx', 'xx')
+        P_N = uvp.generate_noise_spectra(0, polpair, 500, form='Pk', component='real')
         nt.assert_equal(P_N[101102101102].shape, (10, 30))
 
         # test smaller system temp
-        P_N2 = uvp.generate_noise_spectra(0, 1515, 400, form='Pk', component='real')
+        P_N2 = uvp.generate_noise_spectra(0, polpair, 400, form='Pk', component='real')
         nt.assert_true((P_N[101102101102] > P_N2[101102101102]).all())
 
         # test complex
-        P_N2 = uvp.generate_noise_spectra(0, 1515, 500, form='Pk', component='abs')
+        P_N2 = uvp.generate_noise_spectra(0, polpair, 500, form='Pk', component='abs')
         nt.assert_true((P_N[101102101102] < P_N2[101102101102]).all())
 
         # test Dsq
-        Dsq = uvp.generate_noise_spectra(0, 1515, 500, form='DelSq', component='real')
+        Dsq = uvp.generate_noise_spectra(0, polpair, 500, form='DelSq', component='real')
         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, 1515, 500, form='Pk', component='real')
+        # test a blpair selection and int polpair
+        blpairs = uvp.get_blpairs()[:1]
+        P_N = uvp.generate_noise_spectra(0, 1515, 500, form='Pk', blpairs=blpairs, component='real')
+        nt.assert_equal(P_N[101102101102].shape, (10, 30))
+
+        # test as a dictionary of arrays
+        Tsys = dict([(uvp.antnums_to_blpair(k), 500 * np.ones((uvp.Ntimes, uvp.Ndlys)) * np.linspace(1, 2, uvp.Ntimes)[:, None]) for k in uvp.get_blpairs()])
+        P_N = uvp.generate_noise_spectra(0, 1515, Tsys, form='Pk', blpairs=blpairs, component='real')
+        # assert time gradient is captured: 2 * Tsys results in 4 * P_N
+        nt.assert_true(np.isclose(P_N[101102101102][0, 0] * 4, P_N[101102101102][-1, 0]))
 
     def test_average_spectra(self):
         uvp = copy.deepcopy(self.uvp)

hera_pspec/uvpspec.py

@@ -1572,8 +1572,8 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None,
                                little_h=True, form='Pk', num_steps=2000, 
                                component='real'):
         """
-        Generate the expected 1-sigma noise power spectrum given a selection of
-        spectral window, system temp., and polarization. This estimate is
+        Generate the expected RMS noise power spectrum given a selection of
+        spectral window, system temp. [K], and polarization. This estimate is
         constructed as:
 
         P_N = scalar * (Tsys * 1e3)^2 / (integration_time) / sqrt(Nincoherent)
@@ -1588,7 +1588,7 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None,
 
         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 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
         discarding imaginary noise component.
 
@@ -1608,8 +1608,9 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None,
             polpair int. Strings are expanded into polarization pairs, e.g. 
             'XX' becomes ('XX,'XX').
 
-        Tsys : float
-            System temperature in Kelvin.
+        Tsys : dictionary, float or array
+            System temperature in Kelvin for each blpair. Key is blpair-integer,
+            value is Tsys float or ndarray. If fed as an ndarray, shape=(Ntimes,)
 
         blpairs : list
             List of unique blair tuples or i12 integers to calculate noise
@@ -1667,12 +1668,18 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None,
         # Get delays
         dlys = self.get_dlys(spw)
 
+        # handle Tsys
+        if not isinstance(Tsys, (dict, odict)):
+            if not isinstance(Tsys, np.ndarray):
+                Tsys = np.ones(self.Ntimes) * Tsys
+            Tsys = dict([(blp, Tsys) for blp in blpairs])
+
         # Iterate over blpairs to get P_N
         P_N = odict()
         for i, blp in enumerate(blpairs):
             # get indices
             inds = self.blpair_to_indices(blp)
-
+            assert isinstance(Tsys[blp], (float, np.float, int, np.int)) or Tsys[blp].shape[0] == self.Ntimes, "Tsys must be a float or an ndarray with shape[0] == Ntimes"
             P_blp = []
             # iterate over time axis
             for j, ind in enumerate(inds):
@@ -1687,7 +1694,7 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None,
                     k = None
 
                 # Get noise power spectrum
-                pn = noise.calc_P_N(scalar, Tsys, t_int, k=k,
+                pn = noise.calc_P_N(scalar, Tsys[blp][j], t_int, k=k,
                                     Nincoherent=n_samp, form=form, component=component)
 
                 # Put into appropriate form