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

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, SkyProjection, SkyProjectionSerializationModel 

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): 202 ↛ 204line 202 didn't jump to line 204 because the condition on line 202 was always true

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: 280 ↛ 282line 280 didn't jump to line 282 because the condition on line 280 was always true

281 schemas.append(MaskSchema(planes, dtype=dtype)) 

282 if not schemas: 282 ↛ 283line 282 didn't jump to line 283 because the condition on line 282 was never true

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: 289 ↛ 288line 289 didn't jump to line 288 because the condition on line 289 was always true

290 bit = self.bit(plane.name) 

291 if bit.index != 0: 291 ↛ 292line 291 didn't jump to line 292 because the condition on line 291 was never true

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: 305 ↛ 304line 305 didn't jump to line 304 because the condition on line 305 was always true

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 yx0 

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 sky_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 ``yx0`` 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 yx0: Sequence[int] | None = None, 

359 shape: Sequence[int] | None = None, 

360 sky_projection: SkyProjection | 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 isinstance(array_or_fill, np.ndarray): 

367 array = np.array(array_or_fill, dtype=schema.dtype, copy=None) 

368 if array.ndim != 3: 

369 raise ValueError("Mask array must be 3-d.") 

370 if bbox is None: 

371 bbox = Box.from_shape(array.shape[:-1], start=yx0) 

372 elif bbox.shape + (schema.mask_size,) != array.shape: 

373 raise ValueError( 

374 f"Explicit bbox shape {bbox.shape} and schema of size {schema.mask_size} do not " 

375 f"match array with shape {array.shape}." 

376 ) 

377 if shape is not None and shape + (schema.mask_size,) != array.shape: 

378 raise ValueError( 

379 f"Explicit shape {shape} and schema of size {schema.mask_size} do " 

380 f"not match array with shape {array.shape}." 

381 ) 

382 

383 else: 

384 if bbox is None: 

385 if shape is None: 385 ↛ 387line 385 didn't jump to line 387 because the condition on line 385 was always true

386 raise TypeError("No bbox, size, or array provided.") 

387 bbox = Box.from_shape(shape, start=yx0) 

388 array = np.full(bbox.shape + (schema.mask_size,), array_or_fill, dtype=schema.dtype) 

389 self._array = array 

390 self._bbox: Box = bbox 

391 self._schema: MaskSchema = schema 

392 self._sky_projection = sky_projection 

393 

394 @property 

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

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

397 

398 Assigning to this attribute modifies the existing array in place; the 

399 bounding box and underlying data pointer are never changed. 

400 """ 

401 return self._array 

402 

403 @array.setter 

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

405 self._array[:, :] = value 

406 

407 @property 

408 def schema(self) -> MaskSchema: 

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

410 (`MaskSchema`). 

411 """ 

412 return self._schema 

413 

414 @property 

415 def bbox(self) -> Box: 

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

417 

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

419 """ 

420 return self._bbox 

421 

422 @property 

423 def sky_projection(self) -> SkyProjection[Any] | None: 

424 """The projection that maps this mask's pixel grid to the sky 

425 (`SkyProjection` | `None`). 

426 

427 Notes 

428 ----- 

429 The pixel coordinates used by this projection account for the bounding 

430 box ``start`` (i.e. ``yx0``); they are not just array indices. 

