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

93 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-04-18 08:44 +0000

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# We use the old generic syntax in this module because for some reason the new 

41# one confused Sphinx (or one of its plugins), even though it seems fine with 

42# it in other places. We can try again when we're ready to remove types from 

43# the docstrings of annotated functions, in case that matters. 

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

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

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

47 

48 

49class QuantumGraphVisualizer(Generic[_G, _Q, _D]): # noqa: UP046 

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

51 languages. 

52 

53 Notes 

54 ----- 

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

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

57 at least GraphViz and Mermaid. 

58 """ 

59 

60 @abstractmethod 

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

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

63 

64 Parameters 

65 ---------- 

66 qg : `.BaseQuantumGraph` 

67 Quantum graph to visualize. 

68 is_bipartite : `bool` 

69 Whether a bipartite graph visualization is being requested. 

70 

71 Returns 

72 ------- 

73 rendered : `str` 

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

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

76 before the first node. 

77 """ 

78 raise NotImplementedError() 

79 

80 @abstractmethod 

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

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

83 

84 Parameters 

85 ---------- 

86 qg : `.BaseQuantumGraph` 

87 Quantum graph to visualize. 

88 is_bipartite : `bool` 

89 Whether a bipartite graph visualization is being requested. 

90 

91 Returns 

92 ------- 

93 rendered : `str` 

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

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

96 at the end of the file. 

97 """ 

98 raise NotImplementedError() 

99 

100 @abstractmethod 

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

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

103 

104 Parameters 

105 ---------- 

106 quantum_id : `uuid.UUID` 

107 ID of the quantum. 

108 data : `.QuantumInfo` 

109 Mapping with additional information about the quantum. 

110 is_bipartite : `bool` 

111 Whether a bipartite graph visualization is being requested. 

112 

113 Returns 

114 ------- 

115 rendered : `str` 

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

117 additional newline will automatically be added after the string, 

118 before the next node or edge. 

119 """ 

120 raise NotImplementedError() 

121 

122 @abstractmethod 

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

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

125 

126 Parameters 

127 ---------- 

128 dataset_id : `uuid.UUID` 

129 ID of the dataset. 

130 data : `.DatasetInfo` 

131 Mapping with additional information about the dataset. 

132 

133 Returns 

134 ------- 

135 rendered : `str` 

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

137 additional newline will automatically be added after the string, 

138 before the next node or edge. 

139 """ 

140 raise NotImplementedError() 

141 

142 @abstractmethod 

143 def render_edge( 

144 self, 

145 a: uuid.UUID, 

146 b: uuid.UUID, 

147 data: BipartiteEdgeInfo | None, 

148 ) -> str: 

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

150 

151 Parameters 

152 ---------- 

153 a : `uuid.UUID` 

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

155 b : `uuid.UUID` 

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

157 data : `.BipartiteEdgeInfo` 

158 Mapping with additional information about the dataset. 

159 

160 Returns 

161 ------- 

162 rendered : `str` 

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

164 additional newline will automatically be added after the string, 

165 before the next edge. 

166 """ 

167 raise NotImplementedError() 

168 

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

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

171 

172 Parameters 

173 ---------- 

174 qg : `BaseQuantumGraph` 

175 Quantum graph to visualize. 

176 stream : `typing.IO` 

177 File-like object to write to. 

178 """ 

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

180 xgraph = qg.quantum_only_xgraph 

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

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

183 for a, b in xgraph.edges: 

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

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

186 

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

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

189 

190 Parameters 

191 ---------- 

192 qg : `BaseQuantumGraph` 

193 Quantum graph to visualize. 

194 stream : `typing.IO` 

195 File-like object to write to. 

196 """ 

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

198 xgraph = qg.bipartite_xgraph 

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

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

201 case NodeType.TASK: 

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

203 case NodeType.DATASET_TYPE: 

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

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

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

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

208 

209 

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

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

212 

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

214 return "\n".join( 

215 [ 

216 "digraph QuantumGraph {", 

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

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

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

220 ] 

221 ) 

222 

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

224 return "}" 

225 

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

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

228 data_id = data["data_id"] 

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

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

231 

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

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

234 data_id = data["data_id"] 

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

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

237 

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

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

240 

241 _ATTRIBS: ClassVar = dict( 

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

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

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

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

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

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

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

249 ) 

250 

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

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

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

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

255 

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

257 """Render GV node""" 

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

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

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

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

262 attrib = ", ".join( 

263 [ 

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

265 for key, val in attrib_dict.items() 

266 ] 

267 ) 

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

269 

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

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

272 return f"u{node_id.hex}" 

273 

274 

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

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

277 

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

279 return "flowchart TD" 

280 

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

282 return "" 

283 

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

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

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

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

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

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

290 

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

292 label_lines = [ 

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

294 f"ID: {dataset_id}", 

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

296 ] 

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

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

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

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

301 

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

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

304 

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

306 return f"u{node_id.hex}"