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

Make your code as a function:

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]

then you can plot it with different parameters:

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

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

I suggest you use Holoviews to make dynamic graph:

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)])

Here is the result: