diff --git a/.flake8 b/.flake8 index 04fb6718..61006428 100644 --- a/.flake8 +++ b/.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. diff --git a/CHANGELOG.rst b/CHANGELOG.rst index de32cff9..1442f0ac 100644 --- a/CHANGELOG.rst +++ b/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 ---------- diff --git a/hera_sim/components.py b/hera_sim/components.py index 3f158759..79eee7ec 100644 --- a/hera_sim/components.py +++ b/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 diff --git a/hera_sim/eor.py b/hera_sim/eor.py index 263ec7c2..20f05f39 100644 --- a/hera_sim/eor.py +++ b/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 diff --git a/hera_sim/foregrounds.py b/hera_sim/foregrounds.py index f8e54863..12fcb45e 100644 --- a/hera_sim/foregrounds.py +++ b/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 diff --git a/hera_sim/noise.py b/hera_sim/noise.py index c952ca7f..b7dc2bad 100644 --- a/hera_sim/noise.py +++ b/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): diff --git a/hera_sim/rfi.py b/hera_sim/rfi.py index b0d62996..ffe28ea9 100644 --- a/hera_sim/rfi.py +++ b/hera_sim/rfi.py @@ -36,6 +36,8 @@ class RfiStation: are considered "off" and high points are considered "on" (just how high is controlled by ``duty_cycle``). This is the wavelength (in seconds) of that cycle. + rng: np.random.Generator, optional + Random number generator. Notes ----- @@ -53,12 +55,14 @@ def __init__( strength: float = 100.0, std: float = 10.0, timescale: float = 100.0, + rng: np.random.Generator | None = None, ): self.f0 = f0 self.duty_cycle = duty_cycle self.strength = strength self.std = std self.timescale = timescale + self.rng = rng or np.random.default_rng() def __call__(self, lsts, freqs): """Compute the RFI for this station. @@ -96,14 +100,14 @@ def __call__(self, lsts, freqs): ch2 = ch1 + 1 if self.f0 > freqs[ch1] else ch1 - 1 # generate some random phases - phs1, phs2 = np.random.uniform(0, 2 * np.pi, size=2) + phs1, phs2 = self.rng.uniform(0, 2 * np.pi, size=2) # find out when the station is broadcasting is_on = 0.999 * np.cos(lsts * u.sday.to("s") / self.timescale + phs1) is_on = is_on > (1 - 2 * self.duty_cycle) # generate a signal and filter it according to when it's on - signal = np.random.normal(self.strength, self.std, lsts.size) + signal = self.rng.normal(self.strength, self.std, lsts.size) signal = np.where(is_on, signal, 0) * np.exp(1j * phs2) # now add the signal to the rfi array @@ -127,13 +131,16 @@ class Stations(RFI): ---------- stations : list of :class:`RfiStation` The list of stations that produce RFI. + rng: np.random.Generator, optional + Random number generator. """ _alias = ("rfi_stations",) + is_randomized = True return_type = "per_baseline" - def __init__(self, stations=None): - super().__init__(stations=stations) + def __init__(self, stations=None, rng=None): + super().__init__(stations=stations, rng=rng) def __call__(self, lsts, freqs, **kwargs): """Generate the RFI from all stations. @@ -159,7 +166,7 @@ def __call__(self, lsts, freqs, **kwargs): self._check_kwargs(**kwargs) # but this is where the magic comes in (thanks to defaults) - (stations,) = self._extract_kwarg_values(**kwargs) + (stations, rng) = self._extract_kwarg_values(**kwargs) # initialize an array to store the rfi in rfi = np.zeros((lsts.size, freqs.size), dtype=complex) @@ -202,14 +209,19 @@ class Impulse(RFI): Strength of the impulse. This will not be randomized, though a phase offset as a function of frequency will be applied, and will be random for each impulse. + rng: np.random.Generator, optional + Random number generator. """ _alias = ("rfi_impulse",) + is_randomized = True return_type = "per_baseline" - def __init__(self, impulse_chance=0.001, impulse_strength=20.0): + def __init__(self, impulse_chance=0.001, impulse_strength=20.0, rng=None): super().__init__( - impulse_chance=impulse_chance, impulse_strength=impulse_strength + impulse_chance=impulse_chance, + impulse_strength=impulse_strength, + rng=rng, ) def __call__(self, lsts, freqs, **kwargs): @@ -231,18 +243,19 @@ def __call__(self, lsts, freqs, **kwargs): self._check_kwargs(**kwargs) # unpack the kwargs - chance, strength = self._extract_kwarg_values(**kwargs) + chance, strength, rng = self._extract_kwarg_values(**kwargs) + rng = rng or np.random.default_rng() # initialize the rfi array rfi = np.zeros((lsts.size, freqs.size), dtype=complex) # find times when an impulse occurs - impulses = np.where(np.random.uniform(size=lsts.size) <= chance)[0] + impulses = np.where(rng.uniform(size=lsts.size) <= chance)[0] # only do something if there are impulses if impulses.size > 0: # randomly generate some delays for each impulse - dlys = np.random.uniform(-300, 300, impulses.size) # ns + dlys = rng.uniform(-300, 300, impulses.size) # ns # generate the signals signals = strength * np.asarray( @@ -266,16 +279,22 @@ class Scatter(RFI): random strength). scatter_std : float, optional Standard deviation of the RFI strength. + rng: np.random.Generator, optional + Random number generator. """ _alias = ("rfi_scatter",) + is_randomized = True return_type = "per_baseline" - def __init__(self, scatter_chance=0.0001, scatter_strength=10.0, scatter_std=10.0): + def __init__( + self, scatter_chance=0.0001, scatter_strength=10.0, scatter_std=10.0, rng=None + ): super().__init__( scatter_chance=scatter_chance, scatter_strength=scatter_strength, scatter_std=scatter_std, + rng=rng, ) def __call__(self, lsts, freqs, **kwargs): @@ -297,17 +316,18 @@ def __call__(self, lsts, freqs, **kwargs): self._check_kwargs(**kwargs) # now unpack them - chance, strength, std = self._extract_kwarg_values(**kwargs) + chance, strength, std, rng = self._extract_kwarg_values(**kwargs) + rng = rng or np.random.default_rng() # make an empty rfi array rfi = np.zeros((lsts.size, freqs.size), dtype=complex) # find out where to put the rfi - rfis = np.where(np.random.uniform(size=rfi.size) <= chance)[0] + rfis = np.where(rng.uniform(size=rfi.size) <= chance)[0] # simulate the rfi; one random amplitude, all random phases - signal = np.random.normal(strength, std) * np.exp( - 2j * np.pi * np.random.uniform(size=rfis.size) + signal = rng.normal(strength, std) * np.exp( + 2j * np.pi * rng.uniform(size=rfis.size) ) # add the signal to the rfi @@ -333,9 +353,12 @@ class DTV(RFI): Mean strength of RFI. dtv_std : float, optional Standard deviation of RFI strength. + rng: np.random.Generator, optional + Random number generator. """ _alias = ("rfi_dtv",) + is_randomized = True return_type = "per_baseline" def __init__( @@ -345,6 +368,7 @@ def __init__( dtv_chance=0.0001, dtv_strength=10.0, dtv_std=10.0, + rng=None, ): super().__init__( dtv_band=dtv_band, @@ -352,6 +376,7 @@ def __init__( dtv_chance=dtv_chance, dtv_strength=dtv_strength, dtv_std=dtv_std, + rng=rng, ) def __call__(self, lsts, freqs, **kwargs): @@ -379,7 +404,9 @@ def __call__(self, lsts, freqs, **kwargs): dtv_chance, dtv_strength, dtv_std, + rng, ) = self._extract_kwarg_values(**kwargs) + rng = rng or np.random.default_rng() # make an empty rfi array rfi = np.zeros((lsts.size, freqs.size), dtype=complex) @@ -449,12 +476,12 @@ def __call__(self, lsts, freqs, **kwargs): this_rfi = rfi[:, ch1:ch2] # find out which times are affected - rfis = np.random.uniform(size=lsts.size) <= chance + rfis = rng.uniform(size=lsts.size) <= chance # calculate the signal signal = np.atleast_2d( - np.random.normal(strength, std, size=rfis.sum()) - * np.exp(2j * np.pi * np.random.uniform(size=rfis.sum())) + rng.normal(strength, std, size=rfis.sum()) + * np.exp(2j * np.pi * rng.uniform(size=rfis.sum())) ).T # add the signal to the rfi array diff --git a/hera_sim/sigchain.py b/hera_sim/sigchain.py index f5f44b91..ad09ce57 100644 --- a/hera_sim/sigchain.py +++ b/hera_sim/sigchain.py @@ -15,7 +15,6 @@ from pathlib import Path from pyuvdata import UVBeam from pyuvsim import AnalyticBeam -from scipy import stats from scipy.signal.windows import blackmanharris from typing import Callable @@ -58,10 +57,13 @@ class Bandpass(Gain): Taper to apply to the simulated gains. Default is to not apply a taper. taper_kwds Keyword arguments used in generating the taper. + rng + Random number generator. """ _alias = ("gains", "bandpass_gain") is_multiplicative = True + is_randomized = True return_type = "per_antenna" attrs_to_pull = dict( ants="antpos", @@ -74,6 +76,7 @@ def __init__( bp_poly: str | callable | np.ndarray | None = None, taper: str | callable | np.ndarray | None = None, taper_kwds: dict | None = None, + rng: np.random.Generator | None = None, ): super().__init__( gain_spread=gain_spread, @@ -81,6 +84,7 @@ def __init__( bp_poly=bp_poly, taper=taper, taper_kwds=taper_kwds, + rng=rng, ) def __call__(self, freqs, ants, **kwargs): @@ -103,15 +107,21 @@ def __call__(self, freqs, ants, **kwargs): self._check_kwargs(**kwargs) # unpack the kwargs - (gain_spread, dly_rng, bp_poly, taper, taper_kwds) = self._extract_kwarg_values( - **kwargs - ) + ( + gain_spread, + dly_rng, + bp_poly, + taper, + taper_kwds, + rng, + ) = self._extract_kwarg_values(**kwargs) + rng = rng or np.random.default_rng() # get the bandpass gains - bandpass = self._gen_bandpass(freqs, ants, gain_spread, bp_poly) + bandpass = self._gen_bandpass(freqs, ants, gain_spread, bp_poly, rng=rng) # get the delay phases - phase = self._gen_delay_phase(freqs, ants, dly_rng) + phase = self._gen_delay_phase(freqs, ants, dly_rng, rng=rng) if taper is None: taper = np.ones(freqs.size) @@ -138,7 +148,7 @@ def __call__(self, freqs, ants, **kwargs): return {ant: bandpass[ant] * phase[ant] * taper for ant in ants} @_defaults - def _gen_bandpass(self, freqs, ants, gain_spread=0.1, bp_poly=None): + def _gen_bandpass(self, freqs, ants, gain_spread=0.1, bp_poly=None, rng=None): if bp_poly is None: # default to the H1C bandpass bp_poly = np.load(DATA_PATH / "HERA_H1C_BANDPASS.npy") @@ -155,15 +165,16 @@ def _gen_bandpass(self, freqs, ants, gain_spread=0.1, bp_poly=None): gains = {} for ant in ants: delta_bp = np.fft.ifft( - utils.gen_white_noise(freqs.size) * modes * gain_spread + utils.gen_white_noise(freqs.size, rng=rng) * modes * gain_spread ) gains[ant] = bp_base + delta_bp return gains - def _gen_delay_phase(self, freqs, ants, dly_rng=(-20, 20)): + def _gen_delay_phase(self, freqs, ants, dly_rng=(-20, 20), rng=None): phases = {} + rng = rng or np.random.default_rng() for ant in ants: - delay = np.random.uniform(*dly_rng) + delay = rng.uniform(*dly_rng) phases[ant] = np.exp(2j * np.pi * delay * freqs) return phases @@ -187,17 +198,27 @@ class Reflections(Gain): dly_jitter : float, optional Final delays are offset by a normal variable with mean zero and standard deviation ``dly_jitter``. + rng: np.random.Generator, optional + Random number generator. """ _alias = ("reflection_gains", "sigchain_reflections") is_multiplicative = True + is_randomized = True return_type = "per_antenna" attrs_to_pull = dict( ants="antpos", ) def __init__( - self, amp=None, dly=None, phs=None, conj=False, amp_jitter=0, dly_jitter=0 + self, + amp=None, + dly=None, + phs=None, + conj=False, + amp_jitter=0, + dly_jitter=0, + rng=None, ): super().__init__( amp=amp, @@ -206,6 +227,7 @@ def __init__( conj=conj, amp_jitter=amp_jitter, dly_jitter=dly_jitter, + rng=rng, ) def __call__(self, freqs, ants, **kwargs): @@ -228,13 +250,14 @@ def __call__(self, freqs, ants, **kwargs): self._check_kwargs(**kwargs) # unpack the kwargs - amp, dly, phs, conj, amp_jitter, dly_jitter = self._extract_kwarg_values( + amp, dly, phs, conj, amp_jitter, dly_jitter, rng = self._extract_kwarg_values( **kwargs ) + rng = rng or np.random.default_rng() # fill in missing kwargs amp, dly, phs = self._complete_params( - ants, amp, dly, phs, amp_jitter, dly_jitter + ants, amp, dly, phs, amp_jitter, dly_jitter, rng=rng ) # determine gains iteratively @@ -309,7 +332,7 @@ def _type_check(arr): @staticmethod def _complete_params( - ants, amp=None, dly=None, phs=None, amp_jitter=0, dly_jitter=0 + ants, amp=None, dly=None, phs=None, amp_jitter=0, dly_jitter=0, rng=None ): # TODO: docstring isn't exactly accurate, should be updated """ @@ -344,6 +367,8 @@ def _complete_params( For example, setting this to 10 will introduce, on average, delay deviations up to 10 ns. (This is drawn from a normal distribution, so it is possible that delays will exceed the value provided.) + rng: np.random.Generator, optional + Random number generator. Returns ------- @@ -355,16 +380,18 @@ def _complete_params( Phase of each reflection coefficient for each antenna. """ + rng = rng or np.random.default_rng() + def broadcast_param(param, lower_bound, upper_bound, size): if param is None: - return stats.uniform.rvs(lower_bound, upper_bound, size) + return rng.uniform(lower_bound, upper_bound, size) elif np.isscalar(param): return np.ones(size, dtype=float) * param else: if len(param) == size: return np.array(param, dtype=float) else: - return stats.uniform.rvs(*param, size) + return rng.uniform(*param, size) # Transform parameters into arrays. amps = broadcast_param(amp, 0, 1, len(ants)) @@ -372,8 +399,8 @@ def broadcast_param(param, lower_bound, upper_bound, size): phases = broadcast_param(phs, -np.pi, np.pi, len(ants)) # Apply jitter. - amps *= stats.norm.rvs(1, amp_jitter, len(ants)) - dlys += stats.norm.rvs(0, dly_jitter, len(ants)) + amps *= rng.normal(1, amp_jitter, len(ants)) + dlys += rng.normal(0, dly_jitter, len(ants)) return amps, dlys, phases @@ -402,6 +429,8 @@ class ReflectionSpectrum(Gain): Absolute jitter in delay across antennas for each of the reflections. amp_logbase Base of the logarithm to use for generating reflection amplitudes. + rng + Random number generator. Notes ----- @@ -411,6 +440,7 @@ class ReflectionSpectrum(Gain): _alias = ("reflection_spectrum",) is_multiplicative = True + is_randomized = True return_type = "per_antenna" attrs_to_pull = dict( ants="antpos", @@ -425,6 +455,7 @@ def __init__( amp_jitter: float = 0.05, dly_jitter: float = 30, amp_logbase: float = 10, + rng: np.random.Generator | None = None, ): super().__init__( n_copies=n_copies, @@ -434,6 +465,7 @@ def __init__( amp_jitter=amp_jitter, dly_jitter=dly_jitter, amp_logbase=amp_logbase, + rng=rng, ) def __call__( @@ -463,11 +495,13 @@ def __call__( amp_jitter, dly_jitter, amp_logbase, + rng, ) = self._extract_kwarg_values(**kwargs) + rng = rng or np.random.default_rng() amps = np.logspace(*amp_range, n_copies, base=amp_logbase) dlys = np.linspace(*dly_range, n_copies) - phases = np.random.uniform(*phs_range, n_copies) + phases = rng.uniform(*phs_range, n_copies) reflection_gains = {ant: np.ones(freqs.size, dtype=complex) for ant in ants} for amp, dly, phs in zip(amps, dlys, phases): @@ -477,6 +511,7 @@ def __call__( phs=phs, amp_jitter=amp_jitter, dly_jitter=dly_jitter, + rng=rng, ) reflections = reflections(freqs, ants) for ant, reflection in reflections.items(): @@ -511,17 +546,27 @@ class CrossCouplingCrosstalk(Crosstalk, Reflections): dly_jitter : float, optional Final delays are offset by a normal variable with mean zero and standard deviation ``dly_jitter``. + rng : np.random.Generator, optional + Random number generator. """ _alias = ("cross_coupling_xtalk",) is_multiplicative = False + is_randomized = True return_type = "per_baseline" attrs_to_pull = dict( autovis=None, ) def __init__( - self, amp=None, dly=None, phs=None, conj=False, amp_jitter=0, dly_jitter=0 + self, + amp=None, + dly=None, + phs=None, + conj=False, + amp_jitter=0, + dly_jitter=0, + rng=None, ): super().__init__( amp=amp, @@ -530,6 +575,7 @@ def __init__( conj=conj, amp_jitter=amp_jitter, dly_jitter=dly_jitter, + rng=rng, ) def __call__(self, freqs, autovis, **kwargs): @@ -552,13 +598,14 @@ def __call__(self, freqs, autovis, **kwargs): self._check_kwargs(**kwargs) # now unpack them - amp, dly, phs, conj, amp_jitter, dly_jitter = self._extract_kwarg_values( + amp, dly, phs, conj, amp_jitter, dly_jitter, rng = self._extract_kwarg_values( **kwargs ) + rng = rng or np.random.default_rng() # handle the amplitude, phase, and delay amp, dly, phs = self._complete_params( - [1], amp, dly, phs, amp_jitter, dly_jitter + [1], amp, dly, phs, amp_jitter, dly_jitter, rng=rng ) # Make reflection coefficient. @@ -602,6 +649,8 @@ class CrossCouplingSpectrum(Crosstalk): Whether to also produce statistically equivalent cross-talk at negative delays. Note that while the statistics are equivalent, both amplitudes and delays will be different random realizations. + rng : np.random.Generator, optional + Random number generator. Notes ----- @@ -610,6 +659,7 @@ class CrossCouplingSpectrum(Crosstalk): """ _alias = ("cross_coupling_spectrum", "xtalk_spectrum") + is_randomized = True return_type = "per_baseline" attrs_to_pull = dict( autovis=None, @@ -625,6 +675,7 @@ def __init__( dly_jitter=0, amp_logbase=10, symmetrize=True, + rng=None, ): super().__init__( n_copies=n_copies, @@ -635,6 +686,7 @@ def __init__( dly_jitter=dly_jitter, amp_logbase=amp_logbase, symmetrize=symmetrize, + rng=rng, ) def __call__(self, freqs, autovis, **kwargs): @@ -664,6 +716,7 @@ def __call__(self, freqs, autovis, **kwargs): dly_jitter, amp_logbase, symmetrize, + rng, ) = self._extract_kwarg_values(**kwargs) # Construct the arrays of amplitudes and delays. @@ -679,6 +732,7 @@ def __call__(self, freqs, autovis, **kwargs): phs=phs_range, amp_jitter=amp_jitter, dly_jitter=dly_jitter, + rng=rng, ) crosstalk_spectrum += gen_xtalk(freqs, autovis) @@ -1280,12 +1334,15 @@ class OverAirCrossCoupling(Crosstalk): Ratio of the amplitude of the last peak in the cross-coupling spectrum to the first peak. In other words, how much the cross-coupling spectrum decays over the full range of delays it covers. + rng + Random number generator. See Also -------- :class:`CrossCouplingSpectrum` """ + is_randomized = True return_type = "per_baseline" attrs_to_pull = dict( antpair=None, @@ -1306,6 +1363,7 @@ def __init__( dly_jitter: float = 0, max_delay: float = 2000, amp_decay_fac: float = 1e-2, + rng: np.random.Generator | None = None, ): super().__init__( emitter_pos=emitter_pos, @@ -1319,6 +1377,7 @@ def __init__( dly_jitter=dly_jitter, max_delay=max_delay, amp_decay_fac=amp_decay_fac, + rng=rng, ) def __call__( @@ -1365,6 +1424,7 @@ def __call__( dly_jitter, max_delay, amp_decay_fac, + rng, ) = self._extract_kwarg_values(**kwargs) ai, aj = antpair @@ -1396,6 +1456,7 @@ def log(x): dly_jitter=dly_jitter, amp_logbase=amp_decay_base, symmetrize=False, + rng=rng, ) xt_ji = CrossCouplingSpectrum( n_copies=n_copies, @@ -1405,6 +1466,7 @@ def log(x): dly_jitter=dly_jitter, amp_logbase=amp_decay_base, symmetrize=False, + rng=rng, ) return xt_ij(freqs, autovis_i) + xt_ji(freqs, autovis_j) @@ -1417,16 +1479,19 @@ class WhiteNoiseCrosstalk(Crosstalk): ---------- amplitude : float, optional The amplitude of the white noise spectrum (i.e. its standard deviation). + rng : np.random.Generator, optional + Random number generator. """ _alias = ( "whitenoise_xtalk", "white_noise_xtalk", ) + is_randomized = True return_type = "per_baseline" - def __init__(self, amplitude=3.0): - super().__init__(amplitude=amplitude) + def __init__(self, amplitude=3.0, rng=None): + super().__init__(amplitude=amplitude, rng=rng) def __call__(self, freqs, **kwargs): """Compute the cross-correlations. @@ -1446,13 +1511,13 @@ def __call__(self, freqs, **kwargs): self._check_kwargs(**kwargs) # unpack the kwargs - (amplitude,) = self._extract_kwarg_values(**kwargs) + (amplitude, rng) = self._extract_kwarg_values(**kwargs) # why choose this size for the convolving kernel? kernel = np.ones(50 if freqs.size > 50 else int(freqs.size / 2)) # generate the crosstalk - xtalk = np.convolve(utils.gen_white_noise(freqs.size), kernel, "same") + xtalk = np.convolve(utils.gen_white_noise(freqs.size, rng=rng), kernel, "same") # scale the result and return return amplitude * xtalk @@ -1510,6 +1575,7 @@ def vary_gains_in_time( variation_timescale=None, variation_amp=0.05, variation_mode="linear", + rng=None, ): r""" Vary gain amplitudes, phases, or delays in time. @@ -1563,6 +1629,8 @@ def vary_gains_in_time( mode produces a triangle wave variation with period twice the corresponding timescale; this ensures that the gains vary linearly over the entire set of provided times if the default variation timescale is used. + rng: np.random.Generator, optional + Random number generator. Returns ------- @@ -1645,7 +1713,8 @@ def vary_gains_in_time( elif mode == "sinusoidal": envelope *= 1 + amp * np.sin(2 * np.pi * phases) elif mode == "noiselike": - envelope *= stats.norm.rvs(1, amp, times.size) + rng = rng or np.random.default_rng() + envelope *= rng.normal(1, amp, times.size) else: raise NotImplementedError(f"Variation mode {mode!r} not supported.") diff --git a/hera_sim/simulate.py b/hera_sim/simulate.py index 9f10abd4..8b6e6ab8 100644 --- a/hera_sim/simulate.py +++ b/hera_sim/simulate.py @@ -10,7 +10,6 @@ import functools import inspect import numpy as np -import time import warnings import yaml from astropy import constants as const @@ -296,6 +295,14 @@ def add( stacklevel=2, ) + # Record the component simulated and the parameters used. + if defaults._override_defaults: + for param in getattr(model, "kwargs", {}): + if param not in kwargs and param in defaults(): + kwargs[param] = defaults(param) + self._components[model_key] = kwargs.copy() + self._components[model_key]["alias"] = component + # Simulate the effect by iterating over baselines and polarizations. data = self._iteratively_apply( model, @@ -304,25 +311,22 @@ def add( vis_filter=vis_filter, antpairpol_cache=self._antpairpol_cache[model_key], seed=seed, + model_key=model_key, **kwargs, ) # This is None if ret_vis is False if add_vis: - # Record the component simulated and the parameters used. - if defaults._override_defaults: - for param in getattr(model, "kwargs", {}): - if param not in kwargs and param in defaults(): - kwargs[param] = defaults(param) self._update_history(model, **kwargs) if seed: - kwargs["seed"] = seed + self._components[model_key]["seed"] = seed self._update_seeds(model_key) if vis_filter is not None: kwargs["vis_filter"] = vis_filter - self._components[model_key] = kwargs - self._components[model_key]["alias"] = component else: del self._antpairpol_cache[model_key] + del self._components[model_key] + if self._seeds.get(model_key, None): + del self._seeds[model_key] return data @@ -396,6 +400,7 @@ def get( ret_vis=True, seed=seed, vis_filter=vis_filter, + model_key=model_key, **kwargs, ) if ant1 is not None: @@ -417,6 +422,7 @@ def get( seed=seed, vis_filter=vis_filter, antpairpol_cache=None, + model_key=model_key, **kwargs, ) @@ -427,6 +433,7 @@ def get( return data[:, :, pol_ind] # We're only simulating for a particular baseline. + # (The validation check ensures this is the case.) # First, find out if it needs to be conjugated. try: blt_inds = self.data.antpair2ind(ant1, ant2) @@ -437,7 +444,7 @@ def get( blt_inds = self.data.antpair2ind(ant2, ant1) conj_data = True - # We've got three different seeding cases to work out. + # We have three different seeding cases to work out. if seed == "initial": # Initial seeding means we need to do the whole array. data = self._iteratively_apply( @@ -447,6 +454,7 @@ def get( seed=seed, vis_filter=vis_filter, antpairpol_cache=None, + model_key=model_key, **kwargs, )[blt_inds, :, :] if conj_data: # pragma: no cover @@ -455,15 +463,8 @@ def get( return data pol_ind = self.data.get_pols().index(pol) return data[..., pol_ind] - elif seed == "redundant": - if conj_data: - self._seed_rng(seed, model, ant2, ant1, pol) - else: - self._seed_rng(seed, model, ant1, ant2, pol) - elif seed is not None: - self._seed_rng(seed, model, ant1, ant2, pol) - # Prepare the model parameters, then simulate and return the effect. + # Figure out whether we need to do a polarization selection. if pol is None: data_shape = (self.lsts.size, self.freqs.size, len(self.pols)) pols = self.pols @@ -472,15 +473,22 @@ def get( data_shape = (self.lsts.size, self.freqs.size, 1) pols = (pol,) return_slice = (slice(None), slice(None), 0) + + # Prepare the model parameters, then simulate and return the effect. data = np.zeros(data_shape, dtype=complex) for i, _pol in enumerate(pols): args = self._initialize_args_from_model(model) args = self._update_args(args, model, ant1, ant2, pol) args.update(kwargs) if conj_data: - self._seed_rng(seed, model, ant2, ant1, _pol) + _, rng = self._seed_rng( + seed, model, ant2, ant1, _pol, model_key=model_key + ) else: - self._seed_rng(seed, model, ant1, ant2, _pol) + _, rng = self._seed_rng( + seed, model, ant1, ant2, _pol, model_key=model_key + ) + args["rng"] = rng data[..., i] = model(**args) if conj_data: data = np.conj(data) @@ -959,11 +967,12 @@ def _iteratively_apply( *, add_vis: bool = True, ret_vis: bool = False, - seed: Optional[Union[str, int]] = None, - vis_filter: Optional[Sequence] = None, - antpairpol_cache: Optional[Sequence[AntPairPol]] = None, + seed: str | int | None = None, + vis_filter: Sequence | None = None, + antpairpol_cache: Sequence[AntPairPol] | None = None, + model_key: str | None = None, **kwargs, - ) -> Optional[Union[np.ndarray, dict[int, np.ndarray]]]: + ) -> Union[np.ndarray, dict[int, np.ndarray]] | None: """ Simulate an effect for an entire array. @@ -998,6 +1007,10 @@ def _iteratively_apply( List of (ant1, ant2, pol) tuples specifying which antpairpols have already had the effect simulated. Not intended for use by the typical end-user. + model_key + String identifying the model component being computed. This is + handed around to ensure that random number generation schemes using + the "initial" seeding routine can be recovered via ``self.get``. kwargs Extra parameters passed to ``model``. @@ -1014,7 +1027,7 @@ def _iteratively_apply( warnings.warn( "You have chosen to neither add nor return the effect " "you are trying to simulate, so nothing will be " - f"computed. This warning was raised for the model: {model}", + f"computed. This warning was raised for the model: {model_key}", stacklevel=2, ) return @@ -1050,7 +1063,10 @@ def _iteratively_apply( args.update(kwargs) for pol in self.data.get_feedpols(): if seed: - seed = self._seed_rng(seed, model, pol=pol) + seed, rng = self._seed_rng( + seed, model, pol=pol, model_key=model_key + ) + args["rng"] = rng polarized_gains = model(**args) for ant, gain in polarized_gains.items(): gains[(ant, pol)] = gain @@ -1074,6 +1090,19 @@ def _iteratively_apply( if model.return_type == "full_array": args = self._update_args(base_args, model) args.update(kwargs) + if seed: + if seed == "redundant": + warnings.warn( + "You are trying to set the random state once per " + "redundant group while simulating an effect that " + "computes the entire visibility matrix in one go. " + "Any randomness in the simulation component may not " + "come out as expected--please check your settings." + f"This warning was raised for model: {model_key}", + stacklevel=2, + ) + seed, rng = self._seed_rng(model, model_key=model_key) + args["rng"] = rng data_copy += model(**args) else: # Iterate over the array and simulate the effect as-needed. @@ -1091,11 +1120,13 @@ def _iteratively_apply( # Seed the random number generator. key = (ant2, ant1, pol) if conj_in_cache else (ant1, ant2, pol) - seed = self._seed_rng(seed, model, *key) + seed, rng = self._seed_rng(seed, model, *key, model_key=model_key) # Prepare the actual arguments to be used. use_args = self._update_args(base_args, model, ant1, ant2, pol) use_args.update(kwargs) + if model.is_randomized: + use_args["rng"] = rng if use_cached_filters: filter_kwargs = self._get_filters( ant1, @@ -1167,7 +1198,7 @@ def _read_datafile(datafile: Union[str, Path], **kwargs) -> UVData: uvd.read(datafile, read_data=True, **kwargs) return uvd - def _seed_rng(self, seed, model, ant1=None, ant2=None, pol=None): + def _seed_rng(self, seed, model, ant1=None, ant2=None, pol=None, model_key=None): """ Set the random state according to the provided parameters. @@ -1211,6 +1242,11 @@ def _seed_rng(self, seed, model, ant1=None, ant2=None, pol=None): Second antenna in the baseline (for baseline-dependent effects). pol Polarization string. + model_key + Identifier for retrieving the model parameters from the + ``self._components`` attribute. This is only needed for ensuring + that random effects using the "initial" seed can be recovered + with the ``self.get`` method. Returns ------- @@ -1218,6 +1254,8 @@ def _seed_rng(self, seed, model, ant1=None, ant2=None, pol=None): Either the input seed or ``None``, depending on the provided seed. This is just used to ensure that the logic for setting the random state in the :meth:`_iteratively_apply` routine works out. + rng + The random number generator to be used for producing the random effect. Raises ------ @@ -1228,11 +1266,12 @@ def _seed_rng(self, seed, model, ant1=None, ant2=None, pol=None): and a baseline isn't provided; two, the seed is a string, but is not one of the supported seeding modes. """ + model_key = model_key or self._get_model_name(model) if seed is None: - return + rng = self._components[model_key].get("rng", np.random.default_rng()) + return (None, rng) if isinstance(seed, int): - np.random.seed(seed) - return seed + return (seed, np.random.default_rng(seed)) if not isinstance(seed, str): raise TypeError( "The seeding mode must be specified as a string or integer. " @@ -1250,8 +1289,8 @@ def _seed_rng(self, seed, model, ant1=None, ant2=None, pol=None): if pol: key += (pol,) # seed the RNG accordingly - np.random.seed(self._get_seed(model, key)) - return "redundant" + seed = self._get_seed(model_key, key) + return ("redundant", np.random.default_rng(seed)) elif seed == "once": # this option seeds the RNG once per iteration of # _iteratively_apply, using the same seed every time @@ -1260,15 +1299,16 @@ def _seed_rng(self, seed, model, ant1=None, ant2=None, pol=None): # something like PointSourceForeground, where objects on # the sky are being placed randomly key = (pol,) if pol else 0 - np.random.seed(self._get_seed(model, key)) - return "once" + seed = self._get_seed(model_key, key) + return ("once", np.random.default_rng(seed)) elif seed == "initial": # this seeds the RNG once at the very beginning of # _iteratively_apply. this would be useful for something # like ThermalNoise key = (pol,) if pol else -1 - np.random.seed(self._get_seed(model, key)) - return None + rng = np.random.default_rng(self._get_seed(model_key, key)) + self._components[model_key]["rng"] = rng + return (None, rng) else: raise ValueError("Seeding mode not supported.") @@ -1413,20 +1453,44 @@ def _get_component( ) def _generate_seed(self, model, key): - """Generate a random seed based on the current time. + """Generate a random seed and cache it in the ``self._seeds`` attribute. - Populate the ``_seeds`` dictionary appropriately with the result. + Parameters + ---------- + model + The name of the model to retrieve the random seed for, as it would + appear in the ``self._components`` attribute. (This should always + correspond to the ``model_key`` determined in the ``self.add`` method.) + key + The key to use for tracking the random seed. This is only really + used for keeping track of random seeds that are set per polarization + or per redundant group. """ - model = self._get_model_name(model) - # for the sake of randomness - np.random.seed(int(time.time() * 1e6) % 2**32) + # Just to make it extra random. + rng = np.random.default_rng() if model not in self._seeds: self._seeds[model] = {} - self._seeds[model][key] = np.random.randint(2**32) + self._seeds[model][key] = rng.integers(2**32) def _get_seed(self, model, key): - """Retrieve or generate a random seed given a model and key.""" - model = self._get_model_name(model) + """Retrieve or generate a random seed given a model and key. + + Parameters + ---------- + model + The name of the model to retrieve the random seed for, as it would + appear in the ``self._components`` attribute. (This should always + correspond to the ``model_key`` determined in the ``self.add`` method.) + key + The key to use for tracking the random seed. This is only really + used for keeping track of random seeds that are set per polarization + or per redundant group. + + Returns + ------- + seed + The random seed to use for setting the random state. + """ if model not in self._seeds: self._generate_seed(model, key) if key not in self._seeds[model]: diff --git a/hera_sim/tests/test_adjustment.py b/hera_sim/tests/test_adjustment.py index 3a219d73..0d664dea 100644 --- a/hera_sim/tests/test_adjustment.py +++ b/hera_sim/tests/test_adjustment.py @@ -10,7 +10,6 @@ from astropy import units from pyuvdata import UVData from pyuvdata.utils import antnums_to_baseline -from scipy import stats from hera_sim import Simulator, adjustment, antpos, interpolators @@ -132,7 +131,8 @@ def test_match_subarray(): def test_match_translated_array(): # A simple translation should just be undone - translation = np.random.uniform(-1, 1, 3) + rng = np.random.default_rng(0) + translation = rng.uniform(-1, 1, 3) array_1 = {0: [0, 0, 0], 1: [1, 0, 0], 2: [1, 1, 0]} array_2 = {ant: np.array(pos) - translation for ant, pos in array_1.items()} # Won't be an exact match to machine precision, so need some small tolerance. @@ -264,8 +264,9 @@ def test_match_antennas_using_reference_positions_and_labels(base_sim, base_conf jitter_radius = 0.1 # 10 cm new_array = {0: [0, 0, 0], 1: [1, 0, 0], 2: [1, 1, 0], 3: [2, 1, 0], 4: [2, 2, 0]} base_config["array_layout"] = scale_array(new_array) + rng = np.random.default_rng(0) new_array = { - ant: pos + stats.uniform.rvs(-jitter_radius, jitter_radius, 3) + ant: pos + rng.uniform(-jitter_radius, jitter_radius, 3) for ant, pos in base_config["array_layout"].items() } # Mess up the antenna positions by up to 10 cm in each direction, randomly. base_config["array_layout"] = new_array @@ -297,8 +298,9 @@ def test_match_antennas_use_reference_positions_only(base_sim, base_config): jitter_radius = 0.1 # 10 cm new_array = {0: [0, 0, 0], 1: [1, 0, 0], 2: [1, 1, 0], 3: [2, 1, 0], 4: [2, 2, 0]} base_config["array_layout"] = scale_array(new_array) + rng = np.random.default_rng(0) new_array = { - ant: pos + stats.uniform.rvs(-jitter_radius, jitter_radius, 3) + ant: pos + rng.uniform(-jitter_radius, jitter_radius, 3) for ant, pos in base_config["array_layout"].items() } # Mess up the antenna positions by up to 10 cm in each direction, randomly. base_config["array_layout"] = new_array diff --git a/hera_sim/tests/test_compare_pyuvsim.py b/hera_sim/tests/test_compare_pyuvsim.py index 88c243ea..8d1eecb5 100644 --- a/hera_sim/tests/test_compare_pyuvsim.py +++ b/hera_sim/tests/test_compare_pyuvsim.py @@ -28,12 +28,12 @@ def get_uvdata(pol_array=None): obstime = Time("2018-08-31T04:02:30.11", format="isot", scale="utc") if pol_array is None: pol_array = np.array(["XX", "YY", "XY", "YX"]) - np.random.seed(10) + rng = np.random.default_rng(10) # Random antenna locations - x = np.random.random(nants) * 400.0 # Up to 400 metres - y = np.random.random(nants) * 400.0 - z = np.random.random(nants) * 0.0 + x = rng.random(nants) * 400.0 # Up to 400 metres + y = rng.random(nants) * 400.0 + z = rng.random(nants) * 0.0 ants = {i: (x[i], y[i], z[i]) for i in range(nants)} @@ -66,10 +66,11 @@ def get_sky_model(uvdata, nsource): sources = [ [125.7, -30.72, 2, 0], # Fix a single source near zenith ] + rng = np.random.default_rng(0) if nsource > 1: # Add random other sources - ra = np.random.uniform(low=0.0, high=360.0, size=nsource - 1) - dec = -30.72 + np.random.random(nsource - 1) * 10.0 - flux = np.random.random(nsource - 1) * 4 + ra = rng.uniform(low=0.0, high=360.0, size=nsource - 1) + dec = -30.72 + rng.random(nsource - 1) * 10.0 + flux = rng.random(nsource - 1) * 4 for i in range(nsource - 1): sources.append([ra[i], dec[i], flux[i], 0]) sources = np.array(sources) diff --git a/hera_sim/tests/test_defaults.py b/hera_sim/tests/test_defaults.py index 688d6702..759b169f 100644 --- a/hera_sim/tests/test_defaults.py +++ b/hera_sim/tests/test_defaults.py @@ -44,11 +44,9 @@ def test_multiple_param_specification(): def test_bandpass_changes(): defaults.set("h1c", refresh=True) fqs = np.linspace(0.1, 0.2, 100) - np.random.seed(0) - bp = gen_bandpass(fqs, [0])[0] + bp = gen_bandpass(fqs, [0], rng=np.random.default_rng(0))[0] defaults.set("h2c", refresh=True) - np.random.seed(0) - assert not np.all(bp == gen_bandpass(fqs, [0])[0]) + assert not np.all(bp == gen_bandpass(fqs, [0], rng=np.random.default_rng(0))[0]) defaults.deactivate() diff --git a/hera_sim/tests/test_eor.py b/hera_sim/tests/test_eor.py index d7173666..926263d6 100644 --- a/hera_sim/tests/test_eor.py +++ b/hera_sim/tests/test_eor.py @@ -22,21 +22,29 @@ def bl_vec(): @pytest.fixture(scope="function") def base_eor(freqs, lsts, bl_vec, request): - np.random.seed(0) if request.param == "auto": bl_vec = np.array([0, 0, 0]) return eor.noiselike_eor( - lsts=lsts, freqs=freqs, bl_vec=bl_vec, eor_amp=1e-5, fringe_filter_type="tophat" + lsts=lsts, + freqs=freqs, + bl_vec=bl_vec, + eor_amp=1e-5, + fringe_filter_type="tophat", + rng=np.random.default_rng(0), ) @pytest.fixture(scope="function") def scaled_eor(freqs, lsts, bl_vec, request): - np.random.seed(0) if request.param == "auto": bl_vec = np.array([0, 0, 0]) return eor.noiselike_eor( - lsts=lsts, freqs=freqs, bl_vec=bl_vec, eor_amp=1e-3, fringe_filter_type="tophat" + lsts=lsts, + freqs=freqs, + bl_vec=bl_vec, + eor_amp=1e-3, + fringe_filter_type="tophat", + rng=np.random.default_rng(0), ) diff --git a/hera_sim/tests/test_foregrounds.py b/hera_sim/tests/test_foregrounds.py index d8c50586..61dc42a3 100644 --- a/hera_sim/tests/test_foregrounds.py +++ b/hera_sim/tests/test_foregrounds.py @@ -123,8 +123,8 @@ def test_foreground_conjugation(freqs, lsts, Tsky_mdl, omega_p, model): conj_kwargs = kwargs.copy() conj_kwargs["bl_vec"] = -bl_vec - vis = model(**kwargs) - conj_vis = model(**conj_kwargs) + vis = model(**kwargs, rng=np.random.default_rng(0)) + conj_vis = model(**conj_kwargs, rng=np.random.default_rng(0)) assert np.allclose(vis, conj_vis.conj()) # Assert V_ij = V*_ji diff --git a/hera_sim/tests/test_noise.py b/hera_sim/tests/test_noise.py index 9e57a317..7d61aea9 100644 --- a/hera_sim/tests/test_noise.py +++ b/hera_sim/tests/test_noise.py @@ -81,7 +81,6 @@ def test_resample_Tsky_freq_behavior(tsky_powerlaw, tsky_from_model, model): def test_sky_noise_jy( freqs, lsts, tsky_powerlaw, Jy2T, omega_p, channel_width, integration_time, aspect ): - np.random.seed(0) noise_Jy = noise.sky_noise_jy( lsts=lsts, freqs=freqs, @@ -89,6 +88,7 @@ def test_sky_noise_jy( omega_p=omega_p, channel_width=channel_width, integration_time=integration_time, + rng=np.random.default_rng(0), ) # Calculate expected noise level based on radiometer equation. @@ -122,6 +122,7 @@ def test_thermal_noise_with_phase_wrap(freqs, omega_p, autovis, expectation): channel_width = np.mean(np.diff(freqs)) * units.GHz.to("Hz") expected_SNR = np.sqrt(integration_time * channel_width) Trx = 0 + rng = np.random.default_rng(0) if autovis is not None: autovis = np.ones((wrapped_lsts.size, freqs.size), dtype=complex) noise_sim = noise.ThermalNoise( @@ -129,6 +130,7 @@ def test_thermal_noise_with_phase_wrap(freqs, omega_p, autovis, expectation): omega_p=omega_p, Trx=Trx, autovis=autovis, + rng=rng, ) with expectation: vis = noise_sim(lsts=wrapped_lsts, freqs=freqs) diff --git a/hera_sim/tests/test_rfi.py b/hera_sim/tests/test_rfi.py index 3276537d..8b121b6c 100644 --- a/hera_sim/tests/test_rfi.py +++ b/hera_sim/tests/test_rfi.py @@ -18,7 +18,7 @@ def lsts(): @pytest.mark.parametrize("station_freq", [0.150, 0.1505]) def test_rfi_station_strength(freqs, lsts, station_freq): # Generate RFI for a single station. - station = rfi.RfiStation(station_freq, std=0.0) + station = rfi.RfiStation(station_freq, std=0.0, rng=np.random.default_rng(0)) rfi_vis = station(lsts, freqs) # Check that the RFI is inserted where it should be at the correct level. @@ -48,15 +48,16 @@ def test_rfi_station_from_file(freqs, lsts): filename = DATA_PATH / "HERA_H1C_RFI_STATIONS.npy" station_params = np.load(filename) Nstations = station_params.shape[0] - rfi_vis = rfi.rfi_stations(lsts, freqs, stations=filename) + rfi_vis = rfi.rfi_stations( + lsts, freqs, stations=filename, rng=np.random.default_rng(0) + ) assert np.sum(np.sum(np.abs(rfi_vis), axis=0).astype(bool)) >= Nstations @pytest.mark.parametrize("rfi_type", ["scatter", "impulse", "dtv"]) @pytest.mark.parametrize("chance", [0.2, 0.5, 0.8]) def test_rfi_occupancy(freqs, lsts, rfi_type, chance): - np.random.seed(0) - kwargs = {f"{rfi_type}_chance": chance} + kwargs = {f"{rfi_type}_chance": chance, "rng": np.random.default_rng(0)} # Modify expected chance appropriately if simulating DTV. if rfi_type == "dtv": fmin, fmax = (0.15, 0.20) @@ -64,12 +65,11 @@ def test_rfi_occupancy(freqs, lsts, rfi_type, chance): freq_cut = (freqs >= fmin) & (freqs < fmax) chance *= freqs[freq_cut].size / freqs.size rfi_vis = getattr(rfi, f"rfi_{rfi_type}")(lsts, freqs, **kwargs) - assert np.isclose(np.mean(np.abs(rfi_vis).astype(bool)), chance, atol=0.05, rtol=0) + assert np.isclose(np.mean(np.abs(rfi_vis).astype(bool)), chance, atol=0.1, rtol=0) @pytest.mark.parametrize("alignment", ["aligned", "misaligned"]) def test_rfi_dtv_constant_across_subband(freqs, lsts, alignment): - np.random.seed(0) dtv_band = (0.15, 0.20) if alignment == "aligned" else (0.1502, 0.2002) channel_width = 0.01 Nbands = 5 # Just hardcode this to avoid stupid numerical problems. @@ -82,6 +82,7 @@ def test_rfi_dtv_constant_across_subband(freqs, lsts, alignment): dtv_chance=0.5, dtv_strength=10, dtv_std=1, + rng=np.random.default_rng(0), ) for band_edges in zip(subband_edges[:-1], subband_edges[1:]): channels = np.argwhere( @@ -97,7 +98,6 @@ def test_rfi_dtv_constant_across_subband(freqs, lsts, alignment): def test_rfi_dtv_occupancy_variable_chance(freqs, lsts): - np.random.seed(0) dtv_band = (0.15, 0.20) channel_width = 0.01 Nbands = 5 # See above note about hardcoding. @@ -111,6 +111,7 @@ def test_rfi_dtv_occupancy_variable_chance(freqs, lsts): dtv_chance=chances, dtv_strength=10, dtv_std=1, + rng=np.random.default_rng(0), ) expected_occupancy = sum( chance diff --git a/hera_sim/tests/test_sigchain.py b/hera_sim/tests/test_sigchain.py index eb581ad3..af1dd9c9 100644 --- a/hera_sim/tests/test_sigchain.py +++ b/hera_sim/tests/test_sigchain.py @@ -11,8 +11,6 @@ from hera_sim.interpolators import Bandpass, Beam from hera_sim.io import empty_uvdata -np.random.seed(0) - @pytest.fixture(scope="function") def fqs(): @@ -73,7 +71,7 @@ def test_gen_bandpass(): assert 2 in g assert g[1].size == fqs.size assert np.all(g[1] == g[2]) - g = sigchain.gen_bandpass(fqs, list(range(10)), 0.2) + g = sigchain.gen_bandpass(fqs, list(range(10)), 0.2, rng=np.random.default_rng(0)) assert not np.all(g[1] == g[2]) @@ -145,7 +143,10 @@ def test_reflection_gains_correct_delays( dlys, ): # introduce a cable reflection into the autocorrelation - gains = sigchain.gen_reflection_gains(fqs, [0], amp=[1e-1], dly=[300], phs=[1]) + rng = np.random.default_rng(0) + gains = sigchain.gen_reflection_gains( + fqs, [0], amp=[1e-1], dly=[300], phs=[1], rng=rng + ) outvis = sigchain.apply_gains(vis, gains, [0, 0]) ovfft = uvtools.utils.FFT(outvis, axis=1, taper="blackman-harris") @@ -220,7 +221,7 @@ def test_reflection_spectrum(): dlys = np.arange(-1000, 1001, 5) fqs = uvtools.utils.fourier_freqs(dlys) fqs += 0.1 - fqs.min() # Range from 100 MHz to whatever the upper bound is - reflections = reflections(fqs, range(100)) + reflections = reflections(fqs, range(100), rng=np.random.default_rng(0)) reflections = np.vstack(list(reflections.values())) spectra = np.abs(uvtools.utils.FFT(reflections, axis=1)) spectra = spectra / spectra.max(axis=1).reshape(-1, 1) @@ -269,8 +270,9 @@ def test_amp_jitter(): ants = range(10000) amp = 5 amp_jitter = 0.1 + rng = np.random.default_rng(0) jittered_amps = sigchain.Reflections._complete_params( - ants, amp=amp, amp_jitter=amp_jitter + ants, amp=amp, amp_jitter=amp_jitter, rng=rng )[0] assert np.isclose(jittered_amps.mean(), amp, rtol=0.05) assert np.isclose(jittered_amps.std(), amp * amp_jitter, rtol=0.05) @@ -280,8 +282,9 @@ def test_dly_jitter(): ants = range(10000) dly = 500 dly_jitter = 20 + rng = np.random.default_rng(0) jittered_dlys = sigchain.Reflections._complete_params( - ants, dly=dly, dly_jitter=dly_jitter + ants, dly=dly, dly_jitter=dly_jitter, rng=rng )[1] assert np.isclose(jittered_dlys.mean(), dly, rtol=0.05) assert np.isclose(jittered_dlys.std(), dly_jitter, rtol=0.05) @@ -296,6 +299,7 @@ def test_cross_coupling_spectrum(fqs, dlys, Tsky): amp_range=amp_range, dly_range=dly_range, symmetrize=True, + rng=np.random.default_rng(0), ) amplitudes = np.logspace(*amp_range, n_copies) delays = np.linspace(*dly_range, n_copies) @@ -351,6 +355,7 @@ def test_over_air_cross_coupling(Tsky_mdl, lsts): cable_delays=cable_delays, max_delay=max_delay, amp_decay_fac=amp_decay_fac, + rng=np.random.default_rng(0), ) xtalk = gen_xtalk(fqs, (0, 1), antpos, Tsky, Tsky) xt_fft = uvtools.utils.FFT(xtalk, axis=1, taper="bh7") @@ -570,8 +575,8 @@ def uvbeam(tmp_path): def test_mutual_coupling_with_uvbeam(uvbeam, sample_uvdata, sample_coupling): - np.random.seed(0) - data = np.random.normal(size=sample_uvdata.data_array.shape) + 0j + rng = np.random.default_rng(0) + data = rng.normal(size=sample_uvdata.data_array.shape) + 0j sample_uvdata.data_array[...] = data[...] sample_uvdata.data_array += sample_coupling( freqs=sample_uvdata.freq_array.squeeze() / 1e9, @@ -590,9 +595,10 @@ def test_mutual_coupling_input_types( def func(freqs): return np.ones_like(freqs) + rng = np.random.default_rng(0) omega_p = func if omega_p == "callable" else None reflection = func if reflection == "callable" else None - data = np.random.normal(size=sample_uvdata.data_array.shape) + 0j + data = rng.normal(size=sample_uvdata.data_array.shape) + 0j sample_uvdata.data_array += data sample_uvdata.data_array += sample_coupling( freqs=sample_uvdata.freq_array.squeeze() / 1e9, @@ -738,12 +744,14 @@ def test_vary_gain_amp_sinusoidal(gains, times, fringe_rates, fringe_keys): def test_vary_gain_amp_noiselike(gains, times): vary_amp = 0.1 + rng = np.random.default_rng(0) varied_gain = sigchain.vary_gains_in_time( gains=gains, times=times, parameter="amp", variation_mode="noiselike", variation_amp=vary_amp, + rng=rng, )[0] # Check that the mean gain amplitude is the original gain amplitude. @@ -815,6 +823,7 @@ def test_vary_gain_phase_noiselike(gains, times, delay_phases, phase_offsets): parameter="phs", variation_mode="noiselike", variation_amp=vary_amp, + rng=np.random.default_rng(0), )[0] varied_phases = np.angle(varied_gain) @@ -892,6 +901,7 @@ def test_vary_gain_delay_noiselike(gains, times, freqs, delays): parameter="dly", variation_amp=vary_amp, variation_mode="noiselike", + rng=np.random.default_rng(0), )[0] # Determine the bandpass delay at each time. diff --git a/hera_sim/tests/test_simulator.py b/hera_sim/tests/test_simulator.py index 8905164e..b1efef03 100644 --- a/hera_sim/tests/test_simulator.py +++ b/hera_sim/tests/test_simulator.py @@ -112,17 +112,27 @@ def test_nondefault_blt_order_lsts(): def test_add_with_str(base_sim): - base_sim.add("noiselike_eor") + base_sim.add("noiselike_eor", rng=np.random.default_rng(0)) assert not np.all(base_sim.data.data_array == 0) def test_add_with_builtin_class(base_sim): - base_sim.add(DiffuseForeground, Tsky_mdl=Tsky_mdl, omega_p=omega_p) + base_sim.add( + DiffuseForeground, + Tsky_mdl=Tsky_mdl, + omega_p=omega_p, + rng=np.random.default_rng(0), + ) assert not np.all(np.isclose(base_sim.data.data_array, 0)) def test_add_with_class_instance(base_sim): - base_sim.add(diffuse_foreground, Tsky_mdl=Tsky_mdl, omega_p=omega_p) + base_sim.add( + diffuse_foreground, + Tsky_mdl=Tsky_mdl, + omega_p=omega_p, + rng=np.random.default_rng(0), + ) assert not np.all(np.isclose(base_sim.data.data_array, 0)) @@ -301,7 +311,9 @@ def test_get_multiplicative_effect(base_sim, pol, ant1): def test_not_add_vis(base_sim): - vis = base_sim.add("noiselike_eor", add_vis=False, ret_vis=True) + vis = base_sim.add( + "noiselike_eor", add_vis=False, ret_vis=True, rng=np.random.default_rng(0) + ) assert np.all(base_sim.data.data_array == 0) @@ -315,14 +327,16 @@ def test_not_add_vis(base_sim): def test_adding_vis_but_also_returning(base_sim): - vis = base_sim.add("noiselike_eor", ret_vis=True) + vis = base_sim.add("noiselike_eor", ret_vis=True, rng=np.random.default_rng(0)) assert not np.all(vis == 0) assert np.all(np.isclose(vis, base_sim.data.data_array)) # use season defaults for simplicity defaults.set("h1c") - vis += base_sim.add("diffuse_foreground", ret_vis=True) + vis += base_sim.add( + "diffuse_foreground", ret_vis=True, rng=np.random.default_rng(90) + ) # deactivate defaults for good measure defaults.deactivate() assert np.all(np.isclose(vis, base_sim.data.data_array)) @@ -334,7 +348,7 @@ def test_filter(): # only add visibilities for the (0,1) baseline vis_filter = (0, 1, "xx") - sim.add("noiselike_eor", vis_filter=vis_filter) + sim.add("noiselike_eor", vis_filter=vis_filter, rng=np.random.default_rng(10)) assert np.all(sim.data.get_data(0, 0) == 0) assert np.all(sim.data.get_data(1, 1) == 0) assert np.all(sim.data.get_data(0, 1) != 0) @@ -599,13 +613,13 @@ def test_legacy_funcs(component): def test_vis_filter_single_pol(): sim = create_sim(polarization_array=["xx", "yy"]) - sim.add("noiselike_eor", vis_filter=["xx"]) + sim.add("noiselike_eor", vis_filter=["xx"], rng=np.random.default_rng(99)) assert np.all(sim.get_data("xx")) and not np.any(sim.get_data("yy")) def test_vis_filter_two_pol(): sim = create_sim(polarization_array=["xx", "xy", "yx", "yy"]) - sim.add("noiselike_eor", vis_filter=["xx", "yy"]) + sim.add("noiselike_eor", vis_filter=["xx", "yy"], rng=np.random.default_rng(5)) assert all( [ np.all(sim.get_data("xx")), @@ -621,7 +635,7 @@ def test_vis_filter_arbitrary_key(): array_layout=hex_array(2, split_core=False, outriggers=0), polarization_array=["xx", "yy"], ) - sim.add("noiselike_eor", vis_filter=[1, 3, 5, "xx"]) + sim.add("noiselike_eor", vis_filter=[1, 3, 5, "xx"], rng=np.random.default_rng(7)) bls = sim.data.get_antpairs() assert not np.any(sim.get_data("yy")) assert all( @@ -666,7 +680,7 @@ def test_bad_seeds(base_sim, seed): "unsupported": "Seeding mode not supported.", }[seed] with pytest.raises(err, match=match): - base_sim._seed_rng(seed, None) + base_sim._seed_rng(seed, None, model_key="test") def test_get_component_with_function(): diff --git a/hera_sim/tests/test_utils.py b/hera_sim/tests/test_utils.py index c00ad5f7..b23c1454 100644 --- a/hera_sim/tests/test_utils.py +++ b/hera_sim/tests/test_utils.py @@ -124,8 +124,9 @@ def test_rough_filter_noisy_data(freqs, lsts, filter_type): "fringe_filter_type": "gauss", "fr_width": 1e-4, } + rng = np.random.default_rng(0) for i in range(Nrealizations): - data = utils.gen_white_noise((lsts.size, freqs.size)) + data = utils.gen_white_noise((lsts.size, freqs.size), rng=rng) filtered_data = filt(data, *args, **kwargs) filtered_data_mean = np.mean(filtered_data) mean_values[i] = filtered_data_mean.real, filtered_data_mean.imag @@ -153,12 +154,12 @@ def test_rough_delay_filter_missing_param(freqs, lsts, missing_param): def test_delay_filter_norm(freqs): tsky = np.ones(freqs.size) - np.random.seed(1234) # set the seed for reproducibility. + rng = np.random.default_rng(0) # set the seed for reproducibility. out = 0 nreal = 5000 for _ in range(nreal): - _noise = tsky * utils.gen_white_noise(freqs.size) + _noise = tsky * utils.gen_white_noise(freqs.size, rng=rng) outnoise = utils.rough_delay_filter(_noise, freqs, 30, normalize=1) out += np.sum(np.abs(outnoise) ** 2) @@ -276,7 +277,7 @@ def test_fringe_filter_custom(freqs, lsts, fringe_rates): @pytest.mark.parametrize("bl_len_ns", [50, 150]) @pytest.mark.parametrize("fr_width", [1e-4, 3e-4]) def test_rough_fringe_filter_noisy_data(freqs, lsts, fringe_rates, bl_len_ns, fr_width): - data = utils.gen_white_noise((lsts.size, freqs.size)) + data = utils.gen_white_noise((lsts.size, freqs.size), rng=np.random.default_rng(0)) max_fringe_rates = utils.calc_max_fringe_rate(freqs, bl_len_ns) filt_data = utils.rough_fringe_filter( data, lsts, freqs, bl_len_ns, fringe_filter_type="gauss", fr_width=fr_width @@ -333,7 +334,7 @@ def test_gen_white_noise_shape(shape): @pytest.mark.parametrize("shape", [100, (100, 200)]) def test_gen_white_noise_mean(shape): - noise = utils.gen_white_noise(shape) + noise = utils.gen_white_noise(shape, rng=np.random.default_rng(0)) assert np.allclose( [noise.mean().real, noise.mean().imag], 0, rtol=0, atol=5 / np.sqrt(noise.size) ) @@ -341,7 +342,7 @@ def test_gen_white_noise_mean(shape): @pytest.mark.parametrize("shape", [100, (100, 200)]) def test_gen_white_noise_variance(shape): - noise = utils.gen_white_noise(shape) + noise = utils.gen_white_noise(shape, rng=np.random.default_rng(0)) assert np.isclose(np.std(noise), 1, rtol=0, atol=0.1) diff --git a/hera_sim/tests/test_vis.py b/hera_sim/tests/test_vis.py index a094c129..7a3e7f94 100644 --- a/hera_sim/tests/test_vis.py +++ b/hera_sim/tests/test_vis.py @@ -37,7 +37,6 @@ def __init__(self, *args, **kwargs): SIMULATORS = SIMULATORS + (VisGPU,) -np.random.seed(0) NTIMES = 10 NPIX = 12 * 16**2 NFREQ = 5 diff --git a/hera_sim/utils.py b/hera_sim/utils.py index 612f45a7..211a9eb7 100644 --- a/hera_sim/utils.py +++ b/hera_sim/utils.py @@ -420,7 +420,9 @@ def wrap2pipi(a): return res -def gen_white_noise(size: int | tuple[int] = 1) -> np.ndarray: +def gen_white_noise( + size: int | tuple[int] = 1, rng: np.random.Generator | None = None +) -> np.ndarray: """Produce complex Gaussian noise with unity variance. Parameters @@ -428,16 +430,21 @@ def gen_white_noise(size: int | tuple[int] = 1) -> np.ndarray: size Shape of output array. Can be an integer if a single dimension is required, otherwise a tuple of ints. + rng + Random number generator. Returns ------- noise White noise realization with specified shape. """ + # Split power evenly between real and imaginary components. std = 1 / np.sqrt(2) - return np.random.normal(scale=std, size=size) + 1j * np.random.normal( - scale=std, size=size - ) + args = dict(scale=std, size=size) + + # Create a random number generator if needed, then generate noise. + rng = rng or np.random.default_rng() + return rng.normal(**args) + 1j * rng.normal(**args) def jansky_to_kelvin(freqs: np.ndarray, omega_p: Beam | np.ndarray) -> np.ndarray: