Coverage for python/lsst/source/injection/bin/generate_injection_catalog.py: 9%
93 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-23 13:55 +0000
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-23 13:55 +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 "-n",
91 "--number",
92 type=int,
93 help="Number of generated parameter combinations. Ignored if density given.",
94 metavar="VALUE",
95 default=1,
96 )
97 parser_general.add_argument(
98 "-s",
99 "--density",
100 type=int,
101 help="Desired source density (N/deg^2). If given, number option is ignored.",
102 metavar="VALUE",
103 )
104 parser_general.add_argument(
105 "-p",
106 "--parameter",
107 help="An input parameter definition.",
108 metavar=("COLNAME VALUE", "VALUE"),
109 nargs="+",
110 action="append",
111 )
112 parser_general.add_argument(
113 "--seed",
114 type=str,
115 help="Seed override when generating quasi-random RA/Dec positions.",
116 metavar="SEED",
117 )
119 # Butler options.
120 parser_butler = parser.add_argument_group("Butler Options")
121 parser_butler.add_argument(
122 "-b",
123 "--butler-config",
124 type=str,
125 help="Location of the butler/registry config file.",
126 metavar="TEXT",
127 )
128 # WCS options.
129 parser_wcs = parser.add_argument_group("WCS Options")
130 parser_wcs.add_argument(
131 "-w",
132 "--wcs-type-name",
133 help="Dataset type containing a `wcs` component for WCS spatial conversions.",
134 metavar="TEXT",
135 )
136 parser_wcs.add_argument(
137 "-c",
138 "--collections",
139 type=str,
140 help="Collections to query for dataset types containing a WCS component.",
141 metavar="COLL",
142 )
143 parser_wcs.add_argument(
144 "--where",
145 type=str,
146 help="A string expression similar to an SQL WHERE clause to query for datasets with a WCS component.",
147 metavar="COLL",
148 )
149 # Ingestion options.
150 parser_ingest = parser.add_argument_group("Repository Ingestion Options")
151 parser_ingest.add_argument(
152 "-i",
153 "--injection-band",
154 type=str,
155 help="Band(s) associated with the generated table.",
156 metavar=("BAND", "BAND"),
157 nargs="+",
158 )
159 parser_ingest.add_argument(
160 "-o",
161 "--output-collection",
162 type=str,
163 help="Name of the output collection to ingest the injection catalog into.",
164 metavar="COLL",
165 )
166 parser_ingest.add_argument(
167 "-t",
168 "--dataset-type-name",
169 type=str,
170 help="Output dataset type name for the ingested source injection catalog.",
171 metavar="TEXT",
172 default="injection_catalog",
173 )
174 # Options to write table to disk.
175 parser_write = parser.add_argument_group("Write to Disk Options")
176 parser_write.add_argument(
177 "-f",
178 "--filename",
179 help="Output filename for the generated injection catalog.",
180 metavar="TEXT",
181 )
182 parser_write.add_argument(
183 "--format",
184 help="Output injection catalog format, overriding automatic format selection.",
185 metavar="TEXT",
186 )
187 parser_write.add_argument(
188 "--overwrite",
189 help="Overwrite the output file if it already exists.",
190 action="store_true",
191 )
192 # Help.
193 parser_misc = parser.add_argument_group("Miscellaneous Options")
194 parser_misc.add_argument(
195 "-h",
196 "--help",
197 action="help",
198 help="Show this help message and exit.",
199 )
200 return parser
203def main():
204 """Use this as the main entry point when calling from the command line."""
205 # Set up logging.
206 tz = time.strftime("%z")
207 logging.basicConfig(
208 format="%(levelname)s %(asctime)s.%(msecs)03d" + tz + " - %(message)s", datefmt="%Y-%m-%dT%H:%M:%S"
209 )
210 logger = logging.getLogger(__name__)
211 logger.setLevel(logging.DEBUG)
213 args = build_argparser().parse_args()
215 # Validate all butler options are provided.
216 butler_config = vars(args).get("butler_config", "")
217 wcs_type_name = vars(args).get("wcs_type_name", "")
218 collections = vars(args).get("collections", "")
219 num_wcs = sum([bool(wcs_type_name), bool(collections)])
220 if num_wcs == 1 or (num_wcs == 2 and not butler_config):
221 raise RuntimeError("A butler query for WCS requires a butler repo, a dataset type and a collection.")
222 injection_band = vars(args).get("injection_band", "")
223 output_collection = vars(args).get("output_collection", "")
224 dataset_type_name = vars(args).get("dataset_type_name", "") # Defaults to "injection_catalog".
225 num_ingest = sum([bool(injection_band), bool(output_collection)])
226 if num_ingest == 1 or (num_ingest == 2 and not butler_config):
227 raise RuntimeError("Catalog ingestion requires a butler repo, a band and an output collection.")
229 # Parse the input parameters.
230 params = {}
231 if hasattr(args, "parameter"):
232 for param in args.parameter:
233 if len(param) < 2:
234 raise RuntimeError("Each parameter must be associated with at least one value.")
235 name = param[0]
236 try:
237 values = [float(x) for x in param[1:]]
238 except ValueError:
239 values = param[1:]
240 params[name] = values
242 # Get the input WCS.
243 if not wcs_type_name:
244 wcs = None
245 logger.info("No WCS provided, source positions generated using Cartesian geometry.")
246 else:
247 butler = Butler(butler_config)
248 where = vars(args).get("where", "") # Optional where query.
249 try:
250 dataset_ref = list(
251 butler.registry.queryDatasets(
252 datasetType=wcs_type_name,
253 collections=collections,
254 where=where,
255 findFirst=True,
256 )
257 )[0]
258 except IndexError:
259 raise RuntimeError(f"No {wcs_type_name} dataset type found in {args.collections} for: {where}.")
260 ddhandle = butler.getDeferred(wcs_type_name, dataId=dataset_ref.dataId, collections=collections)
261 try:
262 wcs = ddhandle.get(component="wcs")
263 except KeyError:
264 raise RuntimeError(f"No WCS component found for {wcs_type_name} dataset type.")
265 logger.info("Using WCS in %s for %s.", wcs_type_name, dataset_ref.dataId)
267 # Generate the source injection catalog.
268 density = vars(args).get("density", None)
269 seed = vars(args).get("seed", None)
270 table = generate_injection_catalog(
271 ra_lim=args.ra_lim,
272 dec_lim=args.dec_lim,
273 wcs=wcs,
274 number=args.number,
275 density=density,
276 seed=seed,
277 **params,
278 )
280 # Save table to disk.
281 filename = vars(args).get("filename", False)
282 file_format = vars(args).get("format", None)
283 overwrite = vars(args).get("overwrite", False)
284 if filename:
285 table.write(filename, format=file_format, overwrite=overwrite)
286 logger.info("Written injection catalog to '%s'.", filename)
288 # Ingest table into a butler repo.
289 if injection_band:
290 writeable_butler = Butler(butler_config, writeable=True)
291 for band in injection_band:
292 _ = ingest_injection_catalog(
293 writeable_butler=writeable_butler,
294 table=table,
295 band=band,
296 output_collection=output_collection,
297 dataset_type_name=dataset_type_name,
298 )
300 # Print the table to stdout.
301 if not filename and not injection_band:
302 print(table)