ENH: vectorize scalar zero-search-functions

benchmarks/benchmarks/optimize_zeros.py

@@ -1,6 +1,7 @@
 from __future__ import division, print_function, absolute_import
 
 from math import sqrt, exp, cos, sin
+import numpy as np
 
 # Import testing parameters
 try:
@@ -56,3 +57,64 @@ def setup(self, func, meth):
 
     def time_newton(self, func, meth):
         newton(self.f, self.x0, args=(), fprime=self.f_1, fprime2=self.f_2)
+
+
+class NewtonArray(Benchmark):
+    params = [['loop', 'array'], ['newton', 'secant', 'halley']]
+    param_names = ['vectorization', 'solver']
+
+    def setup(self, vec, meth):
+        if vec == 'loop':
+            if meth == 'newton':
+                self.fvec = lambda f, x0, args, fprime, fprime2: [
+                    newton(f, x, args=(a0, a1) + args[2:], fprime=fprime)
+                    for (x, a0, a1) in zip(x0, args[0], args[1])
+                ]
+            elif meth == 'halley':
+                self.fvec = lambda f, x0, args, fprime, fprime2: [
+                    newton(
+                        f, x, args=(a0, a1) + args[2:], fprime=fprime,
+                        fprime2=fprime2
+                    ) for (x, a0, a1) in zip(x0, args[0], args[1])
+                ]
+            else:
+                self.fvec = lambda f, x0, args, fprime, fprime2: [
+                    newton(f, x, args=(a0, a1) + args[2:]) for (x, a0, a1)
+                    in zip(x0, args[0], args[1])
+                ]
+        else:
+            if meth == 'newton':
+                self.fvec = lambda f, x0, args, fprime, fprime2: newton(
+                    f, x0, args=args, fprime=fprime
+                )
+            elif meth == 'halley':
+                self.fvec = newton
+            else:
+                self.fvec = lambda f, x0, args, fprime, fprime2: newton(
+                    f, x0, args=args
+                )
+
+    def time_array_newton(self, vec, meth):
+
+        def f(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 f_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 f_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] * 10
+        self.fvec(f, x0, args=args, fprime=f_1, fprime2=f_2)

scipy/optimize/tests/test_zeros.py

@@ -6,20 +6,22 @@
 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
 
 # Import testing parameters
 from scipy.optimize._tstutils import functions, fstrings
 
+TOL = 4*np.finfo(float).eps  # tolerance
+
+
 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 = rtol = TOL
         for function, fname in zip(functions, fstrings):
             zero, r = method(function, a, b, xtol=xtol, rtol=rtol,
                              full_output=True)
@@ -56,6 +58,72 @@ 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_array_newton(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] * 10
+        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(f1, x0, args=args)
+        assert_allclose(x, x_expected)
+
+    def test_array_secant_active_zero_der(self):
+        """test secant doesn't continue to iterate zero derivatives"""
+        x = zeros.newton(lambda x, *a: x*x - a[0], x0=[4.123, 5],
+                         args=[np.array([17, 25])])
+        assert_allclose(x, (4.123105625617661, 5.0))
+
+    def test_array_newton_integers(self):
+        # test secant with float
+        x = zeros.newton(lambda y, z: z - y ** 2, [4.0] * 2,
+                         args=([15.0, 17.0],))
+        assert_allclose(x, (3.872983346207417, 4.123105625617661))
+        # test integer becomes float
+        x = zeros.newton(lambda y, z: z - y ** 2, [4] * 2, args=([15, 17],))
+        assert_allclose(x, (3.872983346207417, 4.123105625617661))
+
+    def test_array_newton_zero_der_failures(self):
+        # test derivative zero warning
+        assert_warns(RuntimeWarning, zeros.newton,
+                     lambda y: y**2 - 2, [0., 0.], lambda y: 2 * y)
+        # test failures and zero_der
+        with pytest.warns(RuntimeWarning):
+            results = zeros.newton(lambda y: y**2 - 2, [0., 0.],
+                                   lambda y: 2*y, converged=True)
+            assert_allclose(results.root, 0)
+            assert results.zero_der.all()
+            assert not results.converged.any()
+
     def test_newton_full_output(self):
         # Test the full_output capability, both when converging and not.
         # Use simple polynomials, to avoid hitting platform dependencies
@@ -109,8 +177,7 @@ def f(x):
         return x - root
 
     methods = [cc.bisect, cc.ridder]
