Implement distortion for Roman

pyproject.toml

@@ -21,6 +21,7 @@ dependencies = [
     "photutils>=1.0.0",
     "poppy>=1.0.0",
     "pysiaf>=0.19.1",
+    "soc_roman_tools>=0.1.0",
     "synphot>=1.0.0",
 ]
 

webbpsf/distortion.py

@@ -6,6 +6,8 @@
 from scipy.interpolate import RegularGridInterpolator
 from scipy.ndimage import rotate
 
+from soc_roman_tools.siaf.siaf import RomanSiaf
+
 def _get_default_siaf(instrument, aper_name):
     """
     Create instance of pysiaf for the input instrument and aperture
@@ -33,8 +35,12 @@ def _get_default_siaf(instrument, aper_name):
         siaf_name = instrument
 
     # Select a single SIAF aperture
-    siaf = pysiaf.Siaf(siaf_name)
-    aper = siaf.apertures[aper_name]
+    if instrument=='WFI':
+        siaf = RomanSiaf()
+        aper = siaf[aper_name]
+    else:
+        siaf = pysiaf.Siaf(siaf_name)
+        aper = siaf.apertures[aper_name]
 
     return aper
 
@@ -107,11 +113,16 @@ def distort_image(hdulist_or_filename, ext=0, to_frame='sci', fill_value=0,
 
     if aper is None:
         # Log instrument and detector names
-        instrument = hdu_list[0].header["INSTRUME"].upper()
-        aper_name = hdu_list[0].header["APERNAME"].upper()
+        instrument = hdu_list[0].header["INSTRUME"].upper().strip()
+
+        if instrument == 'WFI':
+            aper_name = 'WFI' + hdu_list[0].header["DETECTOR"][-2:] + "_FULL"
+        else:
+            aper_name = hdu_list[0].header["APERNAME"].upper()
+
         # Pull default values
         aper = _get_default_siaf(instrument, aper_name)
-    
+
     # Pixel scale information
     ny, nx = hdu_list[ext].shape
     pixelscale = hdu_list[ext].header["PIXELSCL"]  # the pixel scale carries the over-sample value
@@ -222,8 +233,11 @@ def apply_distortion(hdulist_or_filename=None, fill_value=0):
     ext = 1  # edit the oversampled PSF (OVERDIST extension)
 
     # Log instrument and detector names
-    instrument = hdu_list[0].header["INSTRUME"].upper()
-    aper_name = hdu_list[0].header["APERNAME"].upper()
+    instrument = hdu_list[0].header["INSTRUME"].upper().strip()
+    if instrument == 'WFI':
+        aper_name = 'WFI' + hdu_list[0].header["DETECTOR"][-2:] + "_FULL"
+    else:
+        aper_name = hdu_list[0].header["APERNAME"].upper()
 
     # Pull default values
     aper = _get_default_siaf(instrument, aper_name)
@@ -287,12 +301,17 @@ def apply_rotation(hdulist_or_filename=None, rotate_value=None, crop=True):
     psf = copy.deepcopy(hdu_list)
 
     # Log instrument and detector names
-    instrument = hdu_list[0].header["INSTRUME"].upper()
-    aper_name = hdu_list[0].header["APERNAME"].upper()
+    instrument = hdu_list[0].header["INSTRUME"].upper().strip()
+    if instrument == 'WFI':
+        aper_name = 'WFI' + hdu_list[0].header["DETECTOR"][-2:] + "_FULL"
+    else:
+        aper_name = hdu_list[0].header["APERNAME"].upper()
 
     if instrument in ["MIRI", "NIRSPEC"]:
         raise ValueError("{}'s rotation is already included in WebbPSF and "
                          "shouldn't be added again.".format(instrument))
+    if instrument == "WFI":
+        raise ValueError("Rotation not necessary for {:} as pupil are aligned with SCAs (to confirm).".format(instrument))
 
     # Set rotation value if not already set by a keyword argument
     if rotate_value is None:

webbpsf/gridded_library.py

@@ -141,8 +141,7 @@ def __init__(self, instrument, filter_name, detectors="all", num_psfs=16, psf_lo
         # Set PSF attributes for the 3 kwargs that will be used before the calc_psf() call
         if "add_distortion" in kwargs:
             self.add_distortion = kwargs["add_distortion"]
-            if self.webb.name == "WFI":
-                del kwargs["add_distortion"]
+
         else:
             self.add_distortion = True
             kwargs["add_distortion"] = self.add_distortion

webbpsf/roman.py

@@ -19,9 +19,9 @@
 
 from . import utils
 from . import webbpsf_core
+from . import distortion
 from .optics import _fix_zgrid_NaNs
 
-
 _log = logging.getLogger('webbpsf')
 import pprint
 
@@ -281,35 +281,42 @@ def __init__(self, *args, **kwargs):
     def calc_psf(self, outfile=None, source=None, nlambda=None, monochromatic=None,
                  fov_arcsec=None, fov_pixels=None, oversample=None, detector_oversample=None, fft_oversample=None,
                  overwrite=True, display=False, save_intermediates=False, return_intermediates=False,
-                 normalize='first', add_distortion=False, crop_psf=False):
+                 normalize='first', add_distortion=True, crop_psf=False):
         """
         Compute a PSF
 
         Parameters
         ----------
         add_distortion : bool
