Refactored to decrease cyclomatic complexity

wafo/integrate.py

@@ -212,55 +212,7 @@ def romberg(fun, a, b, releps=1e-3, abseps=1e-3):
     return res, abserr
 
 
-def h_roots(n, method='newton'):
+def _h_roots_newton(n, releps=3e-14, max_iter=10):
-    '''
-    Returns the roots (x) of the nth order Hermite polynomial,
-    H_n(x), and weights (w) to use in Gaussian Quadrature over
-    [-inf,inf] with weighting function exp(-x**2).
-
-    Parameters
-    ----------
-    n : integer
-        number of roots
-    method : 'newton' or 'eigenvalue'
-        uses Newton Raphson to find zeros of the Hermite polynomial (Fast)
-        or eigenvalue of the jacobi matrix (Slow) to obtain the nodes and
-        weights, respectively.
-
-    Returns
-    -------
-    x : ndarray
-        roots
-    w : ndarray
-        weights
-
-    Example
-    -------
-    >>> import numpy as np
-    >>> x, w = h_roots(10)
-    >>> np.allclose(np.sum(x*w), -5.2516042729766621e-19)
-    True
-
-    See also
-    --------
-    qrule, gaussq
-
-    References
-    ----------
-    [1]  Golub, G. H. and Welsch, J. H. (1969)
-    'Calculation of Gaussian Quadrature Rules'
-    Mathematics of Computation, vol 23,page 221-230,
-
-    [2]. Stroud and Secrest (1966), 'gaussian quadrature formulas',
-      prentice-hall, Englewood cliffs, n.j.
-    '''
-
-    if not method.startswith('n'):
-        return ort.h_roots(n)
-    else:
-        sqrt = np.sqrt
-        max_iter = 10
-        releps = 3e-14
     C = [9.084064e-01, 5.214976e-02, 2.579930e-03, 3.986126e-03]
     # PIM4=0.7511255444649425
     PIM4 = np.pi ** (-1. / 4)
@@ -309,28 +261,27 @@ def h_roots(n, method='newton'):
             break
     else:
         warnings.warn('too many iterations!')
+    w = 2. / pp ** 2
+    return _expand_roots(z, w, n, m)
+#     x = np.empty(n)
+#     w = np.empty(n)
+#     x[0:m] = z      # Store the root
+#     x[n - 1:n - m - 1:-1] = -z     # and its symmetric counterpart.
+#     w[0:m] = 2. / pp ** 2    # Compute the weight
+#     w[n - 1:n - m - 1:-1] = w[0:m]  # and its symmetric counterpart.
+#     return x, w
 
-        x = np.empty(n)
-        w = np.empty(n)
-        x[0:m] = z      # Store the root
-        x[n - 1:n - m - 1:-1] = -z     # and its symmetric counterpart.
-        w[0:m] = 2. / pp ** 2    # Compute the weight
-        w[n - 1:n - m - 1:-1] = w[0:m]  # and its symmetric counterpart.
-        return x, w
 
-
+def h_roots(n, method='newton'):
-def j_roots(n, alpha, beta, method='newton'):
     '''
-    Returns the roots of the nth order Jacobi polynomial, P^(alpha,beta)_n(x)
+    Returns the roots (x) of the nth order Hermite polynomial,
-    and weights (w) to use in Gaussian Quadrature over [-1,1] with weighting
+    H_n(x), and weights (w) to use in Gaussian Quadrature over
-    function (1-x)**alpha (1+x)**beta with alpha,beta > -1.
+    [-inf,inf] with weighting function exp(-x**2).
 
     Parameters
     ----------
     n : integer
         number of roots
-    alpha,beta : scalars
-        defining shape of Jacobi polynomial
     method : 'newton' or 'eigenvalue'
         uses Newton Raphson to find zeros of the Hermite polynomial (Fast)
         or eigenvalue of the jacobi matrix (Slow) to obtain the nodes and
@@ -343,20 +294,19 @@ def j_roots(n, alpha, beta, method='newton'):
     w : ndarray
         weights
 
-
     Example
-    --------
+    -------
-    >>> [x,w]= j_roots(10,0,0)
+    >>> import numpy as np
-    >>> sum(x*w)
+    >>> x, w = h_roots(10)
-    2.7755575615628914e-16
+    >>> np.allclose(np.sum(x*w), -5.2516042729766621e-19)
+    True
 
     See also
     --------
     qrule, gaussq
 
-
+    References
-    Reference
+    ----------
-    ---------
     [1]  Golub, G. H. and Welsch, J. H. (1969)
     'Calculation of Gaussian Quadrature Rules'
     Mathematics of Computation, vol 23,page 221-230,
@@ -366,11 +316,11 @@ def j_roots(n, alpha, beta, method='newton'):
     '''
 
     if not method.startswith('n'):
