Coverage for python/lsst/meas/transiNet/rbTransiNetTask.py: 99%
64 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 09:18 +0000
1# This file is part of meas_transiNet.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
22__all__ = ["RBTransiNetTask", "RBTransiNetConfig"]
24import lsst.geom
25from lsst.images import Box
26import lsst.pex.config
27import lsst.pipe.base
28from lsst.utils.timer import timeMethod
29import numpy as np
31from . import rbTransiNetInterface
32from lsst.meas.transiNet.modelPackages.storageAdapterButler import StorageAdapterButler
35class RBTransiNetConnections(lsst.pipe.base.PipelineTaskConnections,
36 dimensions=("instrument", "visit", "detector"),
37 defaultTemplates={"coaddName": "deep", "fakesType": ""}):
38 # The butler converts an afw image component to ``ImageV2``.
39 # NOTE: Do we want the "ready to difference" template, or something
40 # earlier? This one is warped, but not PSF-matched.
41 template = lsst.pipe.base.connectionTypes.Input(
42 doc="Image component of the input warped template to subtract.",
43 dimensions=("instrument", "visit", "detector"),
44 storageClass="ImageV2",
45 name="{fakesType}{coaddName}Diff_templateExp.image"
46 )
47 science = lsst.pipe.base.connectionTypes.Input(
48 doc="Image component of the input science exposure to subtract from.",
49 dimensions=("instrument", "visit", "detector"),
50 storageClass="ImageV2",
51 name="{fakesType}calexp.image"
52 )
53 difference = lsst.pipe.base.connectionTypes.Input(
54 doc="Image component of the result of subtracting the convolved template"
55 " from the science image.",
56 dimensions=("instrument", "visit", "detector"),
57 storageClass="ImageV2",
58 name="{fakesType}{coaddName}Diff_differenceExp.image",
59 )
60 diaSources = lsst.pipe.base.connectionTypes.Input(
61 doc="Detected sources on the difference image.",
62 dimensions=("instrument", "visit", "detector"),
63 storageClass="SourceCatalog",
64 name="{fakesType}{coaddName}Diff_candidateDiaSrc",
65 )
66 pretrainedModel = lsst.pipe.base.connectionTypes.PrerequisiteInput(
67 doc="Pretrained neural network model (-package) for the RBClassifier.",
68 dimensions=(),
69 storageClass="NNModelPackagePayload",
70 name=StorageAdapterButler.dataset_type_name,
71 )
73 # Outputs
74 classifications = lsst.pipe.base.connectionTypes.Output(
75 doc="Catalog of real/bogus classifications for each diaSource, "
76 "element-wise aligned with diaSources.",
77 dimensions=("instrument", "visit", "detector"),
78 storageClass="Catalog",
79 name="{fakesType}{coaddName}RealBogusSources",
80 )
82 def __init__(self, *, config=None):
83 super().__init__(config=config)
85 if self.config.modelPackageStorageMode != "butler": 85 ↛ exitline 85 didn't return from function '__init__' because the condition on line 85 was always true
86 del self.pretrainedModel
89class RBTransiNetConfig(lsst.pipe.base.PipelineTaskConfig, pipelineConnections=RBTransiNetConnections):
90 modelPackageName = lsst.pex.config.Field(
91 optional=True,
92 dtype=str,
93 doc=("A unique identifier of a model package. ")
94 )
95 modelPackageStorageMode = lsst.pex.config.ChoiceField(
96 dtype=str,
97 doc=("A string that indicates _where_ and _how_ the model package is stored."),
98 allowed={'local': 'packages stored in the meas_transiNet repository',
99 'neighbor': 'packages stored in the rbClassifier_data repository',
100 'butler': 'packages stored in the butler repository',
101 },
102 default='neighbor',
103 )
104 cutoutSize = lsst.pex.config.Field(
105 dtype=int,
106 doc="Width/height of square cutouts to send to classifier.",
107 default=51,
108 )
110 def validate(self):
111 # if we are in the butler mode, the user should not set
112 # a modelPackageName as a config field.
113 if self.modelPackageStorageMode == "butler":
114 if self.modelPackageName is not None:
115 raise ValueError("In a _real_ run of a pipeline when the "
116 "modelPackageStorageMode is 'butler', "
117 "the modelPackageName cannot be specified "
118 "as a config field. Pass it as a collection"
119 "name in the command-line instead.")
122class RBTransiNetTask(lsst.pipe.base.PipelineTask):
123 """Task for running TransiNet real/bogus classification on the output of
124 the image subtraction pipeline.
125 """
126 _DefaultName = "rbTransiNet"
127 ConfigClass = RBTransiNetConfig
129 def __init__(self, **kwargs):
130 super().__init__(**kwargs)
132 self.butler_loaded_package = None
134 @timeMethod
135 def run(self, template, science, difference, diaSources, pretrainedModel=None):
137 # Create the TransiNet interface object.
138 # Note: assuming each quanta creates one instance of this task, this is
139 # a proper place for doing this since loading of the model is run only
140 # once. However, if in the future we come up with a design in which one
141 # task instance is used for multiple quanta, this will need to be moved
142 # somewhere else -- e.g. to the __init__ method, or even to runQuantum.
143 self.butler_loaded_package = pretrainedModel # This will be used by the interface
144 self.interface = rbTransiNetInterface.RBTransiNetInterface(self)
146 cutouts = [self._make_cutouts(template, science, difference, source) for source in diaSources]
147 self.log.info("Extracted %d cutouts.", len(cutouts))
148 scores = self.interface.infer(cutouts)
149 self.log.info("Scored %d cutouts.", len(scores))
150 schema = lsst.afw.table.Schema()
151 schema.addField(diaSources.schema["id"].asField())
152 schema.addField("score", doc="real/bogus score of this source", type=np.float32)
153 schema.addField("version", doc="Version of the reliability model", type=str, size=7)
154 classifications = lsst.afw.table.BaseCatalog(schema)
155 classifications.resize(len(scores))
157 classifications["id"] = diaSources["id"]
158 classifications["score"] = scores
159 for record in classifications:
160 record.set("version", self.interface.version)
162 return lsst.pipe.base.Struct(classifications=classifications)
164 def _make_cutouts(self, template, science, difference, source):
165 """Return cutouts of each image centered at the source location.
167 Parameters
168 ----------
169 template : `lsst.images.Image`
170 science : `lsst.images.Image`
171 difference : `lsst.images.Image`
172 Images to cut out of.
173 source : `lsst.afw.table.SourceRecord`
174 Source to make cutouts of.
176 Returns
177 -------
178 cutouts, `lsst.meas.transiNet.CutoutInputs`
179 Cutouts of each of the input images.
180 """
182 # Try to create cutouts, or simply return empty cutouts if
183 # failed (most probably out-of-border box)
184 extent = lsst.geom.Extent2I(self.config.cutoutSize)
185 box = Box.from_legacy(lsst.geom.Box2I.makeCenteredBox(source.getCentroid(), extent))
187 if science.bbox.contains(box):
188 science_cutout = np.nan_to_num(science[box].array)
189 template_cutout = np.nan_to_num(template[box].array)
190 difference_cutout = np.nan_to_num(difference[box].array)
191 else:
192 science_cutout = np.zeros((self.config.cutoutSize, self.config.cutoutSize), dtype=np.float32)
193 template_cutout = np.zeros_like(science_cutout)
194 difference_cutout = np.zeros_like(science_cutout)
196 return rbTransiNetInterface.CutoutInputs(science=science_cutout,
197 template=template_cutout,
198 difference=difference_cutout)