ENH: vectorize scalar zero-search-functions

scipy/optimize/tests/test_zeros.py

@@ -5,7 +5,7 @@
 from numpy.testing import (assert_warns, assert_, 
                            assert_allclose,
                            assert_equal)
-from numpy import finfo
+import numpy as np
 
 from scipy.optimize import zeros as cc
 from scipy.optimize import zeros
@@ -18,8 +18,8 @@ class TestBasic(object):
     def run_check(self, method, name):
         a = .5
         b = sqrt(3)
-        xtol = 4*finfo(float).eps
-        rtol = 4*finfo(float).eps
+        xtol = 4*np.finfo(float).eps
+        rtol = 4*np.finfo(float).eps
         for function, fname in zip(functions, fstrings):
             zero, r = method(function, a, b, xtol=xtol, rtol=rtol,
                              full_output=True)
@@ -56,6 +56,48 @@ def test_newton(self):
             x = zeros.newton(f, 3, fprime=f_1, fprime2=f_2, tol=1e-6)
             assert_allclose(f(x), 0, atol=1e-6)
 
+    def test_newton_array(self):
+        """test newton with array"""
+
+        def f1(x, *a):
+            b = a[0] + x * a[3]
+            return a[1] - a[2] * (np.exp(b / a[5]) - 1.0) - b / a[4] - x
+
+        def f1_1(x, *a):
+            b = a[3] / a[5]
+            return -a[2] * np.exp(a[0] / a[5] + x * b) * b - a[3] / a[4] - 1
+
+        def f1_2(x, *a):
+            b = a[3] / a[5]
+            return -a[2] * np.exp(a[0] / a[5] + x * b) * b ** 2
+
+        a0 = np.array([
+            5.32725221, 5.48673747, 5.49539973,
+            5.36387202, 4.80237316, 1.43764452,
+            5.23063958, 5.46094772, 5.50512718,
+            5.42046290
+        ])
+        a1 = (np.sin(range(10)) + 1.0) * 7.0
+        args = (a0, a1, 1e-09, 0.004, 10, 0.27456)
+        x0 = 7.0
+        x = zeros.newton(f1, x0, f1_1, args)
+        x_expected = (
+            6.17264965, 11.7702805, 12.2219954,
+            7.11017681, 1.18151293, 0.143707955,
+            4.31928228, 10.5419107, 12.7552490,
+            8.91225749
+        )
+        assert_allclose(x, x_expected)
+        # test halley's
+        x = zeros.newton(f1, x0, f1_1, args, fprime2=f1_2)
+        assert_allclose(x, x_expected)
+        # test secant
+        x = zeros.newton(lambda y, z: z - y ** 2, 4.0, args=([15.0, 17.0],))
+        assert_allclose(x, (3.872983346207417, 4.123105625617661))
+        # test derivative zero warning
+        assert_warns(RuntimeWarning, zeros.newton,
+                     lambda y: y ** 2, [0., 0.], lambda y: 2 * y)
+
     def test_deriv_zero_warning(self):
         func = lambda x: x**2
         dfunc = lambda x: 2*x
@@ -69,8 +111,8 @@ def f(x):
         return x - root
 
     methods = [cc.bisect, cc.ridder]
