```"""Univariate polynomial manipulations with Numpy and Sympy.

This code shows how to use univariate polynomials with Numpy and Sympy, based
on a simple problem that appears often in least-squares fitting contexts.

Consider solving the equation

x^2 sum_i(a_i/(a_i x + b_i)) = k

where a, b are arrays and k is a scalar.  This can be rewritten as

x^2 P(x)/Q(x) = k

and outside of the roots of Q(x), the solutions are given by the roots of the
polynomial R(x):

R(x) := x^2 P(x) - k Q(x)

In this example, we compute this polynomial R(x) using both numpy and sympy,
and then find its roots via numpy (sympy has limited support for numerical root
finding).

Notes:

- We always normalize R(x) so its leading coefficient is 1.

- In newer versions of numpy (after February 2010), a new Polynomial class will
be available that provides more sophisticated behavior than the basic poly1d
shown here.
"""

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------
from __future__ import print_function # for doctests

import numpy as np
import sympy as sym

#-----------------------------------------------------------------------------
# Function definitions
#-----------------------------------------------------------------------------

def symarray(shape, prefix=''):
"""Create a numpy ndarray of symbols (as an object array).

The created symbols are named prefix_i1_i2_...  You should thus provide a
non-empty prefix if you want your symbols to be unique for different output
arrays, as Sympy symbols with identical names are the same object.

Parameters
----------

shape : int or tuple
Shape of the created array.  If an int, the array is one-dimensional; for
more than one dimension the shape must be a tuple.

prefix : string
A prefix prepended to the name of every symbol.

Examples
--------

>>> symarray(3)
array([_0, _1, _2], dtype=object)

If you want multiple symarrays to contain distinct symbols, you *must*
provide unique prefixes:
>>> a = symarray(3)
>>> b = symarray(3)
>>> a[0] is b[0]
True
>>> a = symarray(3, 'a')
>>> b = symarray(3, 'b')
>>> a[0] is b[0]
False

Creating symarrays with a prefix:
>>> symarray(3, 'a')
array([a_0, a_1, a_2], dtype=object)

For more than one dimension, the shape must be given as a tuple:
>>> symarray((2,3), 'a')
array([[a_0_0, a_0_1, a_0_2],
[a_1_0, a_1_1, a_1_2]], dtype=object)
>>> symarray((2,3,2), 'a')
array([[[a_0_0_0, a_0_0_1],
[a_0_1_0, a_0_1_1],
[a_0_2_0, a_0_2_1]],
<BLANKLINE>
[[a_1_0_0, a_1_0_1],
[a_1_1_0, a_1_1_1],
[a_1_2_0, a_1_2_1]]], dtype=object)
"""
arr = np.empty(shape, dtype=object)
for index in np.ndindex(shape):
arr[index] = sym.Symbol('%s_%s' % (prefix, '_'.join(map(str, index))))
return arr

def prod(seq):
"""Multiply the elements of a sequence.

Note that the product accumulator is initialized to 1, so prod([])==1.

Parameters
----------

seq : iterable
Any iterable sequence whose elements can be multiplied.

Returns
-------
prod : scalar
The result of the multiplication of all the elements in the input.

Examples
--------
>>> prod([1,2,3])
6
>>> prod(['a',2,3])
'aaaaaa'
>>> prod([])
1
"""
return reduce(lambda x,y: x*y, seq, 1)

def spoly2npoly1d(poly):
"""Convert a univariate Sympy polynomial to a numpy.poly1d.

The coefficients of the

Parameters
----------
poly : sympy.Poly instance
A univariate sympy polynomial object.

Returns
-------
poly1d : a numpy.poly1d instance

Examples
--------
>>> import sympy
>>> x = sympy.Symbol('x')
>>> p = sympy.Poly(x**2 + 3*x + 1, x)
>>> p_num = spoly2npoly1d(p)
>>> p_num
poly1d([ 1.,  3.,  1.])
"""

if poly.is_multivariate:
e = 'multivariate polynomials can not be converted to poly1d'
raise ValueError(e)
return np.poly1d(np.array(map(float, poly.coeffs), dtype=float))

def sym_rpoly(na=None, nb=None, nk=None, nterms=None):
"""Compute the R polynomial as defined above.

Internally the construction of R is done symbolically. If the numerical
variables (na, nb, nk) were supplied, these values are substituted at the
end and a sympy.Poly object with numerical coefficients is returned.  If
(na, nb, nk) are not given, then nterms *must* be given, and a symbolic

Parameters
----------
na : ndarray, optional
Numerical array of 'a' coefficients.

nb : ndarray, optional
Numerical array of 'b' coefficients.

k : float, optional
Numerical value of k.

nterms : int, optional.
Number of terms. This is only used if na, nb and nk are *not* given, in
which case a symbolic answer is returned with nterms total.

Returns
-------
poly : sympy.Poly instance
A univariate polynomial in x.

Examples
--------
With only nterms, a symbolic polynomial is returned:
>>> sym_rpoly(nterms=1)
Poly(x**2 - k*x - b_0*k/a_0, x)

But if numerical values are supplied, the output polynomial has numerical
values:
>>> sym_rpoly([2], [5], 3)
Poly(x**2 - 3*x - 15/2, x)
>>> sym_rpoly([1, 2], [4, 6], 1)
Poly(x**3 + 3*x**2 - 7/2*x - 6, x)
"""
if na is None and nterms is None:
raise ValueError('You must provide either arrays or nterms')

if nterms is None:
na = np.asarray(na)
nb = np.asarray(nb)
# We have input arrays, return numerical answer
if na.ndim > 1:
raise ValueError('Input arrays must be one-dimensional')
nterms = na.shape
mode = 'num'
else:
# We have only a length, return symbolic answer.
mode = 'sym'

# Create symbolic variables
k, x = sym.symbols('k x')
a = symarray(nterms, 'a')
b = symarray(nterms, 'b')

# Construct polynomial symbolically
t = [ ai/(ai*x+bi) for (ai, bi) in zip(a,b)]
P, Q = sym.fraction(sym.together(sum(t)))
Rs = sym.Poly(x**2*P -k*Q, x)

if mode == 'num':
# Substitute coefficients for numerical values
Rs = Rs.subs([(k, nk)] + zip(a, na) + zip(b, nb))

# Normalize
return Rs

def num_rpoly(a, b, k):
"""Compute the R polynomial as defined above using numpy.

Parameters
----------
na : ndarray

nb : ndarray

k : float

Returns
-------
poly : numpy.poly1d instance
A univariate polynomial in x.

Examples
--------
>>> num_rpoly([2.0], [5.0], 3.0)
poly1d([ 1. , -3. , -7.5])

Printing numpy's poly1d objects gives a more familiar readout:
>>> print(num_rpoly([2.0], [5.0], 3.0))
2
1 x - 3 x - 7.5

>>> num_rpoly([1.0, 2.0], [4.0, 6.0], 1.0)
poly1d([ 1. ,  3. , -3.5, -6. ])
>>> print(num_rpoly([1.0, 2.0], [4.0, 6.0], 1.0))
3     2
1 x + 3 x - 3.5 x - 6
"""

a = np.asarray(a)
b = np.asarray(b)
x2 = np.poly1d([1, 0, 0])
Qr = -b/a
Q = np.poly1d(Qr, r=True)
denoms = [ np.poly1d([1, -di]) for di in Qr ]
P = sum( (prod(dj for (j, dj) in enumerate(denoms) if j != i ))
for i in range(len(a)) )
R = x2*P - k*Q
R /= R.coeffs[0]
return R

def num_rpoly2(a, b, k):
"""Compute the R polynomial as defined above using numpy.

This is just an alternate implementation which keeps the a factors in the
numerators, it should be identical to num_rpoly and is kept here only to
illustrate some aspects of numpy poly1d objects.

Parameters
----------
na : ndarray

nb : ndarray

k : float

Returns
-------
poly : numpy.poly1d instance
A univariate polynomial in x.

Examples
--------
>>> num_rpoly([2.0], [5.0], 3.0)
poly1d([ 1. , -3. , -7.5])

Printing numpy's poly1d objects gives a more familiar readout:
>>> print(num_rpoly([2.0], [5.0], 3.0))
2
1 x - 3 x - 7.5

>>> num_rpoly([1.0, 2.0], [4.0, 6.0], 1.0)
poly1d([ 1. ,  3. , -3.5, -6. ])
>>> print(num_rpoly([1.0, 2.0], [4.0, 6.0], 1.0))
3     2
1 x + 3 x - 3.5 x - 6
"""
a = np.asarray(a)
b = np.asarray(b)
x2 = np.poly1d([1, 0, 0])
# When building the polynomial Q out of roots, by default numpy keeps the
# leading coefficient to 1, so we need to multiply back by the product of
# all the a's
Q = np.poly1d(-b/a, r=True)
Q *= a.prod()
# We need the individual components of the denominator Q
denoms = [ np.poly1d([ai, bi]) for (ai, bi) in zip(a, b) ]
# With these, we can compute the numerator P.  Note that we must call
# float() on each ai, because multiplying a poly1d by a numpy scalar
# produces odd results.  I consider this a bug in numpy.
P  = sum( float(ai)* prod(dj for (j, dj) in enumerate(denoms) if j != i)
for (i, ai) in enumerate(a) )
# Now the computation proceeds identically.
R = x2*P - k*Q
R /= R.coeffs[0]
return R

def test_compare_with_sympy():
"""Use the symbolic solution as a reference for the numpy ones."""
import numpy.testing as npt

a = np.array([3.4, 4.5, 3.2])
b = np.array([2.1, 5.5, 4.5])
k = 0.5

Rn = num_rpoly(a, b, k)
Rn2 = num_rpoly2(a, b, k)
Rs = sym_rpoly(a, b, k)
Rns = spoly2npoly1d(Rs)

# Check that both numerical implementations coincide with the symbolic one
npt.assert_almost_equal(Rn.coeffs, Rns.coeffs, 15)
npt.assert_almost_equal(Rn2.coeffs, Rns.coeffs, 15)

#-----------------------------------------------------------------------------
# Main script
#-----------------------------------------------------------------------------

if __name__ == '__main__':
# Simple illustrative example
a = np.array([3.4, 4.5, 3.2])
b = np.array([2.1, 5.5, 4.5])
k = 0.5

Rn = num_rpoly(a, b, k)
Rn2 = num_rpoly2(a, b, k)
Rs = sym_rpoly(a, b, k)
Rns = spoly2npoly1d(Rs)

print('Numpy, R(x) =', Rn, sep='\n')
print('Sympy, R(x) =', Rs, sep='\n')
roots = Rn.roots
roots.sort()
print()
print('Roots:', roots)
sys.exit(1)

```