Merge pull request #108 from lincc-frameworks/apply-noise

apply_noise(flux, flux_err)

docs/notebooks/test_snia.ipynb

@@ -7,8 +7,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2"
+    "# %load_ext autoreload\n",
+    "# %autoreload 2"
    ]
   },
   {
@@ -269,10 +269,9 @@
    "outputs": [],
    "source": [
     "for i in range(0, 3):\n",
-    "    try:\n",
-    "        plt.plot(res[i][\"wavelengths_rest\"] * (1 + z[i]), res[i][\"flux_flam\"][0], color=\"r\")\n",
-    "    except Exception:\n",
+    "    if res[i] is None:\n",
     "        continue\n",
+    "    plt.plot(res[i][\"wavelengths_rest\"] * (1 + z[i]), res[i][\"flux_flam\"][0], color=\"r\")\n",
     "    saltpars = {\"x0\": x0[i], \"x1\": x1[i], \"c\": c[i], \"z\": z[i], \"t0\": t0[i]}\n",
     "    model = sncosmo.Model(\"salt3\")\n",
     "    model.update(saltpars)\n",
@@ -299,12 +298,28 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "colors = \"green\", \"red\"\n",
+    "mag_zp = 31.4\n",
     "for i in range(0, 3):\n",
-    "    times = res[i][\"times\"]\n",
-    "    colors = [\"red\", \"brown\"]\n",
-    "    for f, color in zip(\"ri\", colors):\n",
+    "    _fig, (ax_f, ax_m) = plt.subplots(2, figsize=(5, 8))\n",
+    "    for f, color in zip(\"gr\", colors):\n",
     "        band_name = f\"LSST_{f}\"\n",
-    "        plt.plot(times, res[i][\"bandfluxes\"][band_name], \"-\", label=f, color=color, alpha=0.6, lw=2)\n",
+    "        band_idx = res[i][\"filters\"] == band_name\n",
+    "        t = res[i][\"times\"][band_idx]\n",
+    "        flux_perfect = res[i][\"bandfluxes_perfect\"][band_idx]\n",
+    "        flux = res[i][\"bandfluxes\"][band_idx]\n",
+    "        flux_err = res[i][\"bandfluxes_error\"][band_idx]\n",
+    "        mag = mag_zp - 2.5 * np.log10(flux)\n",
+    "        mag_plus = mag_zp - 2.5 * np.log10(flux - flux_err)\n",
+    "        mag_minus = mag_zp - 2.5 * np.log10(flux + flux_err)\n",
+    "\n",
+    "        ax_f.plot(t, flux_perfect, \"-\", label=f, color=color, alpha=0.6, lw=2)\n",
+    "        ax_f.fill_between(t, flux - flux_err, flux + flux_err, color=\"black\", alpha=0.3)\n",
+    "\n",
+    "        # ax_m.scatter(t, mag, label=f, marker='x', s=8, color=color, lw=2)\n",
+    "        # ax_m.errorbar(t, mag_minus, max_plus, ls='', color=color)\n",
+    "        ax_m.fill_between(t, mag_minus, mag_plus, color=color, alpha=0.6)\n",
+    "\n",
     "        saltpars = {\"x0\": x0[i], \"x1\": x1[i], \"c\": c[i], \"z\": z[i], \"t0\": t0[i]}\n",
     "        model = sncosmo.Model(\"salt3\")\n",
     "        model.update(saltpars)\n",
@@ -313,11 +328,14 @@
     "        sncosmo_band = sncosmo.Bandpass(\n",
     "            *passbands.passbands[band_name].processed_transmission_table.T, name=band_name\n",
     "        )\n",
-    "        flux = model.bandflux(sncosmo_band, times, zpsys=\"ab\", zp=8.9 + 2.5 * 9)  # -48.6)\n",
-    "        plt.plot(times, flux, \"--\", label=f, color=color)\n",
-    "        plt.xlabel(\"MJD\")\n",
-    "        plt.ylabel(\"Flux, nJy\")\n",
-    "        plt.legend()\n",
+    "        flux_sncosmo = model.bandflux(sncosmo_band, t, zpsys=\"ab\", zp=mag_zp)\n",
+    "        ax_f.plot(t, flux_sncosmo, \"--\", label=f, color=color)\n",
+    "    ax_f.set_xlabel(\"MJD\")\n",
+    "    ax_f.set_ylabel(\"Flux, nJy\")\n",
+    "    ax_m.set_xlabel(\"MJD\")\n",
+    "    ax_m.set_ylabel(\"mag\")\n",
+    "    ax_m.invert_yaxis()\n",
+    "    plt.legend()\n",
     "    plt.show()"
    ]
   },
