Merge pull request #432 from timedilatesme/supernovae_microlensing

notebook to generate visit images faster with precomputed backgrounds

Author: sibirrer

notebooks/ERL_lensed_AGN_sims.ipynb …/MainNotebooks/ERL_lensed_AGN_sims.ipynbnotebooks/ERL_lensed_AGN_sims.ipynb renamed to notebooks/MainNotebooks/ERL_lensed_AGN_sims.ipynb notebooks/ERL_lensed_AGN_sims.ipynb renamed to notebooks/MainNotebooks/ERL_lensed_AGN_sims.ipynb

@@ -50,6 +50,7 @@ def simulate_roman_image(
     dec=None,
     date=datetime.datetime(year=2027, month=7, day=7, hour=0, minute=0, second=0),
     psf_directory=None,
+    precomputed_background=None,
 ):
     """Creates a Roman-simulated image of a selected lens with noise, using
     galsim's noise settings and PSFs from STPSF.
@@ -101,6 +102,9 @@ def simulate_roman_image(
         Otherwise, the psf will be generated by stpsf on the fly, which is very slow.
         See the note in the ``get_psf`` method's docstring for details on the PSF file naming convention.
     :type psf_directory: str
+    :param precomputed_background: A 2-D numpy array of background counts, shape (num_pix+6, num_pix+6), to be added to the image before noise realizations. This can be computed by calling precompute_roman_background with the same band, num_pix, oversample, exposure_time, and date parameters as passed into this function. If None, the background will be computed on the fly.
+    :type precomputed_background: numpy.ndarray, optional
+
     :return: simulated image in units of flux per second.
     :rtype: numpy.ndarray
     """
@@ -182,18 +186,25 @@ def simulate_roman_image(
     image_no_noise = convolved.drawImage(im)
 
     if add_noise:
-        # Obtain sky and thermal background corresponding to certain band and add it to the image
-        # Poisson noise realization is not handled until later
-        image_with_background = add_sky_plus_thermal_background(
-            image_no_noise,
-            band,
-            detector,
-            num_pix,
-            _exposure_time,
-            ra,
-            dec,
-            date,
-        )
+        if precomputed_background is not None:
+            # Convert precomputed numpy array to a galsim Image with the same
+            # scale and origin as image_no_noise so the addition is element-wise.
+            bg = galsim.ImageF(precomputed_background.astype(np.float32), scale=0.11)
+            bg.setOrigin(0, 0)
+            image_with_background = image_no_noise + bg
+        else:
+            # Obtain sky and thermal background corresponding to certain band and add it to the image
+            # Poisson noise realization is not handled until later
+            image_with_background = add_sky_plus_thermal_background(
+                image_no_noise,
+                band,
+                detector,
+                num_pix,
+                _exposure_time,
+                ra,
+                dec,
+                date,
+            )
 
         # Add noise realizations and detector effects
         # Need to handle each exposure separately to properly take into account read noise and persistence
@@ -420,3 +431,56 @@ def _get_wcs_dict(ra, dec, date):
 
     # NB targ_pos indicates the position to observe at the center of the focal plane array
     return roman.getWCS(world_pos=targ_pos, date=date)
+
+
+def precompute_roman_background(
+    band,
+    num_pix,
+    exposure_time,
+    detector=None,
+    detector_pos=None,
+    ra=None,
+    dec=None,
+    date=datetime.datetime(year=2027, month=7, day=7),
+    oversample=3,
+):
+    """Precompute the Roman sky + thermal background as a numpy array.
+
+    Calls add_sky_plus_thermal_background with a zero image so no lens
+    computation is involved. Pass the returned array to
+    simulate_roman_image via precomputed_background to skip the WCS /
+    sky-level computation on every call.
+
+    :param band: imaging band (e.g. 'F106').
+    :param num_pix: pixels per axis — must match simulate_roman_image.
+    :param exposure_time: seconds — must match simulate_roman_image
+        internally.
+    :param detector: WFI detector (1–18). Random if None.
+    :param detector_pos: (x, y) pixel position. Random if None.
+    :param ra: RA in degrees. Random if None.
+    :param dec: Dec in degrees. Random if None.
+    :param date: observation date; affects zodiacal light level.
+    :param oversample: must match simulate_roman_image.
+    :return: 2-D numpy array of background counts, shape (num_pix+6,
+        num_pix+6).
+    """
+    if detector is None:
+        detector = random.randint(1, 18)
+    if detector_pos is None:
+        x_pos = random.randint(4 + num_pix * oversample, 4092 - num_pix * oversample)
+        y_pos = random.randint(4 + num_pix * oversample, 4092 - num_pix * oversample)
+        detector_pos = (x_pos, y_pos)
+    if ra is None:
+        ra = random.uniform(5, 60)
+    if dec is None:
+        dec = random.uniform(-40, -20)
+
+    num_pix_buffered = num_pix + 6
+
+    zero_image = galsim.ImageF(num_pix_buffered, num_pix_buffered, scale=0.11)
+    zero_image.setOrigin(0, 0)
+
+    background_image = add_sky_plus_thermal_background(
+        zero_image, band, detector, num_pix_buffered, exposure_time, ra, dec, date
+    )
+    return background_image.array

notebooks/ERL_lensed_SNe_sims.ipynb …/MainNotebooks/ERL_lensed_SNe_sims.ipynbnotebooks/ERL_lensed_SNe_sims.ipynb renamed to notebooks/MainNotebooks/ERL_lensed_SNe_sims.ipynb notebooks/ERL_lensed_SNe_sims.ipynb renamed to notebooks/MainNotebooks/ERL_lensed_SNe_sims.ipynb

@@ -1,9 +1,11 @@
+# test_roman_image_simulation.py
 import astropy.cosmology
 import numpy as np
 import numpy.testing as npt
 from slsim.Lenses.lens import Lens
 from slsim.ImageSimulation.roman_image_simulation import (
     simulate_roman_image,
+    precompute_roman_background,
 )
 from slsim.ImageSimulation import image_quality_lenstronomy
 from slsim.ImageSimulation.image_simulation import simulate_image
@@ -294,6 +296,23 @@ def test_simulate_roman_image_with_psf_without_noise():
         np.sum(galsim_image), np.sum(image_ref), rtol=0, atol=0.1
     )
 
