How to animate a for loop with variable range - python

I get different plots for each value that I wrote in for loop's upper range as I expected. But I would like to animate the plot from range (0,0) an goes to (0,15) as the upper limit changes 1 by 1, by using matplotlib animation function animation.FuncAnimation(). So there will be 16 frames total in the animation. I messed up with the animation part, so I'm pasting the code that gives 1 plot output. Thanks in advance!
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
g = 1.0
def distance(x1,y1,x2,y2):
r = np.sqrt((x2-x1)**2+(y2-y1)**2)
return r
def gravit(m1,m2,r):
f = g*m1*m2/(r**2)
return f
def angle(y2, y1, x2, x1):
ydif = y1-y2
xdif = x1-x2
angle = np.arctan2(ydif,xdif)
return angle
m1 = 100
x1, y1 = 0,0
m2 = 1
x2, y2 = -15,-10
vx1 = 1
vy1 = 0
ax1 = 0
ay1 = 0
vx2 = 2
vy2 = 3.9
ax2 = 0
ay2 = 0
x1coor = [x1]
y1coor = [y1]
x2coor = [x2]
y2coor = [y2]
for t in range(0,10): #This value of 10 should be change from 0 to 15 in the animation
r = distance(x1,y1,x2,y2)
fx1 = gravit(m1, m2, r) * np.cos(angle(y2,y1,x2,x1))
fy1 = gravit(m1, m2, r) * np.sin(angle(y2,y1,x2,x1))
ax2 = fx1/m2
vx2 = vx2 + ax2
ay2 = fy1/m2
vy2 = vy2 + ay2
x2 = x2 + vx2 + 0.5*ax2
y2 = y2 + vy2 + 0.5*ay2
x1 = x1 + vx1 + 0.5*ax1
y1 = y1 + vy1 + 0.5*ay1
x1coor.append(x1)
y1coor.append(y1)
x2coor.append(x2)
y2coor.append(y2)
plt.axes().set_aspect('equal')
plt.axis([-30,30,-30,30])
plt.plot(x1coor,y1coor, '-.', color='blue')
plt.plot(x2coor,y2coor, '-.', color='black')
plt.scatter(x1,y1,s=m1*20, color='blue')
plt.scatter(x2,y2,s=m2*20, color='red')

The FuncAnimation function of matplotlib's animation module requires a figure and a function to draw each frame, so first initialize your figure following :
fig = plt.figure()
ax1 = fig.add_subplot(1, 1, 1)
ax1.set_aspect("equal")
ax1.set_xlim(-30, 30)
ax1.set_ylim(-30, 30)
l1, = ax1.plot([], [], linestyle = "-.", color = "blue")
l2, = ax1.plot([], [], '-.', color = "black")
s1, = ax1.plot([], [], linestyle = None, marker = "o", markersize = 5, color = "blue")
s2, = ax1.plot([], [], linestyle = None, marker = "o", markersize = 5, color = "red")
Note that the different line plots contain no input data and are saved in the variables l1, l2, s1 and s2.
Now define the function that will be called to draw each frame. The first argument is always the frame number. This function should return all the plot objects that are updated using their methods 'set_data':
def update_fig(i, x1coor, y1coor, x2coor, y2coor):
l1.set_data(x1coor[:i+1], y1coor[:i+1])
l2.set_data(x2coor[:i+1], y2coor[:i+1])
s1.set_data(x1coor[i], y1coor[i])
s2.set_data(x2coor[i], y2coor[i])
return l1, l2, s1, s2,
You can now run your animation (note that 'update_fig' arguments, except 'i', are passed to the function using the keyword 'fargs'):
ani = animation.FuncAnimation(fig, update_fig,
frames = len(x1coor),
fargs = (x1coor, y1coor, x2coor, y2coor),
interval = 100,
repeat = True,
)

Related

How to plot circles in polar plot in python

