Reduced cyclomatic complexity

wafo/polynomial.py

@@ -122,6 +122,12 @@ def polyint(p, m=1, k=None):
             msg = "Order of integral must be positive (see polyder)"
             raise ValueError(msg)
 
+    def _check_integration_const(k, m):
+        if len(k) < m:
+            msg = "k must be a scalar or a rank-1 array of length 1 or >m."
+            raise ValueError(msg)
+
+    def _init(m, k):
         m = int(m)
         _check_order(m)
         if k is None:
@@ -129,9 +135,10 @@ def polyint(p, m=1, k=None):
         k = np.atleast_1d(k)
         if len(k) == 1 and m > 1:
             k = k[0] * np.ones(m, float)
-    if len(k) < m:
+        _check_integration_const(k, m)
-        raise ValueError(
+        return m, k
-            "k must be a scalar or a rank-1 array of length 1 or >m.")
+
+    m, k = _init(m, k)
     truepoly = isinstance(p, poly1d)
     p = np.asarray(p)
     if m == 0:
@@ -508,18 +515,23 @@ def polyreloc(p, x, y=0.0):
     >>> wp.polyval(r,0)     # = polyval(p,1)
     array([3, 5, 7])
     """
+    def _reshape(r):
+        if r.ndim > 1 and r.shape[-1] == 1:
+            r.shape = (r.size,)
+        return r
 
-    truepoly = isinstance(p, poly1d)
+    def _relocate_with_horner(p, x, y):
         r = np.atleast_1d(p).copy()
         n = r.shape[0]
-
         # Relocate polynomial using Horner's algorithm
         for ii in range(n, 1, -1):
             for i in range(1, ii):
                 r[i] = r[i] - x * r[i - 1]
         r[-1] = r[-1] + y
-    if r.ndim > 1 and r.shape[-1] == 1:
+        return _reshape(r)
-        r.shape = (r.size,)
+
+    truepoly = isinstance(p, poly1d)
+    r = _relocate_with_horner(p, x, y)
     if truepoly:
         r = poly1d(r)
     return r
@@ -1284,21 +1296,26 @@ def chebfit_dct(f, n=(10, ), domain=None):
     Approximations for Functions of a Single Independent Variable"
     Journal of the ACM (JACM), Vol. 12 ,  Issue 3, pp 295 - 314
     """
-    n = np.atleast_1d(n)
+    def _check(n):
         if np.any(n > 50):
             warnings.warn('CHEBFIT should only be used for n<50')
 
-    if hasattr(f, '__call__'):
+    def _zip(n, domain):
         if domain is None:
             domain = (-1, 1) * len(n)
-        domain = np.atleast_2d(domain).reshape((-1, 2))
+        return zip(n, np.atleast_2d(domain).reshape((-1, 2)))
+
+    def _init_ck(f, n, domain):
+        n = np.atleast_1d(n)
+        _check(n)
+        if hasattr(f, '__call__'):
             xi = [map_to_interval(chebroot(ni), d[0], d[1])
-              for ni, d in zip(n, domain)]
+                  for ni, d in _zip(n, domain)]
             Xi = np.meshgrid(*xi)
-        ck = f(*Xi)
+            return f(*Xi), n
-    else:
+        return f, f.shape
-        ck = f
+
-        n = ck.shape
+    ck, n = _init_ck(f, n, domain)
 
     ndim = len(n)
     for i in range(ndim):
@@ -1868,27 +1885,38 @@ def padefitlsq(fun, m, k, a=-1, b=1, trace=False, x=None, end_points=True):
             warnings.warn('Check the result! Number of function values ' +
                           'should be at least: {0:d}'.format(npt))
 
-    NFAC = 8
+    def _init(fun, a, b, x, end_points, npt):
-    BIG = 1e30
-    MAXIT = 5
-
-    smallest_devmax = BIG
-    ncof = m + k + 1
-    # Number of points where function is evaluated, i.e. fineness of mesh
-    npt = NFAC * ncof
-
         if x is None:
             x = map_to_interval(_points(npt, end_points), a, b)
-
         if hasattr(fun, '__call__'):
             fs = fun(x)
         else:
             fs = fun
             _check_size(fs, npt)
+        return x, fs
 
+    def _cond_plot1(trace, x, fs):
         if trace:
             plt.plot(x, fs, '+')
 
+    def _cond_plot2(x, fs, x, ys, ix, devmax):
+        if trace:
+            print('Iteration=%d,  max error=%g' % (ix, devmax))
+            plt.plot(x, fs, x, ee + fs)
+
+    NFAC = 8
+    BIG = 1e30
+    MAXIT = 5
+
+    smallest_devmax = BIG
+    ncof = m + k + 1
+    # Number of points where function is evaluated, i.e. fineness of mesh
+    npt = NFAC * ncof
+
+    x, fs = _init(fun, a, b, x, end_points, npt)
+
+    _cond_plot1(trace, x, fs)
+
     wt = np.ones((npt))
     ee = np.ones((npt))
     mad = 0
