Coverage for python / lsst / daf / butler / tests / registry_data / spatial.py: 0%
174 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-14 23:36 +0000
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 Pixelization,
68 Q3cPixelization,
69 RangeSet,
70 UnitVector3d,
71 Vector3d,
72)
74# Pixelization for which one pixel defines the overall area of interest.
75PARENT_PIX = Mq3cPixelization(6)
77# Pixelization used as the commonSkyPix in butler (needs to be kept consistent
78# with the dimensions.yaml used in tests).
79COMMON_PIX = HtmPixelization(7)
81# Pixelization used to define visit and visit-detector regions.
82# Doesn't matter what this is, as long as its different from COMMON_PIX and
83# PATCH_GRID_PIX so we get interesting overlaps.
84DETECTOR_GRID_PIX = Mq3cPixelization(10)
86# Pixelization used to define tract and patch regions.
87# Doesn't matter what this is, as long as its different from COMMON_PIX and
88# DETECTOR_GRID_PIX so we get interesting overlaps.
89PATCH_GRID_PIX = Q3cPixelization(9)
91# Name of the instrument; this matches the dimension record in base.yaml,
92# which must be imported before the YAML file created by this script.
93INSTRUMENT_NAME = "Cam1"
95# Name of the skymap.
96SKYMAP_NAME = "SkyMap1"
98# Data used to define the visits and visit-detector regions.
99#
100# The write_yaml function adds fields common to all records and transforms
101# skypix indices into actual regions.
102#
103# Edit this data structure and examine the resulting plots to add more test
104# dimension records; running with ``--show-detector-grid`` may be helpful.
105VISIT_DATA: dict[int, dict[str, Any]] = {
106 1: {
107 "physical_filter": "Cam1-G",
108 "day_obs": 20210909,
109 "exposure_time": 60.0,
110 "target_name": "test_target",
111 "observation_reason": "science",
112 "science_program": "test_survey",
113 "zenith_angle": 5.0,
114 "datetime_begin": Time("2021-09-09T03:00:00", format="isot", scale="tai"),
115 "datetime_end": Time("2021-09-09T03:01:00", format="isot", scale="tai"),
116 "detector_regions": {
117 1: [12058870, 12058871, 12058872, 12058873],
118 2: [12058823, 12058824, 12058818, 12058829],
119 3: [12058848, 12058849, 12058850, 12058851],
120 4: [12058846],
121 },
122 },
123 2: {
124 "physical_filter": "Cam1-R1",
125 "day_obs": 20210909,
126 "exposure_time": 45.0,
127 "target_name": "test_target",
128 "observation_reason": "science",
129 "science_program": "test_survey",
130 "zenith_angle": 10.0,
131 "datetime_begin": Time("2021-09-09T03:02:00", format="isot", scale="tai"),
132 "datetime_end": Time("2021-09-09T03:03:00", format="isot", scale="tai"),
133 "detector_regions": {
134 1: [12058857, 12058854, 12058646, 12058649],
135 2: [12058842, 12058841, 12058661, 12058662],
136 3: [12058642, 12058653, 12058641, 12058654],
137 4: [12058659],
138 },
139 },
140}
142# Data used to define the tract and patch regions.
143#
144# The write_yaml function adds fields common to all records and transforms
145# skypix indices into actual regions.
146#
147# Edit this data structure and examine the resulting plots to add more test
148# dimension records; running with ``--show-patch-grid`` may be helpful.
149TRACT_DATA: dict[int, dict[int, dict[str, Any]]] = {
150 0: {
151 0: {"cell_x": 0, "cell_y": 0, "region": 458787},
152 1: {"cell_x": 1, "cell_y": 0, "region": 458790},
153 2: {"cell_x": 0, "cell_y": 1, "region": 458785},
154 3: {"cell_x": 1, "cell_y": 1, "region": 458788},
155 4: {"cell_x": 0, "cell_y": 2, "region": 458763},
156 5: {"cell_x": 1, "cell_y": 2, "region": 458766},
157 },
158 1: {
159 0: {"cell_x": 0, "cell_y": 0, "region": 458761},
160 1: {"cell_x": 1, "cell_y": 0, "region": 458764},
161 2: {"cell_x": 0, "cell_y": 1, "region": 458755},
162 3: {"cell_x": 1, "cell_y": 1, "region": 458758},
163 4: {"cell_x": 0, "cell_y": 2, "region": 458753},
164 5: {"cell_x": 1, "cell_y": 2, "region": 458756},
165 },
166}
169def main() -> None:
170 """Run script."""
171 parser = argparse.ArgumentParser(description="Create and examine spatial-topology registry test data.")
172 default_filename = os.path.join(os.path.dirname(__file__), "spatial.yaml")
173 parser.add_argument(
174 "--filename", type=str, default=default_filename, help="Filename for YAML export file."
175 )
176 parser.add_argument(
177 "--show-detector-grid",
178 action="store_true",
179 default=False,
180 help="Show the skypix grid used to define visit/detector regions.",
181 )
182 parser.add_argument(
183 "--show-patch-grid",
184 action="store_true",
185 default=False,
186 help="Show the skypix grid used to define patch regions.",
187 )
188 parser.add_argument(
189 "--no-common-skypix-grid",
190 dest="common_skypix_grid",
191 action="store_false",
192 default=True,
193 help="Do not show the common skypix grid.",
194 )
195 parser.add_argument(
196 "--show-healpix-grid",
197 type=int,
198 default=[],
199 help="Show a HEALPIX grid of this level.",
200 action="append",
201 )
202 parser.add_argument(
203 "--no-plot", action="store_false", dest="make_plots", default=True, help="Do not plot the regions."
204 )
205 parser.add_argument(
206 "--no-write",
207 action="store_false",
208 dest="write_yaml",
209 default=True,
210 help="Do not write the YAML export file.",
211 )
212 namespace = parser.parse_args()
213 if namespace.make_plots:
214 make_plots(
215 detector_grid=namespace.show_detector_grid,
216 patch_grid=namespace.show_patch_grid,
217 healpix_grids=namespace.show_healpix_grid,
218 common_skypix_grid=namespace.common_skypix_grid,
219 )
220 if namespace.write_yaml:
221 write_yaml(namespace.filename)
224def make_plots(
225 detector_grid: bool, patch_grid: bool, common_skypix_grid: bool = True, healpix_grids: Iterable[int] = ()
226) -> None:
227 """Plot the regions of the dimension records defined by this script.
229 Parameters
230 ----------
231 detector_grid : `bool`
232 If `True`, show the skypix grid used to define visit and visit-detector
233 regions.
234 patch_grid : `bool`
235 If `True`, show the skypix grid used to define tract and patch regions.
236 common_skypix_grid : `bool`, optional
237 If `True`, show the common skypix grid.
238 healpix_grids : `~collections.abc.Iterable` [`int`], optional
239 Levels of healpix grids to display.
240 """
241 parent_index = PARENT_PIX.index(UnitVector3d(1, 0, 0))
242 parent_pixel = PARENT_PIX.pixel(parent_index)
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 detector_vertices.extend(polygon.getVertices())
386 visit_vertices.extend(detector_vertices)
387 visit_detector_records.append(
388 {
389 "instrument": INSTRUMENT_NAME,
390 "visit": visit_id,
391 "detector": detector_id,
392 "region": ConvexPolygon(detector_vertices),
393 }
394 )
395 visit_record = visit_data.copy()
396 del visit_record["detector_regions"]
397 visit_record["instrument"] = INSTRUMENT_NAME
398 visit_record["id"] = visit_id
399 visit_record["name"] = str(visit_id)
400 visit_record["region"] = ConvexPolygon(visit_vertices)
401 visit_records.append(visit_record)
402 skymap_records = [
403 {
404 "name": SKYMAP_NAME,
405 "hash": b"notreallyahashofanything!",
406 "tract_max": 50,
407 "patch_nx_max": 2,
408 "patch_ny_max": 3,
409 },
410 ]
411 tract_records = []
412 patch_records = []
413 for tract_id, tract_data in TRACT_DATA.items():
414 tract_vertices = []
415 for patch_id, patch_data in tract_data.items():
416 patch_polygon = PATCH_GRID_PIX.pixel(patch_data["region"])
417 tract_vertices.extend(patch_polygon.getVertices())
418 patch_record = patch_data.copy()
419 patch_record["region"] = patch_polygon
420 patch_record["id"] = patch_id
421 patch_record["tract"] = tract_id
422 patch_record["skymap"] = SKYMAP_NAME
423 patch_records.append(patch_record)
424 tract_record: dict[str, Any] = {}
425 tract_record["id"] = tract_id
426 tract_record["skymap"] = SKYMAP_NAME
427 tract_record["region"] = ConvexPolygon(tract_vertices)
428 tract_records.append(tract_record)
429 document = {
430 "description": "Butler Data Repository Export",
431 "version": "1.0.2",
432 "universe_version": 7,
433 "universe_namespace": "daf_butler",
434 "data": [
435 {
436 "type": "dimension",
437 "element": "day_obs",
438 "records": day_obs_records,
439 },
440 {
441 "type": "dimension",
442 "element": "visit",
443 "records": visit_records,
444 },
445 {
446 "type": "dimension",
447 "element": "visit_detector_region",
448 "records": visit_detector_records,
449 },
450 {
451 "type": "dimension",
452 "element": "skymap",
453 "records": skymap_records,
454 },
455 {
456 "type": "dimension",
457 "element": "tract",
458 "records": tract_records,
459 },
460 {
461 "type": "dimension",
462 "element": "patch",
463 "records": patch_records,
464 },
465 ],
466 }
467 with open(filename, mode="w") as file:
468 file.write("# Spatial test data; see spatial.py for more information.\n")
469 yaml.dump(document, file, sort_keys=False)
472def lonlat_tuple(position: LonLat | Vector3d) -> tuple[float, float]:
473 """Transform a `lsst.sphgeom.LonLat` or `lsst.sphgeom.Vector3d` to a
474 2-tuple of `float` degrees.
475 """
476 lonlat = LonLat(position)
477 return (lonlat.getLon().asDegrees(), lonlat.getLat().asDegrees())
480def make_tangent_wcs(position: LonLat | Vector3d) -> WCS:
481 """Create an `astropy.WCS` that maps the sky to a tangent plane with
482 degree-unit pixels at the given point.
484 Notes
485 -----
486 This uses astropy instead of afw just to avoid a new dependency. I suspect
487 that with a tiny bit of math I could just convert directly from
488 `sphgeom.UnitVector3d` to a point in a Euclidean 2-d plane, which is all I
489 want, but this seems fine as-is.
490 """
491 result = WCS(naxis=2)
492 result.wcs.crpix = [0.0, 0.0]
493 result.wcs.crval = lonlat_tuple(position)
494 result.wcs.ctype = ["RA---TAN", "DEC--TAN"]
495 result.wcs.cd = [[1.0, 0.0], [0.0, 1.0]]
496 return result
499def project_polygon_center(wcs: WCS, polygon: ConvexPolygon) -> np.ndarray:
500 """Return the WCS-projected center of the given polygon as a 2-element
501 `float` array.
502 """
503 return wcs.wcs_world2pix(np.array(lonlat_tuple(polygon.getCentroid()))[np.newaxis, :], 0)[0]
506def project_polygon_vertices(wcs: WCS, polygon: ConvexPolygon) -> np.ndarray:
507 """Return the WCS-projected vertices of the given polygon as a `float`
508 array with shape ``(n, 2)``.
509 """
510 vertices_sky = []
511 for vertex in polygon.getVertices():
512 vertices_sky.append(lonlat_tuple(vertex))
513 return wcs.wcs_world2pix(np.array(vertices_sky), 0)
516def plot_pixels(
517 pixelization: Pixelization,
518 wcs: WCS,
519 indices: Iterable[int],
520 *callbacks: Callable[[int, np.ndarray, np.ndarray], None],
521) -> None:
522 """Perform plotting actions defined by callbacks on each of a series of
523 skypix pixels.
525 Parameters
526 ----------
527 pixelization : `lsst.sphgeom.Pixelization`
528 Pixelization that interprets ``indices``.
529 wcs : `WCS`
530 Tangent plane to project spherical polygons onto.
531 indices : `~collections.abc.Iterable` [ `int` ]
532 Pixel indices to plot.
533 *callbacks
534 Callbacks to call for each pixel, passing the pixel index, the
535 projected center, and the projected vertices.
536 """
537 for index in indices:
538 polygon = pixelization.pixel(index)
539 center = project_polygon_center(wcs, polygon)
540 vertices = project_polygon_vertices(wcs, polygon)
541 for callback in callbacks:
542 callback(index, center, vertices)
545def plot_hull(
546 pixelization: Pixelization,
547 wcs: WCS,
548 indices: Iterable[int],
549 *callbacks: Callable[[list[int], np.ndarray, np.ndarray], None],
550) -> None:
551 """Perform plotting actions defined by callbacks on the convex hull of
552 a series of skypix pixels.
554 Parameters
555 ----------
556 pixelization : `lsst.sphgeom.Pixelization`
557 Pixelization that interprets ``indices``.
558 wcs : `WCS`
559 Tangent plane to project spherical polygons onto.
560 indices : `~collections.abc.Iterable` [ `int` ]
561 Pixel indices to plot.
562 *callbacks
563 Callbacks to call passing the list of pixel indices, the
564 projected center of the convex hull, and the projected vertices of the
565 convex hull.
566 """
567 vertices = []
568 indices = list(indices)
569 for index in indices:
570 polygon = pixelization.pixel(index)
571 vertices.extend(polygon.getVertices())
572 polygon = ConvexPolygon(vertices)
573 projected_center = project_polygon_center(wcs, polygon)
574 projected_vertices = project_polygon_vertices(wcs, polygon)
575 for callback in callbacks:
576 callback(indices, projected_center, projected_vertices)
579def polygons(label: str | None = None, **kwargs: Any) -> Callable[[Any, np.ndarray, np.ndarray], None]:
580 """Return a callback for use with `plot_pixels` and `plot_hull` that plots
581 polygon vertices.
583 Parameters
584 ----------
585 label : `str`, optional
586 Legend label for all polygons. Automatically deduplicated so the
587 legend will only contain one entry for each call.
588 **kwargs
589 Forwarded to `matplotlib.pyplot.fill`.
591 Returns
592 -------
593 func : `~collections.abc.Callable`
594 Callable for use with `plot_hull` or `plot_pixels`.
595 """
596 labels_used = set()
598 def func(index: Any, center: np.ndarray, vertices: np.ndarray) -> None:
599 if label is not None and label in labels_used:
600 label_to_use = None
601 else:
602 label_to_use = label
603 labels_used.add(label)
604 pyplot.fill(
605 vertices[:, 0],
606 vertices[:, 1],
607 label=label_to_use,
608 **kwargs,
609 )
611 return func
614def index_labels(**kwargs: Any) -> Callable[[int, np.ndarray, np.ndarray], None]:
615 """Return a callback for use with `plot_pixels` and `plot_hull` that adds
616 text annotations with pixel indices at pixel centers.
618 Parameters
619 ----------
620 **kwargs
621 Forwarded to `matplotlib.pyplot.text`.
622 """
624 def func(index: int, center: np.ndarray, vertices: np.ndarray) -> None:
625 pyplot.text(
626 center[0],
627 center[1],
628 str(index),
629 ha="center",
630 va="center",
631 **kwargs,
632 )
634 return func
637def labels(text: str, **kwargs: Any) -> Callable[[Any, np.ndarray, np.ndarray], None]:
638 """Return a callback for use with `plot_pixels` and `plot_hull` that adds
639 text annotations with the given text.
641 Parameters
642 ----------
643 text : `str`
644 Label text.
645 **kwargs
646 Forwarded to `matplotlib.pyplot.text`.
647 """
649 def func(index: Any, center: np.ndarray, vertices: np.ndarray) -> None:
650 pyplot.text(center[0], center[1], text, ha="center", va="center", **kwargs)
652 return func
655def flatten_ranges(ranges: RangeSet) -> Iterator[int]:
656 """Flatten an `lsst.sphgeom.RangeSet` into an iterator over pixel
657 indices.
658 """
659 for begin, end in ranges:
660 yield from range(begin, end)
663if __name__ == "__main__":
664 main()