Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[WIP] Trying to make MBAR work better adaptively when it fails. #442

Draft
wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

[WIP] Trying to make MBAR work better adaptively when it fails. #442

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2910cf2
Trying to make MBAR work better adaptively.
mrshirts Jun 20, 2022
4be3ac8
More changes in adaptive solving
mrshirts Jun 21, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 27 additions & 9 deletions pymbar/mbar.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
##############################################################################
# pymbar: A Python Library for MBAR
#
# Copyright 2017 University of Colorado Boulder
# Copyright 2010-2017 Memorial Sloan-Kettering Cancer Center
# Portions of this software are Copyright (c) 2010-2016 University of Virginia
# Copyright 2016-2022 University of Colorado Boulder
# Copyright 2010-2022 Memorial Sloan-Kettering Cancer Center
# Portions of this software are Copyright (c) 2010-2015 University of Virginia
# Portions of this software are Copyright (c) 2006-2007 The Regents of the University of California. All Rights Reserved.
# Portions of this software are Copyright (c) 2007-2008 Stanford University and Columbia University.
#
Expand Down Expand Up @@ -43,8 +43,6 @@
from pymbar import mbar_solvers
from pymbar.utils import kln_to_kn, kn_to_n, ParameterError, DataError, logsumexp, check_w_normalized

DEFAULT_SOLVER_PROTOCOL = mbar_solvers.DEFAULT_SOLVER_PROTOCOL

