Merge pull request #308 from HERA-Team/rng_update

Update to new numpy RNG API

.flake8

@@ -11,6 +11,8 @@ ignore =
     D401,
     # allow method names to be the same as python builtins
     A003,
+    # allow module names to be the same as python builtins
+    A005,
     # inline strong start-string without end-string. This is OK in the case of **kwargs in parameters.
     RST210,
     # Logging statement uses f-string.

CHANGELOG.rst

@@ -5,6 +5,23 @@ Changelog
 dev
 ===
 
+Added
+-----
+- Classes subclassed from ``SimulationComponent`` now have a ``is_randomized``
+  class attribute that informs the ``Simulator`` of whether it should provide
+  a ``BitGenerator`` to the class when simulating the component.
+  - Classes which use a random component should now have a ``rng`` attribute,
+    which should be treated in the same manner as other model parameters. In
+    other words, random states are now effectively treated as model parameters.
+
+Changed
+-------
+- All random number generation now uses the new ``numpy`` API.
+  - Rather than seed the global random state, a new ``BitGenerator`` is made
+    with whatever random seed is desired.
+  - The Simulator API has remained virtually unchanged, but the internal logic
+    that handles random state management has received a significant update.
+
 Deprecated
 ----------
 

hera_sim/components.py

@@ -44,6 +44,8 @@ class determine how to apply the effect simulated by
 
     #: Whether this systematic multiplies existing visibilities
     is_multiplicative: bool = False
+    #: Whether this systematic contains a randomized component
+    is_randomized: bool = False
     #: Whether the returned value is per-antenna, per-baseline, or the full array
     return_type: str | None = None
     # This isn't exactly safe, but different instances of a class should

hera_sim/eor.py

@@ -49,6 +49,7 @@ class NoiselikeEoR(EoR):
 
     _alias = ("noiselike_eor",)
     is_smooth_in_freq = False
+    is_randomized = True
     return_type = "per_baseline"
     attrs_to_pull = dict(
         bl_vec=None,
@@ -61,6 +62,7 @@ def __init__(
         max_delay: Optional[float] = None,
         fringe_filter_type: str = "tophat",
         fringe_filter_kwargs: Optional[dict] = None,
+        rng: np.random.Generator | None = None,
     ):
         fringe_filter_kwargs = fringe_filter_kwargs or {}
 
@@ -70,6 +72,7 @@ def __init__(
             max_delay=max_delay,
             fringe_filter_type=fringe_filter_type,
             fringe_filter_kwargs=fringe_filter_kwargs,
+            rng=rng,
         )
 
     def __call__(self, lsts, freqs, bl_vec, **kwargs):
@@ -84,13 +87,11 @@ def __call__(self, lsts, freqs, bl_vec, **kwargs):
             max_delay,
             fringe_filter_type,
             fringe_filter_kwargs,
+            rng,
         ) = self._extract_kwarg_values(**kwargs)
 
-        # make white noise in freq/time (original says in frate/freq, not sure why)
-        data = utils.gen_white_noise(size=(len(lsts), len(freqs)))
-
-        # scale data by EoR amplitude
-        data *= eor_amp
+        # Make white noise in time and frequency with the desired amplitude.
+        data = utils.gen_white_noise(size=(len(lsts), len(freqs)), rng=rng) * eor_amp
 
         # apply delay filter; default does nothing
         # TODO: find out why bl_len_ns is hardcoded as 1e10, also

hera_sim/foregrounds.py

@@ -39,6 +39,8 @@ class DiffuseForeground(Foreground):
         Keyword arguments and associated values to be passed to
         :func:`~hera_sim.utils.rough_fringe_filter`. Default is to use the
         following settings: ``fringe_filter_type : tophat``.
+    rng: np.random.Generator, optional
+        Random number generator.
 
     Notes
     -----
@@ -65,6 +67,7 @@ class DiffuseForeground(Foreground):
 
     _alias = ("diffuse_foreground",)
     is_smooth_in_freq = True
