Merge pull request #280 from choderalab/add_adaptive

Adding back in adaptive solving of the MBAR equations, using the previous algorithm.

appveyor.yml

@@ -12,13 +12,7 @@ environment:
       CONDA_PY: "27"
       CONDA_NPY: "112"
 
-    - PYTHON: "C:\\Python34_64"
-      PYTHON_VERSION: "3.4"
-      PYTHON_ARCH: "64"
-      CONDA_PY: "34"
-      CONDA_NPY: "111"
-
-    - PYTHON: "C:\\Python34_64"
+    - PYTHON: "C:\\Python35_64"
       PYTHON_VERSION: "3.5"
       PYTHON_ARCH: "64"
       CONDA_PY: "35"

pymbar/mbar.py

@@ -278,6 +278,24 @@ def __init__(self, u_kn, N_k, maximum_iterations=10000, relative_tolerance=1.0e-
                 print("f_k = ")
                 print(self.f_k)
 
+        for solver in subsampling_protocol:
+            if 'options' not in solver:
+                solver['options'] = dict()  # make sure there is SOME dictionary in for options.
+                if self.verbose:
+                    # should add in other ways to get information out of the scipy solvers, not just adaptive,
+                    # which might involve passing in different combinations of options, and passing out other strings.
+                    if solver['method'] == 'adaptive':
+                        solver['options']['verbose']=True
+
+        for solver in solver_protocol:
+            if 'options' not in solver:
+                solver['options'] = dict()
+                if self.verbose:
+                    # should add in other ways to get information out of the scipy solvers, not just adaptive,
+                    # which might involve passing in different combinations of options, and passing out other strings.
+                    if solver['method'] == 'adaptive':
+                        solver['options']['verbose']=True
+
         self.f_k = mbar_solvers.solve_mbar_with_subsampling(self.u_kn, self.N_k, self.f_k, solver_protocol, subsampling_protocol, subsampling, x_kindices=self.x_kindices)
         self.Log_W_nk = mbar_solvers.mbar_log_W_nk(self.u_kn, self.N_k, self.f_k)
 

pymbar/mbar_solvers.py

@@ -7,8 +7,11 @@
 
 # Below are the recommended default protocols (ordered sequence of minimization algorithms / NLE solvers) for solving the MBAR equations.
 # Note: we use tuples instead of lists to avoid accidental mutability.
-DEFAULT_SUBSAMPLING_PROTOCOL = (dict(method="L-BFGS-B"), )  # First use BFGS on subsampled data.
-DEFAULT_SOLVER_PROTOCOL = (dict(method="hybr"), )  # Then do fmin hybrid on full dataset.
+#DEFAULT_SUBSAMPLING_PROTOCOL = (dict(method="L-BFGS-B"), )  # First use BFGS on subsampled data.
+#DEFAULT_SOLVER_PROTOCOL = (dict(method="hybr"), )  # Then do fmin hybrid on full dataset.
+# Use Adpative solver as first attempt
+DEFAULT_SUBSAMPLING_PROTOCOL = (dict(method="adaptive"),)
+DEFAULT_SOLVER_PROTOCOL = (dict(method="adaptive",),)
 
 
 def validate_inputs(u_kn, N_k, f_k):
@@ -230,6 +233,117 @@ def mbar_W_nk(u_kn, N_k, f_k):
     """
     return np.exp(mbar_log_W_nk(u_kn, N_k, f_k))
 
+def adaptive(u_kn, N_k, f_k, tol = 1.0e-12, options = None):
+
+    """
+    Determine dimensionless free energies by a combination of Newton-Raphson iteration and self-consistent iteration.
+    Picks whichever method gives the lowest gradient.
+    Is slower than NR since it calculates the log norms twice each iteration.
+
+    OPTIONAL ARGUMENTS
+    tol (float between 0 and 1) - relative tolerance for convergence (default 1.0e-12)
+
+    options: dictionary of options
+        gamma (float between 0 and 1) - incrementor for NR iterations (default 1.0).  Usually not changed now, since adaptively switch.
+        maximum_iterations (int) - maximum number of Newton-Raphson iterations (default 250: either NR converges or doesn't, pretty quickly)
+        verbose (boolean) - verbosity level for debug output
+
+    NOTES
+
+
+    This method determines the dimensionless free energies by
+    minimizing a convex function whose solution is the desired
+    estimator.  The original idea came from the construction of a
+    likelihood function that independently reproduced the work of
+    Geyer (see [1] and Section 6 of [2]).  This can alternatively be
+    formulated as a root-finding algorithm for the Z-estimator.  More
+    details of this procedure will follow in a subsequent paper.  Only
+    those states with nonzero counts are include in the estimation
+    procedure.
+
+    REFERENCES
+    See Appendix C.2 of [1].
+
+    """
+    # put the defaults here in case we get passed an 'options' dictionary that is only partial
+    options.setdefault('verbose',False)
+    options.setdefault('maximum_iterations',250)
+    options.setdefault('print_warning',False)
+    options.setdefault('gamma',1.0)
+
+    gamma = options['gamma']
+    doneIterating = False
+    if options['verbose'] == True:
+        print("Determining dimensionless free energies by Newton-Raphson / self-consistent iteration.")
+
+    if tol < 1.5e-15:
+        print("Tolerance may be too close to machine precision to converge.")
+    # keep track of Newton-Raphson and self-consistent iterations
+    nr_iter = 0
+    sci_iter = 0
+
+    f_sci = np.zeros(len(f_k), dtype=np.float64)
+    f_nr = np.zeros(len(f_k), dtype=np.float64)
+
+    # Perform Newton-Raphson iterations (with sci computed on the way)
+    for iteration in range(0, options['maximum_iterations']):
+        g = mbar_gradient(u_kn, N_k, f_k)  # Objective function gradient
+        H = mbar_hessian(u_kn, N_k, f_k)  # Objective function hessian
+        Hinvg = np.linalg.lstsq(H, g, rcond=-1)[0]
+        Hinvg -= Hinvg[0]
+        f_nr = f_k - gamma * Hinvg
+
+        # self-consistent iteration gradient norm and saved log sums.
+        f_sci = self_consistent_update(u_kn, N_k, f_k)
+        f_sci = f_sci -  f_sci[0]   # zero out the minimum
+        g_sci = mbar_gradient(u_kn, N_k, f_sci)
+        gnorm_sci = np.dot(g_sci, g_sci)
+
+        # newton raphson gradient norm and saved log sums.
+        g_nr = mbar_gradient(u_kn, N_k, f_nr)
+        gnorm_nr = np.dot(g_nr, g_nr)
+
+        # we could save the gradient, for the next round, but it's not too expensive to
+        # compute since we are doing the Hessian anyway.
+
+        if options['verbose']:
+            print("self consistent iteration gradient norm is %10.5g, Newton-Raphson gradient norm is %10.5g" % (gnorm_sci, gnorm_nr))
+        # decide which directon to go depending on size of gradient norm
+        f_old = f_k
+        if (gnorm_sci < gnorm_nr or sci_iter < 2):
+            f_k = f_sci
+            sci_iter += 1
+            if options['verbose']:
+                if sci_iter < 2:
+                    print("Choosing self-consistent iteration on iteration %d" % iteration)
+                else:
+                    print("Choosing self-consistent iteration for lower gradient on iteration %d" % iteration)
+        else:
+            f_k = f_nr
+            nr_iter += 1
+            if options['verbose']:
+                print("Newton-Raphson used on iteration %d" % iteration)
+
+        # routine changes them.
+        max_delta = np.max(np.abs(f_k[1:]-f_old[1:]))/np.max(np.abs(f_k[1:]))
+        if np.isnan(max_delta) or (max_delta < tol):
+            doneIterating = True
+            break
+
+    if doneIterating:
+        if options['verbose']:
+            print('Converged to tolerance of {:e} in {:d} iterations.'.format(max_delta, iteration + 1))
+            print('Of {:d} iterations, {:d} were Newton-Raphson iterations and {:d} were self-consistent iterations'.format(iteration + 1, nr_iter, sci_iter))
+            if np.all(f_k == 0.0):
+                # all f_k appear to be zero
+                print('WARNING: All f_k appear to be zero.')
+    else:
+        print('WARNING: Did not converge to within specified tolerance.')
+        if options['maximum_iterations'] <= 0:
+            print("No iterations ran be cause maximum_iterations was <= 0 ({})!".format(options['maximum_iterations']))
+        else:
+            print('max_delta = {:e}, tol = {:e}, maximum_iterations = {:d}, iterations completed = {:d}'.format(max_delta,tol, options['maximum_iterations'], iteration))
+    return f_k
 
 def precondition_u_kn(u_kn, N_k, f_k):
     """Subtract a sample-dependent constant from u_kn to improve precision
@@ -262,7 +376,7 @@ def precondition_u_kn(u_kn, N_k, f_k):
     return u_kn
 
 
-def solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, method="hybr", tol=1E-20, options=None):
+def solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, method="hybr", tol=1E-12, options=None):
     """Solve MBAR self-consistent equations using some form of equation solver.
 
     Parameters
@@ -277,8 +391,10 @@ def solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, method="hybr", tol=1
     method : str, optional, default="hybr"
         The optimization routine to use.  This can be any of the methods
         available via scipy.optimize.minimize() or scipy.optimize.root().
-    tol : float, optional, default=1E-20
+    tol : float, optional, default=1E-14
         The convergance tolerance for minimize() or root()
+    verbose: bool
+        Whether to print information about the solution method.
     options: dict, optional, default=None
         Optional dictionary of algorithm-specific parameters.  See
         scipy.optimize.root or scipy.optimize.minimize for details.
@@ -319,10 +435,13 @@ def solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, method="hybr", tol=1
             if method in ["L-BFGS-B", "CG"]:
                 hess = None  # To suppress warning from passing a hessian function.
             results = scipy.optimize.minimize(grad_and_obj, f_k_nonzero[1:], jac=True, hess=hess, method=method, tol=tol, options=options)
+            f_k_nonzero = pad(results["x"])
+        elif method == 'adaptive':
+            results = adaptive(u_kn_nonzero, N_k_nonzero, f_k_nonzero, tol=tol, options=options)
+            f_k_nonzero = results # they are the same for adaptive, until we decide to return more.
         else:
             results = scipy.optimize.root(grad, f_k_nonzero[1:], jac=hess, method=method, tol=tol, options=options)
-
-    f_k_nonzero = pad(results["x"])
+            f_k_nonzero = pad(results["x"])
 
     #If there were runtime warnings, show the messages
     if len(w) > 0:
@@ -337,7 +456,7 @@ def solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, method="hybr", tol=1
     return f_k_nonzero, results
 
 
-def solve_mbar(u_kn_nonzero, N_k_nonzero, f_k_nonzero, solver_protocol=None, verbose=False):
+def solve_mbar(u_kn_nonzero, N_k_nonzero, f_k_nonzero, solver_protocol=None):
     """Solve MBAR self-consistent equations using some sequence of equation solvers.
 
     Parameters