-    xtol = 4*finfo(float).eps
-    rtol = 4*finfo(float).eps
+    xtol = 4*np.finfo(float).eps
+    rtol = 4*np.finfo(float).eps
     for method in methods:
         res = method(f, -1e8, 1e7, xtol=xtol, rtol=rtol)
         assert_allclose(root, res, atol=xtol, rtol=rtol,
@@ -94,7 +136,7 @@ def f(x):
             return x - 0.6
 
     atol = 0.51
-    rtol = 4*finfo(float).eps
+    rtol = 4*np.finfo(float).eps
     methods = [cc.brentq, cc.brenth]
     for method in methods:
         res = method(f, 0, 1, xtol=atol, rtol=rtol)

scipy/optimize/zeros.py

@@ -3,11 +3,11 @@
 import warnings
 
 from . import _zeros
-from numpy import finfo, sign, sqrt
+import numpy as np
 
 _iter = 100
 _xtol = 2e-12
-_rtol = 4*finfo(float).eps
+_rtol = 4*np.finfo(float).eps
 
 __all__ = ['newton', 'bisect', 'ridder', 'brentq', 'brenth']
 
@@ -163,54 +163,50 @@ def newton(func, x0, fprime=None, args=(), tol=1.48e-8, maxiter=50,
         # Newton-Rapheson method
         # Multiply by 1.0 to convert to floating point.  We don't use float(x0)
         # so it still works if x0 is complex.
-        p0 = 1.0 * x0
-        fder2 = 0
+        p0 = 1.0 * np.asarray(x0)  # convert to ndarray
         for iter in range(maxiter):
             myargs = (p0,) + args
-            fder = fprime(*myargs)
-            if fder == 0:
+            fder = np.asarray(fprime(*myargs))  # convert to ndarray
+            if (fder == 0).any():
                 msg = "derivative was zero."
                 warnings.warn(msg, RuntimeWarning)
                 return p0
-            fval = func(*myargs)
-            if fprime2 is not None:
-                fder2 = fprime2(*myargs)
-            if fder2 == 0:
+            fval = np.asarray(func(*myargs))  # convert to ndarray
+            if fprime2 is None:
                 # Newton step
                 p = p0 - fval / fder
             else:
-                # Parabolic Halley's method
-                discr = fder ** 2 - 2 * fval * fder2
-                if discr < 0:
-                    p = p0 - fder / fder2
-                else:
-                    p = p0 - 2*fval / (fder + sign(fder) * sqrt(discr))
-            if abs(p - p0) < tol:
+                fder2 = np.asarray(fprime2(*myargs))  # convert to ndarray
+                # Halley's method
+                # https://en.wikipedia.org/wiki/Halley%27s_method
+                p = p0 - 2 * fval * fder / (2 * fder ** 2 - fval * fder2)
+            if np.abs(p - p0).max() < tol:
                 return p
             p0 = p
     else:
         # Secant method
-        p0 = x0
-        if x0 >= 0:
-            p1 = x0*(1 + 1e-4) + 1e-4
-        else:
-            p1 = x0*(1 + 1e-4) - 1e-4
-        q0 = func(*((p0,) + args))
-        q1 = func(*((p1,) + args))
+        p0 = np.asarray(x0)
+        dx = np.finfo(float).eps**0.33
+        dp = np.where(p0 >= 0, dx, -dx)
+        p1 = p0 * (1 + dx) + dp
+        q0 = np.asarray(func(*((p0,) + args)))
+        q1 = np.asarray(func(*((p1,) + args)))
         for iter in range(maxiter):
-            if q1 == q0:
-                if p1 != p0:
-                    msg = "Tolerance of %s reached" % (p1 - p0)
+            divide_by_zero = (q1 == q0)
+            if divide_by_zero.any():
+                tolerance_reached = (p1 != p0)
+                if (divide_by_zero & tolerance_reached).any():
+                    msg = "Tolerance of %s reached" % np.sqrt(sum((p1 - p0)**2))
                     warnings.warn(msg, RuntimeWarning)
                 return (p1 + p0)/2.0
             else:
                 p = p1 - q1*(p1 - p0)/(q1 - q0)
-            if abs(p - p1) < tol:
+            if np.abs(p - p1).max() < tol:
                 return p
             p0 = p1
             q0 = q1
             p1 = p
-            q1 = func(*((p1,) + args))
+            q1 = np.asarray(func(*((p1,) + args)))
     msg = "Failed to converge after %d iterations, value is %s" % (maxiter, p)
     raise RuntimeError(msg)