Statistics
| Branch: | Revision:

iof-tools / networkxMiCe / networkx-master / networkx / readwrite / nx_shp.py @ 5cef0f13

History | View | Annotate | Download (11.3 KB)

1
"""
2
*********
3
Shapefile
4
*********
5

6
Generates a networkx.DiGraph from point and line shapefiles.
7

8
"The Esri Shapefile or simply a shapefile is a popular geospatial vector
9
data format for geographic information systems software. It is developed
10
and regulated by Esri as a (mostly) open specification for data
11
interoperability among Esri and other software products."
12
See https://en.wikipedia.org/wiki/Shapefile for additional information.
13
"""
14
#    Copyright (C) 2004-2019 by
15
#    Ben Reilly <benwreilly@gmail.com>
16
#    Aric Hagberg <hagberg@lanl.gov>
17
#    Dan Schult <dschult@colgate.edu>
18
#    Pieter Swart <swart@lanl.gov>
19
#    All rights reserved.
20
#    BSD license.
21
import networkx as nx
22
__author__ = """Ben Reilly (benwreilly@gmail.com)"""
23
__all__ = ['read_shp', 'write_shp']
24

    
25

    
26
def read_shp(path, simplify=True, geom_attrs=True, strict=True):
27
    """Generates a networkx.DiGraph from shapefiles. Point geometries are
28
    translated into nodes, lines into edges. Coordinate tuples are used as
29
    keys. Attributes are preserved, line geometries are simplified into start
30
    and end coordinates. Accepts a single shapefile or directory of many
31
    shapefiles.
32

33
    "The Esri Shapefile or simply a shapefile is a popular geospatial vector
34
    data format for geographic information systems software [1]_."
35

36
    Parameters
37
    ----------
38
    path : file or string
39
       File, directory, or filename to read.
40

41
    simplify:  bool
42
        If True, simplify line geometries to start and end coordinates.
43
        If False, and line feature geometry has multiple segments, the
44
        non-geometric attributes for that feature will be repeated for each
45
        edge comprising that feature.
46

47
    geom_attrs: bool
48
        If True, include the Wkb, Wkt and Json geometry attributes with
49
        each edge.
50

51
        NOTE:  if these attributes are available, write_shp will use them
52
        to write the geometry.  If nodes store the underlying coordinates for
53
        the edge geometry as well (as they do when they are read via
54
        this method) and they change, your geomety will be out of sync.
55

56
    strict: bool
57
        If True, raise NetworkXError when feature geometry is missing or
58
        GeometryType is not supported.
59
        If False, silently ignore missing or unsupported geometry in features.
60

61
    Returns
62
    -------
63
    G : NetworkX graph
64

65
    Raises
66
    ------
67
    ImportError
68
       If ogr module is not available.
69

70
    RuntimeError
71
       If file cannot be open or read.
72

73
    NetworkXError
74
       If strict=True and feature is missing geometry or GeometryType is
75
       not supported.
76

77
    Examples
78
    --------
79
    >>> G=nx.read_shp('test.shp') # doctest: +SKIP
80

81
    References
82
    ----------
83
    .. [1] https://en.wikipedia.org/wiki/Shapefile
84
    """
85
    try:
86
        from osgeo import ogr
87
    except ImportError:
88
        raise ImportError("read_shp requires OGR: http://www.gdal.org/")
89

    
90
    if not isinstance(path, str):
91
        return
92

    
93
    net = nx.DiGraph()
94
    shp = ogr.Open(path)
95
    if shp is None:
96
        raise RuntimeError("Unable to open {}".format(path))
97
    for lyr in shp:
98
        fields = [x.GetName() for x in lyr.schema]
99
        for f in lyr:
100
            g = f.geometry()
101
            if g is None:
102
                if strict:
103
                    raise nx.NetworkXError("Bad data: feature missing geometry")
104
                else:
105
                    continue
106
            flddata = [f.GetField(f.GetFieldIndex(x)) for x in fields]
107
            attributes = dict(zip(fields, flddata))
108
            attributes["ShpName"] = lyr.GetName()
109
            # Note:  Using layer level geometry type
110
            if g.GetGeometryType() == ogr.wkbPoint:
111
                net.add_node((g.GetPoint_2D(0)), **attributes)
112
            elif g.GetGeometryType() in (ogr.wkbLineString,
113
                                         ogr.wkbMultiLineString):
