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

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

21 from typing import Optional

22

23 from abcd_graph.callbacks.abstract import (

24 ABCDCallback,

25 BuildContext,

26)

27 from abcd_graph.exporter import GraphExporter

28 from abcd_graph.graph.core.abcd_objects import (

29 Community,

30 GraphImpl,

31)

32 from abcd_graph.graph.core.utils import get_community_color_map

33 from abcd_graph.models import Model

34 from abcd_graph.utils import require

35

36

37 class 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()