modified tests for passing pol separately

hera_pspec/pspecdata.py

@@ -917,7 +917,7 @@ def cov_q_hat(self, key1, key2, time_indices=None):
             qc[indnum] = np.trace(np.matmul(Ealphas, np.matmul(N1, np.matmul(Ebetas, N2))), axis1=2, axis2=3)
         return qc/4.
 
-    def q_hat(self, key1, key2, feed_pol, allow_fft=False, exact_norm = False):
+    def q_hat(self, key1, key2, allow_fft=False, exact_norm = False, feed_pol=False):
         """
 
         If exact_norm is False:
@@ -967,9 +967,10 @@ def q_hat(self, key1, key2, feed_pol, allow_fft=False, exact_norm = False):
             If False, Q_alt is used (HERA memo #44, Eq. 16), and the power 
             spectrum is normalized separately.
         
-        feed_pol: str/int
+        feed_pol: str/int/bool, optional
             Used only if exact_norm is True. This argument is passed to get_Q
-            to extract the requested beam polarization.
+            to extract the requested beam polarization. Default is False, in 
+            which case an isotropic beam would be used.
 
         Returns
         -------
@@ -1515,7 +1516,7 @@ def get_Q_alt(self, mode, allow_fft=True):
         Q_alt = np.einsum('i,j', m.conj(), m) # dot it with its conjugate
         return Q_alt
 
-    def get_Q(self, mode, feed_pol):
+    def get_Q(self, mode, feed_pol=False):
         '''
         Computes Q_alpha(i,j), which is the response of the data covariance to the bandpower (dC/dP_alpha). 
         This includes contributions from primary beam.
@@ -1525,8 +1526,8 @@ def get_Q(self, mode, feed_pol):
         mode : int
             Central wavenumber (index) of the bandpower, p_alpha.
 
-        feed_pol : str/int
-            Polarization for the beam 
+        feed_pol : str/int/bool, optional
+            Polarization for the beam. If False, isotropic beam would be returned.
 
         Return
         -------
@@ -2342,7 +2343,7 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None,
 
                     # Calculate unnormalized bandpowers
                     if verbose: print("  Building q_hat...")
-                    qv = self.q_hat(key1, key2, feed_pol, exact_norm=exact_norm)
+                    qv = self.q_hat(key1, key2, exact_norm=exact_norm, feed_pol = feed_pol)
 
                     # Normalize power spectrum estimate
                     if exact_norm:

hera_pspec/tests/test_pspecdata.py

@@ -341,19 +341,19 @@ def test_get_Q(self):
                + 1.j * np.random.normal(size=vect_length)
 
         self.ds.spw_Nfreqs = vect_length
-
+        pol = 'xx' 
         #Test if there is a warning if user does not pass the beam
         key1 = (0, 24, 38)
         key2 = (1, 24, 38)
         uvd = copy.deepcopy(self.uvd)
         ds_t = pspecdata.PSpecData(dsets=[uvd, uvd])
         with warnings.catch_warnings(record=True) as w:
-            ds_t.get_Q(0)
+            ds_t.get_Q(0, pol)
         assert len(w) > 0
 
         for i in range(vect_length):
             try:
-                Q_matrix = self.ds.get_Q(i)
+                Q_matrix = self.ds.get_Q(i, pol)
                 # Test that if the number of delay bins hasn't been set
                 # the code defaults to putting that equal to Nfreqs
                 self.assertEqual(self.ds.spw_Ndlys, self.ds.spw_Nfreqs)
@@ -375,7 +375,7 @@ def test_get_Q(self):
 
         x_vect = np.ones(vect_length)
         try:
-            Q_matrix = self.ds.get_Q(vect_length/2)
+            Q_matrix = self.ds.get_Q(vect_length/2, pol)
         except IndexError:
             Q_matrix = np.ones((vect_length, vect_length))
         xQx = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, x_vect))
@@ -386,7 +386,7 @@ def test_get_Q(self):
         self.ds.set_Ndlys(vect_length-3)
         for i in range(vect_length-3):
             try:
-                Q_matrix = self.ds.get_Q(i)
+                Q_matrix = self.ds.get_Q(i, pol)
             except IndexError:
                 Q_matrix = np.ones((vect_length,vect_length))
             xQy = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, y_vect))
@@ -404,15 +404,15 @@ def test_get_Q(self):
 
         x_vect = np.ones(vect_length)
         try:
-            Q_matrix = self.ds.get_Q((vect_length-2)/2-1)
+            Q_matrix = self.ds.get_Q((vect_length-2)/2-1, pol)
         except IndexError:
             Q_matrix = np.ones((vect_length,vect_length))
         xQx = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, x_vect))
         self.assertAlmostEqual(xQx, np.abs(vect_length**2.))
 
         # Make sure that error is raised when asking for a delay mode outside
         # of the range of delay bins
-        nt.assert_raises(IndexError, self.ds.get_Q, vect_length-1)
+        nt.assert_raises(IndexError, self.ds.get_Q, vect_length-1, pol)
 
     def test_get_unnormed_E(self):
         """
@@ -1113,7 +1113,7 @@ def test_pspec(self):
         bls_Q   = [(24, 25)]
         uvp = ds_Q.pspec(bls_Q, bls_Q, (0, 1), [('xx', 'xx')], input_data_weight='identity',
                                        norm='I', taper='none', verbose=True, exact_norm=False)
-        Q_sample = ds_Q.get_Q((ds_Q.spw_range[1] - ds_Q.spw_range[0])/2) #Get Q matrix for 0th delay mode
+        Q_sample = ds_Q.get_Q((ds_Q.spw_range[1] - ds_Q.spw_range[0])/2, 'xx') #Get Q matrix for 0th delay mode
 
         nt.assert_equal(np.shape(Q_sample), (ds_Q.spw_range[1] - ds_Q.spw_range[0],\
                                              ds_Q.spw_range[1] - ds_Q.spw_range[0])) #Check for the right shape