plot multiple curves on same plot inside function

February 8, 2022

I have a following function with takes 2 arguments psi,lam and returns 1 array y.

lam=np.arange(0,1,0.1)
psi=np.deg2rad(np.arange(0,361,1))

def test(psi,lam):

    y=[]
    
    for i in range(len(lam)):
        sin_psi  = np.sin(psi)
        cos_psi  = np.cos(psi)
        sin_beta = lam*sin_psi
        cos_beta = np.sqrt(1.0 - sin_beta**2)
        ssin_pb  = sin_psi*sin_beta
        y.append((lam*(cos_psi/cos_beta)**2 - ssin_pb)/cos_beta + cos_psi)
        plt.plot(psi,y[i])
    
return y

I would like the function to return range(len(lam))=10 plots of y on the vertical axis against psi on x axis.

However, it seems to be only plotting the same curve multiple times. Not sure what I am missing?

>Solution :

import matplotlib.pyplot as plt
import numpy as np

lam=np.arange(0,1,0.1)
psi=np.deg2rad(np.arange(0,361,1))

def test(angle,var):
    sin_psi  = np.sin(psi)
    cos_psi  = np.cos(psi)
    sin_beta = var*sin_psi
    cos_beta = np.sqrt(1.0 - sin_beta**2)
    ssin_pb  = sin_psi*sin_beta
    return ((var*(cos_psi/cos_beta)**2 - ssin_pb)/cos_beta + cos_psi)

for i in lam:
    plt.plot(psi,test(psi,i))
plt.show()

I moved the variable outside of the function, this way you may also use it for other cases. The only other thing is that you should call plt.show() after you’re done drawing.