Update polarisation strings to new pyuvdata names and try to remove m…

…atplotlib warning in tests
HERA-Team · May 11, 2018 · 66461cb · 66461cb
1 parent 729520a
commit 66461cb
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 54 deletions.
diff --git a/.travis.yml b/.travis.yml
@@ -46,6 +46,7 @@ install:
 before_script:
   - "export DISPLAY=:99.0"
   - "sh -e /etc/init.d/xvfb start"
+  - "export MPLBACKEND=agg"
   - sleep 3
 
 script: nosetests hera_pspec --with-coverage --cover-package=hera_pspec

diff --git a/hera_pspec/pspecbeam.py b/hera_pspec/pspecbeam.py
@@ -117,15 +117,16 @@ def __init__(self, cosmo=None):
         else:
             self.cosmo = conversions.Cosmo_Conversions()
 
-    def compute_pspec_scalar(self, lower_freq, upper_freq, num_freqs, num_steps=5000, 
-                             pol='I', taper='none', little_h=True, noise_scalar=False):
+    def compute_pspec_scalar(self, lower_freq, upper_freq, num_freqs, 
+                             num_steps=5000, pol='pI', taper='none', 
+                             little_h=True, noise_scalar=False):
         """
         Computes the scalar function to convert a power spectrum estimate
         in "telescope units" to cosmological units
 
         See arxiv:1304.4991 and HERA memo #27 for details.
 
-        Currently, only the "I", "XX" and "YY" polarization beams are supported.
+        Currently, only the "pI", "XX" and "YY" polarization beams are supported.
         See Equations 4 and 5 of Moore et al. (2017) ApJ 836, 154
         or arxiv:1502.05072 for details.
 
@@ -148,8 +149,8 @@ def compute_pspec_scalar(self, lower_freq, upper_freq, num_freqs, num_steps=5000
 
         pol: str, optional
                 Which polarization to compute the beam scalar for.
-                'I', 'Q', 'U', 'V', 'XX', 'YY', 'XY', 'YX', although 
-                Default: 'I'
+                'pI', 'pQ', 'pU', 'pV', 'XX', 'YY', 'XY', 'YX'
+                Default: 'pI'
 
         taper : str, optional
             Whether a tapering function (e.g. Blackman-Harris) is being used in 
@@ -187,7 +188,7 @@ def compute_pspec_scalar(self, lower_freq, upper_freq, num_freqs, num_steps=5000
                                        noise_scalar=noise_scalar)
         return scalar
 
-    def Jy_to_mK(self, freqs, pol='I'):
+    def Jy_to_mK(self, freqs, pol='pI'):
         """
         Return the multiplicative factor [mK / Jy], to convert a visibility 
         from Jy -> mK,
@@ -206,8 +207,8 @@ def Jy_to_mK(self, freqs, pol='I'):
 
         pol: str, optional
                 Which polarization to compute the beam scalar for.
-                'I', 'Q', 'U', 'V', 'XX', 'YY', 'XY', 'YX', although 
-                Default: 'I'
+                'pI', 'pQ', 'pU', 'pV', 'XX', 'YY', 'XY', 'YX'
+                Default: 'pI'
 
         Returns
         -------
@@ -300,7 +301,7 @@ def __init__(self, fwhm, beam_freqs, cosmo=None):
         else:
             self.cosmo = conversions.Cosmo_Conversions()
 
-    def power_beam_int(self, pol='I'):
+    def power_beam_int(self, pol='pI'):
         """
         Computes the integral of the beam over solid angle to give
         a beam area (in sr). Uses analytic formula that the answer
@@ -316,8 +317,8 @@ def power_beam_int(self, pol='I'):
         ----------
         pol: str, optional
                 Which polarization to compute the beam scalar for.
-                'I', 'Q', 'U', 'V', 'XX', 'YY', 'XY', 'YX' 
-                Default: 'I'
+                'pI', 'pQ', 'pU', 'pV', 'XX', 'YY', 'XY', 'YX' 
+                Default: 'pI'
 
         Returns
         -------
@@ -327,7 +328,7 @@ def power_beam_int(self, pol='I'):
         return np.ones_like(self.beam_freqs) * 2. * np.pi * self.fwhm**2 \
                / (8. * np.log(2.))
 
-    def power_beam_sq_int(self, pol='I'):
+    def power_beam_sq_int(self, pol='pI'):
         """
         Computes the integral of the beam**2 over solid angle to give
         a beam area (in str). Uses analytic formula that the answer
