Coverage for python/lsst/source/injection/utils/ingest_injection_catalog.py: 26%

23 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-18 12:18 +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/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ["ingest_injection_catalog"] 

25 

26import logging 

27 

28import esutil 

29from astropy.table import Table, vstack 

30from lsst.daf.butler import Butler, CollectionType, DatasetRef, DatasetType, DimensionUniverse 

31 

32 

33def ingest_injection_catalog( 

34 writeable_butler: Butler, 

35 table: Table | list[Table], 

36 band: str, 

37 output_collection: str, 

38 dataset_type_name: str = "injection_catalog", 

39 log_level: int = logging.INFO, 

40) -> list[DatasetRef]: 

41 """Ingest a source table into the butler. 

42 

43 This function ingests either a single astropy Table or list of Tables into 

44 the Butler. If a list of Tables is provided, these will be vertically 

45 stacked together into one single Table for ingestion. Input source tables 

46 are expected to contain the columns ``ra`` and ``dec``, with data in units 

47 of degrees. This spatial information will be used to shard up the source 

48 table on-disk using a depth 7 Hierarchical Triangular Mesh (HTM7) format. 

49 HTM7 trixels have an area of ~0.315 square degreees. 

50 

51 Parameters 

52 ---------- 

53 writeable_butler : `lsst.daf.butler.Butler` 

54 An instantiated writeable butler. 

55 table : `astropy.table.Table` or `list` [`astropy.table.Table`] 

56 Input source table(s). Requires columns ``ra`` and ``dec`` in degrees. 

57 band : `str` 

58 Band associated with the input source table(s). 

59 output_collection : `str` 

60 Name of the output collection to ingest the consolidated table into. 

61 dataset_type_name : `str`, optional 

62 Dataset type name for the ingested consolidated table. 

63 log_level : `int`, optional 

64 The log level to use for logging. 

65 

66 Returns 

67 ------- 

68 dataset_refs : `list` [`lsst.daf.butler.DatasetRef`] 

69 List containing the dataset refs for the ingested source table. 

70 """ 

71 # Instantiate logger. 

72 logger = logging.getLogger(__name__) 

73 logger.setLevel(log_level) 

74 

75 # Concatenate inputs to a single Table if required. 

76 if isinstance(table, list): 

77 table = vstack(table) 

78 

79 # Register the dataset type and collection. 

80 table_dataset_type = DatasetType( 

81 dataset_type_name, 

82 ("htm7", "band"), 

83 "ArrowAstropy", 

84 universe=DimensionUniverse(), 

85 ) 

86 writeable_butler.registry.registerDatasetType(table_dataset_type) 

87 writeable_butler.registry.registerCollection( 

88 name=output_collection, 

89 type=CollectionType.RUN, 

90 ) 

91 

92 # Generate HTM trixel indices. 

93 htm = esutil.htm.HTM(7) 

94 htm_indices = htm.lookup_id(table["ra"], table["dec"]) # type: ignore 

95 

96 # Loop over all unique trixel indices and ingest. 

97 dataset_refs: list[DatasetRef] = [] 

98 for htm_index in set(htm_indices): 

99 table_subset = table[htm_indices == htm_index] 

100 dataset_ref = writeable_butler.put( 

101 table_subset, 

102 table_dataset_type, 

103 {"htm7": htm_index, "band": band}, 

104 run=output_collection, 

105 ) 

106 dataset_refs.append(dataset_ref) 

107 

108 # Log results and return. 

109 logger.info( 

110 "Ingested %d %s-band %s DatasetRef%s into: %s", 

111 len(dataset_refs), 

112 band, 

113 dataset_type_name, 

114 "" if len(dataset_refs) == 1 else "s", 

115 output_collection, 

116 ) 

117 return dataset_refs