Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# This file is part of ctrl_mpexec. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://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 <http://www.gnu.org/licenses/>. 

21 

22"""Module defining few methods to generate GraphViz diagrams from pipelines 

23or quantum graphs. 

24""" 

25 

26__all__ = ["graph2dot", "pipeline2dot"] 

27 

28# ------------------------------- 

29# Imports of standard modules -- 

30# ------------------------------- 

31 

32# ----------------------------- 

33# Imports for other modules -- 

34# ----------------------------- 

35from lsst.daf.butler import DimensionUniverse 

36from lsst.pipe.base import iterConnections, Pipeline 

37 

38# ---------------------------------- 

39# Local non-exported definitions -- 

40# ---------------------------------- 

41 

42 

43def _renderTaskNode(nodeName, taskDef, file, idx=None): 

44 """Render GV node for a task""" 

45 label = [taskDef.taskName.rpartition('.')[-1]] 

46 if idx is not None: 

47 label += ["index: {}".format(idx)] 

48 if taskDef.label: 

49 label += ["label: {}".format(taskDef.label)] 

50 label = r'\n'.join(label) 

51 attrib = dict(shape="box", 

52 style="filled,bold", 

53 fillcolor="gray70", 

54 label=label) 

55 attrib = ['{}="{}"'.format(key, val) for key, val in attrib.items()] 

56 print("{} [{}];".format(nodeName, ", ".join(attrib)), file=file) 

57 

58 

59def _renderDSTypeNode(name, dimensions, file): 

60 """Render GV node for a dataset type""" 

61 label = [name] 

62 if dimensions: 

63 label += ["Dimensions: " + ", ".join(dimensions)] 

64 label = r'\n'.join(label) 

65 attrib = dict(shape="box", 

66 style="rounded,filled", 

67 fillcolor="gray90", 

68 label=label) 

69 attrib = ['{}="{}"'.format(key, val) for key, val in attrib.items()] 

70 print("{} [{}];".format(name, ", ".join(attrib)), file=file) 

71 

72 

73def _renderDSNode(nodeName, dsRef, file): 

74 """Render GV node for a dataset""" 

75 label = [dsRef.datasetType.name] 

76 for key in sorted(dsRef.dataId.keys()): 

77 label += [str(key) + "=" + str(dsRef.dataId[key])] 

78 label = r'\n'.join(label) 

79 attrib = dict(shape="box", 

80 style="rounded,filled", 

81 fillcolor="gray90", 

82 label=label) 

83 attrib = ['{}="{}"'.format(key, val) for key, val in attrib.items()] 

84 print("{} [{}];".format(nodeName, ", ".join(attrib)), file=file) 

85 

86 

87def _datasetRefId(dsRef): 

88 """Make an idetifying string for given ref""" 

89 idStr = str(dsRef.datasetType.name) 

90 for key in sorted(dsRef.dataId.keys()): 

91 idStr += ":" + str(key) + "=" + str(dsRef.dataId[key]) 

92 return idStr 

93 

94 

95def _makeDSNode(dsRef, allDatasetRefs, file): 

96 """Make new node for dataset if it does not exist. 

97 

98 Returns node name. 

99 """ 

100 dsRefId = _datasetRefId(dsRef) 

101 nodeName = allDatasetRefs.get(dsRefId) 

102 if nodeName is None: 

103 idx = len(allDatasetRefs) 

104 nodeName = "dsref_{}".format(idx) 

105 allDatasetRefs[dsRefId] = nodeName 

106 _renderDSNode(nodeName, dsRef, file) 

107 return nodeName 

108 

109# ------------------------ 

110# Exported definitions -- 

111# ------------------------ 

112 

113 

114def graph2dot(qgraph, file): 

115 """Convert QuantumGraph into GraphViz digraph. 

116 

117 This method is mostly for documentation/presentation purposes. 

118 

119 Parameters 

120 ---------- 

121 qgraph: `pipe.base.QuantumGraph` 

122 QuantumGraph instance. 

123 file : str or file object 

124 File where GraphViz graph (DOT language) is written, can be a file name 

125 or file object. 

126 

127 Raises 

128 ------ 

129 `OSError` is raised when output file cannot be open. 

130 `ImportError` is raised when task class cannot be imported. 

131 """ 

132 # open a file if needed 

133 close = False 

134 if not hasattr(file, "write"): 

135 file = open(file, "w") 

136 close = True 

