plot multiple curves on same plot inside function

Question

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?

Answer 1

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