+    is_randomized = True
     return_type = "per_baseline"
     attrs_to_pull = dict(
         bl_vec=None,
@@ -76,6 +79,7 @@ def __init__(
         omega_p=None,
         delay_filter_kwargs=None,
         fringe_filter_kwargs=None,
+        rng=None,
     ):
         if delay_filter_kwargs is None:
             delay_filter_kwargs = {
@@ -91,6 +95,7 @@ def __init__(
             omega_p=omega_p,
             delay_filter_kwargs=delay_filter_kwargs,
             fringe_filter_kwargs=fringe_filter_kwargs,
+            rng=rng,
         )
 
     def __call__(self, lsts, freqs, bl_vec, **kwargs):
@@ -122,6 +127,7 @@ def __call__(self, lsts, freqs, bl_vec, **kwargs):
             omega_p,
             delay_filter_kwargs,
             fringe_filter_kwargs,
+            rng,
         ) = self._extract_kwarg_values(**kwargs)
 
         if Tsky_mdl is None:
@@ -147,7 +153,7 @@ def __call__(self, lsts, freqs, bl_vec, **kwargs):
         if np.isclose(np.linalg.norm(bl_vec), 0):
             return vis
 
-        vis *= utils.gen_white_noise(vis.shape)
+        vis *= utils.gen_white_noise(size=vis.shape, rng=rng)
 
         vis = utils.rough_fringe_filter(
             vis, lsts, freqs, bl_vec[0], **fringe_filter_kwargs
@@ -191,9 +197,12 @@ class PointSourceForeground(Foreground):
     reference_freq : float, optional
         Reference frequency used to make the point source flux densities
         chromatic, in units of GHz.
+    rng: np.random.Generator, optional
+        Random number generator.
     """
 
     _alias = ("pntsrc_foreground",)
+    is_randomized = True
     return_type = "per_baseline"
     attrs_to_pull = dict(
         bl_vec=None,
@@ -208,6 +217,7 @@ def __init__(
         spectral_index_mean=-1,
         spectral_index_std=0.5,
         reference_freq=0.15,
+        rng=None,
     ):
         super().__init__(
             nsrcs=nsrcs,
@@ -217,6 +227,7 @@ def __init__(
             spectral_index_mean=spectral_index_mean,
             spectral_index_std=spectral_index_std,
             reference_freq=reference_freq,
+            rng=rng,
         )
 
     def __call__(self, lsts, freqs, bl_vec, **kwargs):
@@ -259,17 +270,17 @@ def __call__(self, lsts, freqs, bl_vec, **kwargs):
             spectral_index_mean,
             spectral_index_std,
             f0,
+            rng,
         ) = self._extract_kwarg_values(**kwargs)
 
         # get baseline length (it should already be in ns)
         bl_len_ns = np.linalg.norm(bl_vec)
 
-        # randomly generate source RAs
-        ras = np.random.uniform(0, 2 * np.pi, nsrcs)
-
-        # draw spectral indices from normal distribution
-        spec_indices = np.random.normal(
-            spectral_index_mean, spectral_index_std, size=nsrcs
+        # Randomly generate source positions and spectral indices.
+        rng = rng or np.random.default_rng()
+        ras = rng.uniform(0, 2 * np.pi, nsrcs)
+        spec_indices = rng.normal(
+            loc=spectral_index_mean, scale=spectral_index_std, size=nsrcs
         )
 
         # calculate beam width, hardcoded for HERA
@@ -278,44 +289,44 @@ def __call__(self, lsts, freqs, bl_vec, **kwargs):
         # draw flux densities from a power law
         alpha = beta + 1
         flux_densities = (
-            Smax**alpha + Smin**alpha * (1 - np.random.uniform(size=nsrcs))
+            Smax**alpha + Smin**alpha * (1 - rng.uniform(size=nsrcs))
         ) ** (1 / alpha)
 
         # initialize the visibility array
         vis = np.zeros((lsts.size, freqs.size), dtype=complex)
 
-        # iterate over ra, flux, spectral indices
+        # Compute the visibility source-by-source.
         for ra, flux, index in zip(ras, flux_densities, spec_indices):
-            # find which lst index to use?
+            # Figure out when the source crosses the meridian.
             lst_ind = np.argmin(np.abs(utils.compute_ha(lsts, ra)))
 
-            # slight offset in delay? why??
-            dtau = np.random.uniform(-1, 1) * 0.1 * bl_len_ns
+            # This is effectively baking in that up to 10% of the baseline
+            # length is oriented along the North-South direction, since this
+            # is the delay measured when the source transits the meridian.
+            # (Still need to think more carefully about this, but this seems
+            # like the right explanation for this, as well as the factor of
+            # 0.9 in the calculation of w further down.)
+            dtau = rng.uniform(-1, 1) * 0.1 * bl_len_ns
 
-            # fill in the corresponding region of the visibility array
+            # Add the contribution from the source as it transits the meridian.
             vis[lst_ind, :] += flux * (freqs / f0) ** index
-
-            # now multiply in the phase
             vis[lst_ind, :] *= np.exp(2j * np.pi * freqs * dtau)
 
-        # get hour angles for lsts at 0 RA (why?)
+        # Figure out the hour angles to use for computing the beam kernel.
         has = utils.compute_ha(lsts, 0)
 
         # convolve vis with beam at each frequency
         for j, freq in enumerate(freqs):
-            # first calculate the beam, using truncated Gaussian model
+            # Treat the beam as if it's a Gaussian with a sharp horizon.
             beam = np.exp(-(has**2) / (2 * beam_width[j] ** 2))
             beam = np.where(np.abs(has) > np.pi / 2, 0, beam)
 
-            # who the hell knows what this does
+            # Compute the phase evolution as the source transits the sky.
             w = 0.9 * bl_len_ns * np.sin(has) * freq
-
             phase = np.exp(2j * np.pi * w)
 
-            # define the convolving kernel
+            # Now actually apply the mock source transit.
             kernel = beam * phase
-
-            # now actually convolve the kernel and the raw vis
             vis[:, j] = np.fft.ifft(np.fft.fft(kernel) * np.fft.fft(vis[:, j]))
 
         return vis

hera_sim/noise.py

@@ -47,16 +47,17 @@ class ThermalNoise(Noise):
         Antenna numbers for the baseline that noise is being simulated for. This is
         just used to determine whether to simulate noise via the radiometer equation
         or to just add a bias from the receiver temperature.
+    rng: np.random.Generator, optional
+        Random number generator.
 
     Notes
     -----
-    At the time of writing, we're unsure of the correct prescription for
-    autocorrelations, so we only add a receiver temperature bias to baselines that
-    are interpreted as autocorrelations (i.e. where the ``bl_vec`` parameter provided
-    on calling an instance of this class is nearly zero length).
+    Considering the SNR in autocorrelations is typically very high, we only add
+    a receiver temperature bias to the autocorrelations.
     """
 
     _alias = ("thermal_noise",)
+    is_randomized = True
     return_type = "per_baseline"
     attrs_to_pull = dict(
         autovis=None,
@@ -72,6 +73,7 @@ def __init__(
         Trx=0,
         autovis=None,
         antpair=None,
+        rng=None,
     ):
         super().__init__(
             Tsky_mdl=Tsky_mdl,
@@ -81,6 +83,7 @@ def __init__(
             Trx=Trx,
             autovis=autovis,
             antpair=antpair,
+            rng=rng,
         )
 
     def __call__(self, lsts: np.ndarray, freqs: np.ndarray, **kwargs):
@@ -119,6 +122,7 @@ def __call__(self, lsts: np.ndarray, freqs: np.ndarray, **kwargs):
             Trx,
             autovis,
             antpair,
+            rng,
         ) = self._extract_kwarg_values(**kwargs)
 
         # get the channel width in Hz if not specified
@@ -172,7 +176,7 @@ def __call__(self, lsts: np.ndarray, freqs: np.ndarray, **kwargs):
         vis /= utils.jansky_to_kelvin(freqs, omega_p).reshape(1, -1)
 
         # make it noisy
-        return utils.gen_white_noise(size=vis.shape) * vis
+        return utils.gen_white_noise(size=vis.shape, rng=rng) * vis
 
     @staticmethod
     def resample_Tsky(lsts, freqs, Tsky_mdl=None, Tsky=180.0, mfreq=0.18, index=-2.5):