Search code examples
pythonmatplotlibrectangles

How to draw graph stick with rotated rectangle properly in matplotlib?

I am trying to figure out how to draw the following properly.

Lets say I have two nodes, i.e. 2 points in R^2, namely src=(x1,y1) and dst=(x2,y2).

I want to draw an edge/line between them, but starting with a box.

This is an example I created with Tikz: enter image description here

src and dest are the centers of these nodes. Note that they can also be on different heights. The problem is when I try to build the box (using patches.Rectangle), I have to rotate it and then it changes its size.

Thanks in advance

Solution

  • If you're happy to use a filled rectange for the graph with a rotated rectangle as stick, this can be emulated by a simple line, which is as thick as a scatter point at the graph edges. 

    import matplotlib.pyplot as plt
import numpy as np

def plot_spec_line(src, dst, length=0.3, srctext="v", dsttext="w", 
                   ax=None, color="k", textcolor="k"):
    if not ax: ax=plt.gca()
    lend=src+(dst-src)*length
    ax.plot([src[0],lend[0]], [src[1],lend[1]], lw=24,solid_capstyle="butt", zorder=1, color=color )
    ax.plot([src[0],dst[0]], [src[1],dst[1]], lw=2,solid_capstyle="butt", zorder=0, color=color)
    ax.scatter([src[0],dst[0]], [src[1],dst[1]], s=24**2, marker="o", lw=2, edgecolors=color, c="w", zorder=2)
    ax.text(src[0],src[1], srctext, fontsize=12, ha="center", va="center", zorder=3, color=textcolor)
    ax.text(dst[0],dst[1], dsttext, fontsize=12, ha="center", va="center", zorder=3, color=textcolor)

s = np.array([1,1])
d = np.array([3,2])    
plot_spec_line(s, d, length=0.3, srctext="v", dsttext="w")

s = np.array([1.5,0.9])
d = np.array([2.8,1.2])    
plot_spec_line(s, d, length=0.2, srctext="a", dsttext="b", color="gray")

s = np.array([1,2])
d = np.array([2,1.9])    
plot_spec_line(s, d, length=0.7, srctext="X", dsttext="Y", textcolor="limegreen", color="limegreen")

plt.margins(0.2)
plt.show()

    enter image description here

    To obtain a rectangle, which can be filled or not and more imporatantly, has an edge, you can use a Rectangle. You can rotate a rectangle by setting an appropriate transform. It then makes sense to use a plot of equal aspect ratio, such that the Rectangle and circles are not skewed.

    import matplotlib.pyplot as plt
import matplotlib.transforms
import numpy as np

def plot_spec_line(src, dst, length=0.3, radius=0.1, srctext="v", dsttext="w", 
                   ax=None, color="k", textcolor="k", reccolor="w", lw=1 ):
    if not ax: ax=plt.gca()
    lend=np.sqrt(np.sum(((dst-src)*length)**2))
    s = dst-src
    angle = np.rad2deg(np.arctan2(s[1],s[0]))
    delta = np.array([0,radius])
    tr = matplotlib.transforms.Affine2D().rotate_deg_around(src[0],src[1], angle)
    t = tr + ax.transData
    rec = plt.Rectangle(src-delta, width=lend, height=radius*2, ec=color,facecolor=reccolor, transform=t, linewidth=lw)
    ax.add_patch(rec)
    ax.plot([src[0],dst[0]], [src[1],dst[1]], lw=lw,solid_capstyle="butt", zorder=0, color=color)
    circ1= plt.Circle(src, radius=radius, fill=True, facecolor="w", edgecolor=color, lw=lw)
    circ2= plt.Circle(dst, radius=radius, fill=True, facecolor="w", edgecolor=color, lw=lw)
    ax.add_patch(circ1)
    ax.add_patch(circ2)
    ax.text(src[0],src[1], srctext, fontsize=12, ha="center", va="center", zorder=3, color=textcolor)
    ax.text(dst[0],dst[1], dsttext, fontsize=12, ha="center", va="center", zorder=3, color=textcolor)

s = np.array([1,1])
d = np.array([3,2])    
plot_spec_line(s, d, length=0.3, srctext="v", dsttext="w", radius=0.06,
               reccolor="plum", )

s = np.array([1.5,0.9])
d = np.array([2.8,1.2])    
plot_spec_line(s, d, length=0.2, srctext="a", dsttext="b", color="gray", lw=2.5)

s = np.array([1,2])
d = np.array([2,1.9])    
plot_spec_line(s, d, length=0.7, srctext="X", dsttext="Y", textcolor="limegreen",
               reccolor="palegreen", color="limegreen")

plt.margins(0.2)
plt.gca().set_aspect("equal")
plt.show()

    enter image description here