Search code examples
pythonnumpypolynomials

Get Polynomial X at Y? (Python 3.10, NumPy)


I'm attempting to calculate all possible real X-values at a certain Y-value from a polynomial given in descending coefficent order, in Python 3.10. I want the resulting X-values to be provided to me in a list.

I've tried using the roots() function of the numpy library, as shown in one of the answers to this post, however it does not appear to work:

import numpy as np
import matplotlib.pyplot as plt

def main():
    coeffs = np.array([1, 2, 2])
    y = 1.5

    polyDataX = np.linspace(-2, 0)
    polyDataY = np.empty(shape = len(polyDataX), dtype = float)

    for i in range(len(polyDataX)):
        polyDataY[i] = coeffs[0] * pow(polyDataX[i], 2) + coeffs[1] * polyDataX[i] + coeffs[2]

    coeffs[-1] -= y
    x = np.roots(coeffs).tolist()

    plt.axhline(y, color = "orange")
    plt.plot(polyDataX, polyDataY, color = "blue")
    plt.title("X = " + str(x))
    plt.show()
    plt.close()
    plt.clf()

if (__name__ == "__main__"):
    main()

In my example above, I have the coefficents of my polynomial stored in the local variable coeffs, in descending order. I then attempt to gather all the X-values at the Y-value of 0.5, stored within the x and y local variables respectivelly. I then display the gathered X-values as the title of the shown plot.

The script above results in the following plot: enter image description here

With the X-values being shown as [-2.0, 0.0], instead of the correct: enter image description here

What is the proper way to get all real X-values of a polynomial at a certain Y-value in Python?

Thanks for reading my post, any guidance is appreciated.


Solution

  • You should be making use of the numpy.polynomial.Polynomial class that was added in numpy v1.4 (more information here). With that class, you can create a polynomial object. To find your solution, you can subtract y from the Polynomial object and then call the roots method. Another nice feature is that you can directly call the object, which you can use to compute polyDataY.

    Just note that the Polynomial class expects the coefficients to be given backward from np.roots, i.e. a quadratic of x^2 + 2x + 2 should have the coefficients (2, 2, 1). To keep things consistent with what you gave, I just pass the reversed coeffs.

    import numpy as np
    from numpy.polynomial import Polynomial
    import matplotlib.pyplot as plt
    
    plt.close("all")
    
    coeffs = np.array([1, 2, 2])
    poly = Polynomial(coeffs[::-1])
    
    polyDataX = np.linspace(-2, 0)
    polyDataY = poly(polyDataX)
    
    y = 1.5
    x = (poly - y).roots()
    plt.axhline(y, color = "orange")
    plt.plot(polyDataX, polyDataY, color = "blue")
    plt.title("X = " + str(x))
    plt.show()