Coverage for python/lsst/images/_mask.py: 22%
402 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-06 01:43 -0700
1# This file is part of lsst-images.
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# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
12from __future__ import annotations
14__all__ = (
15 "Mask",
16 "MaskPlane",
17 "MaskPlaneBit",
18 "MaskSchema",
19 "MaskSerializationModel",
20 "get_legacy_deep_coadd_mask_planes",
21 "get_legacy_difference_image_mask_planes",
22 "get_legacy_visit_image_mask_planes",
23)
25import dataclasses
26import math
27from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence, Set
28from types import EllipsisType
29from typing import TYPE_CHECKING, Any, ClassVar, cast
31import astropy.io.fits
32import astropy.wcs
33import numpy as np
34import numpy.typing as npt
35import pydantic
37from lsst.resources import ResourcePath, ResourcePathExpression
39from . import fits
40from ._generalized_image import GeneralizedImage
41from ._geom import YX, Box, NoOverlapError
42from ._transforms import Frame, Projection, ProjectionSerializationModel
43from .serialization import (
44 ArchiveReadError,
45 ArchiveTree,
46 ArrayReferenceModel,
47 InlineArrayModel,
48 InputArchive,
49 IntegerType,
50 InvalidParameterError,
51 MetadataValue,
52 NumberType,
53 OutputArchive,
54 is_integer,
55 no_header_updates,
56)
57from .utils import is_none
59if TYPE_CHECKING:
60 try:
61 from lsst.afw.image import Mask as LegacyMask
62 except ImportError:
63 type LegacyMask = Any # type: ignore[no-redef]
66@dataclasses.dataclass(frozen=True)
67class MaskPlane:
68 """Name and description of a single plane in a mask array."""
70 name: str
71 """Unique name for the mask plane (`str`)."""
73 description: str
74 """Human-readable documentation for the mask plane (`str`)."""
76 @classmethod
77 def read_legacy(cls, header: astropy.io.fits.Header) -> dict[str, int]:
78 """Read mask plane descriptions written by
79 `lsst.afw.image.Mask.writeFits`.
81 Parameters
82 ----------
83 header
84 FITS header.
86 Returns
87 -------
88 `dict` [`str`, `int`]
89 A dictionary mapping mask plane name to integer bit index.
90 """
91 result: dict[str, int] = {}
92 for card in list(header.cards):
93 if card.keyword.startswith("MP_"):
94 result[card.keyword.removeprefix("MP_")] = card.value
95 del header[card.keyword]
96 return result
99@dataclasses.dataclass(frozen=True)
100class MaskPlaneBit:
101 """The nested array index and mask value associated with a single mask
102 plane.
103 """
105 index: int
106 """Index into the last dimension of the mask array where this plane's bit
107 is stored.
108 """
110 mask: np.integer
111 """Bitmask that selects just this plane's bit from a mask array value
112 (`numpy.integer`).
113 """
115 @classmethod
116 def compute(cls, overall_index: int, stride: int, mask_type: type[np.integer]) -> MaskPlaneBit:
117 """Construct a `MaskPlaneBit` from the overall index of a plane in a
118 `MaskSchema` and the stride (number of bits per mask array element).
119 """
120 index, bit = divmod(overall_index, stride)
121 return cls(index, mask_type(1 << bit))
124class MaskSchema:
125 """A schema for a bit-packed mask array.
127 Parameters
128 ----------
129 planes
130 Iterable of `MaskPlane` instances that define the schema. `None`
131 values may be included to reserve bits for future use.
132 dtype
133 The numpy data type of the mask arrays that use this schema.
135 Notes
136 -----
137 A `MaskSchema` is a collection of mask planes, which each correspond to a
138 single bit in a mask array. Mask schemas are immutable and associated with
139 a particular array data type, allowing them to safely precompute the index
140 and bitmask for each plane.
142 `MaskSchema` indexing is by integer (the overall index of a plane in the
143 schema). The `descriptions` attribute may be indexed by plane name to get
144 the description for that plane, and the `bitmask` method can be used to
145 obtain an array that can be used to select one or more planes by name in
146 a mask array that uses this schema.
148 If no mask planes are provided, a `None` placeholder is automatically
149 added.
150 """
152 def __init__(self, planes: Iterable[MaskPlane | None], dtype: npt.DTypeLike = np.uint8):
153 self._planes: tuple[MaskPlane | None, ...] = tuple(planes) or (None,)
154 self._dtype = cast(np.dtype[np.integer], np.dtype(dtype))
155 stride = self.bits_per_element(self._dtype)
156 self._descriptions = {plane.name: plane.description for plane in self._planes if plane is not None}
157 self._mask_size = math.ceil(len(self._planes) / stride)
158 self._bits: dict[str, MaskPlaneBit] = {
159 plane.name: MaskPlaneBit.compute(n, stride, self._dtype.type)
160 for n, plane in enumerate(self._planes)
161 if plane is not None
162 }
164 @staticmethod
165 def bits_per_element(dtype: npt.DTypeLike) -> int:
166 """Return the number of mask bits per array element for the given
167 data type.
168 """
169 dtype = np.dtype(dtype)
170 match dtype.kind:
171 case "u":
172 return dtype.itemsize * 8
173 case "i":
174 return dtype.itemsize * 8 - 1
175 case _:
176 raise TypeError(f"dtype for masks must be an integer; got {dtype} with kind={dtype.kind}.")
178 def __iter__(self) -> Iterator[MaskPlane | None]:
179 return iter(self._planes)
181 def __len__(self) -> int:
182 return len(self._planes)
184 def __contains__(self, plane: str | MaskPlane) -> bool:
185 return getattr(plane, "name", plane) in self.names
187 def __getitem__(self, i: int) -> MaskPlane | None:
188 return self._planes[i]
190 def __repr__(self) -> str:
191 return f"MaskSchema({list(self._planes)}, dtype={self._dtype!r})"
193 def __str__(self) -> str:
194 return "\n".join(
195 [
196 f"{name} [{bit.index}@{hex(bit.mask)}]: {self._descriptions[name]}"
197 for name, bit in self._bits.items()
198 ]
199 )
201 def __eq__(self, other: object) -> bool:
202 if isinstance(other, MaskSchema):
203 return self._planes == other._planes and self._dtype == other._dtype
204 return False
206 @property
207 def dtype(self) -> np.dtype:
208 """The numpy data type of the mask arrays that use this schema."""
209 return self._dtype
211 @property
212 def mask_size(self) -> int:
213 """The number of elements in the last dimension of any mask array that
214 uses this schema.
215 """
216 return self._mask_size
218 @property
219 def names(self) -> Set[str]:
220 """The names of the mask planes, in bit order."""
221 return self._bits.keys()
223 @property
224 def descriptions(self) -> Mapping[str, str]:
225 """A mapping from plane name to description."""
226 return self._descriptions
228 def bit(self, plane: str) -> MaskPlaneBit:
229 """Return the last array index and mask for the given mask plane."""
230 return self._bits[plane]
232 def bitmask(self, *planes: str) -> np.ndarray:
233 """Return a 1-d mask array that represents the union (i.e. bitwise OR)
234 of the planes with the given names.
236 Parameters
237 ----------
238 *planes
239 Mask plane names.
241 Returns
242 -------
243 numpy.ndarray
244 A 1-d array with shape ``(mask_size,)``.
245 """
246 result = np.zeros(self.mask_size, dtype=self._dtype)
247 for plane in planes:
248 bit = self._bits[plane]
249 result[bit.index] |= bit.mask
250 return result
252 def split(self, dtype: npt.DTypeLike) -> list[MaskSchema]:
253 """Split the schema into an equivalent series of schemas that each
254 have a `mask_size` of ``1``, dropping all `None` placeholders.
256 Parameters
257 ----------
258 dtype
259 Data type of the new mask pixels.
261 Returns
262 -------
263 `list` [`MaskSchema`]
264 A list of mask schemas that together include all planes in
265 ``self`` and have `mask_size` equal to ``1``. If there are no
266 mask planes (only `None` placeholders) in ``self``, a single mask
267 schema with a `None` placeholder is returned; otherwise `None`
268 placeholders are returned.
269 """
270 dtype = np.dtype(dtype)
271 planes: list[MaskPlane] = []
272 schemas: list[MaskSchema] = []
273 n_planes_per_schema = self.bits_per_element(dtype)
274 for plane in self._planes:
275 if plane is not None:
276 planes.append(plane)
277 if len(planes) == n_planes_per_schema:
278 schemas.append(MaskSchema(planes, dtype=dtype))
279 planes.clear()
280 if planes:
281 schemas.append(MaskSchema(planes, dtype=dtype))
282 if not schemas:
283 schemas.append(MaskSchema([None], dtype=dtype))
284 return schemas
286 def update_header(self, header: astropy.io.fits.Header) -> None:
287 """Add a description of this mask schema to a FITS header."""
288 for n, plane in enumerate(self):
289 if plane is not None:
290 bit = self.bit(plane.name)
291 if bit.index != 0:
292 raise TypeError("Only mask schemas with mask_size==1 can be described in FITS.")
293 header.set(f"MSKN{n:04d}", plane.name, f"Name for mask plane {n}.")
294 header.set(f"MSKM{n:04d}", bit.mask, f"Bitmask for plane n={n}; always 1<<n.")
295 # We don't add a comment to the description card, because it's
296 # likely to overrun a single card and get the CONTINUE
297 # treatment. That will cause Astropy to warn about the comment
298 # being truncated and that's worse than just leaving it
299 # unexplained; it's pretty obvious from context what it is.
300 header.set(f"MSKD{n:04d}", plane.description)
302 def strip_header(self, header: astropy.io.fits.Header) -> None:
303 """Remove all header cards added by `update_header`."""
304 for n, plane in enumerate(self):
305 if plane is not None:
306 header.remove(f"MSKN{n:04d}", ignore_missing=True)
307 header.remove(f"MSKM{n:04d}", ignore_missing=True)
308 header.remove(f"MSKD{n:04d}", ignore_missing=True)
311class Mask(GeneralizedImage):
312 """A 2-d bitmask image backed by a 3-d byte array.
314 Parameters
315 ----------
316 array_or_fill
317 Array or fill value for the mask. If a fill value, ``bbox`` or
318 ``shape`` must be provided.
319 schema
320 Schema that defines the planes and their bit assignments.
321 bbox
322 Bounding box for the mask. This sets the shape of the first two
323 dimensions of the array.
324 start
325 Logical coordinates of the first pixel in the array, ordered ``y``,
326 ``x`` (unless an `XY` instance is passed). Ignored if
327 ``bbox`` is provided. Defaults to zeros.
328 shape
329 Leading dimensions of the array, ordered ``y``, ``x`` (unless an `XY`
330 instance is passed). Only needed if ``array_or_fill`` is not an
331 array and ``bbox`` is not provided. Like the bbox, this does not
332 include the last dimension of the array.
333 projection
334 Projection that maps the pixel grid to the sky.
335 metadata
336 Arbitrary flexible metadata to associate with the mask.
338 Notes
339 -----
340 Indexing the `array` attribute of a `Mask` does not take into account its
341 ``start`` offset, but accessing a subimage mask by indexing a `Mask` with
342 a `Box` does, and the `bbox` of the subimage is set to match its location
343 within the original mask.
345 A mask's ``bbox`` corresponds to the leading dimensions of its backing
346 `numpy.ndarray`, while the last dimension's size is always equal to the
347 `~MaskSchema.mask_size` of its schema, since a schema can in general
348 require multiple array elements to represent all of its planes.
349 """
351 def __init__(
352 self,
353 array_or_fill: np.ndarray | int = 0,
354 /,
355 *,
356 schema: MaskSchema,
357 bbox: Box | None = None,
358 start: Sequence[int] | None = None,
359 shape: Sequence[int] | None = None,
360 projection: Projection | None = None,
361 metadata: dict[str, MetadataValue] | None = None,
362 ):
363 super().__init__(metadata)
364 if shape is not None:
365 shape = tuple(shape)
366 if start is not None:
367 start = tuple(start)
368 if isinstance(array_or_fill, np.ndarray):
369 array = np.array(array_or_fill, dtype=schema.dtype, copy=None)
370 if array.ndim != 3:
371 raise ValueError("Mask array must be 3-d.")
372 if bbox is None:
373 bbox = Box.from_shape(array.shape[:-1], start=start)
374 elif bbox.shape + (schema.mask_size,) != array.shape:
375 raise ValueError(
376 f"Explicit bbox shape {bbox.shape} and schema of size {schema.mask_size} do not "
377 f"match array with shape {array.shape}."
378 )
379 if shape is not None and shape + (schema.mask_size,) != array.shape:
380 raise ValueError(
381 f"Explicit shape {shape} and schema of size {schema.mask_size} do "
382 f"not match array with shape {array.shape}."
383 )
385 else:
386 if bbox is None:
387 if shape is None:
388 raise TypeError("No bbox, size, or array provided.")
389 bbox = Box.from_shape(shape, start=start)
390 array = np.full(bbox.shape + (schema.mask_size,), array_or_fill, dtype=schema.dtype)
391 self._array = array
392 self._bbox: Box = bbox
393 self._schema: MaskSchema = schema
394 self._projection = projection
396 @property
397 def array(self) -> np.ndarray:
398 """The low-level array (`numpy.ndarray`).
400 Assigning to this attribute modifies the existing array in place; the
401 bounding box and underlying data pointer are never changed.
402 """
403 return self._array
405 @array.setter
406 def array(self, value: np.ndarray | int) -> None:
407 self._array[:, :] = value
409 @property
410 def schema(self) -> MaskSchema:
411 """Schema that defines the planes and their bit assignments
412 (`MaskSchema`).
413 """
414 return self._schema
416 @property
417 def bbox(self) -> Box:
418 """2-d bounding box of the mask (`Box`).
420 This sets the shape of the first two dimensions of the array.
421 """
422 return self._bbox
424 @property
425 def projection(self) -> Projection[Any] | None:
426 """The projection that maps this mask's pixel grid to the sky
427 (`Projection` | `None`).
429 Notes
430 -----
431 The pixel coordinates used by this projection account for the bounding
432 box ``start``; they are not just array indices.
433 """
434 return self._projection
436 def __getitem__(self, bbox: Box | EllipsisType) -> Mask:
437 if bbox is ...:
438 return self
439 super().__getitem__(bbox)
440 return self._transfer_metadata(
441 Mask(
442 self.array[bbox.y.slice_within(self._bbox.y), bbox.x.slice_within(self._bbox.x), :],
443 bbox=bbox,
444 schema=self.schema,
445 projection=self._projection,
446 ),
447 bbox=bbox,
448 )
450 def __setitem__(self, bbox: Box | EllipsisType, value: Mask) -> None:
451 subview = self[bbox]
452 subview.clear()
453 subview.update(value)
455 def __str__(self) -> str:
456 return f"Mask({self.bbox!s}, {list(self.schema.names)})"
458 def __repr__(self) -> str:
459 return f"Mask(..., bbox={self.bbox!r}, schema={self.schema!r})"
461 def __eq__(self, other: object) -> bool:
462 if not isinstance(other, Mask):
463 return NotImplemented
464 return (
465 self._bbox == other._bbox
466 and self._schema == other._schema
467 and np.array_equal(self._array, other._array, equal_nan=True)
468 )
470 def copy(self) -> Mask:
471 """Deep-copy the mask and metadata."""
472 return self._transfer_metadata(
473 Mask(self._array.copy(), bbox=self._bbox, schema=self._schema, projection=self._projection),
474 copy=True,
475 )
477 def view(
478 self,
479 *,
480 schema: MaskSchema | EllipsisType = ...,
481 projection: Projection | None | EllipsisType = ...,
482 start: Sequence[int] | EllipsisType = ...,
483 ) -> Mask:
484 """Make a view of the mask, with optional updates.
486 Notes
487 -----
488 This can only be used to make changes to schema descriptions; plane
489 names must remain the same (in the same order).
490 """
491 if schema is ...:
492 schema = self._schema
493 else:
494 if list(schema.names) != list(self.schema.names):
495 raise ValueError("Cannot create a mask view with a schema with different names.")
496 if projection is ...:
497 projection = self._projection
498 if start is ...:
499 start = self._bbox.start
500 return self._transfer_metadata(Mask(self._array, start=start, schema=schema, projection=projection))
502 def update(self, other: Mask) -> None:
503 """Update ``self`` to include all common mask values set in ``other``.
505 Notes
506 -----
507 This only operates on the intersection of the two mask bounding boxes
508 and the mask planes that are present in both. Mask bits are only set,
509 not cleared (i.e. this uses ``|=`` updates, not ``=`` assignments).
510 """
511 lhs = self
512 rhs = other
513 if other.bbox != self.bbox:
514 try:
515 bbox = self.bbox.intersection(other.bbox)
516 except NoOverlapError:
517 return
518 lhs = self[bbox]
519 rhs = other[bbox]
520 for name in self.schema.names & other.schema.names:
521 lhs.set(name, rhs.get(name))
523 def get(self, plane: str) -> np.ndarray:
524 """Return a 2-d boolean array for the given mask plane.
526 Parameters
527 ----------
528 plane
529 Name of the mask plane.
531 Returns
532 -------
533 numpy.ndarray
534 A 2-d boolean array with the same shape as `bbox` that is `True`
535 where the bit for ``plane`` is set and `False` elsewhere.
536 """
537 bit = self.schema.bit(plane)
538 return (self._array[..., bit.index] & bit.mask).astype(bool)
540 def set(self, plane: str, boolean_mask: np.ndarray | EllipsisType = ...) -> None:
541 """Set a mask plane.
543 Parameters
544 ----------
545 plane
546 Name of the mask plane to set
547 boolean_mask
548 A 2-d boolean array with the same shape as `bbox` that is `True`
549 where the bit for ``plane`` should be set and `False` where it
550 should be left unchanged (*not* set to zero). May be ``...`` to
551 set the bit everywhere.
552 """
553 bit = self.schema.bit(plane)
554 if boolean_mask is not ...:
555 boolean_mask = boolean_mask.astype(bool)
556 self._array[boolean_mask, bit.index] |= bit.mask
558 def clear(self, plane: str | None = None, boolean_mask: np.ndarray | EllipsisType = ...) -> None:
559 """Clear one or more mask planes.
561 Parameters
562 ----------
563 plane
564 Name of the mask plane to set. If `None` all mask planes are
565 cleared.
566 boolean_mask
567 A 2-d boolean array with the same shape as `bbox` that is `True`
568 where the bit for ``plane`` should be cleared and `False` where it
569 should be left unchanged. May be ``...`` to clear the bit
570 everywhere.
571 """
572 if boolean_mask is not ...:
573 boolean_mask = boolean_mask.astype(bool)
574 if plane is None:
575 self._array[boolean_mask, :] = 0
576 else:
577 bit = self.schema.bit(plane)
578 self._array[boolean_mask, bit.index] &= ~bit.mask
580 def serialize[P: pydantic.BaseModel](
581 self,
582 archive: OutputArchive[P],
583 *,
584 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
585 save_projection: bool = True,
586 add_offset_wcs: str | None = "A",
587 ) -> MaskSerializationModel[P]:
588 """Serialize the mask to an output archive.
590 Parameters
591 ----------
592 archive
593 Archive to write to.
594 update_header
595 A callback that will be given the FITS header for the HDU
596 containing this mask in order to add keys to it. This callback
597 may be provided but will not be called if the output format is not
598 FITS. As multiple HDUs may be added, this function may be called
599 multiple times.
600 save_projection
601 If `True`, save the `Projection` attached to the image, if there
602 is one. This does not affect whether a FITS WCS corresponding to
603 the projection is written (it always is, if available, and if
604 ``add_offset_wcs`` is not ``" "``).
605 add_offset_wcs
606 A FITS WCS single-character suffix to use when adding a linear
607 WCS that maps the FITS array to the logical pixel coordinates
608 defined by ``bbox.start``. Set to `None` to not write this WCS.
609 If this is set to ``" "``, it will prevent the `Projection` from
610 being saved as a FITS WCS.
611 """
612 if _archive_prefers_native_mask_arrays(archive):
613 # HDS presents array dimensions in Fortran order, which is the
614 # reverse of the h5py dataset shape. Store the in-memory trailing
615 # mask-byte axis first in HDF5 so Starlink tools see HDS axes
616 # (x, y, byte), without changing the bit packing within a pixel.
617 array_model = archive.add_array(np.moveaxis(self._array, -1, 0), update_header=update_header)
618 if not isinstance(array_model, ArrayReferenceModel):
619 raise RuntimeError("Native mask arrays require reference array storage.")
620 array_model.shape = list(self._array.shape)
621 data: list[ArrayReferenceModel | InlineArrayModel] = [array_model]
622 else:
623 data = []
624 for schema_2d in self.schema.split(np.int32):
625 mask_2d = Mask(0, bbox=self.bbox, schema=schema_2d, projection=self._projection)
626 mask_2d.update(self)
627 data.append(
628 mask_2d._serialize_2d(archive, update_header=update_header, add_offset_wcs=add_offset_wcs)
629 )
630 serialized_projection: ProjectionSerializationModel[P] | None = None
631 if save_projection and self.projection is not None:
632 serialized_projection = archive.serialize_direct("projection", self.projection.serialize)
633 serialized_dtype = NumberType.from_numpy(self.schema.dtype)
634 assert is_integer(serialized_dtype), "Mask dtypes should always be integers."
635 return MaskSerializationModel.model_construct(
636 data=data,
637 start=list(self.bbox.start),
638 planes=list(self.schema),
639 dtype=serialized_dtype,
640 projection=serialized_projection,
641 metadata=self.metadata,
642 )
644 def _serialize_2d[P: pydantic.BaseModel](
645 self,
646 archive: OutputArchive[P],
647 *,
648 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
649 add_offset_wcs: str | None = "A",
650 ) -> ArrayReferenceModel | InlineArrayModel:
651 def _update_header(header: astropy.io.fits.Header) -> None:
652 update_header(header)
653 self.schema.update_header(header)
654 if self.projection is not None and add_offset_wcs != " ":
655 if self.fits_wcs:
656 header.update(self.fits_wcs.to_header(relax=True))
657 if add_offset_wcs is not None:
658 fits.add_offset_wcs(header, x=self.bbox.x.start, y=self.bbox.y.start, key=add_offset_wcs)
660 assert self.array.shape[2] == 1, "Mask should be split before calling this method."
661 return archive.add_array(self._array[:, :, 0], update_header=_update_header)
663 @staticmethod
664 def _get_archive_tree_type[P: pydantic.BaseModel](
665 pointer_type: type[P],
666 ) -> type[MaskSerializationModel[P]]:
667 """Return the serialization model type for this object for an archive
668 type that uses the given pointer type.
669 """
670 return MaskSerializationModel[pointer_type] # type: ignore
672 _archive_default_name: ClassVar[str] = "mask"
673 """The name this object should be serialized with when written as the
674 top-level object.
675 """
677 @staticmethod
678 def from_legacy(
679 legacy: Any,
680 plane_map: Mapping[str, MaskPlane] | None = None,
681 ) -> Mask:
682 """Convert from an `lsst.afw.image.Mask` instance.
684 Parameters
685 ----------
686 legacy
687 An `lsst.afw.image.Mask` instance. This will not share pixel
688 data with the new object.
689 plane_map
690 A mapping from legacy mask plane name to the new plane name and
691 description. If not provided, the right legacy mask plane will be
692 guessed, but this can depend on which mask planes the legacy
693 mask actually has set.
694 """
695 return Mask._from_legacy_array(
696 legacy.array,
697 legacy.getMaskPlaneDict(),
698 start=YX(y=legacy.getY0(), x=legacy.getX0()),
699 plane_map=plane_map,
700 )
702 def to_legacy(self, plane_map: Mapping[str, MaskPlane] | None = None) -> Any:
703 """Convert to an `lsst.afw.image.Mask` instance.
705 The pixel data will not be shared between the two objects.
707 Parameters
708 ----------
709 plane_map
710 A mapping from legacy mask plane name to the new plane name and
711 description.
712 """
713 import lsst.afw.image
714 import lsst.geom
716 result = lsst.afw.image.Mask(self.bbox.to_legacy())
717 if plane_map is None:
718 plane_map = {plane.name: plane for plane in self.schema if plane is not None}
719 for old_name, new_plane in plane_map.items():
720 old_bit = result.addMaskPlane(old_name)
721 old_bitmask = 1 << old_bit
722 if new_plane in self.schema:
723 result.array[self.get(new_plane.name)] |= old_bitmask
724 return result
726 @staticmethod
727 def _from_legacy_array(
728 array2d: np.ndarray,
729 old_planes: Mapping[str, int],
730 *,
731 start: YX[int],
732 plane_map: Mapping[str, MaskPlane] | None = None,
733 projection: Projection | None = None,
734 ) -> Mask:
735 if plane_map is None:
736 plane_map = _guess_legacy_plane_map(old_planes)
737 planes: list[MaskPlane] = list(plane_map.values()) if plane_map is not None else []
738 new_name_to_old_bitmask: dict[str, int] = {}
739 for old_name, old_bit in old_planes.items():
740 old_bitmask = 1 << old_bit
741 if plane_map is not None:
742 if new_plane := plane_map.get(old_name):
743 # Already added to 'planes' at initialization.
744 new_name_to_old_bitmask[new_plane.name] = old_bitmask
745 else:
746 if n_orphaned := np.count_nonzero(array2d.astype(np.uint64) & old_bitmask):
747 raise RuntimeError(
748 f"Legacy mask plane {old_name!r} is not remapped, "
749 f"but {n_orphaned} pixels have this bit set."
750 )
751 else:
752 planes.append(MaskPlane(old_name, ""))
753 new_name_to_old_bitmask[old_name] = old_bitmask
754 schema = MaskSchema(planes)
755 mask = Mask(0, schema=schema, start=start, shape=array2d.shape, projection=projection)
756 for new_name, old_bitmask in new_name_to_old_bitmask.items():
757 mask.set(new_name, array2d & old_bitmask)
758 return mask
760 @staticmethod
761 def read_legacy(
762 uri: ResourcePathExpression,
763 *,
764 plane_map: Mapping[str, MaskPlane] | None = None,
765 ext: str | int = 1,
766 fits_wcs_frame: Frame | None = None,
767 ) -> Mask:
768 """Read a FITS file written by `lsst.afw.image.Mask.writeFits`.
770 Parameters
771 ----------
772 uri
773 URI or file name.
774 plane_map
775 A mapping from legacy mask plane name to the new plane name and
776 description. If not provided, the right legacy mask plane will be
777 guessed, but this can depend on which mask planes the legacy
778 mask actually has set.
779 ext
780 Name or index of the FITS HDU to read.
781 fits_wcs_frame
782 If not `None` and the HDU containing the mask has a FITS WCS,
783 attach a `Projection` to the returned mask by converting that WCS.
784 """
785 opaque_metadata = fits.FitsOpaqueMetadata()
786 fs, fspath = ResourcePath(uri).to_fsspec()
787 with fs.open(fspath) as stream, astropy.io.fits.open(stream) as hdu_list:
788 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header)
789 result = Mask._read_legacy_hdu(
790 hdu_list[ext], opaque_metadata, plane_map=plane_map, fits_wcs_frame=fits_wcs_frame
791 )
792 result._opaque_metadata = opaque_metadata
793 return result
795 @staticmethod
796 def _read_legacy_hdu(
797 hdu: astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU | astropy.io.fits.BinTableHDU,
798 opaque_metadata: fits.FitsOpaqueMetadata,
799 plane_map: Mapping[str, MaskPlane] | None = None,
800 fits_wcs_frame: Frame | None = None,
801 ) -> Mask:
802 if isinstance(hdu, astropy.io.fits.BinTableHDU):
803 hdu = astropy.io.fits.CompImageHDU(bintable=hdu)
804 dx: int = hdu.header.pop("LTV1")
805 dy: int = hdu.header.pop("LTV2")
806 start = YX(y=-dy, x=-dx)
807 old_planes = MaskPlane.read_legacy(hdu.header)
808 projection: Projection | None = None
809 if fits_wcs_frame is not None:
810 try:
811 fits_wcs = astropy.wcs.WCS(hdu.header)
812 except KeyError:
813 pass
814 else:
815 projection = Projection.from_fits_wcs(
816 fits_wcs, pixel_frame=fits_wcs_frame, x0=start.x, y0=start.y
817 )
818 mask = Mask._from_legacy_array(
819 hdu.data, old_planes, start=start, plane_map=plane_map, projection=projection
820 )
821 fits.strip_wcs_cards(hdu.header)
822 hdu.header.strip()
823 hdu.header.remove("EXTTYPE", ignore_missing=True)
824 hdu.header.remove("INHERIT", ignore_missing=True)
825 # afw set BUNIT on masks because of limitations in how FITS
826 # metadata is handled there.
827 hdu.header.remove("BUNIT", ignore_missing=True)
828 opaque_metadata.add_header(hdu.header)
829 return mask
832class MaskSerializationModel[P: pydantic.BaseModel](ArchiveTree):
833 """Pydantic model used to represent the serialized form of a `.Mask`."""
835 SCHEMA_NAME: ClassVar[str] = "mask"
836 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
837 MIN_READ_VERSION: ClassVar[int] = 1
838 PUBLIC_TYPE: ClassVar[type] = Mask
840 data: list[ArrayReferenceModel | InlineArrayModel] = pydantic.Field(
841 description="References to pixel data."
842 )
843 start: list[int] = pydantic.Field(
844 description="Coordinate of the first pixels in the array, ordered (y, x)."
845 )
846 planes: list[MaskPlane | None] = pydantic.Field(description="Definitions of the bitplanes in the mask.")
847 dtype: IntegerType = pydantic.Field(description="Data type of the in-memory mask.")
848 projection: ProjectionSerializationModel[P] | None = pydantic.Field(
849 default=None,
850 exclude_if=is_none,
851 description="Projection that maps the logical pixel grid onto the sky.",
852 )
854 @property
855 def bbox(self) -> Box:
856 """The 2-d bounding box of the mask."""
857 shape = self.data[0].shape
858 if len(shape) == 3:
859 shape = shape[:2]
860 return Box.from_shape(shape, start=self.start)
862 def deserialize(
863 self,
864 archive: InputArchive[Any],
865 *,
866 bbox: Box | None = None,
867 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
868 **kwargs: Any,
869 ) -> Mask:
870 """Deserialize a mask from an input archive.
872 Parameters
873 ----------
874 archive
875 Archive to read from.
876 bbox
877 Bounding box of a subimage to read instead.
878 strip_header
879 A callable that strips out any FITS header cards added by the
880 ``update_header`` argument in the corresponding call to
881 `Mask.serialize`.
882 **kwargs
883 Unsupported keyword arguments are accepted only to provide better
884 error messages (raising `serialization.InvalidParameterError`).
885 """
886 if kwargs:
887 raise InvalidParameterError(f"Unrecognized parameters for Mask: {set(kwargs.keys())}.")
888 slices: tuple[slice, ...] | EllipsisType = ...
889 if bbox is not None:
890 slices = bbox.slice_within(self.bbox)
891 else:
892 bbox = self.bbox
893 if not is_integer(self.dtype):
894 raise ArchiveReadError(f"Mask array has a non-integer dtype: {self.dtype}.")
895 schema = MaskSchema(self.planes, dtype=self.dtype.to_numpy())
896 projection = self.projection.deserialize(archive) if self.projection is not None else None
897 if len(self.data) == 1 and tuple(self.data[0].shape) == tuple(self.bbox.shape) + (schema.mask_size,):
898 storage_slices = slices if slices is ... else (slice(None),) + slices
899 array = archive.get_array(self.data[0], strip_header=strip_header, slices=storage_slices)
900 array = np.moveaxis(array, 0, -1)
901 return Mask(array, schema=schema, bbox=bbox, projection=projection)._finish_deserialize(self)
902 result = Mask(0, schema=schema, bbox=bbox, projection=projection)
903 schemas_2d = schema.split(np.int32)
904 if len(schemas_2d) != len(self.data):
905 raise ArchiveReadError(
906 f"Number of mask arrays ({len(self.data)}) does not match expectation ({len(schemas_2d)})."
907 )
908 for array_model, schema_2d in zip(self.data, schemas_2d):
909 mask_2d = self._deserialize_2d(
910 array_model, schema_2d, bbox.start, archive, strip_header=strip_header, slices=slices
911 )
912 result.update(mask_2d)
913 return result._finish_deserialize(self)
915 @staticmethod
916 def _deserialize_2d(
917 ref: ArrayReferenceModel | InlineArrayModel,
918 schema_2d: MaskSchema,
919 start: Sequence[int],
920 archive: InputArchive[Any],
921 *,
922 slices: tuple[slice, ...] | EllipsisType = ...,
923 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
924 ) -> Mask:
925 def _strip_header(header: astropy.io.fits.Header) -> None:
926 strip_header(header)
927 schema_2d.strip_header(header)
928 fits.strip_wcs_cards(header)
930 array_2d = archive.get_array(ref, strip_header=_strip_header, slices=slices)
931 return Mask(array_2d[:, :, np.newaxis], schema=schema_2d, start=start)
933 def deserialize_component(self, component: str, archive: InputArchive[Any], **kwargs: Any) -> Any:
934 if kwargs:
935 raise InvalidParameterError(f"Unsupported parameters for Mask components: {set(kwargs.keys())}.")
936 return super().deserialize_component(component, archive)
939def _archive_prefers_native_mask_arrays(archive: OutputArchive[Any]) -> bool:
940 """Return whether an archive wants masks in their native 3-D layout."""
941 current: Any = archive
942 while current is not None:
943 if getattr(current, "_prefer_native_mask_arrays", False):
944 return True
945 current = getattr(current, "_parent", None)
946 return False
949def get_legacy_visit_image_mask_planes() -> dict[str, MaskPlane]:
950 """Return a mapping from legacy mask plane name to `MaskPlane` instance
951 for LSST visit images, c. DP2.
952 """
953 return {
954 "BAD": MaskPlane("BAD", "Bad pixel in the instrument, including bad amplifiers."),
955 "SAT": MaskPlane(
956 "SATURATED", "Pixel was saturated or affected by saturation in a neighboring pixel."
957 ),
958 "INTRP": MaskPlane("INTERPOLATED", "Original pixel value was interpolated."),
959 "CR": MaskPlane("COSMIC_RAY", "A cosmic ray affected this pixel."),
960 "EDGE": MaskPlane(
961 "DETECTION_EDGE",
962 "Pixel was too close to the edge to be considered for detection, "
963 "due to the finite size of the detection kernel.",
964 ),
965 "DETECTED": MaskPlane("DETECTED", "Pixel was part of a detected source."),
966 "SUSPECT": MaskPlane("SUSPECT", "Pixel was close to the saturation level. "),
967 "NO_DATA": MaskPlane("NO_DATA", "No data was available for this pixel."),
968 "VIGNETTED": MaskPlane("VIGNETTED", "Pixel was vignetted by the optics."),
969 "PARTLY_VIGNETTED": MaskPlane("PARTLY_VIGNETTED", "Pixel was partly vignetted by the optics."),
970 "CROSSTALK": MaskPlane("CROSSTALK", "Pixel was affected by crosstalk and corrected accordingly."),
971 "ITL_DIP": MaskPlane(
972 "ITL_DIP", "Pixel was affected by a dark vertical trail from a bright source, on an ITL CCD."
973 ),
974 "NOT_DEBLENDED": MaskPlane(
975 "NOT_DEBLENDED",
976 "Pixel belonged to a detection that was not deblended, usually due to size limits.",
977 ),
978 "SPIKE": MaskPlane(
979 "SPIKE", "Pixel is in the neighborhood of a diffraction spike from a bright star."
980 ),
981 "UNMASKEDNAN": MaskPlane("UNMASKED_NAN", "Pixel was found to be NaN unexpectedly."),
982 }
985def get_legacy_difference_image_mask_planes() -> dict[str, MaskPlane]:
986 """Return a mapping from legacy mask plane name to `MaskPlane` instance
987 for LSST difference images, c. DP2.
988 """
989 result = get_legacy_visit_image_mask_planes()
990 result["DETECTED_NEGATIVE"] = MaskPlane(
991 "DETECTED_NEGATIVE", "Pixel was part of a detected source with negative flux."
992 )
993 result["SAT_TEMPLATE"] = MaskPlane("SAT_TEMPLATE", "Template pixel was saturated.")
994 result["HIGH_VARIANCE"] = MaskPlane("HIGH_VARIANCE", "TODO[DM-55036]")
995 result["STREAK"] = MaskPlane(
996 "STREAK", "An extended streak (probably an artificial satellite) affected this pixel."
997 )
998 return result
1001def get_legacy_deep_coadd_mask_planes() -> dict[str, MaskPlane]:
1002 """Return a mapping from legacy mask plane name to `MaskPlane` instance
1003 for LSST deep coadds, c. DP2.
1004 """
1005 return {
1006 # TODO: reconcile this with counts from the DP2 coadds.
1007 # BAD, CLIPPED, SUSPECT, PARTLY_VIGNETTED, SPIKE: should be fully
1008 # rejected from (cell) coadds with no propagation.
1009 "NO_DATA": MaskPlane("NO_DATA", "No data was available for this pixel."),
1010 "INTRP": MaskPlane("INTERPOLATED", "Pixel value is the result of interpolating nearby good pixels."),
1011 "CR": MaskPlane(
1012 "COSMIC_RAY",
1013 "A cosmic ray affected this pixel on at least one input image (and was interpolated).",
1014 ),
1015 "SAT": MaskPlane("SATURATED", "More than 10% of the potential input visits."),
1016 "EDGE": MaskPlane(
1017 "DETECTION_EDGE",
1018 "Pixel was too close to the edge to be considered for detection, "
1019 "due to the finite size of the detection kernel.",
1020 ),
1021 "REJECTED": MaskPlane(
1022 "REJECTED", "At least one input visit was left out of the coadd for this pixel due to masking."
1023 ),
1024 "DETECTED": MaskPlane("DETECTED", "Pixel was part of a detected source."),
1025 "INEXACT_PSF": MaskPlane(
1026 "INEXACT_PSF",
1027 "Pixel is on or near a cell boundary and hence its PSF may be (usually slightly) discontinuous.",
1028 ),
1029 }
1032def _guess_legacy_plane_map(old_planes: Mapping[str, int]) -> dict[str, MaskPlane]:
1033 """Guess which of the ``get_legacy_*_plane_map`` created the given mask
1034 plane dictionary and call it.
1035 """
1036 if "SAT_TEMPLATE" in old_planes:
1037 return get_legacy_difference_image_mask_planes()
1038 if "INEXACT_PSF" in old_planes:
1039 return get_legacy_deep_coadd_mask_planes()
1040 return get_legacy_visit_image_mask_planes()