I have a list with x,y,z, r coordinates (cartesians). I need to plot some circles in a polar plot, but I don't know how to do it with cartesians coordinates.
I am trying to do it with this line
circle1 = plt.Circle((x[i], y[i]), r[i], transform=ax3.transData._b, color = 'r', alpha=0.5, fill=False)
but this doesn't seem to work because I obtain the circles too far away from the center of the origin.
any help?
data1 = pd.read_csv('Uchuu_lightcone_0_11.9_voids.txt', sep='\s+', header=None)
data1 = pd.DataFrame(data1)
x = data1[0]
y = data1[1]
r = data1[3]
z = data1[2]
azvoids, elvoids, rvoids = cart2sph(x,y,z)
d = ax3.scatter(azvoids, rvoids, s=3, c='red', alpha=1, marker='.')
for i in range(len(x)):
if elvoids[i] > 35 and elvoids[i] < 45:
circle1 = plt.Circle((x[i], y[i]), r[i], transform=ax3.transData._b, color = 'r', alpha=0.5, fill=False)
ax3.add_artist(circle1)
# The cart2sph function is
def cart2sph(x,y,z):
""" x, y, z : ndarray coordinates
ceval: backend to use:
- eval : pure Numpy
- numexpr.evaluate: Numexpr """
azimuth = arctan2(y,x)*180/math.pi
xy2 = x**2 + y**2
elevation = arctan2(z, sqrt(xy2))*180/math.pi
r = sqrt(xy2 + z**2)
return azimuth, elevation, r
You should use azvoids and rvoids to plot the center of the circle since you use those to show tham in the scatter plot
import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.transforms as transforms
def cart2sph(x, y, z):
""" x, y, z : ndarray coordinates
ceval: backend to use:
- eval : pure Numpy
- numexpr.evaluate: Numexpr """
azimuth = np.arctan2(y, x) * 180 / math.pi
xy2 = x ** 2 + y ** 2
elevation = np.arctan2(z, np.sqrt(xy2)) * 180 / math.pi
r = np.sqrt(xy2 + z ** 2)
return azimuth, elevation, r
#
# data1 = pd.read_csv('Uchuu_lightcone_0_11.9_voids.txt', sep='\s+', header=None)
# data1 = pd.DataFrame(data1)
N=100
x = (np.random.rand(N)-0.5)*100
y = (np.random.rand(N)-0.5)*100
z = (np.random.rand(N)-0.5)*100
r = np.random.rand(N)*10
azvoids, elvoids, rvoids = cart2sph(x, y, z)
fig = plt.figure()
ax3 = plt.subplot(111 )
d = plt.scatter(azvoids, elvoids, s=3 , c='red', alpha=1, marker='.' )
for i in range(len(x)):
if elvoids[i] > 35 and elvoids[i] < 45:
# circle1 = plt.Circle((azvoids[i], elvoids[i]), rvoids[i], color='r', alpha=0.5, fill=False)
x, y = ax3.transData.transform((azvoids[i], elvoids[i]))
trans = (fig.dpi_scale_trans +
transforms.ScaledTranslation(azvoids[i], elvoids[i], ax3.transData))
circle1 = plt.Circle((azvoids[i], elvoids[i]), rvoids[i] , color='r', alpha=0.5, fill=None)
ax3.add_artist(circle1)
plt.axis('equal')
plt.show()
# The cart2sph function is

Matplotlib Annotations disappear during animation when updating xlim

