Rubin deep coadd simulation code

config/imsim-config-deep-coadd.yaml

@@ -0,0 +1,112 @@
+modules:
+    - imsim
+
+eval_variables:
+    itract: { type: DeepCoaddData, field: tract }
+    ipatch: { type: DeepCoaddData, field: patch }
+    sband: { type: DeepCoaddData, field: band }
+    incoadds: { type: DeepCoaddData, field: ncoadds }
+    fpupil_R_outer: 4.18
+    fpupil_R_inner: 2.55  # M1 inner diameter is 2.558, but we need a bit of slack for off-axis rays
+    fpupil_area:
+        type: Eval
+        str: "np.pi * (pupil_R_outer**2 - pupil_R_inner**2) * 100**2"
+
+input:
+    sky_catalog:
+        file_name: default_sky_cat.yaml
+        band: $band
+
+    deep_coadds:
+        repo: dp2
+        collection: dp2
+        skymap_name: lsst_cells_v2
+
+    checkpoint:
+        dir: checkpoint
+        file_name:
+          type: FormattedStr
+          format : checkpoint_%s-%d-%d.hdf
+          items:
+              - $band
+              - $tract
+              - $patch
+
+gal:
+    type: SkyCatObj
+
+stamp:
+    type: RubinDeepCoaddStamp
+    max_flux_simple: 100  # When to switch to simple SED.
+    world_pos:
+        type: SkyCatWorldPos
+    band: $band
+    pupil_area: $pupil_area
+    exptime: 30.0  # Exposure time per visit
+
+image:
+    type: RubinDeepCoaddImage
+    random_seed: 42
+    nproc: 1
+    nbatch: 100
+    nsubbatch: 50
+    nbatch_per_checkpoint: 1
+    band: $band
+    tract: $tract
+    patch: $patch
+
+    wcs:
+        type: RubinDeepCoaddWcs
+        band: $band
+        tract: $tract
+        patch: $patch
+
+    bandpass:
+        type: RubinDeepCoaddBandpass
+        band: $band
+
+    noise:
+        type: RubinDeepCoaddGaussianNoise
+        band: $band
+        tract: $tract
+        patch: $patch
+
+psf:
+    type: RubinDeepCoaddPSF
+
+output:
+   type: RubinDeepCoaddOutput
+   coadd_num:
+      type: Sequence
+      nitems: $ncoadds
+      first: 0
+   nproc: 1
+   nfiles: 1
+   dir: fits
+   file_name:
+       type: FormattedStr
+       format: deep_coadd_%s-%d-%d.fits
+       items:
+           - $band
+           - $tract
+           - $patch
+
+   truth:
+       dir: '@output.dir'
+       file_name:
+           type: FormattedStr
+           format: centroid_%s-%d-%d.txt.gz
+           items:
+               - $band
+               - $tract
+               - $patch
+       columns:
+           object_id: "@object_id"
+           ra: $sky_pos.ra.deg
+           dec: $sky_pos.dec.deg
+           x: $image_pos.x
+           y: $image_pos.y
+           nominal_flux: "@simulated_flux"
+           phot_flux: "@phot_flux"
+           realized_flux: "@realized_flux"
+           stamp_size: "@stamp_size"

imsim/__init__.py

@@ -41,3 +41,10 @@
 from .table_row import *
 from .photon_pooling import *
 from .full_focal_plane import *
+from .deep_coadds import *
+from .deep_coadd_image import *
+from .deep_coadd_wcs import *
+from .deep_coadd_output import *
+from .deep_coadd_stamp import *
+from .deep_coadd_bandpass import *
+from .deep_coadd_psf import *

imsim/deep_coadd_bandpass.py

@@ -0,0 +1,27 @@
+import galsim
+from galsim.config import BandpassBuilder, RegisterBandpassType
+
+
+__all__ = ['RubinDeepCoaddBandpassBuilder']
+
+
+class RubinDeepCoaddBandpassBuilder(BandpassBuilder):
+    """Build the Rubin bandpass to be used with deep coadd sims
+    from standard passbands."""
+    def buildBandpass(self, config, base, logger):
+        """
+        Parameters:
+            config:     The configuration dict for the bandpass type.
+            base:       The base configuration dict.
+            logger:     If provided, a logger for logging debug statements.
+        Returns:
+            the constructed Bandpass object.
+        """
+        req = { 'band': str }
+        params, safe = galsim.config.GetAllParams(config, base, req=req)
+        deep_coadds = galsim.config.GetInputObj('deep_coadds', config, base,
+                                                'RubinDeepCoaddBandpass')
+        return deep_coadds.getBandpass(params['band']), safe
+
+
+RegisterBandpassType('RubinDeepCoaddBandpass', RubinDeepCoaddBandpassBuilder())

imsim/deep_coadd_image.py

