Commit 78dcac84 authored by Tiago Peixoto's avatar Tiago Peixoto
Browse files

Implement BlockState.get_edges_prob()

parent c5e1c6a2
......@@ -167,8 +167,6 @@ struct Layers
size_t _total_B;
bool _is_partition_stats_enabled;
//TODO: remove_vertex and add_vertex, etc.
void move_vertex(size_t v, size_t s)
{
if (BaseState::_vweight[v] == 0)
......@@ -217,6 +215,39 @@ struct Layers
}
}
void remove_vertex(size_t v)
{
size_t r = _b[v];
auto& ls = _vc[v];
auto& vs = _vmap[v];
for (size_t j = 0; j < ls.size(); ++j)
{
int l = ls[j];
size_t u = vs[j];
auto& state = _layers[l];
size_t r_u = state._b[u];
state.remove_vertex(u);
if (state._wr[r_u] == 0)
state.remove_block_map(r);
}
BaseState::remove_vertex(v);
}
void add_vertex(size_t v, size_t r)
{
auto& ls = _vc[v];
auto& vs = _vmap[v];
for (size_t j = 0; j < ls.size(); ++j)
{
int l = ls[j];
size_t u = vs[j];
auto& state = _layers[l];
size_t r_u = state.get_block_map(r);
state.add_vertex(u, r_u);
}
BaseState::add_vertex(v, r);
}
template <class VMap>
void set_partition(VMap&& b)
{
......
......@@ -654,6 +654,70 @@ class BlockState(object):
"""
return self._state.get_move_prob(int(v), self.b[v], s, c, reverse)
def get_edges_prob(self, edge_list, missing=True, entropy_args={}):
"""Compute the log-probability of the missing (or spurious if ``missing=False``)
edges given by ``edge_list`` (a list of ``(source, target)`` tuples, or
:meth:`~graph_tool.Edge` instances). The values in ``entropy_args`` are
passed to :meth:`graph_tool.BlockState.entropy()` to calculate the
log-probability.
"""
pos = {}
for u, v in edge_list:
pos[u] = self.b[u]
pos[v] = self.b[v]
Si = self.entropy(**entropy_args)
for v in pos.keys():
self.remove_vertex(v)
try:
if missing:
new_es = []
for u, v in edge_list:
e = self.g.add_edge(u, v)
new_es.append(e)
self.E += 1
else:
old_es = []
for e in edge_list:
if isinstance(e, tuple):
u, v = e
tmp = self.g.edge(u, v)
if tmp is None:
raise ValueError("edge not found: (%d, %d)" % (int(u),
int(v)))
self.g.remove_edge(tmp)
else:
u, v = e
self.g.remove_edge(e)
old_es.append((u, v))
self.E -= 1
for v in pos.keys():
self.add_vertex(v, pos[v])
Sf = self.entropy(**entropy_args)
for v in pos.keys():
self.remove_vertex(v)
finally:
if missing:
for e in new_es:
self.g.remove_edge(e)
self.E -= 1
else:
for u, v in old_es:
self.g.add_edge(u, v)
self.E += 1
for v in pos.keys():
self.add_vertex(v, pos[v])
if missing:
return Si - Sf
else:
return Sf - Si
def _mcmc_sweep_dispatch(self, mcmc_state):
if (mcmc_state.multigraph and not mcmc_state.dense and
......
......@@ -177,7 +177,8 @@ class NestedBlockState(object):
if _bm_test():
self._consistency_check()
def level_entropy(self, l, dense=False, multigraph=True, bstate=None, **kwargs):
def level_entropy(self, l, dense=False, multigraph=True, bstate=None,
**kwargs):
"""Compute the entropy of level ``l``."""
if bstate is None:
......@@ -198,6 +199,24 @@ class NestedBlockState(object):
**kwargs)
return S
def get_edges_prob(self, edge_list, missing=True, entropy_args={}):
"""Compute the log-probability of the missing (or spurious if ``missing=False``)
edges given by ``edge_list`` (a list of ``(source, target)`` tuples, or
:meth:`~graph_tool.Edge` instances). The values in ``entropy_args`` are
passed to :meth:`graph_tool.NestedBlockState.entropy()` to calculate the
log-probability.
"""
L = 0
for l, state in enumerate(self.levels):
eargs = overlay(entropy_args, dl=True,
edges_dl=(l == (len(self.levels) - 1)))
if l > 0:
eargs = overlay(eargs, dense=True, multigraph=True)
L += state.get_edges_prob(edge_list, missing=missing,
entropy_args=eargs)
edge_list = [(state.b[u], state.b[v]) for u, v in edge_list]
return L
def get_bstack(self):
"""Return the nested levels as individual graphs.
......
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