summaryrefslogtreecommitdiff
path: root/networkx/generators/stochastic.py
blob: adf6be79d52a1765d7f86183c2713c19b143fe44 (plain)
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
"""Stocastic graph."""
import networkx as nx
#    Copyright (C) 2010 by 
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    All rights reserved.
#    BSD license.
__author__ = "Aric Hagberg <hagberg@lanl.gov>"
__all__ = ['stochastic_graph']

def stochastic_graph(G, copy=True, weight='weight'):
    """Return a right-stochastic representation of G.

    A right-stochastic graph is a weighted graph in which all of
    the node (out) neighbors edge weights sum to 1.
    
    Parameters
    -----------
    G : graph
      A NetworkX graph, must have valid edge weights

    copy : boolean, optional
      If True make a copy of the graph, otherwise modify original graph

    weight : key (optional)
      Edge data key used for weight.  If None all weights are set to 1.
    """        
    if type(G) == nx.MultiGraph or type(G) == nx.MultiDiGraph:
        raise Exception("stochastic_graph not implemented for multigraphs")

    if not G.is_directed():
        raise Exception("stochastic_graph not defined for undirected graphs")

    if copy:
        W=nx.DiGraph(G)
    else:
        W=G # reference original graph, no copy

    degree=W.out_degree(weight=weight)
    for (u,v,d) in W.edges(data=True):
        d[weight]=d.get(weight,1.0)/degree[u]
    return W