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

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 functools import cached_property 

27from types import EllipsisType 

28from typing import Any, ClassVar, cast 

29 

30import astropy.io.fits 

31import astropy.wcs 

32import numpy as np 

33import numpy.typing as npt 

34import pydantic 

35from astro_metadata_translator import ObservationInfo 

36 

37from lsst.resources import ResourcePath, ResourcePathExpression 

38 

39from . import fits 

40from ._geom import YX, Box 

41from ._transforms import Frame, Projection, ProjectionAstropyView, 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: 

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 @property 

434 def astropy_wcs(self) -> ProjectionAstropyView | None: 

435 """An Astropy WCS for this mask's pixel array. 

436 

437 Notes 

438 ----- 

439 As expected for Astropy WCS objects, this defines pixel coordinates 

440 such that the first row and column in `array` are ``(0, 0)``, not 

441 ``bbox.start``, as is the case for `projection`. 

442 

443 This object satisfies the `astropy.wcs.wcsapi.BaseHighLevelWCS` and 

444 `astropy.wcs.wcsapi.BaseLowLevelWCS` interfaces, but it is not an 

445 `astropy.wcs.WCS` (use `fits_wcs` for that). 

446 """ 

447 return self._projection.as_astropy(self.bbox) if self._projection is not None else None 

448 

449 @cached_property 

450 def fits_wcs(self) -> astropy.wcs.WCS | None: 

451 """An Astropy FITS WCS for this mask's pixel array. 

452 

453 Notes 

454 ----- 

455 As expected for Astropy WCS objects, this defines pixel coordinates 

456 such that the first row and column in `array` are ``(0, 0)``, not 

457 ``bbox.start``, as is the case for `projection`. 

458 

459 This may be an approximation or absent if `projection` is not 

460 naturally representable as a FITS WCS. 

461 """ 

462 return ( 

463 self._projection.as_fits_wcs(self.bbox, allow_approximation=True) 

464 if self._projection is not None 

465 else None 

466 ) 

467 

468 def __getitem__(self, bbox: Box) -> Mask: 

469 result = Mask( 

470 self.array[bbox.y.slice_within(self._bbox.y), bbox.x.slice_within(self._bbox.x), :], 

471 bbox=bbox, 

472 schema=self.schema, 

473 ) 

474 if self._opaque_metadata is not None: 

475 result._opaque_metadata = self._opaque_metadata.subset(bbox) 

476 return result 

477 

478 def __str__(self) -> str: 

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

480 

481 def __repr__(self) -> str: 

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

483 

484 def __eq__(self, other: object) -> bool: 

485 if not isinstance(other, Mask): 

486 return NotImplemented 

487 return ( 

488 self._bbox == other._bbox 

489 and self._schema == other._schema 

490 and np.array_equal(self._array, other._array, equal_nan=True) 

491 ) 

492 

493 def copy(self) -> Mask: 

494 """Deep-copy the mask.""" 

495 result = Mask( 

496 self._array.copy(), 

497 bbox=self._bbox, 

498 schema=self._schema, 

499 projection=self._projection, 

500 obs_info=self._obs_info, 

501 ) 

502 if self._opaque_metadata is not None: 

503 result._opaque_metadata = self._opaque_metadata.copy() 

504 return result 

505 

506 def view( 

507 self, 

508 *, 

509 schema: MaskSchema | EllipsisType = ..., 

510 projection: Projection | None | EllipsisType = ..., 

511 start: Sequence[int] | EllipsisType = ..., 

512 obs_info: ObservationInfo | None | EllipsisType = ..., 

513 ) -> Mask: 

514 """Make a view of the mask, with optional updates. 

515 

516 Notes 

517 ----- 

518 This can only be used to make changes to schema descriptions; plane 

519 names must remain the same (in the same order). 

520 """ 

521 if schema is ...: 

522 schema = self._schema 

523 else: 

524 if list(schema.names) != list(self.schema.names): 

525 raise ValueError("Cannot create a mask view with a schema with different names.") 

526 if projection is ...: 

527 projection = self._projection 

528 if start is ...: 

529 start = self._bbox.start 

530 if obs_info is ...: 

531 obs_info = self._obs_info 

532 result = Mask(self._array, start=start, schema=schema, projection=projection, obs_info=obs_info) 

533 if self._opaque_metadata is not None: 

534 result._opaque_metadata = self._opaque_metadata.copy() 

535 return result 

536 

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

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

539 

540 Notes 

541 ----- 

542 This only operates on the intersection of the two mask bounding boxes 

543 and the mask planes that are present in both. Mask bits are only set, 

544 not cleared (i.e. this uses ``|=`` updates, not ``=`` assignments). 

545 """ 

546 lhs = self 

547 rhs = other 

548 if other.bbox != self.bbox: 

549 if (bbox := self.bbox.intersection(other.bbox)) is None: 

550 return 

551 lhs = self[bbox] 

552 rhs = other[bbox] 

553 for name in self.schema.names & other.schema.names: 

554 lhs.set(name, rhs.get(name)) 

555 

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

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

558 

559 Parameters 

560 ---------- 

561 plane 

562 Name of the mask plane. 

563 

564 Returns 

565 ------- 

566 numpy.ndarray 

567 A 2-d boolean array with the same shape as `bbox` that is `True` 

568 where the bit for ``plane`` is set and `False` elsewhere. 

569 """ 

