Coverage for python/lsst/daf/butler/tests/registry_data/spatial.py: 0%
177 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-30 03:56 -0700
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28"""A script that creates and displays test data for topological spatial
29operations.
31The test data created by this script is intended to cover the complete set of
32interesting topological relationships between visit, visit_detector_region,
33tract, patch, and commonSkyPix with a relatively small number of actual
34dimension records. We use a different sky pixelization for visit- and
35skymap-based regions as an easy way to get different kinds of regions with
36interesting relationships.
38The data created by this script must be imported after that in base.yaml
39(which defines the instrument and detectors it assumes). It defines visits
40that don't actually correspond to any exposures; those could be added in
41another export file (along with the visit_definition records that relate them)
42in the future.
43"""
45from __future__ import annotations
47__all__ = []
49import argparse
50import os.path
51from collections.abc import Callable, Iterable, Iterator
52from typing import Any
54import numpy as np
55import yaml
56from astropy.time import Time
57from astropy.wcs import WCS
58from matplotlib import pyplot
60import lsst.daf.butler # register Time/YAML conversions. # noqa: F401
61from lsst.sphgeom import (
62 ConvexPolygon,
63 HealpixPixelization,
64 HtmPixelization,
65 LonLat,
66 Mq3cPixelization,
67 PixelizationABC,
68 Q3cPixelization,
69 RangeSet,
70 UnitVector3d,
71)
73# Pixelization for which one pixel defines the overall area of interest.
74PARENT_PIX = Mq3cPixelization(6)
76# Pixelization used as the commonSkyPix in butler (needs to be kept consistent
77# with the dimensions.yaml used in tests).
78COMMON_PIX = HtmPixelization(7)
80# Pixelization used to define visit and visit-detector regions.
81# Doesn't matter what this is, as long as its different from COMMON_PIX and
82# PATCH_GRID_PIX so we get interesting overlaps.
83DETECTOR_GRID_PIX = Mq3cPixelization(10)
85# Pixelization used to define tract and patch regions.
86# Doesn't matter what this is, as long as its different from COMMON_PIX and
87# DETECTOR_GRID_PIX so we get interesting overlaps.
88PATCH_GRID_PIX = Q3cPixelization(9)
90# Name of the instrument; this matches the dimension record in base.yaml,
91# which must be imported before the YAML file created by this script.
92INSTRUMENT_NAME = "Cam1"
94# Name of the skymap.
95SKYMAP_NAME = "SkyMap1"
97# Data used to define the visits and visit-detector regions.
98#
99# The write_yaml function adds fields common to all records and transforms
100# skypix indices into actual regions.
101#
102# Edit this data structure and examine the resulting plots to add more test
103# dimension records; running with ``--show-detector-grid`` may be helpful.
104VISIT_DATA: dict[int, dict[str, Any]] = {
105 1: {
106 "physical_filter": "Cam1-G",
107 "day_obs": 20210909,
108 "exposure_time": 60.0,
109 "target_name": "test_target",
110 "observation_reason": "science",
111 "science_program": "test_survey",
112 "zenith_angle": 5.0,
113 "datetime_begin": Time("2021-09-09T03:00:00", format="isot", scale="tai"),
114 "datetime_end": Time("2021-09-09T03:01:00", format="isot", scale="tai"),
115 "detector_regions": {
116 1: [12058870, 12058871, 12058872, 12058873],
117 2: [12058823, 12058824, 12058818, 12058829],
118 3: [12058848, 12058849, 12058850, 12058851],
119 4: [12058846],
120 },
121 },
122 2: {
123 "physical_filter": "Cam1-R1",
124 "day_obs": 20210909,
125 "exposure_time": 45.0,
126 "target_name": "test_target",
127 "observation_reason": "science",
128 "science_program": "test_survey",
129 "zenith_angle": 10.0,
130 "datetime_begin": Time("2021-09-09T03:02:00", format="isot", scale="tai"),
131 "datetime_end": Time("2021-09-09T03:03:00", format="isot", scale="tai"),
132 "detector_regions": {
133 1: [12058857, 12058854, 12058646, 12058649],
134 2: [12058842, 12058841, 12058661, 12058662],
135 3: [12058642, 12058653, 12058641, 12058654],
136 4: [12058659],
137 },
138 },
139}
141# Data used to define the tract and patch regions.
142#
143# The write_yaml function adds fields common to all records and transforms
144# skypix indices into actual regions.
145#
146# Edit this data structure and examine the resulting plots to add more test
147# dimension records; running with ``--show-patch-grid`` may be helpful.
148TRACT_DATA: dict[int, dict[int, dict[str, Any]]] = {
149 0: {
150 0: {"cell_x": 0, "cell_y": 0, "region": 458787},
151 1: {"cell_x": 1, "cell_y": 0, "region": 458790},
152 2: {"cell_x": 0, "cell_y": 1, "region": 458785},
153 3: {"cell_x": 1, "cell_y": 1, "region": 458788},
154 4: {"cell_x": 0, "cell_y": 2, "region": 458763},
155 5: {"cell_x": 1, "cell_y": 2, "region": 458766},
156 },
157 1: {
158 0: {"cell_x": 0, "cell_y": 0, "region": 458761},
159 1: {"cell_x": 1, "cell_y": 0, "region": 458764},
160 2: {"cell_x": 0, "cell_y": 1, "region": 458755},
161 3: {"cell_x": 1, "cell_y": 1, "region": 458758},
162 4: {"cell_x": 0, "cell_y": 2, "region": 458753},
163 5: {"cell_x": 1, "cell_y": 2, "region": 458756},
164 },
165}
168def main() -> None:
169 """Run script."""
170 parser = argparse.ArgumentParser(description="Create and examine spatial-topology registry test data.")
171 default_filename = os.path.join(os.path.dirname(__file__), "spatial.yaml")
172 parser.add_argument(
173 "--filename", type=str, default=default_filename, help="Filename for YAML export file."
174 )
175 parser.add_argument(
176 "--show-detector-grid",
177 action="store_true",
178 default=False,
179 help="Show the skypix grid used to define visit/detector regions.",
180 )
181 parser.add_argument(
182 "--show-patch-grid",
183 action="store_true",
184 default=False,
185 help="Show the skypix grid used to define patch regions.",
186 )
187 parser.add_argument(
188 "--no-common-skypix-grid",
189 dest="common_skypix_grid",
190 action="store_false",
191 default=True,
192 help="Do not show the common skypix grid.",
193 )
194 parser.add_argument(
195 "--show-healpix-grid",
196 type=int,
197 default=[],
198 help="Show a HEALPIX grid of this level.",
199 action="append",
200 )
201 parser.add_argument(
202 "--no-plot", action="store_false", dest="make_plots", default=True, help="Do not plot the regions."
203 )
204 parser.add_argument(
205 "--no-write",
206 action="store_false",
207 dest="write_yaml",
208 default=True,
209 help="Do not write the YAML export file.",
210 )
211 namespace = parser.parse_args()
212 if namespace.make_plots:
213 make_plots(
214 detector_grid=namespace.show_detector_grid,
215 patch_grid=namespace.show_patch_grid,
216 healpix_grids=namespace.show_healpix_grid,
217 common_skypix_grid=namespace.common_skypix_grid,
218 )
219 if namespace.write_yaml:
220 write_yaml(namespace.filename)
223def make_plots(
224 detector_grid: bool, patch_grid: bool, common_skypix_grid: bool = True, healpix_grids: Iterable[int] = ()
225) -> None:
226 """Plot the regions of the dimension records defined by this script.
228 Parameters
229 ----------
230 detector_grid : `bool`
231 If `True`, show the skypix grid used to define visit and visit-detector
232 regions.
233 patch_grid : `bool`
234 If `True`, show the skypix grid used to define tract and patch regions.
235 common_skypix_grid : `bool`, optional
236 If `True`, show the common skypix grid.
237 healpix_grids : `~collections.abc.Iterable` [`int`], optional
238 Levels of healpix grids to display.
239 """
240 parent_index = PARENT_PIX.index(UnitVector3d(1, 0, 0))
241 parent_pixel = PARENT_PIX.pixel(parent_index)
242 assert isinstance(parent_pixel, ConvexPolygon)
243 common_ranges = COMMON_PIX.envelope(parent_pixel)
244 detector_grid_ranges = DETECTOR_GRID_PIX.interior(parent_pixel)
245 patch_grid_ranges = PATCH_GRID_PIX.envelope(parent_pixel)
246 wcs = make_tangent_wcs(parent_pixel.getCentroid())
247 labels_used = set()
248 pyplot.figure(figsize=(16, 16))
249 pyplot.axis("off")
250 if common_skypix_grid:
251 plot_pixels(
252 COMMON_PIX,
253 wcs,
254 flatten_ranges(common_ranges),
255 polygons(facecolor="none", edgecolor="black", label="htm7"),
256 index_labels(color="black", alpha=0.5),
257 )
258 if detector_grid:
259 plot_pixels(
260 DETECTOR_GRID_PIX,
261 wcs,
262 flatten_ranges(detector_grid_ranges),
263 polygons(
264 edgecolor="black",
265 linewidth=1,
266 alpha=0.5,
267 linestyle=":",
268 facecolor="none",
269 ),
270 index_labels(color="black", alpha=0.5),
271 )
272 if patch_grid:
273 plot_pixels(
274 PATCH_GRID_PIX,
275 wcs,
276 flatten_ranges(patch_grid_ranges),
277 polygons(
278 edgecolor="black",
279 linewidth=1,
280 alpha=0.5,
281 linestyle=":",
282 facecolor="none",
283 ),
284 index_labels(color="black", alpha=0.5),
285 )
286 for healpix_level in healpix_grids:
287 pixelization = HealpixPixelization(healpix_level)
288 healpix_ranges = pixelization.envelope(parent_pixel)
289 plot_pixels(
290 pixelization,
291 wcs,
292 flatten_ranges(healpix_ranges),
293 polygons(
294 edgecolor="magenta",
295 linewidth=1,
296 alpha=0.5,
297 linestyle=":",
298 facecolor="none",
299 label=f"healpix{healpix_level}",
300 ),
301 index_labels(color="magenta", alpha=0.5),
302 )
303 colors = iter(["red", "blue", "cyan", "green"])
304 for (visit_id, visit_data), color in zip(VISIT_DATA.items(), colors, strict=False):
305 for detector_id, pixel_indices in visit_data["detector_regions"].items():
306 label: str | None = f"visit={visit_id}"
307 if label in labels_used:
308 label = None
309 else:
310 labels_used.add(label)
311 plot_hull(
312 DETECTOR_GRID_PIX,
313 wcs,
314 pixel_indices,
315 polygons(
316 edgecolor="none",
317 facecolor=color,
318 alpha=0.25,
319 label=label,
320 ),
321 labels(
322 text=str(detector_id),
323 color=color,
324 ),
325 )
326 for (tract_id, tract_data), color in zip(TRACT_DATA.items(), colors, strict=True):
327 for patch_id, patch_data in tract_data.items():
328 label = f"tract={tract_id}"
329 if label in labels_used:
330 label = None
331 else:
332 labels_used.add(label)
333 plot_pixels(
334 PATCH_GRID_PIX,
335 wcs,
336 [patch_data["region"]],
337 polygons(
338 edgecolor=color,
339 facecolor=color,
340 alpha=0.25,
341 label=label,
342 ),
343 labels(
344 text=str(patch_id),
345 color=color,
346 ),
347 )
348 parent_vertices = wcs.wcs_world2pix(np.array([lonlat_tuple(v) for v in parent_pixel.getVertices()]), 0)
349 pyplot.xlim(parent_vertices[:, 0].min(), parent_vertices[:, 0].max())
350 pyplot.ylim(parent_vertices[:, 1].min(), parent_vertices[:, 1].max())
351 pyplot.legend()
352 pyplot.show()
355def write_yaml(filename: str) -> None:
356 """Write the YAML export script with dimension record definitions.
358 Parameters
359 ----------
360 filename : `str`
361 Name of the file to write.
363 Notes
364 -----
365 This creates a YAML export file that defines records for the following
366 dimensions:
368 - visit_system
369 - visit
370 - visit_detector_region
371 - skymap
372 - tract
373 - patch
375 """
376 day_obs_records = [{"instrument": INSTRUMENT_NAME, "id": 20210909}]
377 visit_records = []
378 visit_detector_records = []
379 for visit_id, visit_data in VISIT_DATA.items():
380 visit_vertices = []
381 for detector_id, pixel_indices in visit_data["detector_regions"].items():
382 detector_vertices = []
383 for index in pixel_indices:
384 polygon = DETECTOR_GRID_PIX.pixel(index)
385 assert isinstance(polygon, ConvexPolygon)
386 detector_vertices.extend(polygon.getVertices())
387 visit_vertices.extend(detector_vertices)
388 visit_detector_records.append(
389 {
390 "instrument": INSTRUMENT_NAME,
391 "visit": visit_id,
392 "detector": detector_id,
393 "region": ConvexPolygon(detector_vertices),
394 }
395 )
396 visit_record = visit_data.copy()
397 del visit_record["detector_regions"]
398 visit_record["instrument"] = INSTRUMENT_NAME
399 visit_record["id"] = visit_id
400 visit_record["name"] = str(visit_id)
401 visit_record["region"] = ConvexPolygon(visit_vertices)
402 visit_records.append(visit_record)
403 skymap_records = [
404 {
405 "name": SKYMAP_NAME,
406 "hash": b"notreallyahashofanything!",
407 "tract_max": 50,
408 "patch_nx_max": 2,
409 "patch_ny_max": 3,
410 },
411 ]
412 tract_records = []
413 patch_records = []
414 for tract_id, tract_data in TRACT_DATA.items():
415 tract_vertices = []
416 for patch_id, patch_data in tract_data.items():
417 patch_polygon = PATCH_GRID_PIX.pixel(patch_data["region"])
418 assert isinstance(patch_polygon, ConvexPolygon)
419 tract_vertices.extend(patch_polygon.getVertices())
420 patch_record = patch_data.copy()
421 patch_record["region"] = patch_polygon
422 patch_record["id"] = patch_id
423 patch_record["tract"] = tract_id
424 patch_record["skymap"] = SKYMAP_NAME
425 patch_records.append(patch_record)
426 tract_record: dict[str, Any] = {}
427 tract_record["id"] = tract_id
428 tract_record["skymap"] = SKYMAP_NAME
429 tract_record["region"] = ConvexPolygon(tract_vertices)
430 tract_records.append(tract_record)
431 document = {
432 "description": "Butler Data Repository Export",
433 "version": "1.0.2",
434 "universe_version": 7,
435 "universe_namespace": "daf_butler",
436 "data": [
437 {
438 "type": "dimension",
439 "element": "day_obs",
440 "records": day_obs_records,
441 },
442 {
443 "type": "dimension",
444 "element": "visit",
445 "records": visit_records,
446 },
447 {
448 "type": "dimension",
449 "element": "visit_detector_region",
450 "records": visit_detector_records,
451 },
452 {
453 "type": "dimension",
454 "element": "skymap",
455 "records": skymap_records,
456 },
457 {
458 "type": "dimension",
459 "element": "tract",
460 "records": tract_records,
461 },
462 {
463 "type": "dimension",
464 "element": "patch",
465 "records": patch_records,
466 },
467 ],
468 }
469 with open(filename, mode="w") as file:
470 file.write("# Spatial test data; see spatial.py for more information.\n")
471 yaml.dump(document, file, sort_keys=False)
474def lonlat_tuple(position: LonLat | UnitVector3d) -> tuple[float, float]:
475 """Transform a `lsst.sphgeom.LonLat` or `lsst.sphgeom.Vector3d` to a
476 2-tuple of `float` degrees.
477 """
478 lonlat = LonLat(position)
479 return (lonlat.getLon().asDegrees(), lonlat.getLat().asDegrees())
482def make_tangent_wcs(position: LonLat | UnitVector3d) -> WCS:
483 """Create an `astropy.WCS` that maps the sky to a tangent plane with
484 degree-unit pixels at the given point.
486 Notes
487 -----
488 This uses astropy instead of afw just to avoid a new dependency. I suspect
489 that with a tiny bit of math I could just convert directly from
490 `sphgeom.UnitVector3d` to a point in a Euclidean 2-d plane, which is all I
491 want, but this seems fine as-is.
492 """
493 result = WCS(naxis=2)
494 result.wcs.crpix = [0.0, 0.0]
495 result.wcs.crval = lonlat_tuple(position)
496 result.wcs.ctype = ["RA---TAN", "DEC--TAN"]
497 result.wcs.cd = [[1.0, 0.0], [0.0, 1.0]]
498 return result
501def project_polygon_center(wcs: WCS, polygon: ConvexPolygon) -> np.ndarray:
502 """Return the WCS-projected center of the given polygon as a 2-element
503 `float` array.
504 """
505 return wcs.wcs_world2pix(np.array(lonlat_tuple(polygon.getCentroid()))[np.newaxis, :], 0)[0]
508def project_polygon_vertices(wcs: WCS, polygon: ConvexPolygon) -> np.ndarray:
509 """Return the WCS-projected vertices of the given polygon as a `float`
510 array with shape ``(n, 2)``.
511 """
512 vertices_sky = []
513 for vertex in polygon.getVertices():
514 vertices_sky.append(lonlat_tuple(vertex))
515 return wcs.wcs_world2pix(np.array(vertices_sky), 0)
518def plot_pixels(
519 pixelization: PixelizationABC,
520 wcs: WCS,
521 indices: Iterable[int],
522 *callbacks: Callable[[int, np.ndarray, np.ndarray], None],
523) -> None:
524 """Perform plotting actions defined by callbacks on each of a series of
525 skypix pixels.
527 Parameters
528 ----------
529 pixelization : `lsst.sphgeom.Pixelization`
530 Pixelization that interprets ``indices``.
531 wcs : `WCS`
532 Tangent plane to project spherical polygons onto.
533 indices : `~collections.abc.Iterable` [ `int` ]
534 Pixel indices to plot.
535 *callbacks
536 Callbacks to call for each pixel, passing the pixel index, the
537 projected center, and the projected vertices.
538 """
539 for index in indices:
540 polygon = pixelization.pixel(index)
541 center = project_polygon_center(wcs, polygon)
542 vertices = project_polygon_vertices(wcs, polygon)
543 for callback in callbacks:
544 callback(index, center, vertices)
547def plot_hull(
548 pixelization: PixelizationABC,
549 wcs: WCS,
550 indices: Iterable[int],
551 *callbacks: Callable[[list[int], np.ndarray, np.ndarray], None],
552) -> None:
553 """Perform plotting actions defined by callbacks on the convex hull of
554 a series of skypix pixels.
556 Parameters
557 ----------
558 pixelization : `lsst.sphgeom.Pixelization`
559 Pixelization that interprets ``indices``.
560 wcs : `WCS`
561 Tangent plane to project spherical polygons onto.
562 indices : `~collections.abc.Iterable` [ `int` ]
563 Pixel indices to plot.
564 *callbacks
565 Callbacks to call passing the list of pixel indices, the
566 projected center of the convex hull, and the projected vertices of the
567 convex hull.
568 """
569 vertices = []
570 indices = list(indices)
571 for index in indices:
572 polygon = pixelization.pixel(index)
573 vertices.extend(polygon.getVertices())
574 polygon = ConvexPolygon(vertices)
575 projected_center = project_polygon_center(wcs, polygon)
576 projected_vertices = project_polygon_vertices(wcs, polygon)
577 for callback in callbacks:
578 callback(indices, projected_center, projected_vertices)
581def polygons(label: str | None = None, **kwargs: Any) -> Callable[[Any, np.ndarray, np.ndarray], None]:
582 """Return a callback for use with `plot_pixels` and `plot_hull` that plots
583 polygon vertices.
585 Parameters
586 ----------
587 label : `str`, optional
588 Legend label for all polygons. Automatically deduplicated so the
589 legend will only contain one entry for each call.
590 **kwargs
591 Forwarded to `matplotlib.pyplot.fill`.
593 Returns
594 -------
595 func : `~collections.abc.Callable`
596 Callable for use with `plot_hull` or `plot_pixels`.
597 """
598 labels_used = set()
600 def func(index: Any, center: np.ndarray, vertices: np.ndarray) -> None:
601 if label is not None and label in labels_used:
602 label_to_use = None
603 else:
604 label_to_use = label
605 labels_used.add(label)
606 pyplot.fill(
607 vertices[:, 0],
608 vertices[:, 1],
609 label=label_to_use,
610 **kwargs,
611 )
613 return func
616def index_labels(**kwargs: Any) -> Callable[[int, np.ndarray, np.ndarray], None]:
617 """Return a callback for use with `plot_pixels` and `plot_hull` that adds
618 text annotations with pixel indices at pixel centers.
620 Parameters
621 ----------
622 **kwargs
623 Forwarded to `matplotlib.pyplot.text`.
624 """
626 def func(index: int, center: np.ndarray, vertices: np.ndarray) -> None:
627 pyplot.text(
628 center[0],
629 center[1],
630 str(index),
631 ha="center",
632 va="center",
633 **kwargs,
634 )
636 return func
639def labels(text: str, **kwargs: Any) -> Callable[[Any, np.ndarray, np.ndarray], None]:
640 """Return a callback for use with `plot_pixels` and `plot_hull` that adds
641 text annotations with the given text.
643 Parameters
644 ----------
645 text : `str`
646 Label text.
647 **kwargs
648 Forwarded to `matplotlib.pyplot.text`.
649 """
651 def func(index: Any, center: np.ndarray, vertices: np.ndarray) -> None:
652 pyplot.text(center[0], center[1], text, ha="center", va="center", **kwargs)
654 return func
657def flatten_ranges(ranges: RangeSet) -> Iterator[int]:
658 """Flatten an `lsst.sphgeom.RangeSet` into an iterator over pixel
659 indices.
660 """
661 for begin, end in ranges:
662 yield from range(begin, end)
665if __name__ == "__main__":
666 main()