Coverage for src/abcd_graph/callbacks/visualizer.py: 99%

67 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2024-12-04 21:31 +0100

1# Copyright (c) 2024 Jordan Barrett & Aleksander Wojnarowicz 

2# 

3# Permission is hereby granted, free of charge, to any person obtaining a copy 

4# of this software and associated documentation files (the "Software"), to deal 

5# in the Software without restriction, including without limitation the rights 

6# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

7# copies of the Software, and to permit persons to whom the Software is 

8# furnished to do so, subject to the following conditions: 

9# 

10# The above copyright notice and this permission notice shall be included in all 

11# copies or substantial portions of the Software. 

12# 

13# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

14# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

15# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

16# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

17# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

18# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

19# SOFTWARE. 

20 

21from typing import Optional 

22 

23from abcd_graph.callbacks.abstract import ( 

24 ABCDCallback, 

25 BuildContext, 

26) 

27from abcd_graph.exporter import GraphExporter 

28from abcd_graph.graph.core.abcd_objects import ( 

29 Community, 

30 GraphImpl, 

31) 

32from abcd_graph.graph.core.utils import get_community_color_map 

33from abcd_graph.models import Model 

34from abcd_graph.utils import require 

35 

36 

37class Visualizer(ABCDCallback): 

38 def __init__(self) -> None: 

39 self._communities: list[Community] = [] 

40 self._model_used: Optional[Model] = None 

41 self._exporter: Optional[GraphExporter] = None 

42 self._graph: Optional[GraphImpl] = None 

43 

44 def after_build(self, graph: GraphImpl, context: BuildContext, exporter: GraphExporter) -> None: 

45 self._communities = graph.communities 

46 self._model_used = context.model_used 

47 self._exporter = exporter 

48 

49 self._graph = graph 

50 

51 @require("matplotlib") 

52 def draw_community_cdf(self) -> None: 

53 import matplotlib.pyplot as plt # type: ignore[import] 

54 

55 assert self._graph is not None 

56 

57 actual_cdf = self._graph.actual_community_cdf 

58 expected_cdf = self._graph.expected_community_cdf 

59 

60 x_actual = list(actual_cdf.keys()) 

61 x_expected = list(expected_cdf.keys()) 

62 y_actual = list(actual_cdf.values()) 

63 y_expected = list(expected_cdf.values()) 

64 

65 plt.plot(x_actual, y_actual, label="Actual") 

66 plt.plot(x_expected, y_expected, label="Expected") 

67 

68 plt.xlabel("Community Size") 

69 plt.ylabel("CDF") 

70 plt.legend() 

71 plt.title("Community size CDF") 

72 plt.show() 

73 

74 @require("matplotlib") 

75 def draw_degree_cdf(self) -> None: 

76 import matplotlib.pyplot as plt 

77 

78 assert self._graph is not None 

79 

80 actual_cdf = self._graph.actual_degree_cdf 

81 expected_cdf = self._graph.expected_degree_cdf 

82 

83 x_actual = list(actual_cdf.keys()) 

84 x_expected = list(expected_cdf.keys()) 

85 y_actual = list(actual_cdf.values()) 

86 y_expected = list(expected_cdf.values()) 

87 

88 plt.plot(x_actual, y_actual, label="Actual") 

89 plt.plot(x_expected, y_expected, label="Expected") 

90 

91 plt.xlabel("Degree") 

92 plt.ylabel("CDF") 

93 plt.legend() 

94 plt.title("Degree CDF") 

95 plt.show() 

96 

97 @require("networkx") 

98 @require("matplotlib") 

99 def draw_communities(self) -> None: 

100 assert self._graph is not None 

101 if len(self._graph.deg_b) > 100: 

102 raise ValueError("Drawing communities is only supported for graphs with at most 100 vertices") 

103 

104 if self._model_used is not None and self._model_used.__name__ != "configuration_model": 

105 raise NotImplementedError("Drawing communities is only supported for the configuration model") 

106 

107 import networkx as nx # type: ignore[import] 

108 from matplotlib import pyplot as plt 

109 

110 assert self._exporter is not None 

111 

112 nx_g = self._exporter.to_networkx() 

113 

114 color_map = get_community_color_map(communities=self._communities) 

115 

116 nx.draw(nx_g, node_color=color_map, with_labels=True, font_weight="bold") 

117 plt.show()