Image simulator upgrades (c3-time-domain#120)

Add galaxies, saturated pixel bleed trails, satellite streaks, and cosmic ray muon tracks.

.gitignore

@@ -7,8 +7,12 @@ tests/local_config.yaml
 tests/local_overrides.yaml
 tests/local_augments.yaml
 tests/improc/cache/*
-data/DECam_examples/c4d_221104_074232_ori.fits.fz
+data/test_data/DECam_examples/c4d_221104_074232_ori.fits.fz
+data/test_data/DECam_examples/*_cached
+data/DECam_default_calibrators
+data/GaiaDR3_excerpt
 .pytest.ini
+tests/plots/*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

docker/application/Dockerfile

@@ -178,6 +178,7 @@ RUN pip install \
       pandas==1.5.3 \
       photutils==1.6.0 \
       psycopg2==2.9.5 \
+      pylandau==2.2.1 \
       pytest==7.2.2 \
       pytest-timeout==2.1.0 \
       python-dateutil==2.8.2 \

requirements.txt

@@ -29,11 +29,13 @@ numexpr==2.8.4
 numpy==1.24.2
 packaging==23.0
 pandas==1.5.3
+photutils==1.9.0
 Pillow==10.0.1
 pluggy==1.0.0
 psycopg2==2.9.5
 py-cpuinfo==9.0.0
 pyerfa==2.0.0.2
+pylandau==2.2.1
 pyparsing==3.0.9
 pytest==7.2.2
 python-dateutil==2.8.2

tests/conftest.py

@@ -70,15 +70,24 @@ def tests_setup_and_teardown():
         session.commit()
 
 
-@pytest.fixture
-def headless_plots():
+@pytest.fixture(scope="session")
+def blocking_plots():
     import matplotlib
-
     backend = matplotlib.get_backend()
-    # ref: https://stackoverflow.com/questions/15713279/calling-pylab-savefig-without-display-in-ipython
-    matplotlib.use("Agg")
 
-    yield None
+    # make sure there's a folder to put the plots in
+    if not os.path.isdir(os.path.join(CODE_ROOT, 'tests/plots')):
+        os.makedirs(os.path.join(CODE_ROOT, 'tests/plots'))
+
+    inter = os.getenv('INTERACTIVE', False)
+    if isinstance(inter, str):
+        inter = inter.lower() in ('true', '1')
+
+    if not inter:  # for non-interactive plots, use headless plots that just save to disk
+        # ref: https://stackoverflow.com/questions/15713279/calling-pylab-savefig-without-display-in-ipython
+        matplotlib.use("Agg")
+
+    yield inter
 
     matplotlib.use(backend)
 

tests/docker-compose.yaml

@@ -103,6 +103,7 @@ services:
     environment:
       SEECHANGE_CONFIG: /seechange/tests/seechange_config_test.yaml
       SEECHANGE_TEST_ARCHIVE_DIR: /archive_storage/base
+      RUN_SLOW_TESTS: 1
     depends_on:
       setuptables:
         condition: service_completed_successfully

tests/improc/test_simulator.py

@@ -0,0 +1,190 @@
+import os
+import pytest
+import time
+
+import numpy as np
+
+import matplotlib.pyplot as plt
+
+from models.base import CODE_ROOT
+from improc.simulator import Simulator, SimGalaxies, SimStreaks, SimCosmicRays
+from improc.sky_flat import sigma_clipping
+
+# uncomment this to run the plotting tests interactively
+# os.environ['INTERACTIVE'] = '1'
+
+
+@pytest.mark.skipif( os.getenv('INTERACTIVE') is None, reason='Set INTERACTIVE to run this test' )
+def test_make_star_field(blocking_plots):
+    s = Simulator( image_size_x=256, star_number=1000, galaxy_number=0)
+    s.make_image()
+    vmin = np.percentile(s.image, 1)
+    vmax = np.percentile(s.image, 99)
+
+    plt.imshow(s.image, vmin=vmin, vmax=vmax)
+    plt.show(block=blocking_plots)
+
+    filename = os.path.join(CODE_ROOT, 'tests/plots/simulator_star_field')
+    plt.savefig(filename+'.png')
+    plt.savefig(filename+'.pdf')
+
+
+@pytest.mark.skipif( os.getenv('INTERACTIVE') is None, reason='Set INTERACTIVE to run this test' )
+def test_make_galaxy_field(blocking_plots):
+    s = Simulator( image_size_x=256, star_number=0, galaxy_number=1000, galaxy_min_width=1, galaxy_min_flux=1000 )
+    t0 = time.time()
+    s.make_image()
+    print(s.truth.oversampling)
+    print(f'Generating image took {time.time() - t0:.1f} seconds')
+
+    mu, sig = sigma_clipping(s.image.astype(float))
+
+    plt.imshow(s.image, vmin=mu-sig, vmax=mu+10*sig)
+
+    plt.show(block=blocking_plots)
+
+    filename = os.path.join(CODE_ROOT, 'tests/plots/simulator_galaxy_field')
+    plt.savefig(filename+'.png')
+    plt.savefig(filename+'.pdf')
+
+
+@pytest.mark.parametrize("exp_scale", [5.0, 500.0])
+def test_making_galaxy_in_image(exp_scale, blocking_plots):
+    center_x = 33.123
+    center_y = 76.456
+    imsize = (256, 256)
+    sersic_scale = 5.0
+    exp_flux = 1000
+    sersic_flux = 0
+    rotation = 33.0
+    cos_i = 0.5
+    cutoff_radius = exp_scale * 5.0
+
+    im1 = SimGalaxies.make_galaxy_image(
+        imsize=imsize,
+        center_x=center_x,
+        center_y=center_y,
+        exp_scale=exp_scale,
+        sersic_scale=sersic_scale,
+        exp_flux=exp_flux,
+        sersic_flux=sersic_flux,
+        cos_i=cos_i,
+        rotation=rotation,
+        cutoff_radius=cutoff_radius,
+    )
+
+    im2 = np.zeros(imsize)
+    SimGalaxies.add_galaxy_to_image(
+        image=im2,
+        center_x=center_x,
+        center_y=center_y,
+        exp_scale=exp_scale,
+        sersic_scale=sersic_scale,
+        exp_flux=exp_flux,
+        sersic_flux=sersic_flux,
+        cos_i=cos_i,
+        rotation=rotation,
+        cutoff_radius=cutoff_radius,
+    )
+
+    diff = im1 - im2
+
+    if blocking_plots:
+        fig, ax = plt.subplots(1, 3)
+
+        ax[0].imshow(im1)
+        ax[0].set_title('make_galaxy_image')
+
+        ax[1].imshow(im2)
+        ax[1].set_title('add_galaxy_to_image')
+
+        h = ax[2].imshow(np.log10(abs(diff)))
+        ax[2].set_title('log10 of difference')
+        plt.colorbar(h, ax=ax[2], orientation='vertical')
+
+        plt.show(block=True)
+
+    peak1 = np.unravel_index(im1.argmax(), im1.shape)
+    peak2 = np.unravel_index(im2.argmax(), im2.shape)
+
+    assert peak1 == peak2
+    assert np.max(diff) < 1e-5
+
+
+def test_bleeding_pixels(blocking_plots):
+    s = Simulator(bleed_fraction_up=0.1, bleed_fraction_down=0.5, saturation_limit=1000)
+    s.make_sensor()
+    im = np.random.normal(0, 1, (256, 256))  # no saturated pixels
+
+    # when no saturated pixels exist, the bleeds should be a no-op
+    im_bleed = s.add_bleeds_to_image(im)
+
+    diff = im_bleed - im
+    assert np.max(diff) < 1e-5
+
+    # add a saturated pixel
+    im[100, 20] = 1000 * 53
+    im_bleed = s.add_bleeds_to_image(im)
+
+    expected_down_electrons = 1000 * 53 * 0.1 - 1000
+    expected_up_electrons = 1000 * 53 * 0.5 - 1000
+    original_electrons_in_pixel = 1000
+    expected_total_electrons = expected_down_electrons + expected_up_electrons + original_electrons_in_pixel
+    standard_deviation = 256  # one per pixel, 256*256 pixels, sqrt
+    # print(expected_total_electrons)
+    # print(np.sum(im_bleed))
+    assert abs(np.sum(im_bleed) - expected_total_electrons) < 3 * standard_deviation
+
+    # add another pixel
+    im[101, 20] = 1000 * 130
+    im_bleed = s.add_bleeds_to_image(im)
+
+    expected_down_electrons = 1000 * 183 * 0.1 - 2000
+    expected_up_electrons = 1000 * 183 * 0.5 - 2000
+    original_electrons_in_pixel = 2000
+    expected_total_electrons = expected_down_electrons + expected_up_electrons + original_electrons_in_pixel
+    # print(expected_total_electrons)
+    # print(np.sum(im_bleed))
+    assert abs(np.sum(im_bleed) - expected_total_electrons) < 3 * standard_deviation
+
+    # test horizontal bleeding
+    s.sensor.bleed_vertical = False
+
+    im_bleed = s.add_bleeds_to_image(im)
+    assert abs(np.sum(im_bleed) - expected_total_electrons) < 3 * standard_deviation
+
+    if blocking_plots:
+        plt.imshow(im_bleed)
+        plt.show(block=True)
+
+
+@pytest.mark.skipif( os.getenv('INTERACTIVE') is None, reason='Set INTERACTIVE to run this test' )
+def test_streak_images(blocking_plots):
+    im = SimStreaks.make_streak_image(center_x=50.3, length=25)
+
+    # if blocking_plots:
+    #     plt.imshow(im)
+    #     plt.show(block=True)
+
+    s = Simulator(streak_number=10, star_number=0, galaxy_number=0)
+    s.make_image()
+
+    if blocking_plots:
+        plt.imshow(s.image)
+        plt.show(block=True)
+
+
+@pytest.mark.skipif( os.getenv('INTERACTIVE') is None, reason='Set INTERACTIVE to run this test' )
+def test_track_images(blocking_plots):
+    im = SimCosmicRays.make_track_image(center_x=50.3, length=25, energy=10)
+
+    # if blocking_plots:
+    #     plt.imshow(im)
+    #     plt.show(block=True)
+
+    s = Simulator(streak_number=0, star_number=0, galaxy_number=0, track_number=30)
+    s.make_image()
+
+    if blocking_plots:
+        plt.imshow(s.image, vmin=100, vmax=200)
+        plt.show(block=True)

tests/improc/test_sky_flat.py

@@ -11,6 +11,7 @@
 from improc.simulator import Simulator
 from improc.sky_flat import calc_sky_flat
 
+
 @pytest.mark.flaky(max_runs=3)
 @pytest.mark.parametrize("num_images", [10, 300])
 def test_simple_sky_flat(num_images):

tests/improc/test_zogy.py

@@ -156,7 +156,8 @@ def test_subtraction_no_new_sources():
     assert zogy_failures == 0
 
 
-def test_subtraction_snr_histograms(headless_plots):
+@pytest.mark.skipif( os.getenv('INTERACTIVE') is None, reason='Set INTERACTIVE to run this test' )
+def test_subtraction_snr_histograms(blocking_plots):
     background = 5.0
     seeing = 3.0
     iterations = 300
@@ -194,33 +195,10 @@ def test_subtraction_snr_histograms(headless_plots):
     measured = np.zeros((iterations, len(fluxes)))
 
     for j in range(iterations):
-
-        # remove the additional stars from the reference image
-        # sim.stars.remove_stars(len(sim.stars.star_mean_fluxes) - sim.pars.star_number)
-
-        # make a reference image
-        # sim.make_image(new_sky=True)
-        # truth1 = sim.truth
-        # im1 = sim.apply_bias_correction(sim.image)
-        # im1 -= truth1.background_instance
-        # psf1 = truth1.psf_downsampled
-        # bkg1 = truth1.total_bkg_var
-
         psf1 = sim.psf_downsampled
         bkg1 = background  # TODO: make this slightly variable?
         im1 = np.random.normal(0, np.sqrt(bkg1), size=transients_overlay.shape)
 
-        # # make a new image, with transients
-        # for f, p in zip(fluxes, pos):
-        #     sim.add_extra_stars(flux=f, x=p, y=p)
-
-        # sim.make_image(new_sky=True)
-        # truth2 = sim.truth
-        # im2 = sim.apply_bias_correction(sim.image)
-        # im2 -= truth2.background_instance
-        # psf2 = truth2.psf_downsampled
-        # bkg2 = truth2.total_bkg_var
-
         psf2 = sim.psf_downsampled
         bkg2 = background  # TODO: make this slightly variable?
         im2 = np.random.normal(transients_overlay, np.sqrt(bkg2 + transients_overlay))
@@ -261,18 +239,16 @@ def test_subtraction_snr_histograms(headless_plots):
         ax.set_xlabel('Measured S/N')
         ax.legend(loc='upper right')
         plt.subplots_adjust(hspace=0.5)
-    plt.show()
+    plt.show(block=blocking_plots)
 
-    plot_path = os.path.join(CODE_ROOT, "tests/plots")
-    if not os.path.isdir(plot_path):
-        os.mkdir(plot_path)
-    plt.savefig(os.path.join(plot_path, "zogy_snr_histograms.pdf"))
-    plt.savefig(os.path.join(plot_path, "zogy_snr_histograms.png"))
+    filename = os.path.join(CODE_ROOT, "tests/plots/zogy_snr_histograms")
+    plt.savefig(filename + ".png")
+    plt.savefig(filename + ".pdf")
 
 
 @pytest.mark.skip( reason="This test frequently fails even with the flaky.  Can we use a random seed?" )
 @pytest.mark.flaky(max_runs=3)
-def test_subtraction_new_sources_snr():
+def test_subtraction_new_sources_snr(blocking_plots):
     num_stars = 300
     sim = Simulator(
         image_size_x=imsize,  # not too big, but allow some space for stars
@@ -349,12 +325,13 @@ def test_subtraction_new_sources_snr():
     assert chi2 / (len(fluxes) - 2) < 2.0  # the fit should be good
     assert p1[0] > 0.8  # should be close to one, but we are losing about 15% S/N and I don't know why
 
-    # plt.errorbar(exp, mean, err, fmt='.', label='measured vs. expected')
-    # plt.plot(exp, np.polyval(p1, exp), 'r-', label=f'fit: {p1[0]:.3f}x + {p1[1]:.3f}')
-    # plt.xlabel('expected S/N')
-    # plt.ylabel('measured S/N')
-    # plt.legend()
-    # plt.show(block=True)
+    if blocking_plots:
+        plt.errorbar(exp, mean, err, fmt='.', label='measured vs. expected')
+        plt.plot(exp, np.polyval(p1, exp), 'r-', label=f'fit: {p1[0]:.3f}x + {p1[1]:.3f}')
+        plt.xlabel('expected S/N')
+        plt.ylabel('measured S/N')
+        plt.legend()
+        plt.show(block=True)
 
 
 @pytest.mark.skip( reason="This test frequently fails even with the flaky.  Can we use a random seed?" )

tests/pipeline/test_extraction.py

@@ -126,6 +126,7 @@ def test_sep_save_source_list(decam_small_image, provenance_base, code_version):
                 session.execute(sa.delete(Image).where(Image.id == image_id))
             session.commit()
 
+
 # This is running sextractor in one particular way that is used by more than one test
 @pytest.fixture
 def run_sextractor( decam_example_reduced_image_ds ):
@@ -144,6 +145,7 @@ def run_sextractor( decam_example_reduced_image_ds ):
 
     yield sourcelist, sourcefile
 
+
 def test_sextractor_extract_once( decam_example_reduced_image_ds, run_sextractor ):
     sourcelist, sourcefile = run_sextractor
 
@@ -198,6 +200,7 @@ def test_sextractor_extract_once( decam_example_reduced_image_ds, run_sextractor
     assert sourcelist.apfluxadu(apnum=0)[0].min() == pytest.approx( 89.445114, rel=1e-5 )
     assert sourcelist.apfluxadu(apnum=0)[0].max() == pytest.approx( 557651.8, rel=1e-5 )
 
+
 def test_run_psfex( decam_example_reduced_image_ds ):
     sourcelist = decam_example_reduced_image_ds.sources
     tempname = ''.join( random.choices( 'abcdefghijklmnopqrstuvwxyz', k=10 ) )

tests/pipeline/test_photo_cal.py

@@ -10,7 +10,7 @@
 from pipeline.photo_cal import PhotCalibrator
 
 
-def test_decam_photo_cal( decam_example_reduced_image_ds_with_zp, headless_plots ):
+def test_decam_photo_cal( decam_example_reduced_image_ds_with_zp, blocking_plots ):
     ds, photomotor = decam_example_reduced_image_ds_with_zp
 
     with SmartSession() as session: