class ArrowedLineCollection(LineCollection):
def __init__(self, *kl, **kw):
'''The same arguments as in matplotlib.collections.LineCollection.
To initiate arrows, additionally use 'add_arrows'.
'''
LineCollection.__init__(self, *kl, **kw)
self.kw = kw # to be reused in add_arrows
self.draw_arrows = True
return
def add_arrows(self, arrow_offsets=10, arrow_length=10, arrow_width=5):
'''Draw arrows at the end of every edge.
Parameters
----------
arrow_offsets: a list of length equal to the number of edges. It stores offstets (in pixels) by which arrows
should be moved back to make space for a node. A single value (the same for all) is also accepted.
arrow_length: the length of the arrow head, in pixels?
arrow_width: the width of the arrow head, in pixels?
'''
self.draw_arrows=True
self.arrows = LineCollection([], **self.kw)
self.arrow_length=arrow_length
self.arrow_width=arrow_width
self.arrow_offsets = []
if type(arrow_offsets)==list:
if len(arrow_offsets) != len(self.get_paths()): raise ValueError('arrow_offsets does not match the number of edges.')
self.arrow_offsets = arrow_offsets
else:
self.arrow_offsets = [arrow_offsets]*len(self.get_paths())
return
def draw(self, renderer):
''' Overrides the matplotlib.collections.LineCollection.draw(). Adds arrows.
'''
LineCollection.draw(self, renderer)
if not self.draw_arrows: return
segments=[]
for i,path in enumerate(self.get_paths()):
L= [self.get_transform().transform(numpy.transpose(j[0])) for j in path.iter_segments()]
V=L[1]-L[0]
scale = min(1.0, vector_length(V)/(4*self.arrow_offsets[i])) #make the arrows smaller if their size is comparable with edge length
node_shift = scale * self.arrow_offsets[i] * norm_vector(V)
head_length_vector = scale * self.arrow_length * norm_vector(V)
head_width_vector = scale * self.arrow_width * perpendicular_vector(norm_vector(V))
segments.append((L[1]-head_length_vector+head_width_vector-node_shift, L[1]-node_shift ,L[1]-head_length_vector-head_width_vector-node_shift))
self.arrows.set_segments(segments)
self.arrows.draw(renderer)
return
class MyCell(matplotlib.table.CustomCell):
""" Extending matplotlib tables.
Adapted from https://stackoverflow.com/a/53573651/505698
"""
def __init__(self, *args, visible_edges, **kwargs):
super().__init__(*args, visible_edges=visible_edges, **kwargs)
seg = np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0],
[0.0, 0.0]]).reshape(-1, 1, 2)
segments = np.concatenate([seg[:-1], seg[1:]], axis=1)
self.edgelines = LineCollection(segments,
edgecolor=kwargs.get("edgecolor"))
self._text.set_zorder(2)
self.set_zorder(1)
def set_transform(self, trans):
self.edgelines.set_transform(trans)
super().set_transform(trans)
def draw(self, renderer):
c = self.get_edgecolor()
self.set_edgecolor((1, 1, 1, 0))
super().draw(renderer)
self.update_segments(c)
self.edgelines.draw(renderer)
self.set_edgecolor(c)
def update_segments(self, color):
x, y = self.get_xy()
w, h = self.get_width(), self.get_height()
seg = np.array([[x, y], [x + w, y], [x + w, y + h], [x, y + h],
[x, y]]).reshape(-1, 1, 2)
segments = np.concatenate([seg[:-1], seg[1:]], axis=1)
self.edgelines.set_segments(segments)
self.edgelines.set_linewidth(self.get_linewidth())
colors = [
color if edge in self._visible_edges else (1, 1, 1, 0)
for edge in self._edges
]
self.edgelines.set_edgecolor(colors)
def get_path(self):
codes = [Path.MOVETO] + [Path.LINETO] * 3 + [Path.CLOSEPOLY]
return Path(
[[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0], [0.0, 0.0]],
codes,
readonly=True,
)
def draw(self, renderer):
''' Overrides the matplotlib.collections.LineCollection.draw(). Adds arrows.
'''
LineCollection.draw(self, renderer)
if not self.draw_arrows: return
segments=[]
for i,path in enumerate(self.get_paths()):
L= [self.get_transform().transform(numpy.transpose(j[0])) for j in path.iter_segments()]
V=L[1]-L[0]
scale = min(1.0, vector_length(V)/(4*self.arrow_offsets[i])) #make the arrows smaller if their size is comparable with edge length
node_shift = scale * self.arrow_offsets[i] * norm_vector(V)
head_length_vector = scale * self.arrow_length * norm_vector(V)
head_width_vector = scale * self.arrow_width * perpendicular_vector(norm_vector(V))
segments.append((L[1]-head_length_vector+head_width_vector-node_shift, L[1]-node_shift ,L[1]-head_length_vector-head_width_vector-node_shift))
self.arrows.set_segments(segments)
self.arrows.draw(renderer)
return
def draw(self, renderer, project=False):
if project:
self.do_3d_projection(renderer)
LineCollection.draw(self, renderer)
def draw(self, renderer, project=False):
if project:
self.do_3d_projection(renderer)
LineCollection.draw(self, renderer)
def draw(self, *argrs, **kargs):
self._transform_path()
if not self._nodraw:
LineCollection.draw(self, *argrs, **kargs)
def draw(self, *argrs, **kargs):
self._transform_path()
if not self._nodraw:
LineCollection.draw(self, *argrs, **kargs)