@@ -343,8 +344,8 @@ def power_beam_sq_int(self, pol='I'):
         ----------
         pol: str, optional
                 Which polarization to compute the beam scalar for.
-                'I', 'Q', 'U', 'V', 'XX', 'YY', 'XY', 'YX' 
-                Default: 'I'
+                'pI', 'pQ', 'pU', 'pV', 'XX', 'YY', 'XY', 'YX' 
+                Default: 'pI'
 
         Returns
         -------
@@ -381,7 +382,7 @@ def __init__(self, beam_fname, cosmo=None):
         else:
             self.cosmo = conversions.Cosmo_Conversions()
 
-    def power_beam_int(self, pol='I'):
+    def power_beam_int(self, pol='pI'):
         """
         Computes the integral of the beam over solid angle to give
         a beam area (in str) as a function of frequency. Uses function
@@ -394,8 +395,8 @@ def power_beam_int(self, pol='I'):
         ----------
         pol: str, optional
                 Which polarization to compute the beam scalar for.
-                'I', 'Q', 'U', 'V', 'XX', 'YY', 'XY', 'YX' 
-                Default: 'I'
+                'pI', 'pQ', 'pU', 'pV', 'XX', 'YY', 'XY', 'YX' 
+                Default: 'pI'
 
         Returns
         -------
@@ -407,7 +408,7 @@ def power_beam_int(self, pol='I'):
         else:
             raise NotImplementedError("Outdated version of pyuvdata.")
 
-    def power_beam_sq_int(self, pol='I'):
+    def power_beam_sq_int(self, pol='pI'):
         """
         Computes the integral of the beam**2 over solid angle to give
         a beam**2 area (in str) as a function of frequency. Uses function
@@ -420,8 +421,8 @@ def power_beam_sq_int(self, pol='I'):
         ----------
         pol: str, optional
                 Which polarization to compute the beam scalar for.
-                'I', 'Q', 'U', 'V', 'XX', 'YY', 'XY', 'YX'
-                Default: 'I'
+                'pI', 'pQ', 'pU', 'pV', 'XX', 'YY', 'XY', 'YX'
+                Default: 'pI'
 
         Returns
         -------
@@ -442,7 +443,7 @@ def __init__(self, OmegaP, OmegaPP, beam_freqs, cosmo=None):
         
         Allowed polarizations are: 
         
-            I, Q, U, V, XX, YY, XY, YX
+            pI, pQ, pU, pV, XX, YY, XY, YX
         
         Other polarizations will be ignored.
         
@@ -471,7 +472,7 @@ def __init__(self, OmegaP, OmegaPP, beam_freqs, cosmo=None):
             Cosmology object. Uses the default cosmology object if not 
             specified. Default: None.
         """
-        self.allowed_pols = ['I', 'Q', 'U', 'V', 
+        self.allowed_pols = ['pI', 'pQ', 'pU', 'pV', 
                              'XX', 'YY', 'XY', 'YX']
         self.OmegaP = {}; self.OmegaPP = {}
 
@@ -480,8 +481,8 @@ def __init__(self, OmegaP, OmegaPP, beam_freqs, cosmo=None):
 
         if isinstance(OmegaP, np.ndarray) and isinstance(OmegaPP, np.ndarray):
             # Only single arrays were specified; assume I
-            OmegaP = {'I': OmegaP}
-            OmegaPP = {'I': OmegaPP}
+            OmegaP = {'pI': OmegaP}
+            OmegaPP = {'pI': OmegaPP}
 
         elif isinstance(OmegaP, np.ndarray) or isinstance(OmegaPP, np.ndarray):
             # Mixed dict and array types are not allowed
@@ -531,7 +532,7 @@ def add_pol(self, pol, OmegaP, OmegaPP):
             Which polarization to add beam solid angle arrays for. Valid 
             options are:
             
-              'I', 'Q', 'U', 'V', 
+              'pI', 'pQ', 'pU', 'pV', 
               'XX', 'YY', 'XY', 'YX' 
             
             If the arrays already exist for the specified polarization, they 
@@ -566,7 +567,7 @@ def add_pol(self, pol, OmegaP, OmegaPP):
         self.OmegaPP[pol] = OmegaPP
 
 
-    def power_beam_int(self, pol='I'):
+    def power_beam_int(self, pol='pI'):
         """
         Computes the integral of the beam over solid angle to give
         a beam area (in str) as a function of frequency.