-            Included for API compatibility with the JWST instrument classes, but has no
-            effect on the results for Roman WFI PSF calculations.
+            If True, will add 2 new extensions to the PSF HDUlist object. The 2nd extension
+            will be a distorted version of the over-sampled PSF and the 3rd extension will
+            be a distorted version of the detector-sampled PSF.
         crop_psf : bool
             Included for API compatibility with the JWST instrument classes, but has no
             effect on the results for Roman WFI PSF calculations.
 
         """
 
-        if add_distortion is not False or crop_psf is not False:
-            raise AttributeError('`add_distortion` and `crop_psf` still under '
-                                 'development for Roman WFI')
-            # return default values to True once implemented
+        # Save new keywords to the options dictionary
+        self.options['add_distortion'] = add_distortion
+        self.options['crop_psf'] = crop_psf
 
+        # add_distortion keyword is not implemented for RomanCoronagraph Class
+        if self.name == "RomanCoronagraph" and add_distortion==True:
+            self.options['add_distortion'] = False
+            self.options['crop_psf'] = False
+            _log.info("Geometric distortions are not implemented in WebbPSF for Roman CGI. The add_distortion keyword must be set to False for this case.")
+
         # Run poppy calc_psf
         psf = webbpsf_core.SpaceTelescopeInstrument.calc_psf(self, outfile=outfile, source=source, nlambda=nlambda,
                                                 monochromatic=monochromatic, fov_arcsec=fov_arcsec,
                                                 fov_pixels=fov_pixels, oversample=oversample,
                                                 detector_oversample=detector_oversample, fft_oversample=fft_oversample,
                                                 overwrite=overwrite, display=display,
-                                                save_intermediates=save_intermediates,
+                                                save_intermediates=save_intermediates, 
                                                 return_intermediates=return_intermediates, normalize=normalize)
 
+
         return psf
 
     # slightly different versions of the following two functions
@@ -352,6 +359,51 @@ def _get_fits_header(self, result, options):
         result[0].header['DETECTOR'] = (self.detector, 'Detector selected')
 
 
+    def _calc_psf_format_output(self, result, options):
+        """
+        Add distortion to the created 1-extension PSF
+
+        Apply desired formatting to output file:
+                 - rebin to detector pixel scale if desired
+                 - set up FITS extensions if desired
+                 - output either the oversampled, rebinned, or both
+        Which image(s) get output depends on the value of the options['output_mode']
+        parameter. It may be set to 'Oversampled image' to output just the oversampled image,
+        'Detector sampled image' to output just the image binned down onto detector pixels, or
+        'Both as FITS extensions' to output the oversampled image as primary HDU and the
+        rebinned image as the first image extension. For convenience, the option can be set
+        to just 'oversampled', 'detector', or 'both'.
+
+        Modifies the 'result' HDUList object.
+
+        """
+       # Pull values from options dictionary
+        add_distortion = options.get('add_distortion', True)
+        crop_psf = options.get('crop_psf', True)
+        # Add distortion if set in calc_psf
+        if add_distortion:
+            _log.debug("Adding PSF distortion(s)")
+
+            # Set up new extensions to add distortion to:
+            n_exts = len(result)
+            for ext in np.arange(n_exts):
+                hdu_new = fits.ImageHDU(result[ext].data, result[ext].header)  # these will be the PSFs that are edited
+                result.append(hdu_new)
+                ext_new = ext + n_exts
+                result[ext_new].header["EXTNAME"] = result[ext].header["EXTNAME"][0:4] + "DIST"  # change extension name
+                _log.debug("Appending new extension {} with EXTNAME = {}".format(ext_new, result[ext_new].header["EXTNAME"]))
+
+            _log.debug("WFI: Adding optical distortion")
+            psf_distorted = distortion.apply_distortion(result)  # apply siaf distortion model
+
+            # Edit the variable to match if input didn't request distortion
+            # (cannot set result = psf_distorted due to return method)
+            [result.append(fits.ImageHDU()) for i in np.arange(len(psf_distorted) - len(result))]
+            for ext in np.arange(len(psf_distorted)): result[ext] = psf_distorted[ext]
+
+        # Rewrite result variable based on output_mode set:
+        webbpsf_core.SpaceTelescopeInstrument._calc_psf_format_output(self, result, options)
+
 class WFIPupilController:
     """
     This is a helper class for the WFI and is used to swap in
