Coverage for python / lsst / images / _mask.py: 23%
376 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-03 00:49 -0800
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-03 00:49 -0800
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_visit_image_mask_planes",
21)
23import dataclasses
24import math
25from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence, Set
26from types import EllipsisType
27from typing import Any, ClassVar, cast
29import astropy.io.fits
30import astropy.wcs
31import numpy as np
32import numpy.typing as npt
33import pydantic
34from astro_metadata_translator import ObservationInfo
36from lsst.resources import ResourcePath, ResourcePathExpression
38from . import fits
39from ._generalized_image import GeneralizedImage
40from ._geom import YX, Box
41from ._transforms import Frame, Projection, ProjectionSerializationModel
42from .serialization import (
43 ArchiveReadError,
44 ArchiveTree,
45 ArrayReferenceModel,
46 InputArchive,
47 IntegerType,
48 NumberType,
49 OpaqueArchiveMetadata,
50 OutputArchive,
51 is_integer,
52 no_header_updates,
53)
54from .utils import is_none
57@dataclasses.dataclass(frozen=True)
58class MaskPlane:
59 """Name and description of a single plane in a mask array."""
61 name: str
62 """Unique name for the mask plane (`str`)."""
64 description: str
65 """Human-readable documentation for the mask plane (`str`)."""
67 @classmethod
68 def read_legacy(cls, header: astropy.io.fits.Header) -> dict[str, int]:
69 """Read mask plane descriptions written by
70 `lsst.afw.image.Mask.writeFits`.
72 Parameters
73 ----------
74 header
75 FITS header.
77 Returns
78 -------
79 `dict` [`str`, `int`]
80 A dictionary mapping mask plane name to integer bit index.
81 """
82 result: dict[str, int] = {}
83 for card in list(header.cards):
84 if card.keyword.startswith("MP_"):
85 result[card.keyword.removeprefix("MP_")] = card.value
86 del header[card.keyword]
87 return result
90@dataclasses.dataclass(frozen=True)
91class MaskPlaneBit:
92 """The nested array index and mask value associated with a single mask
93 plane.
94 """
96 index: int
97 """Index into the last dimension of the mask array where this plane's bit
98 is stored.
99 """
101 mask: np.integer
102 """Bitmask that selects just this plane's bit from a mask array value
103 (`numpy.integer`).
104 """
106 @classmethod
107 def compute(cls, overall_index: int, stride: int, mask_type: type[np.integer]) -> MaskPlaneBit:
108 """Construct a `MaskPlaneBit` from the overall index of a plane in a
109 `MaskSchema` and the stride (number of bits per mask array element).
110 """
111 index, bit = divmod(overall_index, stride)
112 return cls(index, mask_type(1 << bit))
115class MaskSchema:
116 """A schema for a bit-packed mask array.
118 Parameters
119 ----------
120 planes
121 Iterable of `MaskPlane` instances that define the schema. `None`
122 values may be included to reserve bits for future use.
123 dtype
124 The numpy data type of the mask arrays that use this schema.
126 Notes
127 -----
128 A `MaskSchema` is a collection of mask planes, which each correspond to a
129 single bit in a mask array. Mask schemas are immutable and associated with
130 a particular array data type, allowing them to safely precompute the index
131 and bitmask for each plane.
133 `MaskSchema` indexing is by integer (the overall index of a plane in the
134 schema). The `descriptions` attribute may be indexed by plane name to get
135 the description for that plane, and the `bitmask` method can be used to
136 obtain an array that can be used to select one or more planes by name in
137 a mask array that uses this schema.
139 If no mask planes are provided, a `None` placeholder is automatically
140 added.
141 """
143 def __init__(self, planes: Iterable[MaskPlane | None], dtype: npt.DTypeLike = np.uint8):
144 self._planes: tuple[MaskPlane | None, ...] = tuple(planes) or (None,)
145 self._dtype = cast(np.dtype[np.integer], np.dtype(dtype))
146 stride = self.bits_per_element(self._dtype)
147 self._descriptions = {plane.name: plane.description for plane in self._planes if plane is not None}
148 self._mask_size = math.ceil(len(self._planes) / stride)
149 self._bits: dict[str, MaskPlaneBit] = {
150 plane.name: MaskPlaneBit.compute(n, stride, self._dtype.type)
151 for n, plane in enumerate(self._planes)
152 if plane is not None
153 }
155 @staticmethod
156 def bits_per_element(dtype: npt.DTypeLike) -> int:
157 """Return the number of mask bits per array element for the given
158 data type.
159 """
160 dtype = np.dtype(dtype)
161 match dtype.kind:
162 case "u":
163 return dtype.itemsize * 8
164 case "i":
165 return dtype.itemsize * 8 - 1
166 case _:
167 raise TypeError(f"dtype for masks must be an integer; got {dtype} with kind={dtype.kind}.")
169 def __iter__(self) -> Iterator[MaskPlane | None]:
170 return iter(self._planes)
172 def __len__(self) -> int:
173 return len(self._planes)
175 def __getitem__(self, i: int) -> MaskPlane | None:
176 return self._planes[i]
178 def __repr__(self) -> str:
179 return f"MaskSchema({list(self._planes)}, dtype={self._dtype!r})"
181 def __str__(self) -> str:
182 return "\n".join(
183 [
184 f"{name} [{bit.index}@{hex(bit.mask)}]: {self._descriptions[name]}"
185 for name, bit in self._bits.items()
186 ]
187 )
189 def __eq__(self, other: object) -> bool:
190 if isinstance(other, MaskSchema):
191 return self._planes == other._planes and self._dtype == other._dtype
192 return False
194 @property
195 def dtype(self) -> np.dtype:
196 """The numpy data type of the mask arrays that use this schema."""
197 return self._dtype
199 @property
200 def mask_size(self) -> int:
201 """The number of elements in the last dimension of any mask array that
202 uses this schema.
203 """
204 return self._mask_size
206 @property
207 def names(self) -> Set[str]:
208 """The names of the mask planes, in bit order."""
209 return self._bits.keys()
211 @property
212 def descriptions(self) -> Mapping[str, str]:
213 """A mapping from plane name to description."""
214 return self._descriptions
216 def bit(self, plane: str) -> MaskPlaneBit:
217 """Return the last array index and mask for the given mask plane."""
218 return self._bits[plane]
220 def bitmask(self, *planes: str) -> np.ndarray:
221 """Return a 1-d mask array that represents the union (i.e. bitwise OR)
222 of the planes with the given names.
224 Parameters
225 ----------
226 *planes
227 Mask plane names.
229 Returns
230 -------
231 numpy.ndarray
232 A 1-d array with shape ``(mask_size,)``.
233 """
234 result = np.zeros(self.mask_size, dtype=self._dtype)
235 for plane in planes:
236 bit = self._bits[plane]
237 result[bit.index] |= bit.mask
238 return result
240 def split(self, dtype: npt.DTypeLike) -> list[MaskSchema]:
241 """Split the schema into an equivalent series of schemas that each
242 have a `mask_size` of ``1``, dropping all `None` placeholders.
244 Parameters
245 ----------
246 dtype
247 Data type of the new mask pixels.
249 Returns
250 -------
251 `list` [`MaskSchema`]
252 A list of mask schemas that together include all planes in
253 ``self`` and have `mask_size` equal to ``1``. If there are no
254 mask planes (only `None` placeholders) in ``self``, a single mask
255 schema with a `None` placeholder is returned; otherwise `None`
256 placeholders are returned.
257 """
258 dtype = np.dtype(dtype)
259 planes: list[MaskPlane] = []
260 schemas: list[MaskSchema] = []
261 n_planes_per_schema = self.bits_per_element(dtype)
262 for plane in self._planes:
263 if plane is not None:
264 planes.append(plane)
265 if len(planes) == n_planes_per_schema:
266 schemas.append(MaskSchema(planes, dtype=dtype))
267 planes.clear()
268 if planes:
269 schemas.append(MaskSchema(planes, dtype=dtype))
270 if not schemas:
271 schemas.append(MaskSchema([None], dtype=dtype))
272 return schemas
274 def update_header(self, header: astropy.io.fits.Header) -> None:
275 """Add a description of this mask schema to a FITS header."""
276 for n, plane in enumerate(self):
277 if plane is not None:
278 bit = self.bit(plane.name)
279 if bit.index != 0:
280 raise TypeError("Only mask schemas with mask_size==1 can be described in FITS.")
281 header.set(f"MSKN{n + 1:04d}", plane.name, f"Name for mask plane {n + 1}.")
282 header.set(f"MSKM{n + 1:04d}", bit.mask, f"Bitmask for plane n={n + 1}; always 1<<(n-1).")
283 # We don't add a comment to the description card, because it's
284 # likely to overrun a single card and get the CONTINUE
285 # treatment . That will cause Astropy to warn about the comment
286 # being truncated and that's worse than just leaving it
287 # unexplained; it's pretty obvious from context what it is.
288 header.set(f"MSKD{n + 1:04d}", plane.description)
290 def strip_header(self, header: astropy.io.fits.Header) -> None:
291 """Remove all header cards added by `update_header`."""
292 for n, plane in enumerate(self):
293 if plane is not None:
294 header.remove(f"MSKN{n + 1:04d}", ignore_missing=True)
295 header.remove(f"MSKM{n + 1:04d}", ignore_missing=True)
296 header.remove(f"MSKD{n + 1:04d}", ignore_missing=True)
299class Mask(GeneralizedImage):
300 """A 2-d bitmask image backed by a 3-d byte array.
302 Parameters
303 ----------
304 array_or_fill
305 Array or fill value for the mask. If a fill value, ``bbox`` or
306 ``shape`` must be provided.
307 schema
308 Schema that defines the planes and their bit assignments.
309 bbox
310 Bounding box for the mask. This sets the shape of the first two
311 dimensions of the array.
312 start
313 Logical coordinates of the first pixel in the array, ordered ``y``,
314 ``x`` (unless an `XY` instance is passed). Ignored if
315 ``bbox`` is provided. Defaults to zeros.
316 shape
317 Leading dimensions of the array, ordered ``y``, ``x`` (unless an `XY`
318 instance is passed). Only needed if ``array_or_fill`` is not an
319 array and ``bbox`` is not provided. Like the bbox, this does not
320 include the last dimension of the array.
321 projection
322 Projection that maps the pixel grid to the sky.
323 obs_info
324 General information about the associated observation in standardized
325 form.
327 Notes
328 -----
329 Indexing the `array` attribute of a `Mask` does not take into account its
330 ``start`` offset, but accessing a subimage mask by indexing a `Mask` with
331 a `Box` does, and the `bbox` of the subimage is set to match its location
332 within the original mask.
334 A mask's ``bbox`` corresponds to the leading dimensions of its backing
335 `numpy.ndarray`, while the last dimension's size is always equal to the
336 `~MaskSchema.mask_size` of its schema, since a schema can in general
337 require multiple array elements to represent all of its planes.
338 """
340 def __init__(
341 self,
342 array_or_fill: np.ndarray | int = 0,
343 /,
344 *,
345 schema: MaskSchema,
346 bbox: Box | None = None,
347 start: Sequence[int] | None = None,
348 shape: Sequence[int] | None = None,
349 projection: Projection | None = None,
350 obs_info: ObservationInfo | None = None,
351 ):
352 if shape is not None:
353 shape = tuple(shape)
354 if start is not None:
355 start = tuple(start)
356 if isinstance(array_or_fill, np.ndarray):
357 array = np.array(array_or_fill, dtype=schema.dtype)
358 if array.ndim != 3:
359 raise ValueError("Mask array must be 3-d.")
360 if bbox is None:
361 bbox = Box.from_shape(array.shape[:-1], start=start)
362 elif bbox.shape + (schema.mask_size,) != array.shape:
363 raise ValueError(
364 f"Explicit bbox shape {bbox.shape} and schema of size {schema.mask_size} do not "
365 f"match array with shape {array.shape}."
366 )
367 if shape is not None and shape + (schema.mask_size,) != array.shape:
368 raise ValueError(
369 f"Explicit shape {shape} and schema of size {schema.mask_size} do "
370 f"not match array with shape {array.shape}."
371 )
373 else:
374 if bbox is None:
375 if shape is None:
376 raise TypeError("No bbox, size, or array provided.")
377 bbox = Box.from_shape(shape, start=start)
378 array = np.full(bbox.shape + (schema.mask_size,), array_or_fill, dtype=schema.dtype)
379 self._array = array
380 self._bbox: Box = bbox
381 self._schema: MaskSchema = schema
382 self._projection = projection
383 self._opaque_metadata: OpaqueArchiveMetadata | None = None
384 self._obs_info = obs_info
386 @property
387 def array(self) -> np.ndarray:
388 """The low-level array (`numpy.ndarray`).
390 Assigning to this attribute modifies the existing array in place; the
391 bounding box and underlying data pointer are never changed.
392 """
393 return self._array
395 @array.setter
396 def array(self, value: np.ndarray | int) -> None:
397 self._array[:, :] = value
399 @property
400 def schema(self) -> MaskSchema:
401 """Schema that defines the planes and their bit assignments
402 (`MaskSchema`).
403 """
404 return self._schema
406 @property
407 def bbox(self) -> Box:
408 """2-d bounding box of the mask (`Box`).
410 This sets the shape of the first two dimensions of the array.
411 """
412 return self._bbox
414 @property
415 def projection(self) -> Projection[Any] | None:
416 """The projection that maps this mask's pixel grid to the sky
417 (`Projection` | `None`).
419 Notes
420 -----
421 The pixel coordinates used by this projection account for the bounding
422 box ``start``; they are not just array indices.
423 """
424 return self._projection
426 @property
427 def obs_info(self) -> ObservationInfo | None:
428 """General information about the associated observation in standard
429 form. (`~astro_metadata_translator.ObservationInfo` | `None`).
430 """
431 return self._obs_info
433 def __getitem__(self, bbox: Box | EllipsisType) -> Mask:
434 super().__getitem__(bbox)
435 if bbox is ...:
436 return self
437 result = Mask(
438 self.array[bbox.y.slice_within(self._bbox.y), bbox.x.slice_within(self._bbox.x), :],
439 bbox=bbox,
440 schema=self.schema,
441 )
442 if self._opaque_metadata is not None:
443 result._opaque_metadata = self._opaque_metadata.subset(bbox)
444 return result
446 def __setitem__(self, bbox: Box | EllipsisType, value: Mask) -> None:
447 subview = self[bbox]
448 subview.clear()
449 subview.update(value)
451 def __str__(self) -> str:
452 return f"Mask({self.bbox!s}, {list(self.schema.names)})"
454 def __repr__(self) -> str:
455 return f"Mask(..., bbox={self.bbox!r}, schema={self.schema!r})"
457 def __eq__(self, other: object) -> bool:
458 if not isinstance(other, Mask):
459 return NotImplemented
460 return (
461 self._bbox == other._bbox
462 and self._schema == other._schema
463 and np.array_equal(self._array, other._array, equal_nan=True)
464 )
466 def copy(self) -> Mask:
467 """Deep-copy the mask."""
468 result = Mask(
469 self._array.copy(),
470 bbox=self._bbox,
471 schema=self._schema,
472 projection=self._projection,
473 obs_info=self._obs_info,
474 )
475 if self._opaque_metadata is not None:
476 result._opaque_metadata = self._opaque_metadata.copy()
477 return result
479 def view(
480 self,
481 *,
482 schema: MaskSchema | EllipsisType = ...,
483 projection: Projection | None | EllipsisType = ...,
484 start: Sequence[int] | EllipsisType = ...,
485 obs_info: ObservationInfo | None | EllipsisType = ...,
486 ) -> Mask:
487 """Make a view of the mask, with optional updates.
489 Notes
490 -----
491 This can only be used to make changes to schema descriptions; plane
492 names must remain the same (in the same order).
493 """
494 if schema is ...:
495 schema = self._schema
496 else:
497 if list(schema.names) != list(self.schema.names):
498 raise ValueError("Cannot create a mask view with a schema with different names.")
499 if projection is ...:
500 projection = self._projection
501 if start is ...:
502 start = self._bbox.start
503 if obs_info is ...:
504 obs_info = self._obs_info
505 result = Mask(self._array, start=start, schema=schema, projection=projection, obs_info=obs_info)
506 if self._opaque_metadata is not None:
507 result._opaque_metadata = self._opaque_metadata.copy()
508 return result
510 def update(self, other: Mask) -> None:
511 """Update ``self`` to include all common mask values set in ``other``.
513 Notes
514 -----
515 This only operates on the intersection of the two mask bounding boxes
516 and the mask planes that are present in both. Mask bits are only set,
517 not cleared (i.e. this uses ``|=`` updates, not ``=`` assignments).
518 """
519 lhs = self
520 rhs = other
521 if other.bbox != self.bbox:
522 if (bbox := self.bbox.intersection(other.bbox)) is None:
523 return
524 lhs = self[bbox]
525 rhs = other[bbox]
526 for name in self.schema.names & other.schema.names:
527 lhs.set(name, rhs.get(name))
529 def get(self, plane: str) -> np.ndarray:
530 """Return a 2-d boolean array for the given mask plane.
532 Parameters
533 ----------
534 plane
535 Name of the mask plane.
537 Returns
538 -------
539 numpy.ndarray
540 A 2-d boolean array with the same shape as `bbox` that is `True`
541 where the bit for ``plane`` is set and `False` elsewhere.
542 """
543 bit = self.schema.bit(plane)
544 return (self._array[..., bit.index] & bit.mask).astype(bool)
546 def set(self, plane: str, boolean_mask: np.ndarray | EllipsisType = ...) -> None:
547 """Set a mask plane.
549 Parameters
550 ----------
551 plane
552 Name of the mask plane to set
553 boolean_mask
554 A 2-d boolean array with the same shape as `bbox` that is `True`
555 where the bit for ``plane`` should be set and `False` where it
556 should be left unchanged (*not* set to zero). May be ``...`` to
557 set the bit everywhere.
558 """
559 bit = self.schema.bit(plane)
560 if boolean_mask is not ...:
561 boolean_mask = boolean_mask.astype(bool)
562 self._array[boolean_mask, bit.index] |= bit.mask
564 def clear(self, plane: str | None = None, boolean_mask: np.ndarray | EllipsisType = ...) -> None:
565 """Clear one or more mask planes.
567 Parameters
568 ----------
569 plane
570 Name of the mask plane to set. If `None` all mask planes are
571 cleared.
572 boolean_mask
573 A 2-d boolean array with the same shape as `bbox` that is `True`
574 where the bit for ``plane`` should be cleared and `False` where it
575 should be left unchanged. May be ``...`` to clear the bit
576 everywhere.
577 """
578 if boolean_mask is not ...:
579 boolean_mask = boolean_mask.astype(bool)
580 if plane is None:
581 self._array[boolean_mask, :] = 0
582 else:
583 bit = self.schema.bit(plane)
584 self._array[boolean_mask, bit.index] &= ~bit.mask
586 def serialize[P: pydantic.BaseModel](
587 self,
588 archive: OutputArchive[P],
589 *,
590 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
591 save_projection: bool = True,
592 add_offset_wcs: str | None = "A",
593 ) -> MaskSerializationModel[P]:
594 """Serialize the mask to an output archive.
596 Parameters
597 ----------
598 archive
599 Archive to write to.
600 update_header
601 A callback that will be given the FITS header for the HDU
602 containing this mask in order to add keys to it. This callback
603 may be provided but will not be called if the output format is not
604 FITS. As multiple HDUs may be added, this function may be called
605 multiple times.
606 save_projection
607 If `True`, save the `Projection` attached to the mask, if there
608 is one.
609 add_offset_wcs
610 A FITS WCS single-character suffix to use when adding a linear
611 WCS that maps the FITS array to the logical pixel coordinates
612 defined by ``bbox.start``. Set to `None` to not write this WCS.
613 """
614 data: list[ArrayReferenceModel] = []
615 for schema_2d in self.schema.split(np.int32):
616 mask_2d = Mask(
617 0, bbox=self.bbox, schema=schema_2d, projection=self._projection, obs_info=self._obs_info
618 )
619 mask_2d.update(self)
620 data.append(mask_2d._serialize_2d(archive, update_header=update_header))
621 serialized_projection: ProjectionSerializationModel[P] | None = None
622 if save_projection and self.projection is not None:
623 serialized_projection = archive.serialize_direct("projection", self.projection.serialize)
624 serialized_dtype = NumberType.from_numpy(self.schema.dtype)
625 assert is_integer(serialized_dtype), "Mask dtypes should always be integers."
626 return MaskSerializationModel.model_construct(
627 data=data,
628 start=list(self.bbox.start),
629 planes=list(self.schema),
630 dtype=serialized_dtype,
631 projection=serialized_projection,
632 obs_info=self._obs_info,
633 )
635 def _serialize_2d[P: pydantic.BaseModel](
636 self,
637 archive: OutputArchive[P],
638 *,
639 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
640 save_projection: bool = True,
641 add_offset_wcs: str | None = "A",
642 ) -> ArrayReferenceModel:
643 def _update_header(header: astropy.io.fits.Header) -> None:
644 update_header(header)
645 self.schema.update_header(header)
646 if self.projection is not None:
647 if self.fits_wcs:
648 header.update(self.fits_wcs.to_header(relax=True))
649 if add_offset_wcs is not None:
650 fits.add_offset_wcs(header, x=self.bbox.x.start, y=self.bbox.y.start, key=add_offset_wcs)
652 assert self.array.shape[2] == 1, "Mask should be split before calling this method."
653 return archive.add_array(self._array[:, :, 0], update_header=_update_header)
655 @classmethod
656 def deserialize(
657 cls,
658 model: MaskSerializationModel[Any],
659 archive: InputArchive[Any],
660 *,
661 bbox: Box | None = None,
662 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
663 ) -> Mask:
664 """Deserialize a mask from an input archive.
666 Parameters
667 ----------
668 model
669 A Pydantic model representation of the mask, holding references
670 to data stored in the archive.
671 archive
672 Archive to read from.
673 bbox
674 Bounding box of a subimage to read instead.
675 strip_header
676 A callable that strips out any FITS header cards added by the
677 ``update_header`` argument in the corresponding call to
678 `serialize`.
679 """
680 slices: tuple[slice, ...] | EllipsisType = ...
681 if bbox is not None:
682 slices = bbox.slice_within(model.bbox)
683 else:
684 bbox = model.bbox
685 if not is_integer(model.dtype):
686 raise ArchiveReadError(f"Mask array has a non-integer dtype: {model.dtype}.")
687 schema = MaskSchema(model.planes, dtype=model.dtype.to_numpy())
688 projection = (
689 Projection.deserialize(model.projection, archive) if model.projection is not None else None
690 )
691 result = Mask(0, schema=schema, bbox=bbox, projection=projection, obs_info=model.obs_info)
692 schemas_2d = schema.split(np.int32)
693 if len(schemas_2d) != len(model.data):
694 raise ArchiveReadError(
695 f"Number of mask arrays ({len(model.data)}) does not match expectation ({len(schemas_2d)})."
696 )
697 for ref, schema_2d in zip(model.data, schemas_2d):
698 mask_2d = cls._deserialize_2d(
699 ref, schema_2d, bbox.start, archive, strip_header=strip_header, slices=slices
700 )
701 result.update(mask_2d)
702 return result
704 @classmethod
705 def _deserialize_2d(
706 cls,
707 ref: ArrayReferenceModel,
708 schema_2d: MaskSchema,
709 start: Sequence[int],
710 archive: InputArchive[Any],
711 *,
712 slices: tuple[slice, ...] | EllipsisType = ...,
713 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
714 ) -> Mask:
715 def _strip_header(header: astropy.io.fits.Header) -> None:
716 strip_header(header)
717 schema_2d.strip_header(header)
718 fits.strip_wcs_cards(header)
720 array_2d = archive.get_array(ref, strip_header=_strip_header, slices=slices)
721 return Mask(array_2d[:, :, np.newaxis], schema=schema_2d, start=start)
723 @staticmethod
724 def _get_archive_tree_type[P: pydantic.BaseModel](
725 pointer_type: type[P],
726 ) -> type[MaskSerializationModel[P]]:
727 """Return the serialization model type for this object for an archive
728 type that uses the given pointer type.
729 """
730 return MaskSerializationModel[pointer_type] # type: ignore
732 _archive_default_name: ClassVar[str] = "mask"
733 """The name this object should be serialized with when written as the
734 top-level object.
735 """
737 def write_fits(
738 self,
739 filename: str,
740 *,
741 compression: fits.FitsCompressionOptions | None = fits.FitsCompressionOptions.DEFAULT,
742 ) -> None:
743 """Write the mask to a FITS file.
745 Parameters
746 ----------
747 filename
748 Name of the file to write to. Must be a local file.
749 compression
750 Compression options.
751 """
752 compression_options = {}
753 if compression is not fits.FitsCompressionOptions.DEFAULT:
754 compression_options[self._archive_default_name] = compression
755 fits.write(self, filename, compression_options)
757 @staticmethod
758 def read_fits(url: ResourcePathExpression, *, bbox: Box | None = None) -> Mask:
759 """Read an image from a FITS file.
761 Parameters
762 ----------
763 url
764 URL of the file to read; may be any type supported by
765 `lsst.resources.ResourcePath`.
766 bbox
767 Bounding box of a subimage to read instead.
768 """
769 return fits.read(Mask, url, bbox=bbox)
771 @staticmethod
772 def from_legacy(
773 legacy: Any,
774 plane_map: Mapping[str, MaskPlane] | None = None,
775 ) -> Mask:
776 """Convert from an `lsst.afw.image.Mask` instance.
778 Parameters
779 ----------
780 legacy
781 An `lsst.afw.image.Mask` instance. This will not share pixel
782 data with the new object.
783 plane_map
784 A mapping from legacy mask plane name to the new plane name and
785 description.
786 """
787 return Mask._from_legacy_array(
788 legacy.array,
789 legacy.getMaskPlaneDict(),
790 start=YX(y=legacy.getY0(), x=legacy.getX0()),
791 plane_map=plane_map,
792 )
794 def to_legacy(self, plane_map: Mapping[str, MaskPlane] | None = None) -> Any:
795 """Convert to an `lsst.afw.image.Mask` instance.
797 The pixel data will not be shared between the two objects.
799 Parameters
800 ----------
801 plane_map
802 A mapping from legacy mask plane name to the new plane name and
803 description.
804 """
805 import lsst.afw.image
806 import lsst.geom
808 result = lsst.afw.image.Mask(self.bbox.to_legacy())
809 if plane_map is None:
810 plane_map = {plane.name: plane for plane in self.schema if plane is not None}
811 for old_name, new_plane in plane_map.items():
812 old_bit = result.addMaskPlane(old_name)
813 old_bitmask = 1 << old_bit
814 result.array[self.get(new_plane.name)] |= old_bitmask
815 return result
817 @staticmethod
818 def _from_legacy_array(
819 array2d: np.ndarray,
820 old_planes: Mapping[str, int],
821 *,
822 start: YX[int],
823 plane_map: Mapping[str, MaskPlane] | None = None,
824 projection: Projection | None = None,
825 ) -> Mask:
826 planes: list[MaskPlane] = []
827 new_name_to_old_bitmask: dict[str, int] = {}
828 for old_name, old_bit in old_planes.items():
829 old_bitmask = 1 << old_bit
830 if plane_map is not None:
831 if new_plane := plane_map.get(old_name):
832 planes.append(new_plane)
833 new_name_to_old_bitmask[new_plane.name] = old_bitmask
834 else:
835 if n_orphaned := np.count_nonzero(array2d & old_bitmask):
836 raise RuntimeError(
837 f"Legacy mask plane {old_name!r} is not remapped, "
838 f"but {n_orphaned} pixels have this bit set."
839 )
840 else:
841 planes.append(MaskPlane(old_name, ""))
842 new_name_to_old_bitmask[old_name] = old_bitmask
843 schema = MaskSchema(planes)
844 mask = Mask(0, schema=schema, start=start, shape=array2d.shape, projection=projection)
845 for new_name, old_bitmask in new_name_to_old_bitmask.items():
846 mask.set(new_name, array2d & old_bitmask)
847 return mask
849 @staticmethod
850 def read_legacy(
851 uri: ResourcePathExpression,
852 *,
853 plane_map: Mapping[str, MaskPlane] | None = None,
854 ext: str | int = 1,
855 fits_wcs_frame: Frame | None = None,
856 ) -> Mask:
857 """Read a FITS file written by `lsst.afw.image.Mask.writeFits`.
859 Parameters
860 ----------
861 uri
862 URI or file name.
863 plane_map
864 A mapping from legacy mask plane name to the new plane name and
865 description.
866 ext
867 Name or index of the FITS HDU to read.
868 fits_wcs_frame
869 If not `None` and the HDU containing the mask has a FITS WCS,
870 attach a `Projection` to the returned mask by converting that WCS.
871 """
872 opaque_metadata = fits.FitsOpaqueMetadata()
873 fs, fspath = ResourcePath(uri).to_fsspec()
874 with fs.open(fspath) as stream, astropy.io.fits.open(stream) as hdu_list:
875 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header)
876 result = Mask._read_legacy_hdu(
877 hdu_list[ext], opaque_metadata, plane_map=plane_map, fits_wcs_frame=fits_wcs_frame
878 )
879 result._opaque_metadata = opaque_metadata
880 return result
882 @staticmethod
883 def _read_legacy_hdu(
884 hdu: astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU | astropy.io.fits.BinTableHDU,
885 opaque_metadata: fits.FitsOpaqueMetadata,
886 plane_map: Mapping[str, MaskPlane] | None = None,
887 fits_wcs_frame: Frame | None = None,
888 ) -> Mask:
889 if isinstance(hdu, astropy.io.fits.BinTableHDU):
890 hdu = astropy.io.fits.CompImageHDU(bintable=hdu)
891 dx: int = hdu.header.pop("LTV1")
892 dy: int = hdu.header.pop("LTV2")
893 start = YX(y=-dy, x=-dx)
894 old_planes = MaskPlane.read_legacy(hdu.header)
895 projection: Projection | None = None
896 if fits_wcs_frame is not None:
897 try:
898 fits_wcs = astropy.wcs.WCS(hdu.header)
899 except KeyError:
900 pass
901 else:
902 projection = Projection.from_fits_wcs(
903 fits_wcs, pixel_frame=fits_wcs_frame, x0=start.x, y0=start.y
904 )
905 mask = Mask._from_legacy_array(
906 hdu.data, old_planes, start=start, plane_map=plane_map, projection=projection
907 )
908 fits.strip_wcs_cards(hdu.header)
909 hdu.header.strip()
910 hdu.header.remove("EXTTYPE", ignore_missing=True)
911 hdu.header.remove("INHERIT", ignore_missing=True)
912 # afw set BUNIT on masks because of limitations in how FITS
913 # metadata is handled there.
914 hdu.header.remove("BUNIT", ignore_missing=True)
915 opaque_metadata.add_header(hdu.header)
916 return mask
919class MaskSerializationModel[P: pydantic.BaseModel](ArchiveTree):
920 """Pydantic model used to represent the serialized form of a `.Mask`."""
922 data: list[ArrayReferenceModel] = pydantic.Field(description="References to pixel data.")
923 start: list[int] = pydantic.Field(
924 description="Coordinate of the first pixels in the array, ordered (y, x)."
925 )
926 planes: list[MaskPlane | None] = pydantic.Field(description="Definitions of the bitplanes in the mask.")
927 dtype: IntegerType = pydantic.Field(description="Data type of the in-memory mask.")
928 projection: ProjectionSerializationModel[P] | None = pydantic.Field(
929 default=None,
930 exclude_if=is_none,
931 description="Projection that maps the logical pixel grid onto the sky.",
932 )
933 obs_info: ObservationInfo | None = pydantic.Field(
934 default=None,
935 exclude_if=is_none,
936 description="Standardized description of image metadata",
937 )
939 @property
940 def bbox(self) -> Box:
941 """The 2-d bounding box of the mask."""
942 return Box.from_shape(self.data[0].shape, start=self.start)
945def get_legacy_visit_image_mask_planes() -> dict[str, MaskPlane]:
946 """Return a mapping from legacy mask plane name to `MaskPlane` instance
947 for LSST visit images, c. DP2.
948 """
949 return {
950 "BAD": MaskPlane("BAD", "Bad pixel in the instrument, including bad amplifiers."),
951 "SAT": MaskPlane(
952 "SATURATED", "Pixel was saturated or affected by saturation in a neighboring pixel."
953 ),
954 "INTRP": MaskPlane("INTERPOLATED", "Original pixel value was interpolated."),
955 "CR": MaskPlane("COSMIC_RAY", "A cosmic ray affected this pixel."),
956 "EDGE": MaskPlane(
957 "DETECTION_EDGE",
958 "Pixel was too close to the edge to be considered for detection, "
959 "due to the finite size of the detection kernel.",
960 ),
961 "DETECTED": MaskPlane("DETECTED", "Pixel was part of a detected source."),
962 "SUSPECT": MaskPlane("SUSPECT", "Pixel was close to the saturation level. "),
963 "NO_DATA": MaskPlane("NO_DATA", "No data was available for this pixel."),
964 "VIGNETTED": MaskPlane("VIGNETTED", "Pixel was vignetted by the optics."),
965 "PARTLY_VIGNETTED": MaskPlane("PARTLY_VIGNETTED", "Pixel was partly vignetted by the optics."),
966 "CROSSTALK": MaskPlane("CROSSTALK", "Pixel was affected by crosstalk and corrected accordingly."),
967 "ITL_DIP": MaskPlane(
968 "ITL_DIP", "Pixel was affected by a dark vertical trail from a bright source, on an ITL CCD."
969 ),
970 "NOT_DEBLENDED": MaskPlane(
971 "NOT_DEBLENDED",
972 "Pixel belonged to a detection that was not deblended, usually due to size limits.",
973 ),
974 "SPIKE": MaskPlane(
975 "SPIKE", "Pixel is in the neighborhood of a diffraction spike from a bright star."
976 ),
977 }