# =========================================================================
# MBAR class definition
# =========================================================================
Expand Down Expand Up @@ -106,9 +104,9 @@ def __init__(self, u_kn, N_k, maximum_iterations=10000, relative_tolerance=1.0e-
from above, once ``u_kln`` is phased out.

maximum_iterations : int, optional
Set to limit the maximum number of iterations performed (default 1000)
Set to limit the maximum number of iterations performed (default 10000)
relative_tolerance : float, optional
Set to determine the relative tolerance convergence criteria (default 1.0e-6)
Set to determine the relative tolerance convergence criteria (default 1.0e-7)
verbosity : bool, optional
Set to True if verbose debug output is desired (default False)
initial_f_k : np.ndarray, float, shape=(K), optional
Expand Down Expand Up @@ -292,14 +290,34 @@ def __init__(self, u_kn, N_k, maximum_iterations=10000, relative_tolerance=1.0e-
print(self.f_k)

if solver_protocol is None:
solver_protocol = ({'method': None},)
solver_protocol = mbar_solvers.DEFAULT_SOLVER_PROTOCOL

# Add backup option if only one method is given
# The backup method is 'adaptive', unless 'adaptive' is already the only method selected.
if len(solver_protocol) == 1:
#Add backup option.
fallback = dict()
if solver_protocol[0]['method'] != 'adaptive':
fallback['method'] = 'adaptive'
print("Adding backup method 'adaptive'")
else:
# the backup method is already adaptive, try l-bfgs-b
fallback['method'] = 'BFGS'
print("Adding backup method 'BFGS'")
solver_protocol = solver_protocol + (fallback,)


for solver in solver_protocol:
if 'options' not in solver:
solver['options'] = dict()
solver['options']['maximum_iterations'] = maximum_iterations
if 'continuation' not in solver:
solver['continuation'] = None
if 'maxiter' not in solver['options']:
solver['options']['maxiter'] = maximum_iterations
if 'verbose' not in solver['options']:
# 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.
# to be addressed with logging in mbar 4.0.
solver['options']['verbose'] = self.verbose

self.f_k = mbar_solvers.solve_mbar_for_all_states(self.u_kn, self.N_k, self.f_k, solver_protocol)
Expand Down
147 changes: 98 additions & 49 deletions pymbar/mbar_solvers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,17 @@
from pymbar.utils import ensure_type, logsumexp, check_w_normalized
import warnings

# 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.
# Use Adpative solver as first attempt
DEFAULT_SOLVER_METHOD = "adaptive"
DEFAULT_SOLVER_PROTOCOL = (dict(method=DEFAULT_SOLVER_METHOD,),)
# the 'solver_protocol' dictionary is set up to make it possible to wrap around scipy solvers.
# We introduce our our robust solver as well.

# Note: we use tuples instead of lists to avoid accidental mutability.
#DEFAULT_SOLVER_PROTOCOL = (dict(method="adaptive",),dict(method="L-BFGS-B",))
DEFAULT_SOLVER_PROTOCOL = (dict(method="BFGS",continuation=True),dict(method="adaptive",))
# Uses all of the gradient based methods, but not the non-gradient methods
# ["Nelder-Mead", "Powell", "COBYLA"]",
scipy_minimize_options = ["L-BFGS-B", "dogleg", "CG", "BFGS", "Newton-CG", "TNC", "trust-ncg", "trust-krylov", "trust-exact", "SLSQP"]
scipy_nohess_options = ["L-BFGS-B","BFGS","CG","TNC","SLSQP"] # don't pass a hessian to these to avoid warnings to these.
scipy_root_options = ['hybr','lm'] # only use root options with the hessian.

def validate_inputs(u_kn, N_k, f_k):
"""Check types and return inputs for MBAR calculations.
Expand Down Expand Up @@ -242,11 +245,12 @@ def adaptive(u_kn, N_k, f_k, tol = 1.0e-12, options = None):
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)
tol (float between 0 and 1) - relative tolerance for convergence (default 1.0e-8)

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)
maxiter (int) - maximum number of Newton-Raphson/self-consistent iterations (default 10000:
either NR converges or doesn't, pretty quickly, but the number of self-consistent iterations can take a while.)
verbose (boolean) - verbosity level for debug output

NOTES
Expand All @@ -268,16 +272,17 @@ def adaptive(u_kn, N_k, f_k, tol = 1.0e-12, options = None):
"""
# put the defaults here in case we get passed an 'options' dictionary that is only partial
options.setdefault('verbose',False)
options.setdefault('maximum_iterations',10000)
options.setdefault('maxiter',10000)
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:
if tol < 1.5e-8:
print("Tolerance may be too close to machine precision to converge.")
# keep track of Newton-Raphson and self-consistent iterations
nr_iter = 0
Expand All @@ -286,8 +291,9 @@ def adaptive(u_kn, N_k, f_k, tol = 1.0e-12, options = None):
f_sci = np.zeros(len(f_k), dtype=np.float64)
f_nr = np.zeros(len(f_k), dtype=np.float64)

success = False
# Perform Newton-Raphson iterations (with sci computed on the way)
for iteration in range(0, options['maximum_iterations']):
for iteration in range(0, options['maxiter']):
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]
Expand All @@ -298,11 +304,11 @@ def adaptive(u_kn, N_k, f_k, tol = 1.0e-12, options = None):
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)
gnorm_sci = np.sqrt(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)
gnorm_nr = np.sqrt(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.
Expand Down Expand Up @@ -331,23 +337,33 @@ def adaptive(u_kn, N_k, f_k, tol = 1.0e-12, options = None):
max_delta = np.max(np.abs(f_k[1:]-f_old[1:])/div)
if np.isnan(max_delta) or (max_delta < tol):
doneIterating = True
success = True
warn = "Convergence achieved by change in f with respect to previous guess."
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.')
warn = 'WARNING: All f_k appear to be zero.'
message = warn
success = False
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']))
warn = 'WARNING: Did not converge to within specified tolerance.\n'
if options['maxiter'] <= 0:
warn += "No iterations ran because maximum_iterations was <= 0 ({})!".format(options['maxiter'])
else:
print('max_delta = {:e}, tol = {:e}, maximum_iterations = {:d}, iterations completed = {:d}'.format(max_delta,tol, options['maximum_iterations'], iteration))
return f_k
warn += 'max_delta = {:e}, tol = {:e}, maximum_iterations = {:d}, iterations completed = {:d}'.format(max_delta,tol, options['maxiter'], iteration)

results = dict()
results['success'] = success
results['message'] = warn
results['x'] = f_k

return results

def precondition_u_kn(u_kn, N_k, f_k):
"""Subtract a sample-dependent constant from u_kn to improve precision
Expand Down Expand Up @@ -380,7 +396,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-12, options=None):
def solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, method = "adaptive", tol=1E-12, continuation=None, options=None):
"""Solve MBAR self-consistent equations using some form of equation solver.

Parameters
Expand All @@ -392,16 +408,7 @@ def solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, method="hybr", tol=1
The number of samples in each state for the nonempty states
f_k_nonzero : np.ndarray, shape=(n_states), dtype='float'
The reduced free energies for the nonempty states
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-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.
solver: dict of solver options

Returns
-------
Expand All @@ -422,6 +429,7 @@ def solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, method="hybr", tol=1
For fast but precise convergence, we recommend calling this function
multiple times to polish the result. `solve_mbar()` facilitates this.
"""

u_kn_nonzero, N_k_nonzero, f_k_nonzero = validate_inputs(u_kn_nonzero, N_k_nonzero, f_k_nonzero)
f_k_nonzero = f_k_nonzero - f_k_nonzero[0] # Work with reduced dimensions with f_k[0] := 0
u_kn_nonzero = precondition_u_kn(u_kn_nonzero, N_k_nonzero, f_k_nonzero)
Expand All @@ -435,18 +443,21 @@ def solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, method="hybr", tol=1
hess = lambda x: mbar_hessian(u_kn_nonzero, N_k_nonzero, pad(x))[1:][:, 1:] # Hessian of objective function

with warnings.catch_warnings(record=True) as w:
if method in ["L-BFGS-B", "dogleg", "CG", "BFGS", "Newton-CG", "TNC", "trust-ncg", "SLSQP"]:
if method in ["L-BFGS-B", "CG"]:
if method in scipy_minimize_options:
if method in scipy_nohess_options:
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:
f_k_nonzero = results["x"]
elif method in scipy_root_options:
# find the root in the gradient. Not sure how well this has been tested . . . ?
results = scipy.optimize.root(grad, f_k_nonzero[1:], jac=hess, method=method, tol=tol, options=options)
f_k_nonzero = pad(results["x"])

else:
raise ParameterError(f"Method {method} for solution of free energies not recognized")

# If there were runtime warnings, show the messages
if len(w) > 0:
can_ignore = True
Expand All @@ -460,7 +471,7 @@ def solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, method="hybr", tol=1
# Ensure MBAR solved correctly
w_nk_check = mbar_W_nk(u_kn_nonzero, N_k_nonzero, f_k_nonzero)
check_w_normalized(w_nk_check, N_k_nonzero)
print("MBAR weights converged within tolerance, despite the SciPy Warnings. Please validate your results.")
print("MBAR weights converged within tolerance, despite the SciPy Warnings. However, please validate your results.")

return f_k_nonzero, results

Expand Down Expand Up @@ -498,24 +509,62 @@ def solve_mbar(u_kn_nonzero, N_k_nonzero, f_k_nonzero, solver_protocol=None):
by subtracting off the first component of f_k and fixing that component
to be zero.

This function calls `solve_mbar_once()` multiple times to achieve
converged results. Generally, a single call to solve_mbar_once()
will not give fully converged answers because of limited numerical precision.
This function calls `solve_mbar_once()`. Usually, a single call
will converge to the correct answer. However, there are a number
specific cases that may succeed or fail with a specific algorithm.

The default case is to use the adaptive solver, which alternates
between N-R iteration and self-consistent iteration, depending on
which has the the lowest gradient. If the adaptive solver fails,
it then tries the BFGS algorithm. If one starts with BFSG,
then it changes to adaptive if BFGS fails.. Currently, it restarts
with whatever the initialized values are when it switches algorithms,
with the assumption that if it is failing, the current guess
may not be a workable starting point.

Each call to `solve_mbar_once()` re-conditions the nonlinear
equations using the current guess.

"""
if solver_protocol is None:
solver_protocol = DEFAULT_SOLVER_PROTOCOL
for protocol in solver_protocol:
if protocol['method'] is None:
protocol['method'] = DEFAULT_SOLVER_METHOD

import pdb
pdb.set_trace()
all_fks = []
all_results = []
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_gnorms = []
success = False
for k, solver in enumerate(solver_protocol):
print(f"Trying solution with solver {solver['method']}")
f_k_nonzero_result, results = solve_mbar_once(u_kn_nonzero, N_k_nonzero, f_k_nonzero, **solver)
all_fks.append(f_k_nonzero_result)
all_results.append(results)
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
all_gnorms.append(np.linalg.norm(mbar_gradient(u_kn_nonzero, N_k_nonzero, f_k_nonzero_result)))

if results["success"] == True:
success = True
best_gnorm = all_gnorms[-1]
break
else:
print(f"Failed to reach a solution to within tolerance with {solver['method']}: trying next method")
print(all_gnorms[-1])
if solver["continuation"]:
f_k_nonzero = f_k_nonzero_result
print("Will continue with results from previous method")

if success is True:
print("Solution found within tolerance!")
else:
i_best_gnorm = np.argmin(all_gnorms)
print("No solution found to within tolerance.")
best_method = solver_protocol[i_best_gnorm]['method']
print(f"The solution with the smallest gradient norm is {best_method}")
best_gnorm = all_gnorms[i_best_gnorm]
f_k_nonzero_result = all_fks[i_best_gnorm]
print("Please exercise caution with this solution and consider alternative methods or a different tolerance.")

print(f"Final gradient norm: {best_gnorm:.3g}")

return f_k_nonzero_result, all_results


def solve_mbar_for_all_states(u_kn, N_k, f_k, solver_protocol):
Expand Down
8 changes: 4 additions & 4 deletions pymbar/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -356,17 +356,17 @@ def check_w_normalized(W, N_k, tolerance = 1.0e-4):
which_badcolumns = np.arange(K)[badcolumns]
firstbad = which_badcolumns[0]
raise ParameterError(
'Warning: Should have \sum_n W_nk = 1. Actual column sum for state %d was %f. %d other columns have similar problems' %
(firstbad, column_sums[firstbad], np.sum(badcolumns)))
'Warning: Should have \sum_n W_nk = 1. Actual column sum for state %d was %f. %d other columns have similar problems.' %
(firstbad, column_sums[firstbad], np.sum(badcolumns)) + ' This generally indicates the free energies are not converged.')

row_sums = np.sum(W * N_k, axis=1)
badrows = (np.abs(row_sums - 1) > tolerance)
if np.any(badrows):
which_badrows = np.arange(N)[badrows]
firstbad = which_badrows[0]
raise ParameterError(
'Warning: Should have \sum_k N_k W_nk = 1. Actual row sum for sample %d was %f. %d other rows have similar problems' %
(firstbad, row_sums[firstbad], np.sum(badrows)))
'Warning: Should have \sum_k N_k W_nk = 1. Actual row sum for sample %d was %f. %d other rows have similar problems.' %
(firstbad, row_sums[firstbad], np.sum(badrows)) + ' This generally indicates the free energies are not converged.')

return

Expand Down