Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

DM-30895: Move AssembleCoaddTask and its children to drp_tasks #829

Merged
merged 8 commits into from
Sep 20, 2023
323 changes: 15 additions & 308 deletions python/lsst/pipe/tasks/assembleChi2Coadd.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,311 +19,18 @@
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.


import logging
import numpy as np

from lsst.afw.detection import Psf
import lsst.afw.math as afwMath
import lsst.afw.image as afwImage
import lsst.afw.table as afwTable
from lsst.meas.algorithms import SourceDetectionTask
from lsst.meas.base import SkyMapIdGeneratorConfig
import lsst.pex.config as pexConfig
import lsst.pipe.base as pipeBase
import lsst.pipe.base.connectionTypes as cT


log = logging.getLogger(__name__)


def calculateKernelSize(sigma: float, nSigmaForKernel: float = 7) -> int:
"""Calculate the size of the smoothing kernel.

Parameters
----------
sigma:
Gaussian sigma of smoothing kernel.
nSigmaForKernel:
The multiple of `sigma` to use to set the size of the kernel.
Note that that is the full width of the kernel bounding box
(so a value of 7 means 3.5 sigma on either side of center).
The value will be rounded up to the nearest odd integer.

Returns
-------
size:
Size of the smoothing kernel.
"""
return (int(sigma * nSigmaForKernel + 0.5)//2)*2 + 1 # make sure it is odd


def convolveImage(image: afwImage.Image, psf: Psf) -> afwImage.Image:
"""Convolve an image with a psf

This methodm and the docstring, is based off the method in
`~lsst.meas.algorithms.detection.SourceDetectionTask`.

We convolve the image with a Gaussian approximation to the PSF,
because this is separable and therefore fast. It's technically a
correlation rather than a convolution, but since we use a symmetric
Gaussian there's no difference.

Parameters
----------
image:
The image to convovle.
psf:
The PSF to convolve the `image` with.

Returns
-------
convolved:
The result of convolving `image` with the `psf`.
"""
sigma = psf.computeShape(psf.getAveragePosition()).getDeterminantRadius()
bbox = image.getBBox()

# Smooth using a Gaussian (which is separable, hence fast) of width sigma
# Make a SingleGaussian (separable) kernel with the 'sigma'
kWidth = calculateKernelSize(sigma)
gaussFunc = afwMath.GaussianFunction1D(sigma)
gaussKernel = afwMath.SeparableKernel(kWidth, kWidth, gaussFunc, gaussFunc)

convolvedImage = image.Factory(bbox)

afwMath.convolve(convolvedImage, image, gaussKernel, afwMath.ConvolutionControl())

return convolvedImage.Factory(convolvedImage, bbox, afwImage.PARENT, False)


class AssembleChi2CoaddConnections(pipeBase.PipelineTaskConnections,
dimensions=("tract", "patch", "skymap"),
defaultTemplates={"inputCoaddName": "deep",
"outputCoaddName": "deepChi2"}):
inputCoadds = cT.Input(
doc="Exposure on which to run deblending",
name="{inputCoaddName}Coadd_calexp",
storageClass="ExposureF",
multiple=True,
dimensions=("tract", "patch", "band", "skymap")
)
chi2Coadd = cT.Output(
doc="Chi^2 exposure, produced by merging multiband coadds",
name="{outputCoaddName}Coadd_calexp",
storageClass="ExposureF",
dimensions=("tract", "patch", "skymap"),
)


class AssembleChi2CoaddConfig(pipeBase.PipelineTaskConfig,
pipelineConnections=AssembleChi2CoaddConnections):
outputPixelatedVariance = pexConfig.Field(
dtype=bool,
default=False,
doc="Whether to output a pixelated variance map for the generated "
"chi^2 coadd, or to have a flat variance map defined by combining "
"the inverse variance maps of the coadds that were combined."
)

useUnionForMask = pexConfig.Field(
dtype=bool,
default=True,
doc="Whether to calculate the union of the mask plane in each band, "
"or the intersection of the mask plane in each band."
)


class AssembleChi2CoaddTask(pipeBase.PipelineTask):
"""Assemble a chi^2 coadd from a collection of multi-band coadds

References
----------
.. [1] Szalay, A. S., Connolly, A. J., and Szokoly, G. P., “Simultaneous
Multicolor Detection of Faint Galaxies in the Hubble Deep Field”,
The Astronomical Journal, vol. 117, no. 1, pp. 68–74,
1999. doi:10.1086/300689.

.. [2] Kaiser 2001 whitepaper,
http://pan-starrs.ifa.hawaii.edu/project/people/kaiser/imageprocessing/im%2B%2B.pdf # noqa: E501

.. [3] https://dmtn-015.lsst.io/

.. [4] https://project.lsst.org/meetings/law/sites/lsst.org.meetings.law/files/Building%20and%20using%20coadds.pdf # noqa: E501
"""
ConfigClass = AssembleChi2CoaddConfig
_DefaultName = "assembleChi2Coadd"

def __init__(self, initInputs, **kwargs):
super().__init__(initInputs=initInputs, **kwargs)

def combinedMasks(self, masks: list[afwImage.MaskX]) -> afwImage.MaskX:
"""Combine the mask plane in each input coadd

Parameters
----------
mMask:
The MultibandMask in each band.

Returns
-------
result:
The resulting single band mask.
"""
refMask = masks[0]
bbox = refMask.getBBox()
mask = refMask.array
for _mask in masks[1:]:
if self.config.useUnionForMask:
mask = mask | _mask.array
else:
mask = mask & _mask.array
result = refMask.Factory(bbox)
result.array[:] = mask
return result

def runQuantum(self, butlerQC, inputRefs, outputRefs):
inputs = butlerQC.get(inputRefs)
outputs = self.run(**inputs)
butlerQC.put(outputs, outputRefs)

def run(self, inputCoadds: list[afwImage.Exposure]) -> pipeBase.Struct:
"""Assemble the chi2 coadd from the multiband coadds

Parameters
----------
inputCoadds:
The coadds to combine into a single chi2 coadd.

Returns
-------
result:
The chi2 coadd created from the input coadds.
"""
convControl = afwMath.ConvolutionControl()
convControl.setDoNormalize(False)
convControl.setDoCopyEdge(False)

# Set a reference exposure to use for creating the new coadd.
# It doesn't matter which exposure we use, since we just need the
# bounding box information and Factory to create a new expsure with
# the same dtype.
refExp = inputCoadds[0]
bbox = refExp.getBBox()

image = refExp.image.Factory(bbox)
variance_list = []
# Convovle the image in each band and weight by the median variance
for calexp in inputCoadds:
convolved = convolveImage(calexp.image, calexp.getPsf())
_variance = np.median(calexp.variance.array)
convolved.array[:] /= _variance
image += convolved
variance_list.append(_variance)

variance = refExp.variance.Factory(bbox)
if self.config.outputPixelatedVariance:
# Write the per pixel variance to the output coadd
variance.array[:] = np.sum([1/coadd.variance for coadd in inputCoadds], axis=0)
else:
# Use a flat variance in each band
variance.array[:] = np.sum(1/np.array(variance_list))
# Combine the masks planes to calculate the mask plae of the new coadd
mask = self.combinedMasks([coadd.mask for coadd in inputCoadds])
# Create the exposure
maskedImage = refExp.maskedImage.Factory(image, mask=mask, variance=variance)
chi2coadd = refExp.Factory(maskedImage, exposureInfo=refExp.getInfo())
chi2coadd.info.setFilter(None)
return pipeBase.Struct(chi2Coadd=chi2coadd)


class DetectChi2SourcesConnections(
pipeBase.PipelineTaskConnections,
dimensions=("tract", "patch", "skymap"),
defaultTemplates={
"inputCoaddName": "deepChi2",
"outputCoaddName": "deepChi2"
}
):
detectionSchema = cT.InitOutput(
doc="Schema of the detection catalog",
name="{outputCoaddName}Coadd_det_schema",
storageClass="SourceCatalog",
)
exposure = cT.Input(
doc="Exposure on which detections are to be performed",
name="{inputCoaddName}Coadd_calexp",
storageClass="ExposureF",
dimensions=("tract", "patch", "skymap"),
)
outputSources = cT.Output(
doc="Detected sources catalog",
name="{outputCoaddName}Coadd_det",
storageClass="SourceCatalog",
dimensions=("tract", "patch", "skymap"),
)


class DetectChi2SourcesConfig(pipeBase.PipelineTaskConfig, pipelineConnections=DetectChi2SourcesConnections):
detection = pexConfig.ConfigurableField(
target=SourceDetectionTask,
doc="Detect sources in chi2 coadd"
)

idGenerator = SkyMapIdGeneratorConfig.make_field()

def setDefaults(self):
super().setDefaults()
self.detection.reEstimateBackground = False
self.detection.thresholdValue = 3


class DetectChi2SourcesTask(pipeBase.PipelineTask):
_DefaultName = "detectChi2Sources"
ConfigClass = DetectChi2SourcesConfig

def __init__(self, schema=None, **kwargs):
# N.B. Super is used here to handle the multiple inheritance of PipelineTasks, the init tree
# call structure has been reviewed carefully to be sure super will work as intended.
super().__init__(**kwargs)
if schema is None:
schema = afwTable.SourceTable.makeMinimalSchema()
self.schema = schema
self.makeSubtask("detection", schema=self.schema)
self.detectionSchema = afwTable.SourceCatalog(self.schema)

def runQuantum(self, butlerQC, inputRefs, outputRefs):
inputs = butlerQC.get(inputRefs)
idGenerator = self.config.idGenerator.apply(butlerQC.quantum.dataId)
inputs["idFactory"] = idGenerator.make_table_id_factory()
inputs["expId"] = idGenerator.catalog_id
outputs = self.run(**inputs)
butlerQC.put(outputs, outputRefs)

def run(self, exposure: afwImage.Exposure, idFactory: afwTable.IdFactory, expId: int) -> pipeBase.Struct:
"""Run detection on a chi2 exposure.

Parameters
----------
exposure :
Exposure on which to detect (may be backround-subtracted and scaled,
depending on configuration).
idFactory :
IdFactory to set source identifiers.
expId :
Exposure identifier (integer) for RNG seed.

Returns
-------
result : `lsst.pipe.base.Struct`
Results as a struct with attributes:
``outputSources``
Catalog of detections (`lsst.afw.table.SourceCatalog`).
"""
table = afwTable.SourceTable.make(self.schema, idFactory)
# We override `doSmooth` since the chi2 coadd has already had an
# extra PSF convolution applied to decorrelate the images
# accross bands.
detections = self.detection.run(table, exposure, expId=expId, doSmooth=False)
sources = detections.sources
return pipeBase.Struct(outputSources=sources)
# TODO: Remove this entire file in DM-40826.

import warnings

try:
from lsst.drp.tasks.assemble_chi2_coadd import * # noqa: F401, F403
except ImportError as error:
error.msg += ". Please import the chi^2 coaddition task from drp_tasks package."
raise error
finally:
warnings.warn("lsst.pipe.tasks.assembleChi2Coadd is deprecated and will be removed after v27; "
"Please use lsst.drp.tasks.assemble_chi2_coadd instead.",
DeprecationWarning,
stacklevel=2
)
Loading