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

68 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-03-22 04:07 -0700

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 mag_lim: Sequence[float] | None = None, 

42 wcs: SkyWcs = None, 

43 number: int = 1, 

44 density: int | None = None, 

45 seed: Any = None, 

46 log_level: int = logging.INFO, 

47 **kwargs: Any, 

48) -> Table: 

49 """Generate a synthetic source injection catalog. 

50 

51 This function generates a synthetic source injection catalog from user 

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

53 

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

55 sequence. Optional magnitudes may also be generated using the same 

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

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

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

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

60 

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

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

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

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

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

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

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

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

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

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

71 list of integers. 

72 

73 Parameters 

74 ---------- 

75 ra_lim : `Sequence` [`float`] 

76 The right ascension limits of the catalog in degrees. 

77 dec_lim : `Sequence` [`float`] 

78 The declination limits of the catalog in degrees. 

79 mag_lim : `Sequence` [`float`], optional 

80 The magnitude limits of the catalog in magnitudes. 

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

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

83 the catalog will be generated using Cartesian geometry. 

84 number : `int`, optional 

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

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

87 if ``density`` is specified. 

88 density : `int` | None, optional 

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

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

91 parameter combination generations will be calculated to achieve the 

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

93 seed : `Any`, optional 

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

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

96 ascension and declination limit ranges. 

97 log_level : `int`, optional 

98 The log level to use for logging. 

99 **kwargs : `Any` 

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

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

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

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

104 of times specified by ``number``. 

105 

106 Returns 

107 ------- 

108 table : `astropy.table.Table` 

109 The fully populated synthetic source injection catalog. The catalog 

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

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

112 information: the unique source identification number and the repeated 

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

114 """ 

115 # Instantiate logger. 

116 logger = logging.getLogger(__name__) 

117 logger.setLevel(log_level) 

118 

119 # Parse optional keyword input parameters. 

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

121 

122 # Determine the BBox limits and pixel scale. 

123 if wcs: 

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

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

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

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

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

129 else: 

130 xlim, ylim = ra_lim, dec_lim 

131 

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

133 if density: 

134 dec_lim_rad = np.deg2rad(dec_lim) 

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

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

137 native_density = len(rows) / area 

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

139 if number > 0: 

140 logger.info( 

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

142 number, 

143 number * native_density, 

144 ) 

145 else: 

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

147 number = 1 

148 

149 # Generate the fully expanded parameter table. 

150 values.append(range(number)) 

151 keys = list(kwargs.keys()) 

152 keys.append("version_id") 

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

154 

155 # Generate on-sky coordinate pairs. 

156 if not seed: 

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

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

159 # We use hashlib.sha256 for guaranteed repeatability. 

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

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

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

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

164 # Flip RA values if no WCS given. 

165 if not wcs: 

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

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

168 if wcs: 

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

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

171 else: 

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

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

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

175 rng = np.random.default_rng(hashed_seed) 

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

177 

178 # Generate random magnitudes if limits are specified 

179 if mag_lim: 

180 mag_sampler = qmc.Halton(d=1, seed=hashed_seed) 

181 mag_sample = mag_sampler.random(n=len(param_table)) 

182 mags = Table(qmc.scale(mag_sample, mag_lim[0], mag_lim[1]), names=("mag",)) 

183 sky_coords = hstack([sky_coords, mags]) 

184 

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

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

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

188 injection_id.name = "injection_id" 

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

190 table.remove_column("version_id") 

191 

192 # Final logger report and return. 

193 if number == 1: 

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

195 else: 

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

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

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

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

200 return table