Coverage for python/lsst/source/injection/utils/generate_injection_catalog.py: 15%

63 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-11 11:50 +0000

1# This file is part of source_injection. 

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 

22from __future__ import annotations 

23 

24__all__ = ["generate_injection_catalog"] 

25 

26import hashlib 

27import itertools 

28import logging 

29from collections.abc import Sequence 

30from typing import Any 

31 

32import numpy as np 

33from astropy.table import Table, hstack 

34from lsst.afw.geom import SkyWcs 

35from scipy.stats import qmc 

36 

37 

38def generate_injection_catalog( 

39 ra_lim: Sequence[float], 

40 dec_lim: Sequence[float], 

41 wcs: SkyWcs = None, 

42 number: int = 1, 

43 density: int | None = None, 

44 seed: Any = None, 

45 log_level: int = logging.INFO, 

46 **kwargs: Any, 

47) -> Table: 

48 """Generate a synthetic source injection catalog. 

49 

50 This function generates a synthetic source injection catalog from user 

51 supplied input parameters. The catalog is returned as an astropy Table. 

52 

53 On-sky source positions are generated using the quasi-random Halton 

54 sequence. By default, the Halton sequence is seeded using the product of 

55 the right ascension and declination limit ranges. This ensures that the 

56 same sequence is always generated for the same limits. This seed may be 

57 overridden by specifying the ``seed`` parameter. 

58 

59 A unique injection ID is generated for each source. The injection ID 

60 encodes two pieces of information: the unique source identification number 

61 and the version number of the source as specified by the ``number`` 

62 parameter. To achieve this, the unique source ID number is multiplied by 

63 `10**n` such that the sum of the multiplied source ID number with the 

64 unique repeated version number will always be unique. For example, an 

65 injection catalog with `number = 3` versions of each source will have 

66 injection IDs: 0, 1, 2, 10, 11, 12, 20, 21, 22, etc. If `number = 20`, then 

67 the injection IDs will be: 0, 1, 2, ..., 17, 18, 19, 100, 101, 102, etc. 

68 If `number = 1` (default) then the injection ID will be a simple sequential 

69 list of integers. 

70 

71 Parameters 

72 ---------- 

73 ra_lim : `Sequence` [`float`] 

74 The right ascension limits of the catalog in degrees. 

75 dec_lim : `Sequence` [`float`] 

76 The declination limits of the catalog in degrees. 

77 wcs : `lsst.afw.geom.SkyWcs`, optional 

78 The WCS associated with these data. If not given or ``None`` (default), 

79 the catalog will be generated using Cartesian geometry. 

80 number : `int`, optional 

81 The number of times to generate each unique combination of input 

82 parameters. The default is 1 (i.e., no repeats). This will be ignored 

83 if ``density`` is specified. 

84 density : `int` | None, optional 

85 The desired source density in sources per square degree. If given, the 

86 ``number`` parameter will be ignored. Instead, the number of unique 

87 parameter combination generations will be calculated to achieve the 

88 desired density. The default is `None` (i.e., no density calculation). 

89 seed : `Any`, optional 

90 The seed to use for the Halton sequence. If not given or ``None`` 

91 (default), the seed will be set using the product of the right 

92 ascension and declination limit ranges. 

93 log_level : `int`, optional 

94 The log level to use for logging. 

95 **kwargs : `Any` 

96 The input parameters used to generate the catalog. Each parameter key 

97 will be used as a column name in the catalog. The values are the unique 

98 values for that parameter. The output catalog will contain a row for 

99 each unique combination of input parameters and be generated the number 

100 of times specified by ``number``. 

101 

102 Returns 

103 ------- 

104 table : `astropy.table.Table` 

105 The fully populated synthetic source injection catalog. The catalog 

106 will contain an automatically generated ``injection_id`` column that 

107 is unique for each source. The injection ID encodes two pieces of 

108 information: the unique source identification number and the repeated 

109 version number of the source as defined by the ``number`` parameter. 

110 """ 

111 # Instantiate logger. 

112 logger = logging.getLogger(__name__) 

113 logger.setLevel(log_level) 