I am having an issue with FuncAnimation where my annotations are removed after updating xlim. Here is the code with a preview underneath
You can try the code in a google colab here https://colab.research.google.com/drive/1NrM-ZnSQKhADccpjCbNeOC5PU8uXw-Sb?authuser=2#scrollTo=bcYtgNaTYJ3g
import os
import matplotlib.animation as ani
import matplotlib.pyplot as plt
from collections import deque
from typing import List
from IPython.display import Image
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
def create_animation_for_data(qty_lists: List[List[int]], gif_path, add_horizontal_guides=True, dynamic_y_axis=True,
dynamic_x_axis=True):
plt.style.use(['dark_background'])
fruits_to_color = ["red", "orange"]
qtys1, qty2 = qty_lists
times = [i for i in range(len(qtys1))]
artists = []
zoom = 0.5
first_image_path = 'assets/apple.png'
first_image = plt.imread(first_image_path)
first_offset_image = OffsetImage(first_image, zoom=zoom)
second_image_path = 'assets/orange.png'
second_image = plt.imread(second_image_path)
second_offset_image = OffsetImage(second_image, zoom=zoom)
# initializing a figure in
# which the graph will be plotted
my_dpi = 200
figsize_pixels = (650, 1000)
figsize = (int(figsize_pixels[0] / my_dpi), figsize_pixels[1] / my_dpi)
fig = plt.figure(figsize=figsize, dpi=my_dpi)
fig.set_tight_layout(True)
records_per_second = 1
seconds_show_on_screen = 30
max_width_on_screen = records_per_second * seconds_show_on_screen
ticks_every = 20
graph_max_y = 100
max_y_list = list(range(60, graph_max_y + 1, ticks_every * 2))
max_y_index = 0
max_y = max_y_list[max_y_index]
highest_max_y = max_y_list[- 1]
xlim_min = times[0]
if dynamic_x_axis:
xlim_max = times[max_width_on_screen - 1]
else:
xlim_max = times[-1]
axes_xlim = (xlim_min, xlim_max)
if dynamic_y_axis:
ylim_max = max_y
else:
ylim_max = highest_max_y
axes_ylim = (0, ylim_max)
ax1 = plt.axes(xlim=axes_xlim, ylim=axes_ylim)
# Set a title
plt.title('Qty over time', fontsize=20)
# Set axis labels
plt.xlabel('Time', fontsize=18)
plt.ylabel('Qty', fontsize=18)
plotlays, plotcols = [2], fruits_to_color
labels = ['Apple', 'Orange']
lines = []
for index in range(len(qty_lists)):
lobj = ax1.plot([], [], lw=5, color=plotcols[index],
label=labels[index])[0]
lines.append(lobj)
# Make sure your axis ticks are large enough to be easily read.
# You don't want your viewers squinting to read your plot.
plt.yticks(range(0, highest_max_y + 1, ticks_every), [str(x) for x in range(0, highest_max_y + 1, ticks_every)],
fontsize=14)
plt.xticks(fontsize=14)
# Provide tick lines across the plot to help your viewers trace along
# the axis ticks. Make sure that the lines are light and small so they
# don't obscure the primary data lines.
if add_horizontal_guides:
max_x = len(times)
for y in range(0, highest_max_y + 1, ticks_every):
plt.plot(times, [y] * max_x, "--",
lw=1, color="white", alpha=0.7)
# Do this after the plotting done above
ax1.set_ylim(*axes_ylim)
# Remove the tick marks; they are unnecessary with the tick lines we just plotted.
plt.tick_params(axis="both", which="both", bottom="off", top="off",
labelbottom="on", left="off", right="off", labelleft="on")
# empty list to store x and y axis values
xdata = deque()
ydata1 = deque()
ydata2 = deque()
first_image_annotation = AnnotationBbox(
first_offset_image, (times[0], 0), xycoords='data', frameon=False)
artists.append(ax1.add_artist(first_image_annotation))
second_image_annotation = AnnotationBbox(
second_offset_image, (times[0], 0), xycoords='data', frameon=False)
artists.append(ax1.add_artist(second_image_annotation))
ann_list = [
first_image_annotation,
second_image_annotation,
]
# animation function
def animate(i):
nonlocal max_y_index
nonlocal ann_list
# appending new points to x, y axes points list
x1_and_x2 = times[i]
xdata.append(x1_and_x2)
y1 = qty_lists[0][i]
ydata1.append(y1)
y2 = qty_lists[1][i]
ydata2.append(y2)
xlist = [xdata, xdata]
ylist = [ydata1, ydata2]
# If we have passed our max_width_on_screen
if len(xdata) > max_width_on_screen:
# Delete the oldest record
xdata.popleft()
max_x = max(xdata)
min_x = max_x - seconds_show_on_screen
# Update our x axis
if dynamic_x_axis:
ax1.set_xlim(min_x, max_x)
graph_max_y_data = max(ydata1[-1], ydata2[-1])
max_y_data = max_y_list[max_y_index]
while graph_max_y_data > max_y_data:
max_y_index += 1
max_y_data = max_y_list[max_y_index]
# Update our y axis
if dynamic_y_axis:
ax1.set_ylim(0, max_y_data)
ydata1.popleft()
ydata2.popleft()
first_image_annotation_xybox = (x1_and_x2, y1)
first_image_annotation.xybox = first_image_annotation_xybox
second_image_annotation_xybox = (x1_and_x2, y2)
second_image_annotation.xybox = second_image_annotation_xybox
for lnum, line in enumerate(lines):
# set data for each line separately.
line.set_data(xlist[lnum], ylist[lnum])
return lines, ann_list
# call the animator
anim = ani.FuncAnimation(fig, animate, frames=len(times), interval=30, blit=False)
# save the animation as gif file
anim.save(gif_path, writer='imagemagick', fps=2)
return os.path.abspath(gif_path)
static_axes_gif = 'FuncAnimation-annotated-static-axes.gif'
print(static_axes_gif)
animation_path_static_axes = create_animation_for_data(data_to_plot,
static_axes_gif,
dynamic_x_axis=False,
dynamic_y_axis=False)
Image(url=animation_path_static_axes)
static_x_gif = 'FuncAnimation-annotated-static-x.gif'
print(static_x_gif)
animation_path_static_x_axis = create_animation_for_data(data_to_plot,
static_x_gif,
dynamic_x_axis=False,
dynamic_y_axis=True)
Image(url=animation_path_static_x_axis)
static_y_gif = 'FuncAnimation-annotated-static-y.gif'
print(static_y_gif)
animation_path_static_y_axis = create_animation_for_data(data_to_plot,
static_y_gif,
dynamic_x_axis=True,
dynamic_y_axis=False)
Image(url=animation_path_static_y_axis)
dynamic_axes_gif = 'FuncAnimation-annotated-dynamic-axes.gif'
print(dynamic_axes_gif)
animation_path_dynamic_axes = create_animation_for_data(data_to_plot,
dynamic_axes_gif,
dynamic_x_axis=True,
dynamic_y_axis=True)
Image(url=animation_path_dynamic_axes)
I am making 4 graphs with variable dynamic axes (dynamic = updating axis during amination):
FuncAnimation-annotated-static-axes.gif
xlim and ylim are fixed
FuncAnimation-annotated-static-x.gif
xlim is fixed
ylim is dynamic
FuncAnimation-annotated-static-y.gif
xlim is dynamic
ylim is fixed
FuncAnimation-annotated-dynamic-axes.gif
xlim and ylim are dynamic
My annotations disappear in the two cases where the xlim is updated:
FuncAnimation-annotated-static-y.gif
FuncAnimation-annotated-dynamic-axes.gif
Note that when xlim is static this doesn't happen:
FuncAnimation-annotated-static-axes.gif
FuncAnimation-annotated-static-x.gif
Does anyone know why this happens or how to update the xlim without removing annotations?
Please let me know if something is unclear / worded poorly as I really need to solve this.
So the issue is with how I was moving my annotation.
This is the fix:
# Don't do this - updating xlim will causing the annotation do disappear
# first_image_annotation_xybox = (x1_and_x2, y1)
# first_image_annotation.xybox = first_image_annotation_xybox
#
# second_image_annotation_xybox = (x1_and_x2, y2)
# second_image_annotation.xybox = second_image_annotation_xybox
for lnum, line in enumerate(lines):
# set data for each line separately.
line.set_data(xlist[lnum], ylist[lnum])
# Do this - Update our annotations
for ann in ann_list:
ann.remove()
ann_list = []
first_image_annotation = AnnotationBbox(
first_offset_image, (x1_and_x2, y1), xycoords='data', frameon=False)
ann_list.append(ax1.add_artist(first_image_annotation))
second_image_annotation = AnnotationBbox(
second_offset_image, (x1_and_x2, y2), xycoords='data', frameon=False)
ann_list.append(ax1.add_artist(second_image_annotation))
return lines, ann_list
The rest of the code is the same. Wonder why this happens on updating xlim and not on updating ylim ¯\(ツ)/¯

