Add variable PSF option to romanisim and make it the default. (#101)

* Add variable PSF option to romanisim and make it the default.
* Pin pytest at <= 8.0.
* Fix offset treatment.

pyproject.toml

@@ -20,7 +20,7 @@ dependencies = [
     'astropy >=5.3',
     'crds >=10.3.1',
     'defusedxml >=0.5.0',
-    'galsim >=2.5.0',
+    'galsim >=2.5.1',
     'rad >=0.18.0',
     'roman_datamodels >=0.18.0',
     'gwcs >=0.18.1',
@@ -45,7 +45,7 @@ test = [
     'ci-watson >=0.3.0',
     'colorama >=0.4.1',
     'getch >=1.0.0',
-    'pytest >=4.6.0',
+    'pytest >=4.6.0, <=8.0',
     'pytest-openfiles >=0.5.0',
     'pytest-doctestplus >=0.10.0',
     'pytest-cov >=2.9.0',

romanisim/image.py

@@ -278,7 +278,11 @@ def add_objects_to_image(image, objlist, xpos, ypos, psf,
                 raise ValueError('Non-chromatic sources must have specified '
                                  'fluxes!')
             profile = profile.withFlux(obj.flux[filter_name])
-        final = galsim.Convolve(profile * flux_to_counts_factor, psf)
+        if hasattr(psf, 'at_position'):
+            psf0 = psf.at_position(xpos[i], ypos[i])
+        else:
+            psf0 = psf
+        final = galsim.Convolve(profile * flux_to_counts_factor, psf0)
         if chromatic:
             stamp = final.drawImage(
                 bandpass, center=image_pos, wcs=image.wcs.local(image_pos),
@@ -530,7 +534,8 @@ def simulate_counts(metadata, objlist,
             and objlist[0].profile.spectral):
         chromatic = True
     psf = romanisim.psf.make_psf(sca, filter_name, wcs=imwcs,
-                                 chromatic=chromatic, webbpsf=webbpsf)
+                                 chromatic=chromatic, webbpsf=webbpsf,
+                                 variable=True)
     image = galsim.ImageF(roman.n_pix, roman.n_pix, wcs=imwcs, xmin=0, ymin=0)
     SCA_cent_pos = imwcs.toWorld(image.true_center)
     sky_level = roman.getSkyLevel(bandpass, world_pos=SCA_cent_pos,

romanisim/psf.py

@@ -27,9 +27,9 @@
 from romanisim import log
 
 
-def make_psf(sca, filter_name, wcs=None, webbpsf=True, pix=None,
-             chromatic=False, **kw):
-    """Make a PSF profile for Roman.
+def make_one_psf(sca, filter_name, wcs=None, webbpsf=True, pix=None,
+                 chromatic=False, **kw):
+    """Make a PSF profile for Roman at a specific detector location.
 
     Can construct both PSFs using galsim's built-in galsim.roman.roman_psfs
     routine, or can use webbpsf.
@@ -54,7 +54,7 @@ def make_psf(sca, filter_name, wcs=None, webbpsf=True, pix=None,
         galsim profile object for convolution with source profiles when
         rendering scenes.
     """
-    pix = pix if pix is not None else (2044, 2044)
+    pix = pix if pix is not None else (roman.n_pix // 2, roman.n_pix // 2)
     if wcs is None:
         log.warning('wcs is None; unlikely to get orientation of PSF correct.')
     if not webbpsf:
@@ -118,3 +118,92 @@ def make_psf(sca, filter_name, wcs=None, webbpsf=True, pix=None,
     intimg = galsim.InterpolatedImage(
         gimg, normalization='flux', use_true_center=True, offset=centroid)
     return intimg
+
+
+def make_psf(sca, filter_name, wcs=None, webbpsf=True, pix=None,
+             chromatic=False, variable=False, **kw):
+    """Make a PSF profile for Roman.
+
+    Optionally supports spatially variable PSFs via interpolation between
+    the four corners.
+
+    Parameters
+    ----------
+    sca : int
+        SCA number
+    filter_name : str
+        name of filter
+    wcs : callable (optional)
+        function giving mapping from pixels to sky for use in computing local
+        scale of image for webbpsf PSFs
+    pix : tuple (float, float)
+        pixel location of PSF on focal plane
+    variable : bool
+        True if a variable PSF object is desired
+    **kw : dict
+        Additional keywords passed to galsim.roman.getPSF
+
+    Returns
+    -------
+    profile : galsim.gsobject.GSObject
+        galsim profile object for convolution with source profiles when
+        rendering scenes.
+    """
+    if not variable:
+        return make_one_psf(sca, filter_name, wcs=wcs, webbpsf=webbpsf,
+                            pix=pix, chromatic=chromatic, **kw)
+    elif pix is not None:
+        raise ValueError('cannot set both pix and variable')
+    buf = 40
+    # WebbPSF complains if we get too close to (0, 0) for some reason.
+    # For other corners one can go to within a fraction of a pixel.
+    corners = dict(
+        ll=[buf, buf], lr=[roman.n_pix - buf, buf],
+        ul=[buf, roman.n_pix - buf], ur=[roman.n_pix - buf, roman.n_pix - buf])
+    psfs = dict()
+    for corner, pix in corners.items():
+        psfs[corner] = make_one_psf(sca, filter_name, wcs=wcs, webbpsf=webbpsf,
+                                    pix=pix, chromatic=chromatic, **kw)
+    return VariablePSF(corners, psfs)
+
+
+class VariablePSF:
+    """Spatially variable PSF wrapping GalSim profiles.
+
+    Linearly interpolates between four corner PSF profiles by summing
+    weighted GalSim PSF profiles.
+    """
+
+    def __init__(self, corners, psf):
+        self.corners = corners
+        self.psf = psf
+
+    def at_position(self, x, y):
+        """Instantiate a PSF profile at (x, y).
+
+        Linearly interpolate between the four corners to obtain the
+        PSF at this location.
+
+        Parameters
+        ----------
+        x : float
+            x position
+        y : float
+            y position
+
+        Returns
+        -------
+        GalSim profile representing PSF at (x, y).
+        """
+        npix = self.corners['ur'][-1]
+        off = self.corners['ll'][0]
+        wleft = np.clip((npix - x) / (npix - off), 0, 1)
+        wlow = np.clip((npix - y) / (npix - off), 0, 1)
+        # x = [0, off] -> 1
+        # x = [npix, infinity] -> 0
+        # linearly between those, likewise for y.
+        out = (self.psf['ll'] * wleft * wlow +
+               self.psf['lr'] * (1 - wleft) * wlow +
+               self.psf['ul'] * wleft * (1 - wlow) +
+               self.psf['ur'] * (1 - wleft) * (1 - wlow))
+        return out

romanisim/tests/test_psf.py

@@ -33,6 +33,8 @@ def test_make_psf():
     chromatic[3] = True
     bandpass = galsim.roman.getBandpasses(AB_zeropoint=True)['H158']
     vega_sed = galsim.SED('vega.txt', 'nm', 'flambda')
+    varpsf = psf.make_psf(8, 'F087', webbpsf=True, variable=True)
+    psfs.append(varpsf.at_position(100, 100))
     for p, c in zip(psfs, chromatic):
         if not c:
             im = p.drawImage().array