Merge pull request #174 from HERA-Team/noise_sim_validation

Noise sim validation

hera_pspec/plot.py

@@ -593,7 +593,9 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True,
                 red_tol=1.0, center_line=False, horizon_lines=False,
                 title=None, ax=None, cmap='viridis', figsize=(8, 6),
                 deltasq=False, colorbar=False, cbax=None, vmin=None, vmax=None,
-                edgecolor='none', flip_xax=False, flip_yax=False, lw=2, **kwargs):
+                edgecolor='none', flip_xax=False, flip_yax=False, lw=2, set_bl_tick_major=False,
+                set_bl_tick_minor=False, xtick_size=10, xtick_rot=0, ytick_size=10, ytick_rot=0,
+                **kwargs):
     """
     Plot a 2D delay spectrum (or spectra) from a UVPSpec object. Note that
     all integrations and redundant baselines are averaged (unless specifying times)
@@ -693,6 +695,13 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True,
     lw : int, optional
         Line-width of horizon and center lines if plotted. Default: 2.
 
+    set_bl_tick_major : bool, optional
+        If True, use the baseline lengths as major ticks, rather than default uniform
+        grid.
+
+    set_bl_tick_minor : bool, optional
+        If True, use the baseline lengths as minor ticks, which have no labels.
+
     kwargs : dictionary
         Additional keyword arguments to pass to pcolormesh() call.
 
@@ -803,10 +812,16 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True,
                         vmin=vmin, vmax=vmax, **kwargs)
 
     # Configure ticks
-    if rotate:
-        ax.set_xticks(np.around(x_axis, 4))
-    else:
-        ax.set_yticks(np.around(y_axis, 4))
+    if set_bl_tick_major:
+        if rotate:
+            ax.set_xticks(map(lambda x: np.around(x, _get_sigfig(x)+2), x_axis))
+        else:
+            ax.set_yticks(map(lambda x: np.around(x, _get_sigfig(x)+2), y_axis))
+    if set_bl_tick_minor:
+        if rotate:
+            ax.set_xticks(map(lambda x: np.around(x, _get_sigfig(x)+2), x_axis), minor=True)
+        else:
+            ax.set_yticks(map(lambda x: np.around(x, _get_sigfig(x)+2), y_axis), minor=True)
 
     # Add colorbar
     if colorbar:
@@ -878,14 +893,22 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, fold=False, delay=True,
 
     # flip axes
     if flip_xax:
-        plt.gca().invert_xaxis()
+        fig.sca(ax)
+        fig.gca().invert_xaxis()
     if flip_yax:
-        plt.gca().invert_yaxis()
+        fig.sca(ax)
+        fig.gca().invert_yaxis()
 
     # add title
     if title is not None:
         ax.set_title(title, fontsize=12)
 
+    # Configure tick sizes and rotation
+    [tl.set_size(xtick_size) for tl in ax.get_xticklabels()]
+    [tl.set_rotation(xtick_rot) for tl in ax.get_xticklabels()]
+    [tl.set_size(ytick_size) for tl in ax.get_yticklabels()]
+    [tl.set_rotation(ytick_rot) for tl in ax.get_yticklabels()]
+
     # return figure
     if new_plot:
         return fig
@@ -974,3 +997,15 @@ def plot_uvdata_waterfalls(uvd, basename, data='data', plot_mode='log',
         fig.tight_layout()
         fig.savefig(outfile, format=format)
         fig.clf()
+
+def _get_sigfig(x):
+    return -int(np.floor(np.log10(np.abs(x))))
+
+def _round_sigfig(x, up=True):
+    sigfigs = get_sigfig(x)
+    if up:
+        return np.ceil(10**sigfigs * x) / 10**sigfigs
+    else:
+        return np.floor(10**sigfigs * x) / 10**sigfigs
+
+

hera_pspec/testing.py

@@ -5,6 +5,7 @@
 from hera_pspec import uvpspec, pspecdata, conversions, pspecbeam, utils
 from pyuvdata import UVData
 from hera_cal.utils import JD2LST
+from scipy import stats
 
 
 def build_vanilla_uvpspec(beam=None):
@@ -83,7 +84,7 @@ def build_vanilla_uvpspec(beam=None):
                                    2005235.09142983,
                                   -3239928.42475397])
 
-    store_cov=True
+    store_cov = True
     cosmo = conversions.Cosmo_Conversions()
 
     data_array, wgt_array = {}, {}
@@ -220,3 +221,122 @@ def uvpspec_from_data(data, bl_grps, data_std=None, spw_ranges=None,
                    spw_ranges=spw_ranges, taper=taper, verbose=verbose, 
                    store_cov=store_cov, n_dlys=n_dlys)
     return uvp
+
+
+def noise_sim(data, Tsys, beam=None, Nextend=0, seed=None, inplace=False,
+              whiten=False, run_check=True):
+    """
+    Generate a simulated Gaussian noise (visibility) realization given
+    a system temperature Tsys. If a primary beam model is not provided,
+    this is in units of Kelvin-steradians
+
+        Trms = Tsys / sqrt(channel_width * integration_time)
+
+    where Trms is divided by an additional sqrt(2) if the polarization
+    in data is a pseudo-Stokes polarization. If a primary beam model is
+    provided, the output is converted to Jansky.
+
+    Parameters
+    ----------
+    data : str or UVData object
+        A UVData object or path to miriad file.
+
+    Tsys : float
+        System temperature in Kelvin.
+
+    beam : str or PSpecBeam object, optional
+        A PSpecBeam object or path to beamfits file.
+
+    Nextend : int, optional
+        Number of times to extend time axis by default length
+        before creating noise sim. Can be used to increase
+        number statistics before forming noise realization.
+
+    seed : int, optional
+        Seed to set before forming noise realization.
+
+    inplace : bool, optional
+        If True, overwrite input data and return None, else
+        make a copy and return copy.
+
+    whiten : bool, optional
+        If True, clear input data of flags if they exist and
+        set all nsamples to 1.
+
+    run_check : bool, optional
+        If True, run UVData check before return.
+
+    Returns
+    -------
+    data : UVData with noise realizations.
+    """
+    # Read data files
+    if isinstance(data, (str, np.str)):
+        _data = UVData()
+        _data.read_miriad(data)
+        data = _data
+    elif isinstance(data, UVData):
+        if not inplace:
+            data = copy.deepcopy(data)
+    assert isinstance(data, UVData)
+
+    # whiten input data
+    if whiten:
+        data.flag_array[:] = False
+        data.nsample_array[:] = 1.0
+
+    # Configure beam
+    if beam is not None:
+        if isinstance(beam, (str, np.str)):
+            beam = pspecbeam.PSpecBeamUV(beam)
+        assert isinstance(beam, pspecbeam.PSpecBeamBase)    
+
+    # Extend times
+    Nextend = int(Nextend)
+    if Nextend > 0:
+        assert data.phase_type == 'drift', "data must be drift phased in order to extend along time axis"
+        data = copy.deepcopy(data)
+        _data = copy.deepcopy(data)
+        dt = np.median(np.diff(np.unique(_data.time_array)))
+        dl = np.median(np.diff(np.unique(_data.lst_array)))
+        for i in range(Nextend):
+            _data.time_array += dt * _data.Ntimes * (i+1)
+            _data.lst_array += dl * _data.Ntimes * (i+1)
+            _data.lst_array %= 2*np.pi
+            data += _data
+
+    # Get Trms
+    int_time = data.integration_time
+    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)
+
+    # if a pol is pStokes pol, divide by extra sqrt(2)
+    polcorr = np.array([np.sqrt(2) if p in [1, 2, 3, 4] else 1.0 for p in data.polarization_array])
+    Trms /= polcorr
+
+    # Get Vrms in Jy using beam
+    if beam is not None:
+        freqs = np.unique(data.freq_array)[None, 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([np.sqrt(2) if p in [1, 2, 3, 4] else 1.0 for p in data.polarization_array])
+        rms = K_to_Jy * Trms
+    else:
+        rms = Trms
+
+    # Generate noise
+    if seed is not None:
+        np.random.seed(seed)
+    data.data_array = (stats.norm.rvs(0, 1./np.sqrt(2), size=rms.size).reshape(rms.shape) \
+                       + 1j * stats.norm.rvs(0, 1./np.sqrt(2), size=rms.size).reshape(rms.shape) ) * rms
+    f = np.isnan(data.data_array) + np.isinf(data.data_array)
+    data.data_array[f] = np.nan
+    data.flag_array[f] = True
+
+    if run_check:
+        data.check()
+
+    if not inplace:
+        return data
+

hera_pspec/tests/test_grouping.py

@@ -3,11 +3,12 @@
 import numpy as np
 import os
 from hera_pspec.data import DATA_PATH
-from hera_pspec import uvpspec, conversions, parameter, pspecbeam, pspecdata, testing
+from hera_pspec import uvpspec, conversions, parameter, pspecbeam, pspecdata, testing, utils
 from hera_pspec import uvpspec_utils as uvputils
 from hera_pspec import grouping, container
 from pyuvdata import UVData
 from hera_cal import redcal
+import copy
 
 
 class Test_grouping(unittest.TestCase):
@@ -99,7 +100,7 @@ def test_sample_baselines(self):
         # Check that returned length is the same for groups too
         samp = grouping.sample_baselines(g1)
         self.assertEqual(len(g1), len(samp))
-    
+
     def test_bootstrap_average_blpairs(self):
         """
         Test bootstrap averaging over power spectra.