431 """ 

432 return self._sky_projection 

433 

434 def __getitem__(self, bbox: Box | EllipsisType) -> Mask: 

435 if bbox is ...: 

436 return self 

437 super().__getitem__(bbox) 

438 return self._transfer_metadata( 

439 Mask( 

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

441 bbox=bbox, 

442 schema=self.schema, 

443 sky_projection=self._sky_projection, 

444 ), 

445 bbox=bbox, 

446 ) 

447 

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

449 subview = self[bbox] 

450 subview.clear() 

451 subview.update(value) 

452 

453 def __str__(self) -> str: 

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

455 

456 def __repr__(self) -> str: 

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

458 

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

460 if not isinstance(other, Mask): 

461 return NotImplemented 

462 return ( 

463 self._bbox == other._bbox 

464 and self._schema == other._schema 

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

466 ) 

467 

468 def copy(self) -> Mask: 

469 """Deep-copy the mask and metadata.""" 

470 return self._transfer_metadata( 

471 Mask( 

472 self._array.copy(), bbox=self._bbox, schema=self._schema, sky_projection=self._sky_projection 

473 ), 

474 copy=True, 

475 ) 

476 

477 def view( 

478 self, 

479 *, 

480 schema: MaskSchema | EllipsisType = ..., 

481 sky_projection: SkyProjection | None | EllipsisType = ..., 

482 yx0: 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 ...: 491 ↛ 494line 491 didn't jump to line 494 because the condition on line 491 was always true

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 sky_projection is ...: 496 ↛ 497line 496 didn't jump to line 497 because the condition on line 496 was never true

497 sky_projection = self._sky_projection 

498 if yx0 is ...: 498 ↛ 500line 498 didn't jump to line 500 because the condition on line 498 was always true

499 yx0 = self._bbox.start 

500 return self._transfer_metadata( 

501 Mask(self._array, yx0=yx0, schema=schema, sky_projection=sky_projection) 

502 ) 

503 

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

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

506 

507 Notes 

508 ----- 

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

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

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

512 """ 

513 lhs = self 

514 rhs = other 

515 if other.bbox != self.bbox: 515 ↛ 516line 515 didn't jump to line 516 because the condition on line 515 was never true

516 try: 

517 bbox = self.bbox.intersection(other.bbox) 

518 except NoOverlapError: 

519 return 

520 lhs = self[bbox] 

521 rhs = other[bbox] 

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

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

524 

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

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

527 

528 Parameters 

529 ---------- 

530 plane 

531 Name of the mask plane. 

532 

533 Returns 

534 ------- 

535 numpy.ndarray 

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

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

538 """ 

539 bit = self.schema.bit(plane) 

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

541 

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

543 """Set a mask plane. 

544 

545 Parameters 

546 ---------- 

547 plane 

548 Name of the mask plane to set 

549 boolean_mask 

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

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

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

553 set the bit everywhere. 

554 """ 

555 bit = self.schema.bit(plane) 

556 if boolean_mask is not ...: 556 ↛ 558line 556 didn't jump to line 558 because the condition on line 556 was always true

557 boolean_mask = boolean_mask.astype(bool) 

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

559 

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

561 """Clear one or more mask planes. 

562 

563 Parameters 

564 ---------- 

565 plane 

566 Name of the mask plane to set. If `None` all mask planes are 

567 cleared. 

568 boolean_mask 

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

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

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

572 everywhere. 

573 """ 

574 if boolean_mask is not ...: 574 ↛ 575line 574 didn't jump to line 575 because the condition on line 574 was never true

575 boolean_mask = boolean_mask.astype(bool) 

576 if plane is None: 576 ↛ 579line 576 didn't jump to line 579 because the condition on line 576 was always true

577 self._array[boolean_mask, :] = 0 

578 else: 

579 bit = self.schema.bit(plane) 

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

581 

582 def serialize[P: pydantic.BaseModel]( 

583 self, 

584 archive: OutputArchive[P], 

585 *, 

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

587 save_projection: bool = True, 

588 add_offset_wcs: str | None = "A", 

589 tile_shape: tuple[int, ...] | None = None, 

590 options_name: str | None = None, 

591 ) -> MaskSerializationModel[P]: 

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

593 

594 Parameters 

595 ---------- 

596 archive 

597 Archive to write to. 

598 update_header 

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

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

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

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

603 multiple times. 

604 save_projection 

605 If `True`, save the `SkyProjection` attached to the image, if there 

606 is one. This does not affect whether a FITS WCS corresponding to 

607 the projection is written (it always is, if available, and if 

608 ``add_offset_wcs`` is not ``" "``). 

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`` / ``yx0``. Set to `None` to not write 

613 this WCS. If this is set to ``" "``, it will prevent the 

614 `SkyProjection` from being saved as a FITS WCS. 

615 tile_shape 

616 The recommended shape of each tile, if the archive will save 

617 the array in distinct tiles for faster subarray retrieval. 

618 This is a hint; archives are not required to use this value. 

619 options_name 

620 Use this name to look up archive options. 

621 """ 

