roots_sh_jacobi using mpmath ref scipy ticket 12425

jacobi.py

import numpy as np
from mpmath import mp
from scipy import special

def an_func(a, b, k):
    if a + b == 0:
        if k == 0:
            return (b-a)/(2+a+b)
        else:
            return 0
    else:
        if k == 0:
            return (b-a)/(2+a+b)
        else:
            return (b*b - a*a) / ((2*k+a+b)*(2*k+a+b+2))

def bn_func(a, b, k):
    ret = 2 / (2*k+a+b)*mp.sqrt((k+a)*(k+b)/(2*k+a+b+1))
    if k == 1:
        return ret
    else:
        return ret * mp.sqrt(k*(k+a+b)/(2*k+a+b-1))


def roots_sh_jacobi_mpmath(n, p, q, dps=None):
    if dps is not None:
        mp.dps = dps

    # map p and q to a and b

    a = p - q
    b = q -1

    # build the needed matrix
    c = mp.matrix(n)
    for k in range(n):
        c[k, k] = an_func(a, b, k)
        if k > 0:
            bn = bn_func(a, b, k)
            c[k-1,k] = bn
            c[k,k-1] = bn

    # find estimates of the roots
    vals = mp.eigsy(c, eigvals_only=True)

    # clean up roots with find root
    x = np.zeros(n, object)
    f = lambda x: mp.jacobi(n, a, b, x)
    for k, root in enumerate(vals):
        x[k] = mp.findroot(f, root)


    # now find the weights
    mu0 = mp.power(2, a+b+1) * mp.beta(a+1, b+1)

    fm = np.zeros(n, object)
    dy = np.zeros(n, object)

    for k in range(n):
        fm[k] = mp.jacobi(n-1, a, b, x[k])
        dy[k] =  (n + a + b + 1) * mp.jacobi(n-1, a+1, b+1, x[k]) / 2

    fm /= abs(fm).max()
    dy /= abs(dy).max()

    w = 1 / (fm * dy)
    w *= mu0 / sum(w)

    # convert back to shifted
    x = (x + 1)/2
    w /= mp.power(2,p)
    mu0 /= mp.power(2,p)

    return x, w, mu0

def roots_jacobi(n, alpha, beta, mu=False):
    m = int(n)
    if n < 1 or n != m:
        raise ValueError("n must be a positive integer.")
    if alpha <= -1 or beta <= -1:
        raise ValueError("alpha and beta must be greater than -1.")

    if alpha == 0.0 and beta == 0.0:
        return special.roots_legendre(m, mu)
    if alpha == beta:
        return special.roots_gegenbauer(m, alpha+0.5, mu)

    mu0 = 1.0 # 2.0**(alpha+beta+1)*cephes.beta(alpha+1, beta+1)
    a = alpha
    b = beta
    if a + b == 0.0:
        an_func = lambda k: np.where(k == 0, (b-a)/(2+a+b), 0.0)
    else:
        an_func = lambda k: np.where(k == 0, (b-a)/(2+a+b),
                  (b*b - a*a) / ((2.0*k+a+b)*(2.0*k+a+b+2)))

    bn_func = lambda k: 2.0 / (2.0*k+a+b)*np.sqrt((k+a)*(k+b) / (2*k+a+b+1)) \
              * np.where(k == 1, 1.0, np.sqrt(k*(k+a+b) / (2.0*k+a+b-1)))

    f = lambda n, x: special.eval_jacobi(n, a, b, x)
    df = lambda n, x: 0.5 * (n + a + b + 1) \
                      * special.eval_jacobi(n-1, a+1, b+1, x)
    return special.orthogonal._gen_roots_and_weights(m, mu0, an_func, bn_func, f, df, False, mu)

def roots_sh_jacobi(n, p1, q1, mu=False):
    if (p1-q1) <= -1 or q1 <= 0:
        raise ValueError("(p - q) must be greater than -1, and q must be greater than 0.")
    x, w, m = roots_jacobi(n, p1-q1, q1-1, True)
    x = (x + 1) / 2
    scale = 1.0 # 2.0**p1
    w /= scale
    m /= scale
    if mu:
        return x, w, m
    else:
        return x, w