@@ -0,0 +1,146 @@
+from itertools import pairwise
+import numpy as np
+import galsim
+from galsim.config import (GetAllParams, AddNoise, RegisterImageType,
+                           RegisterNoiseType, GetInputObj, GaussianNoiseBuilder)
+from .lsst_image import LSST_ImageBuilderBase
+
+
+__all__ = ['RubinDeepCoaddImageBuilder']
+
+
+class RubinDeepCoaddImageBuilder(LSST_ImageBuilderBase):  # noqa: N801
+    def setup(self, config, base, image_num, obj_num, ignore, logger):
+        """
+        Do the initialization and setup for a Rubin/LSST deep coadd image.
+
+        Parameters:
+            config:    The configuration dict for the image field.
+            base:      The base configuration dict.
+            image_num: The current image number.
+            obj_num:   The first object number in the image.
+            ignore:    A list of parameters that are allowed to be in config
+                       that we can ignore here.
+            logger:    A logger object to log progress.
+
+        Returns:
+            xsize, ysize
+        """
+        if 'noise' not in base['image']:
+            raise RuntimeError(
+                "A noise type needs to be specified for rendering the image. "
+                "The ``RubinDeepCoaddGaussianNoise`` type models the noise "
+                "level in the original coadd image and could be used here."
+            )
+
+        self.nobjects = self.getNObj(config, base, image_num, logger=logger)
+        if 'nobjects' in config:
+            # User specified nobjects.
+            # Make sure it's not more than what any input catalog can
+            # handle (we don't want repeated objects).
+            input_nobj = galsim.config.ProcessInputNObjects(base)
+            if input_nobj is not None:
+                self.nobjects = min(self.nobjects, input_nobj)
+        logger.info('image %d: nobj = %d', image_num, self.nobjects)
+
+        req = { 'tract': int, 'patch': int, 'band': str }
+        opt = { 'nbatch': int, 'nsubbatch': int, 'nbatch_per_checkpoint': int }
+        extra_ignore = ['image_pos', 'world_pos', 'stamp_size',
+                        'stamp_xsize', 'stamp_ysize', 'nobjects' ]
+        params = GetAllParams(config, base, req=req, opt=opt,
+                              ignore=ignore+extra_ignore)[0]
+
+        self.tract = params['tract']
+        self.patch = params['patch']
+        self.band = params['band']
+
+        self.nbatch = params.get('nbatch', 10)
+        self.nbatch_per_checkpoint = params.get('nbatch_per_checkpoint', 1)
+        self.nsubbatch = params.get('nsubbatch', 50)
+        try:
+            self.checkpoint = GetInputObj(
+                'checkpoint', config, base, 'RubinCoaddImageBuilder')
+        except galsim.config.GalSimConfigError:
+            self.checkpoint = None
+
+        deep_coadds = GetInputObj(
+            'deep_coadds', config, base, 'LSSTC_CoaddImageBuilder')
+        bbox = deep_coadds.skymap[self.tract][self.patch].getOuterBBox()
+        xsize, ysize = bbox.width, bbox.height
+
+        data_id = {'band': self.band,
+                   'tract': self.tract,
+                   'patch': self.patch}
+        self.noise_level = np.std(
+            deep_coadds.get(data_id=data_id).noise_realizations[0].array)
+
+        return xsize, ysize
+
+    def addNoise(self, image, config, base, image_num, obj_num, current_var,
+                 logger):
+        """Add image noise
+
+        Parameters:
+            image:          The image onto which to add the noise.
+            config:         The configuration dict for the image field.
+            base:           The base configuration dict.
+            image_num:      The current image number.
+            obj_num:        The first object number in the image.
+            current_var:    The current noise variance in each postage stamps.
+            logger:         A logger object to log progress.
+        """
+        AddNoise(base, image, current_var, logger)
+
+    def buildImage(self, config, base, image_num, obj_num, logger):
+
+        all_stamps = []
+        all_vars = []
+        full_image = self._create_full_image(config, base)
+        min_batch_size = 10
+        if self.nobjects//self.nbatch < min_batch_size:
+            num_batches = 1
+        else:
+            num_batches = self.nbatch
+        bounds = np.linspace(0, self.nobjects, num_batches + 1, dtype=int)
+        for batch, (start_obj_num, end_obj_num) in enumerate(pairwise(bounds)):
+            nobj_batch = end_obj_num - start_obj_num
+            stamps, current_vars = galsim.config.BuildStamps(
+                nobj_batch, base, logger=logger, obj_num=start_obj_num,
+                do_noise=False)
+            base['index_key'] = 'image_num'
+
+            for k in range(nobj_batch):
+                if stamps[k] is None:
+                    continue
+                bounds = stamps[k].bounds & full_image.bounds
+                if not bounds.isDefined():
+                    continue
+                full_image[bounds] += stamps[k][bounds]
+
+            nz_var = np.nonzero(current_vars)[0]
+            all_stamps.extend([stamps[k] for k in nz_var])
+            all_vars.extend([current_vars[k] for k in nz_var])
+
+        current_var = galsim.config.FlattenNoiseVariance(
+            base, full_image, all_stamps, all_vars, logger)
+        return full_image, current_var
+
+
+class RubinDeepCoaddGaussianNoiseBuilder(GaussianNoiseBuilder):
+    """
+    Treat the variance of the coadd noise realization plane
+    as Gaussian noise for a specified deep_coadd image.
+    """
+
+    def getNoiseVariance(self, config, base, full=False):
+        req = { 'tract': int, 'patch': int, 'band': str }
+        params = GetAllParams(config, base, req=req)[0]
+        deep_coadds = galsim.config.GetInputObj('deep_coadds', config, base,
+                                                'RubinDeepCoaddGaussianNoiseBuilder')
+        return deep_coadds.getNoiseVar(params['tract'], params['patch'],
+                                       params['band'])
+
+
+RegisterImageType('RubinDeepCoaddImage', RubinDeepCoaddImageBuilder())
+RegisterNoiseType('RubinDeepCoaddGaussianNoise',
+                  RubinDeepCoaddGaussianNoiseBuilder())