@@ -365,7 +383,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.9"
+   "version": "3.12.6"
   }
  },
  "nbformat": 4,

src/tdastro/astro_utils/mag_flux.py

@@ -8,32 +8,32 @@
 
 
 def mag2flux(mag: npt.ArrayLike) -> npt.ArrayLike:
-    """Convert AB magnitude to flux in nJy
+    """Convert AB magnitude to bandflux in nJy
 
     Parameters
     ----------
     mag : ndarray of `float`
-        The magnitude to convert to flux.
+        The magnitude to convert to bandflux.
 
     Returns
     -------
-    flux : ndarray of `float`
-        The flux corresponding to the input magnitude.
+    bandflux : ndarray of `float`
+        The bandflux corresponding to the input magnitude.
     """
     return np.power(10.0, -0.4 * (mag - MAG_AB_ZP_NJY))
 
 
 def flux2mag(flux_njy: npt.ArrayLike) -> npt.ArrayLike:
-    """Convert flux in nJy to AB magnitude
+    """Convert bandflux in nJy to AB magnitude
 
     Parameters
     ----------
     flux_njy : ndarray of `float`
-        The flux to convert to magnitude.
+        The bandflux to convert to magnitude.
 
     Returns
     -------
     mag : ndarray of `float`
-        The magnitude corresponding to the input flux.
+        The magnitude corresponding to the input bandflux.
     """
     return MAG_AB_ZP_NJY - 2.5 * np.log10(flux_njy)

src/tdastro/astro_utils/noise_model.py

@@ -2,8 +2,8 @@
 import numpy.typing as npt
 
 
