Coverage for python / lsst / pipe / base / quantum_graph / visualization.py: 89%

1# This file is part of pipe_base. 

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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27 

28from __future__ import annotations 

29 

30__all__ = ("QuantumGraphDotVisualizer", "QuantumGraphMermaidVisualizer", "QuantumGraphVisualizer") 

31 

32import html 

33import uuid 

34from abc import abstractmethod 

35from typing import IO, ClassVar, Generic, TypeVar 

36 

37from ..pipeline_graph import NodeType 

38from ._common import BaseQuantumGraph, BipartiteEdgeInfo, DatasetInfo, QuantumInfo 

39 

40_G = TypeVar("_G", bound=BaseQuantumGraph, contravariant=True) 

41_Q = TypeVar("_Q", bound=QuantumInfo, contravariant=True) 

42_D = TypeVar("_D", bound=DatasetInfo, contravariant=True) 

43 

44 

45class QuantumGraphVisualizer(Generic[_G, _Q, _D]): 

46 """A base class for exporting quantum graphs to graph-visualization 

47 languages. 

48 

49 Notes 

50 ----- 

51 This base class is not intended to support implementations for every 

52 possible visualization language, but it does neatly unify common logic for 

53 at least GraphViz and Mermaid. 

54 """ 

55 

56 @abstractmethod 

57 def render_header(self, qg: _G, is_bipartite: bool) -> str: 

58 """Return the beginning of a graph visualization. 

59 

60 Parameters 

61 ---------- 

62 qg : `.BaseQuantumGraph` 

63 Quantum graph to visualize. 

64 is_bipartite : `bool` 

65 Whether a bipartite graph visualization is being requested. 

66 

67 Returns 

68 ------- 

69 rendered : `str` 

70 String that starts the visualization. May contain newlines, but 

71 an additional newline will automatically be added after the string, 

72 before the first node. 

73 """ 

74 raise NotImplementedError() 

75 

76 @abstractmethod 

77 def render_footer(self, qg: _G, is_bipartite: bool) -> str: 

78 """Return the ending of a graph visualization. 

79 

80 Parameters 

81 ---------- 

82 qg : `.BaseQuantumGraph` 

83 Quantum graph to visualize. 

84 is_bipartite : `bool` 

85 Whether a bipartite graph visualization is being requested. 

86 

87 Returns 

88 ------- 

89 rendered : `str` 

90 String that ends the visualization. May contain newlines, but 

91 an additional newline will automatically be added after the string, 

92 at the end of the file. 

93 """ 

94 raise NotImplementedError() 

95 

96 @abstractmethod 

97 def render_quantum(self, quantum_id: uuid.UUID, data: _Q, is_bipartite: bool) -> str: 

98 """Return the representation of a quantum in a graph visualization. 

99 

100 Parameters 

101 ---------- 

102 quantum_id : `uuid.UUID` 

103 ID of the quantum. 

104 data : `.QuantumInfo` 

105 Mapping with additional information about the quantum. 

106 is_bipartite : `bool` 

107 Whether a bipartite graph visualization is being requested. 

108 

109 Returns 

110 ------- 

111 rendered : `str` 

112 String that represents the quantum. May contain newlines, but an 

113 additional newline will automatically be added after the string, 

114 before the next node or edge. 

115 """ 

116 raise NotImplementedError() 

117 

118 @abstractmethod 

119 def render_dataset(self, dataset_id: uuid.UUID, data: _D) -> str: 

120 """Return the representation of a dataset in a graph visualization. 

121 

122 Parameters 

123 ---------- 

124 dataset_id : `uuid.UUID` 

125 ID of the dataset. 

126 data : `.DatasetInfo` 

127 Mapping with additional information about the dataset. 

128 

129 Returns 

130 ------- 

131 rendered : `str` 

132 String that represents the dataset. May contain newlines, but an 

133 additional newline will automatically be added after the string, 

134 before the next node or edge. 

135 """ 

136 raise NotImplementedError() 

137 

138 @abstractmethod 

139 def render_edge( 

140 self, 

141 a: uuid.UUID, 

142 b: uuid.UUID, 

143 data: BipartiteEdgeInfo | None, 

144 ) -> str: 

