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
+from numpy import finfo, exp as np_exp, sin as np_sin, asarray
 
 from scipy.optimize import zeros as cc
 from scipy.optimize import zeros
@@ -56,6 +56,44 @@ 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 f_solarcell(i, v, il, io, rs, rsh, vt):
+            vd = v + i * rs
+            return il - io * (np_exp(vd / vt) - 1.0) - vd / rsh - i
+
+        def fprime(i, v, il, io, rs, rsh, vt):
+            return -io * np_exp((v + i * rs) / vt) * rs / vt - rs / rsh - 1
+
+        def fprime2(i, v, il, io, rs, rsh, vt):
+            return -io * np_exp((v + i * rs) / vt) * (rs / vt)**2
+
+        il = (np_sin(range(10)) + 1.0) * 7.0
+        v = asarray([
+            5.32725221, 5.48673747, 5.49539973,
+            5.36387202, 4.80237316, 1.43764452,
+            5.23063958, 5.46094772, 5.50512718,
+            5.42046290
+        ])
+        args = (v, il, 1e-09, 0.004, 10, 0.27456)
+        x0 = 7.0
+        x = zeros.newton(f_solarcell, x0, fprime, args)
+        y = (6.17264965, 11.7702805, 12.2219954,
+             7.11017681, 1.18151293, 0.143707955,
+             4.31928228, 10.5419107, 12.7552490,
+             8.91225749)
+        assert_allclose(x, y)
+        # test halley's
+        x = zeros.newton(f_solarcell, x0, fprime, args, fprime2=fprime2)
+        assert_allclose(x, y)
+        # test secant
+        x = zeros.newton(lambda x, y: y - x**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 x: x**2, [0., 0.], lambda x: 2*x)
+
     def test_deriv_zero_warning(self):
         func = lambda x: x**2
         dfunc = lambda x: 2*x

scipy/optimize/zeros.py

@@ -3,7 +3,7 @@
 import warnings
 
 from . import _zeros
-from numpy import finfo, sign, sqrt
+from numpy import finfo, sqrt, asarray, abs as np_abs, where
 
 _iter = 100
 _xtol = 2e-12
@@ -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 * asarray(x0)  # convert to ndarray
         for iter in range(maxiter):
             myargs = (p0,) + args
-            fder = fprime(*myargs)
-            if fder == 0:
+            fder = 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 = 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 = 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 = asarray(x0)
+        dx = finfo(float).eps**0.33
+        dp = where(p0 >= 0, dx, -dx)
+        p1 = p0 * (1 + dx) + dp
+        q0 = asarray(func(*((p0,) + args)))
+        q1 = 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" % 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 = asarray(func(*((p1,) + args)))
     msg = "Failed to converge after %d iterations, value is %s" % (maxiter, p)
     raise RuntimeError(msg)