hroots_test, scipy pr 3987

hroots_test.py

import numpy as np
import scipy.special as cephes
import scipy.linalg as linalg
from mpmath import mp

# current scipy routine
def _gen_roots_and_weights(n, mu0, an_func, bn_func, f, df, symmetrize, mu):
    """[x,w] = gen_roots_and_weights(n,an_func,sqrt_bn_func,mu)

    Returns the roots (x) of an nth order orthogonal polynomial,
    and weights (w) to use in appropriate Gaussian quadrature with that
    orthogonal polynomial.

    The polynomials have the recurrence relation
          P_n+1(x) = (x - A_n) P_n(x) - B_n P_n-1(x)

    an_func(n)          should return A_n
    sqrt_bn_func(n)     should return sqrt(B_n)
    mu ( = h_0 )        is the integral of the weight over the orthogonal
                        interval
    """
    k = np.arange(n, dtype='d')
    c = np.zeros((2, n))
    c[0,1:] = bn_func(k[1:])
    c[1,:] = an_func(k)
    x = linalg.eigvals_banded(c, overwrite_a_band=True)

    # improve roots by one application of Newton's method
    y = f(n, x)
    dy = df(n, x)
    x -= y/dy

    fm = f(n-1, x) 
    fm /= np.abs(fm).max()
    dy /= np.abs(dy).max()
    w = 1.0 / (fm * dy)

    if symmetrize:
        w = (w + w[::-1]) / 2
        x = (x - x[::-1]) / 2

    w *= mu0 / w.sum()

    if mu:
        return x, w, mu0
    else:
        return x, w


# current scipy routine
def h_roots(n, mu=False):
    """Gauss-Hermite (physicst's) quadrature

    Computes the sample points and weights for Gauss-Hermite quadrature.
    The sample points are the roots of the `n`th degree Hermite polynomial,
    :math:`H_n(x)`.  These sample points and weights correctly integrate
    polynomials of degree :math:`2*n - 1` or less over the interval
    :math:`[-inf, inf]` with weight function :math:`f(x) = e^{-x^2}`.

    Parameters
    ----------
    n : int
        quadrature order 
    mu : boolean
        If True, return the sum of the weights, optional.

    Returns
    -------
    x : ndarray
        Sample points
    w : ndarray
        Weights
    mu : float
        Sum of the weights
        
    See Also
    --------
    integrate.quadrature
    integrate.fixed_quad
    numpy.polynomial.hermite.hermgauss
    """
    m = int(n)
    if n < 1 or n != m:
        raise ValueError("n must be a positive integer.")

    mu0 = np.sqrt(np.pi)
    an_func = lambda k: 0.0*k
    bn_func = lambda k: np.sqrt(k/2.0)
    f = cephes.eval_hermite
    df = lambda n, x: 2.0 * n * cephes.eval_hermite(n-1, x)
    return _gen_roots_and_weights(m, mu0, an_func, bn_func, f, df, True, mu)

# from scipy PR 3987, essentially.
def h_roots_new(n):
    bn = np.sqrt(np.arange(1, n, dtype='d') / 2.0)
    c  = np.zeros((2, n), 'd', 'F')
    c[0,1:] = bn

    x,v = linalg.eig_banded(c, overwrite_a_band=True)
    w = v[0,:]**2

    w = (w + w[::-1]) / 2
    x = (x - x[::-1]) / 2

    mu0 = np.sqrt(np.pi)
    w *= mu0 / w.sum()

    return x, w

def h_roots_pbdv(n):
    bn = np.sqrt(np.arange(1, n, dtype='d') / 2.0)
    c  = np.zeros((2, n), 'd', 'F')
    c[0,1:] = bn

    x = linalg.eigvals_banded(c, overwrite_a_band=True)

    # improve roots by one application of Newton's method
    
    y, dy = cephes.pbdv(n, np.sqrt(2)*x)
    x -= y/dy

    fm = np.exp(x**2/2)*cephes.pbdv(n-1, np.sqrt(2)*x)[0]
    fm /= np.abs(fm).max()
    w = 1.0 / fm / fm

    w = (w + w[::-1]) / 2
    x = (x - x[::-1]) / 2

    mu0 = np.sqrt(np.pi)
    w *= mu0 / w.sum()

    return x, w

