Coverage for python/lsst/meas/algorithms/stamps.py: 28%
183 statements
« prev ^ index » next coverage.py v7.4.2, created at 2024-02-21 11:20 +0000
« prev ^ index » next coverage.py v7.4.2, created at 2024-02-21 11:20 +0000
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/>.
22"""Collection of small images (stamps)."""
24__all__ = ["Stamp", "Stamps", "StampsBase", "writeFits", "readFitsWithOptions"]
26import abc
27from collections.abc import Sequence
28from dataclasses import dataclass, field, fields
30import numpy as np
31from lsst.afw.fits import Fits, readMetadata
32from lsst.afw.image import ImageFitsReader, MaskedImage, 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
40def writeFits(filename, stamps, metadata, type_name, write_mask, write_variance, write_archive=False):
41 """Write a single FITS file containing all stamps.
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
99def readFitsWithOptions(filename, stamp_factory, options):
100 """Read stamps from FITS file, allowing for only a subregion of the stamps
101 to be read.
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.algorithm.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.
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
124 Notes
125 -----
126 The data are read using the data type expected by the
127 `~lsst.afw.image.MaskedImage` class attached to the `AbstractStamp`
128 dataclass associated with the factory method.
129 """
130 # extract necessary info from metadata
131 metadata = readMetadata(filename, hdu=0)
132 nStamps = metadata["N_STAMPS"]
133 has_archive = metadata["HAS_ARCHIVE"]
134 if has_archive:
135 archive_ids = metadata.getArray("ARCHIVE_IDS")
136 with Fits(filename, "r") as f:
137 nExtensions = f.countHdus()
138 # check if a bbox was provided
139 kwargs = {}
140 if options:
141 # gen3 API
142 if "bbox" in options.keys():
143 kwargs["bbox"] = options["bbox"]
144 # gen2 API
145 elif "llcX" in options.keys():
146 llcX = options["llcX"]
147 llcY = options["llcY"]
148 width = options["width"]
149 height = options["height"]
150 bbox = Box2I(Point2I(llcX, llcY), Extent2I(width, height))
151 kwargs["bbox"] = bbox
152 stamp_parts = {}
154 # Determine the dtype from the factory. This allows a Stamp class
155 # to be defined in terms of MaskedImageD or MaskedImageI without
156 # forcing everything to floats.
157 masked_image_cls = None
158 for stamp_field in fields(stamp_factory.__self__):
159 if issubclass(stamp_field.type, MaskedImage):
160 masked_image_cls = stamp_field.type
161 break
162 else:
163 raise RuntimeError("Stamp factory does not use MaskedImage.")
164 default_dtype = np.dtype(masked_image_cls.dtype)
165 variance_dtype = np.dtype(np.float32) # Variance is always the same type.
167 # We need to be careful because nExtensions includes the primary HDU.
168 for idx in range(nExtensions - 1):
169 dtype = None
170 md = readMetadata(filename, hdu=idx + 1)
171 if md["EXTNAME"] in ("IMAGE", "VARIANCE"):
172 reader = ImageFitsReader(filename, hdu=idx + 1)
173 if md["EXTNAME"] == "VARIANCE":
174 dtype = variance_dtype
175 else:
176 dtype = default_dtype
177 elif md["EXTNAME"] == "MASK":
178 reader = MaskFitsReader(filename, hdu=idx + 1)
179 elif md["EXTNAME"] == "ARCHIVE_INDEX":
180 f.setHdu(idx + 1)
181 archive = InputArchive.readFits(f)
182 continue
183 elif md["EXTTYPE"] == "ARCHIVE_DATA":
184 continue
185 else:
186 raise ValueError(f"Unknown extension type: {md['EXTNAME']}")
187 stamp_parts.setdefault(md["EXTVER"], {})[md["EXTNAME"].lower()] = reader.read(dtype=dtype,
188 **kwargs)
189 if len(stamp_parts) != nStamps:
190 raise ValueError(
191 f"Number of stamps read ({len(stamp_parts)}) does not agree with the "
192 f"number of stamps recorded in the metadata ({nStamps})."
193 )
194 # construct stamps themselves
195 stamps = []
196 for k in range(nStamps):
197 # Need to increment by one since EXTVER starts at 1
198 maskedImage = masked_image_cls(**stamp_parts[k + 1])
199 archive_element = archive.get(archive_ids[k]) if has_archive else None
200 stamps.append(stamp_factory(maskedImage, metadata, k, archive_element))
202 return stamps, metadata
205@dataclass
206class AbstractStamp(abc.ABC):
207 """Single abstract stamp.
209 Parameters
210 ----------
211 Inherit from this class to add metadata to the stamp.
212 """
214 @classmethod
215 @abc.abstractmethod
216 def factory(cls, stamp_im, metadata, index, archive_element=None):
217 """This method is needed to service the FITS reader. We need a standard
218 interface to construct objects like this. Parameters needed to
219 construct this object are passed in via a metadata dictionary and then
220 passed to the constructor of this class.
222 Parameters
223 ----------
224 stamp : `~lsst.afw.image.MaskedImage`
225 Pixel data to pass to the constructor
226 metadata : `dict`
227 Dictionary containing the information
228 needed by the constructor.
229 idx : `int`
230 Index into the lists in ``metadata``
231 archive_element : `~lsst.afw.table.io.Persistable`, optional
232 Archive element (e.g. Transform or WCS) associated with this stamp.
234 Returns
235 -------
236 stamp : `AbstractStamp`
237 An instance of this class
238 """
239 raise NotImplementedError
242def _default_position():
243 # SpherePoint is nominally mutable in C++ so we must use a factory
244 # and return an entirely new SpherePoint each time a Stamps is created.
245 return SpherePoint(Angle(np.nan), Angle(np.nan))
248@dataclass
249class Stamp(AbstractStamp):
250 """Single stamp.
252 Parameters
253 ----------
254 stamp_im : `~lsst.afw.image.MaskedImageF`
255 The actual pixel values for the postage stamp.
256 archive_element : `~lsst.afw.table.io.Persistable` or `None`, optional
257 Archive element (e.g. Transform or WCS) associated with this stamp.
258 position : `~lsst.geom.SpherePoint` or `None`, optional
259 Position of the center of the stamp. Note the user must keep track of
260 the coordinate system.
261 """
263 stamp_im: MaskedImageF
264 archive_element: Persistable | None = None
265 position: SpherePoint | None = field(default_factory=_default_position)
267 @classmethod
268 def factory(cls, stamp_im, metadata, index, archive_element=None):
269 """This method is needed to service the FITS reader. We need a standard
270 interface to construct objects like this. Parameters needed to
271 construct this object are passed in via a metadata dictionary and then
272 passed to the constructor of this class. If lists of values are passed
273 with the following keys, they will be passed to the constructor,
274 otherwise dummy values will be passed: RA_DEG, DEC_DEG. They should
275 each point to lists of values.
277 Parameters
278 ----------
279 stamp : `~lsst.afw.image.MaskedImage`
280 Pixel data to pass to the constructor
281 metadata : `dict`
282 Dictionary containing the information
283 needed by the constructor.
284 idx : `int`
285 Index into the lists in ``metadata``
286 archive_element : `~lsst.afw.table.io.Persistable`, optional
287 Archive element (e.g. Transform or WCS) associated with this stamp.
289 Returns
290 -------
291 stamp : `Stamp`
292 An instance of this class
293 """
294 if "RA_DEG" in metadata and "DEC_DEG" in metadata:
295 return cls(
296 stamp_im=stamp_im,
297 archive_element=archive_element,
298 position=SpherePoint(
299 Angle(metadata.getArray("RA_DEG")[index], degrees),
300 Angle(metadata.getArray("DEC_DEG")[index], degrees),
301 ),
302 )
303 else:
304 return cls(
305 stamp_im=stamp_im,
306 archive_element=archive_element,
307 position=SpherePoint(Angle(np.nan), Angle(np.nan)),
308 )
311class StampsBase(abc.ABC, Sequence):
312 """Collection of stamps and associated metadata.
314 Parameters
315 ----------
316 stamps : iterable
317 This should be an iterable of dataclass objects
318 a la ``~lsst.meas.algorithms.Stamp``.
319 metadata : `~lsst.daf.base.PropertyList`, optional
320 Metadata associated with the objects within the stamps.
321 use_mask : `bool`, optional
322 If ``True`` read and write the mask data. Default ``True``.
323 use_variance : `bool`, optional
324 If ``True`` read and write the variance data. Default ``True``.
325 use_archive : `bool`, optional
326 If ``True``, read and write an Archive that contains a Persistable
327 associated with each stamp, for example a Transform or a WCS.
328 Default ``False``.
330 Notes
331 -----
332 A butler can be used to read only a part of the stamps,
333 specified by a bbox:
335 >>> starSubregions = butler.get(
336 "brightStarStamps",
337 dataId,
338 parameters={"bbox": bbox}
339 )
340 """
342 def __init__(self, stamps, metadata=None, use_mask=True, use_variance=True, use_archive=False):
343 for stamp in stamps:
344 if not isinstance(stamp, AbstractStamp):
345 raise ValueError(f"The entries in stamps must inherit from AbstractStamp. Got {type(stamp)}.")
346 self._stamps = stamps
347 self._metadata = PropertyList() if metadata is None else metadata.deepCopy()
348 self.use_mask = use_mask
349 self.use_variance = use_variance
350 self.use_archive = use_archive
352 @classmethod
353 def readFits(cls, filename):
354 """Build an instance of this class from a file.
356 Parameters
357 ----------
358 filename : `str`
359 Name of the file to read
360 """
362 return cls.readFitsWithOptions(filename, None)
364 @classmethod
365 def readFitsWithOptions(cls, filename, options):
366 """Build an instance of this class with options.
368 Parameters
369 ----------
370 filename : `str`
371 Name of the file to read
372 options : `PropertyList`
373 Collection of metadata parameters
374 """
375 # To avoid problems since this is no longer an abstract method.
376 # TO-DO: Consider refactoring this method. This class check was added
377 # to allow the butler formatter to use a generic type but still end up
378 # giving the correct type back, ensuring that the abstract base class
379 # is not used by mistake. Perhaps this logic can be optimised.
380 if cls is not StampsBase:
381 raise NotImplementedError(f"Please implement specific FITS reader for class {cls}")
383 # Load metadata to get class
384 metadata = readMetadata(filename, hdu=0)
385 type_name = metadata.get("STAMPCLS")
386 if type_name is None:
387 raise RuntimeError(
388 f"No class name in file {filename}. Unable to instantiate correct stamps subclass. "
389 "Is this an old version format Stamps file?"
390 )
392 # Import class and override `cls`
393 stamp_type = doImport(type_name)
394 cls = stamp_type
396 return cls.readFitsWithOptions(filename, options)
398 @abc.abstractmethod
399 def _refresh_metadata(self):
400 """Make sure metadata is up to date, as this object can be extended."""
401 raise NotImplementedError
403 def writeFits(self, filename):
404 """Write this object to a file.
406 Parameters
407 ----------
408 filename : `str`
409 Name of file to write.
410 """
411 self._refresh_metadata()
412 type_name = get_full_type_name(self)
413 writeFits(
414 filename,
415 self._stamps,
416 self._metadata,
417 type_name,
418 self.use_mask,
419 self.use_variance,
420 self.use_archive,
421 )
423 def __len__(self):
424 return len(self._stamps)
426 def __getitem__(self, index):
427 return self._stamps[index]
429 def __iter__(self):
430 return iter(self._stamps)
432 def getMaskedImages(self):
433 """Retrieve star images.
435 Returns
436 -------
437 maskedImages :
438 `list` [`~lsst.afw.image.MaskedImageF`]
439 """
440 return [stamp.stamp_im for stamp in self._stamps]
442 def getArchiveElements(self):
443 """Retrieve archive elements associated with each stamp.
445 Returns
446 -------
447 archiveElements :
448 `list` [`~lsst.afw.table.io.Persistable`]
449 """
450 return [stamp.archive_element for stamp in self._stamps]
452 @property
453 def metadata(self):
454 return self._metadata
457class Stamps(StampsBase):
458 def _refresh_metadata(self):
459 positions = self.getPositions()
460 self._metadata["RA_DEG"] = [p.getRa().asDegrees() for p in positions]
461 self._metadata["DEC_DEG"] = [p.getDec().asDegrees() for p in positions]
463 def getPositions(self):
464 return [s.position for s in self._stamps]
466 def append(self, item):
467 """Add an additional stamp.
469 Parameters
470 ----------
471 item : `Stamp`
472 Stamp object to append.
473 """
474 if not isinstance(item, Stamp):
475 raise ValueError("Objects added must be a Stamp object.")
476 self._stamps.append(item)
477 return None
479 def extend(self, stamp_list):
480 """Extend Stamps instance by appending elements from another instance.
482 Parameters
483 ----------
484 stamps_list : `list` [`Stamp`]
485 List of Stamp object to append.
486 """
487 for s in stamp_list:
488 if not isinstance(s, Stamp):
489 raise ValueError("Can only extend with Stamp objects")
490 self._stamps += stamp_list
492 @classmethod
493 def readFits(cls, filename):
494 """Build an instance of this class from a file.
496 Parameters
497 ----------
498 filename : `str`
499 Name of the file to read.
501 Returns
502 -------
503 object : `Stamps`
504 An instance of this class.
505 """
506 return cls.readFitsWithOptions(filename, None)
508 @classmethod
509 def readFitsWithOptions(cls, filename, options):
510 """Build an instance of this class with options.
512 Parameters
513 ----------
514 filename : `str`
515 Name of the file to read.
516 options : `PropertyList` or `dict`
517 Collection of metadata parameters.
519 Returns
520 -------
521 object : `Stamps`
522 An instance of this class.
523 """
524 stamps, metadata = readFitsWithOptions(filename, Stamp.factory, options)
525 return cls(
526 stamps,
527 metadata=metadata,
528 use_mask=metadata["HAS_MASK"],
529 use_variance=metadata["HAS_VARIANCE"],
530 use_archive=metadata["HAS_ARCHIVE"],
531 )