Coverage for python/lsst/analysis/tools/tasks/sourceObjectTableAnalysis.py: 38%

47 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-24 04:09 -0700

1from __future__ import annotations 

2 

3__all__ = ("SourceObjectTableAnalysisConfig", "SourceObjectTableAnalysisTask") 

4 

5import lsst.pex.config as pexConfig 

6import numpy as np 

7import pandas as pd 

8from astropy.table import vstack 

9from lsst.pipe.base import connectionTypes as ct 

10from smatch import Matcher 

11 

12from ..interfaces import AnalysisBaseConfig, AnalysisBaseConnections, AnalysisPipelineTask 

13 

14 

15class SourceObjectTableAnalysisConnections( 

16 AnalysisBaseConnections, 

17 dimensions=("visit",), 

18 defaultTemplates={ 

19 "inputName": "sourceTable_visit", 

20 "inputCoaddName": "deep", 

21 "associatedSourcesInputName": "isolated_star_sources", 

22 "outputName": "sourceObjectTable", 

23 }, 

24): 

25 data = ct.Input( 

26 doc="Visit based source table to load from the butler", 

27 name="sourceTable_visit", 

28 storageClass="ArrowAstropy", 

29 dimensions=("visit", "band"), 

30 deferLoad=True, 

31 ) 

32 

33 associatedSources = ct.Input( 

34 doc="Table of associated sources", 

35 name="{associatedSourcesInputName}", 

36 storageClass="ArrowAstropy", 

37 multiple=True, 

38 deferLoad=True, 

39 dimensions=("instrument", "skymap", "tract"), 

40 ) 

41 

42 refCat = ct.Input( 

43 doc="Catalog of positions to use as reference.", 

44 name="objectTable", 

45 storageClass="DataFrame", 

46 dimensions=["skymap", "tract", "patch"], 

47 multiple=True, 

48 deferLoad=True, 

49 ) 

50 

51 

52class SourceObjectTableAnalysisConfig( 

53 AnalysisBaseConfig, pipelineConnections=SourceObjectTableAnalysisConnections 

54): 

55 ra_column = pexConfig.Field( 

56 doc="Name of column in refCat to use for right ascension.", 

57 dtype=str, 

58 default="r_ra", 

59 ) 

60 dec_column = pexConfig.Field( 

61 doc="Name of column in refCat to use for declination.", 

62 dtype=str, 

63 default="r_dec", 

64 ) 

65 

66 

67class SourceObjectTableAnalysisTask(AnalysisPipelineTask): 

68 ConfigClass = SourceObjectTableAnalysisConfig 

69 _DefaultName = "sourceTableVisitAnalysis" 

70 

71 def runQuantum(self, butlerQC, inputRefs, outputRefs): 

72 inputs = butlerQC.get(inputRefs) 

73 

74 # Get isolated sources: 

75 visit = inputs["data"].dataId["visit"] 

76 band = inputs["data"].dataId["band"] 

77 names = self.collectInputNames() 

78 names -= {self.config.ra_column, self.config.dec_column} 

79 data = inputs["data"].get(parameters={"columns": names}) 

80 

81 dataId = butlerQC.quantum.dataId 

82 plotInfo = self.parsePlotInfo(inputs, dataId) 

83 

84 isolatedSources = [] 

85 for associatedSourcesRef in inputs["associatedSources"]: 

86 associatedSources = associatedSourcesRef.get(parameters={"columns": ["visit", "source_row"]}) 

87 visit_sources = associatedSources[associatedSources["visit"] == visit] 

88 isolatedSources.append(data[visit_sources["source_row"]]) 

89 isolatedSources = vstack(isolatedSources) 

90 

91 # Get objects: 

92 allRefCats = [] 

93 for refCatRef in inputs["refCat"]: 

94 refCat = refCatRef.get( 

95 parameters={"columns": ["detect_isPrimary", self.config.ra_column, self.config.dec_column]} 

96 ) 

97 goodInds = ( 

98 refCat["detect_isPrimary"] 

99 & np.isfinite(refCat[self.config.ra_column]) 

100 & np.isfinite(refCat[self.config.dec_column]) 

101 ) 

102 allRefCats.append(refCat[goodInds]) 

103 

104 refCat = pd.concat(allRefCats) 

105 

106 with Matcher(isolatedSources["coord_ra"], isolatedSources["coord_dec"]) as m: 

107 idx, isolatedMatchIndices, refMatchIndices, dists = m.query_radius( 

108 refCat[self.config.ra_column].values, 

109 refCat[self.config.dec_column].values, 

110 1 / 3600.0, 

111 return_indices=True, 

112 ) 

113 

114 matchRef = refCat.iloc[refMatchIndices] 

115 matchIS = isolatedSources[isolatedMatchIndices].to_pandas() 

116 

117 allCat = pd.concat([matchRef.reset_index(), matchIS.reset_index()], axis=1) 

118 outputs = self.run(data=allCat, bands=band, plotInfo=plotInfo) 

119 butlerQC.put(outputs, outputRefs)