Coverage for python/lsst/images/_mask.py: 22%

402 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-08 08:45 +0000

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. 

11 

12from __future__ import annotations 

13 

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) 

24 

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 

30 

31import astropy.io.fits 

32import astropy.wcs 

33import numpy as np 

34import numpy.typing as npt 

35import pydantic 

36 

37from lsst.resources import ResourcePath, ResourcePathExpression 

38 

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 

58 

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] 

64 

65 

66@dataclasses.dataclass(frozen=True) 

67class MaskPlane: 

68 """Name and description of a single plane in a mask array.""" 

69 

70 name: str 

71 """Unique name for the mask plane (`str`).""" 

72 

73 description: str 

74 """Human-readable documentation for the mask plane (`str`).""" 

75 

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`. 

80 

81 Parameters 

82 ---------- 

83 header 

84 FITS header. 

85 

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 

97 

98 

99@dataclasses.dataclass(frozen=True) 

100class MaskPlaneBit: 

101 """The nested array index and mask value associated with a single mask 

102 plane. 

103 """ 

104 

105 index: int 

106 """Index into the last dimension of the mask array where this plane's bit 

107 is stored. 

108 """ 

109 

110 mask: np.integer 

111 """Bitmask that selects just this plane's bit from a mask array value 

112 (`numpy.integer`). 

113 """ 

114 

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)) 

122 

123 

124class MaskSchema: 

125 """A schema for a bit-packed mask array. 

126 

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. 

134 

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. 

141 

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. 

147 

148 If no mask planes are provided, a `None` placeholder is automatically 

149 added. 

150 """ 

151 

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 } 

163 

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}.") 

177 

178 def __iter__(self) -> Iterator[MaskPlane | None]: 

179 return iter(self._planes) 

180 

181 def __len__(self) -> int: 

182 return len(self._planes) 

183 

184 def __contains__(self, plane: str | MaskPlane) -> bool: 

185 return getattr(plane, "name", plane) in self.names 

186 

187 def __getitem__(self, i: int) -> MaskPlane | None: 

188 return self._planes[i] 

189 

190 def __repr__(self) -> str: 

191 return f"MaskSchema({list(self._planes)}, dtype={self._dtype!r})" 

192 

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 ) 

200 

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 

205 

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 

210 

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 

217 

218 @property 

219 def names(self) -> Set[str]: 

220 """The names of the mask planes, in bit order.""" 

221 return self._bits.keys() 

222 

223 @property 

224 def descriptions(self) -> Mapping[str, str]: 

225 """A mapping from plane name to description.""" 

226 return self._descriptions 

227 

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] 

231 

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. 

235 

236 Parameters 

237 ---------- 

238 *planes 

239 Mask plane names. 

240 

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 

251 

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. 

255 

256 Parameters 

257 ---------- 

258 dtype 

259 Data type of the new mask pixels. 

260 

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 

285 

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) 

301 

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) 

309 

310 

311class Mask(GeneralizedImage): 

312 """A 2-d bitmask image backed by a 3-d byte array. 

313 

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. 

337 

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. 

344 

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 """ 

350 

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 ) 

384 

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 

395 

396 @property 

397 def array(self) -> np.ndarray: 

398 """The low-level array (`numpy.ndarray`). 

399 

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 

404 

405 @array.setter 

406 def array(self, value: np.ndarray | int) -> None: 

407 self._array[:, :] = value 

408 

409 @property 

410 def schema(self) -> MaskSchema: 

411 """Schema that defines the planes and their bit assignments 

412 (`MaskSchema`). 

413 """ 

414 return self._schema 

415 

416 @property 

417 def bbox(self) -> Box: 

418 """2-d bounding box of the mask (`Box`). 

419 

420 This sets the shape of the first two dimensions of the array. 

421 """ 

422 return self._bbox 

423 

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`). 

428 

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 

435 

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 ) 

449 

450 def __setitem__(self, bbox: Box | EllipsisType, value: Mask) -> None: 

451 subview = self[bbox] 

452 subview.clear() 

453 subview.update(value) 

454 

455 def __str__(self) -> str: 

456 return f"Mask({self.bbox!s}, {list(self.schema.names)})" 

457 

458 def __repr__(self) -> str: 

459 return f"Mask(..., bbox={self.bbox!r}, schema={self.schema!r})" 

460 

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 ) 

469 

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 ) 

476 

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. 

485 

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)) 

501 

502 def update(self, other: Mask) -> None: 

503 """Update ``self`` to include all common mask values set in ``other``. 

504 

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)) 

522 

523 def get(self, plane: str) -> np.ndarray: 

524 """Return a 2-d boolean array for the given mask plane. 

525 

526 Parameters 

527 ---------- 

528 plane 

529 Name of the mask plane. 

530 

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) 

539 

540 def set(self, plane: str, boolean_mask: np.ndarray | EllipsisType = ...) -> None: 

541 """Set a mask plane. 

542 

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 

557 

558 def clear(self, plane: str | None = None, boolean_mask: np.ndarray | EllipsisType = ...) -> None: 

559 """Clear one or more mask planes. 

560 

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 

579 

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. 

589 

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 ) 

643 

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) 

659 

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) 

662 

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 

671 

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 """ 

676 

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. 

683 

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 ) 

701 

702 def to_legacy(self, plane_map: Mapping[str, MaskPlane] | None = None) -> Any: 

703 """Convert to an `lsst.afw.image.Mask` instance. 

704 

705 The pixel data will not be shared between the two objects. 

706 

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 

715 

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 

725 

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 

759 

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`. 

769 

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 

794 

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 

830 

831 

832class MaskSerializationModel[P: pydantic.BaseModel](ArchiveTree): 

833 """Pydantic model used to represent the serialized form of a `.Mask`.""" 

834 

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 

839 

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 ) 

853 

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) 

861 

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. 

871 

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) 

914 

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) 

929 

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) 

932 

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) 

937 

938 

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 

947 

948 

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 } 

983 

984 

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 

999 

1000 

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 } 

1030 

1031 

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()