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

376 statements  

« prev     ^ index     » next       coverage.py v7.13.4, created at 2026-03-04 01:58 -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. 

11 

12from __future__ import annotations 

13 

14__all__ = ( 

15 "Mask", 

16 "MaskPlane", 

17 "MaskPlaneBit", 

18 "MaskSchema", 

19 "MaskSerializationModel", 

20 "get_legacy_visit_image_mask_planes", 

21) 

22 

23import dataclasses 

24import math 

25from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence, Set 

26from types import EllipsisType 

27from typing import Any, ClassVar, cast 

28 

29import astropy.io.fits 

30import astropy.wcs 

31import numpy as np 

32import numpy.typing as npt 

33import pydantic 

34from astro_metadata_translator import ObservationInfo 

35 

36from lsst.resources import ResourcePath, ResourcePathExpression 

37 

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 

55 

56 

57@dataclasses.dataclass(frozen=True) 

58class MaskPlane: 

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

60 

61 name: str 

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

63 

64 description: str 

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

66 

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

71 

72 Parameters 

73 ---------- 

74 header 

75 FITS header. 

76 

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 

88 

89 

90@dataclasses.dataclass(frozen=True) 

91class MaskPlaneBit: 

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

93 plane. 

94 """ 

95 

96 index: int 

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

98 is stored. 

99 """ 

100 

101 mask: np.integer 

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

103 (`numpy.integer`). 

104 """ 

105 

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

113 

114 

115class MaskSchema: 

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

117 

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. 

125 

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. 

132 

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. 

138 

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

140 added. 

141 """ 

142 

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 } 

154 

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

168 

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

170 return iter(self._planes) 

171 

172 def __len__(self) -> int: 

173 return len(self._planes) 

174 

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

176 return self._planes[i] 

177 

178 def __repr__(self) -> str: 

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

180 

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 ) 

188 

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 

193 

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 

198 

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 

205 

206 @property 

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

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

209 return self._bits.keys() 

210 

211 @property 

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

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

214 return self._descriptions 

215 

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] 

219 

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. 

223 

224 Parameters 

225 ---------- 

226 *planes 

227 Mask plane names. 

228 

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 

239 

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. 

243 

244 Parameters 

245 ---------- 

246 dtype 

247 Data type of the new mask pixels. 

248 

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 

273 

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) 

289 

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) 

297 

298 

299class Mask(GeneralizedImage): 

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

301 

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. 

326 

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. 

333 

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

339 

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 ) 

372 

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 

385 

386 @property 

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

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

389 

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 

394 

395 @array.setter 

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

397 self._array[:, :] = value 

398 

399 @property 

400 def schema(self) -> MaskSchema: 

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

402 (`MaskSchema`). 

403 """ 

404 return self._schema 

405 

406 @property 

407 def bbox(self) -> Box: 

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

409 

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

411 """ 

412 return self._bbox 

413 

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

418 

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 

425 

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 

432 

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 

445 

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

447 subview = self[bbox] 

448 subview.clear() 

449 subview.update(value) 

450 

451 def __str__(self) -> str: 

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

453 

454 def __repr__(self) -> str: 

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

456 

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 ) 

465 

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 

478 

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. 

488 

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 

509 

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

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

512 

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

528 

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

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

531 

532 Parameters 

533 ---------- 

534 plane 

535 Name of the mask plane. 

536 

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) 

545 

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

547 """Set a mask plane. 

548 

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 

563 

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

565 """Clear one or more mask planes. 

566 

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 

585 

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. 

595 

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 ) 

634 

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) 

651 

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) 

654 

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. 

665 

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 

703 

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) 

719 

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) 

722 

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 

731 

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

736 

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. 

744 

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) 

756 

757 @staticmethod 

758 def read_fits(url: ResourcePathExpression, *, bbox: Box | None = None) -> Mask: 

759 """Read an image from a FITS file. 

760 

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) 

770 

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. 

777 

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 ) 

793 

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

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

796 

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

798 

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 

807 

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 

816 

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 

848 

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

858 

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 

881 

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 

917 

918 

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

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

921 

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 ) 

938 

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) 

943 

944 

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 }