feat: Issue #48 feature 1

docs/tutorials/lc_noise_compare_tools21cm.ipynb

@@ -648,20 +648,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from tuesday.core import sample_lc_noise"
+    "from tuesday.core import sample_from_rms_noise"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "noise = sample_lc_noise(sigmaN, \n",
+    "noise = sample_from_rms_noise(sigmaN, \n",
     "                        seed=4, \n",
     "                        nsamples=10)"
    ]
@@ -743,9 +743,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "pytorch_env",
+   "display_name": "21cmFAST",
    "language": "python",
-   "name": "pytorch_env"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -757,7 +757,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.11.12"
   }
  },
  "nbformat": 4,

src/tuesday/core/__init__.py

@@ -26,6 +26,7 @@
     "sample_lc_noise",
     "taper2d",
     "thermal_noise_per_voxel",
+    "thermal_noise_uv",
     "validate",
 ]
 from .instrument_models.noise import (
@@ -34,6 +35,7 @@
     sample_lc_noise,
     taper2d,
     thermal_noise_per_voxel,
+    thermal_noise_uv,
 )
 from .plotting.powerspectra import (
     plot_1d_power_spectrum_k,

src/tuesday/core/instrument_models/noise.py

@@ -17,8 +17,8 @@
 def grid_baselines_uv(
     uvws: np.ndarray,
     freq: un.Quantity,
-    boxlength: un.Quantity,
-    lc_shape: tuple[int, int, int],
+    box_len: un.Quantity,
+    box_ncells: int,
     weights: np.ndarray,
     include_mirrored_bls: bool = True,
     avg_mirrored_bls: bool = True,
@@ -31,12 +31,10 @@ def grid_baselines_uv(
         Baselines in uv space with shape (N bls, N time offsets, 3).
     freq : un.Quantity
         Frequency at which the baselines are projected.
-    boxlength : un.Quantity
+    box_len : un.Quantity
         Transverse length of the simulation box.
-    lc_shape : tuple
-        Shape of the lightcone (Nx, Ny, Nz).
-        We assume that Nx = Ny to be sky-plane dimensions,
-        and Nz to be to line-of-sight (frequency) dimension.
+    box_ncells : int
+        Number of voxels Nx = Ny of a lightcone or coeval box.
     weights : np.ndarray
         Weights for each baseline group with shape (N bls).
     include_mirrored_bls : bool, optional
@@ -55,13 +53,13 @@ def grid_baselines_uv(
         of observation with shape (Nu=Nx, Nv=Nx).
 
     """
-    if "littleh" in boxlength.unit.to_string():
-        boxlength = boxlength.to(un.Mpc / littleh)
+    if "littleh" in box_len.unit.to_string():
+        box_len = box_len.to(un.Mpc / littleh)
     else:
-        boxlength = boxlength.to(un.Mpc) * Planck18.h / littleh
-    dx = float(boxlength.value) / float(lc_shape[0])
+        box_len = box_len.to(un.Mpc) * Planck18.h / littleh
+    dx = float(box_len.value) / float(box_ncells)
     ugrid_edges = (
-        np.fft.fftshift(np.fft.fftfreq(lc_shape[0], d=dx)) * 2 * np.pi * boxlength.unit
+        np.fft.fftshift(np.fft.fftfreq(box_ncells, d=dx)) * 2 * np.pi * box_len.unit
     )
 
     du = ugrid_edges[1] - ugrid_edges[0]
@@ -86,8 +84,8 @@ def grid_baselines_uv(
 def thermal_noise_per_voxel(
     observation: Observation,
     freqs: np.ndarray,
-    boxlen: float,
-    lc_shape: tuple[int, int, int],
+    box_len: float,
+    box_ncells: int,
     antenna_effective_area: un.Quantity | None = None,
     beam_area: un.Quantity | None = None,
 ):
@@ -105,12 +103,10 @@ def thermal_noise_per_voxel(
         Instance of `Observation`.
     freqs : astropy.units.Quantity
         Frequencies at which the noise is calculated.
-    boxlen : astropy.units.Quantity
+    box_len : astropy.units.Quantity
         Transverse length of the simulation box.
-    lc_shape : tuple
-        Shape of the lightcone (Nx, Ny, Nz).
-        We assume that Nx = Ny to be sky-plane dimensions,
-        and Nz to be to line-of-sight (frequency) dimension.
+    box_ncells : int
+        Number of voxels Nx = Ny of a lightcone or coeval box.
     antenna_effective_area : astropy.units.Quantity, optional
         Effective area of the antenna with shape (Nfreqs,).
     beam_area : astropy.units.Quantity, optional
@@ -167,8 +163,8 @@ def thermal_noise_per_voxel(
         tsys = obs.Tsys.to(un.mK)
 
         d = Planck18.comoving_distance(f2z(nu)).to(un.Mpc)  # Mpc
-        theta_box = (boxlen.to(un.Mpc) / d) * un.rad
-        omega_pix = theta_box**2 / np.prod(lc_shape[:2])
+        theta_box = (box_len.to(un.Mpc) / d) * un.rad
+        omega_pix = theta_box**2 / box_ncells**2
 
         sqrt = np.sqrt(2.0 * observation.bandwidth.to("Hz") * obs.integration_time).to(
             un.dimensionless_unscaled
@@ -211,23 +207,30 @@ def sample_from_rms_noise(
     seed: int | None = None,
     nsamples: int = 1,
     window_fnc: str = "blackmanharris",
+    return_in_uv: bool = False,
 ):
     """Sample noise for a lightcone slice given the corresponding rms noise in uv space.
 
     Parameters
     ----------
     rms_noise : astropy.units.Quantity
         RMS noise in uv space, shape (Nx, Ny, Nfreqs).
+    seed : int, optional
+        Random seed for reproducibility, by default None.
     nsamples : int, optional
         Number of noise realisations to sample, by default 1.
     window_fnc : str, optional
         Name of window function to be applied to the noise sampled in uv space,
         by default windows.blackmanharris.
+    return_in_uv : bool, optional
+        If True, return the noise sampled in uv space instead of real space,
+        by default False.
 
     Returns
     -------
-    lc_noise : un.Quantity
-        Noise sampled in real space, shape (nsamples, Nx, Ny, Nfreqs
+    noise : un.Quantity
+        Noise sampled in real or uv space, shape
+        (nsamples, Nx or Nu, Ny or Nv, Nfreqs)
 
     """
     if len(rms_noise.shape) == 2:
@@ -246,23 +249,53 @@ def sample_from_rms_noise(
 
     noise *= window_fnc[None, ..., None]
     noise = (noise + np.conj(noise)) / 2.0
-    noise = np.fft.ifft2(np.fft.ifftshift(noise, axes=(1, 2)), axes=(1, 2))
+    noise = np.fft.ifftshift(noise, axes=(1, 2))
+    if not return_in_uv:
+        noise = np.fft.ifft2(noise, axes=(1, 2)).real * rms_noise.unit
+    else:
+        noise = noise * rms_noise.unit
 
-    return noise.real * rms_noise.unit
+    return noise
 
 
-def sample_lc_noise(
+def thermal_noise_uv(
     observation: Observation,
     freqs: un.Quantity,
-    boxlength: un.Quantity,
-    lc_shape: tuple[int, int, int],
+    box_len: un.Quantity,
+    box_ncells: int,
     antenna_effective_area: un.Quantity | None = None,
     beam_area: un.Quantity | None = None,
-    seed: int | None = None,
-    nsamples: int = 1,
-    window_fnc: str = "blackmanharris",
+    min_nbls_per_uv_cell: int = 1,
 ):
-    """Test the grid_baselines function."""
+    """Thermal noise RMS per voxel in uv space.
+
+    Parameters
+    ----------
+    observation : py21cmsense.Observation
+        Instance of `Observation`.
+    freqs : astropy.units.Quantity
+        Frequencies at which the noise is calculated.
+    box_len : astropy.units.Quantity
+        Length of the box in which the noise is calculated.
+    box_ncells : int
+        Number of voxels Nx = Ny of a lightcone or coeval box.
+    antenna_effective_area : astropy.units.Quantity, optional
+        Effective area of the antenna with shape (Nfreqs,).
+    beam_area : astropy.units.Quantity, optional
+        Beam area of the antenna with shape (Nfreqs,).
+    min_nbls_per_uv_cell : int, optional
+        Minimum number of baselines per uv cell to consider
+        the cell to be measured, by default 1.
+        sigma is set to zero for uv cells with less than
+        this number of baselines.
+
+    Returns
+    -------
+    sigma : astropy.units.Quantity
+        Thermal noise RMS per voxel in uv space
+        with shape (Nx, Ny, Nfreqs).
+
+    """
     observatory = observation.observatory
     time_offsets = observatory.time_offsets_from_obs_int_time(
         observation.integration_time, observation.time_per_day
@@ -276,24 +309,130 @@ def sample_lc_noise(
         baselines=baselines, time_offset=time_offsets
     )
 
-    uv_coverage = np.zeros((lc_shape[0], lc_shape[0], len(freqs)))
+    uv_coverage = np.zeros((box_ncells, box_ncells, len(freqs)))
 
     for i, freq in enumerate(freqs):
         uv_coverage[..., i] += grid_baselines_uv(
-            proj_bls[::2] * freq / freqs[0], freq, boxlength, lc_shape, weights[::2]
+            proj_bls[::2] * freq / freqs[0], freq, box_len, box_ncells, weights[::2]
         )
 
     sigma_rms = thermal_noise_per_voxel(
         observation,
         freqs,
-        boxlength,
-        lc_shape,
+        box_len,
+        box_ncells,
         antenna_effective_area=antenna_effective_area,
         beam_area=beam_area,
     )
     sigma = sigma_rms / np.sqrt(uv_coverage * observation.n_days)
-    sigma[uv_coverage == 0.0] = 0.0
+    sigma[uv_coverage < min_nbls_per_uv_cell] = 0.0
+    return sigma
+
+
+def sample_lc_noise(
+    lightcone: un.Quantity,
+    observation: Observation,
+    box_len: un.Quantity,
+    *,
+    freqs: un.Quantity | None = None,
+    lightcone_redshifts: float | None = None,
+    antenna_effective_area: un.Quantity | None = None,
+    beam_area: un.Quantity | None = None,
+    seed: int | None = None,
+    nsamples: int = 1,
+    window_fnc: str = "blackmanharris",
+    min_nbls_per_uv_cell: int = 1,
+    wedge_mu_min: float = 0.0,
+):
+    """Sample thermal noise and add it to a lightcone in Fourier space.
 
-    return sample_from_rms_noise(
-        sigma, seed=seed, nsamples=nsamples, window_fnc=window_fnc
+    Parameters
+    ----------
+    lightcone : astropy.units.Quantity
+        Lightcone slice with shape (Nx, Ny, Nz).
+    observation : py21cmsense.Observation
+        Instance of `Observation`.
+    box_len : astropy.units.Quantity
+        Length of the lightcone box side.
+    freqs : astropy.units.Quantity, optional
+        Frequencies at which the thermal noise is calculated.
+        Must have the same length as the lightcone frequency axis.
+        If not provided, freqs are calculated from lightcone_redshifts.
+    antenna_effective_area : astropy.units.Quantity, optional
+        Effective area of the antenna with shape (Nfreqs,).
+    beam_area : astropy.units.Quantity, optional
+        Beam area of the antenna with shape (Nfreqs,).
+    nsamples : int, optional
+        Number of noise realisations to sample, by default 1.
+    seed : int, optional
+        Random seed for reproducibility, by default None.
+    window_fnc : str, optional
+        Name of window function to be applied to the noise sampled in uv space,
+        by default windows.blackmanharris.
+    min_nbls_per_uv_cell : int, optional
+        Minimum number of baselines per uv cell to consider
+        the cell to be measured, by default 1.
+        Thermal noise in uv space is set to zero for
+        uv cells with less than this number of baselines.
+    wedge_mu_min : float, optional
+        Minimum cosine of the angle between the line of sight
+        and the k vector to consider a mode to be uncontaminated
+        by foregrounds, by default 0.0.
+
+    Returns
+    -------
+    lightcone samples with noise
+    """
+    if freqs is None:
+        if lightcone_redshifts is None:
+            raise ValueError("You must provide either freqs or lightcone_redshifts.")
+        freqs = (1420.0 / (1 + lightcone_redshifts)) * un.MHz
+    if len(freqs) != lightcone.shape[2]:
+        raise ValueError(
+            "The length of freqs must be the same as the "
+            "length of the lightcone frequency axis."
+        )
+
+    wedge_mu_min = np.asarray(wedge_mu_min) + np.zeros(len(freqs))
+
+    if np.min(wedge_mu_min) < 0 or np.max(wedge_mu_min) > 1:
+        raise ValueError("wedge_mu_min must be between 0 and 1.")
+
+    sigma = thermal_noise_uv(
+        observation,
+        freqs,
+        box_len,
+        lightcone.shape[0],
+        antenna_effective_area=antenna_effective_area,
+        beam_area=beam_area,
+        min_nbls_per_uv_cell=min_nbls_per_uv_cell,
+    )
+
+    noise_realisation_uv = sample_from_rms_noise(
+        sigma,
+        seed=seed,
+        nsamples=nsamples,
+        window_fnc=window_fnc,
+        return_in_uv=True,
     )
+    lightcone -= lightcone.mean(axis=(0, 1), keepdims=True)
+    lc_ft = np.fft.fft2(lightcone.value) * lightcone.unit
+    lc_ft[sigma == 0] = 0.0
+    noisy_lc_ft = lc_ft + noise_realisation_uv
+    noisy_lc_ft[..., sigma == 0] = 0.0
+    k = (
+        np.fft.fftshift(
+            np.fft.fftfreq(
+                lightcone.shape[0], d=(box_len / lightcone.shape[0]).to(un.Mpc).value
+            )
+        )
+        * 2
+        * np.pi
+    )
+    kperp1, kperp2 = np.meshgrid(k, k)
+    theta = np.arctan(kperp1 / kperp2)
+    mu = np.cos(theta)
+    for i in range(len(freqs)):
+        noisy_lc_ft[:, np.abs(mu) < wedge_mu_min[i], i] = 0.0
+
+    return np.fft.ifft2(noisy_lc_ft, axes=(1, 2)).real.to(lightcone.unit)