Make the RomanOptics chromatic=True option more effective and more efficient.

piff/__init__.py

@@ -50,6 +50,8 @@
 # Also let piff.version show the version.
 version = __version__
 
+from . import galsim_patch
+
 # We don't have any C functions, but once we do, I recommend using cffi to
 # wrap them.  This is the entire code we need to get C functions into Python.
 if False:

piff/config.py

@@ -159,11 +159,12 @@ def process(config, logger=None):
             config['select'][key] = config['input'].pop(key)
 
     # read in the input images
-    objects, wcs, pointing = Input.process(config['input'], logger=logger)
+    objects, wcs, pointing, bandpass = Input.process(config['input'], logger=logger)
     stars = Select.process(config.get('select',{}), objects, logger=logger)
 
     psf = PSF.process(config['psf'], logger=logger)
-    psf.fit(stars, wcs, pointing, logger=logger)
+    psf.set_context(wcs, pointing, bandpass)
+    psf.fit(stars, logger=logger)
 
     # Attach these for reference
     psf.initial_objects = objects

piff/convolvepsf.py

@@ -106,6 +106,12 @@ def interp_property_names(self):
             names.update(c.interp_property_names)
         return names
 
+    def set_context(self, wcs, pointing, bandpass):
+        super().set_context(wcs, pointing, bandpass)
+        if isinstance(self.components, list):
+            for comp in self.components:
+                comp.set_context(wcs, pointing, bandpass)
+
     @classmethod
     def parseKwargs(cls, config_psf, logger):
         """Parse the psf field of a configuration dict and return the kwargs to use for

piff/galsim_patch.py

@@ -0,0 +1,92 @@
+import numpy as np
+import galsim
+from astropy import units
+from galsim import utilities
+from galsim.table import LookupTable, _LookupTable
+from galsim.sed import SED
+from galsim.errors import GalSimError
+
+# Temporary compatibility patch for GalSim < 2.9.
+# Remove this module once Piff can depend on GalSim 2.9+, which includes thin(bandpass=...).
+
+def _piff_sed_thin(self, rel_err=1.e-4, trim_zeros=True, preserve_range=True, fast_search=True,
+                   bandpass=None):  # pragma: no cover
+    def _bandpass_native_waves(waves):
+        # The thinning algorithm runs on the SED's native tabulation grid, but when a
+        # bandpass target is provided we need to evaluate the throughput at the same physical
+        # observed wavelengths. This helper converts SED-native wavelengths into the native
+        # wavelength units used by bandpass._tp.
+        if self.wave_factor:
+            observed_nm = waves / self.wave_factor
+        else:
+            observed_nm = (waves * self.wave_type).to(units.nm, units.spectral()).value
+        observed_nm *= (1.0 + self.redshift)
+        if bandpass.wave_factor:
+            return observed_nm * bandpass.wave_factor
+        else:
+            return (observed_nm * units.nm).to(bandpass.wave_type, units.spectral()).value
+
+    if bandpass is not None:
+        if self.blue_limit > bandpass.red_limit or self.red_limit < bandpass.blue_limit:
+            raise GalSimError("Bandpass does not overlap the SED wavelength range.")
+        wave_list, _, _ = utilities.combine_wave_list(self, bandpass)
+        if preserve_range and not trim_zeros:
+            # If we want to preserve the range, add back the limits to each end.
+            front = [self.blue_limit] if self.blue_limit < wave_list[0] else []
+            back = [self.red_limit] if self.red_limit > wave_list[-1] else []
+            if front or back:
+                wave_list = np.concatenate((front, wave_list, back))
+    else:
+        wave_list = self.wave_list
+
+    if len(wave_list) > 0:
+        rest_wave_native = self._get_rest_native_waves(wave_list)
+        spec_native = self._spec(rest_wave_native)
+
+        if bandpass is not None:
+            # Identify the overlapping region in the SED's native wavelength units to
+            # determine which portion of the wave_list is within the bandpass limits.
+            band_native_limits = self._get_rest_native_waves(
+                np.array([bandpass.blue_limit, bandpass.red_limit])
+            )
+            native_blue_limit = np.min(band_native_limits)
+            native_red_limit = np.max(band_native_limits)
+            in_band = np.logical_and(
+                rest_wave_native >= native_blue_limit,
+                rest_wave_native <= native_red_limit,
+            )
+
+            # Compute the product SED * bandpass, since that is the quantity whose integral we
+            # want to preserve for this observation.
+            tp_native = np.zeros_like(rest_wave_native, dtype=float)
+            bp_wave_native = _bandpass_native_waves(rest_wave_native[in_band])
+            tp_native[in_band] = bandpass._tp(bp_wave_native) / bandpass.wave_factor
+            spec_native *= tp_native
+
+        # Note that this is thinning in native units, not nm and photons/nm.
+        interpolant = (self.interpolant if not isinstance(self._spec, LookupTable)
+                       else self._spec.interpolant)
+        newx, newf = utilities.thin_tabulated_values(
+                rest_wave_native, spec_native, rel_err=rel_err,
+                trim_zeros=trim_zeros, preserve_range=preserve_range,
+                fast_search=fast_search, interpolant=interpolant)
+
+        if bandpass is not None:
+            # Convert the thinned product back into an SED by dividing out the bandpass
+            # wherever the throughput is non-zero.
+            in_band = np.logical_and(newx >= native_blue_limit, newx <= native_red_limit)
+            tp_native = np.zeros_like(newx, dtype=float)
+            bp_wave_native = _bandpass_native_waves(newx[in_band])
+            tp_native[in_band] = bandpass._tp(bp_wave_native) / bandpass.wave_factor
+            nz = tp_native != 0.  # Don't divide by 0.
+            assert np.all(newf[~nz] == 0.)
+            newf[nz] /= tp_native[nz]
+
+        newspec = _LookupTable(newx, newf, interpolant=interpolant)
+        return SED(newspec, self.wave_type, self.flux_type, redshift=self.redshift,
+                   fast=self.fast)
+    else:
+        return self
+
+
+galsim.SED.thin = _piff_sed_thin

piff/input.py

@@ -47,11 +47,12 @@ def process(cls, config_input, logger=None):
         :param config_input:    The configuration dict.
         :param logger:          A logger object for logging debug info. [default: None]
 
-        :returns: stars, wcs, pointing
+        :returns: stars, wcs, pointing, bandpass
 
         stars is a list of Star instances with the initial data.
         wcs is a dict of WCS solutions indexed by chipnum.
         pointing is either a galsim.CelestialCoord or None.
+        bandpass is either a galsim.Bandpass or None.
         """
         # Get the class to use for handling the input data
         # Default type is 'Files'
