Statistics
| Branch: | Revision:

iof-tools / networkxMiCe / networkx-master / networkx / tests / test_convert_pandas.py @ 5cef0f13

History | View | Annotate | Download (7.83 KB)

1
from nose import SkipTest
2
from nose.tools import assert_raises
3

    
4
import networkx as nx
5
from networkx.testing import assert_nodes_equal, assert_edges_equal, assert_graphs_equal
6

    
7

    
8
class TestConvertPandas(object):
9
    numpy = 1  # nosetests attribute, use nosetests -a 'not numpy' to skip test
10

    
11
    @classmethod
12
    def setupClass(cls):
13
        try:
14
            import pandas as pd
15
        except ImportError:
16
            raise SkipTest('Pandas not available.')
17

    
18
    def __init__(self):
19
        global pd
20
        import pandas as pd
21

    
22
        self.rng = pd.np.random.RandomState(seed=5)
23
        ints = self.rng.randint(1, 11, size=(3, 2))
24
        a = ['A', 'B', 'C']
25
        b = ['D', 'A', 'E']
26
        df = pd.DataFrame(ints, columns=['weight', 'cost'])
27
        df[0] = a  # Column label 0 (int)
28
        df['b'] = b  # Column label 'b' (str)
29
        self.df = df
30
        mdf = pd.DataFrame([[4, 16, 'A', 'D']],
31
                           columns=['weight', 'cost', 0, 'b'])
32
        self.mdf = df.append(mdf)
33

    
34
    def test_exceptions(self):
35
        G = pd.DataFrame(["a"])  # adj
36
        assert_raises(nx.NetworkXError, nx.to_networkx_graph, G)
37
        G = pd.DataFrame(["a", 0.0])  # elist
38
        assert_raises(nx.NetworkXError, nx.to_networkx_graph, G)
39
        df = pd.DataFrame([[1, 1], [1, 0]], dtype=int, index=[1, 2], columns=["a", "b"])
40
        assert_raises(nx.NetworkXError, nx.from_pandas_adjacency, df)
41

    
42
    def test_from_edgelist_all_attr(self):
43
        Gtrue = nx.Graph([('E', 'C', {'cost': 9, 'weight': 10}),
44
                          ('B', 'A', {'cost': 1, 'weight': 7}),
45
                          ('A', 'D', {'cost': 7, 'weight': 4})])
46
        G = nx.from_pandas_edgelist(self.df, 0, 'b', True)
47
        assert_graphs_equal(G, Gtrue)
48
        # MultiGraph
49
        MGtrue = nx.MultiGraph(Gtrue)
50
        MGtrue.add_edge('A', 'D', cost=16, weight=4)
51
        MG = nx.from_pandas_edgelist(self.mdf, 0, 'b', True, nx.MultiGraph())
52
        assert_graphs_equal(MG, MGtrue)
53

    
54
    def test_from_edgelist_multi_attr(self):
55
        Gtrue = nx.Graph([('E', 'C', {'cost': 9, 'weight': 10}),
56
                          ('B', 'A', {'cost': 1, 'weight': 7}),
57
                          ('A', 'D', {'cost': 7, 'weight': 4})])
58
        G = nx.from_pandas_edgelist(self.df, 0, 'b', ['weight', 'cost'])
59
        assert_graphs_equal(G, Gtrue)
60

    
61
    def test_from_edgelist_multi_attr_incl_target(self):
62
        Gtrue = nx.Graph([('E', 'C', {0: 'C', 'b': 'E', 'weight': 10}),
63
                          ('B', 'A', {0: 'B', 'b': 'A', 'weight': 7}),
64
                          ('A', 'D', {0: 'A', 'b': 'D', 'weight': 4})])
65
        G = nx.from_pandas_edgelist(self.df, 0, 'b', [0, 'b', 'weight'])
66
        assert_graphs_equal(G, Gtrue)
67

    
68
    def test_from_edgelist_multidigraph_and_edge_attr(self):
69
        # example from issue #2374
70
        Gtrue = nx.MultiDiGraph([('X1', 'X4', {'Co': 'zA', 'Mi': 0, 'St': 'X1'}),
71
                                 ('X1', 'X4', {'Co': 'zB', 'Mi': 54, 'St': 'X2'}),
72
                                 ('X1', 'X4', {'Co': 'zB', 'Mi': 49, 'St': 'X3'}),
73
                                 ('X1', 'X4', {'Co': 'zB', 'Mi': 44, 'St': 'X4'}),
74
                                 ('Y1', 'Y3', {'Co': 'zC', 'Mi': 0, 'St': 'Y1'}),
75
                                 ('Y1', 'Y3', {'Co': 'zC', 'Mi': 34, 'St': 'Y2'}),
76
                                 ('Y1', 'Y3', {'Co': 'zC', 'Mi': 29, 'St': 'X2'}),
77
                                 ('Y1', 'Y3', {'Co': 'zC', 'Mi': 24, 'St': 'Y3'}),
78
                                 ('Z1', 'Z3', {'Co': 'zD', 'Mi': 0, 'St': 'Z1'}),
79
                                 ('Z1', 'Z3', {'Co': 'zD', 'Mi': 14, 'St': 'X3'}),
80
                                 ('Z1', 'Z3', {'Co': 'zE', 'Mi': 9, 'St': 'Z2'}),
81
                                 ('Z1', 'Z3', {'Co': 'zE', 'Mi': 4, 'St': 'Z3'})])
