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