Commit 42c81a3d authored by Tiago Peixoto's avatar Tiago Peixoto
Browse files

Fix pickling of NestedBlockState

parent ff9fdc7d
...@@ -71,6 +71,7 @@ class NestedBlockState(object): ...@@ -71,6 +71,7 @@ class NestedBlockState(object):
self.g = g self.g = g
self.kwargs = kwargs.copy() self.kwargs = kwargs.copy()
self.hstate_args = overlay(dict(deg_corr=False), **hstate_args) self.hstate_args = overlay(dict(deg_corr=False), **hstate_args)
self.sampling = sampling
if sampling: if sampling:
self.hstate_args = overlay(self.hstate_args, vweight="nonempty", self.hstate_args = overlay(self.hstate_args, vweight="nonempty",
copy_bg=False, B=g.num_vertices()) copy_bg=False, B=g.num_vertices())
...@@ -83,11 +84,14 @@ class NestedBlockState(object): ...@@ -83,11 +84,14 @@ class NestedBlockState(object):
nbs.append(nb) nbs.append(nb)
bs = nbs bs = nbs
self.hentropy_args = overlay(dict(adjacency=True, self.hentropy_args = overlay(dict(adjacency=True,
dl=True, partition_dl=True, dense=True,
multigraph=True,
dl=True,
partition_dl=True,
degree_dl=True, degree_dl=True,
degree_dl_kind="distributed", degree_dl_kind="distributed",
edges_dl=True, dense=True, edges_dl=True,
multigraph=True, exact=True), exact=True),
**hentropy_args) **hentropy_args)
self.levels = [base_type(g, b=bs[0], **self.kwargs)] self.levels = [base_type(g, b=bs[0], **self.kwargs)]
for b in bs[1:]: for b in bs[1:]:
...@@ -129,7 +133,10 @@ class NestedBlockState(object): ...@@ -129,7 +133,10 @@ class NestedBlockState(object):
**overlay(self.kwargs, **kwargs)) **overlay(self.kwargs, **kwargs))
def __getstate__(self): def __getstate__(self):
state = dict(g=self.g, bs=self.get_bs(), kwargs=self.kwargs) state = dict(g=self.g, bs=self.get_bs(), base_type=type(self.levels[0]),
hstate_args=self.hstate_args,
hentropy_args=self.hstate_args, sampling=self.sampling,
kwargs=self.kwargs)
return state return state
def __setstate__(self, state): def __setstate__(self, state):
...@@ -425,10 +432,7 @@ class NestedBlockState(object): ...@@ -425,10 +432,7 @@ class NestedBlockState(object):
for l in range(len(self.levels) - 1): for l in range(len(self.levels) - 1):
eargs = overlay(self.hentropy_args, eargs = overlay(self.hentropy_args,
**overlay(entropy_args, edges_dl=(l + 1 == len(self.levels) - 1))
adjacency=True,
dense=True,
edges_dl=(l + 1 == len(self.levels) - 1)))
self.levels[l]._couple_state(self.levels[l + 1], self.levels[l]._couple_state(self.levels[l + 1],
get_entropy_args(eargs)) get_entropy_args(eargs))
self.levels[l + 1]._state.clear_egroups() self.levels[l + 1]._state.clear_egroups()
...@@ -455,10 +459,7 @@ class NestedBlockState(object): ...@@ -455,10 +459,7 @@ class NestedBlockState(object):
def callback(s): def callback(s):
s = self.levels[l + 1] s = self.levels[l + 1]
S = s.entropy(**overlay(self.hentropy_args, S = s.entropy(**overlay(self.hentropy_args,
**overlay(entropy_args, edges_dl=(l + 1 == len(self.levels) - 1)))
adjacency=True,
dense=True,
edges_dl=(l + 1 == len(self.levels) - 1))))
return S return S
eargs = overlay(eargs, callback=callback) eargs = overlay(eargs, callback=callback)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment