1# Copyright (c) 2012, GPy authors (see AUTHORS.txt).
2# Licensed under the BSD 3-clause license (see LICENSE.txt)
3import numpy as np
4try:
5    from matplotlib import pyplot as pb
6    from matplotlib.patches import Polygon
7    from matplotlib.collections import PatchCollection
8    #from matplotlib import cm
9    try:
10        __IPYTHON__
11        pb.ion()
12    except NameError:
13        pass
14except:
15    pass
16import re
17
18def plot(shape_records,facecolor='w',edgecolor='k',linewidths=.5, ax=None,xlims=None,ylims=None):
19    """
20    Plot the geometry of a shapefile
21
22    :param shape_records: geometry and attributes list
23    :type shape_records: ShapeRecord object (output of a shapeRecords() method)
24    :param facecolor: color to be used to fill in polygons
25    :param edgecolor: color to be used for lines
26    :param ax: axes to plot on.
27    :type ax: axes handle
28    """
29    #Axes handle
30    if ax is None:
31        fig     = pb.figure()
32        ax      = fig.add_subplot(111)
33
34    #Iterate over shape_records
35    for srec in shape_records:
36        points = np.vstack(srec.shape.points)
37        sparts = srec.shape.parts
38        par = list(sparts) + [points.shape[0]]
39
40        polygs = []
41        for pj in range(len(sparts)):
42            polygs.append(Polygon(points[par[pj]:par[pj+1]]))
43        ax.add_collection(PatchCollection(polygs,facecolor=facecolor,edgecolor=edgecolor, linewidths=linewidths))
44
45    #Plot limits
46    _box = np.vstack([srec.shape.bbox for srec in shape_records])
47    minx,miny = np.min(_box[:,:2],0)
48    maxx,maxy = np.max(_box[:,2:],0)
49
50    if xlims is not None:
51        minx,maxx = xlims
52    if ylims is not None:
53        miny,maxy = ylims
54    ax.set_xlim(minx,maxx)
55    ax.set_ylim(miny,maxy)
56
57
58def string_match(sf,regex,field=2):
59    """
60    Return the geometry and attributes of a shapefile whose fields match a regular expression given
61
62    :param sf: shapefile
63    :type sf: shapefile object
64    :regex: regular expression to match
65    :type regex: string
66    :field: field number to be matched with the regex
67    :type field: integer
68    """
69    index = []
70    shape_records = []
71    for rec in enumerate(sf.shapeRecords()):
72        m = re.search(regex,rec[1].record[field])
73        if m is not None:
74            index.append(rec[0])
75            shape_records.append(rec[1])
76    return index,shape_records
77
78def bbox_match(sf,bbox,inside_only=True):
79    """
80    Return the geometry and attributes of a shapefile that lie within (or intersect) a bounding box
81
82    :param sf: shapefile
83    :type sf: shapefile object
84    :param bbox: bounding box
85    :type bbox: list of floats [x_min,y_min,x_max,y_max]
86    :inside_only: True if the objects returned are those that lie within the bbox and False if the objects returned are any that intersect the bbox
87    :type inside_only: Boolean
88    """
89    A,B,C,D = bbox
90    index = []
91    shape_records = []
92    for rec in enumerate(sf.shapeRecords()):
93        a,b,c,d = rec[1].shape.bbox
94        if inside_only:
95            if A <= a and B <= b and C >= c and D >= d:
96                index.append(rec[0])
97                shape_records.append(rec[1])
98        else:
99            cond1 = A <= a and B <= b and C >= a and D >= b
100            cond2 = A <= c and B <= d and C >= c and D >= d
101            cond3 = A <= a and D >= d and C >= a and B <= d
102            cond4 = A <= c and D >= b and C >= c and B <= b
103            cond5 = a <= C and b <= B and d >= D
104            cond6 = c <= A and b <= B and d >= D
105            cond7 = d <= B and a <= A and c >= C
106            cond8 = b <= D and a <= A and c >= C
107            if cond1 or cond2 or cond3 or cond4 or cond5 or cond6 or cond7 or cond8:
108                index.append(rec[0])
109                shape_records.append(rec[1])
110    return index,shape_records
111
112
113def plot_bbox(sf,bbox,inside_only=True):
114    """
115    Plot the geometry of a shapefile within a bbox
116
117    :param sf: shapefile
118    :type sf: shapefile object
119    :param bbox: bounding box
120    :type bbox: list of floats [x_min,y_min,x_max,y_max]
121    :inside_only: True if the objects returned are those that lie within the bbox and False if the objects returned are any that intersect the bbox
122    :type inside_only: Boolean
123    """
124    index,shape_records = bbox_match(sf,bbox,inside_only)
125    A,B,C,D = bbox
126    plot(shape_records,xlims=[bbox[0],bbox[2]],ylims=[bbox[1],bbox[3]])
127
128def plot_string_match(sf,regex,field,**kwargs):
129    """
130    Plot the geometry of a shapefile whose fields match a regular expression given
131
132    :param sf: shapefile
133    :type sf: shapefile object
134    :regex: regular expression to match
135    :type regex: string
136    :field: field number to be matched with the regex
137    :type field: integer
138    """
139    index,shape_records = string_match(sf,regex,field)
140    plot(shape_records,**kwargs)
141
142
143def new_shape_string(sf,name,regex,field=2,type=None):
144    import shapefile
145    if type is None:
146        type = shapefile.POINT
147    newshp = shapefile.Writer(shapeType = sf.shapeType)
148    newshp.autoBalance = 1
149
150    index,shape_records = string_match(sf,regex,field)
151
152    _fi = [sf.fields[j] for j in index]
153    for f in _fi:
154        newshp.field(name=f[0],fieldType=f[1],size=f[2],decimal=f[3])
155
156    _shre = shape_records
157    for sr in _shre:
158        _points = []
159        _parts = []
160        for point in sr.shape.points:
161            _points.append(point)
162        _parts.append(_points)
163
164        newshp.line(parts=_parts)
165        newshp.records.append(sr.record)
166        print(len(sr.record))
167
168    newshp.save(name)
169    print(index)
170
171def apply_bbox(sf,ax):
172    """
173    Use bbox as xlim and ylim in ax
174    """
175    limits = sf.bbox
176    xlim = limits[0],limits[2]
177    ylim = limits[1],limits[3]
178    ax.set_xlim(xlim)
179    ax.set_ylim(ylim)
180