@@ -151,7 +152,7 @@ def test_bootstrap_average_blpairs(self):
             ps_avg = uvp_avg.get_data((0, blpair, 'xx'))
             ps_boot = uvp4[0].get_data((0, blpair, 'xx'))
             np.testing.assert_array_almost_equal(ps_avg, ps_boot)
-        
+
 def test_bootstrap_resampled_error():
     # generate a UVPSpec
     visfile = os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")
@@ -181,6 +182,48 @@ def test_bootstrap_resampled_error():
     if os.path.exists("uvp.h5"):
         os.remove("uvp.h5")
 
+
+def test_validate_bootstrap_errorbar():
+    """ This is used to test the bootstrapping code
+    against the gaussian noise visibility simulator.
+    The basic premise is that, if working properly,
+    gaussian noise pspectra divided by their bootstrapped
+    errorbars should have a standard deviation that
+    converges to 1. """
+    # get simulated noise in K-str
+    uvfile = os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")
+    Tsys = 300.0  # Kelvin
+
+    # generate complex gaussian noise
+    seed = 4
+    uvd1 = testing.noise_sim(uvfile, Tsys, seed=seed, whiten=True, inplace=False, Nextend=0)
+    seed = 5
+    uvd2 = testing.noise_sim(uvfile, Tsys, seed=seed, whiten=True, inplace=False, Nextend=0)
+
+    # form (auto) baseline-pairs from only 14.6m bls
+    reds, lens, angs = utils.get_reds(uvd1, pick_data_ants=True, bl_len_range=(10, 50),
+                                      bl_deg_range=(0, 180))
+    bls1, bls2 = utils.flatten(reds), utils.flatten(reds)
+
+    # setup PSpecData and form power psectra
+    ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd2)], wgts=[None, None])
+    uvp = ds.pspec(bls1, bls2, (0, 1), [('xx', 'xx')], input_data_weight='identity', norm='I',
+                   taper='none', sampling=False, little_h=False, spw_ranges=[(0, 50)], verbose=False)
+
+    # bootstrap resample
+    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))
+
+
 def test_bootstrap_run():
     # generate a UVPSpec and container
     visfile = os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")