137 

138 print("digraph QuantumGraph {", file=file) 

139 

140 allDatasetRefs = {} 

141 for taskId, nodes in enumerate(qgraph): 

142 

143 taskDef = nodes.taskDef 

144 

145 for qId, quantum in enumerate(nodes.quanta): 

146 

147 # node for a task 

148 taskNodeName = "task_{}_{}".format(taskId, qId) 

149 _renderTaskNode(taskNodeName, taskDef, file) 

150 

151 # quantum inputs 

152 for dsRefs in quantum.predictedInputs.values(): 

153 for dsRef in dsRefs: 

154 nodeName = _makeDSNode(dsRef, allDatasetRefs, file) 

155 print("{} -> {};".format(nodeName, taskNodeName), file=file) 

156 

157 # quantum outputs 

158 for dsRefs in quantum.outputs.values(): 

159 for dsRef in dsRefs: 

160 nodeName = _makeDSNode(dsRef, allDatasetRefs, file) 

161 print("{} -> {};".format(taskNodeName, nodeName), file=file) 

162 

163 print("}", file=file) 

164 if close: 

165 file.close() 

166 

167 

168def pipeline2dot(pipeline, file): 

169 """Convert Pipeline into GraphViz digraph. 

170 

171 This method is mostly for documentation/presentation purposes. 

172 Unlike other methods this method does not validate graph consistency. 

173 

174 Parameters 

175 ---------- 

176 pipeline : `pipe.base.Pipeline` 

177 Pipeline description. 

178 file : str or file object 

179 File where GraphViz graph (DOT language) is written, can be a file name 

180 or file object. 

181 

182 Raises 

183 ------ 

184 `OSError` is raised when output file cannot be open. 

185 `ImportError` is raised when task class cannot be imported. 

186 `MissingTaskFactoryError` is raised when TaskFactory is needed but not 

187 provided. 

188 """ 

189 universe = DimensionUniverse() 

190 

191 def expand_dimensions(dimensions): 

192 """Returns expanded list of dimensions, with special skypix treatment. 

193 

194 Parameters 

195 ---------- 

196 dimensions : `list` [`str`] 

197 

198 Returns 

199 ------- 

200 dimensions : `list` [`str`] 

201 """ 

202 dimensions = set(dimensions) 

203 skypix_dim = [] 

204 if "skypix" in dimensions: 

205 dimensions.remove("skypix") 

206 skypix_dim = ["skypix"] 

207 dimensions = universe.extract(dimensions) 

208 return list(dimensions.names) + skypix_dim 

209 

210 # open a file if needed 

211 close = False 

212 if not hasattr(file, "write"): 

213 file = open(file, "w") 

214 close = True 

215 

216 print("digraph Pipeline {", file=file) 

217 

218 allDatasets = set() 

219 if isinstance(pipeline, Pipeline): 

220 pipeline = pipeline.toExpandedPipeline() 

221 for idx, taskDef in enumerate(pipeline): 

222 

223 # node for a task 

224 taskNodeName = "task{}".format(idx) 

225 _renderTaskNode(taskNodeName, taskDef, file, idx) 

226 

227 for attr in iterConnections(taskDef.connections, 'inputs'): 

228 if attr.name not in allDatasets: 

229 dimensions = expand_dimensions(attr.dimensions) 

230 _renderDSTypeNode(attr.name, dimensions, file) 

231 allDatasets.add(attr.name) 

232 print("{} -> {};".format(attr.name, taskNodeName), file=file) 

233 

234 for attr in iterConnections(taskDef.connections, 'prerequisiteInputs'): 

235 if attr.name not in allDatasets: 

236 dimensions = expand_dimensions(attr.dimensions) 

237 _renderDSTypeNode(attr.name, dimensions, file) 

238 allDatasets.add(attr.name) 

239 # use dashed line for prerequisite edges to distinguish them 

240 print("{} -> {} [style = dashed];".format(attr.name, taskNodeName), file=file) 

241 

242 for attr in iterConnections(taskDef.connections, 'outputs'): 

243 if attr.name not in allDatasets: 

244 dimensions = expand_dimensions(attr.dimensions) 

245 _renderDSTypeNode(attr.name, dimensions, file) 

246 allDatasets.add(attr.name) 

247 print("{} -> {};".format(taskNodeName, attr.name), file=file) 

248 

249 print("}", file=file) 

250 if close: 

251 file.close()