pspecbeam.PSpecBeamUV updates for efield beam

examples/DesignReview_WarmUp.ipynb

@@ -86,7 +86,7 @@
    "metadata": {},
    "source": [
     "### 1c: Choose a beam model\n",
-    "Use the `NF_HERA_Beams.beamfits` file in the `DATA_PATH` for the beam model. You can load this in using the `hera_pspec.pspecbeam.PSpecBeamUV` object. *Don't forget to pass the cosmology object you instantiated previously through to the beam via the *`cosmo`* keyword-argument.*"
+    "Use the `HERA_NF_dipole_power.beamfits` file in the `DATA_PATH` for the beam model. You can load this in using the `hera_pspec.pspecbeam.PSpecBeamUV` object. *Don't forget to pass the cosmology object you instantiated previously through to the beam via the *`cosmo`* keyword-argument.*"
    ]
   },
   {

hera_pspec/pspecbeam.py

@@ -358,30 +358,49 @@ def power_beam_sq_int(self, pol='pI'):
 
 class PSpecBeamUV(PSpecBeamBase):
 
-    def __init__(self, beam_fname, cosmo=None):
+    def __init__(self, uvbeam, cosmo=None):
         """
         Object to store the primary beam for a pspec observation.
         This is subclassed from PSpecBeamBase to take in a pyuvdata
-        UVBeam object.
+        UVBeam filepath or object.
+
+        Note: If one wants to use this object for linear dipole
+        polarizations (e.g. XX, XY, YX, YY) then one can feed
+        uvbeam as a dipole power beam or an efield beam. If, however,
+        one wants to use this for pseudo-Stokes polarizations
+        (e.g. pI, pQ, pU, pV), one must feed uvbeam as a pstokes
+        power beam. See pyuvdata.UVBeam for details on forming
+        pstokes power beams from an efield beam.
 
         Parameters
         ----------
-        beam_fname: str
-            Path to a pyuvdata UVBeam file.
+        uvbeam: str or UVBeam object
+            Path to a pyuvdata UVBeam file or a UVBeam object.
 
         cosmo : conversions.Cosmo_Conversions object, optional
             Cosmology object. Uses the default cosmology object if not 
             specified. Default: None.
         """
-        self.primary_beam = UVBeam()
-        self.primary_beam.read_beamfits(beam_fname)
+        # setup uvbeam object
+        if isinstance(uvbeam, str):
+            uvb = UVBeam()
+            uvb.read_beamfits(uvbeam)
+        else:
+            uvb = uvbeam
 
-        self.beam_freqs = self.primary_beam.freq_array[0]
+        # get frequencies and set cosmology
+        self.beam_freqs = uvb.freq_array[0]
         if cosmo is not None:
             self.cosmo = cosmo
         else:
             self.cosmo = conversions.Cosmo_Conversions()
 
+        # setup primary power beam
+        self.primary_beam = uvb
+        if uvb.beam_type == 'efield':
+            self.primary_beam.efield_to_power(inplace=True)
+            self.primary_beam.peak_normalize()
+
     def power_beam_int(self, pol='pI'):
         """
         Computes the integral of the beam over solid angle to give
@@ -404,7 +423,7 @@ def power_beam_int(self, pol='pI'):
             Scalar integral over beam solid angle.
         """
         if hasattr(self.primary_beam, 'get_beam_area'):
-            return self.primary_beam.get_beam_area(pol)
+            return np.real(self.primary_beam.get_beam_area(pol))
         else:
             raise NotImplementedError("Outdated version of pyuvdata.")
 
@@ -429,7 +448,7 @@ def power_beam_sq_int(self, pol='pI'):
         primary_beam_area: float, array-like
         """
         if hasattr(self.primary_beam, 'get_beam_area'):
-            return self.primary_beam.get_beam_sq_area(pol)
+            return np.real(self.primary_beam.get_beam_sq_area(pol))
         else:
             raise NotImplementedError("Outdated version of pyuvdata.")
 

hera_pspec/tests/test_grouping.py

@@ -10,7 +10,7 @@
 class Test_grouping(unittest.TestCase):
 
     def setUp(self):
-        beamfile = os.path.join(DATA_PATH, 'NF_HERA_Beams.beamfits')
+        beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits')
         self.beam = pspecbeam.PSpecBeamUV(beamfile)
         uvp, cosmo = testing.build_vanilla_uvpspec(beam=self.beam)
         uvp.check()

hera_pspec/tests/test_noise.py

@@ -16,7 +16,7 @@ class Test_Sensitivity(unittest.TestCase):
 
     def setUp(self):
         self.cosmo = conversions.Cosmo_Conversions()
-        self.beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, 'NF_HERA_Beams.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):
@@ -57,7 +57,7 @@ def test_calc_P_N(self):
         t_int = 10.7
         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, 906609626029.72791))
+        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')
         nt.assert_equal(Dsq.shape, (10,))

hera_pspec/tests/test_plot.py

@@ -26,9 +26,9 @@ def setUp(self):
         uvd.read_miriad(os.path.join(DATA_PATH, dfiles[0]))
 
         # Load beam file
-        beamfile = os.path.join(DATA_PATH, 'NF_HERA_Beams.beamfits')
+        beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits')
         self.bm = pspecbeam.PSpecBeamUV(beamfile)
-        self.bm.filename = 'NF_HERA_Beams.beamfits'
+        self.bm.filename = 'HERA_NF_dipole_power.beamfits'
 
         # We only actually have 1 data file here, so slide the time axis by one 
         # integration to avoid noise bias

hera_pspec/tests/test_pspecbeam.py

@@ -5,6 +5,7 @@
 from hera_pspec import pspecbeam, conversions
 from hera_pspec.data import DATA_PATH
 
+
 class Example(unittest.TestCase):
     """
     when running tests in this file by-hand in an interactive interpreter, 
@@ -21,10 +22,7 @@ def runTest(self):
 class Test_DataSet(unittest.TestCase):
 
     def setUp(self):
-        self.beamfile = os.path.join(DATA_PATH, 'NF_HERA_Beams.beamfits')
-        self.bm = pspecbeam.PSpecBeamUV(self.beamfile)
-        self.gauss = pspecbeam.PSpecBeamGauss(0.8, 
-                                  np.linspace(115e6, 130e6, 50, endpoint=False))
+        pass
 
     def tearDown(self):
         pass
@@ -33,84 +31,101 @@ def runTest(self):
         pass
 
     def test_init(self):
-        beamfile = os.path.join(DATA_PATH, 'NF_HERA_Beams.beamfits')
+        beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits')
         bm = pspecbeam.PSpecBeamUV(beamfile)
 
     def test_UVbeam(self):
         # Precomputed results in the following tests were done "by hand" using 
         # iPython notebook "Scalar_dev2.ipynb" in tests directory
-        Om_p = self.bm.power_beam_int()
-        Om_pp = self.bm.power_beam_sq_int()
+        pstokes_beamfile = os.path.join(DATA_PATH, "HERA_NF_pstokes_power.beamfits")
+        beam = pspecbeam.PSpecBeamUV(pstokes_beamfile)
+        Om_p = beam.power_beam_int()
+        Om_pp = beam.power_beam_sq_int()
         lower_freq = 120.*10**6
         upper_freq = 128.*10**6
         num_freqs = 20
-        scalar = self.bm.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=2000)
-
-        # Check that user-defined cosmology can be specified
-        bm = pspecbeam.PSpecBeamUV(self.beamfile,
-                                   cosmo=conversions.Cosmo_Conversions())
 