114 

115 # Parse optional keyword input parameters. 

116 values: list[Any] = [np.atleast_1d(x) for x in kwargs.values()] 

117 

118 # Determine the BBox limits and pixel scale. 

119 if wcs: 

120 sky_corners = list(itertools.product(ra_lim, dec_lim)) 

121 ra_corners, dec_corners = np.array(sky_corners).T 

122 x_corners, y_corners = wcs.skyToPixelArray(ra_corners, dec_corners, degrees=True) 

123 xlim: Any = np.percentile(x_corners, [0, 100]) 

124 ylim: Any = np.percentile(y_corners, [0, 100]) 

125 else: 

126 xlim, ylim = ra_lim, dec_lim 

127 

128 # Automatically calculate the number of generations if density is given. 

129 if density: 

130 dec_lim_rad = np.deg2rad(dec_lim) 

131 area = ((180 / np.pi) * np.diff(ra_lim) * (np.sin(dec_lim_rad[1]) - np.sin(dec_lim_rad[0])))[0] 

132 rows = list(itertools.product(*values)) 

133 native_density = len(rows) / area 

134 number = np.round(density / native_density).astype(int) 

135 if number > 0: 

136 logger.info( 

137 "Setting number of generations to %s, equivalent to %.1f sources per square degree.", 

138 number, 

139 number * native_density, 

140 ) 

141 else: 

142 logger.warning("Requested source density would require number < 1; setting number = 1.") 

143 number = 1 

144 

145 # Generate the fully expanded parameter table. 

146 values.append(range(number)) 

147 keys = list(kwargs.keys()) 

148 keys.append("version_id") 

149 param_table = Table(rows=list(itertools.product(*values)), names=keys) 

150 

151 # Generate on-sky coordinate pairs. 

152 if not seed: 

153 seed = str(np.diff(ra_lim)[0] * np.diff(dec_lim)[0]) 

154 # Random seed is the lower 32 bits of the hashed name. 

155 # We use hashlib.sha256 for guaranteed repeatability. 

156 hex_hash = hashlib.sha256(seed.encode("UTF-8")).hexdigest() 

157 hashed_seed = int("0x" + hex_hash, 0) & 0xFFFFFFFF 

158 sampler = qmc.Halton(d=2, seed=hashed_seed) 

159 sample = sampler.random(n=len(param_table)) 

160 # Flip RA values if no WCS given. 

161 if not wcs: 

162 sample[:, 0] = 1 - sample[:, 0] 

163 xy_coords = Table(qmc.scale(sample, [xlim[0], ylim[0]], [xlim[1], ylim[1]]), names=("x", "y")) 

164 if wcs: 

165 ra_coords, dec_coords = wcs.pixelToSkyArray(xy_coords["x"], xy_coords["y"], degrees=True) 

166 sky_coords = Table([ra_coords, dec_coords], names=("ra", "dec")) 

167 else: 

168 sky_coords = Table(xy_coords, names=("ra", "dec")) 

169 # Perform an additional random permutation of the sky coordinate pairs to 

170 # minimize the potential for on-sky parameter correlations. 

171 rng = np.random.default_rng(hashed_seed) 

172 sky_coords = Table(rng.permutation(sky_coords)) 

173 

174 # Generate the unique injection ID and construct the final table. 

175 source_id = np.concatenate([([i] * number) for i in range(int(len(param_table) / number))]) 

176 injection_id = param_table["version_id"] + source_id * int(10 ** np.ceil(np.log10(number))) 

177 injection_id.name = "injection_id" 

178 table = hstack([injection_id, sky_coords, param_table]) 

179 table.remove_column("version_id") 

180 

181 # Final logger report and return. 

182 if number == 1: 

183 extra_info = f"{len(table)} unique sources." 

184 else: 

185 num_combinations = int(len(table) / number) 

186 grammar = "combination" if num_combinations == 1 else "combinations" 

187 extra_info = f"{len(table)} sources: {num_combinations} {grammar} repeated {number} times." 

188 logger.info("Generated an injection catalog containing %s", extra_info) 

189 return table