@@ -78,7 +79,10 @@ def process(cls, config_input, logger=None):
         # Get the pointing (the coordinate center of the field of view)
         pointing = input_handler.getPointing(logger)
 
-        return stars, wcs, pointing
+        # Get the observing bandpass, if any
+        bandpass = input_handler.getBandpass(logger)
+
+        return stars, wcs, pointing, bandpass
 
     @classmethod
     def __init_subclass__(cls):
@@ -148,6 +152,15 @@ def getPointing(self, logger=None):
         """
         return self.pointing
 
+    def getBandpass(self, logger=None):
+        """Get the observing bandpass for the input images.
+
+        :param logger:      A logger object for logging debug info. [default: None]
+
+        :returns: a galsim.Bandpass or None.
+        """
+        return None
+
 
 class InputFiles(Input):
     """An Input handler that just takes a list of image files and catalog files.
@@ -258,6 +271,10 @@ class InputFiles(Input):
                         SEDs from FITS tables. [default: ``WAVE``]
         :sed_flux_key:  FITS-table column name to use for flux when loading SEDs
                         from FITS tables. [default: ``FLUX``]
+        :bandpass:      GalSim bandpass config dict describing the observing bandpass for the
+                        input image(s). Required when using ``sed_col`` or ``sed_file_name``.
+                        SED thinning will target this bandpass via
+                        ``sed.thin(..., bandpass=...)``. [default: None]
 
     .. note::
 
@@ -435,7 +452,7 @@ def __init__(self, config, logger=None):
                 'noise' : str,
                 'nstars' : int,
               }
-        ignore = [ 'nproc', 'nimages', 'ra', 'dec', 'wcs' ]
+        ignore = [ 'nproc', 'nimages', 'ra', 'dec', 'wcs', 'bandpass' ]
 
         # We're going to change the config dict a bit. Make a copy so we don't mess up the
         # user's original dict (in case they care).
@@ -675,6 +692,12 @@ def __init__(self, config, logger=None):
             sed_flux_type = params.get('sed_flux_type', None)
             sed_wave_key = params.get('sed_wave_key', 'WAVE')
             sed_flux_key = params.get('sed_flux_key', 'FLUX')
+            if (sed_col is not None or sed_file_name is not None) and 'bandpass' not in config:
+                raise ValueError("bandpass is required when using sed_col or sed_file_name")
+            if 'bandpass' in config:
+                bandpass = galsim.config.BuildBandpass(config, 'bandpass', base, logger)[0]
+            else:
+                bandpass = None
             sky_col = params.get('sky_col', None)
             gain_col = params.get('gain_col', None)
             gain = params.get('gain', None)
