plot multiple curves on same plot inside function-CodePudding

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?

CodePudding user response：

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.

CodePudding user response：

Your code has several problems the main being that the return function was inside the loop interrupting it after the first iteration. Imitating your code structure as closely as possible, we can rewrite the code as:

import numpy as np
import matplotlib.pyplot as plt

def test(psi,lam):
    y=[]
    
    for curr_lam in lam:
        sin_psi  = np.sin(psi)
        cos_psi  = np.cos(psi)
        sin_beta = curr_lam*sin_psi
        cos_beta = np.sqrt(1.0 - sin_beta**2)
        ssin_pb  = sin_psi*sin_beta
        val = (curr_lam * (cos_psi/cos_beta)**2 - ssin_pb)/cos_beta   cos_psi
        y.append(val)
        plt.plot(psi, val)   
        
    plt.show()
    return y


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

y = test(psi, lam)
print(y)

Sample output:

As Johan mentioned in the comments, you should also directly iterate over list/arrays. If you need to combine arrays, use

for x1, x2 in zip(arr1, arr2):

If you absolutely need the index value, use

for i, x in enumerate(arr):