-        [x, w] = ort.j_roots(n, alpha, beta)
+        return ort.h_roots(n)
-    else:
+    return _h_roots_newton(n)
-        max_iter = 10
+
-        releps = 3e-14
 
+def _j_roots_newton(n, alpha, beta, releps=3e-14, max_iter=10):
     # Initial approximations to the roots go into z.
     alfbet = alpha + beta
 
@@ -415,14 +365,64 @@ def j_roots(n, alpha, beta, method='newton'):
     else:
         warnings.warn('too many iterations in jrule')
 
-        x = z  # %Store the root and the weight.
+    x = z  # Store the root and the weight.
     f = (sp.gammaln(alpha + n) + sp.gammaln(beta + n) -
          sp.gammaln(n + 1) - sp.gammaln(alpha + beta + n + 1))
     w = (np.exp(f) * tmp * 2 ** alfbet / (pp * L[k0, :]))
-
     return x, w
 
 
+def j_roots(n, alpha, beta, method='newton'):
+    '''
+    Returns the roots of the nth order Jacobi polynomial, P^(alpha,beta)_n(x)
+    and weights (w) to use in Gaussian Quadrature over [-1,1] with weighting
+    function (1-x)**alpha (1+x)**beta with alpha,beta > -1.
+
+    Parameters
+    ----------
+    n : integer
+        number of roots
+    alpha,beta : scalars
+        defining shape of Jacobi polynomial
+    method : 'newton' or 'eigenvalue'
+        uses Newton Raphson to find zeros of the Hermite polynomial (Fast)
+        or eigenvalue of the jacobi matrix (Slow) to obtain the nodes and
+        weights, respectively.
+
+    Returns
+    -------
+    x : ndarray
+        roots
+    w : ndarray
+        weights
+
+
+    Example
+    --------
+    >>> [x,w]= j_roots(10,0,0)
+    >>> sum(x*w)
+    2.7755575615628914e-16
+
+    See also
+    --------
+    qrule, gaussq
+
+
+    Reference
+    ---------
+    [1]  Golub, G. H. and Welsch, J. H. (1969)
+     'Calculation of Gaussian Quadrature Rules'
+      Mathematics of Computation, vol 23,page 221-230,
+
+    [2]. Stroud and Secrest (1966), 'gaussian quadrature formulas',
+          prentice-hall, Englewood cliffs, n.j.
+    '''
+
+    if not method.startswith('n'):
+        return ort.j_roots(n, alpha, beta)
+    return _j_roots_newton(n, alpha, beta)
+
+
 def la_roots(n, alpha=0, method='newton'):
     '''
     Returns the roots (x) of the nth order generalized (associated) Laguerre
@@ -471,9 +471,11 @@ def la_roots(n, alpha=0, method='newton'):
 
     if not method.startswith('n'):
         return ort.la_roots(n, alpha)
-    else:
+    return _la_roots_newton(n, alpha)
-        max_iter = 10
+
-        releps = 3e-14
+
+def _la_roots_newton(n, alpha, releps=3e-14, max_iter=10):
+
     C = [9.084064e-01, 5.214976e-02, 2.579930e-03, 3.986126e-03]
 
     # Initial approximations to the roots go into z.
@@ -527,68 +529,57 @@ def la_roots(n, alpha=0, method='newton'):
     return x, w
 
 
-def p_roots(n, method='newton', a=-1, b=1):
+def _p_roots_newton_start(n):
-    '''
+    m = int(np.fix((n + 1) / 2))
-    Returns the roots (x) of the nth order Legendre polynomial, P_n(x),
+    mm = 4 * m - 1
-    and weights (w) to use in Gaussian Quadrature over [-1,1] with weighting
+    t = (np.pi / (4 * n + 2)) * np.arange(3, mm + 1, 4)
-    function 1.
+    nn = 1 - (1 - 1 / n) / (8 * n * n)
-
+    xo = nn * np.cos(t)
-    Parameters
+    return m, xo
-    ----------
-    n : integer
-        number of roots
-    method : 'newton' or 'eigenvalue'
-        uses Newton Raphson to find zeros of the Hermite polynomial (Fast)
-        or eigenvalue of the jacobi matrix (Slow) to obtain the nodes and
-        weights, respectively.
-
-    Returns
-    -------
-    x : ndarray
-        roots
-    w : ndarray
-        weights
-
-
-    Example
-    -------
-    Integral of exp(x) from a = 0 to b = 3 is: exp(3)-exp(0)=
-    >>> import numpy as np
-    >>> x, w = p_roots(11, a=0, b=3)
-    >>> np.allclose(np.sum(np.exp(x)*w), 19.085536923187668)
-    True
-
-    See also
-    --------
-    quadg.
-
-
-    References
-    ----------
-    [1] Davis and Rabinowitz (1975) 'Methods of Numerical Integration',
-        page 365, Academic Press.
 
