Coverage for include/graph.py: 87%
84 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-12-05 17:26 +0000
1import h5py as h5
2import networkx as nx
3import numpy as np
5""" Network generation function """
8def generate_graph(
9 *,
10 N: int,
11 mean_degree: int = None,
12 type: str,
13 seed: int = None,
14 graph_props: dict = None,
15) -> nx.Graph:
16 """Generates graphs of a given topology
18 :param N: number of nodes
19 :param mean_degree: mean degree of the graph
20 :param type: graph topology kind; can be any of 'random', 'BarabasiAlbert' (scale-free undirected), 'BollobasRiordan'
21 (scale-free directed), 'WattsStrogatz' (small-world)
22 :param seed: the random seed to use for the graph generation (ensuring the graphs are always the same)
23 :param graph_props: dictionary containing the type-specific parameters
24 :return: the networkx graph object. All graphs are fully connected
25 """
27 def _connect_isolates(G: nx.Graph) -> nx.Graph:
28 """Connects isolated vertices to main network body."""
29 isolates = list(nx.isolates(G))
30 N = G.number_of_nodes()
31 for v in isolates:
32 new_nb = np.random.randint(0, N)
33 while new_nb in isolates:
34 new_nb = np.random.randint(0, N)
35 G.add_edge(v, new_nb)
37 return G
39 # Random graph
40 if type.lower() == "random":
41 is_directed: bool = graph_props["is_directed"]
43 G = _connect_isolates(
44 nx.fast_gnp_random_graph(
45 N, mean_degree / (N - 1), directed=is_directed, seed=seed
46 )
47 )
49 # Undirected scale-free graph
50 elif type.lower() == "barabasialbert":
51 G = _connect_isolates(nx.barabasi_albert_graph(N, mean_degree, seed))
53 # Directed scale-free graph
54 elif type.lower() == "bollobasriordan":
55 G = nx.DiGraph(
56 _connect_isolates(
57 nx.scale_free_graph(N, **graph_props["BollobasRiordan"], seed=seed)
58 )
59 )
61 # Small-world (Watts-Strogatz) graph
62 elif type.lower() == "wattsstrogatz":
63 p: float = graph_props["WattsStrogatz"]["p_rewire"]
65 G = _connect_isolates(nx.watts_strogatz_graph(N, mean_degree, p, seed))
67 # Star graph
68 elif type.lower() == "star":
69 G = nx.star_graph(N)
71 # Regular graph
72 elif type.lower() == "regular":
73 G = nx.random_regular_graph(mean_degree, N, seed)
75 # Raise error upon unrecognised graph type
76 else:
77 raise ValueError(f"Unrecognised graph type '{type}'!")
79 # Add random uniform weights to the edges
80 if graph_props["is_weighted"]:
81 for e in G.edges():
82 G[e[0]][e[1]]["weight"] = np.random.rand()
84 return G
87def save_nw(
88 network: nx.Graph,
89 nw_group: h5.Group,
90 *,
91 write_adjacency_matrix: bool = False,
92 static: bool = True,
93):
94 """Saves a network to a h5.Group graph group. The network can be either dynamic of static,
95 static in this case meaning that none of the datasets have a time dimension.
97 :param network: the network to save
98 :param nw_group: the h5.Group
99 :param write_adjacency_matrix: whether to write out the entire adjacency matrix
100 :param static: whether the network is dynamic or static
101 """
103 # Vertices
104 vertices = nw_group.create_dataset(
105 "_vertices",
106 (network.number_of_nodes(),) if static else (1, network.number_of_nodes()),
107 maxshape=(network.number_of_nodes(),)
108 if static
109 else (None, network.number_of_nodes()),
110 chunks=True,
111 compression=3,
112 dtype=int,
113 )
114 vertices.attrs["dim_names"] = ["vertex_idx"] if static else ["time", "vertex_idx"]
115 vertices.attrs["coords_mode__vertex_idx"] = "trivial"
117 # Edges; the network size is assumed to remain constant
118 edges = nw_group.create_dataset(
119 "_edges",
120 (network.size(), 2) if static else (1, network.size(), 2),
121 maxshape=(network.size(), 2) if static else (None, network.size(), 2),
122 chunks=True,
123 compression=3,
124 )
125 edges.attrs["dim_names"] = (
126 ["edge_idx", "vertex_idx"] if static else ["time", "edge_idx", "vertex_idx"]
127 )
128 edges.attrs["coords_mode__edge_idx"] = "trivial"
129 edges.attrs["coords_mode__vertex_idx"] = "trivial"
131 # Edge properties
132 edge_weights = nw_group.create_dataset(
133 "_edge_weights",
134 (network.size(),) if static else (1, network.size()),
135 maxshape=(network.size(),) if static else (None, network.size()),
136 chunks=True,
137 compression=3,
138 )
139 edge_weights.attrs["dim_names"] = ["edge_idx"] if static else ["time", "edge_idx"]
140 edge_weights.attrs["coords_mode__edge_idx"] = "trivial"
142 # Topological properties
143 degree = nw_group.create_dataset(
144 "_degree",
145 (network.number_of_nodes(),) if static else (1, network.number_of_nodes()),
146 maxshape=(network.number_of_nodes(),)
147 if static
148 else (None, network.number_of_nodes()),
149 chunks=True,
150 compression=3,
151 dtype=int,
152 )
153 degree.attrs["dim_names"] = ["vertex_idx"] if static else ["time", "vertex_idx"]
154 degree.attrs["coords_mode__vertex_idx"] = "trivial"
156 degree_w = nw_group.create_dataset(
157 "_degree_weighted",
158 (network.number_of_nodes(),) if static else (1, network.number_of_nodes()),
159 maxshape=(network.number_of_nodes(),)
160 if static
161 else (None, network.number_of_nodes()),
162 chunks=True,
163 compression=3,
164 dtype=float,
165 )
166 degree_w.attrs["dim_names"] = ["vertex_idx"] if static else ["time", "vertex_idx"]
167 degree_w.attrs["coords_mode__vertex_idx"] = "trivial"
169 triangles = nw_group.create_dataset(
170 "_triangles",
171 (network.number_of_nodes(),) if static else (1, network.number_of_nodes()),
172 maxshape=(network.number_of_nodes(),)
173 if static
174 else (None, network.number_of_nodes()),
175 chunks=True,
176 compression=3,
177 dtype=float,
178 )
179 triangles.attrs["dim_names"] = ["vertex_idx"] if static else ["time", "vertex_idx"]
180 triangles.attrs["coords_mode__vertex_idx"] = "trivial"
182 triangles_w = nw_group.create_dataset(
183 "_triangles_weighted",
184 (network.number_of_nodes(),) if static else (1, network.number_of_nodes()),
185 maxshape=(network.number_of_nodes(),)
186 if static
187 else (None, network.number_of_nodes()),
188 chunks=True,
189 compression=3,
190 dtype=float,
191 )
192 triangles_w.attrs["dim_names"] = (
193 ["vertex_idx"] if static else ["time", "vertex_idx"]
194 )
195 triangles_w.attrs["coords_mode__vertex_idx"] = "trivial"
197 if not static:
198 for dset in [
199 vertices,
200 edges,
201 edge_weights,
202 degree,
203 degree_w,
204 triangles,
205 triangles_w,
206 ]:
207 dset["coords_mode__time"] = "trivial"
209 # Write network properties
210 if static:
211 vertices[:] = network.nodes()
212 edges[:, :] = network.edges()
213 edge_weights[:] = list(nx.get_edge_attributes(network, "weight").values())
214 degree[:] = [network.degree(i) for i in network.nodes()]
215 degree_w[:] = [deg[1] for deg in network.degree(weight="weight")]
216 triangles[:] = [
217 val
218 for val in np.diagonal(
219 np.linalg.matrix_power(np.ceil(nx.to_numpy_array(network)), 3)
220 )
221 ]
222 triangles_w[:] = [
223 val
224 for val in np.diagonal(
225 np.linalg.matrix_power(nx.to_numpy_array(network), 3)
226 )
227 ]
228 else:
229 vertices[0, :] = network.nodes()
230 edges[0, :, :] = network.edges()
231 edge_weights[0, :] = list(nx.get_edge_attributes(network, "weight").values())
232 degree[0, :] = [network.degree(i) for i in network.nodes()]
233 degree_w[0, :] = [deg[1] for deg in network.degree(weight="weight")]
234 triangles[0, :] = [
235 val
236 for val in np.diagonal(
237 np.linalg.matrix_power(np.ceil(nx.to_numpy_array(network)), 3)
238 )
239 ]
240 triangles_w[0, :] = [
241 val
242 for val in np.diagonal(
243 np.linalg.matrix_power(nx.to_numpy_array(network), 3)
244 )
245 ]
247 if write_adjacency_matrix:
248 adj_matrix = nx.to_numpy_array(network)
250 # Adjacency matrix: only written out if explicity specified
251 adjacency_matrix = nw_group.create_dataset(
252 "_adjacency_matrix",
253 np.shape(adj_matrix) if static else [1] + list(np.shape(adj_matrix)),
254 chunks=True,
255 compression=3,
256 )
257 adjacency_matrix.attrs["dim_names"] = (
258 ["i", "j"] if static else ["time", "i", "j"]
259 )
260 adjacency_matrix.attrs["coords_mode__i"] = "trivial"
261 adjacency_matrix.attrs["coords_mode__j"] = "trivial"
262 if not static:
263 adjacency_matrix.attrs["coords_mode__time"] = "trivial"
265 if static:
266 adjacency_matrix[:, :] = adj_matrix
267 else:
268 adjacency_matrix[0, :, :] = adj_matrix