Makes stochastic_graph work for multigraphs.

This commit allows the `stochastic_graph` generator to operate on
`MultiDiGraph` instances. It also adds some missing unit tests.

2 files changed, 88 insertions, 59 deletions

diff --git a/networkx/generators/stochastic.py b/networkx/generators/stochastic.py
index 9a6fb1af..73f7759f 100644
--- a/networkx/generators/stochastic.py
+++ b/networkx/generators/stochastic.py
@@ -1,60 +1,61 @@
-"""Functions for generating stochastic graphs from a given weighted directed
-graph.
-
-"""
 #    Copyright (C) 2010-2013 by
 #    Aric Hagberg <hagberg@lanl.gov>
 #    Dan Schult <dschult@colgate.edu>
 #    Pieter Swart <swart@lanl.gov>
 #    All rights reserved.
 #    BSD license.
+"""Functions for generating stochastic graphs from a given weighted directed
+graph.
+
+"""
 from __future__ import division
-import warnings
 
-import networkx as nx
+from functools import partial
+
+from networkx.classes import DiGraph
+from networkx.classes import MultiDiGraph
 from networkx.utils import not_implemented_for
 
 __author__ = "Aric Hagberg <aric.hagberg@gmail.com>"
 __all__ = ['stochastic_graph']
 
 
-@not_implemented_for('multigraph')
 @not_implemented_for('undirected')
 def stochastic_graph(G, copy=True, weight='weight'):
-    """Returns a right-stochastic representation of the directed graph ``G``.
+    """Returns a right-stochastic representation of the directed graph
+    ``G``.
 
-    A right-stochastic graph is a weighted digraph in which for each node, the
-    sum of the weights of all the out-edges of that node is 1. If the graph is
-    already weighted (for example, via a ``'weight'`` edge attribute), the
-    reweighting takes that into account.
+    A right-stochastic graph is a weighted digraph in which for each
+    node, the sum of the weights of all the out-edges of that node is
+    1. If the graph is already weighted (for example, via a ``'weight'``
+    edge attribute), the reweighting takes that into account.
 
     Parameters
     -----------
-    G : directed graph
-      A NetworkX DiGraph
+    G : NetworkX graph
+        A :class:`~networkx.DiGraph` or :class:`~networkx.MultiDiGraph`.
 
     copy : boolean, optional
-      If this is ``True``, then this function returns a new instance of
-      :class:`networkx.Digraph`. Otherwise, the original graph is modified
-      in-place (and also returned, for convenience).
+        If this is ``True``, then this function returns a new graph with
+        the stochastic reweighting. Otherwise, the original graph is
+        modified in-place (and also returned, for convenience).
 
     weight : edge attribute key (optional, default='weight')