622 if _archive_prefers_native_mask_arrays(archive): 

623 # HDS presents array dimensions in Fortran order, which is the 

624 # reverse of the h5py dataset shape. Store the in-memory trailing 

625 # mask-byte axis first in HDF5 so Starlink tools see HDS axes 

626 # (x, y, byte), without changing the bit packing within a pixel. 

627 array_model = archive.add_array( 

628 np.moveaxis(self._array, -1, 0), 

629 update_header=update_header, 

630 tile_shape=tile_shape, 

631 options_name=options_name, 

632 ) 

633 if not isinstance(array_model, ArrayReferenceModel): 633 ↛ 634line 633 didn't jump to line 634 because the condition on line 633 was never true

634 raise RuntimeError("Native mask arrays require reference array storage.") 

635 array_model.shape = list(self._array.shape) 

636 data: list[ArrayReferenceModel | InlineArrayModel] = [array_model] 

637 else: 

638 data = [] 

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

640 mask_2d = Mask(0, bbox=self.bbox, schema=schema_2d, sky_projection=self._sky_projection) 

641 mask_2d.update(self) 

642 data.append( 

643 mask_2d._serialize_2d( 

644 archive, 

645 update_header=update_header, 

646 add_offset_wcs=add_offset_wcs, 

647 tile_shape=tile_shape, 

648 options_name=options_name, 

649 ) 

650 ) 

651 serialized_projection: SkyProjectionSerializationModel[P] | None = None 

652 if save_projection and self.sky_projection is not None: 652 ↛ 653line 652 didn't jump to line 653 because the condition on line 652 was never true

653 serialized_projection = archive.serialize_direct("sky_projection", self.sky_projection.serialize) 

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

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

656 return MaskSerializationModel.model_construct( 

657 data=data, 

658 yx0=list(self.bbox.start), 

659 planes=list(self.schema), 

660 dtype=serialized_dtype, 

661 sky_projection=serialized_projection, 

662 metadata=self.metadata, 

663 ) 

664 

665 def _serialize_2d[P: pydantic.BaseModel]( 

666 self, 

667 archive: OutputArchive[P], 

668 *, 

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

670 add_offset_wcs: str | None = "A", 

671 tile_shape: tuple[int, ...] | None = None, 

672 options_name: str | None = None, 

673 ) -> ArrayReferenceModel | InlineArrayModel: 

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

675 update_header(header) 

676 self.schema.update_header(header) 

677 if self.sky_projection is not None and add_offset_wcs != " ": 

678 if self.fits_wcs: 678 ↛ 680line 678 didn't jump to line 680 because the condition on line 678 was always true

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

680 if add_offset_wcs is not None: 680 ↛ exitline 680 didn't return from function '_update_header' because the condition on line 680 was always true

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

682 

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

684 return archive.add_array( 

685 self._array[:, :, 0], 

686 update_header=_update_header, 

687 tile_shape=tile_shape, 

688 options_name=options_name, 

689 ) 

690 

691 @staticmethod 