82
        df = pd.DataFrame.from_dict({
83
            'O': ['X1', 'X1', 'X1', 'X1', 'Y1', 'Y1', 'Y1', 'Y1', 'Z1', 'Z1', 'Z1', 'Z1'],
84
            'D': ['X4', 'X4', 'X4', 'X4', 'Y3', 'Y3', 'Y3', 'Y3', 'Z3', 'Z3', 'Z3', 'Z3'],
85
            'St': ['X1', 'X2', 'X3', 'X4', 'Y1', 'Y2', 'X2', 'Y3', 'Z1', 'X3', 'Z2', 'Z3'],
86
            'Co': ['zA', 'zB', 'zB', 'zB', 'zC', 'zC', 'zC', 'zC', 'zD', 'zD', 'zE', 'zE'],
87
            'Mi': [0,   54,   49,   44,    0,   34,   29,   24,    0,   14,    9,   4]})
88
        G1 = nx.from_pandas_edgelist(df, source='O', target='D',
89
                                     edge_attr=True,
90
                                     create_using=nx.MultiDiGraph)
91
        G2 = nx.from_pandas_edgelist(df, source='O', target='D',
92
                                     edge_attr=['St', 'Co', 'Mi'],
93
                                     create_using=nx.MultiDiGraph)
94
        assert_graphs_equal(G1, Gtrue)
95
        assert_graphs_equal(G2, Gtrue)
96

    
97
    def test_from_edgelist_one_attr(self):
98
        Gtrue = nx.Graph([('E', 'C', {'weight': 10}),
99
                          ('B', 'A', {'weight': 7}),
100
                          ('A', 'D', {'weight': 4})])
101
        G = nx.from_pandas_edgelist(self.df, 0, 'b', 'weight')
102
        assert_graphs_equal(G, Gtrue)
103

    
104
    def test_from_edgelist_int_attr_name(self):
105
        # note: this also tests that edge_attr can be `source`
106
        Gtrue = nx.Graph([('E', 'C', {0: 'C'}),
107
                          ('B', 'A', {0: 'B'}),
108
                          ('A', 'D', {0: 'A'})])
109
        G = nx.from_pandas_edgelist(self.df, 0, 'b', 0)
110
        assert_graphs_equal(G, Gtrue)
111

    
112
    def test_from_edgelist_invalid_attr(self):
113
        assert_raises(nx.NetworkXError, nx.from_pandas_edgelist,
114
                      self.df, 0, 'b', 'misspell')
115
        assert_raises(nx.NetworkXError, nx.from_pandas_edgelist,
116
                      self.df, 0, 'b', 1)
117

    
118
    def test_from_edgelist_no_attr(self):
119
        Gtrue = nx.Graph([('E', 'C', {}),
120
                          ('B', 'A', {}),
121
                          ('A', 'D', {})])
122
        G = nx.from_pandas_edgelist(self.df, 0, 'b',)
123
        assert_graphs_equal(G, Gtrue)
124

    
125
    def test_from_edgelist(self):
126
        # Pandas DataFrame
127
        g = nx.cycle_graph(10)
128
        G = nx.Graph()
129
        G.add_nodes_from(g)
130
        G.add_weighted_edges_from((u, v, u) for u, v in g.edges())
131
        edgelist = nx.to_edgelist(G)
132
        source = [s for s, t, d in edgelist]
133
        target = [t for s, t, d in edgelist]
134
        weight = [d['weight'] for s, t, d in edgelist]
135
        edges = pd.DataFrame({'source': source,
136
                              'target': target,
137
                              'weight': weight})
138
        GG = nx.from_pandas_edgelist(edges, edge_attr='weight')
139
        assert_nodes_equal(G.nodes(), GG.nodes())
140
        assert_edges_equal(G.edges(), GG.edges())
141
        GW = nx.to_networkx_graph(edges, create_using=nx.Graph)
142
        assert_nodes_equal(G.nodes(), GW.nodes())
143
        assert_edges_equal(G.edges(), GW.edges())
144

    
145
    def test_from_adjacency(self):
146
        nodelist = [1, 2]
147
        dftrue = pd.DataFrame([[1, 1], [1, 0]], dtype=int, index=nodelist, columns=nodelist)
148
        G = nx.Graph([(1, 1), (1, 2)])
149
        df = nx.to_pandas_adjacency(G, dtype=int)
150
        pd.testing.assert_frame_equal(df, dftrue)
151

    
152
    def test_roundtrip(self):
153
        # edgelist
154
        Gtrue = nx.Graph([(1, 1), (1, 2)])
155
        df = nx.to_pandas_edgelist(Gtrue)
156
        G = nx.from_pandas_edgelist(df)
157
        assert_graphs_equal(Gtrue, G)
158
        # adjacency
159
        Gtrue = nx.Graph(({1: {1: {'weight': 1}, 2: {'weight': 1}}, 2: {1: {'weight': 1}}}))
160
        df = nx.to_pandas_adjacency(Gtrue, dtype=int)
161
        G = nx.from_pandas_adjacency(df)
162
        assert_graphs_equal(Gtrue, G)
163

    
164
    def test_from_adjacency_named(self):
165
        # example from issue #3105
166
        data = {"A": {"A": 0, "B": 0, "C": 0},
167
                "B": {"A": 1, "B": 0, "C": 0},
168
                "C": {"A": 0, "B": 1, "C": 0}}
169
        dftrue = pd.DataFrame(data)
170
        df = dftrue[["A", "C", "B"]]
171
        G = nx.from_pandas_adjacency(df, create_using=nx.DiGraph())
172
        df = nx.to_pandas_adjacency(G, dtype=int)
173
        pd.testing.assert_frame_equal(df, dftrue)