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