Sign up ×
Stack Overflow is a community of 4.7 million programmers, just like you, helping each other. Join them, it only takes a minute:

I want to plot the two solutions of quadratic equation as a function of a parameter ( function coeff(t) ). I am using function numpy.roots (I am sure that all roots are real in this case), and I am trying to invoke it from within pyplot.plot as below:

import numpy as np
import matplotlib.pyplot as plt

r = 3.74
def coeff(pp):
    return np.array([pp-1,r+1-0.5*pp,-r])

def sroot(t):
    return np.roots(coeff(t))

a = np.linspace(0,0.9,100)

fig = plt.figure()
plt.plot(a,sroot(a)[0,:])
plt.plot(a,sroot(a)[1,:])

plt.show()

I get error message:

  File "quest.py", line 18, in <module>
    plt.plot(a,sroot(a)[0,:])
  File "quest.py", line 10, in sroot
    return np.roots(coeff(t))
  File "/usr/lib64/python2.7/site-packages/numpy/lib/polynomial.py", line 218, in roots
    p = p.astype(float)

I understand that the numpy.roots takes only list of parameters and is unable to recognize a row in array 3xlen(a). Is there a way to do it in one line, preferably inside the pyplot.plot? I would like to avoid using loop.

share|improve this question

1 Answer 1

This is because you transform all of your coefficient at once and try to call the numpy roots solver on all of them. np.roots only accept 1-d array and solves a single polynomial. Here is a script that does what you want:

import numpy as np
import matplotlib.pyplot as plt

# Parameters
r = 3.74
T = np.linspace(0.0,0.9,100)

# Coefficients
C = np.zeros((len(T),3))
C[:,0] = T-1
C[:,1] = r + 1 - 0.5*T
C[:,2] = r

# Roots
R = np.zeros((len(T),2))
for i in range(len(T)):
    R[i] = np.roots(C[i])

# Plot
fig = plt.figure()
plt.plot(T,R[:,0])
plt.plot(T,R[:,1])

plt.show()
share|improve this answer
    
I always thought python to be an elegant language, so I was looking for solutions that are shorter and do not have any for-loops. Any ideas? – Luise Feb 27 at 23:28
    
Yes. Because you've have only second order polynomials, you can directly find the roots (i.e. without using np.roots). – Nicolas Rougier Feb 28 at 5:19

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.