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, *argrs, **kargs): self._transform_path() if not self._nodraw: LineCollection.draw(self, *argrs, **kargs)