570 bit = self.schema.bit(plane) 

571 return (self._array[..., bit.index] & bit.mask).astype(bool) 

572 

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

574 """Set a mask plane. 

575 

576 Parameters 

577 ---------- 

578 plane 

579 Name of the mask plane to set 

580 boolean_mask 

581 A 2-d boolean array with the same shape as `bbox` that is `True` 

582 where the bit for ``plane`` should be set and `False` where it 

583 should be left unchanged (*not* set to zero). May be ``...`` to 

584 set the bit everywhere. 

585 """ 

586 bit = self.schema.bit(plane) 

587 if boolean_mask is not ...: 

588 boolean_mask = boolean_mask.astype(bool) 

589 self._array[boolean_mask, bit.index] |= bit.mask 

590 

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

592 """Clear a mask plane. 

593 

594 Parameters 

595 ---------- 

596 plane 

597 Name of the mask plane to set 

598 boolean_mask 

599 A 2-d boolean array with the same shape as `bbox` that is `True` 

600 where the bit for ``plane`` should be cleared and `False` where it 

601 should be left unchanged. May be ``...`` to clear the bit 

602 everywhere. 

603 """ 

604 bit = self.schema.bit(plane) 

605 if boolean_mask is not ...: 

606 boolean_mask = boolean_mask.astype(bool) 

607 self._array[boolean_mask, bit.index] &= ~bit.mask 

608 

609 def serialize[P: pydantic.BaseModel]( 

610 self, 

611 archive: OutputArchive[P], 

612 *, 

613 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

614 save_projection: bool = True, 

615 add_offset_wcs: str | None = "A", 

616 ) -> MaskSerializationModel[P]: 

617 """Serialize the mask to an output archive. 

618 

619 Parameters 

620 ---------- 

621 archive 

622 Archive to write to. 

623 update_header 

624 A callback that will be given the FITS header for the HDU 

625 containing this mask in order to add keys to it. This callback 

626 may be provided but will not be called if the output format is not 

627 FITS. As multiple HDUs may be added, this function may be called 

628 multiple times. 

629 save_projection 

630 If `True`, save the `Projection` attached to the mask, if there 

631 is one. 

632 add_offset_wcs 

633 A FITS WCS single-character suffix to use when adding a linear 

634 WCS that maps the FITS array to the logical pixel coordinates 

635 defined by ``bbox.start``. Set to `None` to not write this WCS. 

636 """ 

637 data: list[ArrayReferenceModel] = [] 

638 for schema_2d in self.schema.split(np.int32): 

639 mask_2d = Mask( 

640 0, bbox=self.bbox, schema=schema_2d, projection=self._projection, obs_info=self._obs_info 

641 ) 

642 mask_2d.update(self) 

643 data.append(mask_2d._serialize_2d(archive, update_header=update_header)) 

644 serialized_projection: ProjectionSerializationModel[P] | None = None 

645 if save_projection and self.projection is not None: 

646 serialized_projection = archive.serialize_direct("projection", self.projection.serialize) 

647 serialized_dtype = NumberType.from_numpy(self.schema.dtype) 

648 assert is_integer(serialized_dtype), "Mask dtypes should always be integers." 

649 return MaskSerializationModel.model_construct( 

650 data=data, 

651 start=list(self.bbox.start), 

652 planes=list(self.schema), 

653 dtype=serialized_dtype, 

654 projection=serialized_projection, 

655 obs_info=self._obs_info, 

656 ) 

657 

658 def _serialize_2d[P: pydantic.BaseModel]( 

659 self, 

660 archive: OutputArchive[P], 

661 *, 

662 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

663 save_projection: bool = True, 

664 add_offset_wcs: str | None = "A", 

665 ) -> ArrayReferenceModel: 

666 def _update_header(header: astropy.io.fits.Header) -> None: 

667 update_header(header) 

668 self.schema.update_header(header) 

669 if self.projection is not None: 

670 if self.fits_wcs: 

671 header.update(self.fits_wcs.to_header(relax=True)) 

672 if add_offset_wcs is not None: 

673 fits.add_offset_wcs(header, x=self.bbox.x.start, y=self.bbox.y.start, key=add_offset_wcs) 

674 

675 assert self.array.shape[2] == 1, "Mask should be split before calling this method." 

676 return archive.add_array(self._array[:, :, 0], update_header=_update_header) 

677 

678 @classmethod 

679 def deserialize( 

680 cls, 

681 model: MaskSerializationModel[Any], 

682 archive: InputArchive[Any], 

683 *, 

684 bbox: Box | None = None, 

685 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

686 ) -> Mask: 