imsim/deep_coadd_output.py

@@ -0,0 +1,47 @@
+import galsim
+from galsim.config import (OutputBuilder, RegisterOutputType, ParseValue,
+                           GetAllParams, GetInputObj, BuildImage,
+                           eval_base_variables)
+
+
+__all__ = ['RubinDeepCoaddOutputBuilder']
+
+
+class RubinDeepCoaddOutputBuilder(OutputBuilder):
+
+    _added_eval_base_variables = False
+
+    def setup(self, config, base, file_num, logger):
+        coadd_num = ParseValue(config, 'coadd_num', base, int)[0]
+        if not self._added_eval_base_variables:
+            eval_base_variables.append('coadd_num')
+            self._added_eval_base_variables = True
+        base['coadd_num'] = coadd_num
+        deep_coadds = GetInputObj('deep_coadds', config, base,
+                                  'RubinDeepCoaddOutputBuilder')
+        data_id = deep_coadds.data_ids[coadd_num]
+        patch = deep_coadds.skymap[data_id['tract']][data_id['patch']]
+        bbox = patch.getOuterBBox()
+        base['det_xsize'] = bbox.width
+        base['det_ysize'] = bbox.height
+
+    def getNFiles(self, config, base, logger=None):
+        if 'nfiles' in config:
+            return ParseValue(config, 'nfiles', base, int)[0]
+        else:
+            deep_coadds = GetInputObj('deep_coadds', config, base,
+                                      'RubinDeepCoaddOutputBuilder')
+            return len(deep_coadds.data_ids)
+
+    def buildImages(self, config, base, file_num, image_num, obj_num,
+                    ignore, logger):
+        ignore += [ 'file_name', 'dir', 'nfiles', 'coadd_num' ]
+        params, _ = GetAllParams(config, base, ignore=ignore)
+
+        image = BuildImage(base, image_num, obj_num, logger=logger)
+        image.header = galsim.FitsHeader()
+        image.wcs.header = {}
+        return [ image ]
+
+
+RegisterOutputType('RubinDeepCoaddOutput', RubinDeepCoaddOutputBuilder())

imsim/deep_coadd_psf.py

@@ -0,0 +1,33 @@
+import galsim
+from galsim.config import GetInputObj, RegisterObjectType
+
+
+__all__ = ["BuildRubinDeepCoaddPSF"]
+
+
+def BuildRubinDeepCoaddPSF(config, base, ignore, gsparams, logger):
+    """
+    Build PSFs from Rubin deep_coadds.  The deep_coadd image will be
+    retrieved from the data repository using the butler for the
+    tract-patch combination corresponding to the object's sky position.
+
+    Assuming the input.deep_coadd object is defined, to use this
+    PSF, add the following to the config yaml:
+
+    input.atm_psf: ""  # disable the atmospheric PSF
+    psf:
+        type: RubinDeepCoaddPSF
+    """
+    deep_coadds = GetInputObj('deep_coadds', config, base, 'RubinDeepCoaddPSF')
+    coadd_num = base['coadd_num']
+    assert (coadd_num >= 0 and coadd_num < len(deep_coadds.data_ids))
+    data_id = deep_coadds.data_ids[coadd_num]
+    image_pos = base['image_pos']
+    celestial_coord = base['wcs'].toWorld(image_pos)
+    band = data_id['band']
+    safe = False
+    return deep_coadds.getPSF(celestial_coord, band), safe
+
+
+RegisterObjectType('RubinDeepCoaddPSF', BuildRubinDeepCoaddPSF,
+                   input_type='deep_coadds')