Add test that uses PixelGrid to clean up optical part of PSF

tests/test_sumpsf.py

@@ -500,8 +500,214 @@ def test_mixed_pixel_sum2():
         print('image max diff = ',np.max(np.abs(test_image[b].array - test_star.image[b].array)))
         np.testing.assert_allclose(test_image[b].array, test_star.image[b].array, atol=1.e-7)
 
+@timer
+def test_atmpsf():
+    # Let reality be atmospheric + optical PSF with realistic variation.
+    # Use Kolmogorov with GP for primary model and PixelGrid to mop up the residuals.
+    # The residuals are mostly just in the center of the profile, so the PixelGrid doesn't need
+    # to be very large to fix the main effects of the optical component.
+
+    # Some parameters copied from psf_wf_movie.py in GalSim repo.
+    Ellerbroek_alts = [0.0, 2.58, 5.16, 7.73, 12.89, 15.46]  # km
+    Ellerbroek_weights = [0.652, 0.172, 0.055, 0.025, 0.074, 0.022]
+    Ellerbroek_interp = galsim.LookupTable(Ellerbroek_alts, Ellerbroek_weights,
+                                           interpolant='linear')
+
+    # Use given number of uniformly spaced altitudes
+    alts = np.max(Ellerbroek_alts)*np.arange(6)/5.
+    weights = Ellerbroek_interp(alts)  # interpolate the weights
+    weights /= sum(weights)  # and renormalize
+
+    spd = []  # Wind speed in m/s
+    dirn = [] # Wind direction in radians
+    r0_500 = [] # Fried parameter in m at a wavelength of 500 nm.
+    u = galsim.UniformDeviate(1234)
+    for i in range(6):
+        spd.append(u() * 20)
+        dirn.append(u() * 360*galsim.degrees)
+        r0_500.append(0.1 * weights[i]**(-3./5))
+
+    screens = galsim.Atmosphere(r0_500=r0_500, speed=spd, direction=dirn, altitude=alts, rng=u,
+                                screen_size=102.4, screen_scale=0.1)
+
+    # Add in an optical screen
+    screens.append(galsim.OpticalScreen(diam=4, defocus=1.2, astig1=-0.8, astig2=0.7,
+                                        coma1=0.5, coma2=0.4, spher=-0.9, obscuration=0.4))
+
+    aper = galsim.Aperture(diam=4, obscuration=0.4)
+
+    if __name__ == '__main__':
+        size = 2048
+        nstars = 200
+    else:
+        size = 1024
+        nstars = 50
+
+    pixel_scale = 0.2
+    im = galsim.ImageF(size, size, scale=pixel_scale)
+
+    x = u.np.uniform(25, size-25, size=nstars)
+    y = u.np.uniform(25, size-25, size=nstars)
+
+    for k in range(nstars):
+        flux = 100000
+        theta = ((x[k] - size/2) * pixel_scale * galsim.arcsec,
+                 (y[k] - size/2) * pixel_scale * galsim.arcsec)
+
+        psf = screens.makePSF(lam=500, aper=aper, exptime=100, flux=flux, theta=theta)
+        psf.drawImage(image=im, center=(x[k],y[k]), method='phot', rng=u, add_to_image=True)
+        bounds = galsim.BoundsI(int(x[k]-33), int(x[k]+33), int(y[k]-33), int(y[k]+33))
+
+    # Add a little noise
+    noise = 10
+    im.addNoise(galsim.GaussianNoise(rng=u, sigma=noise))
+    image_file = os.path.join('output', 'atmpsf_im.fits')
+    im.write(image_file)
+
+    # Write out the catalog to a file
+    dtype = [ ('x','f8'), ('y','f8') ]
+    data = np.empty(len(x), dtype=dtype)
+    data['x'] = x
+    data['y'] = y
+    cat_file = os.path.join('output','atmpsf_cat.fits')
+    fitsio.write(cat_file, data, clobber=True)
+
+    psf_file = os.path.join('output','atmpsf.fits')
+    stamp_size = 48
+    config = {
+        'input' : {
+            'image_file_name' : image_file,
+            'cat_file_name' : cat_file,
+            'stamp_size' : 32,
+            'noise': noise**2,
+        },
+        'select' : {
+            'max_snr': 1.e6,
+            'max_edge_frac': 0.1,
+            'hsm_size_reject': True,
+        },
+        'psf' : {
+            'type' : 'Sum',
+            'components': [
+                {
+                    'model' : { 'type' : 'Kolmogorov',
+                                'fastfit': True,
+                              },
+                    'interp' : { 'type' : 'GP',
+                               },
+                },
+                {
+                    'model' : { 'type' : 'PixelGrid',
+                                'init': 'zero',
+                                'scale': pixel_scale,
+                                'size': 12,
+                                'fit_flux': True,
+                              },
+                    'interp' : { 'type' : 'Mean', },
+                }
+            ],
+            'outliers': {
+                'type' : 'Chisq',
+                'nsigma' : 5,
+                'max_remove' : 3,
+            }
+        },
+        'output' : {
+            'file_name' : psf_file,
+            'stats' : [
+                {
+                    'type': 'StarImages',
+                    'file_name': os.path.join('output','test_atmpsf_stars.png'),
+                    'nplot': 10,
+                    'adjust_stars': True,
+                }
+            ],
+        },
+    }
+
+    if __name__ == '__main__':
+        logger = piff.config.setup_logger(verbose=1)
+    else:
+        logger = piff.config.setup_logger(log_file='output/test_atmpsf.log')
+        logger = piff.config.setup_logger(verbose=1)
+
+    psf = piff.process(config, logger)
+
+    if __name__ == '__main__':
+        output = piff.Output.process(config['output'], logger=logger)
+        output.write(psf, logger=logger)
+
+    assert type(psf) is piff.SumPSF
+    assert len(psf.components) == 2
+    assert type(psf.components[0]) is piff.SimplePSF
+    assert type(psf.components[0].model) is piff.Kolmogorov
+    assert type(psf.components[0].interp) is piff.GPInterp
+    assert type(psf.components[1]) is piff.SimplePSF
+    assert type(psf.components[1].model) is piff.PixelGrid
+    assert type(psf.components[1].interp) is piff.Mean
+
+    mean_max_rel_err = 0
+    mean_mse = 0
+    count = 0
+    for i, star in enumerate(psf.stars):
+        if star.is_flagged:
+            continue
+        test_star = psf.drawStar(star)
+
+        b = star.image.bounds.withBorder(-8)
+        max_diff = np.max(np.abs(test_star.image[b].array - star.image[b].array))
+        max_val = np.max(np.abs(star.image[b].array))
+        #print('max_diff / max_val = ',max_diff, max_val, max_diff / max_val)
+        mse = np.sum((test_star.image[b].array - star.image[b].array)**2) / flux**2
+        #print('mse = ',mse)
+        mean_max_rel_err += max_diff/max_val
+        mean_mse += mse
+        count += 1
+
+    mean_max_rel_err /= count
+    mean_mse /= count
+    print('mean maximum relative error = ',mean_max_rel_err)
+    print('mean mean-squared error = ',mean_mse)
+    assert mean_max_rel_err < 0.05
+    assert mean_mse < 2.e-5
+
+    # Without the PixelGrid clean up, it's quite a bit worse.
+    config['psf'] = config['psf']['components'][0]
+    psf = piff.process(config, logger)
+
+    if __name__ == '__main__':
+        config['output']['stats'][0]['file_name'] = os.path.join('output', 'test_atmpsf_stars2.png')
+        output = piff.Output.process(config['output'], logger=logger)
+        output.write(psf, logger=logger)
+
+    mean_max_rel_err = 0
+    mean_mse = 0
+    count = 0
+    for i, star in enumerate(psf.stars):
+        if star.is_flagged:
+            continue
+        test_star = psf.drawStar(star)
+
+        b = star.image.bounds.withBorder(-8)
+        max_diff = np.max(np.abs(test_star.image[b].array - star.image[b].array))
+        max_val = np.max(np.abs(star.image[b].array))
+        #print('max_diff / max_val = ',max_diff / max_val)
+        mse = np.sum((test_star.image[b].array - star.image[b].array)**2) / flux**2
+        #print('mse = ',mse)
+        mean_max_rel_err += max_diff/max_val
+        mean_mse += mse
+        count += 1
+
+    mean_max_rel_err /= count
+    mean_mse /= count
+    print('mean maximum relative error = ',mean_max_rel_err)
+    print('mean mean-squared error = ',mean_mse)
+    assert mean_max_rel_err > 0.07
+    assert mean_mse > 4.e-5
+
 
 if __name__ == '__main__':
     test_trivial_sum1()
     test_easy_sum2()
     test_mixed_pixel_sum2()
+    test_atmpsf()