-      Edge attribute key used for reading the existing weight and setting the
-      new weight.  If no attribute with this key is found for an edge, then the
-      edge weight is assumed to be 1. If an edge has a weight, it must be a
-      a positive number.
+        Edge attribute key used for reading the existing weight and
+        setting the new weight.  If no attribute with this key is found
+        for an edge, then the edge weight is assumed to be 1. If an edge
+        has a weight, it must be a a positive number.
 
     """
     if copy:
-        W = nx.DiGraph(G)
-    else:
-        # Reference the original graph, don't make a copy.
-        W = G
-    degree = dict(W.out_degree(weight=weight))
-    for (u, v, d) in W.edges(data=True):
+        G = MultiDiGraph(G) if G.is_multigraph() else DiGraph(G)
+    # There is a tradeoff here: the dictionary of node degrees may
+    # require a lot of memory, whereas making a call to `G.out_degree`
+    # inside the loop may be costly in computation time.
+    degree = dict(G.out_degree(weight=weight))
+    for u, v, d in G.edges(data=True):
         if degree[u] == 0:
-            warnings.warn('zero out-degree for node %s' % u)
             d[weight] = 0
         else:
             d[weight] = d.get(weight, 1) / degree[u]
-    return W
+    return G
diff --git a/networkx/generators/tests/test_stochastic.py b/networkx/generators/tests/test_stochastic.py
index a0c6c5e8..1c629154 100644
--- a/networkx/generators/tests/test_stochastic.py
+++ b/networkx/generators/tests/test_stochastic.py
@@ -1,33 +1,61 @@
+# test_stochastic.py - unit tests for the stochastic module
+#
+# Copyright 2010, 2011, 2012, 2013, 2014, 2015 NetworkX developers.
+#
+# This file is part of NetworkX.
+#
+# NetworkX is distributed under a BSD license; see LICENSE.txt for more
+# information.
+"""Unit tests for the :mod:`networkx.generators.stochastic` module."""
 from nose.tools import assert_true, assert_equal, raises
 import networkx as nx
 
-def test_stochastic():
-    G=nx.DiGraph()
-    G.add_edge(0,1)
-    G.add_edge(0,2)
-    S=nx.stochastic_graph(G)
-    assert_true(nx.is_isomorphic(G,S))
-    assert_equal(sorted(S.edges(data=True)),
-                 [(0, 1, {'weight': 0.5}), 
-                  (0, 2, {'weight': 0.5})])
-    S=nx.stochastic_graph(G,copy=True)
-    assert_equal(sorted(S.edges(data=True)),
-                 [(0, 1, {'weight': 0.5}), 
-                  (0, 2, {'weight': 0.5})])
-
-def test_stochastic_ints():
-    G=nx.DiGraph()
-    G.add_edge(0,1,weight=1)
-    G.add_edge(0,2,weight=1)
-    S=nx.stochastic_graph(G)
-    assert_equal(sorted(S.edges(data=True)),
-                 [(0, 1, {'weight': 0.5}), 
-                  (0, 2, {'weight': 0.5})])
-
-@raises(nx.NetworkXNotImplemented)
-def test_stochastic_graph_input():
-    S = nx.stochastic_graph(nx.Graph())
-
-@raises(nx.NetworkXNotImplemented)
-def test_stochastic_multigraph_input():
-    S = nx.stochastic_graph(nx.MultiGraph())
+
+class TestStochasticGraph(object):
+    """Unit tests for the :func:`~networkx.stochastic_graph` function.
+
+    """
+
+    def test_default_weights(self):
+        G = nx.DiGraph()
+        G.add_edge(0, 1)
+        G.add_edge(0, 2)
+        S = nx.stochastic_graph(G)
+        assert_true(nx.is_isomorphic(G, S))
+        assert_equal(sorted(S.edges(data=True)),
+                     [(0, 1, {'weight': 0.5}), (0, 2, {'weight': 0.5})])
+
+    def test_in_place(self):
+        """Tests for an in-place reweighting of the edges of the graph.
+
+        """
+        G = nx.DiGraph()
+        G.add_edge(0, 1, weight=1)
+        G.add_edge(0, 2, weight=1)
+        nx.stochastic_graph(G, copy=False)
+        assert_equal(sorted(G.edges(data=True)),
+                     [(0, 1, {'weight': 0.5}), (0, 2, {'weight': 0.5})])
+
+    def test_arbitrary_weights(self):
+        G = nx.DiGraph()
+        G.add_edge(0, 1, weight=1)
+        G.add_edge(0, 2, weight=1)
+        S = nx.stochastic_graph(G)
+        assert_equal(sorted(S.edges(data=True)),
+                     [(0, 1, {'weight': 0.5}), (0, 2, {'weight': 0.5})])
+
+    def test_multidigraph(self):
+        G = nx.MultiDiGraph()
+        G.add_edges_from([(0, 1), (0, 1), (0, 2), (0, 2)])
+        S = nx.stochastic_graph(G)
+        d = dict(weight=0.25)
+        assert_equal(sorted(S.edges(data=True)),
+                     [(0, 1, d), (0, 1, d), (0, 2, d), (0, 2, d)])
+
+    @raises(nx.NetworkXNotImplemented)
+    def test_graph_disallowed(self):
+        nx.stochastic_graph(nx.Graph())
+
+    @raises(nx.NetworkXNotImplemented)
+    def test_multigraph_disallowed(self):
+        nx.stochastic_graph(nx.MultiGraph())