Coverage for python/lsst/meas/algorithms/stamps.py : 27%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# This file is part of meas_algorithms.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://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 <https://www.gnu.org/licenses/>.
21#
22"""Collection of small images (stamps).
23"""
25__all__ = ["Stamp", "Stamps", "StampsBase", "writeFits", "readFitsWithOptions"]
27from collections.abc import Sequence
28import abc
29from dataclasses import dataclass
30import numpy
31from typing import Optional
33import lsst.afw.image as afwImage
34import lsst.afw.fits as afwFits
35from lsst.geom import Box2I, Point2I, Extent2I, Angle, degrees, SpherePoint
36from lsst.daf.base import PropertyList
37from lsst.daf.butler.core.utils import getFullTypeName
38from lsst.utils import doImport
41def writeFits(filename, stamp_ims, metadata, type_name, write_mask, write_variance):
42 """Write a single FITS file containing all stamps.
44 Parameters
45 ----------
46 filename : `str`
47 A string indicating the output filename
48 stamps_ims : iterable of `lsst.afw.image.MaskedImageF`
49 An iterable of masked images
50 metadata : `PropertyList`
51 A collection of key, value metadata pairs to be
52 written to the primary header
53 type_name : `str`
54 Python type name of the StampsBase subclass to use
55 write_mask : `bool`
56 Write the mask data to the output file?
57 write_variance : `bool`
58 Write the variance data to the output file?
59 """
60 metadata['HAS_MASK'] = write_mask
61 metadata['HAS_VARIANCE'] = write_variance
62 metadata['N_STAMPS'] = len(stamp_ims)
63 metadata['STAMPCLS'] = type_name
64 # Record version number in case of future code changes
65 metadata['VERSION'] = 1
66 # create primary HDU with global metadata
67 fitsPrimary = afwFits.Fits(filename, "w")
68 fitsPrimary.createEmpty()
69 fitsPrimary.writeMetadata(metadata)
70 fitsPrimary.closeFile()
72 # add all pixel data optionally writing mask and variance information
73 for i, stamp in enumerate(stamp_ims):
74 metadata = PropertyList()
75 # EXTVER should be 1-based, the index from enumerate is 0-based
76 metadata.update({'EXTVER': i+1, 'EXTNAME': 'IMAGE'})
77 stamp.getImage().writeFits(filename, metadata=metadata, mode='a')
78 if write_mask:
79 metadata = PropertyList()
80 metadata.update({'EXTVER': i+1, 'EXTNAME': 'MASK'})
81 stamp.getMask().writeFits(filename, metadata=metadata, mode='a')
82 if write_variance:
83 metadata = PropertyList()
84 metadata.update({'EXTVER': i+1, 'EXTNAME': 'VARIANCE'})
85 stamp.getVariance().writeFits(filename, metadata=metadata, mode='a')
86 return None
89def readFitsWithOptions(filename, stamp_factory, options):
90 """Read stamps from FITS file, allowing for only a
91 subregion of the stamps to be read.
93 Parameters
94 ----------
95 filename : `str`
96 A string indicating the file to read
97 stamp_factory : classmethod
98 A factory function defined on a dataclass for constructing
99 stamp objects a la `lsst.meas.alrogithm.Stamp`
100 options : `PropertyList` or `dict`
101 A collection of parameters. If it contains a bounding box
102 (``bbox`` key), or if certain other keys (``llcX``, ``llcY``,
103 ``width``, ``height``) are available for one to be constructed,
104 the bounding box is passed to the ``FitsReader`` in order to
105 return a sub-image.
107 Returns
108 -------
109 stamps : `list` of dataclass objects like `Stamp`, PropertyList
110 A tuple of a list of `Stamp`-like objects
111 metadata : `PropertyList`
112 The metadata
113 """
114 # extract necessary info from metadata
115 metadata = afwFits.readMetadata(filename, hdu=0)
116 f = afwFits.Fits(filename, 'r')
117 nExtensions = f.countHdus()
118 nStamps = metadata["N_STAMPS"]
119 # check if a bbox was provided
120 kwargs = {}
121 if options:
122 # gen3 API
123 if "bbox" in options.keys():
124 kwargs["bbox"] = options["bbox"]
125 # gen2 API
126 elif "llcX" in options.keys():
127 llcX = options["llcX"]
128 llcY = options["llcY"]
129 width = options["width"]
130 height = options["height"]
131 bbox = Box2I(Point2I(llcX, llcY), Extent2I(width, height))
132 kwargs["bbox"] = bbox
133 stamp_parts = {}
134 # We need to be careful because nExtensions includes the primary
135 # header data unit
136 for idx in range(nExtensions-1):
137 md = afwFits.readMetadata(filename, hdu=idx+1)
138 if md['EXTNAME'] in ('IMAGE', 'VARIANCE'):
139 reader = afwImage.ImageFitsReader(filename, hdu=idx+1)
140 elif md['EXTNAME'] == 'MASK':
141 reader = afwImage.MaskFitsReader(filename, hdu=idx+1)
142 else:
143 raise ValueError(f"Unknown extension type: {md['EXTNAME']}")
144 stamp_parts.setdefault(md['EXTVER'], {})[md['EXTNAME'].lower()] = reader.read(**kwargs)
145 if len(stamp_parts) != nStamps:
146 raise ValueError(f'Number of stamps read ({len(stamp_parts)}) does not agree with the '
147 f'number of stamps recorded in the metadata ({nStamps}).')
148 # construct stamps themselves
149 stamps = []
150 for k in range(nStamps):
151 # Need to increment by one since EXTVER starts at 1
152 maskedImage = afwImage.MaskedImageF(**stamp_parts[k+1])
153 stamps.append(stamp_factory(maskedImage, metadata, k))
155 return stamps, metadata
158@dataclass
159class AbstractStamp(abc.ABC):
160 """Single abstract stamp
162 Parameters
163 ----------
164 Inherit from this class to add metadata to the stamp
165 """
167 @classmethod
168 @abc.abstractmethod
169 def factory(cls, stamp_im, metadata, index):
170 """This method is needed to service the FITS reader.
171 We need a standard interface to construct objects like this.
172 Parameters needed to construct this object are passed in via
173 a metadata dictionary and then passed to the constructor of
174 this class.
176 Parameters
177 ----------
178 stamp : `lsst.afw.image.MaskedImage`
179 Pixel data to pass to the constructor
180 metadata : `dict`
181 Dictionary containing the information
182 needed by the constructor.
183 idx : `int`
184 Index into the lists in ``metadata``
186 Returns
187 -------
188 stamp : `AbstractStamp`
189 An instance of this class
190 """
191 raise NotImplementedError
194@dataclass
195class Stamp(AbstractStamp):
196 """Single stamp
198 Parameters
199 ----------
200 stamp_im : `lsst.afw.image.MaskedImageF`
201 The actual pixel values for the postage stamp
202 position : `lsst.geom.SpherePoint`
203 Position of the center of the stamp. Note the user
204 must keep track of the coordinate system
205 """
206 stamp_im: afwImage.maskedImage.MaskedImageF
207 position: Optional[SpherePoint] = SpherePoint(Angle(numpy.nan), Angle(numpy.nan))
209 @classmethod
210 def factory(cls, stamp_im, metadata, index):
211 """This method is needed to service the FITS reader.
212 We need a standard interface to construct objects like this.
213 Parameters needed to construct this object are passed in via
214 a metadata dictionary and then passed to the constructor of
215 this class. If lists of values are passed with the following
216 keys, they will be passed to the constructor, otherwise dummy
217 values will be passed: RA_DEG, DEC_DEG. They should
218 each point to lists of values.
220 Parameters
221 ----------
222 stamp : `lsst.afw.image.MaskedImage`
223 Pixel data to pass to the constructor
224 metadata : `dict`
225 Dictionary containing the information
226 needed by the constructor.
227 idx : `int`
228 Index into the lists in ``metadata``
230 Returns
231 -------
232 stamp : `Stamp`
233 An instance of this class
234 """
235 if 'RA_DEG' in metadata and 'DEC_DEG' in metadata:
236 return cls(stamp_im=stamp_im,
237 position=SpherePoint(Angle(metadata.getArray('RA_DEG')[index], degrees),
238 Angle(metadata.getArray('DEC_DEG')[index], degrees)))
239 else:
240 return cls(stamp_im=stamp_im, position=SpherePoint(Angle(numpy.nan), Angle(numpy.nan)))
243class StampsBase(abc.ABC, Sequence):
244 """Collection of stamps and associated metadata.
246 Parameters
247 ----------
248 stamps : iterable
249 This should be an iterable of dataclass objects
250 a la ``lsst.meas.algorithms.Stamp``.
251 metadata : `lsst.daf.base.PropertyList`, optional
252 Metadata associated with the bright stars.
253 use_mask : `bool`, optional
254 If ``True`` read and write the mask data. Default ``True``.
255 use_variance : `bool`, optional
256 If ``True`` read and write the variance data. Default ``True``.
258 Notes
259 -----
260 A butler can be used to read only a part of the stamps,
261 specified by a bbox:
263 >>> starSubregions = butler.get("brightStarStamps_sub", dataId, bbox=bbox)
264 """
266 def __init__(self, stamps, metadata=None, use_mask=True, use_variance=True):
267 if not hasattr(stamps, '__iter__'):
268 raise ValueError('The stamps parameter must be iterable.')
269 for stamp in stamps:
270 if not isinstance(stamp, AbstractStamp):
271 raise ValueError('The entries in stamps must inherit from AbstractStamp. '
272 f'Got {type(stamp)}.')
273 self._stamps = stamps
274 self._metadata = PropertyList() if metadata is None else metadata.deepCopy()
275 self.use_mask = use_mask
276 self.use_variance = use_variance
278 @classmethod
279 def readFits(cls, filename):
280 """Build an instance of this class from a file.
282 Parameters
283 ----------
284 filename : `str`
285 Name of the file to read
286 """
288 return cls.readFitsWithOptions(filename, None)
290 @classmethod
291 def readFitsWithOptions(cls, filename, options):
292 """Build an instance of this class with options.
294 Parameters
295 ----------
296 filename : `str`
297 Name of the file to read
298 options : `PropertyList`
299 Collection of metadata parameters
300 """
301 # To avoid problems since this is no longer an abstract method
302 if cls is not StampsBase:
303 raise NotImplementedError(
304 f"Please implement specific FITS reader for class {cls}"
305 )
307 # Load metadata to get class
308 metadata = afwFits.readMetadata(filename, hdu=0)
309 type_name = metadata.get("STAMPCLS")
310 if type_name is None:
311 raise RuntimeError(
312 f"No class name in file {filename}. Unable to instantiate correct"
313 " stamps subclass. Is this an old version format Stamps file?"
314 )
316 # Import class and override `cls`
317 stamp_type = doImport(type_name)
318 cls = stamp_type
320 return cls.readFitsWithOptions(filename, options)
322 @abc.abstractmethod
323 def _refresh_metadata(self):
324 """Make sure metadata is up to date since this object
325 can be extende
326 """
327 raise NotImplementedError
329 def writeFits(self, filename):
330 """Write this object to a file.
332 Parameters
333 ----------
334 filename : `str`
335 Name of file to write
336 """
337 self._refresh_metadata()
338 stamps_ims = self.getMaskedImages()
339 type_name = getFullTypeName(self)
340 writeFits(filename, stamps_ims, self._metadata, type_name, self.use_mask, self.use_variance)
342 def __len__(self):
343 return len(self._stamps)
345 def __getitem__(self, index):
346 return self._stamps[index]
348 def __iter__(self):
349 return iter(self._stamps)
351 def getMaskedImages(self):
352 """Retrieve star images.
354 Returns
355 -------
356 maskedImages :
357 `list` [`lsst.afw.image.maskedImage.maskedImage.MaskedImageF`]
358 """
359 return [stamp.stamp_im for stamp in self._stamps]
361 @property
362 def metadata(self):
363 return self._metadata
366class Stamps(StampsBase):
367 def _refresh_metadata(self):
368 positions = self.getPositions()
369 self._metadata['RA_DEG'] = [p.getRa().asDegrees() for p in positions]
370 self._metadata['DEC_DEG'] = [p.getDec().asDegrees() for p in positions]
372 def getPositions(self):
373 return [s.position for s in self._stamps]
375 def append(self, item):
376 """Add an additional stamp.
378 Parameters
379 ----------
380 item : `Stamp`
381 Stamp object to append.
382 """
383 if not isinstance(item, Stamp):
384 raise ValueError("Objects added must be a Stamp object.")
385 self._stamps.append(item)
386 return None
388 def extend(self, stamp_list):
389 """Extend Stamps instance by appending elements from another instance.
391 Parameters
392 ----------
393 stamps_list : `list` [`Stamp`]
394 List of Stamp object to append.
395 """
396 for s in stamp_list:
397 if not isinstance(s, Stamp):
398 raise ValueError('Can only extend with Stamp objects')
399 self._stamps += stamp_list
401 @classmethod
402 def readFits(cls, filename):
403 """Build an instance of this class from a file.
405 Parameters
406 ----------
407 filename : `str`
408 Name of the file to read
410 Returns
411 -------
412 object : `Stamps`
413 An instance of this class
414 """
415 return cls.readFitsWithOptions(filename, None)
417 @classmethod
418 def readFitsWithOptions(cls, filename, options):
419 """Build an instance of this class with options.
421 Parameters
422 ----------
423 filename : `str`
424 Name of the file to read
425 options : `PropertyList` or `dict`
426 Collection of metadata parameters
428 Returns
429 -------
430 object : `Stamps`
431 An instance of this class
432 """
433 stamps, metadata = readFitsWithOptions(filename, Stamp.factory, options)
434 return cls(stamps, metadata=metadata, use_mask=metadata['HAS_MASK'],
435 use_variance=metadata['HAS_VARIANCE'])