-    [2]  Golub, G. H. and Welsch, J. H. (1969)
-        'Calculation of Gaussian Quadrature Rules'
-        Mathematics of Computation, vol 23,page 221-230,
 
-    [3] Stroud and Secrest (1966), 'gaussian quadrature formulas',
+def _p_roots_newton(n, ):
-        prentice-hall, Englewood cliffs, n.j.
+    """
-    '''
+    Algorithm given by Davis and Rabinowitz in 'Methods
+    of Numerical Integration', page 365, Academic Press, 1975.
+    """
+    m, xo = _p_roots_newton_start(n)
 
-    if not method.startswith('n'):
+    e1 = n * (n + 1)
-        x, w = ort.p_roots(n)
+    for _j in range(2):
-    else:
+        pkm1 = 1
+        pk = xo
+        for k in range(2, n + 1):
+            t1 = xo * pk
+            pkp1 = t1 - pkm1 - (t1 - pkm1) / k + t1
+            pkm1 = pk
+            pk = pkp1
 
-        m = int(np.fix((n + 1) / 2))
+        den = 1. - xo * xo
+        d1 = n * (pkm1 - xo * pk)
+        dpn = d1 / den
+        d2pn = (2. * xo * dpn - e1 * pk) / den
+        d3pn = (4. * xo * d2pn + (2 - e1) * dpn) / den
+        d4pn = (6. * xo * d3pn + (6 - e1) * d2pn) / den
+        u = pk / dpn
+        v = d2pn / dpn
+        h = -u * (1 + (.5 * u) * (v + u * (v * v - u * d3pn / (3 * dpn))))
+        p = pk + h * (dpn + (.5 * h) * (d2pn +
+                                        (h / 3) * (d3pn + .25 * h * d4pn)))
+        dp = dpn + h * (d2pn + (.5 * h) * (d3pn + h * d4pn / 3))
+        h = h - p / dp
+        xo = xo + h
 
-        mm = 4 * m - 1
+    x = -xo - h
-        t = (np.pi / (4 * n + 2)) * np.arange(3, mm + 1, 4)
+    fx = d1 - h * e1 * (pk + (h / 2) * (dpn + (h / 3) *
-        nn = (1 - (1 - 1 / n) / (8 * n * n))
+                                        (d2pn + (h / 4) *
-        xo = nn * np.cos(t)
+                                         (d3pn + (.2 * h) * d4pn))))
+    w = 2 * (1 - x ** 2) / (fx ** 2)
+    return _expand_roots(x, w, n, m)
 
-        if method.endswith('1'):
 
+def _p_roots_newton1(n, releps=1e-15, max_iter=100):
+    m, xo = _p_roots_newton_start(n)
     # Compute the zeros of the N+1 Legendre Polynomial
     # using the recursion relation and the Newton-Raphson method
 
@@ -599,8 +590,6 @@ def p_roots(n, method='newton', a=-1, b=1):
     Lp = zeros((m,))
     dx = zeros((m,))
 
