0

I'm basing this question on a useful, yet closed question by Ian Roberts, and awesome answer by hayk.

How to make a wavy arrow that goes between specified points on a plt plot, with specified number of waves? Method should be compatible with many subplots to make it universal.

I give the code in the answer. It took me some time to figure out how to make it work on subplots, so you don't have to.

user46147
  • 232
  • 1
  • 16

1 Answers1

0
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.path as mpath
import matplotlib.patches as mpatches
import matplotlib as mpl

def curly_arrow(start, end, ax,arr_size=1, n=5, col='gray', linew=1., width=0.1):
    xmin, ymin = start
    xmax, ymax = end
    dist = np.sqrt((xmin - xmax) ** 2 + (ymin - ymax) ** 2)
    n0 = dist / (2 * np.pi)

    x = np.linspace(0, dist, 151) + xmin
    y = width * np.sin(n * x / n0) + ymin
    line = plt.Line2D(x, y, color=col, lw=linew)

    del_x = xmax - xmin
    del_y = ymax - ymin
    ang = np.arctan2(del_y, del_x)

    line.set_transform(mpl.transforms.Affine2D().rotate_around(xmin, ymin, ang) + ax.transData)
    ax.add_line(line)

    verts = np.array([[0, 1], [0, -1], [2, 0], [0, 1]]).astype(float) * arr_size
    verts[:, 1] += ymax
    verts[:, 0] += xmax
    path = mpath.Path(verts)
    patch = mpatches.PathPatch(path, fc=col, ec=col)

    patch.set_transform(mpl.transforms.Affine2D().rotate_around(xmax, ymax, ang) + ax.transData)
    return patch

then you can use it like this:

fig, (ax1, ax2) = plt.subplots(1,2)

ax1.add_patch(curly_arrow((0, 0.5), (0.6, 0.5),ax1, n=4, arr_size=0.1))
ax2.add_patch(curly_arrow((0, 0.5), (0.6, 0.5),ax2, n=3, arr_size=0.1))
user46147
  • 232
  • 1
  • 16