@@ -709,6 +732,7 @@ def __init__(self, config, logger=None):
                     'sed_flux_type': sed_flux_type,
                     'sed_wave_key': sed_wave_key,
                     'sed_flux_key': sed_flux_key,
+                    'bandpass': bandpass,
                     'sky_col' : sky_col,
                     'gain_col' : gain_col,
                     'sky' : sky,
@@ -719,6 +743,7 @@ def __init__(self, config, logger=None):
                     'stamp_size' : self.stamp_size})
 
         self.use_partial = config.get('use_partial', False)
+        self.bandpass = bandpass
 
         # Read all the wcs's, since we'll need this for the pointing, which in turn we'll
         # need for when we make the stars.
@@ -730,6 +755,9 @@ def __init__(self, config, logger=None):
         self.setPointing(ra, dec, logger)
         self.config = galsim.config.CleanConfig(config)
 
+    def getBandpass(self, logger=None):
+        return self.bandpass
+
     def load_images(self, stars, logger=None):
         """Load the image data into a list of Stars.
 
@@ -1126,7 +1154,7 @@ def _parse_flux_type(flux_type):
 
     @staticmethod
     def _read_sed_file(sed_file_name, sed_wave_type, sed_flux_type,
-                       sed_wave_key, sed_flux_key):
+                       sed_wave_key, sed_flux_key, bandpass):
         cache_key = (sed_file_name, sed_wave_type, sed_flux_type, sed_wave_key, sed_flux_key)
         if cache_key in InputFiles._sed_cache:
             return InputFiles._sed_cache[cache_key]
@@ -1187,6 +1215,7 @@ def readStarCatalog(cat_file_name, cat_hdu, x_col, y_col,
                         ra_col, dec_col, ra_units, dec_units, image,
                         flag_col, skip_flag, use_flag, property_cols, sed_col, sed_wave_type,
                         sed_file_name, sed_flux_type, sed_wave_key, sed_flux_key,
+                        bandpass,
                         properties, image_num, sky_col, gain_col, sky, gain, satur,
                         trust_pos, nstars, stamp_size, config, logger):
         """Read in the star catalogs and return lists of positions for each star in each image.
@@ -1214,6 +1243,7 @@ def readStarCatalog(cat_file_name, cat_hdu, x_col, y_col,
         :param sed_flux_type:   Flux type for SED files.
         :param sed_wave_key:    FITS-table wavelength column name for SED files.
         :param sed_flux_key:    FITS-table flux column name for SED files.
+        :param bandpass:        galsim.Bandpass for SED thinning and unit-flux normalization.
         :param sky_col:         A column with sky (background) levels.
         :param gain_col:        A column with gain values.
         :param sky:             Either a float value for the sky to use for all objects or a str
@@ -1369,6 +1399,7 @@ def safe_to_image(wcs, ra, dec):
                     sed_flux_type=sed_flux_type,
                     sed_wave_key=sed_wave_key,
                     sed_flux_key=sed_flux_key,
+                    bandpass=bandpass,
                 ))
             extra_props['sed'] = np.array(sed_values, dtype=object)
             extra_props['sed_file_name'] = np.array(sed_file_name_values, dtype=object)
@@ -1381,7 +1412,8 @@ def safe_to_image(wcs, ra, dec):
             sed = InputFiles._read_sed_file(
                 sed_file_name, sed_wave_type=sed_wave_type,
                 sed_flux_type=sed_flux_type, sed_wave_key=sed_wave_key,
-                sed_flux_key=sed_flux_key
+                sed_flux_key=sed_flux_key,
+                bandpass=bandpass
             )
             extra_props['sed'] = np.array([sed] * len(cat), dtype=object)
             extra_props['sed_file_name'] = np.array([sed_file_name] * len(cat), dtype=object)

piff/psf.py

@@ -41,6 +41,9 @@ class PSF(object):
     # Normally overridden by subclasses.  But this gives a reasonable default for
     # tests that bypass places where it would be potentially set in normal operations.
     degenerate_points = False
+    wcs = None
+    pointing = None
+    bandpass = None
 
     @classmethod
     def process(cls, config_psf, logger=None):
@@ -54,14 +57,21 @@ def process(cls, config_psf, logger=None):
         several components.
 
         This function merely creates a "blank" PSF object.  It does not actually do any
-        part of the solution yet.  Typically this will be followed by fit:
+        part of the solution yet.  Typically this will be followed by set_context and fit:
 
+            >>> stars, wcs, pointing, bandpass = piff.Input.process(config['input'])
             >>> psf = piff.PSF.process(config['psf'])
-            >>> stars, wcs, pointing = piff.Input.process(config['input'])
-            >>> psf.fit(stars, wcs, pointing)
+            >>> psf.set_context(wcs, pointing, bandpass)
+            >>> psf.fit(stars)
 
         at which point, the ``psf`` instance would have a solution to the PSF model.
 
