My following code can't add a common title, common legend also a common x and y label. When I add title, legend in between the for loop it works but its not working for the whole figure?
I want to add a common title, common legend, and shared x and y label in this plot. What is the best way to do that?
My code
#position vs Time
h = 0.05
nrow, ncol, count = 2,4,1
plt.figure(figsize=(20,10))
plt.title("Common Title") #does not working
with plt.style.context('ggplot'):
for inv in range(8):
plt.subplot(nrow, ncol, count)
x = [1]
v = [0]
rk4_x = [1] #position
rk4_v = [0] #velocity
t = np.arange(0, 30, 0.1)
for i in range(len(t)-1):
k1 = h*v[i]
l1 = h*-(x[i])
k2 = h*(v[i]+l1)
l2 = h*-(x[i]+k1)
x_value = x[i]+0.5*(k1+k2)
x.append(x_value)
v_value = v[i]+0.5*(l1+l2)
v.append(v_value)
k1 = h*rk4_v[i]
l1 = h*-(rk4_x[i])
k2 = h*(rk4_v[i]+0.5*l1)
l2 = h*-(rk4_x[i]+0.5*k1)
k3 = h*(rk4_v[i]+0.5*l2)
l3 = h*-(rk4_x[i]+0.5*k2)
k4 = h*(rk4_v[i]+l3)
l4 = h*-(rk4_x[i]+k3)
x_value = rk4_x[i]+(1/6)*(k1+2*k2+2*k3+k4)
rk4_x.append(x_value)
v_value = rk4_v[i]+(1/6)*(l1+2*l2+2*l3+l4)
rk4_v.append(v_value)
# with plt.style.context('ggplot'):
# plt.figure()
plt.plot(t, x,'r', label= "Position vs Time RK2")
plt.plot(t, rk4_x, 'k--', label = "Position vs Time RK4")
plt.ylim(-1.5, 1.5)
h = round(h,3)
plt.title("h="+str(h), fontsize= '11')
plt.grid('on')
plt.tick_params(labelcolor='none', which='both', top=False, bottom=False, left=False, right=False)
plt.xlabel('Time')
plt.ylabel('Position')
plt.legend()
h = h + 0.05
count = count+1
Use subplots with fig, axs. See the document on matplotlib.pyplot.subplots
And you should use Numpy array not the python list to speed up the numerical calculation. (Not this example, but generally).
Also, If the amount of calculation increases, it seems necessary to make the formula simple. Like x[i+1] = (1 - 0.5 * h**2) * x[i]+ h * v[i]
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
plt.style.use('ggplot')
def get_x_v(h, t):
x = np.ones_like(t)
v = np.zeros_like(t)
rk4_x = np.ones_like(t)
rk4_v = np.zeros_like(t)
for i, _ in enumerate(t[:-1]):
k1 = h * (v[i])
l1 = h * (x[i]) * -1
k2 = h * (v[i] + l1)
l2 = h * (x[i] + k1) * -1
x[i+1] = x[i] + 0.5 * (k1 + k2)
v[i+1] = v[i] + 0.5 * (l1 + l2)
rk1 = h * (rk4_v[i])
rl1 = h * (rk4_x[i]) * -1
rk2 = h * (rk4_v[i] + 0.5 * rl1)
rl2 = h * (rk4_x[i] + 0.5 * rk1) * -1
rk3 = h * (rk4_v[i] + 0.5 * rl2)
rl3 = h * (rk4_x[i] + 0.5 * rk2) * -1
rk4 = h * (rk4_v[i] + rl3)
rl4 = h * (rk4_x[i] + rk3) * -1
rk4_x[i+1] = rk4_x[i]+(1/6)*(rk1+2*rk2+2*rk3+rk4)
rk4_v[i+1] = rk4_v[i]+(1/6)*(rl1+2*rl2+2*rl3+rl4)
return x, v, rk4_x, rk4_v
h = 0.05
ncol, nrow = 4, 2
fig, axs = plt.subplots(nrows=nrow, ncols=ncol, figsize=(20, 10))
fig.suptitle('Common Title', y=0.94, fontsize=20)
fig.supxlabel('Common X Label', y=0.05, fontsize=14)
t = np.arange(0, 30, 0.1)
for ax in axs.flat:
x, v, rk4_x, rk_v = get_x_v(h, t)
ax.plot(t, x, 'r', label="Position vs Time RK2")
ax.plot(t, rk4_x, 'k--', label="Position vs Time RK4")
ax.set_ylim(-1.5, 1.5)
ax.set_title('h={0}'.format(str(round(h,3))), fontsize='11')
ax.grid('on')
ax.tick_params(labelcolor='none', which='both',
top=False, bottom=False, left=False, right=False)
ax.set_xlabel('Time')
ax.set_ylabel('Position')
ax.legend(loc=1)
ax.set_facecolor('gray')
h += 0.05
plt.savefig('test.png')
Related
Trying to plot two separate animations, i.e. in different windows as separate figures. Running this code for me rightly creates two windows, but animates the data on the second figure at the same time. Closing figure 1 results in only the intended data for figure 2 being animated, removing the overlap from the data intended for figure 1. Closing figure 2 results in only the intended data for figure 1 being animated, removing the overlap from the data intended for figure 2.
Minimum code below:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
dx, dv, N, Nb, decp = 2, 1.5, 100, 12, int(1)
Pd = np.zeros([N + 1, 2 * Nb])
Vd = np.zeros([N + 1, 2 * Nb])
Pd[:, 1] = 4
Vd[:, 3] = 2
t = np.zeros(N + 1)
t[0] = 0
for i in range(0, N):
t[i + 1] = (i + 1) * 0.1
Px = []
for i in range(0, (2 * Nb)):
PX = dx * (-Nb + i) / 4
Px.append(PX)
lblx = []
for i in range(0, int((Nb / 2) + 1)):
if i == (Nb / 4):
LBL = r"$\mu_x$"
lblx.append(LBL)
else:
LBL = r"${0}\sigma_x$".format(-(Nb / 4) + i)
lblx.append(LBL)
Pv = []
for i in range(0, (2 * Nb)):
PV = dv * (-Nb + i) / 4
Pv.append(PV)
lblv = []
for i in range(0, int((Nb / 2) + 1)):
if i == (Nb / 4):
LBL = r"$\mu_v$"
lblv.append(LBL)
else:
LBL = r"${0}\sigma_v$".format(-(Nb / 4) + i)
lblv.append(LBL)
fig1 = plt.figure(figsize=(8,6))
def animatex(i):
fig1.clear()
plt.bar(Px, Pd[i, :], width = dx / 4, align = 'edge', color = 'b', \
label = 't = {} seconds'.format(round(t[i], decp)))
s_ticks = np.arange(-3 * dx, (3 + 1) * dx, dx)
plt.xticks(s_ticks, lblx)
plt.ylim(0, np.max(Pd))
plt.xlim(-3 * dx, 3 * dx)
plt.legend()
plt.draw()
anix = FuncAnimation(fig1, animatex, repeat = True, interval = 200, frames = N + 1)
fig2 = plt.figure(figsize=(8,6))
def animatev(i):
fig2.clear()
plt.bar(Pv, Vd[i, :], width = dv / 4, align = 'edge', color = 'b', \
label = 't = {} seconds'.format(round(t[i], decp)))
s_ticks = np.arange(-3 * dv, (3 + 1) * dv, dv)
plt.xticks(s_ticks, lblv)
plt.ylim(0, np.max(Vd))
plt.xlim(-3 * dv, 3 * dv)
plt.legend()
plt.draw()
aniv = FuncAnimation(fig2, animatev, repeat = True, interval = 200, frames = N + 1)
plt.show()
As is probably clear, they are two bar plots, with different vertical and horizontal dimensions. I've seen some solutions for these kinds of problems where the data shares an axis through a shared variable, but here they are not (as can be seen).
For this minimum code, the solution involves having the two bars, one in Pd and the other in Vd, being on their respective intended figures, not both on the second figure.
Let me know if there are any issues with the information here i.e. minimal code requirements not met, more information etc. and I will update.
Ignore any wayward writing style, it is not relevant.
Simplifying your code:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
dx, dv, N, Nb, decp = 2, 1.5, 10, 12, int(1)
Px = np.arange(Nb)
Pd = np.random.randn(N, Nb)
Vd = np.random.randn(N, Nb)
fig1, ax1 = plt.subplots(figsize=(8, 6))
def animatex(i):
ax1.clear()
ax1.bar(Px, Pd[i, :], width=dx / 4, align='edge', color='b')
anix = FuncAnimation(fig1, animatex, repeat=True, interval=200, frames=N)
fig2, ax2 = plt.subplots(figsize=(8, 6))
def animatev(i):
ax2.clear()
ax2.bar(Px, Vd[i, :], width = dv / 4, align='edge', color='b')
aniv = FuncAnimation(fig2, animatev, repeat=True, interval=200, frames=N)
plt.show()
works fine for me. You can add the esthetic/data details back in...
I made an animation of an orbit that plots an orbit using an initial altitude and velocity.
Here is the code:
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import ode
import matplotlib.animation as animation
plt.style.use('dark_background')
mu_earth = 398600
# Planet radius, km
R = 6371
fig = plt.figure()
ax1 = fig.add_subplot(111)
orbit_alt = 500
v_mag = 7.7 # km/s
# time step (delta-time)
dt = 10
satellite_size = 100
r_mag = R + orbit_alt
v_esc = np.sqrt(2*mu_earth*1e9/(R*1000 + orbit_alt*1000))
earth = plt.Circle((0, 0), 6371, color='blue')
ax1.add_artist(earth)
def diff_eq(t, y):
rx, ry, vx, vy = y # state vectors
r = np.array([rx, ry])
radial_vector = np.linalg.norm(r)
ax, ay = -r * mu_earth / radial_vector ** 3
return [vx, vy, ax, ay]
# Initial states
r0 = [r_mag, 0]
v0 = [0, v_mag]
specific_energy = (v_mag * 1000 * v_mag * 1000) / 2 - ((mu_earth * 1e9) / ((R + orbit_alt) * 1000))
semi_major = -mu_earth * 1e9 / (2 * specific_energy)
orbit_period = 2 * np.pi * np.sqrt(semi_major ** 3 / (mu_earth * 1e9))
if v_mag * 1000 < v_esc:
tspan = 5 * orbit_period
else:
tspan = 1500000
eccen = ((v_mag * 1000) ** 2 * (r_mag)*1000)/(mu_earth*1e9) - 1
perigee = (semi_major * (1 - np.abs(eccen)))/1000
apogee = (semi_major * (1 + np.abs(eccen)))/1000
n_steps = int(np.ceil(tspan/dt))
ts = np.zeros((n_steps, 1))
ys = np.zeros((n_steps, 4))
ys0 = r0 + v0
ts[0] = 0
ys[0] = ys0
solver = ode(diff_eq)
solver.set_integrator('lsoda')
solver.set_initial_value(ys0, 0)
rs = ys[:, :2]
def animate(i):
solver.integrate(solver.t + dt)
ts[i] = solver.t
ys[i] = solver.y
point = plt.Circle((ys[i][0], ys[i][1]), satellite_size, facecolor=(1, 1, 1))
ax1.add_artist(point)
satellite, = ax1.plot(rs[:, 0], rs[:, 1], color='r', alpha = 1)
return satellite, point
ax1.set_xlim([-apogee, apogee])
ax1.set_ylim([-apogee, apogee])
ax1.set_xlabel('X (km)')
ax1.set_ylabel('Y (km)')
ax1.set_title("ORBIT SIMULATION")
plt.gca().set_aspect('equal', adjustable='box')
ani = animation.FuncAnimation(fig, animate, interval=0.01, blit=True)
plt.show()
And here is a screen shot of the output:
I want to omit the line that goes to the origin (0, 0). That line is a part of the satellite line object. What I believe is happening is that the origin is considered the first point, and the satellite is the second point, and matplotlib is connecting them. So how could that line leading to the middle be omitted?
The main problem is that your data is a full vector of zeros and your code plots the full dataset. Change from rs[:, 0] to rs[:, 0][:i] to only plot the points that have been simulated so far.
Second consideration is that FuncAnimation is usually used to update data already plotted. See small change below in animate(i).
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import ode
import matplotlib.animation as animation
plt.style.use('dark_background')
mu_earth = 398600
# Planet radius, km
R = 6371
fig = plt.figure()
ax1 = fig.add_subplot(111)
orbit_alt = 500
v_mag = 7.7 # km/s
# time step (delta-time)
dt = 10
satellite_size = 100
r_mag = R + orbit_alt
v_esc = np.sqrt(2*mu_earth*1e9/(R*1000 + orbit_alt*1000))
earth = plt.Circle((0, 0), 6371, color='blue')
ax1.add_artist(earth)
def diff_eq(t, y):
rx, ry, vx, vy = y # state vectors
r = np.array([rx, ry])
radial_vector = np.linalg.norm(r)
ax, ay = -r * mu_earth / radial_vector ** 3
return [vx, vy, ax, ay]
# Initial states
r0 = [r_mag, 0]
v0 = [0, v_mag]
specific_energy = (v_mag * 1000 * v_mag * 1000) / 2 - ((mu_earth * 1e9) / ((R + orbit_alt) * 1000))
semi_major = -mu_earth * 1e9 / (2 * specific_energy)
orbit_period = 2 * np.pi * np.sqrt(semi_major ** 3 / (mu_earth * 1e9))
if v_mag * 1000 < v_esc:
tspan = 5 * orbit_period
else:
tspan = 1500000
eccen = ((v_mag * 1000) ** 2 * (r_mag)*1000)/(mu_earth*1e9) - 1
perigee = (semi_major * (1 - np.abs(eccen)))/1000
apogee = (semi_major * (1 + np.abs(eccen)))/1000
n_steps = int(np.ceil(tspan/dt))
ts = np.zeros((n_steps, 1))
ys = np.zeros((n_steps, 4))
ys0 = r0 + v0
ts[0] = 0
ys[0] = ys0
solver = ode(diff_eq)
solver.set_integrator('lsoda')
solver.set_initial_value(ys0, 0)
rs = ys[:, :2]
orbit, = ax1.plot(rs[:, 0], rs[:, 1], color='r', alpha=1)
point = plt.Circle((ys[0][0], ys[0][1]), satellite_size, facecolor=(1, 1, 1))
ax1.add_artist(point)
def animate(i):
solver.integrate(solver.t + dt)
ts[i] = solver.t
ys[i] = solver.y
orbit.set_data(rs[:, 0][:i], rs[:, 1][:i])
point.set_center((ys[i][0], ys[i][1]))
return orbit, point
ax1.set_xlim([-apogee, apogee])
ax1.set_ylim([-apogee, apogee])
ax1.set_xlabel('X (km)')
ax1.set_ylabel('Y (km)')
ax1.set_title("ORBIT SIMULATION")
plt.gca().set_aspect('equal', adjustable='box')
ani = animation.FuncAnimation(fig, animate, interval=10, blit=True)
plt.show()
I'm trying to plot a diagram of Coulomb damping mass-spring model, see image bellow. mi*N is the damping friction force.
Becuase the damping force switches direction depending on the direction of speed vector, we have 2 different equations of motion x(t), I named them x_a(t) and x_b(t). If you look closely on the diagram above, the use of x(t) depends on the period, for the first period which is from 0 to pi/omega_n x_a(t) is used, for the second period which is from pi/omega_n to 2*pi/omega_n we use x_b(t) and so on.
I was thinking to use a piece of code which would go something like
t = i * (pi/omega_n)
for i = 1, 3, 5, 7,...
x_a(t)
for i = 2, 4, 6, 8,...
x_b(t)
But I have no clue how to implement this.
I managed to define the x_a(t) part of the code and plot it with the undampled model, see bellow.
import numpy as np
import matplotlib.pyplot as plt
#constants
k = 2 #(N/m), spring coef
m = 0.04 #(kg), mass
x0 = -0.1 #(m), preload
mi = 0.3 #(), dry dynamic friction coef. ABS-ABS
N = 0.3 #(N), normal contact force
f_tr = mi * N / k #friction force/pring coef - equivalent distance
omega_0 = np.sqrt(k/m)
#time
t = np.linspace(0,5,100)
#undamped model
x_undamp = x0*np.cos(omega_0*t)
dx_undamp = -omega_0*x0*np.sin(omega_0*t)
#damped model
x_damp = (x0+f_tr)*np.cos(omega_0*t)-f_tr
dx_damp = -omega_0*(x0+f_tr)*np.sin(omega_0*t)
#time to x=0
t0 = np.arccos(f_tr/(x0+f_tr))/omega_0
print t0
#plotting
fig, (ax1, ax2) = plt.subplots(1, 2)
fig.suptitle('position on left, velocity on right')
ax1.plot(t, x_undamp,'r', label='x_undamp')
ax1.plot(t, x_damp, 'b', label = 'x_damp')
ax2.plot(t, dx_undamp, 'r', label = 'dx_undamp')
ax2.plot(t, dx_damp, 'b', label = 'dx_damp')
#grids, titles, legends, axis labels
ax1.grid()
ax2.grid()
ax1.set_title('Position vs time')
ax2.set_title('Velocity vs time')
ax1.legend()
ax2.legend()
ax1.set_xlabel('t(s)')
ax1.set_ylabel('x(m)')
ax2.set_xlabel('t(s)')
ax2.set_ylabel('dx(m/s)')
plt.show()
I'm gonna put plot a diagram of Coulomb damping mass-spring model solution here till someone comes with a better one:
from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt
m = 0.2
k = 2.0
c = 0.1
mus = 0.3
muk = 0.2
g = 9.8
vf = 0.01
v0 = 0.0
t1 = 10
sign = lambda x: np.tanh(100*x)
def Xi(t):
if t < 1 :
return 0
else:
return 1
vXi = np.vectorize(Xi)
def eq(X, t, Xi):
Ff = k * (Xi(t) - X[0])
if np.abs(X[1]) < vf and np.abs(Ff) < mus * m * g :
Ff = k * (Xi(t) - X[0])
else:
Ff = sign(X[1]) * muk * m * g
d2x = (k * (Xi(t) - X[0]) - Ff) / m
return [X[1], d2x]
t = np.linspace(0, t1, 1000)
X0 = [v0, 0]
sol = odeint(func = eq, y0 = X0, t = t, args = (Xi, ), mxstep = 50000, atol = 1e-5)
plt.plot(t, sol[:, 0], 'r-', label = 'Output (x(t))')
plt.plot(t, vXi(t), 'b-', label = 'Input (xi(t))')
plt.ylabel('values')
plt.xlabel('time')
plt.legend(loc='best')
plt.show()
I don't understand why this code (reference):
from numpy import zeros, linspace
import matplotlib.pyplot as plt
# Time unit: 1 h
beta = 10./(40*8*24)
gamma = 3./(15*24)
dt = 0.1 # 6 min
D = 30 # Simulate for D days
N_t = int(D*24/dt) # Corresponding no of hours
t = linspace(0, N_t*dt, N_t+1)
S = zeros(N_t+1)
I = zeros(N_t+1)
R = zeros(N_t+1)
# Initial condition
S[0] = 50
I[0] = 1
R[0] = 0
# Step equations forward in time
for n in range(N_t):
S[n+1] = S[n] - dt*beta*S[n]*I[n]
I[n+1] = I[n] + dt*beta*S[n]*I[n] - dt*gamma*I[n]
R[n+1] = R[n] + dt*gamma*I[n]
fig = plt.figure()
l1, l2, l3 = plt.plot(t, S, t, I, t, R)
fig.legend((l1, l2, l3), ('S', 'I', 'R'), 'upper left')
plt.xlabel('hours')
plt.show()
doesn't produce the same results as this code I made:
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
beta = 10. / (40 * 8 * 24)
gamma = 3. / (15 * 24)
def flu(y, t):
S = y[0]
P = y[1]
R = y[2]
S = - beta * S * P
P = beta * S * P - gamma * P
R = gamma * P
return [S, P, R]
C_I = [50, 1, 0]
t = np.linspace(0, 1000, 1000)
y = odeint(flu, C_I, t)
S = y[:, 0]
P = y[:, 1]
R = y[:, 2]
fig, ax = plt.subplots()
ax.plot(t, S, 'b--', label='S')
ax.plot(t, P, 'r--', label='I')
ax.plot(t, R, 'g--', label='R')
legend = ax.legend(loc='upper right', shadow=True, fontsize='x-large')
legend.get_frame().set_facecolor('#FFFCCC')
plt.show()
I used P instead of I to avoid confusion.
The equations solved with odeint should be the same as the ones provided in the reference link above. And if the equations I use are correct, which I am convinced they are, I don't understand where the mistake(s) lie(s).
Thank you for your help
You set S=y[0] then you set S=- beta * S * P. This overwrites y[0]!!! Similar problems for P and R
Try this:
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
beta = 10. / (40 * 8 * 24)
gamma = 3. / (15 * 24)
def flu(y, t):
S = y[0]
P = y[1]
R = y[2]
dS = - beta * S * P
dP = beta * S * P - gamma * P
dR = gamma * P
return [dS, dP, dR]
C_I = [50, 1, 0]
t = np.linspace(0, 1000, 1000)
y = odeint(flu, C_I, t)
S = y[:, 0]
P = y[:, 1]
R = y[:, 2]
fig, ax = plt.subplots()
ax.plot(t, S, 'b--', label='S')
ax.plot(t, P, 'r--', label='I')
ax.plot(t, R, 'g--', label='R')
legend = ax.legend(loc='upper right', shadow=True, fontsize='x-large')
legend.get_frame().set_facecolor('#FFFCCC')
plt.show()
I'd like to plot pulse propagation in such a way at each step, it plots the pulse shape. In other words, I want a serie of x-z plots, for each values of y. Something like this (without color):
How can I do this using matplotlib (or Mayavi)? Here is what I did so far:
def drawPropagation(beta2, C, z):
""" beta2 in ps / km
C is chirp
z is an array of z positions """
T = numpy.linspace(-10, 10, 100)
sx = T.size
sy = z.size
T = numpy.tile(T, (sy, 1))
z = numpy.tile(z, (sx, 1)).T
U = 1 / numpy.sqrt(1 - 1j*beta2*z * (1 + 1j * C)) * numpy.exp(- 0.5 * (1 + 1j * C) * T * T / (1 - 1j*beta2*z*(1 + 1j*C)))
fig = pyplot.figure()
ax = fig.add_subplot(1,1,1, projection='3d')
surf = ax.plot_wireframe(T, z, abs(U))
Change to:
ax.plot_wireframe(T, z, abs(U), cstride=1000)
and call:
drawPropagation(1.0, 1.0, numpy.linspace(-2, 2, 10))
will create the following graph:
If you need the curve been filled with white color:
import numpy
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import pyplot
from matplotlib.collections import PolyCollection
def drawPropagation(beta2, C, z):
""" beta2 in ps / km
C is chirp
z is an array of z positions """
T = numpy.linspace(-10, 10, 100)
sx = T.size
sy = z.size
T = numpy.tile(T, (sy, 1))
z = numpy.tile(z, (sx, 1)).T
U = 1 / numpy.sqrt(1 - 1j*beta2*z * (1 + 1j * C)) * numpy.exp(- 0.5 * (1 + 1j * C) * T * T / (1 - 1j*beta2*z*(1 + 1j*C)))
fig = pyplot.figure()
ax = fig.add_subplot(1,1,1, projection='3d')
U = numpy.abs(U)
verts = []
for i in xrange(T.shape[0]):
verts.append(zip(T[i, :], U[i, :]))
poly = PolyCollection(verts, facecolors=(1,1,1,1), edgecolors=(0,0,1,1))
ax.add_collection3d(poly, zs=z[:, 0], zdir='y')
ax.set_xlim3d(numpy.min(T), numpy.max(T))
ax.set_ylim3d(numpy.min(z), numpy.max(z))
ax.set_zlim3d(numpy.min(U), numpy.max(U))
drawPropagation(1.0, 1.0, numpy.linspace(-2, 2, 10))
pyplot.show()