114
                for edge in edges_from_line(g, attributes, simplify,
115
                                            geom_attrs):
116
                    e1, e2, attr = edge
117
                    net.add_edge(e1, e2)
118
                    net[e1][e2].update(attr)
119
            else:
120
                if strict:
121
                    raise nx.NetworkXError("GeometryType {} not supported".
122
                                           format(g.GetGeometryType()))
123

    
124
    return net
125

    
126

    
127
def edges_from_line(geom, attrs, simplify=True, geom_attrs=True):
128
    """
129
    Generate edges for each line in geom
130
    Written as a helper for read_shp
131

132
    Parameters
133
    ----------
134

135
    geom:  ogr line geometry
136
        To be converted into an edge or edges
137

138
    attrs:  dict
139
        Attributes to be associated with all geoms
140

141
    simplify:  bool
142
        If True, simplify the line as in read_shp
143

144
    geom_attrs:  bool
145
        If True, add geom attributes to edge as in read_shp
146

147

148
    Returns
149
    -------
150
     edges:  generator of edges
151
        each edge is a tuple of form
152
        (node1_coord, node2_coord, attribute_dict)
153
        suitable for expanding into a networkx Graph add_edge call
154
    """
155
    try:
156
        from osgeo import ogr
157
    except ImportError:
158
        raise ImportError("edges_from_line requires OGR: http://www.gdal.org/")
159

    
160
    if geom.GetGeometryType() == ogr.wkbLineString:
161
        if simplify:
162
            edge_attrs = attrs.copy()
163
            last = geom.GetPointCount() - 1
164
            if geom_attrs:
165
                edge_attrs["Wkb"] = geom.ExportToWkb()
166
                edge_attrs["Wkt"] = geom.ExportToWkt()
167
                edge_attrs["Json"] = geom.ExportToJson()
168
            yield (geom.GetPoint_2D(0), geom.GetPoint_2D(last), edge_attrs)
169
        else:
170
            for i in range(0, geom.GetPointCount() - 1):
171
                pt1 = geom.GetPoint_2D(i)
172
                pt2 = geom.GetPoint_2D(i + 1)
173
                edge_attrs = attrs.copy()
174
                if geom_attrs:
175
                    segment = ogr.Geometry(ogr.wkbLineString)
176
                    segment.AddPoint_2D(pt1[0], pt1[1])
177
                    segment.AddPoint_2D(pt2[0], pt2[1])
178
                    edge_attrs["Wkb"] = segment.ExportToWkb()
179
                    edge_attrs["Wkt"] = segment.ExportToWkt()
180
                    edge_attrs["Json"] = segment.ExportToJson()
181
                    del segment
182
                yield (pt1, pt2, edge_attrs)
183

    
184
    elif geom.GetGeometryType() == ogr.wkbMultiLineString:
185
        for i in range(geom.GetGeometryCount()):
186
            geom_i = geom.GetGeometryRef(i)
187
            for edge in edges_from_line(geom_i, attrs, simplify, geom_attrs):
188
                yield edge
189

    
190

    
191
def write_shp(G, outdir):
192
    """Writes a networkx.DiGraph to two shapefiles, edges and nodes.
193
    Nodes and edges are expected to have a Well Known Binary (Wkb) or
194
    Well Known Text (Wkt) key in order to generate geometries. Also
195
    acceptable are nodes with a numeric tuple key (x,y).
196

197
    "The Esri Shapefile or simply a shapefile is a popular geospatial vector
198
    data format for geographic information systems software [1]_."
199

200
    Parameters
201
    ----------
202
    outdir : directory path
203
       Output directory for the two shapefiles.
204

205
    Returns
206
    -------
207
    None
208

209
    Examples
210
    --------
211
    nx.write_shp(digraph, '/shapefiles') # doctest +SKIP
212

213
    References
214
    ----------
215
    .. [1] https://en.wikipedia.org/wiki/Shapefile
216
    """
217
    try:
218
        from osgeo import ogr
219
    except ImportError:
220
        raise ImportError("write_shp requires OGR: http://www.gdal.org/")
221
    # easier to debug in python if ogr throws exceptions
222
    ogr.UseExceptions()
223

    
224
    def netgeometry(key, data):
225
        if 'Wkb' in data:
226
            geom = ogr.CreateGeometryFromWkb(data['Wkb'])