-    xtol = 4*finfo(float).eps
-    rtol = 4*finfo(float).eps
+    xtol = rtol = TOL
     for method in methods:
         res = method(f, -1e8, 1e7, xtol=xtol, rtol=rtol)
         assert_allclose(root, res, atol=xtol, rtol=rtol,
@@ -134,7 +201,7 @@ def f(x):
             return x - 0.6
 
     atol = 0.51
-    rtol = 4*finfo(float).eps
+    rtol = TOL
     methods = [cc.brentq, cc.brenth]
     for method in methods:
         res = method(f, 0, 1, xtol=atol, rtol=rtol)
@@ -162,13 +229,78 @@ def f_1(x, *a):
         return 2 * a[0] * x + a[1]
 
     def f_2(x, *a):
-        return 2 * a[0]
+        retval = 2 * a[0]
+        try:
+            size = len(x)
+        except TypeError:
+            return retval
+        else:
+            return [retval] * size
 
     z = complex(1.0, 2.0)
     coeffs = (2.0, 3.0, 4.0)
     y = zeros.newton(f, z, args=coeffs, fprime=f_1, fprime2=f_2, tol=1e-6)
     # (-0.75000000000000078+1.1989578808281789j)
     assert_allclose(f(y, *coeffs), 0, atol=1e-6)
+    z = [z] * 10
+    coeffs = (2.0, 3.0, 4.0)
+    y = zeros.newton(f, z, args=coeffs, fprime=f_1, fprime2=f_2, tol=1e-6)
+    assert_allclose(f(y, *coeffs), 0, atol=1e-6)
+
+
+def test_zero_der_nz_dp():
+    """Test secant method with a non-zero dp, but an infinite newton step"""
+    # pick a symmetrical functions and choose a point on the side that with dx
+    # makes a secant that is a flat line with zero slope, EG: f = (x - 100)**2,
+    # which has a root at x = 100 and is symmetrical around the line x = 100
+    # we have to pick a really big number so that it is consistently true
+    # now find a point on each side so that the secant has a zero slope
+    dx = np.finfo(float).eps ** 0.33
+    # 100 - p0 = p1 - 100 = p0 * (1 + dx) + dx - 100
+    # -> 200 = p0 * (2 + dx) + dx
+    p0 = (200.0 - dx) / (2.0 + dx)
+    x = zeros.newton(lambda y: (y - 100.0)**2, x0=[p0] * 10)
+    assert_allclose(x, [100] * 10)
+    # test scalar cases too
+    p0 = (2.0 - 1e-4) / (2.0 + 1e-4)
+    x = zeros.newton(lambda y: (y - 1.0) ** 2, x0=p0)
+    assert_allclose(x, 1)
+    p0 = (-2.0 + 1e-4) / (2.0 + 1e-4)
+    x = zeros.newton(lambda y: (y + 1.0) ** 2, x0=p0)
+    assert_allclose(x, -1)
+
+
+def test_array_newton_failures():
+    """Test that array newton fails as expected"""
+    # p = 0.68  # [MPa]
+    # dp = -0.068 * 1e6  # [Pa]
+    # T = 323  # [K]
+    diameter = 0.10  # [m]
+    # L = 100  # [m]
+    roughness = 0.00015  # [m]
+    rho = 988.1  # [kg/m**3]
+    mu = 5.4790e-04  # [Pa*s]
+    u = 2.488  # [m/s]
+    reynolds_number = rho * u * diameter / mu  # Reynolds number
+
+    def colebrook_eqn(darcy_friction, re, dia):
+        return (1 / np.sqrt(darcy_friction) +
+                2 * np.log10(roughness / 3.7 / dia +
+                             2.51 / re / np.sqrt(darcy_friction)))
+
+    # only some failures
+    with pytest.warns(RuntimeWarning):
+        result = zeros.newton(
+            colebrook_eqn, x0=[0.01, 0.2, 0.02223, 0.3], maxiter=2,
+            args=[reynolds_number, diameter], converged=True
+        )
+        assert not result.converged.all()
+    # they all fail
+    with pytest.raises(RuntimeError):
+        result = zeros.newton(
+            colebrook_eqn, x0=[0.01] * 2, maxiter=2,
+            args=[reynolds_number, diameter], converged=True
+        )
 
 
 # this test should **not** raise a RuntimeWarning
@@ -182,16 +314,29 @@ def f_zeroder_root(x):
     # should work with secant
     r = zeros.newton(f_zeroder_root, x0=0)
     assert_allclose(r, 0, atol=zeros._xtol, rtol=zeros._rtol)
+    # test again with array
+    r = zeros.newton(f_zeroder_root, x0=[0]*10)
+    assert_allclose(r, 0, atol=zeros._xtol, rtol=zeros._rtol)
 
     # 1st derivative
     def fder(x):
         return 3 * x**2 - 2 * x
 
+    # 2nd derivative
+    def fder2(x):
+        return 6*x - 2
+
     # should work with newton and halley
     r = zeros.newton(f_zeroder_root, x0=0, fprime=fder)
     assert_allclose(r, 0, atol=zeros._xtol, rtol=zeros._rtol)
     r = zeros.newton(f_zeroder_root, x0=0, fprime=fder,
-                     fprime2=lambda x: 6*x - 2)
+                     fprime2=fder2)
+    assert_allclose(r, 0, atol=zeros._xtol, rtol=zeros._rtol)
+    # test again with array
+    r = zeros.newton(f_zeroder_root, x0=[0]*10, fprime=fder)
+    assert_allclose(r, 0, atol=zeros._xtol, rtol=zeros._rtol)
+    r = zeros.newton(f_zeroder_root, x0=[0]*10, fprime=fder,
+                     fprime2=fder2)
     assert_allclose(r, 0, atol=zeros._xtol, rtol=zeros._rtol)
 
     # also test that if a root is found we do not raise RuntimeWarning even if
@@ -201,4 +346,7 @@ def fder(x):
     # a zero derivative, so it should return the root w/o RuntimeWarning
     r = zeros.newton(f_zeroder_root, x0=0.5, fprime=fder)
     assert_allclose(r, 0, atol=zeros._xtol, rtol=zeros._rtol)
+    # test again with array
+    r = zeros.newton(f_zeroder_root, x0=[0.5]*10, fprime=fder)
+    assert_allclose(r, 0, atol=zeros._xtol, rtol=zeros._rtol)
     # doesn't apply to halley