-        # Check array dimensionality
-        self.assertEqual(Om_p.ndim, 1)
-        self.assertEqual(Om_pp.ndim, 1)
-
-        # Check that errors are raised for other Stokes parameters
-        for pol in ['pQ', 'pU', 'pV',]:
-            nt.assert_raises(NotImplementedError, self.bm.power_beam_int, pol=pol)
-            nt.assert_raises(NotImplementedError, self.bm.power_beam_sq_int, pol=pol)
-            nt.assert_raises(NotImplementedError, self.bm.compute_pspec_scalar, 
-                             lower_freq, upper_freq, num_freqs, pol=pol)
+        # check pI polarization
+        scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=2000)
 
-        # CHeck that invalid polarizations raise an error
-        pol = 'pZ'
-        nt.assert_raises(KeyError, self.bm.power_beam_int, pol=pol)
-        nt.assert_raises(KeyError, self.bm.power_beam_sq_int, pol=pol)
-        nt.assert_raises(KeyError, self.bm.compute_pspec_scalar, 
-                         lower_freq, upper_freq, num_freqs, pol=pol)
-
-
-        self.assertAlmostEqual(Om_p[0], 0.078694909518866998)
-        self.assertAlmostEqual(Om_p[18], 0.065472512282419112)
-        self.assertAlmostEqual(Om_p[-1], 0.029484832405240326)
+        nt.assert_almost_equal(Om_p[0], 0.080082680885906782)
+        nt.assert_almost_equal(Om_p[18], 0.031990943334017245)
+        nt.assert_almost_equal(Om_p[-1], 0.03100215028171072)
+
+        nt.assert_almost_equal(Om_pp[0], 0.036391945229980432)
+        nt.assert_almost_equal(Om_pp[15], 0.018159280192894631)
+        nt.assert_almost_equal(Om_pp[-1], 0.014528100116719534)
 
