[Matrix-SIG] Polynomial root finding and handy ndgrid construction.

[Travis Oliphant]

Here are some handy routines I've been using that others may have interest
in:

A polynomial root finder.

def roots(p):
        """r = roots(p) returns the roots of the polynomial defined in the
        sequence p: p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1] * x + p[n]
        """
        p = asarray(p)
        N = len(p.shape)
        if (len(p.shape) > 1):
                raise ValueError, "Input must be 1-D."

        nz = nonzero(p)
        if len(nz) == 0 :
                return zeros(p.shape)

        # Strip leading zeros and trailing zeros
        # Save trailing zeros as roots
        N = p.shape[0]
        p = p[nz[0]:nz[-1]+1]
        r = zeros(N-nz[-1]-1)

        # Form companion matrix (eigenvalues are roots)
        n = p.shape[0]
        if (n > 1):
                a = diag(ones(p.shape[0]-2),-1)
                a[0,:] = -p[1:]/p[0]
                r = concatenate((r,eig(a)[0]))

        return r


Here's a class that makes constructing N-D grids very easy from an
interactive session.  It overloads the slice selection mechanism so that
the output is the equivalent of what MATLAB's ndgrid function produces.

class NDgrid:
    def __getitem__(self,key):
        try:
	    size = []
	    for k in range(len(key)):
	        step = key[k].step
                if step == None:
                    step = 1.0
                size.append((key[k].stop - key[k].start)/(step*1.0))
	    nn = indices(size,'d')
	    for k in range(len(size)):
                step = key[k].step
                if step == None:
                    step = 1.0
                nn[k] = (nn[k]+key[k].start/(1.0*step))*step
	    return nn
        except:
            return arange(key.start,key.stop,key.step)
	    
    def __getslice__(self,i,j):
        return arange(i,j)


Usage:

>>> nd = NDgrid()   # I put this in my .pythonrc file

>>> nd[0:4:0.5]
array([ 0. ,  0.5,  1. ,  1.5,  2. ,  2.5,  3. ,  3.5])

>>> nd[0:4:0.5,0:5]
array([[[ 0. ,  0. ,  0. ,  0. ,  0. ],
        [ 0.5,  0.5,  0.5,  0.5,  0.5],
        [ 1. ,  1. ,  1. ,  1. ,  1. ],
        [ 1.5,  1.5,  1.5,  1.5,  1.5],
        [ 2. ,  2. ,  2. ,  2. ,  2. ],
        [ 2.5,  2.5,  2.5,  2.5,  2.5],
        [ 3. ,  3. ,  3. ,  3. ,  3. ],
        [ 3.5,  3.5,  3.5,  3.5,  3.5]],
       [[ 0. ,  1. ,  2. ,  3. ,  4. ],
        [ 0. ,  1. ,  2. ,  3. ,  4. ],
        [ 0. ,  1. ,  2. ,  3. ,  4. ],
        [ 0. ,  1. ,  2. ,  3. ,  4. ],
        [ 0. ,  1. ,  2. ,  3. ,  4. ],
        [ 0. ,  1. ,  2. ,  3. ,  4. ],
        [ 0. ,  1. ,  2. ,  3. ,  4. ],
        [ 0. ,  1. ,  2. ,  3. ,  4. ]]])


Regards,

Travis