plot multiple curves on same plot inside function

user_9106985 0
Viewed 65

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?

matplotlib function python
2 Answers
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.

enter image description here

user_175172860

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:
enter image description here

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