hera_pspec/tests/test_noise.py

@@ -4,7 +4,7 @@
 import os
 import sys
 from hera_pspec.data import DATA_PATH
-from hera_pspec import uvpspec, conversions, parameter, pspecbeam, noise
+from hera_pspec import uvpspec, conversions, pspecdata, pspecbeam, noise, testing, utils
 import copy
 import h5py
 from collections import OrderedDict as odict
@@ -64,4 +64,50 @@ def test_calc_P_N(self):
         nt.assert_true(Dsq[1] < P_N)
 
 
+def test_noise_validation():
+    """
+    make sure that the noise.py code produces
+    correct noise 1-sigma amplitude using a
+    noise simulation.
+    """
+    # get simulated noise in Jy
+    bfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits')
+    beam = pspecbeam.PSpecBeamUV(bfile)
+    uvfile = os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")
+    Tsys = 300.0  # Kelvin
+
+    # generate noise
+    seed = 0
+    uvd = testing.noise_sim(uvfile, Tsys, beam, seed=seed, whiten=True, inplace=False, Nextend=9)
+
+    # get redundant baseline group
+    reds, lens, angs = utils.get_reds(uvd, pick_data_ants=True, bl_len_range=(10, 20),
+                                      bl_deg_range=(0, 1))
+    bls1, bls2, blps = utils.construct_blpairs(reds[0], exclude_auto_bls=True, exclude_permutations=True)
+
+    # setup PSpecData
+    ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None], beam=beam)
+    ds.Jy_to_mK()
+
+    # get pspec
+    uvp = ds.pspec(bls1, bls2, (0, 1), [('xx', 'xx')], input_data_weight='identity', norm='I',
+                   taper='none', sampling=False, little_h=True, spw_ranges=[(0, 50)], verbose=False)
+
+    # get noise spectra from one of the blpairs
+    P_N = uvp.generate_noise_spectra(0, 'xx', Tsys, blpairs=uvp.get_blpairs()[:1], num_steps=2000).values()[0][0, 0]
+
+    # get P_std of real spectra for each baseline across time axis
+    P_stds = np.array([np.std(uvp.get_data((0, bl, 'xx')).real, axis=1) for bl in uvp.get_blpairs()])
+
+    # get average P_std_avg and its standard error
+    P_std_avg = np.mean(P_stds)
+
+    # assert close to P_N: 2%
+    # This should be updated to be within standard error on P_std_avg
+    # when the spw_range-variable pspec amplitude bug is resolved
+    nt.assert_true(np.abs(P_std_avg - P_N) / P_N < 0.02)
+
+
+
+
 