+    # precomputed_background is only used inside the add_noise branch, so a dummy
+    # array must not affect the noiseless result
+    dummy_background = np.ones((45 + 6, 45 + 6), dtype=np.float32)
+    galsim_image_with_dummy_bg = simulate_roman_image(
+        lens_class=LENS,
+        band=BAND,
+        num_pix=45,
+        oversample=3,
+        seed=42,
+        add_noise=False,
+        psf_directory=PSF_DIRECTORY,
+        detector=1,
+        detector_pos=(2000, 2000),
+        precomputed_background=dummy_background,
+    )
+    npt.assert_array_equal(galsim_image_with_dummy_bg, galsim_image)
+
 
 def test_simulate_roman_image_with_time_variable_source():
     detector_kwargs = {
@@ -350,5 +369,79 @@ def test_simulate_roman_image_with_time_variable_source():
     assert not np.allclose(lens_image, lens_image3, atol=1e-16, rtol=1e-16)
 
 
+def test_precompute_roman_background():
+    kwargs_single_band = image_quality_lenstronomy.kwargs_single_band(
+        observatory="Roman", band=BAND
+    )
+
+    background = precompute_roman_background(
+        band=BAND,
+        num_pix=45,
+        exposure_time=kwargs_single_band["exposure_time"],
+        **DETECTOR_KWARGS,
+    )
+
+    assert isinstance(background, np.ndarray)
+    assert background.shape == (45 + 6, 45 + 6)
+    assert np.all(np.isfinite(background))
+    assert np.all(background >= 0)
+
+    # testing the randomised-defaults (detector, detector_pos, ra, dec all None)
+    background_random = precompute_roman_background(
+        band=BAND,
+        num_pix=45,
+        exposure_time=kwargs_single_band["exposure_time"],
+    )
+    assert background_random.shape == (45 + 6, 45 + 6)
+
+
+def test_simulate_roman_image_with_precomputed_background():
+    kwargs_single_band = image_quality_lenstronomy.kwargs_single_band(
+        observatory="Roman", band=BAND
+    )
+    exposure_time = kwargs_single_band["exposure_time"]
+
+    background = precompute_roman_background(
+        band=BAND,
+        num_pix=45,
+        exposure_time=exposure_time,
+        **DETECTOR_KWARGS,
+    )
+
+    image_precomputed = simulate_roman_image(
+        lens_class=LENS,
+        band=BAND,
+        num_pix=45,
+        oversample=3,
+        add_noise=True,
+        subtract_mean_background=True,
+        seed=42,
+        exposure_time=exposure_time,
+        psf_directory=PSF_DIRECTORY,
+        precomputed_background=background,
+        **DETECTOR_KWARGS,
+    )
+
+    image_direct = simulate_roman_image(
+        lens_class=LENS,
+        band=BAND,
+        num_pix=45,
+        oversample=3,
+        add_noise=True,
+        subtract_mean_background=True,
+        seed=42,
+        exposure_time=exposure_time,
+        psf_directory=PSF_DIRECTORY,
+        **DETECTOR_KWARGS,
+    )
+
+    assert image_precomputed.shape == (45, 45)
+    assert image_direct.shape == (45, 45)
+    diff = np.abs(image_precomputed - image_direct) / (
+        np.abs(image_precomputed) + np.abs(image_direct) + 1e-10
+    )
+    npt.assert_array_less(np.mean(diff), 0.1)
+
+
 if __name__ == "__main__":
     pytest.main()