diff --git a/config/imsim-config-deep-coadd.yaml b/config/imsim-config-deep-coadd.yaml new file mode 100644 index 00000000..b786f970 --- /dev/null +++ b/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" diff --git a/imsim/__init__.py b/imsim/__init__.py index 69965189..11c4ef47 100644 --- a/imsim/__init__.py +++ b/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 * diff --git a/imsim/deep_coadd_bandpass.py b/imsim/deep_coadd_bandpass.py new file mode 100644 index 00000000..00e9f131 --- /dev/null +++ b/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()) diff --git a/imsim/deep_coadd_image.py b/imsim/deep_coadd_image.py new file mode 100644 index 00000000..d7f7e8cf --- /dev/null +++ b/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()) diff --git a/imsim/deep_coadd_output.py b/imsim/deep_coadd_output.py new file mode 100644 index 00000000..7741ce2b --- /dev/null +++ b/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()) diff --git a/imsim/deep_coadd_psf.py b/imsim/deep_coadd_psf.py new file mode 100644 index 00000000..b5df6cdf --- /dev/null +++ b/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') diff --git a/imsim/deep_coadd_stamp.py b/imsim/deep_coadd_stamp.py new file mode 100644 index 00000000..5c300bbe --- /dev/null +++ b/imsim/deep_coadd_stamp.py @@ -0,0 +1,186 @@ +import astropy.units as u +import galsim +from galsim.config import StampBuilder, RegisterStampType +from .stamp import LSST_SiliconBuilder +from .stamp_utils import get_stamp_size + + +__all__ = ["RubinDeepCoaddStampBuilder"] + + +class DeepCoaddWeighting(galsim.PhotonOp): + """ + Weight photons so that final object fluxes are in nJy. + """ + def __init__(self, mag, nphotons): + """ + Parameters + ---------- + mag: float + The bandpass magnitude for the object being rendered. + nphotons: float + The number of photons to simulate. This should be + the total number over all visits contributing to the + object's counts. Nominally, this is given by + nphotons = photon_flux * num_visits * exptime_per_visit * pupil_area + """ + # Convert magnitude to nJy: + fnu = 1e9 * (mag*u.ABmag).to(u.Jy).to_value() + self.weight = fnu / nphotons + + def applyTo(self, photon_array, local_wcs=None, rng=None): + photon_array.flux *= self.weight + + +class RubinDeepCoaddStampBuilder(LSST_SiliconBuilder): + + def setup(self, config, base, xsize, ysize, ignore, logger): + # Always disable fft rendering: + self.use_fft = False + + # Use StampBuilder.setup(...) to set the object position. + ignore = ignore + ['max_flux_simple', 'band', 'pupil_area', 'exptime'] + _, _, image_pos, world_pos = StampBuilder.setup( + self, config, base, xsize, ysize, ignore, logger) + + obj = galsim.config.BuildGSObject(base, 'gal', logger=logger)[0] + if obj is None: + raise galsim.config.SkipThisObject('obj is None') + self.obj = obj + + self.rng = galsim.config.GetRNG(config, base, logger, "RubinDeepCoadd") + self.image = base['current_image'] + + deep_coadds = galsim.config.GetInputObj('deep_coadds', config, base, + 'RubinDeepCoaddStampBuilder') + band = galsim.config.ParseValue(config, 'band', base, str)[0] + num_visits = deep_coadds.getNumVisits(world_pos, band) + + bandpass = base['bandpass'] + if not hasattr(obj, 'flux'): + obj.flux = obj.calculateFlux(bandpass) + + # To get the AB magnitude, we need to divide out the pupil + # area and single visit exposure time from the object SED as + # provided from skyCatalogs. + pupil_area = galsim.config.ParseValue(config, 'pupil_area', base, float)[0] + exptime = galsim.config.ParseValue(config, 'exptime', base, float)[0] + sed = obj.sed / pupil_area / exptime + mag = sed.calculateMagnitude(bandpass) + + # Scale the single-visit object flux by the number of visits + # in order to get the correct signal variance. The deep coadd + # weighting ensures that the final flux is in units of nJy. + self.simulated_flux = obj.flux * num_visits + self.deep_coadd_phot_op = DeepCoaddWeighting(mag, self.simulated_flux) + self.phot_flux = galsim.PoissonDeviate(self.rng, mean=self.simulated_flux)() + + # Save values for output truth catalog + base['simulated_flux'] = self.simulated_flux + base['phot_flux'] = self.phot_flux + base['realized_flux'] = 0 # This will be updated by .drawImage + + if self.phot_flux == 0: + raise galsim.config.SkipThisObject('phot_flux==0') + + if xsize > 0 and ysize > 0: + # Use the stamp size from the config. + pass + elif 'size' in config: + # Get the stamp size from the size config entry + xsize = ysize = galsim.config.ParseValue( + config, 'size', base, int)[0] + elif self.simulated_flux < self._tiny_flux: + xsize = ysize = 32 + else: + # Determine the stamp size from the object flux. + noise_var = galsim.config.CalculateNoiseVariance(base) + obj_achrom = obj.evaluateAtWavelength(bandpass.effective_wavelength) + # We're dividing the nominal noise variance by 3.0 when + # passing the value to get_stamp_size in order to suppress + # apparent stamp edges for high aspect ratio galaxies. + stamp_size = get_stamp_size( + obj_achrom=obj_achrom, + nominal_flux=self.simulated_flux, + noise_var=noise_var/3.0, + Nmax=self._Nmax, + pixel_scale=self._pixel_scale, + logger=logger + ) + xsize = ysize = stamp_size + + base['stamp_size'] = xsize + logger.info('Object %d will use stamp size %s, %s and simulated flux %s', + base.get('obj_num',0), xsize, ysize, self.simulated_flux) + return xsize, ysize, image_pos, world_pos + + def buildPSF(self, config, base, gsparams, logger): + psf = galsim.config.BuildGSObject( + base, 'psf', gsparams=gsparams, logger=logger)[0] + return psf + + def draw(self, prof, image, method, offset, config, base, logger): + if prof is None: + image.photons = galsim.PhotonArray(0) + return image + + obj_num = base.get('obj_num',0) + bandpass = base['bandpass'] + + max_flux_simple = config.get('max_flux_simple', 100) + faint = self.simulated_flux < max_flux_simple + + if faint: + logger.info("Flux = %.0f Using trivial sed", self.obj.flux) + for profile_wl in (bandpass.effective_wavelength, + bandpass.red_limit, + bandpass.blue_limit): + sed_value = prof.sed(profile_wl) + if sed_value != 0: + break + if sed_value == 0: + # We can't evaluate the profile for this object, so skip it. + obj_num = base.get('obj_num') + object_id = base.get('object_id') + logger.warning("Zero-valued SED for faint object %d, " + "object_id %s. Skipping.", obj_num, object_id) + return image + prof = prof.evaluateAtWavelength(profile_wl) + prof = prof * self._trivial_sed + else: + self._fix_seds(prof, bandpass, logger) + + image.wcs = base['wcs'] + + photon_ops = [self.deep_coadd_phot_op] + + # Set limit on the size of photons batches to consider when + # calling gsobject.drawImage. + maxN = int(1e6) + if 'maxN' in config: + maxN = galsim.config.ParseValue(config, 'maxN', base, int)[0] + + prof = prof.withFlux(self.phot_flux, bandpass) + if not faint and 'photon_ops' in config: + photon_ops = photon_ops + galsim.config.BuildPhotonOps( + config, 'photon_ops', base, logger) + + image = prof.drawImage(bandpass, + method='phot', + offset=offset, + rng=self.rng, + maxN=maxN, + n_photons=self.phot_flux, + image=image, + sensor=None, + photon_ops=photon_ops, + add_to_image=True, + poisson_flux=False) + base['realized_flux'] = image.added_flux + logger.debug('After .drawImage(...), image.added_flux %s', + image.added_flux) + + return image + + +RegisterStampType('RubinDeepCoaddStamp', RubinDeepCoaddStampBuilder()) diff --git a/imsim/deep_coadd_wcs.py b/imsim/deep_coadd_wcs.py new file mode 100644 index 00000000..a12bbf4e --- /dev/null +++ b/imsim/deep_coadd_wcs.py @@ -0,0 +1,35 @@ +""" +WCS builder for Rubin deep_coadds +""" +from galsim.config import WCSBuilder, GetAllParams, GetInputObj, RegisterWCSType + + +__all__ = ["RubinDeepCoaddWCSBuilder"] + + +class RubinDeepCoaddWCSBuilder(WCSBuilder): + + def buildWCS(self, config, base, logger): + """Build an AstropyWCS from a deep_coadd + + Parameters: + config: The configuration dict for the wcs type. + base: The base configuration dict. + logger: If provided, a logger for logging debug statements. + + Returns: + The constructed WCS object (a galsim.AstropyWCS instance). + """ + req = { + "tract": int, + "patch": int, + "band": str, + } + params, _ = GetAllParams(config, base, req=req) + data_id = { key: params[key] for key in req } + deep_coadds = GetInputObj("deep_coadds", config, base, + "RubinDeepCoaddWCSBuilder") + return deep_coadds.getWcs(data_id) + + +RegisterWCSType("RubinDeepCoaddWcs", RubinDeepCoaddWCSBuilder()) diff --git a/imsim/deep_coadds.py b/imsim/deep_coadds.py new file mode 100644 index 00000000..5063b5dd --- /dev/null +++ b/imsim/deep_coadds.py @@ -0,0 +1,239 @@ +""" +Interface to Rubin deep_coadds +""" +from collections import namedtuple +import numpy as np +import pandas as pd +import galsim +from galsim.config import (InputLoader, RegisterInputType, RegisterValueType, + GetAllParams, GetInputObj) +import lsst.daf.butler as daf_butler +import lsst.geom + + +__all__ = ["DeepCoadds", "DeepCoaddLoader", "DeepCoaddData"] + + +GridKey = namedtuple('GridKey', ['tract', 'patch', 'band']) + + +class DeepCoadds: + """ + This class serves as the access point for deep_coadd information + obtained via butler queries. It caches PSF, wcs, visit, and noise + information in order to minimize the number of butler.get calls. + """ + def __init__(self, butler, skymap_name, data_ids=None, dstype="deep_coadd"): + """ + Parameters + ---------- + butler : lsst.daf.butler.Butler + Butler for the data repo and collection containing the deep_coadds. + skymap_name : str + Name of the skymap to use, e.g., "lsst_cells_v2". + data_ids: list + List of data_ids (dicts) for identifying deep_coadds to simulate. + dstype : str + The dataset type of the cell-based coadds. Default: "deep_coadd". + """ + self.butler = butler + self.skymap_name = skymap_name + self.skymap = butler.get("skyMap", skymap=skymap_name) + self.data_ids = data_ids + self.dstype = dstype + self._psf_cache = {} + self._grid_cache = {} + self._wcs_cache = {} + self._num_visits = {} + self._noise_var = {} + + def get(self, index=None, data_id=None): + """Return the deep_coadd using the index of the self.data_ids list, + or if provided, using the data_id. + """ + if (data_id, index) == (None, None): + return None + if data_id is None and index >= 0 and index < len(self.data_ids): + data_id = self.data_ids[index] + elif data_id not in self.data_ids: + return None + my_data_id = data_id.copy() + my_data_id['skymap'] = self.skymap_name + deep_coadd = self.butler.get(self.dstype, **my_data_id) + + # Fill in various cached quantities to avoid repeated butler + # gets. + wcs_key = data_id['tract'], data_id['patch'] + if wcs_key not in self._wcs_cache: + self._wcs_cache[wcs_key] = galsim.AstropyWCS(wcs=deep_coadd.fits_wcs) + + grid_key = GridKey(data_id['tract'], data_id['patch'], data_id['band']) + self._num_visits[grid_key] = self.num_visits_per_cell(deep_coadd) + self._grid_cache[grid_key] = deep_coadd.grid, deep_coadd.psf + self._noise_var[grid_key] \ + = np.var(deep_coadd.noise_realizations[0].array) + + return deep_coadd + + def getBandpass(self, band): + passband = self.butler.get('standard_passband', band=band) + lut = galsim.LookupTable(passband['wavelength'], + passband['throughput'], + interpolant='linear') + bp = galsim.Bandpass(lut, wave_type='nm').withZeropoint('AB').thin() + return bp + + def getWcs(self, data_id): + key = data_id['tract'], data_id['patch'] + if key not in self._wcs_cache: + self.get(data_id=data_id) # This will fill the cache. + return self._wcs_cache[key] + + def getPSF(self, celestial_coord, band): + """Return the cell coadd PSF, evaluated at the center of + the cell containing this sky position. + """ + ra = celestial_coord.ra / galsim.degrees + dec = celestial_coord.dec / galsim.degrees + grid_key, cell_index, x, y = self._get_cache_keys(ra, dec, band) + data_id = dict(tract=grid_key.tract, patch=grid_key.patch, + band=grid_key.band) + grid, psf = self._grid_cache[grid_key] + + # xy offsets for the current grid. + x_offset = grid.bbox.start.x + y_offset = grid.bbox.start.y + + # Access the cached cell PSFs. + psf_key = grid_key, cell_index + if psf_key not in self._psf_cache: + # Evaluate PSF at cell center. This is based on + # https://github.com/lsst/source_injection/blob/w.2026.38/python/lsst/source/injection/inject_engine.py#L705 + x0 = ( (x // grid.cell_shape.x) * grid.cell_shape.x + + grid.cell_shape.x/2 + x_offset ) + y0 = ( (y // grid.cell_shape.y) * grid.cell_shape.y + + grid.cell_shape.y/2 + y_offset ) + psf_array = psf.compute_kernel_image(x=x0, y=y0).array + wcs = self.getWcs(data_id) + self._psf_cache[psf_key] = galsim.InterpolatedImage( + galsim.Image(psf_array), wcs=wcs) + return self._psf_cache[psf_key] + + def getNumVisits(self, celestial_coord, band): + ra = celestial_coord.ra / galsim.degrees + dec = celestial_coord.dec / galsim.degrees + grid_key, cell_index, _, _ = self._get_cache_keys(ra, dec, band) + i = int(cell_index.i) + j = int(cell_index.j) + return self._num_visits[grid_key][(i, j)] + + def getNoiseVar(self, tract, patch, band): + grid_key = GridKey(tract, patch, band) + return self._noise_var[grid_key] + + def _get_cache_keys(self, ra, dec, band): + # Find the tract, patch for this location. + sky_coords = lsst.geom.SpherePoint( + lsst.geom.Angle(ra*lsst.geom.degrees), + lsst.geom.Angle(dec*lsst.geom.degrees) + ) + tract_info = self.skymap.findTract(sky_coords) + tract = tract_info.getId() + patch = tract_info.findPatch(sky_coords).getSequentialIndex() + grid_key = GridKey(tract, patch, band) + data_id = dict(tract=grid_key.tract, patch=grid_key.patch, + band=grid_key.band) + if grid_key not in self._grid_cache: + self.get(data_id=data_id) + grid, _ = self._grid_cache[grid_key] + + # xy offsets for the current grid. + x_offset = grid.bbox.start.x + y_offset = grid.bbox.start.y + + # Get cell_index for the requested location. + wcs = self.getWcs(data_id) + x, y = wcs.toImage(ra, dec, units='degrees') + cell_index = grid.index_of(x=x + x_offset, y=y + y_offset) + return grid_key, cell_index, x, y + + @staticmethod + def num_visits_per_cell(deep_coadd): + df0 = deep_coadd.provenance.contributions.to_pandas() + num_visits = {} + for i, j in sorted(set(zip(df0['cell_i'], df0['cell_j']))): + df = df0.query(f"cell_i == {i} and cell_j == {j}") + num_visits[(i, j)] = sum(df['overlap_fraction']) + return num_visits + + +class DeepCoaddLoader(InputLoader): + """ + Load the deep_coadds input object. Here's an example yaml entry: + + input.deep_coadds: + repo: dp2 + collection: dp2 + """ + def __init__(self): + super().__init__(init_func=DeepCoadds, takes_logger=True, + use_proxy=False) + self.butler = None + self.deep_coadd_list = None + + def getKwargs(self, config, base, logger): + logger.debug("Get kwargs for DeepCoadds") + req = { + "repo": str, + "collection": str, + } + opt = { + "skymap_name": str, + "dstype": str, + "deep_coadd_list_file": str, + } + params, _ = GetAllParams(config, base, req=req, opt=opt) + if self.butler is None: + self.butler = daf_butler.Butler(params["repo"], + collections=[params["collection"]]) + + deep_coadd_list_file = params.get("deep_coadd_list_file", None) + + if self.deep_coadd_list is None and deep_coadd_list_file is not None: + df = pd.read_parquet(deep_coadd_list_file) + columns = ["band", "tract", "patch"] + self.deep_coadd_list = [dict(zip(columns, row)) + for row in zip(*[df[_] for _ in columns])] + + kwargs = { + "butler": self.butler, + "skymap_name": params.get("skymap", "lsst_cells_v2"), + "data_ids": self.deep_coadd_list, + "dstype": params.get("dstype", "deep_coadd"), + } + safe = True + return kwargs, safe + + +def DeepCoaddData(config, base, value_type): + deep_coadds = GetInputObj('deep_coadds', config, base, 'DeepCoaddData') + num_coadds = len(deep_coadds.data_ids) + + req = { 'field': str } + params, safe = GetAllParams(config, base, req=req) + field = params['field'] + if field == 'ncoadds': + return num_coadds, safe + + coadd_num = base['coadd_num'] + assert (coadd_num >= 0 and coadd_num < num_coadds) + data_id = deep_coadds.data_ids[coadd_num] + + val = value_type(data_id.get(field, None)) + + return val, safe + + +RegisterInputType('deep_coadds', DeepCoaddLoader()) +RegisterValueType('DeepCoaddData', DeepCoaddData, [int, str], + input_type='deep_coadds') diff --git a/imsim/skycat.py b/imsim/skycat.py index dbc2931c..670a8a2f 100644 --- a/imsim/skycat.py +++ b/imsim/skycat.py @@ -15,7 +15,7 @@ class SkyCatalogInterface: """Interface to skyCatalogs package.""" - def __init__(self, file_name, wcs, band, mjd, fiducial_wcs=None, + def __init__(self, file_name, wcs, band, mjd=None, fiducial_wcs=None, fiducial_det_name=None, xsize=4096, ysize=4096, obj_types=None, skycatalog_root=None, edge_pix=200, camera=None, det_name=None, diff --git a/imsim/templates.py b/imsim/templates.py index 9cec8731..71d865fb 100644 --- a/imsim/templates.py +++ b/imsim/templates.py @@ -15,3 +15,4 @@ galsim.config.RegisterTemplate('imsim-config-photon-pooling', os.path.join(config_dir, 'imsim-config-photon-pooling.yaml')) galsim.config.RegisterTemplate('imsim-config-instcat-comcam', os.path.join(config_dir, 'imsim-config-instcat-comcam.yaml')) galsim.config.RegisterTemplate('imsim-config-skycat-comcam', os.path.join(config_dir, 'imsim-config-skycat-comcam.yaml')) +galsim.config.RegisterTemplate('imsim-config-deep-coadd', os.path.join(config_dir, 'imsim-config-deep-coadd.yaml'))