@@ -591,6 +643,7 @@ def __init__(self):
         self.pupilopd = self.opd_list[-1]
 
     def _addAdditionalOptics(self, optsys, **kwargs):
+        _log.debug("   No optics added for WFI")
         return optsys, False, None
 
     def _load_detector_aberrations(self, path):

webbpsf/tests/test_psfgrid.py

@@ -239,14 +239,13 @@ def test_2d_to_griddedpsfmodel():
 def test_wfi():
     """Test that the psf_grid method works for the WFI class"""
 
-    # Check add_distortion not specified defaults to false
     oversample = 2
     fov_pixels = 10
     nlambda = 1
 
     # Create PSF grid
     wfi = roman.WFI()
-    grid = wfi.psf_grid(all_detectors=False, num_psfs=4, fov_pixels=fov_pixels, oversample=oversample, nlambda=nlambda, verbose=False)
+    grid = wfi.psf_grid(all_detectors=False, add_distortion=False, num_psfs=4, fov_pixels=fov_pixels, oversample=oversample, nlambda=nlambda, verbose=False)
 
     # Pull one of the PSFs out of the grid
     psfnum = 1
@@ -267,10 +266,3 @@ def test_wfi():
     assert np.allclose(gridpsf, convpsf*scalefactor), "Data values not as expected"
 
 
-def test_wfi_error():
-    """Check add_distortion=True raises an error"""
-
-    with pytest.raises(NotImplementedError) as excinfo:
-        wfi = roman.WFI()
-        wfi.psf_grid(add_distortion=True, num_psfs=1, fov_pixels=1, detector_oversample=2)
-    assert "NotImplementedError" in str(excinfo)

webbpsf/webbpsf_core.py

@@ -704,13 +704,6 @@ def psf_grid(self, num_psfs=16, all_detectors=True, save=False,
         else:
             psf_location = self.detector_position[::-1]  # (y,x)
 
-        # add_distortion keyword is not implemented for WFI Class
-        if self.name == "WFI" and "add_distortion" not in kwargs:
-            kwargs["add_distortion"] = False
-        elif self.name == "WFI" and kwargs["add_distortion"] == True:
-            raise NotImplementedError("Geometric distortions are not implemented in WebbPSF for WFI Instrument. "
-                                      "The add_distortion keyword must be set to False for this case.")
-
         # Call CreatePSFLibrary class
         inst = gridded_library.CreatePSFLibrary(instrument=self, filter_name=filt, detectors=detectors,
                                                 num_psfs=num_psfs, psf_location=psf_location,