Coverage for tests/bench_stencils.py: 0%

1# This file is part of dax_images_cutout. 

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# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22"""Opt-in benchmark comparing the AST and sphgeom stencil masking backends. 

23 

24This module is intentionally not named ``test_*`` so that pytest does not 

25collect it. Run it directly:: 

26 

27 python tests/bench_stencils.py 

28""" 

29 

30from __future__ import annotations 

31 

32import timeit 

33 

34import astropy.units as u 

35import astropy.wcs 

36 

37from lsst.dax.images.cutout.stencils import MaskBackend, SkyCircle 

38from lsst.images import Box, GeneralFrame, Mask, MaskPlane, MaskSchema, SkyProjection 

39from lsst.sphgeom import Angle, LonLat 

40 

41CENTER = LonLat.fromDegrees(12.0, 13.0) 

42PIXEL_SCALE = 0.2 # arcsec/pixel 

43 

44# (label, circle radius in arcsec, half-size of the square bbox in pixels). 

45CASES = ( 

46 ("small", 5.0, 64), 

47 ("medium", 30.0, 256), 

48 ("large", 120.0, 1024), 

49) 

50REPEATS = 20 

51 

52 

53def _arcsec(value: float) -> Angle: 

54 """Return a `lsst.sphgeom.Angle` for ``value`` arcseconds.""" 

55 return Angle((value * u.arcsec).to_value(u.rad)) 

56 

57 

58def _projection() -> SkyProjection: 

59 """Build a gnomonic `SkyProjection` centered on ``CENTER``.""" 

60 wcs = astropy.wcs.WCS(naxis=2) 

61 wcs.wcs.crpix = [1.0, 1.0] 

62 wcs.wcs.crval = [CENTER.getLon().asDegrees(), CENTER.getLat().asDegrees()] 

63 scale = PIXEL_SCALE / 3600.0 

64 wcs.wcs.cd = [[-scale, 0.0], [0.0, scale]] 

65 wcs.wcs.ctype = ["RA---TAN", "DEC--TAN"] 

66 return SkyProjection.from_fits_wcs(wcs, GeneralFrame(unit=u.pix)) 

67 

68 

69def _run(projection: SkyProjection, radius_arcsec: float, half: int, backend: MaskBackend) -> float: 

70 stencil = SkyCircle(CENTER, _arcsec(radius_arcsec)) 

71 box = Box.factory[-half : half + 1, -half : half + 1] 

72 schema = MaskSchema([MaskPlane("STENCIL", "stencil coverage")]) 

73 

74 def once() -> None: 

75 pixel_stencil = stencil.to_pixels(projection, box, backend=backend) 

76 mask = Mask(schema=schema, bbox=box) 

77 pixel_stencil.set_mask(mask, "STENCIL") 

78 

79 return min(timeit.repeat(once, number=1, repeat=REPEATS)) 

80 

81 

82def main() -> None: 

83 """Print a table of best-of-N timings for each backend and cutout size.""" 

84 projection = _projection() 

85 print(f"{'case':>8} {'bbox':>12} {'AST (ms)':>12} {'SPHGEOM (ms)':>14}") 

86 for label, radius_arcsec, half in CASES: 

87 side = 2 * half + 1 

88 ast_ms = _run(projection, radius_arcsec, half, MaskBackend.AST) * 1e3 

89 sph_ms = _run(projection, radius_arcsec, half, MaskBackend.SPHGEOM) * 1e3 

90 print(f"{label:>8} {f'{side}x{side}':>12} {ast_ms:>12.3f} {sph_ms:>14.3f}") 

91 

92 

93if __name__ == "__main__": 

94 main()