145 """Return the representation of an edge a graph visualization. 

146 

147 Parameters 

148 ---------- 

149 a : `uuid.UUID` 

150 ID of the quantum or dataset for which this is an outgoing edge. 

151 b : `uuid.UUID` 

152 ID of the quantum or dataset for which this is an incoming edge. 

153 data : `.BipartiteEdgeInfo` 

154 Mapping with additional information about the dataset. 

155 

156 Returns 

157 ------- 

158 rendered : `str` 

159 String that represents the edge. May contain newlines, but an 

160 additional newline will automatically be added after the string, 

161 before the next edge. 

162 """ 

163 raise NotImplementedError() 

164 

165 def write_quantum_only(self, qg: _G, stream: IO[str]) -> None: 

166 """Write a visualization for graph with only quantum nodes. 

167 

168 Parameters 

169 ---------- 

170 qg : `BaseQuantumGraph` 

171 Quantum graph to visualize. 

172 stream : `typing.IO` 

173 File-like object to write to. 

174 """ 

175 print(self.render_header(qg, is_bipartite=False), file=stream) 

176 xgraph = qg.quantum_only_xgraph 

177 for node_id, node_data in xgraph.nodes.items(): 

178 print(self.render_quantum(node_id, node_data, is_bipartite=False), file=stream) 

179 for a, b in xgraph.edges: 

180 print(self.render_edge(a, b, data=None), file=stream) 

181 print(self.render_footer(qg, is_bipartite=False), file=stream) 

182 

183 def write_bipartite(self, qg: _G, stream: IO[str]) -> None: 

184 """Write a visualization for graph with both quantum and dataset nodes. 

185 

186 Parameters 

187 ---------- 

188 qg : `BaseQuantumGraph` 

189 Quantum graph to visualize. 

190 stream : `typing.IO` 

191 File-like object to write to. 

192 """ 

193 print(self.render_header(qg, is_bipartite=True), file=stream) 

194 xgraph = qg.bipartite_xgraph 

195 for node_id, node_data in xgraph.nodes.items(): 

196 match node_data["pipeline_node"].key.node_type: 

197 case NodeType.TASK: 

198 print(self.render_quantum(node_id, node_data, is_bipartite=True), file=stream) 

199 case NodeType.DATASET_TYPE: 199 ↛ 195line 199 didn't jump to line 195 because the pattern on line 199 always matched

200 print(self.render_dataset(node_id, node_data), file=stream) 

201 for a, b, edge_data in xgraph.edges(data=True): 

202 print(self.render_edge(a, b, edge_data), file=stream) 

203 print(self.render_footer(qg, is_bipartite=True), file=stream) 

204 

205 

206class QuantumGraphDotVisualizer(QuantumGraphVisualizer[BaseQuantumGraph, QuantumInfo, DatasetInfo]): 

207 """A visualizer for quantum graphs in the GraphViz dot language.""" 

208 

209 def render_header(self, qg: BaseQuantumGraph, is_bipartite: bool) -> str: 

210 return "\n".join( 

211 [ 

212 "digraph QuantumGraph {", 

213 self._render_default("graph", self._ATTRIBS["defaultGraph"]), 

214 self._render_default("node", self._ATTRIBS["defaultNode"]), 

215 self._render_default("edge", self._ATTRIBS["defaultEdge"]), 

216 ] 

217 ) 

218 

219 def render_footer(self, qg: BaseQuantumGraph, is_bipartite: bool) -> str: 

220 return "}" 

221 

222 def render_quantum(self, quantum_id: uuid.UUID, data: QuantumInfo, is_bipartite: bool) -> str: 

223 labels = [f"{quantum_id}", html.escape(data["task_label"])] 

224 data_id = data["data_id"] 

225 labels.extend(f"{key} = {value}" for key, value in data_id.required.items()) 

226 return self._render_node(self._make_node_id(quantum_id), "quantum", labels) 

227 

228 def render_dataset(self, dataset_id: uuid.UUID, data: DatasetInfo) -> str: 

229 labels = [html.escape(data["dataset_type_name"]), f"run: {data['run']!r}"] 

230 data_id = data["data_id"] 