@@ -575,9 +576,9 @@ def power_beam_int(self, pol='I'):
         ----------
         pol: str, optional
             Which polarization to compute the beam scalar for. 
-                'I', 'Q', 'U', 'V', 
+                'pI', 'pQ', 'pU', 'pV', 
                 'XX', 'YY', 'XY', 'YX'
-            Default: I.
+            Default: pI.
 
         Returns
         -------
@@ -592,7 +593,7 @@ def power_beam_int(self, pol='I'):
                            "Available polarizations are: %s" \
                            % (pol, available_pols))
 
-    def power_beam_sq_int(self, pol='I'):
+    def power_beam_sq_int(self, pol='pI'):
         """
         Computes the integral of the beam**2 over solid angle to give
         a beam**2 area (in str) as a function of frequency.
@@ -601,9 +602,9 @@ def power_beam_sq_int(self, pol='I'):
         ----------
         pol: str, optional
             Which polarization to compute the beam scalar for.
-              'I', 'Q', 'U', 'V', 
+              'pI', 'pQ', 'pU', 'pV', 
               'XX', 'YY', 'XY', 'YX' 
-            Default: I.
+            Default: pI.
 
         Returns
         -------

diff --git a/hera_pspec/tests/test_noise.py b/hera_pspec/tests/test_noise.py
@@ -43,14 +43,14 @@ def test_set(self):
 
     def test_scalar(self):
         freqs = np.linspace(150e6, 160e6, 100, endpoint=False)
-        self.sense.calc_scalar(freqs, 'I', num_steps=5000, little_h=True)
+        self.sense.calc_scalar(freqs, 'pI', num_steps=5000, little_h=True)
         nt.assert_true(np.isclose(freqs, self.sense.subband).all())
-        nt.assert_true(self.sense.pol, 'I')
+        nt.assert_true(self.sense.pol, 'pI')
 
     def test_calc_P_N(self):
         # calculate scalar
         freqs = np.linspace(150e6, 160e6, 100, endpoint=False)
-        self.sense.calc_scalar(freqs, 'I', num_steps=5000, little_h=True)
+        self.sense.calc_scalar(freqs, 'pI', num_steps=5000, little_h=True)
         # basic execution 
         k = np.linspace(0, 3, 10)
         Tsys = 500.0

diff --git a/hera_pspec/tests/test_plot.py b/hera_pspec/tests/test_plot.py
@@ -1,9 +1,8 @@
 import unittest
 import nose.tools as nt
+import matplotlib
 import numpy as np
 import os, copy, sys
-import matplotlib
-matplotlib.use('Agg') # Do not assume we have access to an X server
 from hera_pspec import pspecdata, pspecbeam, conversions, plot
 from hera_pspec.data import DATA_PATH
 from pyuvdata import UVData

diff --git a/hera_pspec/tests/test_pspecbeam.py b/hera_pspec/tests/test_pspecbeam.py
@@ -44,7 +44,7 @@ def test_UVbeam(self):
         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='I', num_steps=2000)
+        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,
@@ -55,12 +55,20 @@ def test_UVbeam(self):
         self.assertEqual(Om_pp.ndim, 1)
 
         # Check that errors are raised for other Stokes parameters
-        for pol in ['Q', 'U', 'V', 'Z']:
+        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 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)
@@ -72,7 +80,7 @@ def test_UVbeam(self):
         self.assertAlmostEqual(scalar/567871703.75268996, 1.0, delta=1e-4)
 
         # convergence of integral
-        scalar_large_Nsteps = self.bm.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='I', num_steps=10000) 
+        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)
 
         # test taper execution
@@ -93,7 +101,7 @@ def test_UVbeam(self):
         nt.assert_raises(TypeError, self.bm.Jy_to_mK, np.array([1]))
 
         # test noise scalar
-        sclr = self.bm.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='I', num_steps=2000, noise_scalar=True)
+        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)
 
     def test_Gaussbeam(self):
@@ -102,7 +110,7 @@ def test_Gaussbeam(self):
         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='I', num_steps=2000)
+        scalar = self.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, 
@@ -148,8 +156,8 @@ def test_BeamFromArray(self):
                                               beam_freqs=beam_freqs)
 
         psbeampol = pspecbeam.PSpecBeamFromArray(
-                                OmegaP={'I': Om_P, 'Q': Om_P},
-                                OmegaPP={'I': Om_PP, 'Q': Om_PP},
+                                OmegaP={'pI': Om_P, 'pQ': Om_P},
+                                OmegaPP={'pI': Om_PP, 'pQ': Om_PP},
                                 beam_freqs=beam_freqs)
 
         # Check that user-defined cosmology can be specified
@@ -159,15 +167,15 @@ def test_BeamFromArray(self):
 
         # Compare scalar calculation with Gaussian case
         scalar = psbeam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, 
-                                             pol='I', num_steps=2000)
+                                             pol='pI', num_steps=2000)
         g_scalar = self.gauss.compute_pspec_scalar(lower_freq, upper_freq, 
-                                                   num_freqs, pol='I', 
+                                                   num_freqs, pol='pI', 
                                                    num_steps=2000)
         np.testing.assert_array_almost_equal(scalar, g_scalar)
 
         # Check that polarizations are recognized and invalid ones rejected
         scalarp = psbeampol.compute_pspec_scalar(lower_freq, upper_freq, 
-                                                 num_freqs, pol='Q', 
+                                                 num_freqs, pol='pQ', 
                                                  num_steps=2000)
 
         # Test taper execution (same as Gaussian case)
@@ -177,10 +185,10 @@ def test_BeamFromArray(self):
 
         # Check that invalid init args raise errors
         nt.assert_raises(TypeError, pspecbeam.PSpecBeamFromArray, OmegaP=Om_P, 
-                         OmegaPP={'I': Om_PP}, beam_freqs=beam_freqs)
+                         OmegaPP={'pI': Om_PP}, beam_freqs=beam_freqs)
         nt.assert_raises(KeyError, pspecbeam.PSpecBeamFromArray,
-                         OmegaP={'I': Om_P, 'Q': Om_P},
-                         OmegaPP={'I': Om_PP,},
+                         OmegaP={'pI': Om_P, 'pQ': Om_P},
+                         OmegaPP={'pI': Om_PP,},
                          beam_freqs=beam_freqs)
 
         nt.assert_raises(KeyError, pspecbeam.PSpecBeamFromArray,
@@ -189,13 +197,13 @@ def test_BeamFromArray(self):
                          beam_freqs=beam_freqs)
 
         nt.assert_raises(TypeError, pspecbeam.PSpecBeamFromArray,
-                         OmegaP={'I': Om_P,},
-                         OmegaPP={'I': 'string',},
+                         OmegaP={'pI': Om_P,},
+                         OmegaPP={'pI': 'string',},
                          beam_freqs=beam_freqs)
 
         nt.assert_raises(ValueError, pspecbeam.PSpecBeamFromArray,
-                         OmegaP={'I': Om_P}, 
-                         OmegaPP={'I': Om_PP[:-2],},
+                         OmegaP={'pI': Om_P}, 
+                         OmegaPP={'pI': Om_PP[:-2],},
                          beam_freqs=beam_freqs)
 
         # Check that invalid method args raise errors