I have the following code that should draw a cycloid with animation and save it to a gif
but after running the program, a white square appears that covers everything, I can't find the reason cycloid_animation
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.animation import FuncAnimation, PillowWriter
plt.rcParams['animation.html'] = 'html5'
R = 1
def circle(a, b, r):
# (a,b): the center of the circle
# r: the radius of the circle
# T: The number of the segments
T = 100
x, y = [0]*T, [0]*T
for i,theta in enumerate(np.linspace(0,2*np.pi,T)):
x[i] = a + r*np.cos(theta)
y[i] = b + r*np.sin(theta)
return x, y
# Calculate the cycloid line
thetas = np.linspace(0,4*np.pi,100)
cycloid_x = R*(thetas-np.sin(thetas))
cycloid_y = R*(1-np.cos(thetas))
cycloid_c = R*thetas
fig = plt.figure()
lns = []
trans = plt.axes().transAxes
for i in range(len(thetas)):
x,y = circle(cycloid_c[i], R, R)
ln1, = plt.plot(x, y, 'g-', lw=2)
ln2, = plt.plot(cycloid_x[:i+1] ,cycloid_y[:i+1], 'r-', lw=2)
ln3, = plt.plot(cycloid_x[i], cycloid_y[i], 'bo', markersize=4)
ln4, = plt.plot([cycloid_c[i], cycloid_x[i]], [R,cycloid_y[i]], 'y-', lw=2)
tx1 = plt.text(0.05, 0.8, r'$\theta$ = %.2f $\pi$' % (thetas[i]/np.pi), transform=trans)
lns.append([ln1,ln2,ln3,ln4,tx1])
plt.xlim(0,15)
plt.ylim(0,3)
plt.xlabel('x')
plt.ylabel('y')
plt.grid()
plt.axes().set_aspect('equal')
ani = animation.ArtistAnimation(fig, lns, interval=50)
#ani.save('cycloid_ArtistAnimation.mp4',writer='ffmpeg')
ani.save('cycloid_ArtistAnimation.gif',writer='pillow')
ani
Each time you call plt.axis() you are creating a new axis on top of the figure. Since what you want is to get the current axis and then apply the transformations, after creating the figure you should call plt.gca() to get the current axis and use that instead.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.animation import FuncAnimation, PillowWriter
plt.rcParams['animation.html'] = 'html5'
R = 1
def circle(a, b, r):
# (a,b): the center of the circle
# r: the radius of the circle
# T: The number of the segments
T = 100
x, y = [0]*T, [0]*T
for i,theta in enumerate(np.linspace(0,2*np.pi,T)):
x[i] = a + r*np.cos(theta)
y[i] = b + r*np.sin(theta)
return x, y
# Calculate the cycloid line
thetas = np.linspace(0,4*np.pi,100)
cycloid_x = R*(thetas-np.sin(thetas))
cycloid_y = R*(1-np.cos(thetas))
cycloid_c = R*thetas
fig = plt.figure()
lns = []
trans = plt.gca().transAxes #<=== HERE
for i in range(len(thetas)):
x,y = circle(cycloid_c[i], R, R)
ln1, = plt.plot(x, y, 'g-', lw=2)
ln2, = plt.plot(cycloid_x[:i+1] ,cycloid_y[:i+1], 'r-', lw=2)
ln3, = plt.plot(cycloid_x[i], cycloid_y[i], 'bo', markersize=4)
ln4, = plt.plot([cycloid_c[i], cycloid_x[i]], [R,cycloid_y[i]], 'y-', lw=2)
tx1 = plt.text(0.05, 0.8, r'$\theta$ = %.2f $\pi$' % (thetas[i]/np.pi), transform=trans)
lns.append([ln1,ln2,ln3,ln4,tx1])
plt.xlim(0,15)
plt.ylim(0,3)
plt.xlabel('x')
plt.ylabel('y')
plt.grid()
plt.gca().set_aspect('equal') #<=== And HERE
ani = animation.ArtistAnimation(fig, lns, interval=50)
#ani.save('cycloid_ArtistAnimation.mp4',writer='ffmpeg')
ani.save('cycloid_ArtistAnimation.gif',writer='pillow')
I have a set of xy cooridnates that generate a contour. For the code below, these cooridnates are from groups A and B in the df. I have also created a separate xy cooridnate that is called from C1_X and C1_Y. However this isn't used in generating the contour itself. It is a separate xy coordinate.
Question: Is it possible to return the z-value of the contour at the C1_X C1_Y coordinate?
I have found a separate question that is similar: multivariate spline interpolation in python scipy?. The figure in that question displays what I'm hoping to return but I just want the z-value for one xy coordinate.
The contour in this question is normalised so values fall between -1 and 1. I'm hoping to return the z-value for C1_X and C1_Y, which is the white scatter point seen in the figure beneath the code.
I have attempted to return the z-value for this point using:
# Attempt at returning the z-value for C1
f = RectBivariateSpline(X, Y, normPDF)
z = f(d['C1_X'], d['C1_Y'])
print(z)
But I'm returning an error: raise TypeError('x must be strictly increasing')
TypeError: x must be strictly increasing
I have commented out this function so the code runs.
Side note: This code is written for an animation.
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.stats as sts
import matplotlib.animation as animation
from mpl_toolkits.axes_grid1 import make_axes_locatable
from scipy.interpolate import RectBivariateSpline
DATA_LIMITS = [0, 15]
def datalimits(*data):
return DATA_LIMITS
def mvpdf(x, y, xlim, ylim, radius=1, velocity=0, scale=0, theta=0):
X,Y = np.meshgrid(np.linspace(*xlim), np.linspace(*ylim))
XY = np.stack([X, Y], 2)
PDF = sts.multivariate_normal([x, y]).pdf(XY)
return X, Y, PDF
def mvpdfs(xs, ys, xlim, ylim, radius=None, velocity=None, scale=None, theta=None):
PDFs = []
for i,(x,y) in enumerate(zip(xs,ys)):
X, Y, PDF = mvpdf(x, y, xlim, ylim)
PDFs.append(PDF)
return X, Y, np.sum(PDFs, axis=0)
fig, ax = plt.subplots(figsize = (10,6))
ax.set_xlim(DATA_LIMITS)
ax.set_ylim(DATA_LIMITS)
line_a, = ax.plot([], [], 'o', c='red', alpha = 0.5, markersize=5,zorder=3)
line_b, = ax.plot([], [], 'o', c='blue', alpha = 0.5, markersize=5,zorder=3)
scat = ax.scatter([], [], s=5**2,marker='o', c='white', alpha = 1,zorder=3)
lines=[line_a,line_b]
scats=[scat]
cfs = None
def plotmvs(tdf, xlim=datalimits(df['X']), ylim=datalimits(df['Y']), fig=fig, ax=ax):
global cfs
if cfs:
for tp in cfs.collections:
tp.remove()
df = tdf[1]
PDFs = []
for (group, gdf), group_line in zip(df.groupby('group'), (line_a, line_b)):
group_line.set_data(*gdf[['X','Y']].values.T)
X, Y, PDF = mvpdfs(gdf['X'].values, gdf['Y'].values, xlim, ylim)
PDFs.append(PDF)
for (group, gdf), group_line in zip(df.groupby('group'), lines+scats):
if group in ['A','B']:
group_line.set_data(*gdf[['X','Y']].values.T)
kwargs = {
'xlim': xlim,
'ylim': ylim
}
X, Y, PDF = mvpdfs(gdf['X'].values, gdf['Y'].values, **kwargs)
PDFs.append(PDF)
#plot white scatter point from C1_X, C1_Y
elif group in ['C']:
gdf['X'].values, gdf['Y'].values
scat.set_offsets(gdf[['X','Y']].values)
# normalize PDF by shifting and scaling, so that the smallest value is -1 and the largest is 1
normPDF = (PDFs[0]-PDFs[1])/max(PDFs[0].max(),PDFs[1].max())
''' Attempt at returning z-value for C1_X, C1_Y '''
''' This is the function that I am trying to write that will '''
''' return the contour value '''
#f = RectBivariateSpline(X[::-1, :], Y[::-1, :], normPDF[::-1, :])
#z = f(d['C1_X'], d['C1_Y'])
#print(z)
cfs = ax.contourf(X, Y, normPDF, cmap='jet', alpha = 1, levels=np.linspace(-1,1,10),zorder=1)
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.1)
cbar = fig.colorbar(cfs, ax=ax, cax=cax)
cbar.set_ticks([-1,-0.8,-0.6,-0.4,-0.2,0,0.2,0.4,0.6,0.8,1])
return cfs.collections + [scat] + [line_a,line_b]
''' Sample Dataframe '''
n = 1
time = range(n)
d = ({
'A1_X' : [3],
'A1_Y' : [6],
'A2_X' : [6],
'A2_Y' : [10],
'B1_X' : [12],
'B1_Y' : [2],
'B2_X' : [14],
'B2_Y' : [4],
'C1_X' : [4],
'C1_Y' : [6],
})
# a list of tuples of the form ((time, group_id, point_id, value_label), value)
tuples = [((t, k.split('_')[0][0], int(k.split('_')[0][1:]), k.split('_')[1]), v[i])
for k,v in d.items() for i,t in enumerate(time) ]
df = pd.Series(dict(tuples)).unstack(-1)
df.index.names = ['time', 'group', 'id']
#Code will eventually operate with multiple frames
interval_ms = 1000
delay_ms = 2000
ani = animation.FuncAnimation(fig, plotmvs, frames=df.groupby('time'), interval=interval_ms, repeat_delay=delay_ms,)
plt.show()
I am hoping to return the z value for the white scatter point. Intended Output will display the normalised z value (-1,1) for C1_X,C1_Y.
Upon visual inspection this would be between0.6 and 0.8
Edit 2:
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.stats as sts
import matplotlib.animation as animation
from mpl_toolkits.axes_grid1 import make_axes_locatable
from scipy.interpolate import RectBivariateSpline
import matplotlib.transforms as transforms
DATA_LIMITS = [-85, 85]
def datalimits(*data):
return DATA_LIMITS # dmin - spad, dmax + spad
def rot(theta):
theta = np.deg2rad(theta)
return np.array([
[np.cos(theta), -np.sin(theta)],
[np.sin(theta), np.cos(theta)]
])
def getcov(radius=1, scale=1, theta=0):
cov = np.array([
[radius*(scale + 1), 0],
[0, radius/(scale + 1)]
])
r = rot(theta)
return r # cov # r.T
def mvpdf(x, y, xlim, ylim, radius=1, velocity=0, scale=0, theta=0):
X,Y = np.meshgrid(np.linspace(*xlim), np.linspace(*ylim))
XY = np.stack([X, Y], 2)
x,y = rot(theta) # (velocity/2, 0) + (x, y)
cov = getcov(radius=radius, scale=scale, theta=theta)
PDF = sts.multivariate_normal([x, y], cov).pdf(XY)
return X, Y, PDF
def mvpdfs(xs, ys, xlim, ylim, radius=None, velocity=None, scale=None, theta=None):
PDFs = []
for i,(x,y) in enumerate(zip(xs,ys)):
kwargs = {
'radius': radius[i] if radius is not None else 0.5,
'velocity': velocity[i] if velocity is not None else 0,
'scale': scale[i] if scale is not None else 0,
'theta': theta[i] if theta is not None else 0,
'xlim': xlim,
'ylim': ylim
}
X, Y, PDF = mvpdf(x, y,**kwargs)
PDFs.append(PDF)
return X, Y, np.sum(PDFs, axis=0)
fig, ax = plt.subplots(figsize = (10,6))
ax.set_xlim(DATA_LIMITS)
ax.set_ylim(DATA_LIMITS)
line_a, = ax.plot([], [], 'o', c='red', alpha = 0.5, markersize=3,zorder=3)
line_b, = ax.plot([], [], 'o', c='blue', alpha = 0.5, markersize=3,zorder=3)
lines=[line_a,line_b] ## this is iterable!
offset = lambda p: transforms.ScaledTranslation(p/82.,0, plt.gcf().dpi_scale_trans)
trans = plt.gca().transData
scat = ax.scatter([], [], s=5,marker='o', c='white', alpha = 1,zorder=3,transform=trans+offset(+2) )
scats=[scat]
cfs = None
def plotmvs(tdf, xlim=None, ylim=None, fig=fig, ax=ax):
global cfs
if cfs:
for tp in cfs.collections:
tp.remove()
df = tdf[1]
if xlim is None: xlim = datalimits(df['X'])
if ylim is None: ylim = datalimits(df['Y'])
PDFs = []
for (group, gdf), group_line in zip(df.groupby('group'), lines+scats):
if group in ['A','B']:
group_line.set_data(*gdf[['X','Y']].values.T)
kwargs = {
'radius': gdf['Radius'].values if 'Radius' in gdf else None,
'velocity': gdf['Velocity'].values if 'Velocity' in gdf else None,
'scale': gdf['Scaling'].values if 'Scaling' in gdf else None,
'theta': gdf['Rotation'].values if 'Rotation' in gdf else None,
'xlim': xlim,
'ylim': ylim
}
X, Y, PDF = mvpdfs(gdf['X'].values, gdf['Y'].values, **kwargs)
PDFs.append(PDF)
elif group in ['C']:
gdf['X'].values, gdf['Y'].values
scat.set_offsets(gdf[['X','Y']].values)
normPDF = (PDFs[0]-PDFs[1])/max(PDFs[0].max(),PDFs[1].max())
def get_contour_value_of_point(point_x, point_y, X, Y, Z, precision=10000):
CS = ax.contour(X, Y, Z, 100)
containing_levels = []
for cc, lev in zip(CS.collections, CS.levels):
for pp in cc.get_paths():
if pp.contains_point((point_x, point_y)):
containing_levels.append(lev)
if max(containing_levels) == 0:
return 0
else:
if max(containing_levels) > 0:
lev = max(containing_levels)
adj = 1. / precision
elif max(containing_levels) < 0:
lev = min(containing_levels)
adj = -1. / precision
is_inside = True
while is_inside:
CS = ax.contour(X, Y, Z, [lev])
for pp in CS.collections[0].get_paths():
if not pp.contains_point((point_x, point_y)):
is_inside = False
if is_inside:
lev += adj
return lev - adj
print(get_contour_value_of_point(d['C1_X'], d['C1_Y'], X, Y, normPDF))
cfs = ax.contourf(X, Y, normPDF, cmap='viridis', alpha = 1, levels=np.linspace(-1,1,10),zorder=1)
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.1)
cbar = fig.colorbar(cfs, ax=ax, cax=cax)
cbar.set_ticks([-1,-0.8,-0.6,-0.4,-0.2,0,0.2,0.4,0.6,0.8,1])
return cfs.collections + [scat] + [line_a,line_b]
''' Sample Dataframe '''
n = 10
time = range(n)
d = ({
'A1_X' : [3],
'A1_Y' : [6],
'A2_X' : [6],
'A2_Y' : [10],
'B1_X' : [12],
'B1_Y' : [2],
'B2_X' : [14],
'B2_Y' : [4],
'C1_X' : [4],
'C1_Y' : [6],
})
# a list of tuples of the form ((time, group_id, point_id, value_label), value)
tuples = [((t, k.split('_')[0][0], int(k.split('_')[0][1:]), k.split('_')[1]), v[i])
for k,v in d.items() for i,t in enumerate(time) ]
df = pd.Series(dict(tuples)).unstack(-1)
df.index.names = ['time', 'group', 'id']
#Code will eventually operate with multiple frames
interval_ms = 1000
delay_ms = 2000
ani = animation.FuncAnimation(fig, plotmvs, frames=df.groupby('time'), interval=interval_ms, repeat_delay=delay_ms,)
plt.show()
If you have an arbitrary cloud of (X, Y, Z) points and you want to interpolate the z-coordinate of some (x, y) point, you have a number of different options. The simplest is probably to just use scipy.interpolate.interp2d to get the z-value:
f = interp2d(X.T, Y.T, Z.T)
z = f(x, y)
Since the grid you have appears to be regular, you may be better off using scipy.interpolate.RectBivariateSpline, which has a very similar interface, but is specifically made for regular grids:
f = RectBivariateSpline(X.T, Y.T, Z.T)
z = f(x, y)
Since you have a regular meshgrid, you can also do
f = RectBivariateSpline(X[0, :], Y[:, 0], Z.T)
z = f(x, y)
Notice that the dimensions are flipped between the plotting arrays and the interpolation arrays. Plotting treats axis 0 as rows, i.e. Y, while the interpolation functions treat axis 0 as X. Rather than transposing, you could also switch the X and Y inputs, leaving Z intact for a similar end result, e.g.:
f = RectBivariateSpline(Y, X, Z)
z = f(y, x)
Alternatively, you could change all your plotting code to swap the inputs as well, but that would be too much work at this point. Whatever you do, pick an approach and stick with it. As long as you do it consistently, they should all work.
If you use one of the scipy approaches (recommended), keep the object f around to interpolate any further points you might want.
If you want a more manual approach, you can do something like find the three closest (X, Y, Z) points to (x, y), and find the value of the plane between them at (x, y). For example:
def interp_point(x, y, X, Y, Z):
"""
x, y: scalar coordinates to interpolate at
X, Y, Z: arrays of coordinates corresponding to function
"""
X = X.ravel()
Y = Y.ravel()
Z = Z.ravel()
# distances from x, y to all X, Y points
dist = np.hypot(X - x, Y - y)
# indices of the nearest points
nearest3 = np.argpartition(dist, 2)[:3]
# extract the coordinates
points = np.stack((X[nearest3], Y[nearest3], Z[nearest3]))
# compute 2 vectors in the plane
vecs = np.diff(points, axis=0)
# compute normal to plane
plane = np.cross(vecs[0], vecs[1])
# rhs of plane equation
d = np.dot(plane, points [:, 0])
# The final result:
z = (d - np.dot(plane[:2], [x, y])) / plane[-1]
return z
print(interp_point(x, y, X.T, Y.T, Z.T))
Since your data is on a regular grid, it might be easier to do something like bilinear interpolation on the quad surrounding (x, y):
def interp_grid(x, y, X, Y, Z):
"""
x, y: scalar coordinates to interpolate at
X, Y, Z: arrays of coordinates corresponding to function
"""
X, Y = X[:, 0], Y[0, :]
# find matching element
r, c = np.searchsorted(Y, y), np.searchsorted(X, x)
if r == 0: r += 1
if c == 0: c += 1
# interpolate
z = (Z[r - 1, c - 1] * (X[c] - x) * (Y[r] - y) +
Z[r - 1, c] * (x - X[c - 1]) * (Y[r] - y) +
Z[r, c - 1] * (X[c] - x) * (y - Y[r - 1]) +
Z[r, c] * (x - X[c - 1]) * (y - Y[r - 1])
) / ((X[c] - X[c - 1]) * (Y[r] - Y[r - 1]))
return z
print(interpolate_grid(x, y, X.T, Y.T, Z.T))
Here's an inelegant, brute force approach.* Assuming we have X, Y, and Z values, let's define a function that draws custom contour lines over and over until they intersect with the point at a user-defined level of precision (in your data, make Z = normPDF).
def get_contour_value_of_point(point_x, point_y, X, Y, Z, precision=10000):
fig, ax = plt.subplots()
CS = ax.contour(X, Y, Z, 100)
containing_levels = []
for cc, lev in zip(CS.collections, CS.levels):
for pp in cc.get_paths():
if pp.contains_point((point_x, point_y)):
containing_levels.append(lev)
if max(containing_levels) == 0:
return 0
else:
if max(containing_levels) > 0:
lev = max(containing_levels)
adj = 1. / precision
elif max(containing_levels) < 0:
lev = min(containing_levels)
adj = -1. / precision
is_inside = True
while is_inside:
CS = ax.contour(X, Y, Z, [lev])
for pp in CS.collections[0].get_paths():
if not pp.contains_point((point_x, point_y)):
is_inside = False
if is_inside:
lev += adj
return lev - adj
In more detail: what this is doing is drawing an initial contour map with 100 levels, then finding the list of contour levels whose polygons contain the point in question. We then find the narrowest level (either the highest if the levels are positive or the lowest if the levels are negative). From there, we tighten the level by small steps (corresponding to your desired precision level), checking if the point is still within the polygons. When the point is no longer within the contour polygon, we know that we've found the right level (the last one to contain the point).
As an example, we can use a contour in Matplotlib's library:
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
delta = 0.025
x = np.arange(-3.0, 3.0, delta)
y = np.arange(-2.0, 2.0, delta)
X, Y = np.meshgrid(x, y)
Z1 = np.exp(-X**2 - Y**2)
Z2 = np.exp(-(X - 1)**2 - (Y - 1)**2)
Z = (Z1 - Z2) * 2
With this setup, get_contour_value_of_point(0, -0.6) returns 1.338399999999998, which on a visual examination seems to match. get_contour_value_of_point(0, -0.6) returns -1.48, which also seems to match. Plots below for visual verification.
*I can't guarantee this will cover all use cases. It covered the ones I tried. I would test this fairly rigorously before getting it near any kind of production environment. I would expect there to be more elegant solutions than this (such as Mad Physicist's answer), but this was the one that occurred to me and seemed to work in straightforward, if brute-force, way.
I use matplot3d(ax.bar3d() and ax.plot) to plot two lines around a cuboid, when looking in bird viewbird view, this figure shows the read 3d relationship of all the objects. When looking from sideside view, some part of blue line that behind the green cylinder should be hidden, but this part still can be seen. I tried using parameter alpha, and it failed too. Does any body know how to deal with this ?
My matplotlib version is 2.1.0, and my python version is 3.4.7
# import
import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
# global set
XMAX = 4
ZMAX = 15
mpl.rcParams['legend.fontsize'] = 10
fig = plt.figure()
ax = fig.gca(projection='3d')
print(mpl.__version__)
z = np.linspace(0, ZMAX, 5)
if True:
arrowstartt = np.zeros((3))
arrowstartx = np.zeros((3))
arrowstarty = np.zeros((3))
arrowendt = np.zeros((3))
arrowendx = np.zeros((3))
arrowendy = np.zeros((3))
arrowcolor= 'black'
fontsizes = 25
# x cors
x = (XMAX + 1) / np.max(z) * z
y = 0. * z
zz = 0. * z
ax.plot(x, y, zz, color=arrowcolor, linewidth=1)
arrowstartx[0] = x[-1]
arrowstarty[0] = y[-1]
arrowstartt[0] = zz[-1]
ax.text(x[-1], y[-1], zz[-1], "x", color='k', fontsize=fontsizes)
# y cors
x = 0. * z
y = 2 * XMAX / np.max(z) * z
zz = 0. * z
ax.plot(x, y, zz, color=arrowcolor, linewidth=1)
arrowstartx[1] = x[-1]
arrowstarty[1] = y[-1]
arrowstartt[1] = zz[-1]
ax.text(x[-1], y[-1], zz[-1]-1, "y", color='k', fontsize=fontsizes)
#z cor
x = 0. * z
y = 0. * z
zz = (XMAX) / np.max(z) * z
ax.plot(x, y, zz, color=arrowcolor, linewidth=1)
arrowstartx[2] = x[-1]
arrowstarty[2] = y[-1]
arrowstartt[2] = zz[-1]
ax.text(x[-1], y[-1], zz[-1], "z", color='k', fontsize=fontsizes)
# arrow end
arrowendx = arrowstartx + [1, 0, 0]
arrowendy = arrowstarty + [0, 1, 0]
arrowendt = arrowstartt + [0, 0, 1]
ax.quiver(arrowstartx, arrowstarty, arrowstartt, arrowendx, arrowendy, arrowendt,2, color=arrowcolor, linewidth=1)
# ax.set_xlabel('x')
# ax.set_ylabel('y')
# ax.set_zlabe('z')
# ax.arrow()
''' draw bar as a cylinder '''
if True:
# draw bar
xpos = 0.
ypos = 30.
dx = 1.4
dy = 2
ax.bar3d(xpos, ypos, 0, dx, dy, ZMAX, color='g', zsort='average')
#ax.bar3d(xpos, ypos, 0, dx, dy, ZMAX, alpha=0.8, color='g',zsort='average')
''' draw two lines'''
if True:
# the blue line
y = np.arange(0, 50, 1)
x = np.ones(len(y)) * (-2)
z = np.linspace(0, ZMAX, len(y))
ax.plot(x, y, z, 'b')
#ax.plot(x, y, z, 'b',alpha=0.8)
# the red line
y = np.arange(0, 50, 1)
x = np.ones(len(y)) * 2
z = np.linspace(0, ZMAX, len(y))
ax.plot(x, y, z, 'r')
plt.axis('off')
ax.legend()
plt.show()
I've been fiddling with this bit of Python code to simualate a spring-pendulum system. I altered the equation slightly and it plots fine. However, I also want to add a persistent trace after it like in this gif.
Here is my full code (I can't trim it down any more since you need the ODE solved to generate the plotted data), the relevant bit is near the end:
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from numpy import sin, cos, pi, array
spring_constant = 22.93
length = 0.16
mass = 0.1
# initial conditions
init = array([-0.35, 0, 0.08, 1]) # initial values
#array([theta, theta_dot, x, x_dot])
#Return derivatives of the array z (= [theta, theta_dot, x, x_dot])
def deriv(z, t, spring_k, spring_l, bob_mass):
k = spring_k
l = spring_l
m = bob_mass
g = 9.8
theta = z[0]
thetadot = z[1]
x = z[2]
xdot= z[3]
return array([
thetadot,
(-1.0/(l+x)) * (2*xdot*thetadot + g*sin(theta)),
xdot,
g*cos(theta) + (l+x)*thetadot**2 - (k/m)*x
])
#Create time steps
time = np.linspace(0.0,10.0,1000)
#Numerically solve ODE
y = odeint(deriv,init,time, args = (spring_constant, length, mass))
l = length
r = l+y[:,2]
theta = y[:,0]
dt = np.mean(np.diff(time))
x = r*sin(theta)
y = -r*cos(theta)
##MATPLOTLIB BEGINS HERE##
fig = plt.figure()
ax = fig.add_subplot(111, autoscale_on=False,
xlim=(-1.2*r.max(), 1.2*r.max()),
ylim=(-1.2*r.max(), 0.2*r.max()), aspect = 1.0)
ax.grid()
##ANIMATION STUFF BEGINS HERE##
line, = ax.plot([], [], 'o-', lw=2)
time_template = 'time = %.1fs'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
def init():
line.set_data([], [])
time_text.set_text('')
return line, time_text
def animate(i):
thisx = [0, x[i]]
thisy = [0, y[i]]
line.set_data(thisx, thisy)
time_text.set_text(time_template%(i*dt))
return line, time_text
ani = animation.FuncAnimation(fig, animate, np.arange(1, len(y)),
interval=25, blit=True, init_func=init)
plt.show()
I tried making a list of points that gets appended to every time the animation loop calls, and then drawing all of those points so far each frame:
time_template = 'time = %.1fs'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
foox = []
fooy = []
def init():
line.set_data([], [])
foo.set_data([], [])
time_text.set_text('')
return line, time_text, foo
def animate(i):
thisx = [0, x[i]]
thisy = [0, y[i]]
foox += [x[i]]
fooy += [y[i]]
line.set_data(thisx, thisy)
foo.set_data(foox, fooy)
time_text.set_text(time_template%(i*dt))
return line, time_text, foo
But I get
UnboundLocalError: local variable 'foox' referenced before assignment
Which I guess means it doesn't like it when you use a global variable? I'm not sure how to keep a history of which points have been drawn without using a variable outside of the animate() scope. Anyone know how?
Thank you.
EDIT:
I solved it. I was using += instead of .append() by mistake. Now I feel like an idiot.
For posterity it should be:
def animate(i):
thisx = [0, x[i]]
thisy = [0, y[i]]
foox.append(x[i])
fooy.append(y[i])
line.set_data(thisx, thisy)
foo.set_data(foox, fooy)
time_text.set_text(time_template%(i*dt))
return line, time_text, foo
You are modifying global variables in your animate function, without declaring them as global
foo and line are also redundant
Other than that, your animation works well; you can run the following code to see it:
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from numpy import sin, cos, pi, array
spring_constant = 22.93
length = 0.16
mass = 0.1
# initial conditions
init = array([-0.35, 0, 0.08, 1]) # initial values
#array([theta, theta_dot, x, x_dot])
#Return derivatives of the array z (= [theta, theta_dot, x, x_dot])
def deriv(z, t, spring_k, spring_l, bob_mass):
k = spring_k
l = spring_l
m = bob_mass
g = 9.8
theta = z[0]
thetadot = z[1]
x = z[2]
xdot= z[3]
return array([
thetadot,
(-1.0/(l+x)) * (2*xdot*thetadot + g*sin(theta)),
xdot,
g*cos(theta) + (l+x)*thetadot**2 - (k/m)*x
])
#Create time steps
time = np.linspace(0.0,10.0,1000)
#Numerically solve ODE
y = odeint(deriv,init,time, args = (spring_constant, length, mass))
l = length
r = l+y[:,2]
theta = y[:,0]
dt = np.mean(np.diff(time))
x = r*sin(theta)
y = -r*cos(theta)
##MATPLOTLIB BEGINS HERE##
fig = plt.figure()
ax = fig.add_subplot(111, autoscale_on=False,
xlim=(-1.2*r.max(), 1.2*r.max()),
ylim=(-1.2*r.max(), 0.2*r.max()), aspect = 1.0)
ax.grid()
##ANIMATION STUFF BEGINS HERE##
line, = ax.plot([], [], 'o-', lw=2)
time_template = 'time = %.1fs'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
foox = []
fooy = []
#foo.set_data(foox, fooy)
def init():
global line, time_text, foo
line.set_data([], [])
# foo.set_data([], [])
time_text.set_text('')
return line, time_text#, foo
def animate(i):
global foox, fooy, foo
thisx = [0, x[i]]
thisy = [0, y[i]]
foox += [x[i]]
fooy += [y[i]]
line.set_data(thisx, thisy)
# foo.set_data(foox, fooy)
time_text.set_text(time_template%(i*dt))
return line, time_text#, foo
ani = animation.FuncAnimation(fig, animate, np.arange(1, len(y)), interval=25, blit=False, init_func=init)
plt.show()
I've set blit=False because last I checked, blit was not working on OSX
I solved it. I was using += instead of .append() by mistake. Now I feel like an idiot.
For posterity it should be:
def animate(i):
thisx = [0, x[i]]
thisy = [0, y[i]]
foox.append(x[i])
fooy.append(y[i])
line.set_data(thisx, thisy)
foo.set_data(foox, fooy)
time_text.set_text(time_template%(i*dt))
return line, time_text, foo