692 def _get_archive_tree_type[P: pydantic.BaseModel]( 

693 pointer_type: type[P], 

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

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

696 type that uses the given pointer type. 

697 """ 

698 return MaskSerializationModel[pointer_type] # type: ignore 

699 

700 _archive_default_name: ClassVar[str] = "mask" 

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

702 top-level object. 

703 """ 

704 

705 @staticmethod 

706 def from_legacy( 

707 legacy: Any, 

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

709 ) -> Mask: 

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

711 

712 Parameters 

713 ---------- 

714 legacy 

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

716 data with the new object. 

717 plane_map 

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

719 description. If not provided, the right legacy mask plane will be 

720 guessed, but this can depend on which mask planes the legacy 

721 mask actually has set. 

722 """ 

723 return Mask._from_legacy_array( 

724 legacy.array, 

725 legacy.getMaskPlaneDict(), 

726 yx0=YX(y=legacy.getY0(), x=legacy.getX0()), 

727 plane_map=plane_map, 

728 ) 

729 

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

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

732 

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

734 

735 Parameters 

736 ---------- 

737 plane_map 

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

739 description. 

740 """ 

741 import lsst.afw.image 

742 import lsst.geom 

743 

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

745 if plane_map is None: 745 ↛ 747line 745 didn't jump to line 747 because the condition on line 745 was always true

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

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

748 old_bit = result.addMaskPlane(old_name) 

749 old_bitmask = 1 << old_bit 

750 if old_bitmask == 2147483648: 750 ↛ 753line 750 didn't jump to line 753 because the condition on line 750 was never true

751 # afw uses int32 masks, but relies on overflow wrapping, which 

752 # numpy doesn't like. 

753 old_bitmask = -2147483648 

754 if new_plane in self.schema: 754 ↛ 747line 754 didn't jump to line 747 because the condition on line 754 was always true

755 result.array[self.get(new_plane.name)] |= old_bitmask 

756 return result 

757 

758 @staticmethod 

759 def _from_legacy_array( 

760 array2d: np.ndarray, 

761 old_planes: Mapping[str, int], 

762 *, 

763 yx0: YX[int], 

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

765 sky_projection: SkyProjection | None = None, 

766 ) -> Mask: 

767 if plane_map is None: 

768 plane_map = _guess_legacy_plane_map(old_planes) 

769 planes: list[MaskPlane] = list(plane_map.values()) if plane_map is not None else [] 

770 new_name_to_old_bitmask: dict[str, int] = {} 

771 for old_name, old_bit in old_planes.items(): 

772 old_bitmask = 1 << old_bit 

773 if old_bitmask == 2147483648: 

774 # afw uses int32 masks, but relies on overflow wrapping, which 

775 # numpy doesn't like. 

776 old_bitmask = -2147483648 

777 if new_plane := plane_map.get(old_name): 

778 # Already added to 'planes' at initialization. 

779 new_name_to_old_bitmask[new_plane.name] = old_bitmask 

780 else: 

781 if n_orphaned := np.count_nonzero(array2d & old_bitmask): 

782 raise RuntimeError( 

783 f"Legacy mask plane {old_name!r} is not remapped, " 

784 f"but {n_orphaned} pixels have this bit set." 

785 ) 

786 schema = MaskSchema(planes) 

787 mask = Mask(0, schema=schema, yx0=yx0, shape=array2d.shape, sky_projection=sky_projection) 

788 for new_name, old_bitmask in new_name_to_old_bitmask.items(): 

789 mask.set(new_name, array2d & old_bitmask) 

790 return mask 

791 

792 @staticmethod 

793 def read_legacy( 

794 uri: ResourcePathExpression, 

795 *, 

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

797 ext: str | int = 1, 

798 fits_wcs_frame: Frame | None = None, 

799 ) -> Mask: 

800 """Read a FITS file written by `lsst.afw.image.Mask.writeFits`. 

801 

802 Parameters 

803 ---------- 

804 uri 

805 URI or file name. 

806 plane_map 

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

808 description. If not provided, the right legacy mask plane will be 

809 guessed, but this can depend on which mask planes the legacy 

810 mask actually has set. 

811 ext 

812 Name or index of the FITS HDU to read. 

813 fits_wcs_frame 

814 If not `None` and the HDU containing the mask has a FITS WCS, 

815 attach a `SkyProjection` to the returned mask by converting that 

816 WCS. 

817 """ 

818 opaque_metadata = fits.FitsOpaqueMetadata() 

819 fs, fspath = ResourcePath(uri).to_fsspec() 

820 with fs.open(fspath) as stream, astropy.io.fits.open(stream) as hdu_list: 

821 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header) 

822 result = Mask._read_legacy_hdu( 

823 hdu_list[ext], opaque_metadata, plane_map=plane_map, fits_wcs_frame=fits_wcs_frame 

824 ) 

825 result._opaque_metadata = opaque_metadata 

826 return result 

827 

828 @staticmethod 

829 def _read_legacy_hdu( 

830 hdu: astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU | astropy.io.fits.BinTableHDU,