def hfunc(n, x):
    y0 = np.exp(-x**2/2) / np.sqrt(np.sqrt(np.pi))
    y1 = 2*x * np.exp(-x**2/2) / np.sqrt(2*np.sqrt(np.pi))

    if n < 0:
        return 0.0
    elif n == 0:
        return y0
    elif n == 1:
        return y1
    else:
        for k in range(2,n+1):
            y2 = np.sqrt(2.0/k)*x*y1 - np.sqrt((k-1.0)/k)*y0
            y0 = y1
            y1 = y2
        return y2

def h_roots_hf(n):
    bn = np.sqrt(np.arange(1, n, dtype='d') / 2.0)
    c  = np.zeros((2, n), 'd', 'F')
    c[0,1:] = bn

    x = linalg.eigvals_banded(c, overwrite_a_band=True)

    # improve roots by one application of Newton's method
    
    y = hfunc(n, x)
    dy = -x*y + np.sqrt(2*n)*hfunc(n-1,x)
    x -= y/dy

    fm = np.exp(x**2/2) * hfunc(n-1, x)
    fm /= np.abs(fm).max()
    w = 1.0 / fm / fm

    w = (w + w[::-1]) / 2
    x = (x - x[::-1]) / 2

    mu0 = np.sqrt(np.pi)
    w *= mu0 / w.sum()

    return x, w

def h_roots_mpmath(n, prec=False):
# doing it all here is unnecessary and slow...
#    bn = np.arange(1,n)
#    c = np.diag(bn, -1) + np.diag(bn, +1)
#    c = mp.matrix(c)
#
#    # mpmath isn't vectorized..
#    for k in range(n-1):
#        c[k,k+1] = mp.sqrt(c[k,k+1]/2)
#        c[k+1,k] = mp.sqrt(c[k+1,k]/2)
#
#    xm, vm = mp.eigh(c)
#
#    mu = mp.sqrt(mp.pi())
#    x = []
#    w = []
#    for k in range(n):
#        print k
#        x.append(xm[k])
#        w.append(vm[0,k]**2 * mu)
    bn = np.sqrt(np.arange(1, n, dtype='d') / 2.0)
    c  = np.zeros((2, n), 'd', 'F')
    c[0,1:] = bn

    x0 = linalg.eigvals_banded(c, overwrite_a_band=True)

    # now use mpmath from here on.
    x = np.empty(x0.shape, object) 
    for k in range(n):
        x[k] = mp.mpf(x0[k])

    f = np.vectorize(mp.hermite, 'O')
    df = np.vectorize(lambda n, x: mp.hermite(n-1, x) * 2 * n, 'O')
    
    k = 0
    xm1old = f(n, x[-1])
    #print xm1old
    xm1 = 2*xm1old
    #print xm1old != xm1
    while float(xm1) != 0.0 and xm1old != xm1 and k < 25:
        y = f(n, x)
        dy = df(n, x)
        x -= y/dy
        k += 1
        xm1old = xm1
        xm1 = y[-1]

    fm = f(n-1, x) 
    fm /= np.abs(fm).max()
    dy /= np.abs(dy).max()
    w = 1.0 / (fm * dy)

    w = (w + w[::-1]) / 2
    x = (x - x[::-1]) / 2

    mu0 = mp.sqrt(mp.pi())
    w *= mu0 / w.sum()

    if prec:
        return x, w, xm1
    else:
        return x, w

def find_max_n():
    """
    Determine the maximum order of gauss-hermite quadrature for which all of
    the weight are nonzero when represented as doubles.
    """
    # the first and last weight are the smallest.
    n = 350
    dps0 = mp.dps
    
    mp.dps = 400
    x, w, p = h_roots_mpmath(n, True)
    while w[-1] > np.finfo('d').tiny:
        n += 1
        x, w, p = h_roots_mpmath(n, True)
        print n, mp.dps, float(p), float(w[-1])
        while abs(p) > 1e-20:
            mp.dps += 25
            x, w, p = h_roots_mpmath(n, True)
            print n, mp.dps, float(p), float(w[-1])

    mp.dps = dps0
    return n-1
        
#mp.dps = 400
#xs, ws = h_roots(205)
#xn, wn = h_roots_new(205)
#xm, wm = h_roots_mpmath(205)
#xmd = xm.astype('d')
#wmd = wm.astype('d')