231 labels.extend(f"{key} = {value}" for key, value in data_id.required.items()) 

232 return self._render_node(self._make_node_id(dataset_id), "dataset", labels) 

233 

234 def render_edge(self, a: uuid.UUID, b: uuid.UUID, data: BipartiteEdgeInfo | None) -> str: 

235 return f'"{self._make_node_id(a)}" -> "{self._make_node_id(b)}";' 

236 

237 _ATTRIBS: ClassVar = dict( 

238 defaultGraph=dict(splines="ortho", nodesep="0.5", ranksep="0.75", pad="0.5"), 

239 defaultNode=dict(shape="box", fontname="Monospace", fontsize="14", margin="0.2,0.1", penwidth="3"), 

240 defaultEdge=dict(color="black", arrowsize="1.5", penwidth="1.5"), 

241 task=dict(style="filled", color="black", fillcolor="#B1F2EF"), 

242 quantum=dict(style="filled", color="black", fillcolor="#B1F2EF"), 

243 dsType=dict(style="rounded,filled,bold", color="#00BABC", fillcolor="#F5F5F5"), 

244 dataset=dict(style="rounded,filled,bold", color="#00BABC", fillcolor="#F5F5F5"), 

245 ) 

246 

247 def _render_default(self, type: str, attribs: dict[str, str]) -> str: 

248 """Set default attributes for a given type.""" 

249 default_attribs = ", ".join([f'{key}="{val}"' for key, val in attribs.items()]) 

250 return f"{type} [{default_attribs}];" 

251 

252 def _render_node(self, name: str, style: str, labels: list[str]) -> str: 

253 """Render GV node""" 

254 label = r"</TD></TR><TR><TD>".join(labels) 

255 attrib_dict = dict(self._ATTRIBS[style], label=label) 

256 pre = '<<TABLE BORDER="0" CELLPADDING="5"><TR><TD>' 

257 post = "</TD></TR></TABLE>>" 

258 attrib = ", ".join( 

259 [ 

260 f'{key}="{val}"' if key != "label" else f"{key}={pre}{val}{post}" 

261 for key, val in attrib_dict.items() 

262 ] 

263 ) 

264 return f'"{name}" [{attrib}];' 

265 

266 def _make_node_id(self, node_id: uuid.UUID) -> str: 

267 """Return a GV node ID from a quantum or dataset UUID.""" 

268 return f"u{node_id.hex}" 

269 

270 

271class QuantumGraphMermaidVisualizer(QuantumGraphVisualizer[BaseQuantumGraph, QuantumInfo, DatasetInfo]): 

272 """A visualizer for quantum graphs in the Mermaid language.""" 

273 

274 def render_header(self, qg: BaseQuantumGraph, is_bipartite: bool) -> str: 

275 return "flowchart TD" 

276 

277 def render_footer(self, qg: BaseQuantumGraph, is_bipartite: bool) -> str: 

278 return "" 

279 

280 def render_quantum(self, quantum_id: uuid.UUID, data: QuantumInfo, is_bipartite: bool) -> str: 

281 label_lines = [f"**{data['task_label']}**", f"ID: {quantum_id}"] 

282 for k, v in data["data_id"].required.items(): 

283 label_lines.append(f"{k}={v}") 

284 label = "<br>".join(label_lines) 

285 return f'{self._make_node_id(quantum_id)}["{label}"]' 

286 

287 def render_dataset(self, dataset_id: uuid.UUID, data: DatasetInfo) -> str: 

288 label_lines = [ 

289 f"**{data['dataset_type_name']}**", 

290 f"ID: {dataset_id}", 

291 f"run: {data['run']}", 

292 ] 

293 for k, v in data["data_id"].required.items(): 

294 label_lines.append(f"{k}={v}") 

295 label = "<br>".join(label_lines) 

296 return f'{self._make_node_id(dataset_id)}["{label}"]' 

297 

298 def render_edge(self, a: uuid.UUID, b: uuid.UUID, data: BipartiteEdgeInfo | None) -> str: 

299 return f"{self._make_node_id(a)} --> {self._make_node_id(b)}" 

300 

301 def _make_node_id(self, node_id: uuid.UUID) -> str: 

302 return f"u{node_id.hex}"