def network_to_cluster_graph(network, can_cross_boundary=None, use_weights=True, merge_nengo_nodes=True):
""" Create a cluster graph from a nengo Network.
A cluster graph is a graph wherein the nodes are maximally large
NengoObjectClusters, and edges are nengo Connections that are permitted to
cross component boundaries.
Parameters
----------
network: nengo.Network
The network whose cluster graph we want to construct.
can_cross_boundary: function
A function which accepts a Connection, and returns a boolean specifying
whether the Connection is allowed to cross component boundaries.
use_weights: boolean
Whether edges in the cluster graph should be weighted by the
``size_mid`` attribute of the connection. If not, then all connections
are weighted equally.
merge_nengo_nodes: boolean
If True, then clusters which would consist entirely of nengo Nodes are
merged with a neighboring cluster. This is done because it is typically
not useful to have a processor simulating only Nodes, as it will only
add extra communication without easing the computational burden.
Returns
-------
component0: NengoObjectCluster
A NengoObjectCluster containing all nengo objects which must be
simulated on the master process in the nengo_mpi simulator. If there
are no such objects, then this has value None.
cluster_graph: networkx.Graph
A graph wherein the nodes are instances of NengoObjectCluster.
Importantly, if ``component0`` is not None, then it is included
in ``cluster_graph``.
"""
def merge_clusters(cluster_map, a, b, conn=None):
if a.merge(b):
for obj in b.objects:
cluster_map[obj] = a
del b
if conn is not None:
a.connections.append(conn)
return a
if can_cross_boundary is None:
can_cross_boundary = make_boundary_predicate(network)
# A mapping from each nengo object to the cluster it is a member of
cluster_map = {obj: NengoObjectCluster(obj) for obj in network.all_nodes + network.all_ensembles}
for conn in network.all_connections:
pre_obj = neurons2ensemble(conn.pre_obj)
pre_cluster = cluster_map[pre_obj]
post_obj = neurons2ensemble(conn.post_obj)
post_cluster = cluster_map[post_obj]
if can_cross_boundary(conn):
pre_cluster.add_output(conn)
post_cluster.add_input(conn)
else:
merge_clusters(cluster_map, pre_cluster, post_cluster, conn)
all_clusters = list(set(cluster_map.values()))
_, outputs = find_all_io(network.all_connections)
# merge together all clusters that have to go on component 0
component0 = filter(lambda x: for_component0(x, outputs), all_clusters)
if component0:
all_clusters = filter(lambda x: x not in component0[1:], all_clusters)
component0 = reduce(lambda u, v: merge_clusters(cluster_map, u, v), component0)
else:
component0 = None
any_neurons = any(cluster.n_neurons > 0 for cluster in all_clusters)
if merge_nengo_nodes and any_neurons:
# For each cluster which does not contain any neurons, merge the
# cluster with another cluster which *does* contain neurons,
# and which the original cluster communicates strongly with.
without_neurons = filter(lambda c: c.n_neurons == 0, all_clusters)
for cluster in without_neurons:
# figure out which cluster would be most beneficial to merge with.
counts = defaultdict(int)
for i in cluster.inputs:
pre_obj = neurons2ensemble(i.pre_obj)
pre_cluster = cluster_map[pre_obj]
if pre_cluster.n_neurons > 0:
counts[pre_cluster] += i.size_mid
for o in cluster.outputs:
post_obj = neurons2ensemble(o.post_obj)
post_cluster = cluster_map[post_obj]
if post_cluster.n_neurons > 0:
counts[post_cluster] += o.size_mid
if counts:
best_cluster = max(counts, key=counts.__getitem__)
else:
best_cluster = (n for n in all_clusters if n.n_neurons > 0).next()
merge_clusters(cluster_map, best_cluster, cluster)
all_clusters.remove(cluster)
G = nx.Graph()
G.add_nodes_from(all_clusters)
boundary_connections = filter(can_cross_boundary, network.all_connections)
for conn in boundary_connections:
pre_cluster = cluster_map[neurons2ensemble(conn.pre_obj)]
post_cluster = cluster_map[neurons2ensemble(conn.post_obj)]
if pre_cluster != post_cluster:
weight = conn.size_mid if use_weights else 1.0
if G.has_edge(pre_cluster, post_cluster):
G[pre_cluster][post_cluster]["weight"] += weight
G[pre_cluster][post_cluster]["connections"].append(conn)
else:
G.add_edge(pre_cluster, post_cluster, weight=weight, connections=[conn])
return component0, G
def propogate_assignments(network, assignments, can_cross_boundary):
""" Assign every object in ``network`` to a component.
Propogates the component assignments stored in the dict ``assignments``
(which only needs to contain assignments for top level Networks, Nodes and
Ensembles) down to objects that are contained in those top-level objects.
If assignments is empty, then all objects will be assigned to the 1st
component, which has index 0. The intent is to have some partitioning
algorithm determine some of the assignments before this function is called,
and then have this function propogate those assignments.
Also does validation, making sure that connections that cross component
boundaries have certain properties (see ``can_cross_boundary``) and making
sure that certain types of objects are assigned to component 0.
Objects that must be simulated on component 0 are:
1. Nodes with callable outputs.
2. Ensembles of Direct neurons.
3. Any Node that is the source for a Connection that has a function.
Parameters
----------
network: nengo.Network
The network we are partitioning.
assignments: dict
A dictionary mapping from nengo objects to component indices.
This dictionary will be altered to contain assignments for all objects
in the network. If a network appears in assignments, then all objects
in that network which do not also appear in assignments will be given
the same assignment as the network.
can_cross_boundary: function
A function which accepts a Connection, and returns a boolean specifying
whether the Connection is allowed to cross component boundaries.
Returns
-------
Nothing, but ``assignments`` is modified.
"""
def helper(network, assignments, outputs):
for node in network.nodes:
if callable(node.output):
if node in assignments and assignments[node] != 0:
warnings.warn(
"Found Node with callable output that was assigned to "
"a component other than component 0. Overriding "
"previous assignment."
)
assignments[node] = 0
else:
if any([conn.function is not None for conn in outputs[node]]):
if node in assignments and assignments[node] != 0:
warnings.warn(
"Found Node with an output connection whose "
"function is not None, which is assigned to a "
"component other than component 0. Overriding "
"previous assignment."
)
assignments[node] = 0
elif node not in assignments:
assignments[node] = assignments[network]
for ensemble in network.ensembles:
if isinstance(ensemble.neuron_type, Direct):
if ensemble in assignments and assignments[ensemble] != 0:
warnings.warn(
"Found Direct-mode ensemble that was assigned to a "
"component other than component 0. Overriding "
"previous assignment."
)
assignments[ensemble] = 0
elif ensemble not in assignments:
assignments[ensemble] = assignments[network]
assignments[ensemble.neurons] = assignments[ensemble]
for n in network.networks:
if n not in assignments:
assignments[n] = assignments[network]
helper(n, assignments, outputs)
assignments[network] = 0
_, outputs = find_all_io(network.all_connections)
try:
helper(network, assignments, outputs)
# Assign learning rules
for conn in network.all_connections:
if conn.learning_rule is not None:
rule = conn.learning_rule
if is_iterable(rule):
rule = itervalues(rule) if isinstance(rule, dict) else rule
for r in rule:
assignments[r] = assignments[conn.pre_obj]
elif rule is not None:
assignments[rule] = assignments[conn.pre_obj]
# Check for connections erroneously crossing component boundaries
non_crossing = [conn for conn in network.all_connections if not can_cross_boundary(conn)]
for conn in non_crossing:
pre_component = assignments[conn.pre_obj]
post_component = assignments[conn.post_obj]
if pre_component != post_component:
raise PartitionError(
"Connection %s crosses a component "
"boundary, but it is not permitted to. "
"Pre-object assigned to %d, post-object "
"assigned to %d." % (conn, pre_component, post_component)
)
# Assign probes
for probe in network.all_probes:
target = probe.target.obj if isinstance(probe.target, ObjView) else probe.target
if isinstance(target, Connection):
target = target.pre_obj
assignments[probe] = assignments[target]
except KeyError as e:
# Nengo tests require a value error to be raised in these cases.
msg = ("Invalid Partition. KeyError: %s" % e.message,)
raise ValueError(msg)
nodes = network.all_nodes
nodes_in = all([node in assignments for node in nodes])
assert nodes_in, "Assignments incomplete, missing some nodes."
ensembles = network.all_ensembles
ensembles_in = all([ensemble in assignments for ensemble in ensembles])
assert ensembles_in, "Assignments incomplete, missing some ensembles."
def network_to_cluster_graph(
network, can_cross_boundary,
use_weights=True, merge_nengo_nodes=True):
""" Create a cluster graph from a nengo Network.
A cluster graph is a graph wherein the nodes are maximally large
NengoObjectClusters, and edges are nengo Connections that are permitted to
cross component boundaries.
Parameters
----------
network: nengo.Network
The network whose cluster graph we want to construct.
can_cross_boundary: function
A function which accepts a Connection, and returns a bool specifying
whether the Connection is allowed to cross component boundaries.
use_weights: boolean
Whether edges in the cluster graph should be weighted by the
``size_mid`` attribute of the connection. Otherwise, all connections
are weighted equally.
merge_nengo_nodes: boolean
If True, then clusters which would consist entirely of nengo Nodes are
merged with a neighboring cluster. This is done because it is typically
not useful to have a processor simulating only Nodes, as it will only
add extra communication without easing the computational burden.
Returns
-------
component0: NengoObjectCluster
A NengoObjectCluster containing all nengo objects which must be
simulated on the master process in the nengo_mpi simulator. If there
are no such objects, then this has value None.
cluster_graph: networkx.Graph
A graph wherein the nodes are instances of NengoObjectCluster.
Importantly, if ``component0`` is not None, then it is included
in ``cluster_graph``.
"""
cluster_graph = ClusterGraph(network, can_cross_boundary)
deferred = []
for conn in network.all_connections:
if (isinstance(conn.pre_obj, LearningRule) or
isinstance(conn.post_obj, LearningRule)):
deferred.append(conn)
continue
cluster_graph.process_conn(conn)
for conn in deferred:
cluster_graph.process_conn(conn)
_, outputs = find_all_io(network.all_connections)
# merge together all clusters that have to go on component 0
component0 = (
x for x in cluster_graph.clusters
if for_component0(x, outputs))
first = next(component0, None)
if first is not None:
for c in component0:
cluster_graph.merge_clusters(first, c)
component0 = first
else:
component0 = None
# Check whether there are any neurons in the network.
any_neurons = any(
cluster.n_neurons > 0 for cluster in cluster_graph.clusters)
if merge_nengo_nodes and any_neurons:
# For each cluster which does not contain any neurons, merge the
# cluster with another cluster which *does* contain neurons,
# and which the original cluster communicates strongly with.
clusters = cluster_graph.clusters
without_neurons = [c for c in clusters if c.n_neurons == 0]
with_neurons = [c for c in clusters if c.n_neurons > 0]
for cluster in without_neurons:
# figure out which cluster would be most beneficial to merge with.
counts = defaultdict(int)
for i in cluster.inputs:
pre_cluster = cluster_graph[i.pre_obj]
if pre_cluster.n_neurons > 0:
counts[pre_cluster] += i.size_mid
for o in cluster.outputs:
post_cluster = cluster_graph[o.post_obj]
if post_cluster.n_neurons > 0:
counts[post_cluster] += o.size_mid
if counts:
best_cluster = max(counts, key=counts.__getitem__)
else:
best_cluster = with_neurons[0]
cluster_graph.merge_clusters(best_cluster, cluster)
if cluster is component0:
component0 = best_cluster
G = cluster_graph.as_nx_graph(use_weights)
return component0, G
