-
Notifications
You must be signed in to change notification settings - Fork 6
/
chip_local_search.py
278 lines (218 loc) · 12.9 KB
/
chip_local_search.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
# -*- coding: utf-8 -*-
"""
@author: Makan Arastuie
"""
import time
import numpy as np
import multiprocessing
from joblib import Parallel, delayed
import generative_model_utils as utils
import parameter_estimation as estimate_utils
from sklearn.metrics import adjusted_rand_score
from spectral_clustering import spectral_cluster
import model_fitting_utils as fit_utils
def calc_node_neigh_solutions(event_dict, n_classes, duration, node_membership, agg_adj, beta, log_lik_init, node_batch):
"""
Calculates the log-likelihood of neighboring solutions of a batch of nodes by changing their membership. If a higher
log-likelihood was achieved the best solution will be returned, else a tuple of three np.nan is returned.
:param event_dict: Edge dictionary of events between all node pair. Output of the generative models.
:param n_classes: (int) total number of classes/blocks
:param duration: (int) Duration of the network
:param node_membership: (list) membership of every node to one of K classes
:param agg_adj: aggregated/weighted adjacency of the network
:param beta: K x K np array of block pairs beta. This is fixed for every solution to lower time complexity. Only mu
and m are estimated for each neighboring solution
:param log_lik_init: (float) base log-likelihood
:param node_batch: (list) nodes in the current batch
:return: (node index, best class index, log_likelihood)
"""
best_neigh = (np.nan, np.nan, np.nan)
log_lik = log_lik_init
# node_membership = node_membership.copy()
for n_i in node_batch:
n_i_class = node_membership[n_i]
# Adding a constraint to maintain the number of blocks.
if np.sum(node_membership == n_i_class) <= 2:
continue
for c_i in range(n_classes):
if c_i == n_i_class:
continue
# update node_membership temporarily
node_membership[n_i] = c_i
# Eval the aprox log_lik of this neighbor, by est its mu and alpha/beta and using previous beta.
neigh_mu, neigh_alpha_beta_ratio = estimate_utils.estimate_hawkes_from_counts(agg_adj, node_membership,
duration,
default_mu=1e-10 / duration)
neigh_alpha = neigh_alpha_beta_ratio * beta
block_pair_events = utils.event_dict_to_block_pair_events(event_dict, node_membership, n_classes)
neigh_log_lik = fit_utils.calc_full_log_likelihood(block_pair_events, node_membership,
neigh_mu, neigh_alpha, beta,
duration, n_classes, add_com_assig_log_prob=False)
# if log_lik if this neighbor is better than the "so far" best neighbor, use this neighbors as the best.
if log_lik < neigh_log_lik:
log_lik = neigh_log_lik
best_neigh = (n_i, c_i, log_lik)
node_membership[n_i] = n_i_class
return best_neigh
def chip_local_search(event_dict, n_classes, node_membership_init, duration, max_iter=100, n_cores=-1,
return_fitted_param=False, verbose=True):
"""
Performs local search / hill climbing to increase log-likelihood of the model by switching the community of a single
node at a time. For every neighboring solution only mu and m are estimated, beta is fixed to the base solution to
lower time complexity.
:param event_dict: Edge dictionary of events between all node pair. Output of the generative models.
:param n_classes: (int) total number of classes/blocks
:param node_membership_init: (list) initial membership of every node to one of K classes. Usually output of the
spectral clustering
:param duration: (int) Duration of the network
:param max_iter: (int) maximum number of iterations to be performed by local search.
:param n_cores: (int) number of cores to be used to parallelize the search. If -1, use all available cores.
:param return_fitted_param: if True, return the Hawkes parameters for the model as well.
:param verbose: If True, prints more information on local search.
:return: local optimum node_membership if `return_fitted_param` is false.
"""
n_nodes = len(node_membership_init)
nodes = np.arange(n_nodes)
node_membership = node_membership_init
agg_adj = utils.event_dict_to_aggregated_adjacency(n_nodes, event_dict, dtype=np.int)
# estimate initial params of CHIP and its log-likelihood
(mu,
alpha,
beta,
alpha_beta_ratio) = fit_utils.estimate_bp_hawkes_params(event_dict, node_membership, duration, n_classes)
block_pair_events = utils.event_dict_to_block_pair_events(event_dict, node_membership, n_classes)
init_log_lik = fit_utils.calc_full_log_likelihood(block_pair_events, node_membership, mu, alpha, beta,
duration, n_classes, add_com_assig_log_prob=False)
log_lik = init_log_lik
n_cores = n_cores if n_cores > 0 else multiprocessing.cpu_count()
batch_size = np.int(n_nodes / n_cores) + 1
for iter in range(max_iter):
if verbose:
print(f"Iteration {iter}...", end='\r')
# for each of the (k-1)*n neighboring solutions
possible_solutions = Parallel(n_jobs=n_cores)(delayed(calc_node_neigh_solutions)
(event_dict, n_classes, duration, node_membership, agg_adj,
beta, log_lik, nodes[batch_size * ii: batch_size * (ii + 1)])
for ii in range(n_cores))
possible_solutions = np.array(possible_solutions)
# if all returned log-likelihoods are np.nan, break. You're at a local optima.
if np.all(np.isnan(possible_solutions[:, 2])):
if verbose:
print(f"Local solution found with {iter} iterations.")
break
max_ll_neigh_idx = np.nanargmax(possible_solutions[:, 2])
# if a good neighbor was found, update all CHIP params, and go for the next iteration.
node_membership[int(possible_solutions[max_ll_neigh_idx, 0])] = int(possible_solutions[max_ll_neigh_idx, 1])
(mu,
alpha,
beta,
alpha_beta_ratio) = fit_utils.estimate_bp_hawkes_params(event_dict, node_membership, duration, n_classes)
block_pair_events = utils.event_dict_to_block_pair_events(event_dict, node_membership, n_classes)
log_lik = fit_utils.calc_full_log_likelihood(block_pair_events, node_membership, mu, alpha, beta,
duration, n_classes, add_com_assig_log_prob=False)
if iter == max_iter - 1:
print("Warning: Max iter reached!")
if verbose:
print(f"likelihood went from {init_log_lik:.4f} to {log_lik:.4f}. "
f"{100 * np.abs((log_lik - init_log_lik) / init_log_lik):.2f}% increase.")
if return_fitted_param:
return node_membership, mu, alpha, beta
return node_membership
def chip_local_search_single_core(event_dict, n_classes, node_membership_init, duration, max_iter=100, verbose=True):
"""
This function is only here for speed comparisons against the multi-core version. All parameters are the same as
`chip_local_search`.
"""
n_nodes = len(node_membership_init)
node_membership = node_membership_init
agg_adj = utils.event_dict_to_aggregated_adjacency(n_nodes, event_dict, dtype=np.int)
# estimate initial params of CHIP and its log-likelihood
(mu,
alpha,
beta,
alpha_beta_ratio) = fit_utils.estimate_bp_hawkes_params(event_dict, node_membership, duration, n_classes)
block_pair_events = utils.event_dict_to_block_pair_events(event_dict, node_membership, n_classes)
init_log_lik = fit_utils.calc_full_log_likelihood(block_pair_events, node_membership, mu, alpha, beta,
duration, n_classes, add_com_assig_log_prob=False)
log_lik = init_log_lik
for iter in range(max_iter):
if verbose:
print(f"Iteration {iter}...", end='\r')
# best neighbor will hold the best node_membership update in the form of (node_index, updated_class_membership)
best_neigh = None
# for each of the (k-1)*n neighboring solutions
for n_i in range(n_nodes):
n_i_class = node_membership[n_i]
for c_i in range(n_classes):
if c_i == n_i_class:
continue
# update node_membership temporarily
node_membership[n_i] = c_i
# Eval the aprox log_lik of this neighbor, by est its mu and alpha/beta and using previous beta.
neigh_mu, neigh_alpha_beta_ratio = estimate_utils.estimate_hawkes_from_counts(agg_adj, node_membership,
duration,
default_mu=1e-10 / duration)
neigh_alpha = neigh_alpha_beta_ratio * beta
block_pair_events = utils.event_dict_to_block_pair_events(event_dict, node_membership, n_classes)
neigh_log_lik = fit_utils.calc_full_log_likelihood(block_pair_events, node_membership,
neigh_mu, neigh_alpha, beta,
duration, n_classes, add_com_assig_log_prob=False)
# if log_lik if this neighbor is better than the "so far" best neighbor, use this neighbors as the best.
if log_lik < neigh_log_lik:
log_lik = neigh_log_lik
best_neigh = (n_i, c_i)
node_membership[n_i] = n_i_class
# if no neighbor seem to increase log_lik, break. You're at a local optima.
if best_neigh is None:
if verbose:
print(f"Local solution found with {iter} iterations.")
break
# if a good neighbor was found, update all CHIP params, and go for the next iteration.
node_membership[best_neigh[0]] = best_neigh[1]
(mu,
alpha,
beta,
alpha_beta_ratio) = fit_utils.estimate_bp_hawkes_params(event_dict, node_membership, duration, n_classes)
block_pair_events = utils.event_dict_to_block_pair_events(event_dict, node_membership, n_classes)
log_lik = fit_utils.calc_full_log_likelihood(block_pair_events, node_membership, mu, alpha, beta,
duration, n_classes, add_com_assig_log_prob=False)
if verbose:
print(f"likelihood went from {init_log_lik:.4f} to {log_lik:.4f}. "
f"{100 * np.abs((log_lik - init_log_lik) / init_log_lik):.2f}% increase.")
return node_membership
# Examples of running local search on CHIP model.
if __name__ == '__main__':
n_classes = 4
n_nodes = 1024
duration = 50
params = {'number_of_nodes': n_nodes,
'alpha': 0.6,
'beta': 0.8,
'mu_off_diag': 0.8,
'mu_diag': 1.6,
'end_time': duration,
'class_probabilities': np.ones(n_classes) / n_classes,
'n_cores': -1}
event_dict, true_class_assignments = utils.simulate_community_hawkes(params,
network_name="local_seach_test_networks",
load_if_exists=False)
agg_adj = utils.event_dict_to_aggregated_adjacency(n_nodes, event_dict)
spectral_node_membership = spectral_cluster(agg_adj, num_classes=n_classes)
sc_rand = adjusted_rand_score(true_class_assignments, spectral_node_membership)
print(f"SC Rand index: {sc_rand:.3f}")
print("Parallel")
tic = time.time()
local_search_node_membership = chip_local_search(event_dict, n_classes, spectral_node_membership, duration,
max_iter=10, n_cores=34, verbose=True)
toc = time.time()
print(f"local search took {toc - tic:.2f}s.")
sc_rand = adjusted_rand_score(true_class_assignments, local_search_node_membership)
print(f"Local search Rand index: {sc_rand:.3f}")
print("Single core")
tic = time.time()
local_search_node_membership = chip_local_search_single_core(event_dict, n_classes, spectral_node_membership,
duration, max_iter=10, verbose=True)
toc = time.time()
print(f"local search took {toc - tic:.2f}s.")
sc_rand = adjusted_rand_score(true_class_assignments, local_search_node_membership)
print(f"Local search Rand index: {sc_rand:.3f}")