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?

Maybe you want to use sin_beta = lam[i]*sin_psi? And also y.append((lam[i[*(cos_psi/cos_beta)**2 ...))? A more "pythonic" approach would use for lam_i in lam: — JohanC
– JohanC, Commented Feb 8, 2022 at 9:40

Tempman383838 · Accepted Answer · 2022-02-08 09:56:03Z

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.

Mr. T · Accepted Answer · 2022-02-08 10:17:14Z

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

Although this is closer to the example from the question, wouldn't it still be a better option to extract the lam iteration and get rid of the extra list? I just rewrote the program to something that I would consider more "useful", but if the desired result was both the plot and the list, your answer would be more correct.
@Tempman383838 Well, we don't know. The OP hasn't specified what the expected outcome is. We both had a guess what they intended to achieve.

Collectives™ on Stack Overflow

plot multiple curves on same plot inside function

2 Answers 2

Comments

2 Comments

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

2 Comments

Related