Plotting energy transitions using matplotlib

Question

I want to reproduce plots like this one in matplotlib. I'm performing simulations where the displacement of a particle whose energy is defined by a function, e.g. E(X)=1/2kx^2. I have a list with the locations of the particle and want to plot the transitions over a plot of the energy function, preferably using arrows between successive locations. How do I do this in matplotlib?

I have attached an image of a sample plot (a gradient descent plot - my requirement is quite similar).

So given a list E made from E = [(x*0.1-2)**2+1 for x in range(40)] and a list of points X = [0, 3, 1.5, 2.2, 1.9, 1.95, 2.05] I would like a plot like the one I've attached

Answer 1

You can use plt.quiver with the diff of your points.

import numpy as np
import matplotlib.pyplot as plt

xp = [x*.1 for x in range(41)]
E = [(x-2)**2+1 for x in xp]

X = [0, 3, 1.5, 2.2, 1.9, 1.95, 2.05]
y = [(x-2)**2+1 for x in X]

# computing arrow vectors as diff of x and y
u = [u-v for u, v in zip(X[1:], X)]
v = [u-v for u, v in zip(y[1:], y)]

plt.figure(figsize=(5, 7))
plt.plot(xp, E)
plt.quiver(X[:-1], y[:-1], u, v, angles='xy', scale_units='xy', scale=1)
plt.scatter(X, y, edgecolors='blue', marker='o', color='white')
plt.show()

Output