hera_pspec/tests/test_plot.py

@@ -331,8 +331,8 @@ def test_delay_wedge(self):
         f4 = plot.delay_wedge(uvp, 0, 'xx', blpairs=None, times=None, fold=False, delay=False, component='abs', 
                               rotate=False, log10=True, loglog=False, red_tol=1.0,
                               center_line=True, horizon_lines=True, title='hello', ax=None, cmap='viridis',
-                              figsize=(8, 6), deltasq=False, colorbar=True, cbax=None, vmin=6, vmax=15,
-                              edgecolor='grey', flip_xax=True, flip_yax=True, lw=2)
+                              figsize=(8, 6), deltasq=True, colorbar=True, cbax=None, vmin=6, vmax=15,
+                              edgecolor='grey', flip_xax=True, flip_yax=True, lw=2, set_bl_tick_minor=True)
         plt.close()
 
         # feed axes, red_tol
@@ -343,7 +343,7 @@ def test_delay_wedge(self):
                         rotate=True, log10=True, loglog=False, red_tol=10.0,
                         center_line=False, horizon_lines=False, ax=ax, cmap='viridis',
                         figsize=(8, 6), deltasq=False, colorbar=True, cbax=cbax, vmin=None, vmax=None,
-                        edgecolor='none', flip_xax=False, flip_yax=False, lw=2)
+                        edgecolor='none', flip_xax=False, flip_yax=False, lw=2, set_bl_tick_major=True)
         plt.close()
 
         # test exceptions