label Taylor diagram with number by Python

My Python code here:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.projections import PolarAxes
import mpl_toolkits.axisartist.grid_finder as gf
import mpl_toolkits.axisartist.floating_axes as fa
class TaylorDiagram(object):
def __init__(self, STD ,fig=None, rect=111, label='_'):
self.STD = STD
tr = PolarAxes.PolarTransform()
# Correlation labels
rlocs = np.concatenate(((np.arange(11.0) / 10.0), [0.95, 0.99]))
tlocs = np.arccos(rlocs) # Conversion to polar angles
gl1 = gf.FixedLocator(tlocs) # Positions
tf1 = gf.DictFormatter(dict(zip(tlocs, map(str, rlocs))))
# Standard deviation axis extent
self.smin = 0
self.smax = 1.6 * self.STD
gh = fa.GridHelperCurveLinear(tr,extremes=(0,(np.pi/2),self.smin,self.smax),grid_locator1=gl1,tick_formatter1=tf1,)
if fig is None:
fig = plt.figure()
ax = fa.FloatingSubplot(fig, rect, grid_helper=gh)
fig.add_subplot(ax)
# Angle axis
ax.axis['top'].set_axis_direction('bottom')
ax.axis['top'].label.set_text("Correlation coefficient")
ax.axis['top'].toggle(ticklabels=True, label=True)
ax.axis['top'].major_ticklabels.set_axis_direction('top')
ax.axis['top'].label.set_axis_direction('top')
# X axis
ax.axis['left'].set_axis_direction('bottom')
ax.axis['left'].label.set_text("Standard deviation")
ax.axis['left'].toggle(ticklabels=True, label=True)
ax.axis['left'].major_ticklabels.set_axis_direction('bottom')
ax.axis['left'].label.set_axis_direction('bottom')
# Y axis
ax.axis['right'].set_axis_direction('top')
ax.axis['right'].label.set_text("Standard deviation")
ax.axis['right'].toggle(ticklabels=True, label=True)
ax.axis['right'].major_ticklabels.set_axis_direction('left')
ax.axis['right'].label.set_axis_direction('top')
# Useless
ax.axis['bottom'].set_visible(False)
# Contours along standard deviations
ax.grid()
self._ax = ax # Graphical axes
self.ax = ax.get_aux_axes(tr) # Polar coordinates
# Add reference point and STD contour
l , = self.ax.plot([0], self.STD, 'k*', ls='', ms=12, label=label)
l1 , = self.ax.plot([0], self.STD, 'k*', ls='', ms=12, label=label)
t = np.linspace(0, (np.pi / 2.0))
t1 = np.linspace(0, (np.pi / 2.0))
r = np.zeros_like(t) + self.STD
r1 = np.zeros_like(t) + self.STD
self.ax.plot(t, r, 'k--', label='_')
# Collect sample points for latter use (e.g. legend)
self.samplePoints = [l]
self.samplePoints = [l1]
def add_sample(self,STD,r,*args,**kwargs):
l,= self.ax.plot(np.arccos(r), STD, *args, **kwargs) # (theta, radius)
self.samplePoints.append(l)
return l
def add_sample(self,STD,r1,*args,**kwargs):
l1,= self.ax.plot(np.arccos(r1), STD, *args, **kwargs) # (theta, radius)
self.samplePoints.append(l1)
return l1
def add_contours(self,levels=5,**kwargs):
rs, ts = np.meshgrid(np.linspace(self.smin, self.smax), np.linspace(0, (np.pi / 2.0)))
RMSE=np.sqrt(np.power(self.STD, 2) + np.power(rs, 2) - (2.0 * self.STD * rs *np.cos(ts)))
contours = self.ax.contour(ts, rs, RMSE, levels, **kwargs)
return contours
def srl(obsSTD, s, s1, r, r1, l, l1, fname):
fig=plt.figure(figsize=(8,8))
dia=TaylorDiagram(obsSTD, fig=fig, rect=111, label='ref')
plt.clabel(dia.add_contours(colors='#808080'), inline=1, fontsize=10)
srlc = zip(s, r, l)
srlc1 = zip(s1, r1, l1)
for i in srlc:
dia.add_sample(i[0], i[1], label=i[2], marker='o',mec = 'red', mfc = 'none', mew=1.6)
for i in srlc1:
dia.add_sample(i[0], i[1], label=i[2], marker='^', mec = 'blue', mfc = 'none', mew=1.6)
spl = [p.get_label() for p in dia.samplePoints]
fig.legend(dia.samplePoints, spl, numpoints=1, prop=dict(size='small'), loc=[0.7,0.3])
plt.show()
obsSTD = 1
s = [1.275412605,1.391302157,1.424314937]
s1 = [0.980035327,0.997244197, 1.003002031]
r = [0.572272,0.533529,0.477572]
r1 = [0.82,0.72,0.8]
l = ['A', 'B','C']
l1 = ['A', 'B','C']
fname = 'TaylorDiagram.jpg'
srl(obsSTD, s, s1, r, r1, l,l1, fname)
And give a result that: taylor_dig that the red is before processing, the blue for after one. How could I label the position of A, B, C models as 1,2,3 (enumerate) with the position of before and after processing as well as legends without repeating?