@@ -1924,10 +1952,7 @@ def padefitlsq(fun, m, k, a=-1, b=1, trace=False, x=None, end_points=True):
             smallest_devmax = devmax
             c1 = cof[m::-1]
             c2 = cof[ncof:m:-1].tolist() + [1, ]
-
+        _cond_plot2(x, fs, x, ee + fs, ix, devmax)
-        if trace:
-            print('Iteration=%d,  max error=%g' % (ix, devmax))
-            plt.plot(x, fs, x, ee + fs)
     return poly1d(c1), poly1d(c2)
 
 
@@ -2050,15 +2075,18 @@ def chebvandernd(deg, *xi):
     --------
     chebvander, chebvalnd, chebfitnd
     """
-    ideg = [int(d) for d in deg]
+    def _check_deg(ideg, is_valid, ndim):
-    is_valid = np.array([di == d and di >= 0 for di, d in zip(ideg, deg)])
         if np.any(is_valid != 1):
             raise ValueError("degrees must be non-negative integers")
-    ndim = len(xi)
         if len(ideg) != ndim:
             msg = 'length of deg must be the same as number of dimensions'
             raise ValueError(msg)
 
+    ideg = [int(d) for d in deg]
+    is_valid = np.array([di == d and di >= 0 for di, d in zip(ideg, deg)])
+    ndim = len(xi)
+    _check_deg(ideg, is_valid, ndim)
+
     xi = np.array(xi, copy=0) + 0.0
     chebvander = np.polynomial.chebyshev.chebvander
     shape0 = xi[0].shape
@@ -2161,7 +2189,9 @@ def chebfitnd(xi, f, deg, rcond=None, full=False, w=None):
     Examples
     --------
     """
-    def _check_shapes(z, ndims, sizes):
+    def _check_shapes(z, xi):
+        ndims = np.array([np.ndim(x) for x in xi])
+        sizes = np.array([np.size(x) for x in xi])
         ndim = len(ndims)
         if np.any(ndims != ndim) or z.ndim != ndim:
             msg = "expected {0:d}D array for x1, x2,...,xn and f".format(ndim)
@@ -2173,16 +2203,16 @@ def chebfitnd(xi, f, deg, rcond=None, full=False, w=None):
         if n != len(w):
             raise TypeError("expected x and w to have same length")
 
-    # xi = np.array(xi, copy=0) + 0.0
+    def _scale(lhs):
-    z = np.array(f)
+        if issubclass(lhs.dtype.type, np.complexfloating):
-    degrees = np.asarray(deg, dtype=int)
+            scl = np.sqrt((np.square(lhs.real) +
-    orders = degrees + 1
+                           np.square(lhs.imag)).sum(axis=0))
-    order = np.product(orders)
+        else:
-
+            scl = np.sqrt(np.square(lhs).sum(axis=0))
-    ndims = np.array([x.ndim for x in xi])
+        scl[scl == 0] = 1
-    sizes = np.array([x.size for x in xi])
+        return scl
-    _check_shapes(z, ndims, sizes)
 
+    def _init(xi, z, w, degrees, order):
         lhs = chebvandernd(degrees, *xi).reshape((-1, order))
         rhs = z.ravel()
         if w is not None:
@@ -2190,15 +2220,21 @@ def chebfitnd(xi, f, deg, rcond=None, full=False, w=None):
             _check_size(w, len(lhs))
             lhs = lhs * w
             rhs = rhs * w
+        scl = _scale(lhs)
+        return lhs, scl, rhs
 
-    if rcond is None:
+    # xi = np.array(xi, copy=0) + 0.0
-        rcond = xi[0].size * np.finfo(x.dtype).eps
+    z = np.array(f)
+    _check_shapes(z, xi)
 
-    if issubclass(lhs.dtype.type, np.complexfloating):
+    degrees = np.asarray(deg, dtype=int)
-        scl = np.sqrt((np.square(lhs.real) + np.square(lhs.imag)).sum(axis=0))
+    orders = degrees + 1
-    else:
+    order = np.product(orders)
-        scl = np.sqrt(np.square(lhs).sum(axis=0))
+
-    scl[scl == 0] = 1
+    lhs, rhs, scl = _init(xi, z, w, degrees, order)
+
+    if rcond is None:
+        rcond = xi[0].size * np.finfo(xi[0].dtype).eps
 
     # Solve the least squares problem.
     c, resids, rank, s = np.linalg.lstsq(lhs/scl, rhs, rcond)