how do i make my function returns multiple values within my for-loop? (Beginner)

user_11682513 0
Viewed 41

I made a function that uses the np.roots() to calculate the roots of a cubic function after I solved for coefficients p1,p2, and p3. I'm trying to make it return multiple roots.

Here is my code:

def Q():
     .
     .
     .
    coeff = [1,p1, p2, p3]
    roots = np.roots(coeff)

    #prints out the real positive root

    for i in range(0, len(roots)):
        if np.isreal(roots[i]):
            return np.real(roots[i])
                 
    return 

root = Q(T1,P1,Tc,Pc)
print(root)

output:

root 1

what I want:

root 1
root 2
.
.
.
root n
jupyter-notebook numpy python
1 Answers

You can't exactly return values in a loop: return stops the execution of the function and returns control to the caller. You have to compute all the stuff you want to return up-front, and return it all at once.

In this case, you can select elements from your array using a mask. This is especially easy because np.isreal is vectorized: it's meant to operate on your whole array at once, not just individual elements. So you can get rid of the entire loop and write

return roots[np.isreal(roots)]

The yield keyword is an alternative to return that does sort of allow you to return multiple values. Using yield turns your function into a generator, which you can iterate over until it stops yielding values. To make a generator, just replace return with yield in your loop.

Both versions of the function will return an iterable rather than a single value. You can then either print in a loop or all at once. The loop looks like this:

for root in Q():
    print(root)

To print all at once, you can join the strings with newlines:

print('\n'.join(str(root) for root in Q()))
user_29887300
Related