+        .. note::
+
+            Shared runtime context such as wcs, pointing, and bandpass should be attached
+            with set_context before calling fit.  The old pattern of providing wcs and
+            pointing to fit is still supported, but deprecated.
+
         :param config_psf:  A dict specifying what type of PSF to build along with the
                             appropriate kwargs for building it.
         :param logger:      A logger object for logging debug info. [default: None]
@@ -95,6 +105,24 @@ def process(cls, config_psf, logger=None):
 
         return psf
 
+    def set_context(self, wcs, pointing=None, bandpass=None):
+        """Attach shared runtime context to this PSF instance.
+
+        It is permissible for pointing to be None if the WCS is not Celestial.
+        Also, bandpass may be None if the psf to be fit is not chromatic.
+        If pointing or bandpass is None, the existing value is left unchanged.
+
+        :param wcs:         A dict of WCS solutions indexed by chipnum.
+        :param pointing:    A galsim.CelestialCoord object giving the telescope pointing.
+                            [default: None]
+        :param bandpass:    A galsim.Bandpass giving the observing bandpass. [default: None]
+        """
+        self.wcs = wcs
+        if pointing is not None:
+            self.pointing = pointing
+        if bandpass is not None:
+            self.bandpass = bandpass
+
     def set_num(self, num):
         """If there are multiple components involved in the fit, set the number to use
         for this model.
@@ -373,13 +401,14 @@ def remove_outliers(self, stars, iteration, logger):
             nremoved = 0
         return stars, nremoved
 
-    def fit(self, stars, wcs, pointing, logger=None, convert_funcs=None, draw_method=None):
+    def fit(self, stars, wcs=None, pointing=None, logger=None,
+            convert_funcs=None, draw_method=None):
         """Fit interpolated PSF model to star data using standard sequence of operations.
 
         :param stars:           A list of Star instances.
-        :param wcs:             A dict of WCS solutions indexed by chipnum.
+        :param wcs:             A dict of WCS solutions indexed by chipnum. [deprecated]
         :param pointing:        A galsim.CelestialCoord object giving the telescope pointing.
-                                [Note: pointing should be None if the WCS is not a CelestialWCS]
+                                [deprecated]
         :param logger:          A logger object for logging debug info. [default: None]
         :param convert_funcs:   An optional list of function to apply to the profiles being fit
                                 before drawing it onto the image.  This is used by composite PSFs
@@ -391,8 +420,9 @@ def fit(self, stars, wcs, pointing, logger=None, convert_funcs=None, draw_method
         from .config import LoggerWrapper
         logger = LoggerWrapper(logger)
 
-        self.wcs = wcs
-        self.pointing = pointing
+        if self.wcs is None:
+            logger.error("WARNING: wcs should now be set with set_context before calling fit.")
+            self.set_context(wcs, pointing, self.bandpass)
 
         # Initialize stars as needed by the PSF modeling class.
         stars = self.initialize_flux_center(stars, logger=logger)
@@ -766,9 +796,12 @@ def _write(self, writer, name, logger):
             if hasattr(self, 'stars'):
                 Star.write(self.stars, w, 'stars')
                 logger.verbose("Wrote the PSF stars to name %s", w.get_full_name('stars'))
-            if hasattr(self, 'wcs'):
+            if self.wcs is not None:
                 w.write_wcs_map('wcs', self.wcs, self.pointing)
                 logger.verbose("Wrote the PSF WCS to name %s", w.get_full_name('wcs'))
+            if self.bandpass is not None:
+                w.write_bandpass('bandpass', self.bandpass)
+                logger.verbose("Wrote the PSF bandpass to name %s", w.get_full_name('bandpass'))
             self._finish_write(w, logger=logger)
 
     @classmethod
@@ -822,15 +855,15 @@ def _read(cls, reader, name, logger):
 
         with reader.nested(name) as r:
             # Read the stars, wcs, pointing values
+            wcs, pointing = r.read_wcs_map('wcs', logger=logger)
+            bandpass = r.read_bandpass('bandpass')
             stars = Star.read(r, 'stars')
             if stars is not None:
                 logger.debug("stars = %s", stars)
                 psf.stars = stars
-            wcs, pointing = r.read_wcs_map('wcs', logger=logger)
             if wcs is not None:
-                logger.debug("wcs = %s, pointing = %s",wcs,pointing)
-                psf.wcs = wcs
-                psf.pointing = pointing
+                logger.debug("wcs = %s, pointing = %s, bandpass = %s", wcs, pointing, bandpass)
+                psf.set_context(wcs, pointing, bandpass)
 
             # Just in case the class needs to do something else at the end.
             psf._finish_read(r, logger)