-        self.assertAlmostEqual(Om_pp[0], 0.035171688022986113)
-        self.assertAlmostEqual(Om_pp[24], 0.024137903003171767)
-        self.assertAlmostEqual(Om_pp[-1], 0.013178952686690554)
+        nt.assert_almost_equal(scalar/568847837.72586381, 1.0, delta=1e-4)
 
-        self.assertAlmostEqual(scalar/567871703.75268996, 1.0, delta=1e-4)
+        # Check array dimensionality
+        nt.assert_equal(Om_p.ndim, 1)
+        nt.assert_equal(Om_pp.ndim, 1)
 
         # convergence of integral
-        scalar_large_Nsteps = self.bm.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=10000) 
-        self.assertAlmostEqual(scalar / scalar_large_Nsteps, 1.0, delta=1e-5)
+        scalar_large_Nsteps = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=10000) 
+        nt.assert_almost_equal(scalar / scalar_large_Nsteps, 1.0, delta=1e-5)
+
+        # Check that user-defined cosmology can be specified
+        beam = pspecbeam.PSpecBeamUV(pstokes_beamfile,
+                                     cosmo=conversions.Cosmo_Conversions())
+
+        # Check that errors are not raised for other Stokes parameters
+        for pol in ['pQ', 'pU', 'pV',]:
+            scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol=pol, num_steps=2000)
 
         # test taper execution
-        scalar = self.bm.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, num_steps=5000, taper='blackman')
-        self.assertAlmostEqual(scalar / 1986172241.1760113, 1.0, delta=1e-8)
+        scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, num_steps=5000, taper='blackman')
+        nt.assert_almost_equal(scalar / 1989353792.1765163, 1.0, delta=1e-8)
 
         # test Jy_to_mK
-        M = self.bm.Jy_to_mK(np.linspace(100e6, 200e6, 11))
+        M = beam.Jy_to_mK(np.linspace(100e6, 200e6, 11))
         nt.assert_equal(len(M), 11)
-        nt.assert_almost_equal(M[0], 41.360105524572283)
+        nt.assert_almost_equal(M[0], 40.643366654821904)
 
         # Extrapolation will fail
-        nt.assert_raises(ValueError, self.bm.Jy_to_mK, 99e6)
-        nt.assert_raises(ValueError, self.bm.Jy_to_mK, 201e6)
+        nt.assert_raises(ValueError, beam.Jy_to_mK, 99e6)
+        nt.assert_raises(ValueError, beam.Jy_to_mK, 201e6)
 
         # test exception
-        nt.assert_raises(TypeError, self.bm.Jy_to_mK, [1])
-        nt.assert_raises(TypeError, self.bm.Jy_to_mK, np.array([1]))
+        nt.assert_raises(TypeError, beam.Jy_to_mK, [1])
+        nt.assert_raises(TypeError, beam.Jy_to_mK, np.array([1]))
 
         # test noise scalar
-        sclr = self.bm.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=2000, noise_scalar=True)
-        nt.assert_almost_equal(sclr, 70.983962969086235)
+        sclr = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=2000, noise_scalar=True)
+        nt.assert_almost_equal(sclr, 71.105979715733)
+
+        # Check that invalid polarizations raise an error
+        pol = 'pZ'
+        nt.assert_raises(KeyError, beam.compute_pspec_scalar, 
+                         lower_freq, upper_freq, num_freqs, pol=pol)
+        pol = 'XX'
+        nt.assert_raises(ValueError, beam.compute_pspec_scalar, 
+                         lower_freq, upper_freq, num_freqs, pol=pol)
+
+        # check dipole beams work
+        dipole_beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits")
+        beam = pspecbeam.PSpecBeamUV(dipole_beamfile)
+        scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='XX')
+        nt.assert_raises(ValueError, beam.compute_pspec_scalar,
+                         lower_freq, upper_freq, num_freqs, pol='pI')
+
+        # check efield beams work
+        efield_beamfile = os.path.join(DATA_PATH, "HERA_NF_efield.beamfits")
+        beam = pspecbeam.PSpecBeamUV(efield_beamfile)
+        scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='XX')
+        nt.assert_raises(ValueError, beam.compute_pspec_scalar,
+                         lower_freq, upper_freq, num_freqs, pol='pI')
 
     def test_Gaussbeam(self):