Animating Damped Oscillator

I was able to simulate a mass-spring system under damped oscillations. However, I wanted to add a subplot of position vs time and another subplot velocity vs position (phase path) so that I will be having three animations. How can I add them? The source code that I used is shown below
Update: I tried adding the first subplot position vs time but I cannot get the desired curve for it.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
#Constants
w = np.pi #angular frequency
b = np.pi * 0.2 #damping parameter
#Function that implements rk4 integration
def rk4(t, x, f, dt):
dx1 = f(t, x)*dt
dx2 = f(t+0.5*dt, x+0.5*dx1)*dt
dx3 = f(t+0.5*dt, x+0.5*dx2)*dt
dx4 = f(t+dt, x+dx3)*dt
return x+dx1/6.0+dx2/3.0+dx3/3.0+dx4/6.0
#Function that returns dX/dt for the linearly damped oscillator
def dXdt(t, X):
x = X[0]
vx = X[1]
ax = -2*b*vx - w**2*x
return np.array([vx, ax])
#Initialize Variables
x0 = 5.0 #Initial x position
vx0 = 1.0 #Initial x Velocity
#Setting time array for graph visualization
dt = 0.05
N = 200
x = np.zeros(N)
vx = np.zeros(N)
y = []
# integrate equations of motion using rk4;
# X is a vector that contains the positions and velocities being integrated
X = np.array([x0, vx0])
for i in range(N):
x[i] = X[0]
vx[i] = X[1]
y.append(0)
# update the vector X to the next time step
X = rk4(i*dt, X, dXdt, dt)
fig, (ax1, ax2) = plt.subplots(2,1, figsize=(8,6))
fig.suptitle(r' Damped Oscillation with $\beta$$\approx$' + str(round(b,2)) + r' and $\omega$$\approx$'
+ str(round(w,2)), fontsize=16)
line1, = ax1.plot([], [], lw=10,c="blue",ls="-",ms=50,marker="s",mfc="gray",fillstyle="none",mec="black",markevery=2)
line2, = ax2.plot([], [], lw=2, color='r')
time_template = '\nTime = %.1fs'
time_text = ax1.text(0.1, 0.9, '', transform=ax1.transAxes)
for ax in [ax1, ax2]:
ax1.set_xlim(1.2*min(x), 1.2*max(x))
ax2.set_ylim(1.2*min(x), 1.2*max(x),1)
ax2.set_xlim(0, N*dt)
ax1.set_yticklabels([])
def init():
line1.set_data([], [])
line2.set_data([], [])
time_text.set_text('')
return line1, line2, time_text
def animate(i):
thisx1 = [x[i],0]
thisy1 = [y[i],0]
thisx2 = [i*dt,0]
thisy2 = [x[i],0]
line1.set_data(thisx1, thisy1)
line2.set_data(thisx2, thisy2)
time_text.set_text(time_template % (i*dt))
return line1, line2, time_text
ani = animation.FuncAnimation(fig, animate, np.arange(1, N),
interval=50, blit=True, init_func=init,repeat=False)
After some minors changes to your initial code, the most noteworthy being:
thisx1 = [x[i],0]
thisy1 = [y[i],0]
thisx2 = [i*dt,0]
thisy2 = [x[i],0]
line1.set_data(thisx1, thisy1)
line2.set_data(thisx2, thisy2)
# should be written like this
line1.set_data([x[i],0], [y[i],0])
line2.set_data(t[:i], x[:i])
line3.set_data(x[:i], vx[:i])
The working version, with phase space plot in green, is as follows:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
#Constants
w = np.pi #angular frequency
b = np.pi * 0.2 #damping parameter
#Function that implements rk4 integration
def rk4(t, x, f, dt):
dx1 = f(t, x)*dt
dx2 = f(t+0.5*dt, x+0.5*dx1)*dt
dx3 = f(t+0.5*dt, x+0.5*dx2)*dt
dx4 = f(t+dt, x+dx3)*dt
return x+dx1/6.0+dx2/3.0+dx3/3.0+dx4/6.0
#Function that returns dX/dt for the linearly damped oscillator
def dXdt(t, X):
x = X[0]
vx = X[1]
ax = -2*b*vx - w**2*x
return np.array([vx, ax])
#Initialize Variables
x0 = 5.0 #Initial x position
vx0 = 1.0 #Initial x Velocity
#Setting time array for graph visualization
dt = 0.05
N = 200
t = np.linspace(0,N*dt,N,endpoint=False)
x = np.zeros(N)
vx = np.zeros(N)
y = np.zeros(N)
# integrate equations of motion using rk4;
# X is a vector that contains the positions and velocities being integrated
X = np.array([x0, vx0])
for i in range(N):
x[i] = X[0]
vx[i] = X[1]
# update the vector X to the next time step
X = rk4(i*dt, X, dXdt, dt)
fig, (ax1, ax2, ax3) = plt.subplots(3,1, figsize=(8,12))
fig.suptitle(r' Damped Oscillation with $\beta$$\approx$' + str(round(b,2)) + r' and $\omega$$\approx$'
+ str(round(w,2)), fontsize=16)
line1, = ax1.plot([], [], lw=10,c="blue",ls="-",ms=50,marker="s",mfc="gray",fillstyle="none",mec="black",markevery=2)
line2, = ax2.plot([], [], lw=1, color='r')
line3, = ax3.plot([], [], lw=1, color='g')
time_template = '\nTime = %.1fs'
time_text = ax1.text(0.1, 0.9, '', transform=ax1.transAxes)
ax1.set_xlim(1.2*min(x), 1.2*max(x))
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax2.set_ylim(1.2*min(x), 1.2*max(x),1)
ax2.set_xlim(0, N*dt)
ax2.set_xlabel('t')
ax2.set_ylabel('x')
ax3.set_xlim(1.2*min(x), 1.2*max(x),1)
ax3.set_ylim(1.2*min(vx), 1.2*max(vx),1)
ax3.set_xlabel('x')
ax3.set_ylabel('vx')
def init():
line1.set_data([], [])
line2.set_data([], [])
line3.set_data([], [])
time_text.set_text('')
return line1, line2, line3, time_text
def animate(i):
line1.set_data([x[i],0], [y[i],0])
line2.set_data(t[:i], x[:i])
line3.set_data(x[:i], vx[:i])
time_text.set_text(time_template % (i*dt))
return line1, line2, line3, time_text
ani = animation.FuncAnimation(fig, animate, np.arange(1, N),
interval=50, blit=True, init_func=init,repeat=False)
ani.save('anim.gif')
and gives:

How do I get this to show the legend on the plot?

I am trying to get this code to show a legend on it, but everything I try is not working. Here is my code. I have tried put.legend() in the past and it has worked for me and I am confused why this is not working.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
#declaring my plot
fig1 = plt.figure()
#declaring xvalues
xes = np.arange(-10, 10, 0.01)
xlen = len(xes)
#zeros for yvalues along the axis
yes = np.zeros(xlen)
#declaring my variables
Efieldx = np.zeros((xlen, 1))
Efieldy = np.zeros((xlen, 1))
#locations of my two particles
p1x = 0;
p1y = 1;
p2x = 0;
p2y = -1
q = 1;
Efieldx1 = q/((xes-p1x)*(xes-p1x) + (yes-p1y)*(yes-p1y))**(1.5)*(xes-p1x)
Efieldy1 = q/((xes-p1x)*(xes-p1x) + (yes-p1y)*(yes-p1y))**(1.5)*(yes-p1y)
Efieldx2 = q/((xes-p2x)*(xes-p2x) + (yes-p2y)*(yes-p2y))**(1.5)*(xes-p2x)
Efieldy2 = q/((xes-p1x)*(xes-p1x) + (yes-p1y)*(yes-p1y))**(1.5)*(yes-p2y)
Efieldx = Efieldx1 + Efieldx2
Efieldy = Efieldy1 + Efieldy2
#Efieldx = -1/(xs * xs + ys * ys)^(0.5)
#let's define a function instead:
def f_Efield(q, x, y, xs, ys):
Ex = q*((xs-x)*(xs-x) + (ys-y)*(ys-y))**(-1.5)*(xs-x)
Ey = q/((xs-x)*(xs-x) + (ys-y)*(ys-y))**(1.5)*(ys-y)
return Ex, Ey
#using my new function
Exhere, Eyhere = f_Efield(2, 0, 0,xes, yes)
#plotting:
l, = plt.plot(xes, Efieldx, 'g-')
l, = plt.plot(xes, Exhere, 'r--')
plt.xlim(-10, 10)
plt.ylim(-2, 2)
plt.xlabel('x')
plt.title('Electric field along x-direction \n Andrew Richardson')
#adding a legend
plt.legend()
#displaying the plot
plt.show()
#saving the plot
fig1.savefig('Efield.pdf')
Exhere, Eyhere = f_Efield(-1, 0, 0, xes, yes)
You need to either specify the label property for your plots or pass handles (optional but recommended) and labels to your call to legend otherwise matplotlib has no way of knowing what text to put in the legend
# Using label kwarg
plt.plot(xes, Efieldx, 'g-', label='Efieldx')
plt.plot(xes, Exhere, 'r--', label='Exhere')
plt.legend()
# Using explicit plot handles and labels
p1 = plt.plot(xes, Efieldx, 'g-')
p2 = plt.plot(xes, Exhere, 'r--')
plt.legend([p1, p2], ['Efieldx', 'Exhere'])
# Using just the labels (not recommended)
plt.plot(xes, Efieldx, 'g-')
plt.plot(xes, Exhere, 'r--')
plt.legend(['Efieldx', 'Exhere'])

Categories