Coverage for python/lsst/source/injection/utils/generate_injection_catalog.py: 14%
68 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-05-01 13:10 +0000
« prev ^ index » next coverage.py v7.5.0, created at 2024-05-01 13:10 +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/>.
22from __future__ import annotations
24__all__ = ["generate_injection_catalog"]
26import hashlib
27import itertools
28import logging
29from collections.abc import Sequence
30from typing import Any
32import numpy as np
33from astropy.table import Table, hstack
34from lsst.afw.geom import SkyWcs
35from scipy.stats import qmc
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.
51 This function generates a synthetic source injection catalog from user
52 supplied input parameters. The catalog is returned as an astropy Table.
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.
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.
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``.
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)
119 # Parse optional keyword input parameters.
120 values: list[Any] = [np.atleast_1d(x) for x in kwargs.values()]
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
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
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)
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))
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])
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")
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