python - How to plot same function with many different values in subplots in numpy/matplotlib python?

Question

I have following python code, and would like to:

1. Plot the same function in 1 (only one) figure with many different (lets say 4) 'v0' and 'theta' values, each trajectory in a different color.
2. Make 4 plots in 4 different figures, so that it looks like a square with 4 plots of 4 different 'v0' and 'theta' values
3. Make a widget to vary the v0 and theta values as the user wants with the mouse.

---

import numpy as np 
import scipy.integrate as integrate 
import matplotlib.pyplot as plt 
%matplotlib inline


theta = 45.                   
theta = theta * np.pi/180.   
v0 = 20.0

g = 9.81         
R = 0.035             
m = 0.057         
rho = 1.2041           
C = 0.5                


k = (0.5*np.pi*R**2*C*rho)/m    


x0=0                 
y0=10     
vx0 = v0*np.sin(theta)      
vy0 =
v0*np.cos(theta)     
print(vx0) 
print(vy0)

def f_func(X_vek,time):
f = np.zeros(4)    
f[0] = X_vek[2]    
f[1] = X_vek[3]    
f[2] = - k*(f[0]**2 + f[1]**2)**(0.5)*f[0]         
f[3] = -g - k*(f[0]**2 + f[1]**2)**(0.5)*f[1]      
return f

X0 = [ x0, y0, vx0, vy0]         
t0 = 0. tf = 10  
tau = 0.05    

t = np.arange(t0,tf,tau)   

X = integrate.odeint(f_func,X0,t)    

x = X[:,0]      
y = X[:,1]  
vx = X[:,2]  
vy = X[:,3]

mask = y >= 0    

plt.scatter(x[mask],y[mask]) 
plt.scatter(x[mask],y[mask])
plt.xlabel('x') plt.ylabel('y') plt.show()

I could do point 1 and 2 of my question with changing the values after plotting, then calculate vx0 and vy0 again and then call the integrate function and finally plot again, but that's kinda weird and not clean. Is there any better way to do that? like an array of different v0 and theta values or something?

Thanks!

Answer 1

将您的代码作为一个函数：

def func(theta=45, v0=20):
    theta = theta * np.pi/180.   

    g = 9.81         
    R = 0.035             
    m = 0.057         
    rho = 1.2041           
    C = 0.5                

    k = (0.5*np.pi*R**2*C*rho)/m    

    x0=0                 
    y0=10     
    vx0 = v0*np.sin(theta)      
    vy0 = v0*np.cos(theta)     

    def f_func(X_vek,time):
        f0, f1 = X_vek[2:4].tolist()
        f2 = - k*(f0**2 + f1**2)**(0.5)*f0         
        f3 = -g - k*(f0**2 + f1**2)**(0.5)*f1      
        return [f0, f1, f2, f3]

    X0 = [ x0, y0, vx0, vy0]         
    t0 = 0. 
    tf = 10  
    tau = 0.05    

    t = np.arange(t0,tf,tau)   
    X = integrate.odeint(f_func,X0,t)    

    x = X[:,0]      
    y = X[:,1]  
    vx = X[:,2]  
    vy = X[:,3]
    mask = y >= 0    
    return x[mask], y[mask]

然后你可以用不同的参数绘制它：

plt.plot(*func()) 
plt.plot(*func(theta=30)) 

plt.xlabel('x')
plt.ylabel('y')
plt.show()

我建议你使用 Holoviews 制作动态图：

import holoviews as hv
hv.extension("bokeh")

hv.DynamicMap(
    lambda theta, v0:hv.Curve(func(theta, v0)).redim.range(x=(0, 50), y=(0, 50)), 
    kdims=[hv.Dimension("theta", range=(0, 80), default=40), 
           hv.Dimension("v0", range=(1, 40), default=20)])

结果如下：