-def poisson_flux_std(
-    flux: npt.ArrayLike,
+def poisson_bandflux_std(
+    bandflux: npt.ArrayLike,
     *,
     pixel_scale: npt.ArrayLike,
     total_exposure_time: npt.ArrayLike,
@@ -14,12 +14,12 @@ def poisson_flux_std(
     readout_noise: npt.ArrayLike,
     dark_current: npt.ArrayLike,
 ) -> npt.ArrayLike:
-    """Simulate photon noise for flux measurements.
+    """Simulate photon noise for bandflux measurements.
 
     Parameters
     ----------
-    flux : array_like of float
-        Source flux in energy units, e.g. nJy.
+    bandflux : array_like of float
+        Source bandflux in energy units, e.g. nJy.
     pixel_scale : array_like of float
         Pixel scale of the detector,
         in angular units (e.g. arcsec) per pixel.
@@ -36,7 +36,7 @@ def poisson_flux_std(
         Point spread function effective area,
         in squared angular units, e.g. arcsec^2.
     zp : array_like of float
-        Zero point flux for the observation, e.g. flux
+        Zero point bandflux for the observation, i.e. bandflux
         giving a single electron during the total exposure time.
     readout_noise : array_like of float
         Standard deviation of the readout electrons per pixel per exposure.
@@ -46,7 +46,7 @@ def poisson_flux_std(
     Returns
     -------
     array_like
-        Simulated flux noise, in the same units as the input flux.
+        Simulated bandflux noise, in the same units as the input bandflux.
 
     Notes
     -----
@@ -68,11 +68,34 @@ def poisson_flux_std(
     area_px = footprint / pixel_scale**2
 
     # Get variances, in electrons^2
-    source_variance = flux / zp
+    source_variance = bandflux / zp
     sky_variance = sky * footprint / zp
     readout_variance = readout_noise**2 * area_px * exposure_count
     dark_variance = dark_current * total_exposure_time * area_px
 
     total_variance = source_variance + sky_variance + readout_variance + dark_variance
 
     return np.sqrt(total_variance) * zp
+
+
+def apply_noise(bandflux, bandflux_err, rng=None):
+    """Apply Gaussian noise to a bandflux measurement.
+
+    Parameters
+    ----------
+    bandflux : ndarray of float
+        The bandflux measurement.
+    bandflux_err : ndarray of float
+        The bandflux measurement error.
+    rng : np.random.Generator, optional
+        The random number generator.
+
+    Returns
+    -------
+    ndarray of float
+        The noisy bandflux measurement.
+    """
+    if rng is None:
+        rng = np.random.default_rng()
+
+    return rng.normal(loc=bandflux, scale=bandflux_err)

src/tdastro/astro_utils/opsim.py

@@ -8,7 +8,7 @@
 from scipy.spatial import KDTree
 
 from tdastro.astro_utils.mag_flux import mag2flux
-from tdastro.astro_utils.noise_model import poisson_flux_std
+from tdastro.astro_utils.noise_model import poisson_bandflux_std
 from tdastro.astro_utils.zeropoint import (
     _lsstcam_extinction_coeff,
     _lsstcam_zeropoint_per_sec_zenith,
@@ -363,20 +363,20 @@ def get_observations(self, query_ra, query_dec, radius, cols=None):
             results[col] = self.table[table_col][neighbors].to_numpy()
         return results
 
-    def flux_err_point_source(self, flux, index):
-        """Compute observational flux error for a point source
+    def bandflux_error_point_source(self, bandflux, index):
+        """Compute observational bandflux error for a point source
 
         Parameters
         ----------
-        flux : array_like of float
-            Flux of the point source in nJy.
+        bandflux : array_like of float
+            Band bandflux of the point source in nJy.
         index : array_like of int
             The index of the observation in the OpSim table.
 
         Returns
         -------
         flux_err : array_like of float
-            Simulated flux noise in nJy.
+            Simulated bandflux noise in nJy.
         """
         observations = self.table.iloc[index]
 
@@ -387,8 +387,8 @@ def flux_err_point_source(self, flux, index):
         # table value is in mag/arcsec^2
         sky_njy = mag2flux(observations["skyBrightness"])
 
-        return poisson_flux_std(
-            flux,
+        return poisson_bandflux_std(
+            bandflux,
             pixel_scale=self.pixel_scale,
             total_exposure_time=observations["visitExposureTime"],
             exposure_count=observations["numExposures"],

src/tdastro/astro_utils/passbands.py

@@ -577,7 +577,7 @@ def fluxes_to_bandflux(
 
         F_b = ∫ f(λ)φ_b(λ) dλ
 
-        where f(λ) is the specific flux of an object at the top of the atmosphere, and φ_b(λ) is the
+        where f(λ) is the flux density of an object at the top of the atmosphere, and φ_b(λ) is the
         normalized system response for given band b."
 
         Parameters
@@ -602,7 +602,7 @@ def fluxes_to_bandflux(
             )
 
         # Calculate the bandflux as ∫ f(λ)φ_b(λ) dλ,
-        # where f(λ) is the in-band flux density and φ_b(λ) is the normalized system response
+        # where f(λ) is the flux density and φ_b(λ) is the normalized system response
         integrand = flux_density_matrix * self.processed_transmission_table[:, 1]
         bandfluxes = scipy.integrate.trapezoid(integrand, x=self.waves)
 

src/tdastro/sources/physical_model.py

@@ -18,7 +18,7 @@ class PhysicalModel(ParameterizedNode):
     or constants and are stored in the graph's (external) graph_state dictionary.
 
     Physical models can also have special background pointers that link to another PhysicalModel
-    producing flux. We can chain these to have a supernova in front of a star in front
+    producing flux density. We can chain these to have a supernova in front of a star in front
     of a static background.
 
     Physical models also support adding and applying a variety of effects, such as redshift.
@@ -306,7 +306,7 @@ def get_band_fluxes(self, passband_or_group, times, filters, state) -> np.ndarra
                     "If passband_or_group is a Passband, "
                     "filters must be None or a list of the same filter repeated."
                 )
-            spectral_fluxes = self._evaluate(times, passband_or_group.waves, state)
+            spectral_fluxes = self.evaluate(times, passband_or_group.waves, state)
             return passband_or_group.fluxes_to_bandflux(spectral_fluxes)
 
         if filters is None:
@@ -316,6 +316,6 @@ def get_band_fluxes(self, passband_or_group, times, filters, state) -> np.ndarra
         for filter_name in np.unique(filters):
             passband = passband_or_group.passbands[filter_name]
             filter_mask = filters == filter_name
-            spectral_fluxes = self._evaluate(times[filter_mask], passband.waves, state)
+            spectral_fluxes = self.evaluate(times[filter_mask], passband.waves, state)
             band_fluxes[filter_mask] = passband.fluxes_to_bandflux(spectral_fluxes)
         return band_fluxes

tests/tdastro/astro_utils/test_noise_model.py

@@ -1,6 +1,6 @@
 import numpy as np
 from numpy.testing import assert_allclose
-from tdastro.astro_utils.noise_model import poisson_flux_std
+from tdastro.astro_utils.noise_model import poisson_bandflux_std
 
 
 def test_poisson_flux_std_flux():
@@ -11,8 +11,8 @@ def test_poisson_flux_std_flux():
     flux = 10 ** rng.uniform(0.0, 5.0, 100)
     expected_flux_err = np.sqrt(flux)
 
-    flux_err = poisson_flux_std(
-        flux=flux,
+    flux_err = poisson_bandflux_std(
+        bandflux=flux,
         pixel_scale=rng.uniform(),
         total_exposure_time=rng.uniform(),
         exposure_count=rng.integers(1, 100),
@@ -37,8 +37,8 @@ def test_poisson_flux_std_sky():
     footprint = 10 ** rng.uniform(0.0, 2.0, n)
     expected_flux_err = np.sqrt(sky * footprint)
 
-    flux_err = poisson_flux_std(
-        flux=0.0,
+    flux_err = poisson_bandflux_std(
+        bandflux=0.0,
         pixel_scale=rng.uniform(),
         total_exposure_time=rng.uniform(),
         exposure_count=rng.integers(1, 100),
@@ -65,8 +65,8 @@ def test_poisson_flux_std_readout():
     exposure_count = rng.integers(1, 100, n)
     expected_flux_err = readout_noise * np.sqrt(footprint / pixel_scale**2) * np.sqrt(exposure_count)
 
-    flux_err = poisson_flux_std(
-        flux=0.0,
+    flux_err = poisson_bandflux_std(
+        bandflux=0.0,
         pixel_scale=pixel_scale,
         total_exposure_time=rng.uniform(),
         exposure_count=exposure_count,
@@ -95,8 +95,8 @@ def test_poisson_flux_std_dark():
     dark_current_total = dark_current * total_exposure_time * footprint / pixel_scale**2
     expected_flux_err = np.sqrt(dark_current_total)
 
-    flux_err = poisson_flux_std(
-        flux=0.0,
+    flux_err = poisson_bandflux_std(
+        bandflux=0.0,
         pixel_scale=pixel_scale,
         total_exposure_time=total_exposure_time,
         exposure_count=rng.integers(1, 100, n),

tests/tdastro/astro_utils/test_opsim.py

@@ -248,7 +248,7 @@ def test_opsim_flux_err_point_source(opsim_shorten):
     flux = mag2flux(ops_data.table["fiveSigmaDepth"])
     expected_flux_err = flux / 5.0
 
-    flux_err = ops_data.flux_err_point_source(flux, index=np.arange(len(ops_data)))
+    flux_err = ops_data.bandflux_error_point_source(flux, index=np.arange(len(ops_data)))
 
     # Tolerance is very high, we should investigate why the values are so different.
     np.testing.assert_allclose(flux_err, expected_flux_err, rtol=0.2)

tests/tdastro/astro_utils/test_zeropoint.py

@@ -66,7 +66,7 @@ def test_magnitude_electron_zeropoint_docstring():
 
 
 def test_flux_electron_zeropoint():
-    """Test that flux zeropoints are correct"""
+    """Test that bandflux zeropoints are correct"""
     # Here we just check that magnitude-flux conversion is correct
     airmass = np.array([1, 1.5, 2]).reshape(-1, 1, 1)
     exptime = np.array([30, 38, 45]).reshape(1, -1, 1)