227
        elif 'Wkt' in data:
228
            geom = ogr.CreateGeometryFromWkt(data['Wkt'])
229
        elif type(key[0]).__name__ == 'tuple':  # edge keys are packed tuples
230
            geom = ogr.Geometry(ogr.wkbLineString)
231
            _from, _to = key[0], key[1]
232
            try:
233
                geom.SetPoint(0, *_from)
234
                geom.SetPoint(1, *_to)
235
            except TypeError:
236
                # assume user used tuple of int and choked ogr
237
                _ffrom = [float(x) for x in _from]
238
                _fto = [float(x) for x in _to]
239
                geom.SetPoint(0, *_ffrom)
240
                geom.SetPoint(1, *_fto)
241
        else:
242
            geom = ogr.Geometry(ogr.wkbPoint)
243
            try:
244
                geom.SetPoint(0, *key)
245
            except TypeError:
246
                # assume user used tuple of int and choked ogr
247
                fkey = [float(x) for x in key]
248
                geom.SetPoint(0, *fkey)
249

    
250
        return geom
251

    
252
    # Create_feature with new optional attributes arg (should be dict type)
253
    def create_feature(geometry, lyr, attributes=None):
254
        feature = ogr.Feature(lyr.GetLayerDefn())
255
        feature.SetGeometry(g)
256
        if attributes is not None:
257
            # Loop through attributes, assigning data to each field
258
            for field, data in attributes.items():
259
                feature.SetField(field, data)
260
        lyr.CreateFeature(feature)
261
        feature.Destroy()
262

    
263
    # Conversion dict between python and ogr types
264
    OGRTypes = {int: ogr.OFTInteger, str: ogr.OFTString, float: ogr.OFTReal}
265

    
266
    # Check/add fields from attribute data to Shapefile layers
267
    def add_fields_to_layer(key, value, fields, layer):
268
        # Field not in previous edges so add to dict
269
        if type(value) in OGRTypes:
270
            fields[key] = OGRTypes[type(value)]
271
        else:
272
            # Data type not supported, default to string (char 80)
273
            fields[key] = ogr.OFTString
274
        # Create the new field
275
        newfield = ogr.FieldDefn(key, fields[key])
276
        layer.CreateField(newfield)
277

    
278

    
279
    drv = ogr.GetDriverByName("ESRI Shapefile")
280
    shpdir = drv.CreateDataSource(outdir)
281
    # delete pre-existing output first otherwise ogr chokes
282
    try:
283
        shpdir.DeleteLayer("nodes")
284
    except:
285
        pass
286
    nodes = shpdir.CreateLayer("nodes", None, ogr.wkbPoint)
287

    
288
    # Storage for node field names and their data types
289
    node_fields = {}
290

    
291
    def create_attributes(data, fields, layer):
292
        attributes = {}  # storage for attribute data (indexed by field names)
293
        for key, value in data.items():
294
            # Reject spatial data not required for attribute table
295
            if (key != 'Json' and key != 'Wkt' and key != 'Wkb'
296
                    and key != 'ShpName'):
297
                # Check/add field and data type to fields dict
298
                if key not in fields:
299
                    add_fields_to_layer(key, value, fields, layer)
300
                # Store the data from new field to dict for CreateLayer()
301
                attributes[key] = value
302
        return attributes, layer
303

    
304
    for n in G:
305
        data = G.nodes[n]
306
        g = netgeometry(n, data)
307
        attributes, nodes = create_attributes(data, node_fields, nodes)
308
        create_feature(g, nodes, attributes)
309

    
310
    try:
311
        shpdir.DeleteLayer("edges")
312
    except:
313
        pass
314
    edges = shpdir.CreateLayer("edges", None, ogr.wkbLineString)
315

    
316
    # New edge attribute write support merged into edge loop
317
    edge_fields = {}      # storage for field names and their data types
318

    
319
    for e in G.edges(data=True):
320
        data = G.get_edge_data(*e)
321
        g = netgeometry(e, data)
322
        attributes, edges = create_attributes(e[2], edge_fields, edges)
323
        create_feature(g, edges, attributes)
324

    
325
    nodes, edges = None, None
326

    
327

    
328
# fixture for nose tests
329
def setup_module(module):
330
    from nose import SkipTest
331
    try:
332
        import ogr
333
    except:
334
        raise SkipTest("OGR not available")