Coverage for python / lsst / source / injection / bin / generate_injection_catalog.py: 8%
94 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-17 09:38 +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
24import logging
25import time
26from argparse import SUPPRESS, ArgumentParser
28from lsst.daf.butler import Butler
30from ..utils import generate_injection_catalog, ingest_injection_catalog
31from .source_injection_help_formatter import SourceInjectionHelpFormatter
34def build_argparser():
35 """Build an argument parser for this script."""
36 parser = ArgumentParser(
37 description="""Generate a synthetic source injection catalog.
39This script generates a synthetic source injection catalog from user supplied
40input parameters. The catalog may be printed to screen, ingested into a butler
41repository, or written to disk using the astropy Table API.
43On-sky source positions are generated using the quasi-random Halton sequence.
44By default, the Halton sequence is seeded using the product of the right
45ascension and declination limit ranges. This ensures that the same sequence is
46always generated for the same limits. This seed may be overridden by the user.
48A unique injection ID is generated for each source. The injection ID encodes
49two pieces of information: the unique source identification number and the
50version number of the source as specified by the ``number`` parameter. To
51achieve this, the unique source ID number is multiplied by `10**n` such that
52the sum of the multiplied source ID number with the unique repeated version
53number will always be unique. For example, an injection catalog with 3 versions
54of each source will have injection IDs = 0, 1, 2, 10, 11, 12, 20, 21, 22, etc.
55For number = 20, injection IDs = 0, 1, 2, ..., 17, 18, 19, 100, 101, 102, etc.
56If number = 1 (the default) then the injection ID will be a simple sequential
57list of integers.
59An optional butler query for a WCS dataset type may be provided for use in
60generating on-sky source positions. If no WCS is provided, the source positions
61will be generated using Cartesian geometry.
62""",
63 formatter_class=SourceInjectionHelpFormatter,
64 epilog="More information is available at https://pipelines.lsst.io.",
65 add_help=False,
66 argument_default=SUPPRESS,
67 )
68 # General options.
69 parser_general = parser.add_argument_group("General Options")
70 parser_general.add_argument(
71 "-a",
72 "--ra-lim",
73 type=float,
74 help="Right ascension limits of the catalog in degrees.",
75 required=True,
76 metavar="VALUE",
77 nargs=2,
78 )
79 parser_general.add_argument(
80 "-d",
81 "--dec-lim",
82 type=float,
83 help="Declination limits of the catalog in degrees.",
84 required=True,
85 metavar="VALUE",
86 nargs=2,
87 )
88 parser_general.add_argument(
89 "-m",
90 "--mag-lim",
91 type=float,
92 help="The magnitude limits of the catalog in magnitudes.",
93 required=False,
94 metavar="VALUE",
95 nargs=2,
96 )
97 parser_general.add_argument(
98 "-n",
99 "--number",
100 type=int,
101 help="Number of generated parameter combinations. Ignored if density given.",
102 metavar="VALUE",
103 default=1,
104 )
105 parser_general.add_argument(
106 "-s",
107 "--density",
108 type=int,
109 help="Desired source density (N/deg^2). If given, number option is ignored.",
110 metavar="VALUE",
111 )
112 parser_general.add_argument(
113 "-p",
114 "--parameter",
115 help="An input parameter definition.",
116 metavar=("COLNAME VALUE", "VALUE"),
117 nargs="+",
118 action="append",
119 )
120 parser_general.add_argument(
121 "--seed",
122 type=str,
123 help="Seed override when generating quasi-random RA/Dec positions.",
124 metavar="SEED",
125 )
127 # Butler options.
128 parser_butler = parser.add_argument_group("Butler Options")
129 parser_butler.add_argument(
130 "-b",
131 "--butler-config",
132 type=str,
133 help="Location of the butler/registry config file.",
134 metavar="TEXT",
135 )
136 # WCS options.
137 parser_wcs = parser.add_argument_group("WCS Options")
138 parser_wcs.add_argument(
139 "-w",
140 "--wcs-type-name",
141 help="Dataset type containing a `wcs` component for WCS spatial conversions.",
142 metavar="TEXT",
143 )
144 parser_wcs.add_argument(
145 "-c",
146 "--collections",
147 type=str,
148 help="Collections to query for dataset types containing a WCS component.",
149 metavar="COLL",
150 )
151 parser_wcs.add_argument(
152 "--where",
153 type=str,
154 help="A string expression similar to an SQL WHERE clause to query for datasets with a WCS component.",
155 metavar="COLL",
156 )
157 # Ingestion options.
158 parser_ingest = parser.add_argument_group("Repository Ingestion Options")
159 parser_ingest.add_argument(
160 "-i",
161 "--injection-band",
162 type=str,
163 help="Band(s) associated with the generated table.",
164 metavar=("BAND", "BAND"),
165 nargs="+",
166 )
167 parser_ingest.add_argument(
168 "-o",
169 "--output-collection",
170 type=str,
171 help="Name of the output collection to ingest the injection catalog into.",
172 metavar="COLL",
173 )
174 parser_ingest.add_argument(
175 "-t",
176 "--dataset-type-name",
177 type=str,
178 help="Output dataset type name for the ingested source injection catalog.",
179 metavar="TEXT",
180 default="injection_catalog",
181 )
182 # Options to write table to disk.
183 parser_write = parser.add_argument_group("Write to Disk Options")
184 parser_write.add_argument(
185 "-f",
186 "--filename",
187 help="Output filename for the generated injection catalog.",
188 metavar="TEXT",
189 )
190 parser_write.add_argument(
191 "--format",
192 help="Output injection catalog format, overriding automatic format selection.",
193 metavar="TEXT",
194 )
195 parser_write.add_argument(
196 "--overwrite",
197 help="Overwrite the output file if it already exists.",
198 action="store_true",
199 )
200 # Help.
201 parser_misc = parser.add_argument_group("Miscellaneous Options")
202 parser_misc.add_argument(
203 "-h",
204 "--help",
205 action="help",
206 help="Show this help message and exit.",
207 )
208 return parser
211def main():
212 """Use this as the main entry point when calling from the command line."""
213 # Set up logging.
214 tz = time.strftime("%z")
215 logging.basicConfig(
216 format="%(levelname)s %(asctime)s.%(msecs)03d" + tz + " - %(message)s", datefmt="%Y-%m-%dT%H:%M:%S"
217 )
218 logger = logging.getLogger(__name__)
219 logger.setLevel(logging.DEBUG)
221 args = build_argparser().parse_args()
223 # Validate all butler options are provided.
224 butler_config = vars(args).get("butler_config", "")
225 wcs_type_name = vars(args).get("wcs_type_name", "")
226 collections = vars(args).get("collections", "")
227 num_wcs = sum([bool(wcs_type_name), bool(collections)])
228 if num_wcs == 1 or (num_wcs == 2 and not butler_config):
229 raise RuntimeError("A butler query for WCS requires a butler repo, a dataset type and a collection.")
230 injection_band = vars(args).get("injection_band", "")
231 output_collection = vars(args).get("output_collection", "")
232 dataset_type_name = vars(args).get("dataset_type_name", "") # Defaults to "injection_catalog".
233 num_ingest = sum([bool(injection_band), bool(output_collection)])
234 if num_ingest == 1 or (num_ingest == 2 and not butler_config):
235 raise RuntimeError("Catalog ingestion requires a butler repo, a band and an output collection.")
237 # Parse the input parameters.
238 params = {}
239 if hasattr(args, "parameter"):
240 for param in args.parameter:
241 if len(param) < 2:
242 raise RuntimeError("Each parameter must be associated with at least one value.")
243 name = param[0]
244 try:
245 values = [float(x) for x in param[1:]]
246 except ValueError:
247 values = param[1:]
248 params[name] = values
250 # Get the input WCS.
251 if not wcs_type_name:
252 wcs = None
253 logger.info("No WCS provided, source positions generated using Cartesian geometry.")
254 else:
255 butler = Butler.from_config(butler_config)
256 where = vars(args).get("where", "") # Optional where query.
257 try:
258 dataset_ref = list(
259 butler.registry.queryDatasets(
260 datasetType=wcs_type_name,
261 collections=collections,
262 where=where,
263 findFirst=True,
264 )
265 )[0]
266 except IndexError:
267 raise RuntimeError(f"No {wcs_type_name} dataset type found in {args.collections} for: {where}.")
268 ddhandle = butler.getDeferred(wcs_type_name, dataId=dataset_ref.dataId, collections=collections)
269 try:
270 wcs = ddhandle.get(component="wcs")
271 except KeyError:
272 raise RuntimeError(f"No WCS component found for {wcs_type_name} dataset type.")
273 logger.info("Using WCS in %s for %s.", wcs_type_name, dataset_ref.dataId)
275 # Generate the source injection catalog.
276 mag_lim = vars(args).get("mag_lim", None)
277 density = vars(args).get("density", None)
278 seed = vars(args).get("seed", None)
279 table = generate_injection_catalog(
280 ra_lim=args.ra_lim,
281 dec_lim=args.dec_lim,
282 mag_lim=mag_lim,
283 wcs=wcs,
284 number=args.number,
285 density=density,
286 seed=seed,
287 **params,
288 )
290 # Save table to disk.
291 filename = vars(args).get("filename", False)
292 file_format = vars(args).get("format", None)
293 overwrite = vars(args).get("overwrite", False)
294 if filename:
295 table.write(filename, format=file_format, overwrite=overwrite)
296 logger.info("Written injection catalog to '%s'.", filename)
298 # Ingest table into a butler repo.
299 if injection_band:
300 writeable_butler = Butler.from_config(butler_config, writeable=True)
301 for band in injection_band:
302 _ = ingest_injection_catalog(
303 writeable_butler=writeable_butler,
304 table=table,
305 band=band,
306 output_collection=output_collection,
307 dataset_type_name=dataset_type_name,
308 )
310 # Print the table to stdout.
311 if not filename and not injection_band:
312 print(table)