improve line search

Author: yoelcortes

flexsolve/line_search.py

@@ -5,26 +5,31 @@
 """
 import numpy as np
 from typing import NamedTuple
+from scipy.optimize import fminbound
 
-__all__ = ('inexact_line_search', 'LineSearchResult')
+__all__ = (
+    'exact_line_search',
+    'inexact_line_search', 
+    'LineSearchResult'
+)
 
 
 class LineSearchResult(NamedTuple):
-    t: float
-    x: np.ndarray
-    f: float
+    t: float #: step length
+    x: np.ndarray # data point
+    r: float # residual
 
 
-# def exact_line_search(f, x, correction, t0=1e-9, t1=2, ttol=1e-6, maxiter=20,
-#                       full_output=True):
-#     return fminbound(
-#         lambda t: f(x + t * correction),
-#         x1=t0,
-#         x2=t1,
-#         xtol=ttol,
-#         maxfun=maxiter,
-#         full_output=full_output,
-#     )
+def exact_line_search(f, x, correction, t0=1e-9, t1=2, ttol=1e-6, maxiter=20):
+    t, f, *_ = fminbound(
+        lambda t: f(x + t * correction),
+        x1=t0,
+        x2=t1,
+        xtol=ttol,
+        maxfun=maxiter,
+        full_output=True,
+    )
+    return LineSearchResult(t, x + t * correction, f)
 
 def inexact_line_search(
         f, x, correction, t0=1e-6, t1=1.1, ttol=1e-3, maxiter=10,
@@ -40,7 +45,9 @@ def inexact_line_search(
     if ft1 > fx:
         xguess = x + tguess * correction
         ftguess = f(xguess)
-        if ft1 > ft0:
+        if ftguess < fx: # Done!
+            best = LineSearchResult(tguess, xguess, ftguess)
+        elif ft1 > ft0: # Move towards x0.
             best = LineSearchResult(t0, x0, ft0)
             for i in range(maxiter):
                 t = rho * tguess + (1 - rho) * t0
@@ -49,29 +56,16 @@ def inexact_line_search(
                 if ft < best.f: best = LineSearchResult(t, xt, ft)
                 if abs(t - tguess) < ttol: break
                 tguess = t
-        else:
+        elif ftguess < ft1: # No where to go; moving forward is risky
             best = LineSearchResult(tguess, xguess, ftguess)
-            for i in range(maxiter):
-                t = (1 - rho) * tguess + rho * t1
-                xt = x + t * correction
-                ft = f(xt)
-                if ft < ftguess: best = LineSearchResult(t, xt, ft)
-                if abs(t - tguess) < ttol: break
-                t1 = t
+        else: # Hope for the best
+            best = LineSearchResult(t1, x1, ft1)
     else:
         xguess = x + tguess * correction
         ftguess = f(xguess)
-        if ftguess < ft1:
-            tnext = tguess
+        if ftguess < ft1: # Stick with guess
             best = LineSearchResult(tguess, xguess, ftguess)
-            for i in range(maxiter):
-                t = (1 - rho) * tnext + rho * t1
-                xt = x + t * correction
-                ft = f(xt)
-                if ft < best.f: best = LineSearchResult(t, xt, ft)
-                if abs(t - tguess) < ttol: break
-                t1 = t
-        else:
+        else: # Accelerate
             best = LineSearchResult(t1, x1, ft1)
     return best