687 """Deserialize a mask from an input archive. 

688 

689 Parameters 

690 ---------- 

691 model 

692 A Pydantic model representation of the mask, holding references 

693 to data stored in the archive. 

694 archive 

695 Archive to read from. 

696 bbox 

697 Bounding box of a subimage to read instead. 

698 strip_header 

699 A callable that strips out any FITS header cards added by the 

700 ``update_header`` argument in the corresponding call to 

701 `serialize`. 

702 """ 

703 slices: tuple[slice, ...] | EllipsisType = ... 

704 if bbox is not None: 

705 slices = bbox.slice_within(model.bbox) 

706 else: 

707 bbox = model.bbox 

708 if not is_integer(model.dtype): 

709 raise ArchiveReadError(f"Mask array has a non-integer dtype: {model.dtype}.") 

710 schema = MaskSchema(model.planes, dtype=model.dtype.to_numpy()) 

711 projection = ( 

712 Projection.deserialize(model.projection, archive) if model.projection is not None else None 

713 ) 

714 result = Mask(0, schema=schema, bbox=bbox, projection=projection, obs_info=model.obs_info) 

715 schemas_2d = schema.split(np.int32) 

716 if len(schemas_2d) != len(model.data): 

717 raise ArchiveReadError( 

718 f"Number of mask arrays ({len(model.data)}) does not match expectation ({len(schemas_2d)})." 

719 ) 

720 for ref, schema_2d in zip(model.data, schemas_2d): 

721 mask_2d = cls._deserialize_2d( 

722 ref, schema_2d, bbox.start, archive, strip_header=strip_header, slices=slices 

723 ) 

724 result.update(mask_2d) 

725 return result 

726 

727 @classmethod 

728 def _deserialize_2d( 

729 cls, 

730 ref: ArrayReferenceModel, 

731 schema_2d: MaskSchema, 

732 start: Sequence[int], 

733 archive: InputArchive[Any], 

734 *, 

735 slices: tuple[slice, ...] | EllipsisType = ..., 

736 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

737 ) -> Mask: 

738 def _strip_header(header: astropy.io.fits.Header) -> None: 

739 strip_header(header) 

740 schema_2d.strip_header(header) 

741 fits.strip_wcs_cards(header) 

742 

743 array_2d = archive.get_array(ref, strip_header=_strip_header, slices=slices) 

744 return Mask(array_2d[:, :, np.newaxis], schema=schema_2d, start=start) 

745 

746 @staticmethod 

747 def _get_archive_tree_type[P: pydantic.BaseModel]( 

748 pointer_type: type[P], 

749 ) -> type[MaskSerializationModel[P]]: 

750 """Return the serialization model type for this object for an archive 

751 type that uses the given pointer type. 

752 """ 

753 return MaskSerializationModel[pointer_type] # type: ignore 

754 

755 _archive_default_name: ClassVar[str] = "mask" 

756 """The name this object should be serialized with when written as the 

757 top-level object. 

758 """ 

759 

760 def write_fits( 

761 self, 

762 filename: str, 

763 *, 

764 compression: fits.FitsCompressionOptions | None = fits.FitsCompressionOptions.DEFAULT, 

765 ) -> None: 

766 """Write the mask to a FITS file. 

767 

768 Parameters 

769 ---------- 

770 filename 

771 Name of the file to write to. Must be a local file. 

772 compression 

773 Compression options. 

774 """ 

775 compression_options = {} 

776 if compression is not fits.FitsCompressionOptions.DEFAULT: 

777 compression_options[self._archive_default_name] = compression 

778 fits.write(self, filename, compression_options) 

779 

780 @staticmethod 

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

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

783 

784 Parameters 

785 ---------- 

786 url 

787 URL of the file to read; may be any type supported by 

788 `lsst.resources.ResourcePath`. 

789 bbox 

790 Bounding box of a subimage to read instead. 

791 """ 

792 return fits.read(Mask, url, bbox=bbox) 

793 

794 @staticmethod 

795 def from_legacy( 

796 legacy: Any, 

797 plane_map: Mapping[str, MaskPlane] | None = None, 

798 ) -> Mask: 

799 """Convert from an `lsst.afw.image.Mask` instance. 

800 

801 Parameters 

802 ---------- 

803 legacy 

804 An `lsst.afw.image.Mask` instance. This will not share pixel 

805 data with the new object. 

806 plane_map 

807 A mapping from legacy mask plane name to the new plane name and 

808 description. 

809 """ 

810 return Mask._from_legacy_array( 

811 legacy.array, 

812 legacy.getMaskPlaneDict(), 

813 start=YX(y=legacy.getY0(), x=legacy.getX0()), 

814 plane_map=plane_map, 

815 ) 

816 

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

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

819 

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

821 

822 Parameters 

823 ---------- 

824 plane_map 

825 A mapping from legacy mask plane name to the new plane name and 

826 description. 

827 """ 

828 import lsst.afw.image 

829 import lsst.geom 

830 

831 result = lsst.afw.image.Mask(self.bbox.to_legacy()) 

832 if plane_map is None: 

833 plane_map = {plane.name: plane for plane in self.schema if plane is not None} 

834 for old_name, new_plane in plane_map.items():