def colorbar_legend(ax, values, cmap, vis=True):
"""
Add a vertical colorbar legend to a plot
"""
x_range = ax.get_xlim()[1]-ax.get_xlim()[0]
y_range = ax.get_ylim()[1]-ax.get_ylim()[0]
x = [ax.get_xlim()[0]+x_range*0.05]
y = [ax.get_ylim()[1]-(y_range * 0.25), ax.get_ylim()[1]-(y_range*0.05)]
segs = []
vals=[]
p = (x[0], y[0]+((y[1]-y[0])/256.0))
for i in range(2, 257):
n = (x[0], y[0]+((y[1]-y[0])/256.0)*i)
segs.append((p, n))
p = segs[-1][-1]
vals.append(min(values)+((max(values)-min(values))/256.0)*(i-1))
lcbar = LineCollection(segs, cmap=cmap, lw=15)
lcbar.set_visible(vis)
lcbar.set_array(np.array(vals))
ax.add_collection(lcbar)
lcbar.set_zorder(1)
minlab = str(min(values))[:6]
maxlab = str(max(values))[:6]
ax.text(x[0]+x_range*.02, y[0], minlab, verticalalignment="bottom", visible=vis)
ax.text(x[0]+x_range*.02, y[1], maxlab, verticalalignment="top", visible=vis)
def make_range_frame(self):
xminf, xmaxf = data_bounds_on_axis(self.axes.viewLim.intervalx, self.xbounds)
yminf, ymaxf = data_bounds_on_axis(self.axes.viewLim.intervaly, self.ybounds)
xline = [(xminf, 0), (xmaxf, 0)]
yline = [(0, yminf), (0, ymaxf)]
range_lines = LineCollection(segments=[xline, yline], linewidths=[self.linewidth], colors=[self.color])
range_lines.set_transform(self.axes.transAxes)
range_lines.set_zorder(10)
return range_lines
def _draw_arrows(self, G, pos, ax):
# Matplotlib's directed graph hack
# draw thick line segments at head end of edge
edge_pos = [(pos[n1], pos[n2]) for n1, n2 in self.G.edges()]
arrow_pos = self._calc_arrow_pos(edge_pos)
arrow_collection = LineCollection(arrow_pos,
colors = 'k',
linewidths = 4,
antialiaseds = (1,),
transOffset = ax.transData,
)
arrow_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(arrow_collection)
return arrow_collection
def show_localized(net, estimated, scale=False, align=True,\
display_loc_err=True, show_labels=True):
"""
Display estimated positions.
estimated should be a list of dictionaries.
"""
from matplotlib.pylab import gca
from matplotlib.collections import LineCollection
truePos = Positions.create(net.pos)
estimated = Positions.create(estimated)
# copy estimated so that passed estimated remains unchanged
estimated = deepcopy(estimated)
if align:
# rotate, translate and if needed scale estimated w.r.t. true positions
align_clusters(truePos, estimated, scale)
#TODO: implement display of all subclusters
if len(estimated)>1:
raise(NotImplementedError)
else:
estimated_sc = estimated[0]
#net.show(positions=estimated_sc, show_labels=show_labels)
fig = net.get_fig(positions=estimated_sc, show_labels=show_labels)
ax = fig.gca()
minpos = min(estimated_sc.values(), axis=0)
maxpos = max(estimated_sc.values(), axis=0)
minpos -= (maxpos-minpos)*0.1
maxpos += (maxpos-minpos)*0.1
ax.set_xlim(minpos[0], maxpos[0])
ax.set_ylim(minpos[1], maxpos[1])
#fig.show()
if display_loc_err:
#TODO: not working in ipython notepad
ax = gca()
ax.set_title('Localized positions')
ax.set_title('Localization error display')
edge_pos = asarray([(net.pos[n], estimated_sc[n])
for n in estimated_sc.keys()])
errorCollection = LineCollection(edge_pos, colors='r',
transOffset=ax.transData)
errorCollection.set_zorder(1) # errors go behind nodes
ax.add_collection(errorCollection)
ax.figure.canvas.draw()
def make_range_frame (self):
rx = self.axes.get_xlim()
ry = self.axes.get_ylim()
px = pl.prctile ( self.x )
py = pl.prctile ( self.y )
if self.trim:
if px[2]-px[0]>1.5*(px[3]-px[1]):
px[0] = self.x[self.x>px[2]-1.5*(px[3]-px[1])].min()
if px[4]-px[2]>1.5*(px[3]-px[1]):
px[4] = self.x[self.x<px[2]+1.5*(px[3]-px[1])].min()
x = px-rx[0]
x /= rx[1]-rx[0]
y = py-ry[0]
y /= ry[1]-ry[0]
ex = .003
ey = .003
xline = [
[(x[0],0),(x[1],0)],
[(x[1],ey),(x[2]-ex,ey)],
[(x[2]+ex,ey),(x[3],ey)],
[(x[3],0),(x[4],0)]
]
yline = [
[(0,y[0]),(0,y[1])],
[(ex,y[1]),(ex,y[2]-ey)],
[(ex,y[2]+ey),(ex,y[3])],
[(0,y[3]),(0,y[4])]
]
widths = [1,1,1,1]
range_lines = LineCollection(
segments=pl.clip(xline+yline,0,1),
linewidths=widths+widths,
colors=[[0]*3]*2*len(widths) )
range_lines.set_transform ( self.axes.transAxes )
range_lines.set_zorder(10)
self.axes.get_xaxis().tick_bottom()
self.axes.get_yaxis().tick_left()
self.axes.set_xticks(px)
self.axes.set_yticks(py)
self.axes.tick_params ( width=0 )
return range_lines
def make_box_plot (self):
x = self.x
rx = self.axes.get_xlim()
ry = self.axes.get_ylim()
ex = self.offset*(rx[1]-rx[0])
ey = self.offset*(ry[1]-ry[0])
if self.vert == 1:
ex,ey = ey,ex
p = self.boxstats['main']
n = self.boxstats['notch']
f_lo,f_hi = self.boxstats['fliers']
if self.notch:
lines = [ [(x,p[0]),(x,p[1])],
[(x+ex,p[1]),(x+ex,n[0])],
[(x+ex,n[0]),(x-ex,p[2]-ey)],
[(x-ex,p[2]+ey),(x+ex,n[1])],
[(x+ex,n[1]),(x+ex,p[3])],
[(x,p[3]),(x,p[4])] ]
else:
lines = [ [(x,p[0]),(x,p[1])],
[(x+ex,p[1]),(x+ex,p[2]-ey)],
[(x+ex,p[2]+ey),(x+ex,p[3])],
[(x,p[3]),(x,p[4])] ]
lines = pl.array(lines)
if self.vert==1:
lines = pl.array([ pl.c_[l[:,1],l[:,0]] for l in lines ])
pt = self.axes.plot ( f_lo, [x]*len(f_lo), '.', color=self.color,
markersize=self.lw ) + \
self.axes.plot ( f_hi, [x]*len(f_hi), '.', color=self.color,
markersize=self.lw )
dummy = self.axes.plot ( [p[0],p[-1]],[x-ex,x+ex], '.', markersize=0 )
else:
pt = self.axes.plot ( [x]*len(f_lo), f_lo, '.', color=self.color,
markersize=1 ) + \
self.axes.plot ( [x]*len(f_hi), f_hi, '.', color=self.color,
markersize=1 )
dummy = self.axes.plot ( [x-ex,x+ex], [p[0],p[-1]], '.', markersize=0 )
box = LineCollection (
segments=lines,
linewidths=[self.lw]*lines.shape[0],
colors=[self.color]*lines.shape[0] )
box.set_transform ( self.axes.transData )
box.set_zorder(10)
return box, pt[0],dummy[0]
def draw_networkx_edges_with_arrows(G,
pos,
width,
edge_color,
plt,
alpha=0.5,
ax=None):
ec = colorConverter.to_rgba(edge_color, alpha)
if ax is None:
ax = plt.gca()
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in G.edges()])
edge_collection = LineCollection(edge_pos,
colors=ec,
linewidths=width,
antialiaseds=(1, ),
linestyle='solid',
transOffset=ax.transData)
edge_collection.set_zorder(1)
ax.add_collection(edge_collection)
edge_collection.set_alpha(alpha)
p = .8 # length of edge apart from the arrow part
for (src, dst), lwi in zip(edge_pos, width):
x1, y1 = src
x2, y2 = dst
dx = x2 - x1 # x offset
dy = y2 - y1 # y offset
d = np.sqrt(float(dx**2 + dy**2)) # length of edge
if d == 0:
continue
if dx == 0: # vertical edge
xa = x2
ya = dy * p + y1
if dy == 0: # horizontal edge
ya = y2
xa = dx * p + x1
else:
theta = np.arctan2(dy, dx)
xa = p * d * np.cos(theta) + x1
ya = p * d * np.sin(theta) + y1
dx, dy = x2 - xa, y2 - ya
patch = mpatches.Arrow(xa,
ya,
dx,
dy,
width=lwi / 55,
alpha=lwi * alpha / 5,
color=ec,
transform=ax.transData)
ax.add_patch(patch)
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
return edge_collection
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=None,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
**kwds):
"""Draw the edges of the graph G
This draws only the edges of the graph G.
pos is a dictionary keyed by vertex with a two-tuple
of x-y positions as the value.
See networkx.layout for functions that compute node positions.
edgelist is an optional list of the edges in G to be drawn.
If provided, only the edges in edgelist will be drawn.
edgecolor can be a list of matplotlib color letters such as 'k' or
'b' that lists the color of each edge; the list must be ordered in
the same way as the edge list. Alternatively, this list can contain
numbers and those number are mapped to a color scale using the color
map edge_cmap. Finally, it can also be a list of (r,g,b) or (r,g,b,a)
tuples, in which case these will be used directly to color the edges. If
the latter mode is used, you should not provide a value for alpha, as it
would be applied globally to all lines.
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
See draw_networkx for the list of other optional parameters.
"""
try:
import matplotlib
import matplotlib.pylab as pylab
import numpy as np
from matplotlib.colors import colorConverter,Colormap
from matplotlib.collections import LineCollection
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
pass # unable to open display
if ax is None:
ax=pylab.gca()
if edgelist is None:
edgelist=G.edges()
if not edgelist or len(edgelist)==0: # no edges!
return None
# set edge positions
edge_pos=np.asarray([(pos[e[0]],pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color)==len(edge_pos):
if np.alltrue([cb.is_string_like(c)
for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c,alpha)
for c in edge_color])
elif np.alltrue([not cb.is_string_like(c)
for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if np.alltrue([cb.iterable(c) and len(c) in (3,4)
for c in edge_color]):
edge_colors = tuple(edge_color)
alpha=None
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must consist of either color names or numbers')
else:
if len(edge_color)==1:
edge_colors = ( colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError('edge_color must be a single color or list of exactly m colors where m is the number or edges')
edge_collection = LineCollection(edge_pos,
colors = edge_colors,
linewidths = lw,
antialiaseds = (1,),
linestyle = style,
transOffset = ax.transData,
)
# Note: there was a bug in mpl regarding the handling of alpha values for
# each line in a LineCollection. It was fixed in matplotlib in r7184 and
# r7189 (June 6 2009). We should then not set the alpha value globally,
# since the user can instead provide per-edge alphas now. Only set it
# globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
# need 0.87.7 or greater for edge colormaps. No checks done, this will
# just not work with an older mpl
if edge_colors is None:
if edge_cmap is not None: assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(np.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
pylab.sci(edge_collection)
arrow_collection=None
if G.is_directed() and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = ( colorConverter.to_rgba('k', alpha), )
a_pos=[]
p=1.0-0.25 # make head segment 25 percent of edge length
for src,dst in edge_pos:
x1,y1=src
x2,y2=dst
dx=x2-x1 # x offset
dy=y2-y1 # y offset
d=np.sqrt(float(dx**2+dy**2)) # length of edge
if d==0: # source and target at same position
continue
if dx==0: # vertical edge
xa=x2
ya=dy*p+y1
if dy==0: # horizontal edge
ya=y2
xa=dx*p+x1
else:
theta=np.arctan2(dy,dx)
xa=p*d*np.cos(theta)+x1
ya=p*d*np.sin(theta)+y1
a_pos.append(((xa,ya),(x2,y2)))
arrow_collection = LineCollection(a_pos,
colors = arrow_colors,
linewidths = [4*ww for ww in lw],
antialiaseds = (1,),
transOffset = ax.transData,
)
# update view
minx = np.amin(np.ravel(edge_pos[:,:,0]))
maxx = np.amax(np.ravel(edge_pos[:,:,0]))
miny = np.amin(np.ravel(edge_pos[:,:,1]))
maxy = np.amax(np.ravel(edge_pos[:,:,1]))
w = maxx-minx
h = maxy-miny
padx, pady = 0.05*w, 0.05*h
corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady)
ax.update_datalim( corners)
ax.autoscale_view()
edge_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(edge_collection)
if arrow_collection:
arrow_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(arrow_collection)
return edge_collection
def __init__(self,
ax,
lons,
lats,
basemap_args=None,
zval=None,
alongtrack_args=None,
colorbar=None,
calval=None):
""" Generate a simple map with orbit location """
grid = {
'coarse': {
'color': '1.0',
'dashes': [],
'linewidth': 0.25,
'fontsize': 8,
'zorder': 50
},
'fine': {
'color': '1.0',
'dashes': [],
'linewidth': 0.1,
'fontsize': 8,
'zorder': 50
}
}
# label = "$source$:%s, $orbit$:%s" % (
# dataset.source_id, dataset.orbit)
if basemap_args is None:
basemap_args = get_basemap_args_from_positions(lons,
lats,
scale=1.5)
m = Basemap(ax=ax, **basemap_args)
m.drawmapboundary(linewidth=0.1, fill_color='#CCE5FF', zorder=200)
m.drawcoastlines(linewidth=0.25, color="0.5")
m.fillcontinents(color="1.0", lake_color="0.96", zorder=100)
for type in ['coarse']:
parallels, keyw = self._get_parallels(grid, type)
m.drawparallels(parallels, **keyw)
meridians, keyw = self._get_meridians(grid, type)
m.drawmeridians(meridians, **keyw)
px, py = m(lons, lats)
if zval is not None:
points = np.array([px, py]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = LineCollection(segments, **alongtrack_args)
lc.set_array(zval)
lc.set_linewidth(3)
lc.set_zorder(200)
ax.add_collection(lc)
else:
m.scatter(px, py, marker=".", zorder=200)
if zval is not None and colorbar is not None:
plt.colorbar(lc,
ax=ax,
shrink=0.7,
aspect=15,
drawedges=False,
orientation="horizontal",
**colorbar)
if calval is not None:
for calval_dset in calval:
lon, lat, props = calval_dset
m.scatter(lon, lat, zorder=199, latlon=True, **props)
def draw_networkx_edges(G,
pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
arrowstyle='thick',
label=None,
**kwds):
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
import matplotlib.patches as patches
from matplotlib.colors import colorConverter, Colormap
from matplotlib.collections import LineCollection
from matplotlib.path import Path
import numpy
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
# print "drawing_edges"
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = G.edges()
if not edgelist or len(edgelist) == 0: # no edges!
return None
# set edge positions
edge_pos = numpy.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
# for e in edge_pos:
# print e
if not cb.iterable(width):
lw = (width, )
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color) == len(edge_pos):
if numpy.alltrue([cb.is_string_like(c) for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple(
[colorConverter.to_rgba(c, alpha) for c in edge_color])
elif numpy.alltrue([not cb.is_string_like(c) for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if numpy.alltrue(
[cb.iterable(c) and len(c) in (3, 4) for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError(
'edge_color must consist of either color names or numbers')
else:
if cb.is_string_like(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError(
'edge_color must be a single color or list of exactly m colors where m is the number or edges'
)
edge_collection = LineCollection(
edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
)
# print type(edge_collection)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
# Note: there was a bug in mpl regarding the handling of alpha values for
# each line in a LineCollection. It was fixed in matplotlib in r7184 and
# r7189 (June 6 2009). We should then not set the alpha value globally,
# since the user can instead provide per-edge alphas now. Only set it
# globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert (isinstance(edge_cmap, Colormap))
edge_collection.set_array(numpy.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
arrow_collection = None
if G.is_directed() and arrows:
# a directed graph hack-fix
# draws arrows at each
# waiting for someone else to implement arrows that will work
arrow_colors = edge_colors
a_pos = []
p = .1 # make arrows 10% of total length
angle = 2.7 #angle for arrows
for src, dst in edge_pos:
x1, y1 = src
x2, y2 = dst
dx = x2 - x1 # x offset
dy = y2 - y1 # y offset
d = numpy.sqrt(float(dx**2 + dy**2)) # length of edge
theta = numpy.arctan2(dy, dx)
if d == 0: # source and target at same position
continue
if dx == 0: # vertical edge
xa = x2
ya = dy + y1
if dy == 0: # horizontal edge
ya = y2
xa = dx + x1
else:
# xa = p*d*numpy.cos(theta)+x1
# ya = p*d*numpy.sin(theta)+y1
#corrects the endpoints to better draw
x2 -= .04 * numpy.cos(theta)
y2 -= .04 * numpy.sin(theta)
lx1 = p * d * numpy.cos(theta + angle) + (x2)
lx2 = p * d * numpy.cos(theta - angle) + (x2)
ly1 = p * d * numpy.sin(theta + angle) + (y2)
ly2 = p * d * numpy.sin(theta - angle) + (y2)
a_pos.append(((lx1, ly1), (x2, y2)))
a_pos.append(((lx2, ly2), (x2, y2)))
arrow_collection = LineCollection(
a_pos,
colors=arrow_colors,
linewidths=[1 * ww for ww in lw],
antialiaseds=(1, ),
transOffset=ax.transData,
)
arrow_collection.set_zorder(1) # edges go behind nodes
arrow_collection.set_label(label)
# print type(ax)
ax.add_collection(arrow_collection)
#drawing self loops
d = 1
c = 0.0707
selfedges = []
verts = [
(0.1 * d - 0.1 * d, 0.0), # P0
(c * d - 0.1 * d, c * d), # P0
(0.0 - 0.1 * d, 0.1 * d), # P0
(-c * d - 0.1 * d, c * d), # P0
(-0.1 * d - 0.1 * d, 0.0), # P0
(-c * d - 0.1 * d, -c * d), # P0
(0.0 - 0.1 * d, -0.1 * d), # P0
(c * d - 0.1 * d, -c * d), # P0
(0.1 * d - 0.1 * d, 0.0)
]
# print verts
codes = [
Path.MOVETO,
Path.LINETO,
Path.LINETO,
Path.LINETO,
Path.LINETO,
Path.LINETO,
Path.LINETO,
Path.LINETO,
Path.LINETO,
]
for e in edge_pos:
if (numpy.array_equal(e[0], e[1])):
nodes = verts[:]
for i in range(len(nodes)):
nodes[i] += e[0]
# print nodes
path = Path(nodes, codes)
patch = patches.PathPatch(path,
color=None,
facecolor=None,
edgecolor=edge_colors[0],
fill=False,
lw=4)
ax.add_patch(patch)
# update view
minx = numpy.amin(numpy.ravel(edge_pos[:, :, 0]))
maxx = numpy.amax(numpy.ravel(edge_pos[:, :, 0]))
miny = numpy.amin(numpy.ravel(edge_pos[:, :, 1]))
maxy = numpy.amax(numpy.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
# print ax
ax.update_datalim(corners)
ax.autoscale_view()
# if arrow_collection:
return edge_collection
def odd_import(self, scenario):
try:
lines1 = scenario['lines']['boundary_lines']
lines2 = scenario['lines']['solid_lines']
lines3 = scenario['lines']['dashed_lines']
lines4 = scenario['lines']['thick_lines']
objects = scenario['objects']
common_rotation = scenario['common_rotation']
common_offset = scenario['common_offset']
except Exception:
lines1 = []
lines2 = []
lines3 = []
lines4 = []
objects = []
common_rotation = 0.0
common_offset = [0, 0]
QMessageBox.warning(self, 'Warning', 'Scenario file format error!')
boundary_lines = self.line_rotation(lines1, common_rotation)
solid_lines = self.line_rotation(lines2, common_rotation)
dashed_lines = self.line_rotation(lines3, common_rotation)
thick_lines = self.line_rotation(lines4, common_rotation)
for l in boundary_lines:
for m in l:
m[0] = (m[0] + common_offset[0]) * self.zoom_factor
m[1] = (m[1] + common_offset[1]) * self.zoom_factor
for l in solid_lines:
for m in l:
m[0] = (m[0] + common_offset[0]) * self.zoom_factor
m[1] = (m[1] + common_offset[1]) * self.zoom_factor
for l in dashed_lines:
for m in l:
m[0] = (m[0] + common_offset[0]) * self.zoom_factor
m[1] = (m[1] + common_offset[1]) * self.zoom_factor
for l in thick_lines:
for m in l:
m[0] = (m[0] + common_offset[0]) * self.zoom_factor
m[1] = (m[1] + common_offset[1]) * self.zoom_factor
line_segments1 = LineCollection(boundary_lines, linewidths=2, colors=(0, 0, 0), linestyles='solid')
line_segments1.set_zorder(0)
line_segments2 = LineCollection(solid_lines, linewidths=1, colors=(0.1, 0.1, 0.1), linestyles='solid')
line_segments2.set_zorder(0)
line_segments3 = LineCollection(dashed_lines, linewidths=1, colors=(0.2, 0.2, 0.2), linestyles='dashed')
line_segments3.set_zorder(0)
line_segments4 = LineCollection(thick_lines, linewidths=4.5, colors=(0.8, 0.8, 0.8), linestyles='solid')
line_segments4.set_zorder(0)
self.ax.add_collection(line_segments1)
self.ax.add_collection(line_segments2)
self.ax.add_collection(line_segments3)
self.ax.add_collection(line_segments4)
patch = []
for obj in range(len(objects)):
objects_xy = self.object_rotation(objects[obj][1], common_rotation)
objects_xy[0] = objects_xy[0] + common_offset[0]
objects_xy[1] = objects_xy[1] + common_offset[1]
objects_r = objects[obj][2] + common_rotation
patch.append(self.add_object(objects[obj][0], objects_xy, objects_r))
self.ax.text((objects_xy[0]) * self.zoom_factor, (objects_xy[1] + 2.5) * self.zoom_factor,
'obj' + str(obj + 1), fontsize=10)
patch_segment = PatchCollection(patch, facecolors='gray')
patch_segment.set_zorder(4)
self.ax.add_collection(patch_segment)
return patch
def plot(n,
margin=0.05,
ax=None,
geomap=True,
projection=None,
bus_colors='b',
bus_alpha=1,
line_colors={
'Line': 'g',
'Link': 'cyan'
},
bus_sizes=1e-2,
line_widths={
'Line': 2,
'Link': 2
},
flow=None,
layouter=None,
title="",
line_cmap=None,
bus_cmap=None,
boundaries=None,
geometry=False,
branch_components=['Line', 'Link'],
jitter=None,
color_geomap=None):
"""
Plot the network buses and lines using matplotlib and cartopy.
Parameters
----------
margin : float
Margin at the sides as proportion of distance between max/min x,y
ax : matplotlib ax, defaults to plt.gca()
Axis to which to plot the network
geomap: bool/str, default True
Switch to use Cartopy and draw geographical features.
If string is passed, it will be used as a resolution argument,
valid options are '10m', '50m' and '110m'.
projection: cartopy.crs.Projection, defaults to None
Define the projection of your geomap, only valid if cartopy is
installed. If None (default) is passed the projection for cartopy
is set to cartopy.crs.PlateCarree
bus_colors : dict/pandas.Series
Colors for the buses, defaults to "b". If bus_sizes is a pandas.Series
with a Multiindex, bus_colors defaults to the n.carriers['color']
column.
bus_sizes : dict/pandas.Series
Sizes of bus points, defaults to 1e-2. If a multiindexed Series is passed,
the function will draw pies for each bus (first index level) with
segments of different color (second index level). Such a Series is ob-
tained by e.g. n.generators.groupby(['bus', 'carrier']).p_nom.sum()
bus_alpha : float
Adds alpha channel to buses, defaults to 1.
line_colors : dict/pandas.Series
Colors for the lines, defaults to "g" for Lines and "cyan" for
Links. Colors for branches other than Lines can be
specified using a pandas Series with a MultiIndex.
line_widths : dict/pandas.Series
Widths of lines, defaults to 2. Widths for branches other
than Lines can be specified using a pandas Series with a
MultiIndex.
flow : snapshot/pandas.Series/function/string
Flow to be displayed in the plot, defaults to None. If an element of
n.snapshots is given, the flow at this timestamp will be
displayed. If an aggregation function is given, is will be applied
to the total network flow via pandas.DataFrame.agg (accepts also
function names). Otherwise flows can be specified by passing a pandas
Series with MultiIndex including all necessary branch components.
Use the line_widths argument to additionally adjust the size of the
flow arrows.
layouter : networkx.drawing.layout function, default None
Layouting function from `networkx <https://networkx.github.io/>`_ which
overrules coordinates given in ``n.buses[['x','y']]``. See
`list <https://networkx.github.io/documentation/stable/reference/drawing.html#module-networkx.drawing.layout>`_
of available options.
title : string
Graph title
line_cmap : plt.cm.ColorMap/str|dict
If line_colors are floats, this color map will assign the colors.
Use a dict to specify colormaps for more than one branch type.
bus_cmap : plt.cm.ColorMap/str
If bus_colors are floats, this color map will assign the colors
boundaries : list of four floats
Boundaries of the plot in format [x1,x2,y1,y2]
branch_components : list of str
Branch components to be plotted, defaults to Line and Link.
jitter : None|float
Amount of random noise to add to bus positions to distinguish
overlapping buses
color_geomap : dict or bool
Specify colors to paint land and sea areas in.
If True, it defaults to `{'ocean': 'lightblue', 'land': 'whitesmoke'}`.
If no dictionary is provided, colors are white.
Returns
-------
bus_collection, branch_collection1, ... : tuple of Collections
Collections for buses and branches.
"""
defaults_for_branches = pd.Series({
'Link':
dict(color="cyan", width=2),
'Line':
dict(color="b", width=2),
'Transformer':
dict(color='green', width=2)
}).rename_axis('component')
x, y = _get_coordinates(n, layouter=layouter)
if geomap:
if not cartopy_present:
logger.warning(
"Cartopy needs to be installed to use `geomap=True`.")
geomap = False
if projection is None:
projection = get_projection_from_crs(n.srid)
if ax is None:
ax = plt.gca(projection=projection)
else:
assert isinstance(ax, cartopy.mpl.geoaxes.GeoAxesSubplot), (
'The passed axis is not a GeoAxesSubplot. You can '
'create one with: \nimport cartopy.crs as ccrs \n'
'fig, ax = plt.subplots('
'subplot_kw={"projection":ccrs.PlateCarree()})')
transform = draw_map_cartopy(n, x, y, ax, boundaries, margin, geomap,
color_geomap)
x, y, z = ax.projection.transform_points(transform, x.values,
y.values).T
x, y = pd.Series(x, n.buses.index), pd.Series(y, n.buses.index)
elif ax is None:
ax = plt.gca()
if jitter is not None:
x = x + np.random.uniform(low=-jitter, high=jitter, size=len(x))
y = y + np.random.uniform(low=-jitter, high=jitter, size=len(y))
if isinstance(bus_sizes, pd.Series) and isinstance(bus_sizes.index,
pd.MultiIndex):
# We are drawing pies to show all the different shares
assert len(bus_sizes.index.levels[0].difference(n.buses.index)) == 0, \
"The first MultiIndex level of bus_sizes must contain buses"
if isinstance(bus_colors, dict):
bus_colors = pd.Series(bus_colors)
# case bus_colors isn't a series or dict: look in n.carriers for existent colors
if not isinstance(bus_colors, pd.Series):
bus_colors = n.carriers.color.dropna()
assert bus_sizes.index.levels[1].isin(bus_colors.index).all(), (
"Colors not defined for all elements in the second MultiIndex "
"level of bus_sizes, please make sure that all the elements are "
"included in bus_colors or in n.carriers.color")
bus_sizes = bus_sizes.sort_index(level=0, sort_remaining=False)
if geomap:
bus_sizes *= projected_area_factor(ax, n.srid)**2
patches = []
for b_i in bus_sizes.index.levels[0]:
s = bus_sizes.loc[b_i]
radius = s.sum()**0.5
if radius == 0.0:
ratios = s
else:
ratios = s / s.sum()
start = 0.25
for i, ratio in ratios.iteritems():
patches.append(
Wedge((x.at[b_i], y.at[b_i]),
radius,
360 * start,
360 * (start + ratio),
facecolor=bus_colors[i],
alpha=bus_alpha))
start += ratio
bus_collection = PatchCollection(patches, match_original=True)
ax.add_collection(bus_collection)
else:
c = pd.Series(bus_colors, index=n.buses.index)
s = pd.Series(bus_sizes, index=n.buses.index, dtype="float")
if geomap:
s *= projected_area_factor(ax, n.srid)**2
if bus_cmap is not None and c.dtype is np.dtype('float'):
if isinstance(bus_cmap, str):
bus_cmap = plt.cm.get_cmap(bus_cmap)
norm = plt.Normalize(vmin=c.min(), vmax=c.max())
c = c.apply(lambda cval: bus_cmap(norm(cval)))
patches = []
for b_i in s.index:
radius = s.at[b_i]**0.5
patches.append(
Circle((x.at[b_i], y.at[b_i]),
radius,
facecolor=c.at[b_i],
alpha=bus_alpha))
bus_collection = PatchCollection(patches, match_original=True)
ax.add_collection(bus_collection)
def as_branch_series(ser):
# ensure that this function always return a multiindexed series
if isinstance(ser, dict) and set(ser).issubset(branch_components):
return pd.concat(
{
c.name: pd.Series(s, index=c.df.index)
for c, s in zip(n.iterate_components(ser.keys()),
ser.values())
},
names=['component', 'name'])
elif isinstance(ser, pd.Series) and isinstance(ser.index,
pd.MultiIndex):
return ser.rename_axis(index=['component', 'name'])
else:
ser = pd.Series(ser, n.lines.index)
return pd.concat([ser],
axis=0,
keys=['Line'],
names=['component', 'name']).fillna(0)
line_colors = as_branch_series(line_colors)
line_widths = as_branch_series(line_widths)
if not isinstance(line_cmap, dict):
line_cmap = {'Line': line_cmap}
branch_collections = []
if flow is not None:
flow = (_flow_ds_from_arg(
flow, n, branch_components).pipe(as_branch_series).div(
sum(
len(t.df)
for t in n.iterate_components(branch_components)) + 100))
flow = flow.mul(line_widths[flow.index], fill_value=1)
# update the line width, allows to set line widths separately from flows
line_widths.update((5 * flow.abs()).pipe(np.sqrt))
arrows = directed_flow(n,
flow,
x=x,
y=y,
ax=ax,
geomap=geomap,
branch_colors=line_colors,
branch_comps=branch_components,
cmap=line_cmap['Line'])
branch_collections.append(arrows)
for c in n.iterate_components(branch_components):
l_defaults = defaults_for_branches[c.name]
l_widths = line_widths.get(c.name, l_defaults['width'])
l_nums = None
l_colors = line_colors.get(c.name, l_defaults['color'])
if isinstance(l_colors, pd.Series):
if issubclass(l_colors.dtype.type, np.number):
l_nums = l_colors
l_colors = None
else:
l_colors.fillna(l_defaults['color'], inplace=True)
if not geometry:
segments = (np.asarray(
((c.df.bus0.map(x), c.df.bus0.map(y)),
(c.df.bus1.map(x), c.df.bus1.map(y)))).transpose(2, 0, 1))
else:
from shapely.wkt import loads
from shapely.geometry import LineString
linestrings = c.df.geometry[lambda ds: ds != ''].map(loads)
assert all(isinstance(ls, LineString) for ls in linestrings), (
"The WKT-encoded geometry in the 'geometry' column must be "
"composed of LineStrings")
segments = np.asarray(list(linestrings.map(np.asarray)))
l_collection = LineCollection(segments,
linewidths=l_widths,
antialiaseds=(1, ),
colors=l_colors,
transOffset=ax.transData)
if l_nums is not None:
l_collection.set_array(np.asarray(l_nums))
l_collection.set_cmap(line_cmap.get(c.name, None))
l_collection.autoscale()
ax.add_collection(l_collection)
l_collection.set_zorder(3)
branch_collections.append(l_collection)
bus_collection.set_zorder(4)
ax.update_datalim(compute_bbox_with_margins(margin, x, y))
ax.autoscale_view()
if geomap:
ax.outline_patch.set_visible(False)
ax.axis('off')
else:
ax.set_aspect('equal')
ax.set_title(title)
return (bus_collection, ) + tuple(branch_collections)
def plot(network, margin=0.05, ax=None, basemap=True, bus_colors='b',
line_colors='g', bus_sizes=10, line_widths=2, title="",
line_cmap=None, bus_cmap=None, boundaries=None,
geometry=False, branch_components=['Line', 'Link'], jitter=None):
"""
Plot the network buses and lines using matplotlib and Basemap.
Parameters
----------
margin : float
Margin at the sides as proportion of distance between max/min x,y
ax : matplotlib ax, defaults to plt.gca()
Axis to which to plot the network
basemap : bool, default True
Switch to use Basemap
bus_colors : dict/pandas.Series
Colors for the buses, defaults to "b"
bus_sizes : dict/pandas.Series
Sizes of bus points, defaults to 10
line_colors : dict/pandas.Series
Colors for the lines, defaults to "g" for Lines and "cyan" for
Links. Colors for branches other than Lines can be
specified using a pandas Series with a MultiIndex.
line_widths : dict/pandas.Series
Widths of lines, defaults to 2. Widths for branches other
than Lines can be specified using a pandas Series with a
MultiIndex.
title : string
Graph title
line_cmap : plt.cm.ColorMap/str|dict
If line_colors are floats, this color map will assign the colors.
Use a dict to specify colormaps for more than one branch type.
bus_cmap : plt.cm.ColorMap/str
If bus_colors are floats, this color map will assign the colors
boundaries : list of four floats
Boundaries of the plot in format [x1,x2,y1,y2]
branch_components : list of str
Branch components to be plotted, defaults to Line and Link.
jitter : None|float
Amount of random noise to add to bus positions to distinguish
overlapping buses
Returns
-------
bus_collection, branch_collection1, ... : tuple of Collections
Collections for buses and branches.
"""
defaults_for_branches = {
'Link': dict(color="cyan", width=2),
'Line': dict(color="b", width=2),
'Transformer': dict(color='green', width=2)
}
if not plt_present:
logger.error("Matplotlib is not present, so plotting won't work.")
return
if ax is None:
ax = plt.gca()
def compute_bbox_with_margins(margin, x, y):
#set margins
pos = np.asarray((x, y))
minxy, maxxy = pos.min(axis=1), pos.max(axis=1)
xy1 = minxy - margin*(maxxy - minxy)
xy2 = maxxy + margin*(maxxy - minxy)
return tuple(xy1), tuple(xy2)
x = network.buses["x"]
y = network.buses["y"]
if jitter is not None:
x = x + np.random.uniform(low=-jitter, high=jitter, size=len(x))
y = y + np.random.uniform(low=-jitter, high=jitter, size=len(y))
if basemap and basemap_present:
if boundaries is None:
(x1, y1), (x2, y2) = compute_bbox_with_margins(margin, x, y)
else:
x1, x2, y1, y2 = boundaries
bmap = Basemap(resolution='l', epsg=network.srid,
llcrnrlat=y1, urcrnrlat=y2, llcrnrlon=x1,
urcrnrlon=x2, ax=ax)
bmap.drawcountries()
bmap.drawcoastlines()
x, y = bmap(x.values, y.values)
x = pd.Series(x, network.buses.index)
y = pd.Series(y, network.buses.index)
if isinstance(bus_sizes, pd.Series) and isinstance(bus_sizes.index, pd.MultiIndex):
# We are drawing pies to show all the different shares
assert len(network.buses.index.difference(bus_sizes.index.levels[0])) == 0, \
"The first MultiIndex level of bus_sizes must contain buses"
assert isinstance(bus_colors, dict) and set(bus_colors).issuperset(bus_sizes.index.levels[1]), \
"bus_colors must be a dictionary defining a color for each element " \
"in the second MultiIndex level of bus_sizes"
bus_sizes = bus_sizes.sort_index(level=0, sort_remaining=False)
patches = []
for b_i in bus_sizes.index.levels[0]:
s = bus_sizes.loc[b_i]
radius = s.sum()**0.5
ratios = s/s.sum()
start = 0.25
for i, ratio in ratios.iteritems():
patches.append(Wedge((x.at[b_i], y.at[b_i]), radius,
360*start, 360*(start+ratio),
facecolor=bus_colors[i]))
start += ratio
bus_collection = PatchCollection(patches, match_original=True)
ax.add_collection(bus_collection)
else:
c = pd.Series(bus_colors, index=network.buses.index)
if c.dtype == np.dtype('O'):
c.fillna("b", inplace=True)
c = list(c.values)
s = pd.Series(bus_sizes, index=network.buses.index, dtype="float").fillna(10)
bus_collection = ax.scatter(x, y, c=c, s=s, cmap=bus_cmap)
def as_branch_series(ser):
if isinstance(ser, dict) and set(ser).issubset(branch_components):
return pd.Series(ser)
elif isinstance(ser, pd.Series):
if isinstance(ser.index, pd.MultiIndex):
return ser
index = ser.index
ser = ser.values
else:
index = network.lines.index
return pd.Series(ser,
index=pd.MultiIndex(levels=(["Line"], index),
labels=(np.zeros(len(index)),
np.arange(len(index)))))
line_colors = as_branch_series(line_colors)
line_widths = as_branch_series(line_widths)
if not isinstance(line_cmap, dict):
line_cmap = {'Line': line_cmap}
branch_collections = []
for c in network.iterate_components(branch_components):
l_defaults = defaults_for_branches[c.name]
l_widths = line_widths.get(c.name, l_defaults['width'])
l_nums = None
l_colors = line_colors.get(c.name, l_defaults['color'])
if isinstance(l_colors, pd.Series):
if issubclass(l_colors.dtype.type, np.number):
l_nums = l_colors
l_colors = None
else:
l_colors.fillna(l_defaults['color'], inplace=True)
if not geometry:
segments = (np.asarray(((c.df.bus0.map(x),
c.df.bus0.map(y)),
(c.df.bus1.map(x),
c.df.bus1.map(y))))
.transpose(2, 0, 1))
else:
from shapely.wkt import loads
from shapely.geometry import LineString
linestrings = c.df.geometry.map(loads)
assert all(isinstance(ls, LineString) for ls in linestrings), \
"The WKT-encoded geometry in the 'geometry' column must be composed of LineStrings"
segments = np.asarray(list(linestrings.map(np.asarray)))
if basemap and basemap_present:
segments = np.transpose(bmap(*np.transpose(segments, (2, 0, 1))), (1, 2, 0))
l_collection = LineCollection(segments,
linewidths=l_widths,
antialiaseds=(1,),
colors=l_colors,
transOffset=ax.transData)
if l_nums is not None:
l_collection.set_array(np.asarray(l_nums))
l_collection.set_cmap(line_cmap.get(c.name, None))
l_collection.autoscale()
ax.add_collection(l_collection)
l_collection.set_zorder(1)
branch_collections.append(l_collection)
bus_collection.set_zorder(2)
ax.update_datalim(compute_bbox_with_margins(margin, x, y))
ax.autoscale_view()
ax.set_title(title)
return (bus_collection,) + tuple(branch_collections)
def draw_networkx_multi_edges(G,
pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
arrowstyle='-|>',
arrowsize=10,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
node_size=300,
nodelist=None,
node_shape="o",
rad=None,
**kwds):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color string, or array of floats
Edge color. Can be a single color format string (default='r'),
or a sequence of colors with the same length as edgelist.
If numeric values are specified they will be mapped to
colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=1.0)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
Note: Arrows will be the same color as edges.
arrowstyle : str, optional (default='-|>')
For directed graphs, choose the style of the arrow heads.
See :py:class: `matplotlib.patches.ArrowStyle` for more
options.
arrowsize : int, optional (default=10)
For directed graphs, choose the size of the arrow head head's length and
width. See :py:class: `matplotlib.patches.FancyArrowPatch` for attribute
`mutation_scale` for more info.
label : [None| string]
Label for legend
rad: courbure de la liaison
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
list of matplotlib.patches.FancyArrowPatch
`FancyArrowPatch` instances of the directed edges
Depending whether the drawing includes arrows or not.
Notes
-----
For directed graphs, arrows are drawn at the head end. Arrows can be
turned off with keyword arrows=False. Be sure to include `node_size' as a
keyword argument; arrows are drawn considering the size of nodes.
Examples
--------
>>> G = nx.dodecahedral_graph()
>>> edges = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> G = nx.DiGraph()
>>> G.add_edges_from([(1, 2), (1, 3), (2, 3)])
>>> arcs = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> alphas = [0.3, 0.4, 0.5]
>>> for i, arc in enumerate(arcs): # change alpha values of arcs
... arc.set_alpha(alphas[i])
Also see the NetworkX drawing examples at
https://networkx.github.io/documentation/latest/auto_examples/index.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap, Normalize
from matplotlib.collections import LineCollection
from matplotlib.patches2 import FancyArrowPatch
import numpy as np
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
if nodelist is None:
nodelist = list(G.nodes())
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width, )
else:
lw = width
if not is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color) == len(edge_pos):
if np.alltrue([is_string_like(c) for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple(
[colorConverter.to_rgba(c, alpha) for c in edge_color])
elif np.alltrue([not is_string_like(c) for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if np.alltrue(
[cb.iterable(c) and len(c) in (3, 4) for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must contain color names or numbers')
else:
if is_string_like(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
msg = 'edge_color must be a color or list of one color per edge'
raise ValueError(msg)
if (not G.is_directed() or not arrows):
edge_collection = LineCollection(
edge_pos, #ici
colors=edge_colors,
rotation=100,
linewidths=lw,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
# Note: there was a bug in mpl regarding the handling of alpha values
# for each line in a LineCollection. It was fixed in matplotlib by
# r7184 and r7189 (June 6 2009). We should then not set the alpha
# value globally, since the user can instead provide per-edge alphas
# now. Only set it globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert (isinstance(edge_cmap, Colormap))
edge_collection.set_array(np.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
return edge_collection
arrow_collection = None
if G.is_directed() and arrows:
# Note: Waiting for someone to implement arrow to intersection with
# marker. Meanwhile, this works well for polygons with more than 4
# sides and circle.
def to_marker_edge(marker_size, marker):
if marker in "s^>v<d": # `large` markers need extra space
return np.sqrt(2 * marker_size) / 2
else:
return np.sqrt(marker_size) / 2
# Draw arrows with `matplotlib.patches.FancyarrowPatch`
arrow_collection = []
mutation_scale = arrowsize # scale factor of arrow head
arrow_colors = edge_colors
if arrow_colors is None:
if edge_cmap is not None:
assert (isinstance(edge_cmap, Colormap))
else:
edge_cmap = plt.get_cmap() # default matplotlib colormap
if edge_vmin is None:
edge_vmin = min(edge_color)
if edge_vmax is None:
edge_vmax = max(edge_color)
color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
for i, (src, dst) in enumerate(edge_pos):
x1, y1 = src
x2, y2 = dst
arrow_color = None
line_width = None
shrink_source = 0 # space from source to tail
shrink_target = 0 # space from head to target
if cb.iterable(node_size): # many node sizes
src_node, dst_node = edgelist[i]
index_node = nodelist.index(dst_node)
marker_size = node_size[index_node]
shrink_target = to_marker_edge(marker_size, node_shape)
else:
shrink_target = to_marker_edge(node_size, node_shape)
if arrow_colors is None:
arrow_color = edge_cmap(color_normal(edge_color[i]))
elif len(arrow_colors) > 1:
arrow_color = arrow_colors[i]
else:
arrow_color = arrow_colors[0]
if len(lw) > 1:
line_width = lw[i]
else:
line_width = lw[0]
arrow = FancyArrowPatch((x1, y1), (x2, y2),
arrowstyle=arrowstyle,
shrinkA=shrink_source,
shrinkB=shrink_target,
mutation_scale=mutation_scale,
color=arrow_color,
linewidth=line_width,
zorder=1,
rad=rad) # arrows go behind nodes
# There seems to be a bug in matplotlib to make collections of
# FancyArrowPatch instances. Until fixed, the patches are added
# individually to the axes instance.
arrow_collection.append(arrow)
ax.add_patch(arrow)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
return arrow_collection
def odd_import(self, ax, odd):
if odd[1] == -1:
odd[1] = 0
self.odd_tool.update_odd()
file_name = self.odd_tool.odd_list[odd[0]]
odd_path = 'scenarios/' + str(odd[0]) + '/' + file_name[odd[1]] + '.json'
with open(odd_path, 'r') as load_f:
load_dict = load(load_f)
list_dict = list(load_dict.keys())
scenario = load_dict[list_dict[0]]
try:
lines1 = scenario['lines']['boundary_lines']
lines2 = scenario['lines']['solid_lines']
lines3 = scenario['lines']['dashed_lines']
lines4 = scenario['lines']['thick_lines']
objects = scenario['objects']
common_rotation = scenario['common_rotation']
common_offset = scenario['common_offset']
except Exception:
lines1 = []
lines2 = []
lines3 = []
lines4 = []
objects = []
common_rotation = 0.0
common_offset = [0, 0]
QMessageBox.warning(self, 'Warning', 'Scenario file format error!')
boundary_lines = self.line_rotation(lines1, common_rotation)
solid_lines = self.line_rotation(lines2, common_rotation)
dashed_lines = self.line_rotation(lines3, common_rotation)
thick_lines = self.line_rotation(lines4, common_rotation)
for l in boundary_lines:
for m in l:
m[0] = (m[0] + common_offset[0]) * self.zoom_factor
m[1] = (m[1] + common_offset[1]) * self.zoom_factor
for l in solid_lines:
for m in l:
m[0] = (m[0] + common_offset[0]) * self.zoom_factor
m[1] = (m[1] + common_offset[1]) * self.zoom_factor
for l in dashed_lines:
for m in l:
m[0] = (m[0] + common_offset[0]) * self.zoom_factor
m[1] = (m[1] + common_offset[1]) * self.zoom_factor
for l in thick_lines:
for m in l:
m[0] = (m[0] + common_offset[0]) * self.zoom_factor
m[1] = (m[1] + common_offset[1]) * self.zoom_factor
line_segments1 = LineCollection(boundary_lines, linewidths=2, colors=(0, 0, 0), linestyles='solid')
line_segments1.set_zorder(0)
line_segments2 = LineCollection(solid_lines, linewidths=1, colors=(0.1, 0.1, 0.1), linestyles='solid')
line_segments2.set_zorder(0)
line_segments3 = LineCollection(dashed_lines, linewidths=1, colors=(0.2, 0.2, 0.2), linestyles=(0, (20, 20)))
line_segments3.set_zorder(0)
line_segments4 = LineCollection(thick_lines, linewidths=4.5, colors=(0.8, 0.8, 0.8), linestyles='solid')
line_segments4.set_zorder(0)
self.ax.add_collection(line_segments1)
self.ax.add_collection(line_segments2)
self.ax.add_collection(line_segments3)
self.ax.add_collection(line_segments4)
patch = []
for obj in range(len(objects)):
objects_xy = self.object_rotation(objects[obj][1], common_rotation)
objects_xy[0] = objects_xy[0] + common_offset[0]
objects_xy[1] = objects_xy[1] + common_offset[1]
objects_r = objects[obj][2] + common_rotation
patch.append(self.add_object(objects[obj][0], objects_xy, objects_r))
self.ax.text((objects_xy[0]) * self.zoom_factor, (objects_xy[1] + 2.5) * self.zoom_factor,
'obj' + str(obj + 1), fontsize=10)
patch_segment = PatchCollection(patch, facecolors='gray')
patch_segment.set_zorder(2)
ax.add_collection(patch_segment)
return patch
def add_phylorate(treeplot, rates, nodeidx, vis=True):
"""
Add phylorate plot generated from data analyzed with BAMM
(http://bamm-project.org/introduction.html)
Args:
rates (array): Array of rates along branches
created by r_funcs.phylorate
nodeidx (array): Array of node indices matching rates (also created
by r_funcs.phylorate)
WARNING:
Ladderizing the tree can cause incorrect assignment of Ape node index
numbers. To prevent this, call this function or root.ape_node_idx()
before ladderizing the tree to assign correct Ape node index numbers.
"""
if not treeplot.root.apeidx:
treeplot.root.ape_node_idx()
segments = []
values = []
if treeplot.plottype == "radial":
radpatches = [] # For use in drawing arcs for radial plots
for n in treeplot.root.descendants():
n.rates = rates[nodeidx==n.apeidx]
c = treeplot.n2c[n]
pc = treeplot._path_to_parent(n)[0][1]
xd = c.x - pc[0]
yd = c.y - pc[1]
xseg = xd/len(n.rates)
yseg = yd/len(n.rates)
for i, rate in enumerate(n.rates):
x0 = pc[0] + i*xseg
y0 = pc[1] + i*yseg
x1 = x0 + xseg
y1 = y0 + yseg
segments.append(((x0, y0), (x1, y1)))
values.append(rate)
curverts = treeplot._path_to_parent(n)[0][2:]
curcodes = treeplot._path_to_parent(n)[1][2:]
curcol = RdYlBu(n.rates[0])
radpatches.append(PathPatch(
Path(curverts, curcodes), lw=2, edgecolor = curcol,
fill=False))
else:
for n in treeplot.root.descendants():
n.rates = rates[nodeidx==n.apeidx]
c = treeplot.n2c[n]
pc = treeplot.n2c[n.parent]
seglen = (c.x-pc.x)/len(n.rates)
for i, rate in enumerate(n.rates):
x0 = pc.x + i*seglen
x1 = x0 + seglen
segments.append(((x0, c.y), (x1, c.y)))
values.append(rate)
segments.append(((pc.x, pc.y), (pc.x, c.y)))
values.append(n.rates[0])
lc = LineCollection(segments, cmap=RdYlBu, lw=2)
lc.set_array(np.array(values))
treeplot.add_collection(lc)
lc.set_zorder(1)
if treeplot.plottype == "radial":
arccol = matplotlib.collections.PatchCollection(radpatches,
match_original=True)
treeplot.add_collection(arccol)
arccol.set_visible(vis)
arccol.set_zorder(1)
lc.set_visible(vis)
colorbar_legend(treeplot, values, RdYlBu, vis=vis)
treeplot.figure.canvas.draw_idle()
def draw_networkx_edges(G,
pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
ax=None,
**kwds):
"""Draw the edges of the graph G
This draws only the edges of the graph G.
pos is a dictionary keyed by vertex with a two-tuple
of x-y positions as the value.
See networkx.layout for functions that compute node positions.
edgelist is an optional list of the edges in G to be drawn.
If provided, only the edges in edgelist will be drawn.
See draw_networkx for the list of other optional parameters.
"""
from matplotlib.pylab import gca, hold, draw_if_interactive
if ax is None:
ax = gca()
if edgelist is None:
edgelist = G.edges()
if not edgelist: # no edges!
return None
# set edge positions
head = []
tail = []
for e in edgelist:
# edge e can be a 2-tuple (Graph) or a 3-tuple (Xgraph)
u = e[0]
v = e[1]
head.append(pos[u])
tail.append(pos[v])
edge_pos = asarray(zip(head, tail))
if not cb.iterable(width):
lw = (width, )
else:
lw = width
# edge colors specified with floats won't work here
# since LineCollection doesn't use ScalarMappable.
# You can use an array of RGBA or text labels
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color)==len(edge_pos):
edge_colors = None
else:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
if G.is_directed():
edge_collection = LineCollection(
edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
)
edge_collection.set_alpha(alpha)
else:
arrows = []
for src, dst in edge_pos:
x1, y1 = src
x2, y2 = dst
dx = x2 - x1 # x offset
dy = y2 - y1 # y offset
arrows.append(Arrow(x1, y1, dx, dy, width=lw).get_verts())
edge_collection = PolyCollection(
arrows,
facecolors=edge_colors,
antialiaseds=(1, ),
transOffset=ax.transData,
)
# update view
xx = [x for (x, y) in head + tail]
yy = [y for (x, y) in head + tail]
minx = amin(xx)
maxx = amax(xx)
miny = amin(yy)
maxy = amax(yy)
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
edge_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(edge_collection)
return edge_collection
fig = plt.figure()
ax = fig.add_subplot(111)
ax.invert_yaxis()
ax.axis('equal')
for obs in obstacles:
ax.add_patch(plt.Polygon(obs, fc='#404040'))
ax.add_patch(plt.Circle(target_pos, target_rad, ec='None', fc='red'))
# Draw graph
edges = LineCollection(([dataset[v1][:2], dataset[v2][:2]] for v1, v2 in graph.get_edgelist()))
edges.set_color('Lavender')
edges.set_zorder(1)
ax.add_collection(edges)
# Draw points
params = {}
if args.clustering:
# Open the vertex dendogram
print 'Loading clustering...'
membership = cPickle.load(open(args.clustering, 'rb'))
params['c'] = membership
params['cmap'] = plt.get_cmap('hsv')
# Highlight boundary points
boundary = [v.index for v in graph.vs
if sum((membership[nid] != membership[v.index]
for nid in graph.neighbors(v.index))) > 7]
def draw_networkx_edges(G,
pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
**kwds):
"""Draw the edges of the graph G
This draws only the edges of the graph G.
pos is a dictionary keyed by vertex with a two-tuple
of x-y positions as the value.
See networkx.layout for functions that compute node positions.
edgelist is an optional list of the edges in G to be drawn.
If provided, only the edges in edgelist will be drawn.
edgecolor can be a list of matplotlib color letters such as 'k' or
'b' that lists the color of each edge; the list must be ordered in
the same way as the edge list. Alternatively, this list can contain
numbers and those number are mapped to a color scale using the color
map edge_cmap.
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
See draw_networkx for the list of other optional parameters.
"""
if ax is None:
ax = matplotlib.pylab.gca()
if edgelist is None:
edgelist = G.edges()
if not edgelist or len(edgelist) == 0: # no edges!
return None
# set edge positions
edge_pos = asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width, )
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color)==len(edge_pos):
if matplotlib.numerix.alltrue(
[cb.is_string_like(c) for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple(
[colorConverter.to_rgba(c, alpha) for c in edge_color])
elif matplotlib.numerix.alltrue(
[not cb.is_string_like(c) for c in edge_color]):
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError(
'edge_color must consist of either color names or numbers')
else:
if len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError(
'edge_color must be a single color or list of exactly m colors where m is the number or edges'
)
edge_collection = LineCollection(
edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
)
edge_collection.set_alpha(alpha)
# need 0.87.7 or greater for edge colormaps
mpl_version = matplotlib.__version__
if mpl_version.endswith('svn'):
mpl_version = matplotlib.__version__[0:-3]
if mpl_version.endswith('pre'):
mpl_version = matplotlib.__version__[0:-3]
if map(int, mpl_version.split('.')) >= [0, 87, 7]:
if edge_colors is None:
if edge_cmap is not None: assert (isinstance(edge_cmap, Colormap))
edge_collection.set_array(asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
matplotlib.pylab.sci(edge_collection)
# else:
# sys.stderr.write(\
# """matplotlib version >= 0.87.7 required for colormapped edges.
# (version %s detected)."""%matplotlib.__version__)
# raise UserWarning(\
# """matplotlib version >= 0.87.7 required for colormapped edges.
# (version %s detected)."""%matplotlib.__version__)
arrow_collection = None
if G.directed and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = (colorConverter.to_rgba('k', alpha), )
a_pos = []
p = 1.0 - 0.25 # make head segment 25 percent of edge length
for src, dst in edge_pos:
x1, y1 = src
x2, y2 = dst
dx = x2 - x1 # x offset
dy = y2 - y1 # y offset
d = sqrt(float(dx**2 + dy**2)) # length of edge
if d == 0: # source and target at same position
continue
if dx == 0: # vertical edge
xa = x2
ya = dy * p + y1
if dy == 0: # horizontal edge
ya = y2
xa = dx * p + x1
else:
theta = arctan2(dy, dx)
xa = p * d * cos(theta) + x1
ya = p * d * sin(theta) + y1
a_pos.append(((xa, ya), (x2, y2)))
arrow_collection = LineCollection(
a_pos,
colors=arrow_colors,
linewidths=[4 * ww for ww in lw],
antialiaseds=(1, ),
transOffset=ax.transData,
)
# update view
minx = amin(ravel(edge_pos[:, :, 0]))
maxx = amax(ravel(edge_pos[:, :, 0]))
miny = amin(ravel(edge_pos[:, :, 1]))
maxy = amax(ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
edge_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(edge_collection)
if arrow_collection:
arrow_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(arrow_collection)
return edge_collection
def draw_networkx_edges(
G,
pos,
edgelist=None,
width=1.0,
edge_color="k",
style="solid",
alpha=None,
arrowstyle="-|>",
arrowsize=3,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
node_size=300,
nodelist=None,
node_shape="o",
connectionstyle=None,
min_source_margin=0,
min_target_margin=0,
):
"""Draw the edges of the graph G. Adjusted from networkx.
"""
try:
import matplotlib.pyplot as plt
from matplotlib.colors import colorConverter, Colormap, Normalize
from matplotlib.collections import LineCollection
from matplotlib.patches import FancyArrowPatch
from numbers import Number
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
if nodelist is None:
nodelist = list(G.nodes())
# FancyArrowPatch handles color=None different from LineCollection
if edge_color is None:
edge_color = "k"
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
# Check if edge_color is an array of floats and map to edge_cmap.
# This is the only case handled differently from matplotlib
if (np.iterable(edge_color) and (len(edge_color) == len(edge_pos))
and np.alltrue([isinstance(c, Number) for c in edge_color])):
if edge_cmap is not None:
assert isinstance(edge_cmap, Colormap)
else:
edge_cmap = plt.get_cmap()
if edge_vmin is None:
edge_vmin = min(edge_color)
if edge_vmax is None:
edge_vmax = max(edge_color)
color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
edge_color = [edge_cmap(color_normal(e)) for e in edge_color]
if not G.is_directed() or not arrows:
edge_collection = LineCollection(
edge_pos,
colors=edge_color,
linewidths=width,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
alpha=alpha,
)
edge_collection.set_cmap(edge_cmap)
edge_collection.set_clim(edge_vmin, edge_vmax)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
return edge_collection
arrow_collection = None
if G.is_directed() and arrows:
# Note: Waiting for someone to implement arrow to intersection with
# marker. Meanwhile, this works well for polygons with more than 4
# sides and circle.
def to_marker_edge(marker_size, marker):
if marker in "s^>v<d": # `large` markers need extra space
return np.sqrt(2 * marker_size) / 2
else:
return np.sqrt(marker_size) / 2
# Draw arrows with `matplotlib.patches.FancyarrowPatch`
arrow_collection = []
mutation_scale = arrowsize # scale factor of arrow head
# FancyArrowPatch doesn't handle color strings
arrow_colors = colorConverter.to_rgba_array(edge_color, alpha)
for i, (src, dst) in enumerate(edge_pos):
x1, y1 = src
x2, y2 = dst
shrink_source = 0 # space from source to tail
shrink_target = 0 # space from head to target
if np.iterable(node_size): # many node sizes
source, target = edgelist[i][:2]
source_node_size = node_size[nodelist.index(source)]
target_node_size = node_size[nodelist.index(target)]
shrink_source = to_marker_edge(source_node_size, node_shape)
shrink_target = to_marker_edge(target_node_size, node_shape)
else:
shrink_source = shrink_target = to_marker_edge(
node_size, node_shape)
if shrink_source < min_source_margin:
shrink_source = min_source_margin
if shrink_target < min_target_margin:
shrink_target = min_target_margin
if len(arrow_colors) == len(edge_pos):
arrow_color = arrow_colors[i]
elif len(arrow_colors) == 1:
arrow_color = arrow_colors[0]
else: # Cycle through colors
arrow_color = arrow_colors[i % len(arrow_colors)]
if np.iterable(width):
if len(width) == len(edge_pos):
line_width = width[i]
else:
line_width = width[i % len(width)]
else:
line_width = width
arrow = FancyArrowPatch(
(x1, y1),
(x2, y2),
arrowstyle=arrowstyle,
shrinkA=shrink_source,
shrinkB=shrink_target,
mutation_scale=mutation_scale,
color=arrow_color,
linewidth=line_width,
connectionstyle=connectionstyle,
linestyle=style,
zorder=1,
) # arrows go behind nodes
# There seems to be a bug in matplotlib to make collections of
# FancyArrowPatch instances. Until fixed, the patches are added
# individually to the axes instance.
arrow_collection.append(arrow)
ax.add_patch(arrow)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
ax.tick_params(
axis="both",
which="both",
bottom=False,
left=False,
labelbottom=False,
labelleft=False,
)
return arrow_collection
class GraphDrawing:
"""
This class implements all the drawing related methods for a GraphLocal
instance. These methods include changing colors, highlighting a set etc.
It is not designed to be used individually. Its purpose is to change all
kinds of drawing properties after calling standard drawing functions,
"draw" and "draw_groups" in GraphLocal.
Attributes
----------
G : GraphLocal instance
coords : a n-by-2 or n-by-3 array with coordinates for each node of the graph.
ax,fig : None,None (default)
by default it will create a new figure, or this will plot in axs if not None.
is_3d : True when it is a 3D graph
nodes_collection : a matplotlib PathCollection instance containing all nodes
edge_collection : a matplotlib LineCollection instance containing all edges
groups : list[list] or list, for the first case, each sublist represents a cluster
for the second case, list must have the same length as the number of nodes and
nodes with the number are in the same cluster
"""
def __init__(self, G, coords, ax=None, groups=None, figsize=None):
self.G = G
self.coords = coords
self.is_3d = (len(coords[0]) == 3)
self.edge_pos, self.edge_mapping = self._plotting_build_edgepos(
G, coords)
if ax is None:
fig = plt.figure(figsize=figsize)
if len(coords[0]) == 3:
ax = fig.add_subplot(111, projection='3d')
else:
ax = fig.add_subplot(111)
else:
fig = ax.get_figure()
ax.set_axis_off()
self.fig = fig
self.ax = ax
self.nodes_collection = None
self.edge_collection = None
self.groups = groups
@staticmethod
def _plotting_push_edges_for_node(center, points, pos, edge_pos,
edge_mapping):
for i, p in enumerate(points):
if p >= center:
edge_mapping[(center, p)] = len(edge_pos)
edge_pos.append([pos[center], pos[p]])
@staticmethod
def _plotting_build_edgepos(G, pos):
edge_pos = []
edge_mapping = {}
for i in range(G._num_vertices):
GraphDrawing._plotting_push_edges_for_node(
i, G.aj[G.ai[i]:G.ai[i + 1]], pos, edge_pos, edge_mapping)
edge_pos = np.asarray(edge_pos)
return edge_pos, edge_mapping
def show(self):
"""
show the graph
"""
return self.fig
def highlight(self,
nodelist,
othernodes=False,
otheredges=False,
circled=True,
alpha=0.1):
"""
highlight a set of nodes
Parameters
----------
nodelist: a list of nodes to be highlighted
Optional parameters
------------------
othernodes: bool (False by default)
whether to hide nodes that is not in the nodelist
otheredges: bool (False by default)
whether to hide edges that doesn't connect two nodes in the nodelist
circled: bool (False by default)
set to True to circle nodes in the nodelist
alpha: float (1.0 by default)
change alpha for nodes that are not in the nodelist
"""
nodeset = set(nodelist)
if not othernodes or circled:
node_out = list(set(range(self.G._num_vertices)) - nodeset)
if not othernodes:
self.nodecolor(node_out, alpha=alpha)
if circled:
self.nodecolor(nodelist, facecolor='r', edgecolor='b', alpha=1)
curr_size = self.nodes_collection.get_sizes()[0]
self.nodesize(nodelist, [(curr_size * 1.5)**2] * len(nodelist))
curr_width = self.nodes_collection.get_linewidths()[0]
self.nodewidth(nodelist,
[(curr_width * 1.5)**2] * len(nodelist))
#self.only_circle_nodes(nodelist)
if not otheredges:
for (i, j) in self.edge_mapping.keys():
if i not in nodeset or j not in nodeset:
self.edgecolor(i, j, alpha=alpha)
def only_circle_nodes(self, nodeset):
"""
only circle the nodes in nodeset
"""
facecolors = self.nodes_collection.get_facecolor()
facecolors[nodeset] = [0, 0, 0, 0]
def between_group_alpha(self, alpha):
"""
change the edge alpha value for edges that connect nodes from different groups
"""
if self.groups is not None:
if self.groups.ndim == 2:
node_mapping = np.zeros(self.G._num_vertices,
dtype=self.G.aj.dtype)
for idx, grp in enumerate(self.groups):
node_mapping[grp] = idx
else:
node_mapping = self.groups
for edge in self.edge_mapping.keys():
if node_mapping[i] != node_mapping[j]:
self.edgecolor(edge[0], edge[1], alpha=alpha)
def nodecolor(self,
node,
c=None,
edgecolor=None,
facecolor=None,
alpha=None):
"""
change node color
Parameters
----------
node: integer or list[integer]
c: string or rgb or rgba (None by default)
when set to be None, this function just returns the current colors for the node
edgecolor,facecolor: (None by default)
used when you want different edgecolor and facecolor for the node
when set to be None, it will be same as "c"
alpha: float (None by default)
when set to be None, alpha will not be changed
Returns
-------
list of two lists, where the first is new face color and the second is new edge color, if face color
is not changed, the first is None, if edge color is not changed, the second is None
"""
if c is not None:
edgecolor = c
facecolor = c
ret_facecolor, ret_edgecolor = None, None
if facecolor is not None or alpha is not None:
colors = self.nodes_collection.get_facecolor()
# This means right now, all nodes have the same facecolor
if colors.shape[0] == 1:
# Firstly, we need to expand the color array so that every node has an independant facecolor
self.nodes_collection.set_facecolor(
[colors[0] for i in range(self.G._num_vertices)])
colors = self.nodes_collection.get_facecolor()
ret_facecolor = self._plotting_update_color(
colors, node, facecolor, alpha)
if edgecolor is not None or alpha is not None:
colors = self.nodes_collection.get_edgecolor()
if colors.shape[0] <= 1:
# This means right now, all nodes have hidden edges
if colors.shape[0] == 0:
# Use facecolor as edgecolor
colors = self.nodes_collection.get_facecolor()
# This means right now, all nodes have the same edgecolor
if colors.shape[0] == 1:
# Firstly, we need to expand the color array so that every node has an independant edgecolor
colors = [colors[0] for i in range(self.G._num_vertices)]
self.nodes_collection.set_edgecolor(colors)
colors = self.nodes_collection.get_edgecolor()
ret_edgecolor = self._plotting_update_color(
colors, node, edgecolor, alpha)
return [ret_facecolor, ret_edgecolor]
# The better way here might be diectly modifying self.edge_collection._paths
def edgecolor(self, i, j, c=None, alpha=None):
"""
change edge color
Parameters
----------
i,j: integer, start and end node of the edge
c: string or rgb or rgba (None by default)
when set to be None, this function just returns the current colors for the edge
alpha: float (None by default)
when set to be None, alpha will not be changed
Returns
-------
current edge color
"""
colors = self.edge_collection.get_edgecolor()
if len(colors) == 1:
colors = np.array([colors[0]] * self.G._num_edges)
self.edge_collection.set_edgecolor(c=colors)
idx = self.edge_mapping[(i, j)]
return self._plotting_update_color(colors, idx, c, alpha)
def nodesize(self, node, nodesize):
"""
change node size
Parameters
----------
node: integer or list[integer]
nodesize: float, int, list[int] or list[float]
in the latter two cases, the length of nodesize must
be the same as the length of node
Returns
-------
current node size
"""
sizes = self.nodes_collection.get_sizes()
if len(sizes) == 1:
sizes = np.array([sizes[0]] * self.G._num_vertices)
if isinstance(nodesize, float) or isinstance(nodesize, int):
sizes[node] = nodesize
else:
sizes[node] = np.reshape(nodesize, len(nodesize))
self.nodes_collection.set_sizes(sizes)
return sizes[node]
def nodewidth(self, node, width):
"""
change line width of node
Parameters
----------
node: integer or list[integer]
width: float, int, list[int] or list[float]
in the latter two cases, the length of nodesize must
be the same as the length of node
"""
widths = np.asarray(self.nodes_collection.get_linewidths())
if len(widths) == 1:
widths = np.array([widths[0]] * self.G._num_vertices)
if isinstance(width, float) or isinstance(width, int):
widths[node] = width
else:
widths[node] = np.reshape(width, len(width))
self.nodes_collection.set_linewidths(widths)
return widths[node]
@staticmethod
def _plotting_update_color(container, key, c, alpha):
if c is not None:
if c == "none":
# only circle the nodes
container[key] = c
else:
if alpha is not None:
c = to_rgba(c, alpha)
container[key] = c
else:
c = to_rgb(c)
container[key, 0:3] = c
else:
if alpha is not None:
container[key, 3] = alpha
return container[key]
def scatter(self, **kwargs):
"""
a wrapper of standard matplotlib scatter function
Parameters
----------
**kwargs: same as the parameters in matplotlib scatter
"""
coords = self.coords
if len(self.coords[0]) == 2:
self.nodes_collection = self.ax.scatter([p[0] for p in coords],
[p[1] for p in coords],
**kwargs)
else:
self.nodes_collection = self.ax.scatter([p[0] for p in coords],
[p[1] for p in coords],
[p[2] for p in coords],
**kwargs)
self.ax.add_collection(self.nodes_collection)
self.ax._sci(self.nodes_collection)
def plot(self, **kwargs):
"""
a wrapper of standard matplotlib plot function
Parameters
----------
**kwargs: same as the parameters in matplotlib plot
"""
if len(self.coords[0]) == 2:
self.edge_collection = LineCollection(self.edge_pos, **kwargs)
else:
self.edge_collection = Line3DCollection(self.edge_pos, **kwargs)
#make sure edges are at the bottom
self.edge_collection.set_zorder(1)
self.ax.add_collection(self.edge_collection)
self.ax._sci(self.edge_collection)
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=None,
arrowstyle='-|>',
arrowsize=10,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
node_size=300,
nodelist=None,
node_shape="o",
connectionstyle=None,
**kwds):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color or array of colors (default='k')
Edge color. Can be a single color or a sequence of colors with the same
length as edgelist. Color can be string, or rgb (or rgba) tuple of
floats from 0-1. If numeric values are specified they will be
mapped to colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=None)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
Note: Arrows will be the same color as edges.
arrowstyle : str, optional (default='-|>')
For directed graphs, choose the style of the arrow heads.
See :py:class: `matplotlib.patches.ArrowStyle` for more
options.
arrowsize : int, optional (default=10)
For directed graphs, choose the size of the arrow head head's length and
width. See :py:class: `matplotlib.patches.FancyArrowPatch` for attribute
`mutation_scale` for more info.
connectionstyle : str, optional (default=None)
Pass the connectionstyle parameter to create curved arc of rounding
radius rad. For example, connectionstyle='arc3,rad=0.2'.
See :py:class: `matplotlib.patches.ConnectionStyle` and
:py:class: `matplotlib.patches.FancyArrowPatch` for more info.
label : [None| string]
Label for legend
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
list of matplotlib.patches.FancyArrowPatch
`FancyArrowPatch` instances of the directed edges
Depending whether the drawing includes arrows or not.
Notes
-----
For directed graphs, arrows are drawn at the head end. Arrows can be
turned off with keyword arrows=False. Be sure to include `node_size` as a
keyword argument; arrows are drawn considering the size of nodes.
Examples
--------
>>> G = nx.dodecahedral_graph()
>>> edges = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> G = nx.DiGraph()
>>> G.add_edges_from([(1, 2), (1, 3), (2, 3)])
>>> arcs = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> alphas = [0.3, 0.4, 0.5]
>>> for i, arc in enumerate(arcs): # change alpha values of arcs
... arc.set_alpha(alphas[i])
Also see the NetworkX drawing examples at
https://networkx.github.io/documentation/latest/auto_examples/index.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap, Normalize
from matplotlib.collections import LineCollection
from matplotlib.patches import FancyArrowPatch
import numpy as np
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
if nodelist is None:
nodelist = list(G.nodes())
# FancyArrowPatch handles color=None different from LineCollection
if edge_color is None:
edge_color = 'k'
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
# Check if edge_color is an array of floats and map to edge_cmap.
# This is the only case handled differently from matplotlib
if cb.iterable(edge_color) and (len(edge_color) == len(edge_pos)) \
and np.alltrue([isinstance(c,Number) for c in edge_color]):
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
else:
edge_cmap = plt.get_cmap()
if edge_vmin is None:
edge_vmin = min(edge_color)
if edge_vmax is None:
edge_vmax = max(edge_color)
color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
edge_color = [edge_cmap(color_normal(e)) for e in edge_color]
if (not G.is_directed() or not arrows):
edge_collection = LineCollection(edge_pos,
colors=edge_color,
linewidths=width,
antialiaseds=(1,),
linestyle=style,
transOffset=ax.transData,
alpha=alpha
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
return edge_collection
arrow_collection = None
if G.is_directed() and arrows:
# Note: Waiting for someone to implement arrow to intersection with
# marker. Meanwhile, this works well for polygons with more than 4
# sides and circle.
def to_marker_edge(marker_size, marker):
if marker in "s^>v<d": # `large` markers need extra space
return np.sqrt(2 * marker_size) / 2
else:
return np.sqrt(marker_size) / 2
# Draw arrows with `matplotlib.patches.FancyarrowPatch`
arrow_collection = []
mutation_scale = arrowsize # scale factor of arrow head
# FancyArrowPatch doesn't handle color strings
arrow_colors = colorConverter.to_rgba_array(edge_color,alpha)
for i, (src, dst) in enumerate(edge_pos):
x1, y1 = src
x2, y2 = dst
shrink_source = 0 # space from source to tail
shrink_target = 0 # space from head to target
if cb.iterable(node_size): # many node sizes
src_node, dst_node = edgelist[i][:2]
index_node = nodelist.index(dst_node)
marker_size = node_size[index_node]
shrink_target = to_marker_edge(marker_size, node_shape)
else:
shrink_target = to_marker_edge(node_size, node_shape)
if cb.iterable(arrow_colors):
if len(arrow_colors) == len(edge_pos):
arrow_color = arrow_colors[i]
elif len(arrow_colors)==1:
arrow_color = arrow_colors[0]
else: # Cycle through colors
arrow_color = arrow_colors[i%len(arrow_colors)]
else:
arrow_color = edge_color
if cb.iterable(width):
if len(width) == len(edge_pos):
line_width = width[i]
else:
line_width = width[i%len(width)]
else:
line_width = width
arrow = FancyArrowPatch((x1, y1), (x2, y2),
arrowstyle=arrowstyle,
shrinkA=shrink_source,
shrinkB=shrink_target,
mutation_scale=mutation_scale,
color=arrow_color,
linewidth=line_width,
connectionstyle=connectionstyle,
zorder=1) # arrows go behind nodes
# There seems to be a bug in matplotlib to make collections of
# FancyArrowPatch instances. Until fixed, the patches are added
# individually to the axes instance.
arrow_collection.append(arrow)
ax.add_patch(arrow)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
plt.tick_params(
axis='both',
which='both',
bottom=False,
left=False,
labelbottom=False,
labelleft=False)
return arrow_collection
def ani_frame(pos, t1, t2):
fig, ax = plt.subplots()
fig.patch.set_facecolor('green')
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
# Push axis slightly off teh edge of the screen to remove any border
ax.set_position([-0.01, -0.01, 1.02, 1.02])
ax.set_axis_bgcolor('green')
ax.xaxis.label.set_color('green')
ax.yaxis.label.set_color('green')
# Make it nice and square, the final resolution is this times the DPI
fig.set_size_inches([3,3])
east = pos['East'].squeeze()
north = pos['North'].squeeze()
time = pos['time']
# Plot the whole path as a faint color
bg_idx = np.where((time >= t1) & (time <= t2))
col = (0.2, 0.2, 0.2)
background1, = ax.plot(east[bg_idx], north[bg_idx], color=col, linewidth=3)
col = (0.8, 0.8, 0.8)
background2, = ax.plot(east[bg_idx], north[bg_idx], color=col)
# Get the samples that are in this time window
def getIdx(t):
idx = np.where((time > (t - lag)) & (time <= t))
return idx
# Create the plot of the recent data
idx = getIdx(t1)
col = (0.1, 0.1, 0.8)
x = east[idx]
y = north[idx]
t = (t1 - time[idx]) / lag
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = LineCollection(segments, array=t, cmap=plt.get_cmap('copper'), norm=plt.Normalize(0,1), lw=2)
lc.set_array(t)
lc.set_linewidth(3)
lc.set_zorder(20) # make sure trail is on top
ax.add_collection(lc)
# Mask out the path
def init():
lc.set_segments([])
lc.set_array([])
lc.set_linewidth(0)
return lc,
# Plot a segment of the path
def update_img(t):
idx = getIdx(t)
x = east[idx]
y = north[idx]
col = (double) (t - time[idx]) / lag
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc.set_segments(segments)
lc.set_array(col)
lc.set_linewidth(3)
ax.add_collection(lc)
ax.autoscale()
plt.draw()
return lc,
ani = animation.FuncAnimation(fig,update_img,np.arange(t1,t2,step=1000/fps),interval=0,blit=False)
writer = animation.writers['ffmpeg'](fps=30)
ani.save('demo.mp4',dpi=dpi,fps=30,writer=writer)
return lc
def construct_haldane_eigvect_DOS_plot(xy,
fig,
DOS_ax,
eig_ax,
eigval,
eigvect,
en,
NL,
KL,
marker_num=0,
color_scheme='default',
sub_lattice=-1,
normalization=None):
"""puts together lattice and DOS plots and draws normal mode ellipsoids on top
Parameters
----------
xy: array 2N x 3
Equilibrium position of the gyroscopes
fig :
figure with lattice and DOS drawn
DOS_ax:
axis for the DOS plot
eig_ax
axis for the eigenvalue plot
eigval : array of dimension 2nx1
Eigenvalues of matrix for system
eigvect : array of dimension 2nx2n
Eigenvectors of matrix for system.
Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like:
x0, y0, x1, y1, ... xNP, yNP.
en: int
Number of the eigenvalue you are plotting
Returns
----------
fig :
completed figure for normal mode
[scat_fg, p, f_mark] :
things to be cleared before next normal mode is drawn
"""
s = leplt.absolute_sizer()
plt.sca(DOS_ax)
ev = eigval[en]
ev1 = ev
# Show where current eigenvalue is in DOS plot
(f_mark, ) = plt.plot([ev, ev], plt.ylim(), '-r')
NP = len(xy)
im1 = np.imag(ev)
re1 = np.real(ev)
plt.sca(eig_ax)
plt.title('Mode %d; $\Omega=( %0.6f + %0.6f i)$' % (en, re1, im1))
# Preallocate ellipsoid plot vars
angles_arr = np.zeros(NP)
patch = []
colors = np.zeros(NP + 2)
x0s = np.zeros(NP)
y0s = np.zeros(NP)
mag1 = eigvect[en]
if normalization is None:
mag1 /= np.max(np.abs(mag1))
else:
mag1 *= normalization * float(len(xy))
# Pick a series of times to draw out the ellipsoid
time_arr = np.arange(81.0) * 2. * np.pi / float(abs(ev1) * 80)
exp1 = np.exp(1j * ev1 * time_arr)
cw = []
ccw = []
lines_1 = []
for i in range(NP):
x_disps = 0.5 * (exp1 * mag1[i]).real
y_disps = 0.5 * (exp1 * mag1[i]).imag
x_vals = xy[i, 0] + x_disps
y_vals = xy[i, 1] + y_disps
poly_points = np.array([x_vals, y_vals]).T
polygon = Polygon(poly_points, True)
# x0 is the marker_num^th element of x_disps
x0 = x_disps[marker_num]
y0 = y_disps[marker_num]
x0s[i] = x_vals[marker_num]
y0s[i] = y_vals[marker_num]
# These are the black lines protruding from pivot point to current position
lines_1.append([[xy[i, 0], x_vals[marker_num]],
[xy[i, 1], y_vals[marker_num]]])
mag = np.sqrt(x0**2 + y0**2)
if mag > 0:
anglez = np.arccos(x0 / mag)
else:
anglez = 0
if y0 < 0:
anglez = 2 * np.pi - anglez
angles_arr[i] = anglez
patch.append(polygon)
if color_scheme == 'default':
colors[i] = anglez
else:
if sub_lattice[i] == 0:
colors[i] = 0
else:
colors[i] = np.pi
ccw.append(i)
colors[NP] = 0
colors[NP + 1] = 2 * np.pi
plt.yticks([])
plt.xticks([])
# this is the part that puts a dot a t=0 point
scat_fg = eig_ax.scatter(x0s[cw], y0s[cw], s=s(.02), c='DodgerBlue')
scat_fg2 = eig_ax.scatter(x0s[ccw], y0s[ccw], s=s(.02), c='Red', zorder=3)
NP = len(xy)
try:
NN = np.shape(NL)[1]
except IndexError:
NN = 0
z = np.zeros(NP)
Rnorm = np.array([x0s, y0s, z]).T
# Bond Stretches
inc = 0
stretches = np.zeros(4 * len(xy))
for i in range(len(xy)):
if NN > 0:
for j, k in zip(NL[i], KL[i]):
if i < j and abs(k) > 0:
n1 = float(linalg.norm(Rnorm[i] - Rnorm[j]))
n2 = linalg.norm(xy[i] - xy[j])
stretches[inc] = (n1 - n2)
inc += 1
# For particles with neighbors, get list of bonds to draw by stretches
test = list(np.zeros([inc, 1]))
inc = 0
xy = np.array([x0s, y0s, z]).T
for i in range(len(xy)):
if NN > 0:
for j, k in zip(NL[i], KL[i]):
if i < j and abs(k) > 0:
test[inc] = [xy[(i, j), 0], xy[(i, j), 1]]
inc += 1
stretch = np.array(stretches[0:inc])
# lines connect sites (bonds), while lines_12 draw the black lines from the pinning to location sites
lines = [zip(x, y) for x, y in test]
lines_12 = [zip(x, y) for x, y in lines_1]
lines_st = LineCollection(lines,
array=stretch,
cmap='seismic',
linewidth=8)
lines_st.set_clim([-1. * 0.25, 1 * 0.25])
lines_st.set_zorder(2)
lines_12_st = LineCollection(lines_12, linewidth=0.8)
lines_12_st.set_color('k')
p = PatchCollection(patch, cmap='hsv', alpha=0.6)
p.set_array(np.array(colors))
p.set_clim([0, 2 * np.pi])
p.set_zorder(1)
# eig_ax.add_collection(lines_st)
eig_ax.add_collection(lines_12_st)
eig_ax.add_collection(p)
eig_ax.set_aspect('equal')
# erased ev/(2*pi) here npm 2016
cw_ccw = [cw, ccw, ev]
# print cw_ccw[1]
return fig, [scat_fg, scat_fg2, p, f_mark, lines_12_st], cw_ccw
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
arrowstyle='-|>',
arrowsize=10,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
node_size=300,
nodelist=None,
node_shape="o",
**kwds):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color string, or array of floats
Edge color. Can be a single color format string (default='r'),
or a sequence of colors with the same length as edgelist.
If numeric values are specified they will be mapped to
colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=1.0)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
Note: Arrows will be the same color as edges.
arrowstyle : str, optional (default='-|>')
For directed graphs, choose the style of the arrow heads.
See :py:class: `matplotlib.patches.ArrowStyle` for more
options.
arrowsize : int, optional (default=10)
For directed graphs, choose the size of the arrow head head's length and
width. See :py:class: `matplotlib.patches.FancyArrowPatch` for attribute
`mutation_scale` for more info.
label : [None| string]
Label for legend
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
list of matplotlib.patches.FancyArrowPatch
`FancyArrowPatch` instances of the directed edges
Depending whether the drawing includes arrows or not.
Notes
-----
For directed graphs, arrows are drawn at the head end. Arrows can be
turned off with keyword arrows=False. Be sure to include `node_size' as a
keyword argument; arrows are drawn considering the size of nodes.
Examples
--------
>>> G = nx.dodecahedral_graph()
>>> edges = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> G = nx.DiGraph()
>>> G.add_edges_from([(1, 2), (1, 3), (2, 3)])
>>> arcs = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> alphas = [0.3, 0.4, 0.5]
>>> for i, arc in enumerate(arcs): # change alpha values of arcs
... arc.set_alpha(alphas[i])
Also see the NetworkX drawing examples at
https://networkx.github.io/documentation/latest/auto_examples/index.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap, Normalize
from matplotlib.collections import LineCollection
from matplotlib.patches import FancyArrowPatch
import numpy as np
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
if nodelist is None:
nodelist = list(G.nodes())
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color) == len(edge_pos):
if np.alltrue([is_string_like(c) for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c, alpha)
for c in edge_color])
elif np.alltrue([not is_string_like(c) for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if np.alltrue([cb.iterable(c) and len(c) in (3, 4)
for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must contain color names or numbers')
else:
if is_string_like(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
msg = 'edge_color must be a color or list of one color per edge'
raise ValueError(msg)
if (not G.is_directed() or not arrows):
edge_collection = LineCollection(edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1,),
linestyle=style,
transOffset=ax.transData,
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
# Note: there was a bug in mpl regarding the handling of alpha values
# for each line in a LineCollection. It was fixed in matplotlib by
# r7184 and r7189 (June 6 2009). We should then not set the alpha
# value globally, since the user can instead provide per-edge alphas
# now. Only set it globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(np.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
return edge_collection
arrow_collection = None
if G.is_directed() and arrows:
# Note: Waiting for someone to implement arrow to intersection with
# marker. Meanwhile, this works well for polygons with more than 4
# sides and circle.
def to_marker_edge(marker_size, marker):
if marker in "s^>v<d": # `large` markers need extra space
return np.sqrt(2 * marker_size) / 2
else:
return np.sqrt(marker_size) / 2
# Draw arrows with `matplotlib.patches.FancyarrowPatch`
arrow_collection = []
mutation_scale = arrowsize # scale factor of arrow head
arrow_colors = edge_colors
if arrow_colors is None:
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
else:
edge_cmap = plt.get_cmap() # default matplotlib colormap
if edge_vmin is None:
edge_vmin = min(edge_color)
if edge_vmax is None:
edge_vmax = max(edge_color)
color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
for i, (src, dst) in enumerate(edge_pos):
x1, y1 = src
x2, y2 = dst
arrow_color = None
line_width = None
shrink_source = 0 # space from source to tail
shrink_target = 0 # space from head to target
if cb.iterable(node_size): # many node sizes
src_node, dst_node = edgelist[i]
index_node = nodelist.index(dst_node)
marker_size = node_size[index_node]
shrink_target = to_marker_edge(marker_size, node_shape)
else:
shrink_target = to_marker_edge(node_size, node_shape)
if arrow_colors is None:
arrow_color = edge_cmap(color_normal(edge_color[i]))
elif len(arrow_colors) > 1:
arrow_color = arrow_colors[i]
else:
arrow_color = arrow_colors[0]
if len(lw) > 1:
line_width = lw[i]
else:
line_width = lw[0]
arrow = FancyArrowPatch((x1, y1), (x2, y2),
arrowstyle=arrowstyle,
shrinkA=shrink_source,
shrinkB=shrink_target,
mutation_scale=mutation_scale,
color=arrow_color,
linewidth=line_width,
zorder=1) # arrows go behind nodes
# There seems to be a bug in matplotlib to make collections of
# FancyArrowPatch instances. Until fixed, the patches are added
# individually to the axes instance.
arrow_collection.append(arrow)
ax.add_patch(arrow)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
return arrow_collection
def plot(network, margin=0.05, ax=None, geomap=True, projection=None,
bus_colors='b', line_colors='g', bus_sizes=10, line_widths=2,
title="", line_cmap=None, bus_cmap=None, boundaries=None,
geometry=False, branch_components=['Line', 'Link'], jitter=None,
basemap=None, basemap_parameters=None, color_geomap=None):
"""
Plot the network buses and lines using matplotlib and Basemap.
Parameters
----------
margin : float
Margin at the sides as proportion of distance between max/min x,y
ax : matplotlib ax, defaults to plt.gca()
Axis to which to plot the network
geomap: bool/str, default True
Switch to use Basemap or Cartopy (depends on what is installed).
If string is passed, it will be used as a resolution argument.
For Basemap users 'c' (crude), 'l' (low), 'i' (intermediate),
'h' (high), 'f' (full) are valid resolutions options.
For Cartopy users '10m', '50m', '110m' are valid resolutions options.
projection: cartopy.crs.Projection, defaults to None
Define the projection of your geomap, only valid if cartopy is
installed. If None (default) is passed the projection for cartopy
is set to cartopy.crs.PlateCarree
bus_colors : dict/pandas.Series
Colors for the buses, defaults to "b"
bus_sizes : dict/pandas.Series
Sizes of bus points, defaults to 10
line_colors : dict/pandas.Series
Colors for the lines, defaults to "g" for Lines and "cyan" for
Links. Colors for branches other than Lines can be
specified using a pandas Series with a MultiIndex.
line_widths : dict/pandas.Series
Widths of lines, defaults to 2. Widths for branches other
than Lines can be specified using a pandas Series with a
MultiIndex.
title : string
Graph title
line_cmap : plt.cm.ColorMap/str|dict
If line_colors are floats, this color map will assign the colors.
Use a dict to specify colormaps for more than one branch type.
bus_cmap : plt.cm.ColorMap/str
If bus_colors are floats, this color map will assign the colors
boundaries : list of four floats
Boundaries of the plot in format [x1,x2,y1,y2]
branch_components : list of str
Branch components to be plotted, defaults to Line and Link.
jitter : None|float
Amount of random noise to add to bus positions to distinguish
overlapping buses
basemap_parameters : dict
Specify a dict with additional constructor parameters for the
Basemap. Will disable Cartopy.
Use this feature to set a custom projection.
(e.g. `{'projection': 'tmerc', 'lon_0':10.0, 'lat_0':50.0}`)
color_geomap : dict or bool
Specify colors to paint land and sea areas in.
If True, it defaults to `{'ocean': 'lightblue', 'land': 'whitesmoke'}`.
If no dictionary is provided, colors are white.
Returns
-------
bus_collection, branch_collection1, ... : tuple of Collections
Collections for buses and branches.
"""
defaults_for_branches = {
'Link': dict(color="cyan", width=2),
'Line': dict(color="b", width=2),
'Transformer': dict(color='green', width=2)
}
if not plt_present:
logger.error("Matplotlib is not present, so plotting won't work.")
return
if basemap is not None:
logger.warning("argument `basemap` is deprecated, "
"use `geomap` instead.")
geomap = basemap
if geomap:
if not (cartopy_present or basemap_present):
# Not suggesting Basemap since it is being deprecated
logger.warning("Cartopy needs to be installed to use `geomap=True`.")
geomap = False
# Use cartopy by default, fall back on basemap
use_basemap = False
use_cartopy = cartopy_present
if not use_cartopy:
use_basemap = basemap_present
# If the user specifies basemap parameters, they prefer
# basemap over cartopy.
# (This means that you can force the use of basemap by
# setting `basemap_parameters={}`)
if basemap_present:
if basemap_parameters is not None:
logger.warning("Basemap is being deprecated, consider "
"switching to Cartopy.")
use_basemap = True
use_cartopy = False
if use_cartopy:
if projection is None:
projection = get_projection_from_crs(network.srid)
if ax is None:
ax = plt.gca(projection=projection)
else:
assert isinstance(ax, cartopy.mpl.geoaxes.GeoAxesSubplot), (
'The passed axis is not a GeoAxesSubplot. You can '
'create one with: \nimport cartopy.crs as ccrs \n'
'fig, ax = plt.subplots('
'subplot_kw={"projection":ccrs.PlateCarree()})')
elif ax is None:
ax = plt.gca()
x, y = network.buses["x"], network.buses["y"]
axis_transform = ax.transData
if geomap:
if use_cartopy:
axis_transform = draw_map_cartopy(network, x, y, ax,
boundaries, margin, geomap, color_geomap)
new_coords = pd.DataFrame(
ax.projection.transform_points(axis_transform,
x.values, y.values),
index=network.buses.index, columns=['x', 'y', 'z'])
x, y = new_coords['x'], new_coords['y']
elif use_basemap:
basemap_transform = draw_map_basemap(network, x, y, ax,
boundaries, margin, geomap, basemap_parameters, color_geomap)
# A non-standard projection might be used; the easiest way to
# support this is to tranform the bus coordinates.
x, y = basemap_transform(x.values, y.values)
x = pd.Series(x, network.buses.index)
y = pd.Series(y, network.buses.index)
if jitter is not None:
x = x + np.random.uniform(low=-jitter, high=jitter, size=len(x))
y = y + np.random.uniform(low=-jitter, high=jitter, size=len(y))
if isinstance(bus_sizes, pd.Series) and isinstance(bus_sizes.index, pd.MultiIndex):
# We are drawing pies to show all the different shares
assert len(bus_sizes.index.levels[0].difference(network.buses.index)) == 0, \
"The first MultiIndex level of bus_sizes must contain buses"
assert (isinstance(bus_colors, dict) and
set(bus_colors).issuperset(bus_sizes.index.levels[1])), \
"bus_colors must be a dictionary defining a color for each element " \
"in the second MultiIndex level of bus_sizes"
bus_sizes = bus_sizes.sort_index(level=0, sort_remaining=False)\
* projected_area_factor(ax, network.srid)**2
patches = []
for b_i in bus_sizes.index.levels[0]:
s = bus_sizes.loc[b_i]
radius = s.sum()**0.5
if radius == 0.0:
ratios = s
else:
ratios = s/s.sum()
start = 0.25
for i, ratio in ratios.iteritems():
patches.append(Wedge((x.at[b_i], y.at[b_i]), radius,
360*start, 360*(start+ratio),
facecolor=bus_colors[i]))
start += ratio
bus_collection = PatchCollection(patches, match_original=True)
ax.add_collection(bus_collection)
else:
c = pd.Series(bus_colors, index=network.buses.index)
s = pd.Series(bus_sizes, index=network.buses.index, dtype="float").fillna(10)
bus_collection = ax.scatter(x, y, c=c, s=s, cmap=bus_cmap, edgecolor='face')
def as_branch_series(ser):
if isinstance(ser, dict) and set(ser).issubset(branch_components):
return pd.Series(ser)
elif isinstance(ser, pd.Series):
if isinstance(ser.index, pd.MultiIndex):
return ser
index = ser.index
ser = ser.values
else:
index = network.lines.index
return pd.Series(ser,
index=pd.MultiIndex(levels=(["Line"], index),
codes=(np.zeros(len(index)),
np.arange(len(index)))))
line_colors = as_branch_series(line_colors)
line_widths = as_branch_series(line_widths)
if not isinstance(line_cmap, dict):
line_cmap = {'Line': line_cmap}
branch_collections = []
for c in network.iterate_components(branch_components):
l_defaults = defaults_for_branches[c.name]
l_widths = line_widths.get(c.name, l_defaults['width'])
l_nums = None
l_colors = line_colors.get(c.name, l_defaults['color'])
if isinstance(l_colors, pd.Series):
if issubclass(l_colors.dtype.type, np.number):
l_nums = l_colors
l_colors = None
else:
l_colors.fillna(l_defaults['color'], inplace=True)
if not geometry:
segments = (np.asarray(((c.df.bus0.map(x),
c.df.bus0.map(y)),
(c.df.bus1.map(x),
c.df.bus1.map(y))))
.transpose(2, 0, 1))
else:
from shapely.wkt import loads
from shapely.geometry import LineString
linestrings = c.df.geometry[lambda ds: ds != ''].map(loads)
assert all(isinstance(ls, LineString) for ls in linestrings), (
"The WKT-encoded geometry in the 'geometry' column must be "
"composed of LineStrings")
segments = np.asarray(list(linestrings.map(np.asarray)))
l_collection = LineCollection(segments,
linewidths=l_widths,
antialiaseds=(1,),
colors=l_colors,
transOffset=ax.transData)
if l_nums is not None:
l_collection.set_array(np.asarray(l_nums))
l_collection.set_cmap(line_cmap.get(c.name, None))
l_collection.autoscale()
ax.add_collection(l_collection)
l_collection.set_zorder(3)
branch_collections.append(l_collection)
bus_collection.set_zorder(4)
ax.update_datalim(compute_bbox_with_margins(margin, x, y))
ax.autoscale_view()
if geomap:
if use_cartopy:
ax.outline_patch.set_visible(False)
ax.axis('off')
ax.set_title(title)
return (bus_collection,) + tuple(branch_collections)
self.collection.set_facecolors(self.fc)
self.canvas.draw_idle()
if __name__ == '__main__':
G = nx.read_gpickle(GRAPH_PATH)
graph_data = nx.get_node_attributes(G, 'pos')
data = np.array(tuple(graph_data.values()))
labels = np.zeros(len(data))
instances = np.zeros(len(data))
edgelist = list(G.edges())
edge_pos = np.asarray([(graph_data[e[0]], graph_data[e[1]])
for e in edgelist])
edge_collection = LineCollection(edge_pos)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(None)
fig, ax = plt.subplots()
pts = ax.scatter(data[:, 0], data[:, 1], s=80, c='b')
ax.add_collection(edge_collection)
selector = SelectFromCollection(ax, pts)
fc = pts.get_facecolors()
mode = 'SELECT'
instance = 1
fig.suptitle("Mode: " + mode + " | Instance: " + str(instance),
x=0.5,
y=0.05)
ann_list = []
def draw_networkx_edges(G,
pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
**kwds):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color string, or array of floats
Edge color. Can be a single color format string (default='r'),
or a sequence of colors with the same length as edgelist.
If numeric values are specified they will be mapped to
colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=1.0)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
label : [None| string]
Label for legend
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
Notes
-----
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
Yes, it is ugly but drawing proper arrows with Matplotlib this
way is tricky.
Examples
--------
>>> G = nx.dodecahedral_graph()
>>> edges = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
Also see the NetworkX drawing examples at
https://networkx.github.io/documentation/latest/auto_examples/index.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap
from matplotlib.collections import LineCollection
import numpy as np
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width, )
else:
lw = width
if not is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color) == len(edge_pos):
if np.alltrue([is_string_like(c) for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple(
[colorConverter.to_rgba(c, alpha) for c in edge_color])
elif np.alltrue([not is_string_like(c) for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if np.alltrue(
[cb.iterable(c) and len(c) in (3, 4) for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError(
'edge_color must consist of either color names or numbers')
else:
if is_string_like(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError(
'edge_color must be a single color or list of exactly m colors where m is the number or edges'
)
edge_collection = LineCollection(
edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
# Note: there was a bug in mpl regarding the handling of alpha values for
# each line in a LineCollection. It was fixed in matplotlib in r7184 and
# r7189 (June 6 2009). We should then not set the alpha value globally,
# since the user can instead provide per-edge alphas now. Only set it
# globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert (isinstance(edge_cmap, Colormap))
edge_collection.set_array(np.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
arrow_collection = None
if G.is_directed() and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = edge_colors
a_pos = []
p = 1.0 - 0.25 # make head segment 25 percent of edge length
for src, dst in edge_pos:
x1, y1 = src
x2, y2 = dst
dx = x2 - x1 # x offset
dy = y2 - y1 # y offset
d = np.sqrt(float(dx**2 + dy**2)) # length of edge
if d == 0: # source and target at same position
continue
if dx == 0: # vertical edge
xa = x2
ya = dy * p + y1
if dy == 0: # horizontal edge
ya = y2
xa = dx * p + x1
else:
theta = np.arctan2(dy, dx)
xa = p * d * np.cos(theta) + x1
ya = p * d * np.sin(theta) + y1
a_pos.append(((xa, ya), (x2, y2)))
arrow_collection = LineCollection(
a_pos,
colors=arrow_colors,
linewidths=[4 * ww for ww in lw],
antialiaseds=(1, ),
transOffset=ax.transData,
)
arrow_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(arrow_collection)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
# if arrow_collection:
return edge_collection
def draw_edges(G,
segments,
pos=None,
edgelist=None,
width=1.0,
color='k',
style='solid',
alpha=None,
ax=None,
**kwds):
"""Draw the edges of the graph G.
This draws the edge segments given by a separation of the links in
`data` of the graph G.
Parameters
----------
G : graph
A networkx graph
segments : L x M array
The segmentation of each link. (segments.sum(axis=1) == 1).all()
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
(default=nx.get_node_attributes(G, 'pos'))
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float or array of floats
Line width of edges (default =1.0)
color : tuple of color strings
Edge Segments color. Can be a single color format string (default='r'),
or a sequence of colors with the same length as data.shape[1].
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=1.0)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edge segments
"""
if not np.allclose(segments.sum(axis=1), 1):
segments = segments / segments.sum(axis=1, keepdims=True)
if ax is None:
ax = plt.gca()
if pos is None:
pos = nx.get_node_attributes(G, 'pos')
if edgelist is None:
edgelist = G.edges()
if not edgelist or len(edgelist) == 0: # no edges!
return None
if not cb.iterable(width):
lw = (width, )
else:
lw = width
if cb.iterable(color) \
and len(color) == segments.shape[1]:
if np.alltrue([_is_string_like(c) for c in color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple(
[colorConverter.to_rgba(c, alpha) for c in color])
elif (np.alltrue([not _is_string_like(c) for c in color])
and np.alltrue(
[cb.iterable(c) and len(c) in (3, 4) for c in color])):
edge_colors = tuple(color)
else:
raise ValueError(
'color must consist of either color names or numbers')
else:
if _is_string_like(color) or len(color) == 1:
edge_colors = (colorConverter.to_rgba(color, alpha), )
else:
raise ValueError(
'color must be a single color or list of exactly m colors where m is the number of segments'
)
assert len(edgelist) == segments.shape[
0], "Number edges and segments have to line up"
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
src = edge_pos[:, 0]
dest = edge_pos[:, 1]
positions = src[:, np.newaxis] + np.cumsum(
np.hstack((np.zeros((len(segments), 1)), segments)),
axis=1)[:, :, np.newaxis] * (dest - src)[:, np.newaxis]
linecolls = []
for s in range(segments.shape[1]):
coll = LineCollection(positions[:, s:s + 2],
colors=edge_colors[s:s + 1],
linewidths=lw,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData)
coll.set_zorder(1) # edges go behind nodes
# coll.set_label(label)
if cb.is_numlike(alpha):
coll.set_alpha(alpha)
ax.add_collection(coll)
linecolls.append(coll)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
return linecolls
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
**kwds):
"""Draw the edges of the graph G
This draws only the edges of the graph G.
pos is a dictionary keyed by vertex with a two-tuple
of x-y positions as the value.
See networkx_v099.layout for functions that compute node positions.
edgelist is an optional list of the edges in G to be drawn.
If provided, only the edges in edgelist will be drawn.
edgecolor can be a list of matplotlib color letters such as 'k' or
'b' that lists the color of each edge; the list must be ordered in
the same way as the edge list. Alternatively, this list can contain
numbers and those number are mapped to a color scale using the color
map edge_cmap.
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
See draw_networkx_v099 for the list of other optional parameters.
"""
if ax is None:
ax=matplotlib.pylab.gca()
if edgelist is None:
edgelist=G.edges()
if not edgelist or len(edgelist)==0: # no edges!
return None
# set edge positions
edge_pos=asarray([(pos[e[0]],pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color)==len(edge_pos):
if matplotlib.numerix.alltrue([cb.is_string_like(c)
for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c,alpha)
for c in edge_color])
elif matplotlib.numerix.alltrue([not cb.is_string_like(c)
for c in edge_color]):
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must consist of either color names or numbers')
else:
if len(edge_color)==1:
edge_colors = ( colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError('edge_color must be a single color or list of exactly m colors where m is the number or edges')
edge_collection = LineCollection(edge_pos,
colors = edge_colors,
linewidths = lw,
antialiaseds = (1,),
linestyle = style,
transOffset = ax.transData,
)
edge_collection.set_alpha(alpha)
# need 0.87.7 or greater for edge colormaps
mpl_version=matplotlib.__version__
if mpl_version.endswith('svn'):
mpl_version=matplotlib.__version__[0:-3]
if mpl_version.endswith('pre'):
mpl_version=matplotlib.__version__[0:-3]
if map(int,mpl_version.split('.'))>=[0,87,7]:
if edge_colors is None:
if edge_cmap is not None: assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
matplotlib.pylab.sci(edge_collection)
# else:
# sys.stderr.write(\
# """matplotlib version >= 0.87.7 required for colormapped edges.
# (version %s detected)."""%matplotlib.__version__)
# raise UserWarning(\
# """matplotlib version >= 0.87.7 required for colormapped edges.
# (version %s detected)."""%matplotlib.__version__)
arrow_collection=None
if G.directed and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = ( colorConverter.to_rgba('k', alpha), )
a_pos=[]
p=1.0-0.25 # make head segment 25 percent of edge length
for src,dst in edge_pos:
x1,y1=src
x2,y2=dst
dx=x2-x1 # x offset
dy=y2-y1 # y offset
d=sqrt(float(dx**2+dy**2)) # length of edge
if d==0: # source and target at same position
continue
if dx==0: # vertical edge
xa=x2
ya=dy*p+y1
if dy==0: # horizontal edge
ya=y2
xa=dx*p+x1
else:
theta=arctan2(dy,dx)
xa=p*d*cos(theta)+x1
ya=p*d*sin(theta)+y1
a_pos.append(((xa,ya),(x2,y2)))
arrow_collection = LineCollection(a_pos,
colors = arrow_colors,
linewidths = [4*ww for ww in lw],
antialiaseds = (1,),
transOffset = ax.transData,
)
# update view
minx = amin(ravel(edge_pos[:,:,0]))
maxx = amax(ravel(edge_pos[:,:,0]))
miny = amin(ravel(edge_pos[:,:,1]))
maxy = amax(ravel(edge_pos[:,:,1]))
w = maxx-minx
h = maxy-miny
padx, pady = 0.05*w, 0.05*h
corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady)
ax.update_datalim( corners)
ax.autoscale_view()
edge_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(edge_collection)
if arrow_collection:
arrow_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(arrow_collection)
return edge_collection
def readshapefile(shapefile, name, is_web_merc=False, drawbounds=True, zorder=None,
linewidth=0.5, linestyle=(0, ()), color='k', antialiased=1, ax=None,
default_encoding='utf-8'):
"""
Read in shape file, optionally draw boundaries on map.
.. note::
- Assumes shapes are 2D
- only works for Point, MultiPoint, Polyline and Polygon shapes.
- vertices/points must be in geographic (lat/lon) coordinates.
Mandatory Arguments:
.. tabularcolumns:: |l|L|
============== ====================================================
Argument Description
============== ====================================================
shapefile path to shapefile components. Example:
shapefile='/home/jeff/esri/world_borders' assumes
that world_borders.shp, world_borders.shx and
world_borders.dbf live in /home/jeff/esri.
name name for Basemap attribute to hold the shapefile
vertices or points in map projection
coordinates. Class attribute name+'_info' is a list
of dictionaries, one for each shape, containing
attributes of each shape from dbf file, For
example, if name='counties', self.counties
will be a list of x,y vertices for each shape in
map projection coordinates and self.counties_info
will be a list of dictionaries with shape
attributes. Rings in individual Polygon
shapes are split out into separate polygons, and
additional keys 'RINGNUM' and 'SHAPENUM' are added
to the shape attribute dictionary.
============== ====================================================
The following optional keyword arguments are only relevant for Polyline
and Polygon shape types, for Point and MultiPoint shapes they are
ignored.
.. tabularcolumns:: |l|L|
============== ====================================================
Keyword Description
============== ====================================================
drawbounds draw boundaries of shapes (default True).
zorder shape boundary zorder (if not specified,
default for mathplotlib.lines.LineCollection
is used).
linewidth shape boundary line width (default 0.5)
color shape boundary line color (default black)
antialiased antialiasing switch for shape boundaries
(default True).
ax axes instance (overrides default axes instance)
============== ====================================================
A tuple (num_shapes, type, min, max) containing shape file info
is returned.
num_shapes is the number of shapes, type is the type code (one of
the SHPT* constants defined in the shapelib module, see
http://shapelib.maptools.org/shp_api.html) and min and
max are 4-element lists with the minimum and maximum values of the
vertices. If ``drawbounds=True`` a
matplotlib.patches.LineCollection object is appended to the tuple.
"""
import shapefile as shp
from shapefile import Reader
shp.default_encoding = default_encoding
if not os.path.exists('%s.shp' % shapefile):
raise IOError('cannot locate %s.shp' % shapefile)
if not os.path.exists('%s.shx' % shapefile):
raise IOError('cannot locate %s.shx' % shapefile)
if not os.path.exists('%s.dbf' % shapefile):
raise IOError('cannot locate %s.dbf' % shapefile)
# open shapefile, read vertices for each object, convert
# to map projection coordinates (only works for 2D shape types).
try:
shf = Reader(shapefile)
except:
raise IOError('error reading shapefile %s.shp' % shapefile)
fields = shf.fields
coords = []
attributes = []
msg = dedent("""
shapefile must have lat/lon vertices - it looks like this one has vertices
in map projection coordinates. You can convert the shapefile to geographic
coordinates using the shpproj utility from the shapelib tools
(http://shapelib.maptools.org/shapelib-tools.html)""")
shapes = shf.shapes()
if len(shapes) == 0:
raise IndexError('%s shapes is null' % shapefile)
shptype = shapes[0].shapeType
bbox = shf.bbox.tolist()
info = dict()
info['info'] = (shf.numRecords, shptype, bbox[0:2] + [0., 0.], bbox[2:] + [0., 0.])
npoly = 0
for shprec in shf.shapeRecords():
shp = shprec.shape
rec = shprec.record
npoly = npoly + 1
if shptype != shp.shapeType:
raise ValueError('readshapefile can only handle a single shape type per file')
if shptype not in [1, 3, 5, 8]:
raise ValueError('readshapefile can only handle 2D shape types')
verts = shp.points
if shptype in [1, 8]: # a Point or MultiPoint shape.
lons, lats = list(zip(*verts))
if max(lons) > 721. or min(lons) < -721. or max(lats) > 90.01 or min(lats) < -90.01:
raise ValueError(msg)
# if latitude is slightly greater than 90, truncate to 90
lats = [max(min(lat, 90.0), -90.0) for lat in lats]
if len(verts) > 1: # MultiPoint
if is_web_merc:
# x, y = Projection.ToMerc(lons, lats)
pass
else:
x, y = lons, lats
coords.append(list(zip(x, y)))
else: # single Point
if is_web_merc:
pass
# x, y = Projection.ToMerc(lons[0], lats[0])
else:
x, y = lons[0], lats[0]
coords.append((x, y))
attdict = {}
for r, key in zip(rec, fields[1:]):
attdict[key[0]] = r
attributes.append(attdict)
else: # a Polyline or Polygon shape.
parts = shp.parts.tolist()
ringnum = 0
for indx1, indx2 in zip(parts, parts[1:] + [len(verts)]):
ringnum = ringnum + 1
lons, lats = list(zip(*verts[indx1:indx2]))
if max(lons) > 721. or min(lons) < -721. or max(lats) > 90.01 or min(lats) < -90.01:
raise ValueError(msg)
# if latitude is slightly greater than 90, truncate to 90
lats = [max(min(lat, 90.0), -90.0) for lat in lats]
if is_web_merc:
# x, y = Projection.ToMerc(lons, lats)
pass
else:
x, y = lons, lats
coords.append(list(zip(x, y)))
attdict = {}
for r, key in zip(rec, fields[1:]):
attdict[key[0]] = r
# add information about ring number to dictionary.
attdict['RINGNUM'] = ringnum
attdict['SHAPENUM'] = npoly
attributes.append(attdict)
# draw shape boundaries for polylines, polygons using LineCollection.
if shptype not in [1, 8] and drawbounds:
# get current axes instance (if none specified).
import matplotlib.pyplot as plt
ax = ax or plt.gca()
# make LineCollections for each polygon.
lines = LineCollection(coords, antialiaseds=(antialiased,))
lines.set_color(color)
lines.set_linewidth(linewidth)
lines.set_linestyle(linestyle)
lines.set_label('_nolabel_')
if zorder is not None:
lines.set_zorder(zorder)
ax.add_collection(lines)
# set axes limits to fit map region.
# self.set_axes_limits(ax=ax)
# # clip boundaries to map limbs
# lines,c = self._cliplimb(ax,lines)
# info = info + (lines,)
info['lines'] = lines
info[name] = coords
info[name + '_info'] = attributes
return info
def add_phylorate(treeplot, rates, nodeidx, vis=True):
"""
Add phylorate plot generated from data analyzed with BAMM
(http://bamm-project.org/introduction.html)
Args:
rates (array): Array of rates along branches
created by r_funcs.phylorate
nodeidx (array): Array of node indices matching rates (also created
by r_funcs.phylorate)
WARNING:
Ladderizing the tree can cause incorrect assignment of Ape node index
numbers. To prevent this, call this function or root.ape_node_idx()
before ladderizing the tree to assign correct Ape node index numbers.
"""
if not treeplot.root.apeidx:
treeplot.root.ape_node_idx()
segments = []
values = []
if treeplot.plottype == "radial":
radpatches = [] # For use in drawing arcs for radial plots
for n in treeplot.root.descendants():
n.rates = rates[nodeidx==n.apeidx]
c = treeplot.n2c[n]
pc = treeplot._path_to_parent(n)[0][1]
xd = c.x - pc[0]
yd = c.y - pc[1]
xseg = xd/len(n.rates)
yseg = yd/len(n.rates)
for i, rate in enumerate(n.rates):
x0 = pc[0] + i*xseg
y0 = pc[1] + i*yseg
x1 = x0 + xseg
y1 = y0 + yseg
segments.append(((x0, y0), (x1, y1)))
values.append(rate)
curverts = treeplot._path_to_parent(n)[0][2:]
curcodes = treeplot._path_to_parent(n)[1][2:]
curcol = RdYlBu(n.rates[0])
radpatches.append(PathPatch(
Path(curverts, curcodes), lw=2, edgecolor = curcol,
fill=False))
else:
for n in treeplot.root.descendants():
n.rates = rates[nodeidx==n.apeidx]
c = treeplot.n2c[n]
pc = treeplot.n2c[n.parent]
seglen = (c.x-pc.x)/len(n.rates)
for i, rate in enumerate(n.rates):
x0 = pc.x + i*seglen
x1 = x0 + seglen
segments.append(((x0, c.y), (x1, c.y)))
values.append(rate)
segments.append(((pc.x, pc.y), (pc.x, c.y)))
values.append(n.rates[0])
lc = LineCollection(segments, cmap=RdYlBu, lw=2)
lc.set_array(np.array(values))
treeplot.add_collection(lc)
lc.set_zorder(1)
if treeplot.plottype == "radial":
arccol = matplotlib.collections.PatchCollection(radpatches,
match_original=True)
treeplot.add_collection(arccol)
arccol.set_visible(vis)
arccol.set_zorder(1)
lc.set_visible(vis)
colorbar_legend(treeplot, values, RdYlBu, vis=vis)
treeplot.figure.canvas.draw_idle()
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
node_size=300, # Added by qrancik
**kwds):
"""Draw the edges of the graph G
This draws only the edges of the graph G.
pos is a dictionary keyed by vertex with a two-tuple
of x-y positions as the value.
See networkx.layout for functions that compute node positions.
edgelist is an optional list of the edges in G to be drawn.
If provided, only the edges in edgelist will be drawn.
edgecolor can be a list of matplotlib color letters such as 'k' or
'b' that lists the color of each edge; the list must be ordered in
the same way as the edge list. Alternatively, this list can contain
numbers and those number are mapped to a color scale using the color
map edge_cmap.
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
See draw_networkx for the list of other optional parameters.
"""
if ax is None:
ax=matplotlib.pylab.gca()
if edgelist is None:
edgelist=G.edges()
if not edgelist or len(edgelist)==0: # no edges!
return None
# set edge positions
edge_pos=asarray([(pos[e[0]],pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color)==len(edge_pos):
if matplotlib.numerix.alltrue([cb.is_string_like(c)
for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c,alpha)
for c in edge_color])
elif matplotlib.numerix.alltrue([not cb.is_string_like(c)
for c in edge_color]):
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must consist of either color names or numbers')
else:
if len(edge_color)==1:
edge_colors = ( colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError('edge_color must be a single color or list of exactly m colors where m is the number or edges')
# edge_collection = LineCollection(edge_pos, # Deleted by qrancik
edge_collection = ArrowedLineCollection(edge_pos, # Added by qrancik
colors = edge_colors,
linewidths = lw,
antialiaseds = (1,),
linestyle = style,
transOffset = ax.transData,
)
edge_collection.set_alpha(alpha)
# Added by Satyam
arrow_collection=None
if G.is_directed() and arrows:
edge_collection.add_arrows(math.sqrt(node_size)/1.5) # Found emprically to work well. Maybe more systematic approach needed?
arrow_colors = ( colorConverter.to_rgba('k', alpha), )
arrow_collection = LineCollection(edge_pos, # Added by qrancik
colors = arrow_colors,
linewidths = lw,
antialiaseds = (1,),
linestyle = style,
transOffset = ax.transData,
)
# need 0.87.7 or greater for edge colormaps
mpl_version=matplotlib.__version__
if mpl_version.endswith('svn'):
mpl_version=matplotlib.__version__[0:-3]
if mpl_version.endswith('pre'):
mpl_version=matplotlib.__version__[0:-3]
if map(int,mpl_version.split('.'))>=[0,87,7]:
if edge_colors is None:
if edge_cmap is not None: assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
matplotlib.pylab.sci(edge_collection)
# update view
minx = amin(ravel(edge_pos[:,:,0]))
maxx = amax(ravel(edge_pos[:,:,0]))
miny = amin(ravel(edge_pos[:,:,1]))
maxy = amax(ravel(edge_pos[:,:,1]))
w = maxx-minx
h = maxy-miny
padx, pady = 0.05*w, 0.05*h
corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady)
ax.update_datalim( corners)
ax.autoscale_view()
edge_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(edge_collection)
if arrow_collection:
arrow_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(arrow_collection)
return edge_collection
def plot(network, margin=0.05, ax=None, basemap=True, bus_colors='b',
line_colors='g', bus_sizes=10, line_widths=2, title="",
line_cmap=None, bus_cmap=None, boundaries=None,
geometry=False, branch_types=['Line', 'TransportLink', 'Link']):
"""
Plot the network buses and lines using matplotlib and Basemap.
Parameters
----------
margin : float
Margin at the sides as proportion of distance between max/min x,y
ax : matplotlib ax, defaults to plt.gca()
Axis to which to plot the network
basemap : bool, default True
Switch to use Basemap
bus_colors : dict/pandas.Series
Colors for the buses, defaults to "b"
bus_sizes : dict/pandas.Series
Sizes of bus points, defaults to 10
line_colors : dict/pandas.Series
Colors for the lines, defaults to "g" for Lines and "cyan" for
TransportLinks and Links. Colors for branches other than Lines can be
specified using a pandas Series with a MultiIndex.
line_widths : dict/pandas.Series
Widths of lines, defaults to 2. Widths for branches other
than Lines can be specified using a pandas Series with a
MultiIndex.
title : string
Graph title
line_cmap : plt.cm.ColorMap/str|dict
If line_colors are floats, this color map will assign the colors.
Use a dict to specify colormaps for more than one branch type.
bus_cmap : plt.cm.ColorMap/str
If bus_colors are floats, this color map will assign the colors
boundaries : list of four floats
Boundaries of the plot in format [x1,x2,y1,y2]
branch_types : list of str or pypsa.component
Branch types to be plotted, defaults to Line, TransportLink and Link.
Returns
-------
bus_collection, branch_collection1, ... : tuple of Collections
Collections for buses and branches.
"""
defaults_for_branches = {
'TransportLink': dict(color="cyan", width=2),
'Link': dict(color="cyan", width=2),
'Line': dict(color="b", width=2)
}
if not plt_present:
print("Matplotlib is not present, so plotting won't work.")
return
if ax is None:
ax = plt.gca()
def compute_bbox_with_margins(margin, x, y):
#set margins
pos = np.asarray((x, y))
minxy, maxxy = pos.min(axis=1), pos.max(axis=1)
xy1 = minxy - margin*(maxxy - minxy)
xy2 = maxxy + margin*(maxxy - minxy)
return tuple(xy1), tuple(xy2)
x = network.buses["x"]
y = network.buses["y"]
if basemap and basemap_present:
if boundaries is None:
(x1, y1), (x2, y2) = compute_bbox_with_margins(margin, x, y)
else:
x1, x2, y1, y2 = boundaries
bmap = Basemap(resolution='l', epsg=network.srid,
llcrnrlat=y1, urcrnrlat=y2, llcrnrlon=x1,
urcrnrlon=x2, ax=ax)
bmap.drawcountries()
bmap.drawcoastlines()
x, y = bmap(x.values, y.values)
x = pd.Series(x, network.buses.index)
y = pd.Series(y, network.buses.index)
c = pd.Series(bus_colors, index=network.buses.index)
if c.dtype == np.dtype('O'):
c.fillna("b", inplace=True)
s = pd.Series(bus_sizes, index=network.buses.index, dtype="float").fillna(10)
bus_collection = ax.scatter(x, y, c=c, s=s, cmap=bus_cmap)
def as_branch_series(ser):
if isinstance(ser, pd.Series):
if isinstance(ser.index, pd.MultiIndex):
return ser
index = ser.index
ser = ser.values
else:
index = network.lines.index
return pd.Series(ser,
index=pd.MultiIndex(levels=(["Line"], index),
labels=(np.zeros(len(index)),
np.arange(len(index)))))
line_colors = as_branch_series(line_colors)
line_widths = as_branch_series(line_widths)
if not isinstance(line_cmap, dict):
line_cmap = {'Line': line_cmap}
branch_collections = []
for t in network.iterate_components(branch_types):
l_defaults = defaults_for_branches[t.name]
l_widths = line_widths.get(t.name, l_defaults['width'])
l_nums = None
if t.name in line_colors:
l_colors = line_colors[t.name]
if issubclass(l_colors.dtype.type, np.number):
l_nums = l_colors
l_colors = None
else:
l_colors.fillna(l_defaults['color'], inplace=True)
else:
l_colors = l_defaults['color']
if not geometry:
segments = (np.asarray(((t.df.bus0.map(x),
t.df.bus0.map(y)),
(t.df.bus1.map(x),
t.df.bus1.map(y))))
.transpose(2, 0, 1))
else:
from shapely.wkt import loads
from shapely.geometry import LineString
linestrings = t.df.geometry.map(loads)
assert all(isinstance(ls, LineString) for ls in linestrings), \
"The WKT-encoded geometry in the 'geometry' column must be composed of LineStrings"
segments = np.asarray(list(linestrings.map(np.asarray)))
if basemap and basemap_present:
segments = np.transpose(bmap(*np.transpose(segments, (2, 0, 1))), (1, 2, 0))
l_collection = LineCollection(segments,
linewidths=l_widths,
antialiaseds=(1,),
colors=l_colors,
transOffset=ax.transData)
if l_nums is not None:
l_collection.set_array(np.asarray(l_nums))
l_collection.set_cmap(line_cmap.get(t.name, None))
l_collection.autoscale()
ax.add_collection(l_collection)
l_collection.set_zorder(1)
branch_collections.append(l_collection)
bus_collection.set_zorder(2)
ax.update_datalim(compute_bbox_with_margins(margin, x, y))
ax.autoscale_view()
ax.set_title(title)
return (bus_collection,) + tuple(branch_collections)
def plot(network,
margin=0.05,
ax=None,
basemap=True,
bus_colors='b',
line_colors='g',
bus_sizes=10,
line_widths=2,
title="",
line_cmap=None,
bus_cmap=None,
boundaries=None,
geometry=False,
branch_components=['Line', 'Link'],
jitter=None):
"""
Plot the network buses and lines using matplotlib and Basemap.
Parameters
----------
margin : float
Margin at the sides as proportion of distance between max/min x,y
ax : matplotlib ax, defaults to plt.gca()
Axis to which to plot the network
basemap : bool, default True
Switch to use Basemap
bus_colors : dict/pandas.Series
Colors for the buses, defaults to "b"
bus_sizes : dict/pandas.Series
Sizes of bus points, defaults to 10
line_colors : dict/pandas.Series
Colors for the lines, defaults to "g" for Lines and "cyan" for
Links. Colors for branches other than Lines can be
specified using a pandas Series with a MultiIndex.
line_widths : dict/pandas.Series
Widths of lines, defaults to 2. Widths for branches other
than Lines can be specified using a pandas Series with a
MultiIndex.
title : string
Graph title
line_cmap : plt.cm.ColorMap/str|dict
If line_colors are floats, this color map will assign the colors.
Use a dict to specify colormaps for more than one branch type.
bus_cmap : plt.cm.ColorMap/str
If bus_colors are floats, this color map will assign the colors
boundaries : list of four floats
Boundaries of the plot in format [x1,x2,y1,y2]
branch_components : list of str
Branch components to be plotted, defaults to Line and Link.
jitter : None|float
Amount of random noise to add to bus positions to distinguish
overlapping buses
Returns
-------
bus_collection, branch_collection1, ... : tuple of Collections
Collections for buses and branches.
"""
defaults_for_branches = {
'Link': dict(color="cyan", width=2),
'Line': dict(color="b", width=2),
'Transformer': dict(color='green', width=2)
}
if not plt_present:
logger.error("Matplotlib is not present, so plotting won't work.")
return
if ax is None:
ax = plt.gca()
def compute_bbox_with_margins(margin, x, y):
#set margins
pos = np.asarray((x, y))
minxy, maxxy = pos.min(axis=1), pos.max(axis=1)
xy1 = minxy - margin * (maxxy - minxy)
xy2 = maxxy + margin * (maxxy - minxy)
return tuple(xy1), tuple(xy2)
x = network.buses["x"]
y = network.buses["y"]
if jitter is not None:
x = x + np.random.uniform(low=-jitter, high=jitter, size=len(x))
y = y + np.random.uniform(low=-jitter, high=jitter, size=len(y))
if basemap and basemap_present:
resolution = 'l' if isinstance(basemap, bool) else basemap
if boundaries is None:
(x1, y1), (x2, y2) = compute_bbox_with_margins(margin, x, y)
else:
x1, x2, y1, y2 = boundaries
bmap = Basemap(resolution=resolution,
epsg=network.srid,
llcrnrlat=y1,
urcrnrlat=y2,
llcrnrlon=x1,
urcrnrlon=x2,
ax=ax)
bmap.drawcountries()
bmap.drawcoastlines()
x, y = bmap(x.values, y.values)
x = pd.Series(x, network.buses.index)
y = pd.Series(y, network.buses.index)
if isinstance(bus_sizes, pd.Series) and isinstance(bus_sizes.index,
pd.MultiIndex):
# We are drawing pies to show all the different shares
assert len(bus_sizes.index.levels[0].difference(network.buses.index)) == 0, \
"The first MultiIndex level of bus_sizes must contain buses"
assert isinstance(bus_colors, dict) and set(bus_colors).issuperset(bus_sizes.index.levels[1]), \
"bus_colors must be a dictionary defining a color for each element " \
"in the second MultiIndex level of bus_sizes"
bus_sizes = bus_sizes.sort_index(level=0, sort_remaining=False)
patches = []
for b_i in bus_sizes.index.levels[0]:
s = bus_sizes.loc[b_i]
radius = s.sum()**0.5
ratios = s / s.sum()
start = 0.25
for i, ratio in ratios.iteritems():
patches.append(
Wedge((x.at[b_i], y.at[b_i]),
radius,
360 * start,
360 * (start + ratio),
facecolor=bus_colors[i]))
start += ratio
bus_collection = PatchCollection(patches, match_original=True)
ax.add_collection(bus_collection)
else:
c = pd.Series(bus_colors, index=network.buses.index)
if c.dtype == np.dtype('O'):
c.fillna("b", inplace=True)
c = list(c.values)
s = pd.Series(bus_sizes, index=network.buses.index,
dtype="float").fillna(10)
bus_collection = ax.scatter(x, y, c=c, s=s, cmap=bus_cmap)
def as_branch_series(ser):
if isinstance(ser, dict) and set(ser).issubset(branch_components):
return pd.Series(ser)
elif isinstance(ser, pd.Series):
if isinstance(ser.index, pd.MultiIndex):
return ser
index = ser.index
ser = ser.values
else:
index = network.lines.index
return pd.Series(ser,
index=pd.MultiIndex(levels=(["Line"], index),
labels=(np.zeros(len(index)),
np.arange(len(index)))))
line_colors = as_branch_series(line_colors)
line_widths = as_branch_series(line_widths)
if not isinstance(line_cmap, dict):
line_cmap = {'Line': line_cmap}
branch_collections = []
for c in network.iterate_components(branch_components):
l_defaults = defaults_for_branches[c.name]
l_widths = line_widths.get(c.name, l_defaults['width'])
l_nums = None
l_colors = line_colors.get(c.name, l_defaults['color'])
if isinstance(l_colors, pd.Series):
if issubclass(l_colors.dtype.type, np.number):
l_nums = l_colors
l_colors = None
else:
l_colors.fillna(l_defaults['color'], inplace=True)
if not geometry:
segments = (np.asarray(
((c.df.bus0.map(x), c.df.bus0.map(y)),
(c.df.bus1.map(x), c.df.bus1.map(y)))).transpose(2, 0, 1))
else:
from shapely.wkt import loads
from shapely.geometry import LineString
linestrings = c.df.geometry.map(loads)
assert all(isinstance(ls, LineString) for ls in linestrings), \
"The WKT-encoded geometry in the 'geometry' column must be composed of LineStrings"
segments = np.asarray(list(linestrings.map(np.asarray)))
if basemap and basemap_present:
segments = np.transpose(
bmap(*np.transpose(segments, (2, 0, 1))), (1, 2, 0))
l_collection = LineCollection(segments,
linewidths=l_widths,
antialiaseds=(1, ),
colors=l_colors,
transOffset=ax.transData)
if l_nums is not None:
l_collection.set_array(np.asarray(l_nums))
l_collection.set_cmap(line_cmap.get(c.name, None))
l_collection.autoscale()
ax.add_collection(l_collection)
l_collection.set_zorder(1)
branch_collections.append(l_collection)
bus_collection.set_zorder(2)
ax.update_datalim(compute_bbox_with_margins(margin, x, y))
ax.autoscale_view()
ax.set_title(title)
return (bus_collection, ) + tuple(branch_collections)
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
**kwds):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color string, or array of floats
Edge color. Can be a single color format string (default='r'),
or a sequence of colors with the same length as edgelist.
If numeric values are specified they will be mapped to
colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=1.0)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
label : [None| string]
Label for legend
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
Notes
-----
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
Yes, it is ugly but drawing proper arrows with Matplotlib this
way is tricky.
Examples
--------
>>> G=nx.dodecahedral_graph()
>>> edges=nx.draw_networkx_edges(G,pos=nx.spring_layout(G))
Also see the NetworkX drawing examples at
http://networkx.github.io/documentation/latest/gallery.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap
from matplotlib.collections import LineCollection
import numpy
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
# set edge positions
edge_pos = numpy.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color) == len(edge_pos):
if numpy.alltrue([cb.is_string_like(c)
for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c, alpha)
for c in edge_color])
elif numpy.alltrue([not cb.is_string_like(c)
for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if numpy.alltrue([cb.iterable(c) and len(c) in (3, 4)
for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must consist of either color names or numbers')
else:
if cb.is_string_like(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError('edge_color must be a single color or list of exactly m colors where m is the number or edges')
edge_collection = LineCollection(edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1,),
linestyle=style,
transOffset = ax.transData,
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
# Note: there was a bug in mpl regarding the handling of alpha values for
# each line in a LineCollection. It was fixed in matplotlib in r7184 and
# r7189 (June 6 2009). We should then not set the alpha value globally,
# since the user can instead provide per-edge alphas now. Only set it
# globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(numpy.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
arrow_collection = None
if G.is_directed() and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = edge_colors
a_pos = []
p = 1.0-0.25 # make head segment 25 percent of edge length
for src, dst in edge_pos:
x1, y1 = src
x2, y2 = dst
dx = x2-x1 # x offset
dy = y2-y1 # y offset
d = numpy.sqrt(float(dx**2 + dy**2)) # length of edge
if d == 0: # source and target at same position
continue
if dx == 0: # vertical edge
xa = x2
ya = dy*p+y1
if dy == 0: # horizontal edge
ya = y2
xa = dx*p+x1
else:
theta = numpy.arctan2(dy, dx)
xa = p*d*numpy.cos(theta)+x1
ya = p*d*numpy.sin(theta)+y1
a_pos.append(((xa, ya), (x2, y2)))
arrow_collection = LineCollection(a_pos,
colors=arrow_colors,
linewidths=[4*ww for ww in lw],
antialiaseds=(1,),
transOffset = ax.transData,
)
arrow_collection.set_zorder(1) # edges go behind nodes
arrow_collection.set_label(label)
ax.add_collection(arrow_collection)
# update view
minx = numpy.amin(numpy.ravel(edge_pos[:, :, 0]))
maxx = numpy.amax(numpy.ravel(edge_pos[:, :, 0]))
miny = numpy.amin(numpy.ravel(edge_pos[:, :, 1]))
maxy = numpy.amax(numpy.ravel(edge_pos[:, :, 1]))
w = maxx-minx
h = maxy-miny
padx, pady = 0.05*w, 0.05*h
corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady)
ax.update_datalim(corners)
ax.autoscale_view()
# if arrow_collection:
return edge_collection
def plot(network,
margin=0.05,
ax=None,
basemap=True,
bus_colors='b',
line_colors='g',
bus_sizes=10,
line_widths=2,
title="",
line_cmap=None,
bus_cmap=None,
boundaries=None,
geometry=False,
branch_types=['Line', 'Link']):
"""
Plot the network buses and lines using matplotlib and Basemap.
Parameters
----------
margin : float
Margin at the sides as proportion of distance between max/min x,y
ax : matplotlib ax, defaults to plt.gca()
Axis to which to plot the network
basemap : bool, default True
Switch to use Basemap
bus_colors : dict/pandas.Series
Colors for the buses, defaults to "b"
bus_sizes : dict/pandas.Series
Sizes of bus points, defaults to 10
line_colors : dict/pandas.Series
Colors for the lines, defaults to "g" for Lines and "cyan" for
Links. Colors for branches other than Lines can be
specified using a pandas Series with a MultiIndex.
line_widths : dict/pandas.Series
Widths of lines, defaults to 2. Widths for branches other
than Lines can be specified using a pandas Series with a
MultiIndex.
title : string
Graph title
line_cmap : plt.cm.ColorMap/str|dict
If line_colors are floats, this color map will assign the colors.
Use a dict to specify colormaps for more than one branch type.
bus_cmap : plt.cm.ColorMap/str
If bus_colors are floats, this color map will assign the colors
boundaries : list of four floats
Boundaries of the plot in format [x1,x2,y1,y2]
branch_types : list of str or pypsa.component
Branch types to be plotted, defaults to Line and Link.
Returns
-------
bus_collection, branch_collection1, ... : tuple of Collections
Collections for buses and branches.
"""
defaults_for_branches = {
'Link': dict(color="cyan", width=2),
'Line': dict(color="b", width=2)
}
if not plt_present:
logger.error("Matplotlib is not present, so plotting won't work.")
return
if ax is None:
ax = plt.gca()
def compute_bbox_with_margins(margin, x, y):
#set margins
pos = np.asarray((x, y))
minxy, maxxy = pos.min(axis=1), pos.max(axis=1)
xy1 = minxy - margin * (maxxy - minxy)
xy2 = maxxy + margin * (maxxy - minxy)
return tuple(xy1), tuple(xy2)
x = network.buses["x"]
y = network.buses["y"]
if basemap and basemap_present:
if boundaries is None:
(x1, y1), (x2, y2) = compute_bbox_with_margins(margin, x, y)
else:
x1, x2, y1, y2 = boundaries
bmap = Basemap(resolution='l',
epsg=network.srid,
llcrnrlat=y1,
urcrnrlat=y2,
llcrnrlon=x1,
urcrnrlon=x2,
ax=ax)
bmap.drawcountries()
bmap.drawcoastlines()
x, y = bmap(x.values, y.values)
x = pd.Series(x, network.buses.index)
y = pd.Series(y, network.buses.index)
c = pd.Series(bus_colors, index=network.buses.index)
if c.dtype == np.dtype('O'):
c.fillna("b", inplace=True)
c = list(c.values)
s = pd.Series(bus_sizes, index=network.buses.index,
dtype="float").fillna(10)
bus_collection = ax.scatter(x, y, c=c, s=s, cmap=bus_cmap)
def as_branch_series(ser):
if isinstance(ser, pd.Series):
if isinstance(ser.index, pd.MultiIndex):
return ser
index = ser.index
ser = ser.values
else:
index = network.lines.index
return pd.Series(ser,
index=pd.MultiIndex(levels=(["Line"], index),
labels=(np.zeros(len(index)),
np.arange(len(index)))))
line_colors = as_branch_series(line_colors)
line_widths = as_branch_series(line_widths)
if not isinstance(line_cmap, dict):
line_cmap = {'Line': line_cmap}
branch_collections = []
for t in network.iterate_components(branch_types):
l_defaults = defaults_for_branches[t.name]
l_widths = line_widths.get(t.name, l_defaults['width'])
l_nums = None
if t.name in line_colors:
l_colors = line_colors[t.name]
if issubclass(l_colors.dtype.type, np.number):
l_nums = l_colors
l_colors = None
else:
l_colors.fillna(l_defaults['color'], inplace=True)
else:
l_colors = l_defaults['color']
if not geometry:
segments = (np.asarray(
((t.df.bus0.map(x), t.df.bus0.map(y)),
(t.df.bus1.map(x), t.df.bus1.map(y)))).transpose(2, 0, 1))
else:
from shapely.wkt import loads
from shapely.geometry import LineString
linestrings = t.df.geometry.map(loads)
assert all(isinstance(ls, LineString) for ls in linestrings), \
"The WKT-encoded geometry in the 'geometry' column must be composed of LineStrings"
segments = np.asarray(list(linestrings.map(np.asarray)))
if basemap and basemap_present:
segments = np.transpose(
bmap(*np.transpose(segments, (2, 0, 1))), (1, 2, 0))
l_collection = LineCollection(segments,
linewidths=l_widths,
antialiaseds=(1, ),
colors=l_colors,
transOffset=ax.transData)
if l_nums is not None:
l_collection.set_array(np.asarray(l_nums))
l_collection.set_cmap(line_cmap.get(t.name, None))
l_collection.autoscale()
ax.add_collection(l_collection)
l_collection.set_zorder(1)
branch_collections.append(l_collection)
bus_collection.set_zorder(2)
ax.update_datalim(compute_bbox_with_margins(margin, x, y))
ax.autoscale_view()
ax.set_title(title)
return (bus_collection, ) + tuple(branch_collections)
def draw_networkx_edges(G,
pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=None,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
**kwds):
"""Draw the edges of the graph G
This draws only the edges of the graph G.
pos is a dictionary keyed by vertex with a two-tuple
of x-y positions as the value.
See networkx.layout for functions that compute node positions.
edgelist is an optional list of the edges in G to be drawn.
If provided, only the edges in edgelist will be drawn.
edgecolor can be a list of matplotlib color letters such as 'k' or
'b' that lists the color of each edge; the list must be ordered in
the same way as the edge list. Alternatively, this list can contain
numbers and those number are mapped to a color scale using the color
map edge_cmap. Finally, it can also be a list of (r,g,b) or (r,g,b,a)
tuples, in which case these will be used directly to color the edges. If
the latter mode is used, you should not provide a value for alpha, as it
would be applied globally to all lines.
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
See draw_networkx for the list of other optional parameters.
"""
try:
import matplotlib
import matplotlib.pylab as pylab
import numpy as np
from matplotlib.colors import colorConverter, Colormap
from matplotlib.collections import LineCollection
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
pass # unable to open display
if ax is None:
ax = pylab.gca()
if edgelist is None:
edgelist = G.edges()
if not edgelist or len(edgelist) == 0: # no edges!
return None
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width, )
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color)==len(edge_pos):
if np.alltrue([cb.is_string_like(c) for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple(
[colorConverter.to_rgba(c, alpha) for c in edge_color])
elif np.alltrue([not cb.is_string_like(c) for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if np.alltrue(
[cb.iterable(c) and len(c) in (3, 4) for c in edge_color]):
edge_colors = tuple(edge_color)
alpha = None
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError(
'edge_color must consist of either color names or numbers')
else:
if len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError(
'edge_color must be a single color or list of exactly m colors where m is the number or edges'
)
edge_collection = LineCollection(
edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
)
# Note: there was a bug in mpl regarding the handling of alpha values for
# each line in a LineCollection. It was fixed in matplotlib in r7184 and
# r7189 (June 6 2009). We should then not set the alpha value globally,
# since the user can instead provide per-edge alphas now. Only set it
# globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
# need 0.87.7 or greater for edge colormaps. No checks done, this will
# just not work with an older mpl
if edge_colors is None:
if edge_cmap is not None: assert (isinstance(edge_cmap, Colormap))
edge_collection.set_array(np.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
pylab.sci(edge_collection)
arrow_collection = None
if G.is_directed() and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = (colorConverter.to_rgba('k', alpha), )
a_pos = []
p = 1.0 - 0.25 # make head segment 25 percent of edge length
for src, dst in edge_pos:
x1, y1 = src
x2, y2 = dst
dx = x2 - x1 # x offset
dy = y2 - y1 # y offset
d = np.sqrt(float(dx**2 + dy**2)) # length of edge
if d == 0: # source and target at same position
continue
if dx == 0: # vertical edge
xa = x2
ya = dy * p + y1
if dy == 0: # horizontal edge
ya = y2
xa = dx * p + x1
else:
theta = np.arctan2(dy, dx)
xa = p * d * np.cos(theta) + x1
ya = p * d * np.sin(theta) + y1
a_pos.append(((xa, ya), (x2, y2)))
arrow_collection = LineCollection(
a_pos,
colors=arrow_colors,
linewidths=[4 * ww for ww in lw],
antialiaseds=(1, ),
transOffset=ax.transData,
)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
edge_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(edge_collection)
if arrow_collection:
arrow_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(arrow_collection)
return edge_collection
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=None,
arrowstyle='-|>',
arrowsize=10,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
node_size=300,
nodelist=None,
node_shape="o",
connectionstyle=None,
**kwds):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color or array of colors (default='k')
Edge color. Can be a single color or a sequence of colors with the same
length as edgelist. Color can be string, or rgb (or rgba) tuple of
floats from 0-1. If numeric values are specified they will be
mapped to colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=None)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
Note: Arrows will be the same color as edges.
arrowstyle : str, optional (default='-|>')
For directed graphs, choose the style of the arrow heads.
See :py:class: `matplotlib.patches.ArrowStyle` for more
options.
arrowsize : int, optional (default=10)
For directed graphs, choose the size of the arrow head head's length and
width. See :py:class: `matplotlib.patches.FancyArrowPatch` for attribute
`mutation_scale` for more info.
connectionstyle : str, optional (default=None)
Pass the connectionstyle parameter to create curved arc of rounding
radius rad. For example, connectionstyle='arc3,rad=0.2'.
See :py:class: `matplotlib.patches.ConnectionStyle` and
:py:class: `matplotlib.patches.FancyArrowPatch` for more info.
label : [None| string]
Label for legend
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
list of matplotlib.patches.FancyArrowPatch
`FancyArrowPatch` instances of the directed edges
Depending whether the drawing includes arrows or not.
Notes
-----
For directed graphs, arrows are drawn at the head end. Arrows can be
turned off with keyword arrows=False. Be sure to include `node_size` as a
keyword argument; arrows are drawn considering the size of nodes.
Examples
--------
>>> G = nx.dodecahedral_graph()
>>> edges = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> G = nx.DiGraph()
>>> G.add_edges_from([(1, 2), (1, 3), (2, 3)])
>>> arcs = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> alphas = [0.3, 0.4, 0.5]
>>> for i, arc in enumerate(arcs): # change alpha values of arcs
... arc.set_alpha(alphas[i])
Also see the NetworkX drawing examples at
https://networkx.github.io/documentation/latest/auto_examples/index.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap, Normalize
from matplotlib.collections import LineCollection
from matplotlib.patches import FancyArrowPatch
import numpy as np
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
if nodelist is None:
nodelist = list(G.nodes())
# FancyArrowPatch handles color=None different from LineCollection
if edge_color is None:
edge_color = 'k'
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
# Check if edge_color is an array of floats and map to edge_cmap.
# This is the only case handled differently from matplotlib
if cb.iterable(edge_color) and (len(edge_color) == len(edge_pos)) \
and np.alltrue([isinstance(c,Number) for c in edge_color]):
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
else:
edge_cmap = plt.get_cmap()
if edge_vmin is None:
edge_vmin = min(edge_color)
if edge_vmax is None:
edge_vmax = max(edge_color)
color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
edge_color = [edge_cmap(color_normal(e)) for e in edge_color]
if (not G.is_directed() or not arrows):
edge_collection = LineCollection(edge_pos,
colors=edge_color,
linewidths=width,
antialiaseds=(1,),
linestyle=style,
transOffset=ax.transData,
alpha=alpha
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
return edge_collection
arrow_collection = None
if G.is_directed() and arrows:
# Note: Waiting for someone to implement arrow to intersection with
# marker. Meanwhile, this works well for polygons with more than 4
# sides and circle.
def to_marker_edge(marker_size, marker):
if marker in "s^>v<d": # `large` markers need extra space
return np.sqrt(2 * marker_size) / 2
else:
return np.sqrt(marker_size) / 2
# Draw arrows with `matplotlib.patches.FancyarrowPatch`
arrow_collection = []
mutation_scale = arrowsize # scale factor of arrow head
# FancyArrowPatch doesn't handle color strings
arrow_colors = colorConverter.to_rgba_array(edge_color,alpha)
for i, (src, dst) in enumerate(edge_pos):
x1, y1 = src
x2, y2 = dst
shrink_source = 0 # space from source to tail
shrink_target = 0 # space from head to target
if cb.iterable(node_size): # many node sizes
src_node, dst_node = edgelist[i][:2]
index_node = nodelist.index(dst_node)
marker_size = node_size[index_node]
shrink_target = to_marker_edge(marker_size, node_shape)
else:
shrink_target = to_marker_edge(node_size, node_shape)
if cb.iterable(arrow_colors):
if len(arrow_colors) == len(edge_pos):
arrow_color = arrow_colors[i]
elif len(arrow_colors)==1:
arrow_color = arrow_colors[0]
else: # Cycle through colors
arrow_color = arrow_colors[i%len(arrow_colors)]
else:
arrow_color = edge_color
if cb.iterable(width):
if len(width) == len(edge_pos):
line_width = width[i]
else:
line_width = width[i%len(width)]
else:
line_width = width
arrow = FancyArrowPatch((x1, y1), (x2, y2),
arrowstyle=arrowstyle,
shrinkA=shrink_source,
shrinkB=shrink_target,
mutation_scale=mutation_scale,
color=arrow_color,
linewidth=line_width,
connectionstyle=connectionstyle,
zorder=1) # arrows go behind nodes
# There seems to be a bug in matplotlib to make collections of
# FancyArrowPatch instances. Until fixed, the patches are added
# individually to the axes instance.
arrow_collection.append(arrow)
ax.add_patch(arrow)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
plt.tick_params(
axis='both',
which='both',
bottom=False,
left=False,
labelbottom=False,
labelleft=False)
return arrow_collection
