lsst.meas.algorithms g41a93ceb71+a3dc6b54d0
Loading...
Searching...
No Matches
stamps.py
Go to the documentation of this file.
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
24__all__ = ["Stamp", "Stamps", "StampsBase", "writeFits", "readFitsWithOptions"]
25
26import abc
27from collections.abc import Sequence
28from dataclasses import dataclass
29
30import numpy as np
31from lsst.afw.fits import Fits, readMetadata
32from lsst.afw.image import ImageFitsReader, MaskedImageF, MaskFitsReader
33from lsst.afw.table.io import InputArchive, OutputArchive, Persistable
34from lsst.daf.base import PropertyList
35from lsst.geom import Angle, Box2I, Extent2I, Point2I, SpherePoint, degrees
36from lsst.utils import doImport
37from lsst.utils.introspection import get_full_type_name
38
39
40def writeFits(filename, stamps, metadata, type_name, write_mask, write_variance, write_archive=False):
41 """Write a single FITS file containing all stamps.
42
43 Parameters
44 ----------
45 filename : `str`
46 A string indicating the output filename
47 stamps : iterable of `BaseStamp`
48 An iterable of Stamp objects
49 metadata : `PropertyList`
50 A collection of key, value metadata pairs to be
51 written to the primary header
52 type_name : `str`
53 Python type name of the StampsBase subclass to use
54 write_mask : `bool`
55 Write the mask data to the output file?
56 write_variance : `bool`
57 Write the variance data to the output file?
58 write_archive : `bool`, optional
59 Write an archive to store Persistables along with each stamp?
60 Default: ``False``.
61 """
62 metadata["HAS_MASK"] = write_mask
63 metadata["HAS_VARIANCE"] = write_variance
64 metadata["HAS_ARCHIVE"] = write_archive
65 metadata["N_STAMPS"] = len(stamps)
66 metadata["STAMPCLS"] = type_name
67 # Record version number in case of future code changes
68 metadata["VERSION"] = 1
69 # create primary HDU with global metadata
70 fitsFile = Fits(filename, "w")
71 fitsFile.createEmpty()
72 # Store Persistables in an OutputArchive and write it
73 if write_archive:
74 oa = OutputArchive()
75 archive_ids = [oa.put(stamp.archive_element) for stamp in stamps]
76 metadata["ARCHIVE_IDS"] = archive_ids
77 fitsFile.writeMetadata(metadata)
78 oa.writeFits(fitsFile)
79 else:
80 fitsFile.writeMetadata(metadata)
81 fitsFile.closeFile()
82 # add all pixel data optionally writing mask and variance information
83 for i, stamp in enumerate(stamps):
84 metadata = PropertyList()
85 # EXTVER should be 1-based, the index from enumerate is 0-based
86 metadata.update({"EXTVER": i + 1, "EXTNAME": "IMAGE"})
87 stamp.stamp_im.getImage().writeFits(filename, metadata=metadata, mode="a")
88 if write_mask:
89 metadata = PropertyList()
90 metadata.update({"EXTVER": i + 1, "EXTNAME": "MASK"})
91 stamp.stamp_im.getMask().writeFits(filename, metadata=metadata, mode="a")
92 if write_variance:
93 metadata = PropertyList()
94 metadata.update({"EXTVER": i + 1, "EXTNAME": "VARIANCE"})
95 stamp.stamp_im.getVariance().writeFits(filename, metadata=metadata, mode="a")
96 return None
97
98
99def readFitsWithOptions(filename, stamp_factory, options):
100 """Read stamps from FITS file, allowing for only a subregion of the stamps
101 to be read.
102
103 Parameters
104 ----------
105 filename : `str`
106 A string indicating the file to read
107 stamp_factory : classmethod
108 A factory function defined on a dataclass for constructing
109 stamp objects a la `~lsst.meas.alrogithm.Stamp`
110 options : `PropertyList` or `dict`
111 A collection of parameters. If it contains a bounding box
112 (``bbox`` key), or if certain other keys (``llcX``, ``llcY``,
113 ``width``, ``height``) are available for one to be constructed,
114 the bounding box is passed to the ``FitsReader`` in order to
115 return a sub-image.
116
117 Returns
118 -------
119 stamps : `list` of dataclass objects like `Stamp`, PropertyList
120 A tuple of a list of `Stamp`-like objects
121 metadata : `PropertyList`
122 The metadata
123 """
124 # extract necessary info from metadata
125 metadata = readMetadata(filename, hdu=0)
126 nStamps = metadata["N_STAMPS"]
127 has_archive = metadata["HAS_ARCHIVE"]
128 if has_archive:
129 archive_ids = metadata.getArray("ARCHIVE_IDS")
130 with Fits(filename, "r") as f:
131 nExtensions = f.countHdus()
132 # check if a bbox was provided
133 kwargs = {}
134 if options:
135 # gen3 API
136 if "bbox" in options.keys():
137 kwargs["bbox"] = options["bbox"]
138 # gen2 API
139 elif "llcX" in options.keys():
140 llcX = options["llcX"]
141 llcY = options["llcY"]
142 width = options["width"]
143 height = options["height"]
144 bbox = Box2I(Point2I(llcX, llcY), Extent2I(width, height))
145 kwargs["bbox"] = bbox
146 stamp_parts = {}
147 # We need to be careful because nExtensions includes the primary HDU.
148 for idx in range(nExtensions - 1):
149 md = readMetadata(filename, hdu=idx + 1)
150 if md["EXTNAME"] in ("IMAGE", "VARIANCE"):
151 reader = ImageFitsReader(filename, hdu=idx + 1)
152 elif md["EXTNAME"] == "MASK":
153 reader = MaskFitsReader(filename, hdu=idx + 1)
154 elif md["EXTNAME"] == "ARCHIVE_INDEX":
155 f.setHdu(idx + 1)
156 archive = InputArchive.readFits(f)
157 continue
158 elif md["EXTTYPE"] == "ARCHIVE_DATA":
159 continue
160 else:
161 raise ValueError(f"Unknown extension type: {md['EXTNAME']}")
162 stamp_parts.setdefault(md["EXTVER"], {})[md["EXTNAME"].lower()] = reader.read(**kwargs)
163 if len(stamp_parts) != nStamps:
164 raise ValueError(
165 f"Number of stamps read ({len(stamp_parts)}) does not agree with the "
166 f"number of stamps recorded in the metadata ({nStamps})."
167 )
168 # construct stamps themselves
169 stamps = []
170 for k in range(nStamps):
171 # Need to increment by one since EXTVER starts at 1
172 maskedImage = MaskedImageF(**stamp_parts[k + 1])
173 archive_element = archive.get(archive_ids[k]) if has_archive else None
174 stamps.append(stamp_factory(maskedImage, metadata, k, archive_element))
175
176 return stamps, metadata
177
178
179@dataclass
180class AbstractStamp(abc.ABC):
181 """Single abstract stamp.
182
183 Parameters
184 ----------
185 Inherit from this class to add metadata to the stamp.
186 """
187
188 @classmethod
189 @abc.abstractmethod
190 def factory(cls, stamp_im, metadata, index, archive_element=None):
191 """This method is needed to service the FITS reader. We need a standard
192 interface to construct objects like this. Parameters needed to
193 construct this object are passed in via a metadata dictionary and then
194 passed to the constructor of this class.
195
196 Parameters
197 ----------
199 Pixel data to pass to the constructor
200 metadata : `dict`
201 Dictionary containing the information
202 needed by the constructor.
203 idx : `int`
204 Index into the lists in ``metadata``
205 archive_element : `~lsst.afw.table.io.Persistable`, optional
206 Archive element (e.g. Transform or WCS) associated with this stamp.
207
208 Returns
209 -------
210 stamp : `AbstractStamp`
211 An instance of this class
212 """
213 raise NotImplementedError
214
215
216@dataclass
218 """Single stamp.
219
220 Parameters
221 ----------
222 stamp_im : `~lsst.afw.image.MaskedImageF`
223 The actual pixel values for the postage stamp.
224 archive_element : `~lsst.afw.table.io.Persistable` or `None`, optional
225 Archive element (e.g. Transform or WCS) associated with this stamp.
226 position : `~lsst.geom.SpherePoint` or `None`, optional
227 Position of the center of the stamp. Note the user must keep track of
228 the coordinate system.
229 """
230
231 stamp_im: MaskedImageF
232 archive_element: Persistable | None = None
233 position: SpherePoint | None = SpherePoint(Angle(np.nan), Angle(np.nan))
234
235 @classmethod
236 def factory(cls, stamp_im, metadata, index, archive_element=None):
237 """This method is needed to service the FITS reader. We need a standard
238 interface to construct objects like this. Parameters needed to
239 construct this object are passed in via a metadata dictionary and then
240 passed to the constructor of this class. If lists of values are passed
241 with the following keys, they will be passed to the constructor,
242 otherwise dummy values will be passed: RA_DEG, DEC_DEG. They should
243 each point to lists of values.
244
245 Parameters
246 ----------
248 Pixel data to pass to the constructor
249 metadata : `dict`
250 Dictionary containing the information
251 needed by the constructor.
252 idx : `int`
253 Index into the lists in ``metadata``
254 archive_element : `~lsst.afw.table.io.Persistable`, optional
255 Archive element (e.g. Transform or WCS) associated with this stamp.
256
257 Returns
258 -------
259 stamp : `Stamp`
260 An instance of this class
261 """
262 if "RA_DEG" in metadata and "DEC_DEG" in metadata:
263 return cls(
264 stamp_im=stamp_im,
265 archive_element=archive_element,
266 position=SpherePoint(
267 Angle(metadata.getArray("RA_DEG")[index], degrees),
268 Angle(metadata.getArray("DEC_DEG")[index], degrees),
269 ),
270 )
271 else:
272 return cls(
273 stamp_im=stamp_im,
274 archive_element=archive_element,
275 position=SpherePoint(Angle(np.nan), Angle(np.nan)),
276 )
277
278
279class StampsBase(abc.ABC, Sequence):
280 """Collection of stamps and associated metadata.
281
282 Parameters
283 ----------
284 stamps : iterable
285 This should be an iterable of dataclass objects
287 metadata : `~lsst.daf.base.PropertyList`, optional
288 Metadata associated with the objects within the stamps.
289 use_mask : `bool`, optional
290 If ``True`` read and write the mask data. Default ``True``.
291 use_variance : `bool`, optional
292 If ``True`` read and write the variance data. Default ``True``.
293 use_archive : `bool`, optional
294 If ``True``, read and write an Archive that contains a Persistable
295 associated with each stamp, for example a Transform or a WCS.
296 Default ``False``.
297
298 Notes
299 -----
300 A butler can be used to read only a part of the stamps,
301 specified by a bbox:
302
303 >>> starSubregions = butler.get(
304 "brightStarStamps",
305 dataId,
306 parameters={"bbox": bbox}
307 )
308 """
309
310 def __init__(self, stamps, metadata=None, use_mask=True, use_variance=True, use_archive=False):
311 for stamp in stamps:
312 if not isinstance(stamp, AbstractStamp):
313 raise ValueError(f"The entries in stamps must inherit from AbstractStamp. Got {type(stamp)}.")
314 self._stamps = stamps
315 self._metadata = PropertyList() if metadata is None else metadata.deepCopy()
316 self.use_mask = use_mask
317 self.use_variance = use_variance
318 self.use_archive = use_archive
319
320 @classmethod
321 def readFits(cls, filename):
322 """Build an instance of this class from a file.
323
324 Parameters
325 ----------
326 filename : `str`
327 Name of the file to read
328 """
329
330 return cls.readFitsWithOptions(filename, None)
331
332 @classmethod
333 def readFitsWithOptions(cls, filename, options):
334 """Build an instance of this class with options.
335
336 Parameters
337 ----------
338 filename : `str`
339 Name of the file to read
340 options : `PropertyList`
341 Collection of metadata parameters
342 """
343 # To avoid problems since this is no longer an abstract method.
344 # TO-DO: Consider refactoring this method. This class check was added
345 # to allow the butler formatter to use a generic type but still end up
346 # giving the correct type back, ensuring that the abstract base class
347 # is not used by mistake. Perhaps this logic can be optimised.
348 if cls is not StampsBase:
349 raise NotImplementedError(f"Please implement specific FITS reader for class {cls}")
350
351 # Load metadata to get class
352 metadata = readMetadata(filename, hdu=0)
353 type_name = metadata.get("STAMPCLS")
354 if type_name is None:
355 raise RuntimeError(
356 f"No class name in file {filename}. Unable to instantiate correct stamps subclass. "
357 "Is this an old version format Stamps file?"
358 )
359
360 # Import class and override `cls`
361 stamp_type = doImport(type_name)
362 cls = stamp_type
363
364 return cls.readFitsWithOptions(filename, options)
365
366 @abc.abstractmethod
367 def _refresh_metadata(self):
368 """Make sure metadata is up to date, as this object can be extended."""
369 raise NotImplementedError
370
371 def writeFits(self, filename):
372 """Write this object to a file.
373
374 Parameters
375 ----------
376 filename : `str`
377 Name of file to write.
378 """
379 self._refresh_metadata()
380 type_name = get_full_type_name(self)
381 writeFits(
382 filename,
383 self._stamps,
384 self._metadata,
385 type_name,
386 self.use_mask,
387 self.use_variance,
388 self.use_archive,
389 )
390
391 def __len__(self):
392 return len(self._stamps)
393
394 def __getitem__(self, index):
395 return self._stamps[index]
396
397 def __iter__(self):
398 return iter(self._stamps)
399
401 """Retrieve star images.
402
403 Returns
404 -------
405 maskedImages :
406 `list` [`~lsst.afw.image.MaskedImageF`]
407 """
408 return [stamp.stamp_im for stamp in self._stamps]
409
411 """Retrieve archive elements associated with each stamp.
412
413 Returns
414 -------
415 archiveElements :
417 """
418 return [stamp.archive_element for stamp in self._stamps]
419
420 @property
421 def metadata(self):
422 return self._metadata
423
424
426 def _refresh_metadata(self):
427 positions = self.getPositions()
428 self._metadata["RA_DEG"] = [p.getRa().asDegrees() for p in positions]
429 self._metadata["DEC_DEG"] = [p.getDec().asDegrees() for p in positions]
430
431 def getPositions(self):
432 return [s.position for s in self._stamps]
433
434 def append(self, item):
435 """Add an additional stamp.
436
437 Parameters
438 ----------
439 item : `Stamp`
440 Stamp object to append.
441 """
442 if not isinstance(item, Stamp):
443 raise ValueError("Objects added must be a Stamp object.")
444 self._stamps.append(item)
445 return None
446
447 def extend(self, stamp_list):
448 """Extend Stamps instance by appending elements from another instance.
449
450 Parameters
451 ----------
452 stamps_list : `list` [`Stamp`]
453 List of Stamp object to append.
454 """
455 for s in stamp_list:
456 if not isinstance(s, Stamp):
457 raise ValueError("Can only extend with Stamp objects")
458 self._stamps += stamp_list
459
460 @classmethod
461 def readFits(cls, filename):
462 """Build an instance of this class from a file.
463
464 Parameters
465 ----------
466 filename : `str`
467 Name of the file to read.
468
469 Returns
470 -------
471 object : `Stamps`
472 An instance of this class.
473 """
474 return cls.readFitsWithOptionsreadFitsWithOptions(filename, None)
475
476 @classmethod
477 def readFitsWithOptions(cls, filename, options):
478 """Build an instance of this class with options.
479
480 Parameters
481 ----------
482 filename : `str`
483 Name of the file to read.
484 options : `PropertyList` or `dict`
485 Collection of metadata parameters.
486
487 Returns
488 -------
489 object : `Stamps`
490 An instance of this class.
491 """
492 stamps, metadata = readFitsWithOptions(filename, Stamp.factory, options)
493 return cls(
494 stamps,
495 metadata=metadata,
496 use_mask=metadata["HAS_MASK"],
497 use_variance=metadata["HAS_VARIANCE"],
498 use_archive=metadata["HAS_ARCHIVE"],
499 )
def factory(cls, stamp_im, metadata, index, archive_element=None)
Definition: stamps.py:190
def factory(cls, stamp_im, metadata, index, archive_element=None)
Definition: stamps.py:236
def readFitsWithOptions(cls, filename, options)
Definition: stamps.py:333
def __init__(self, stamps, metadata=None, use_mask=True, use_variance=True, use_archive=False)
Definition: stamps.py:310
def extend(self, stamp_list)
Definition: stamps.py:447
def readFits(cls, filename)
Definition: stamps.py:461
def readFitsWithOptions(cls, filename, options)
Definition: stamps.py:477
def writeFits(filename, stamps, metadata, type_name, write_mask, write_variance, write_archive=False)
Definition: stamps.py:40