-        Om_p = self.gauss.power_beam_int()
-        Om_pp = self.gauss.power_beam_sq_int()
+        gauss = pspecbeam.PSpecBeamGauss(0.8, np.linspace(115e6, 130e6, 50, endpoint=False))
+
+        Om_p = gauss.power_beam_int()
+        Om_pp = gauss.power_beam_sq_int()
         lower_freq = 120.*10**6
         upper_freq = 128.*10**6
         num_freqs = 20
-        scalar = self.gauss.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=2000)
+        scalar = gauss.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=2000)
 
         # Check that user-defined cosmology can be specified
         bgauss = pspecbeam.PSpecBeamGauss(0.8, 
@@ -130,21 +145,22 @@ def test_Gaussbeam(self):
         self.assertAlmostEqual(scalar/6392750120.8657961, 1.0, delta=1e-4)
 
         # convergence of integral
-        scalar_large_Nsteps = self.gauss.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, num_steps=5000)
+        scalar_large_Nsteps = gauss.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, num_steps=5000)
         self.assertAlmostEqual(scalar / scalar_large_Nsteps, 1.0, delta=1e-5)
 
         # test taper execution
-        scalar = self.gauss.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, num_steps=5000, taper='blackman')
+        scalar = gauss.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, num_steps=5000, taper='blackman')
         self.assertAlmostEqual(scalar / 22123832163.072491, 1.0, delta=1e-8)    
 
     def test_BeamFromArray(self):
         """
         Test PSpecBeamFromArray
         """
         # Get Gaussian beam to use as a reference
-        Om_P = self.gauss.power_beam_int()
-        Om_PP = self.gauss.power_beam_sq_int()
-        beam_freqs = self.gauss.beam_freqs
+        gauss = pspecbeam.PSpecBeamGauss(0.8, np.linspace(115e6, 130e6, 50, endpoint=False))
+        Om_P = gauss.power_beam_int()
+        Om_PP = gauss.power_beam_sq_int()
+        beam_freqs = gauss.beam_freqs
 
         # Array specs for tests
         lower_freq = 120.*10**6
@@ -168,7 +184,7 @@ def test_BeamFromArray(self):
         # Compare scalar calculation with Gaussian case
         scalar = psbeam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, 
                                              pol='pI', num_steps=2000)
-        g_scalar = self.gauss.compute_pspec_scalar(lower_freq, upper_freq, 
+        g_scalar = gauss.compute_pspec_scalar(lower_freq, upper_freq, 
                                                    num_freqs, pol='pI', 
                                                    num_steps=2000)
         np.testing.assert_array_almost_equal(scalar, g_scalar)
@@ -215,7 +231,6 @@ def test_BeamFromArray(self):
         # Check that string works
         self.assert_(len(str(psbeam)) > 0)
 
-
     def test_PSpecBeamBase(self):
         """
         Test that base class can be instantiated.
@@ -226,11 +241,13 @@ def test_PSpecBeamBase(self):
         bm2 = pspecbeam.PSpecBeamBase(cosmo=conversions.Cosmo_Conversions())
 
     def test_get_Omegas(self):
-        OP, OPP = self.bm.get_Omegas('xx')
-        nt.assert_equal(OP.shape, (101, 1))
-        nt.assert_equal(OPP.shape, (101, 1))
-        OP, OPP = self.bm.get_Omegas([-5, -6])
-        nt.assert_equal(OP.shape, (101, 2))
-        nt.assert_equal(OPP.shape, (101, 2))
+        beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits')
+        beam = pspecbeam.PSpecBeamUV(beamfile)
+        OP, OPP = beam.get_Omegas('xx')
+        nt.assert_equal(OP.shape, (26, 1))
+        nt.assert_equal(OPP.shape, (26, 1))
+        OP, OPP = beam.get_Omegas([-5, -6])
+        nt.assert_equal(OP.shape, (26, 2))
+        nt.assert_equal(OPP.shape, (26, 2))