@@ -359,7 +478,7 @@ def solve_mbar(u_kn_nonzero, N_k_nonzero, f_k_nonzero, solver_protocol=None, ver
         The converged reduced free energies.
     all_results : list(dict())
         List of results from each step of solver_protocol.  Each element in
-        list contains the results dictionary from solve_mbar_single()
+        list contains the results dictionary from solve_mbar_once()
         for the corresponding step.
 
     Notes
@@ -371,7 +490,7 @@ def solve_mbar(u_kn_nonzero, N_k_nonzero, f_k_nonzero, solver_protocol=None, ver
     to be zero.
 
     This function calls `solve_mbar_once()` multiple times to achieve
-    converged results.  Generally, a single call to solve_mbar_single()
+    converged results.  Generally, a single call to solve_mbar_once()
     will not give fully converged answers because of limited numerical precision.
     Each call to `solve_mbar_once()` re-conditions the nonlinear
     equations using the current guess.
@@ -383,10 +502,7 @@ def solve_mbar(u_kn_nonzero, N_k_nonzero, f_k_nonzero, solver_protocol=None, ver
     for k, options in enumerate(solver_protocol):
         f_k_nonzero, results = solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, **options)
         all_results.append(results)
-
-    if verbose:
-        print(("Final gradient norm: %.3g" % np.linalg.norm(mbar_gradient(u_kn_nonzero, N_k_nonzero, f_k_nonzero))))
-
+        all_results.append(("Final gradient norm: %.3g" % np.linalg.norm(mbar_gradient(u_kn_nonzero, N_k_nonzero, f_k_nonzero))))
     return f_k_nonzero, all_results
 
 
@@ -497,7 +613,7 @@ def solve_mbar_with_subsampling(u_kn, N_k, f_k, solver_protocol, subsampling_pro
 
         f_k[states_with_samples] = f_k_nonzero
         f_k_nonzero, all_results = solve_mbar(u_kn[states_with_samples], N_k[states_with_samples], f_k[states_with_samples], solver_protocol=solver_protocol)
-    
+
     f_k[states_with_samples] = f_k_nonzero
 
     # Update all free energies because those from states with zero samples are not correctly computed by Newton-Raphson.

pymbar/tests/test_mbar_solvers.py

@@ -74,8 +74,8 @@ def test_protocols():
     fa = fa[1:] - fa[0]
 
     #scipy.optimize.minimize methods, same ones that are checked for in mbar_solvers.py
-    subsampling_protocols = ["L-BFGS-B", "dogleg", "CG", "BFGS", "Newton-CG", "TNC", "trust-ncg", "SLSQP"]
-    solver_protocols = ['hybr', 'lm'] #scipy.optimize.root methods. Omitting methods which do not use the Jacobian
+    subsampling_protocols = ['adaptive', 'L-BFGS-B', 'dogleg', 'CG', 'BFGS', 'Newton-CG', 'TNC', 'trust-ncg', 'SLSQP']
+    solver_protocols = ['adaptive', 'hybr', 'lm'] #scipy.optimize.root methods. Omitting methods which do not use the Jacobian. Adding the custom adaptive protocol.
     for subsampling_protocol in subsampling_protocols:
         for solver_protocol in solver_protocols:
             #Solve MBAR with zeros for initial weights