-            releps = 1e-15
-            max_iter = 100
     # Compute the zeros of the N+1 Legendre Polynomial
     # using the recursion relation and the Newton-Raphson method
 
@@ -632,51 +621,75 @@ def p_roots(n, method='newton', a=-1, b=1):
 
     x = -xo
     w = 2. / ((1 - x ** 2) * (Lp ** 2))
-        else:
+    return _expand_roots(x, w, n, m)
-            # Algorithm given by Davis and Rabinowitz in 'Methods
-            # of Numerical Integration', page 365, Academic Press, 1975.
-
-            e1 = n * (n + 1)
 
-            for _j in range(2):
-                pkm1 = 1
-                pk = xo
-                for k in range(2, n + 1):
-                    t1 = xo * pk
-                    pkp1 = t1 - pkm1 - (t1 - pkm1) / k + t1
-                    pkm1 = pk
-                    pk = pkp1
-
-                den = 1. - xo * xo
-                d1 = n * (pkm1 - xo * pk)
-                dpn = d1 / den
-                d2pn = (2. * xo * dpn - e1 * pk) / den
-                d3pn = (4. * xo * d2pn + (2 - e1) * dpn) / den
-                d4pn = (6. * xo * d3pn + (6 - e1) * d2pn) / den
-                u = pk / dpn
-                v = d2pn / dpn
-                h = (-u * (1 + (.5 * u) * (v + u *
-                                           (v * v - u * d3pn / (3 * dpn)))))
-                p = (pk + h * (dpn + (.5 * h) * (d2pn + (h / 3) *
-                                                 (d3pn + .25 * h * d4pn))))
-                dp = dpn + h * (d2pn + (.5 * h) * (d3pn + h * d4pn / 3))
-                h = h - p / dp
-                xo = xo + h
-
-            x = -xo - h
-            fx = (d1 - h * e1 * (pk + (h / 2) *
-                                 (dpn + (h / 3) * (d2pn + (h / 4) *
-                                                   (d3pn + (.2 * h) * d4pn)))))
-            w = 2 * (1 - x ** 2) / (fx ** 2)
 
+def _expand_roots(x, w, n, m):
     if (m + m) > n:
         x[m - 1] = 0.0
-
     if not ((m + m) == n):
         m = m - 1
-
     x = np.hstack((x, -x[m - 1::-1]))
     w = np.hstack((w, w[m - 1::-1]))
+    return x, w
+
+
+def p_roots(n, method='newton', a=-1, b=1):
+    '''
+    Returns the roots (x) of the nth order Legendre polynomial, P_n(x),
+    and weights (w) to use in Gaussian Quadrature over [-1,1] with weighting
+    function 1.
+
+    Parameters
+    ----------
+    n : integer
+        number of roots
+    method : 'newton' or 'eigenvalue'
+        uses Newton Raphson to find zeros of the Hermite polynomial (Fast)
+        or eigenvalue of the jacobi matrix (Slow) to obtain the nodes and
+        weights, respectively.
+
+    Returns
+    -------
+    x : ndarray
+        roots
+    w : ndarray
+        weights
+
+
+    Example
+    -------
+    Integral of exp(x) from a = 0 to b = 3 is: exp(3)-exp(0)=
+    >>> import numpy as np
+    >>> x, w = p_roots(11, a=0, b=3)
+    >>> np.allclose(np.sum(np.exp(x)*w), 19.085536923187668)
+    True
+
+    See also
+    --------
+    quadg.
+
+
+    References
+    ----------
+    [1] Davis and Rabinowitz (1975) 'Methods of Numerical Integration',
+        page 365, Academic Press.
+
+    [2]  Golub, G. H. and Welsch, J. H. (1969)
+        'Calculation of Gaussian Quadrature Rules'
+        Mathematics of Computation, vol 23,page 221-230,
+
+    [3] Stroud and Secrest (1966), 'gaussian quadrature formulas',
+        prentice-hall, Englewood cliffs, n.j.
+    '''
+
+    if not method.startswith('n'):
+        x, w = ort.p_roots(n)
+    else:
+        if method.endswith('1'):
+            x, w = _p_roots_newton1(n)
+        else:
+            x, w = _p_roots_newton(n)
 
     if (a != -1) | (b != 1):
         # Linear map from[-1,1] to [a,b]