############################################################################## # pymbar: A Python Library for MBAR (FES module) # # Copyright 2019 University of Colorado Boulder # # Authors: Michael Shirts # # pymbar is free software: you can redistribute it and/or modify # it under the terms of the MIT License as # # This library is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # MIT License for more details. # # You should have received a copy of the MIT License along with pymbar. ############################################################################## """ A module implementing calculation of free energy surfaces (profiles) from biased simulations. """ import logging import math import numpy as np import pymbar from pymbar.utils import kln_to_kn, ParameterError, DataError, logsumexp from pymbar import timeseries # bunch of imports needed for doing newton optimization of B-splines from scipy.interpolate import BSpline, make_lsq_spline # imports needed for scipy minimizations from scipy.integrate import quad from scipy.optimize import minimize from timeit import default_timer as timer # may remove timing? logger = logging.getLogger(__name__) # ========================================================================= # FES class definition # ========================================================================= class FES: """ Methods for generating free energy surfaces (profile) with statistical uncertainties. Notes ----- Note that this method assumes the data are uncorrelated. Correlated data must be subsampled to extract uncorrelated (effectively independent) samples. References ---------- [1] Shirts MR and Chodera JD. Statistically optimal analysis of samples from multiple equilibrium states. J. Chem. Phys. 129:124105, 2008 http://dx.doi.org/10.1063/1.2978177 [2] Shirts MR and Ferguson AF. Statistically optimal continuous free energy surfaces from umbrella sampling and multistate reweighting https://arxiv.org/abs/2001.01170 """ # ========================================================================= def __init__(self, u_kn, N_k, verbose=False, mbar_options=None, timings=True, **kwargs): """Initialize a free energy surface calculation by performing multistate Bennett acceptance ratio (MBAR) on a set of simulation data from umbrella sampling at K states. Upon initialization, the dimensionless free energies for all states are computed. This may take anywhere from seconds to minutes, depending upon the quantity of data. This also creates an internal mbar object that is used to create the free energy surface. Parameters ---------- u_kn : np.ndarray, float, shape=(K, N_max) ``u_kn[k,n]`` is the reduced potential energy of uncorrelated configuration n evaluated at state ``k``. N_k : np.ndarray, int, shape=(K) ``N_k[k]`` is the number of uncorrelated snapshots sampled from state ``k``. Some may be zero, indicating that there are no samples from that state. We assume that the states are ordered such that the first ``N_k`` are from the first state, the 2nd ``N_k`` the second state, and so forth. This only becomes important for bar -- MBAR does not care which samples are from which state. We should eventually allow this assumption to be overwritten by parameters passed from above, once ``u_kln`` is phased out. mbar_options : dict The following options supported by mbar (see MBAR documentation) maximum_iterations : int, optional relative_tolerance : float, optional verbosity : bool, optional initial_f_k : np.ndarray, float, shape=(K), optional solver_protocol : list(dict) or None, optional, default=None initialize : 'zeros' or 'BAR', optional, Default: 'zeros' x_kindices : array of ints, shape=(K), which state index each sample is from. Examples -------- >>> from pymbar import testsystems >>> (x_n, u_kn, N_k, s_n) = testsystems.HarmonicOscillatorsTestCase().sample(mode='u_kn') >>> fes = FES(u_kn, N_k) """ for key, val in kwargs.items(): logging.warning( "Warning: parameter {:s}={:s} is unrecognized and unused.".format(key, val) ) # Store local copies of necessary data. # N_k[k] is the number of samples from state k, some of which might be # zero. self.N_k = np.array(N_k, dtype=np.int64) self.N = np.sum(self.N_k) # for now, still need to convert from 3 to 2 dim # Get dimensions of reduced potential energy matrix, and convert to KxN # form if needed. if len(np.shape(u_kn)) == 3: # need to set self.K, and it's the second index self.K = np.shape(u_kn)[1] u_kn = kln_to_kn(u_kn, N_k=self.N_k) # u_kn[k,n] is the reduced potential energy of sample n evaluated at # state k self.u_kn = np.array(u_kn, dtype=np.float64) K, N = np.shape(u_kn) if np.sum(self.N_k) != N: raise ParameterError( "The sum of all N_k must equal the total number of samples (length of second dimension of u_kn." ) # Store local copies of other data self.K = K # number of thermodynamic states energies are evaluated at # N = \sum_{k=1}^K N_k is the total number of samples self.N = N # maximum number of configurations # verbosity level -- if True, will print extra debug information self.verbose = verbose if timings: self.timings = True if mbar_options == None: fes_mbar = pymbar.MBAR(u_kn, N_k) else: # if the dictionary does not define the option, add it in required_mbar_options = ( "maximum_iterations", "relative_tolerance", "verbose", "initial_f_k", "solver_protocol", "initialize", "x_kindices", ) for o in required_mbar_options: if o not in mbar_options: mbar_options[o] = None # reset the options that might be none to the default value if mbar_options["maximum_iterations"] is None: mbar_options["maximum_iterations"] = 10000 if mbar_options["relative_tolerance"] is None: mbar_options["relative_tolerance"] = 1.0e-7 if mbar_options["initialize"] is None: mbar_options["initialize"] = "zeros" fes_mbar = pymbar.MBAR( u_kn, N_k, maximum_iterations=mbar_options["maximum_iterations"], relative_tolerance=mbar_options["relative_tolerance"], verbose=mbar_options["verbose"], initial_f_k=mbar_options["initial_f_k"], solver_protocol=mbar_options["solver_protocol"], initialize=mbar_options["initialize"], x_kindices=mbar_options["x_kindices"], ) self.mbar = fes_mbar # TODO: eliminate this call - it's causing problems # with deepcopy, needed in some cases, since you can't # copy the np.random module # self._random = np.random # self._seed = None if self.verbose: logger.info("FES initialized") # TODO: see above about not storing np.random # @property # def seed(self): # return self._seed # # def reset_random(self): # self._random = np.random # self._seed = None def generate_fes( self, u_n, x_n, fes_type="histogram", histogram_parameters=None, kde_parameters=None, spline_parameters=None, n_bootstraps=0, seed=-1, ): """ Given an intialized MBAR object, a set of points, the desired energies at that point, and a method, generate an object that contains the FES information. Parameters ---------- u_n : np.ndarray, float, shape=(N) u_n[n] is the reduced potential energy of snapshot n of state for which the FES is to be computed. Often, it will be one of the states in of u_kn, used in initializing the FES object, but we want to allow more generality. x_n : np.ndarray, float, shape=(N,D) x_n[n] is the d-dimensional coordinates of the samples, where D is the reduced dimensional space. fes_type : str options = 'histogram', 'kde', 'spline' histogram_parameters: dictionary Input dictionary with the following keys: bin_edges: list of ndim np.ndarray, each array shaped ndum+1 The bin edges. Compatible with `bin_edges` output of np.histogram. kde_parameters all the parameters from sklearn (https://scikit-learn.org/stable/modules/generated/sklearn.neighbors.KernelDensity.html). Defaults will be used if nothing changed. spline_parameters 'spline_weights' which type of fit to use: 'biasedstates' - sum of log likelihood over all weighted states 'unbiasedstate' - log likelihood of the single unbiased state 'simplesum' - sum of log likelihoods from the biased simulation. Essentially equivalent to vFEP (York et al.) 'optimization_algorithm': 'Custom-NR' - a custom Newton-Raphson that is particularly fast for close data, but can fail 'Newton-CG' - scipy Newton-CG, only Hessian based method that works correctly because of data ordering. ' ' - scipy gradient based methods that work, but are generally slower (CG, BFGS, L-LBFGS-B, TNC, SLSQP) 'fkbias' : array of functions Return the Kth bias potential for each function 'nspline' : int Number of spline points 'kdegree' : int Degree of the spline. Default is cubic ('3') 'objective' : string 'ml','map' # whether to fit the maximum likelihood or the maximum a posteriori n_bootstraps : int, 0 or > 1, Default: 0 Number of bootstraps to create an uncertainty estimate. If 0, no bootstrapping is done. Required if one uses uncertainty_method = 'bootstrap' in get_fes seed : int, Default = -1 Set the randomization seed. Settting should get the randomization (assuming the same calls are made in the same order) to return the same numbers. This is local to this class and will not change any other random objects. Returns ------- dict, optional if 'timings' is set to True in __init__, returns the time taken to generate the FES Notes ----- * fes_type = 'histogram': * This method works by computing the free energy of localizing the system to each bin for the given potential by aggregating the log weights for the given potential. * To estimate uncertainties, the NxK weight matrix W_nk is augmented to be Nx(K+nbins) in order to accomodate the normalized weights of states . . . * the potential is given by u_n within each bin and infinite potential outside the bin. The uncertainties with respect to the bin of lowest free energy are then computed in the standard way. Examples -------- >>> # Generate some test data >>> from pymbar import testsystems >>> from pymbar import FES >>> x_n, u_kn, N_k, s_n = testsystems.HarmonicOscillatorsTestCase().sample(mode='u_kn',seed=0) >>> # Select the potential we want to compute the FES for (here, condition 0). >>> u_n = u_kn[0, :] >>> # Sort into nbins equally-populated bins >>> nbins = 10 # number of equally-populated bins to use >>> import numpy as np >>> N_tot = N_k.sum() >>> x_n_sorted = np.sort(x_n) # unroll to n-indices >>> bins = np.append(x_n_sorted[0::int(N_tot/nbins)], x_n_sorted.max()+0.1) >>> bin_widths = bins[1:] - bins[0:-1] >>> # Compute FES for these unequally-sized bins. >>> fes = FES(u_kn, N_k) >>> histogram_parameters = dict() >>> histogram_parameters['bin_edges'] = [bins] >>> _ = fes.generate_fes(u_n, x_n, fes_type='histogram', histogram_parameters = histogram_parameters) >>> results = fes.get_fes(x_n) >>> f_i = results['f_i'] >>> for i, x_n in enumerate(x_n): # doctest: +SKIP >>> print(x_n, f_i[i]) # doctest: +SKIP >>> mbar = fes.get_mbar() >>> print(mbar.f_k) # doctest: +SKIP >>> print(N_k) # doctest: +SKIP """ result_vals = dict() # for results we may want to return. self.fes_type = fes_type # eventually, we just want the desired energy of each sample. For now, we allow conversion # from older 2d format (K,Nmax instead of N); this is data SAMPLED from # each k, not the energy at different K. if len(np.shape(u_n)) == 2: u_n = pymbar.mbar.kn_to_n(u_n, N_k=self.N_k) self.u_n = u_n if seed >= 0: np.random.seed(seed) # TODO: see above about storing np.random # if seed >= 0: # # Set a better seeded random state # self._random = np.random.RandomState(seed=seed) # self._seed = seed # we need to save this for calculating uncertainties. if not np.issubdtype(type(n_bootstraps), np.integer) or n_bootstraps == 1: raise ValueError( f"n_bootstraps must be an integer of 0 or >=2, it was set to {n_bootstraps}" ) self.n_bootstraps = n_bootstraps if self.timings: start = timer() self.fes_function = list() # set some variables before bootstrapping loop. self.mc_data = None # we have not sampled MC data yet. if self.fes_type == "histogram": self._setup_fes_histogram(histogram_parameters) elif fes_type == "kde": self._setup_fes_kde(kde_parameters) elif fes_type == "spline": self._setup_fes_spline(spline_parameters) else: raise ParameterError("fes_type {:s} is not defined!".format(fes_type)) N_k = self.mbar.N_k K = self.mbar.K N = np.sum(N_k) for b in range(n_bootstraps + 1): # generate bootstrap samples. # we bootstrap from each simulation separately. if b == 0: # the default bootstrap_indices = np.arange(0, N) mbar = self.mbar x_nb = x_n else: index = 0 for k in range(K): # TODO: address issue with storinig np.random in self bootstrap_indices[index : index + N_k[k]] = index + np.random.randint( 0, N_k[k], size=N_k[k] ) index += N_k[k] # recompute MBAR. mbar = pymbar.MBAR( self.u_kn[:, bootstrap_indices], self.N_k, initial_f_k=self.mbar.f_k ) x_nb = x_n[bootstrap_indices] # Compute unnormalized log weights for the given reduced potential # u_n, needed for all methods. log_w_nb = mbar._computeUnnormalizedLogWeights(self.u_n[bootstrap_indices]) # calculate a few other things used for multiple methods max_log_w_nb = np.max(log_w_nb) # we need to solve underflow. w_nb = np.exp(log_w_nb - max_log_w_nb) w_nb = w_nb / np.sum(w_nb) # normalize the weights # normalized weights for all states. w_knb = np.exp(mbar.Log_W_nk) if b == 0: self.w_n = w_nb self.w_kn = w_knb if self.fes_type == "histogram": # not clear if need to pass both w_nb and log_w_nb, but saves some processing self._generate_fes_histogram(b, x_nb, w_nb, log_w_nb) elif self.fes_type == "kde": self._generate_fes_kde(b, x_nb, w_nb) elif self.fes_type == "spline": self._generate_fes_spline(b, x_nb, w_nb) # we put the timings outside, since the switch / common stuff is really # low. if self.timings: end = timer() result_vals["timing"] = end - start return result_vals # should we return results under some other conditions? def _setup_fes_histogram(self, histogram_parameters): """ Does initial processsing of histogram_parameters Parameters ---------- histogram_parameters : dict() A options:values dictonary for parameters to create the FES using histogramss Returns ------- None Internally, creates histogram_data object and histogram_datas to store information generated by generate_fes """ if "bin_edges" not in histogram_parameters: raise ParameterError( "histogram_parameters['bin_edges'] cannot be undefined with fes_type = histogram" ) # code expects that the bin edges consist in an list of # arrays. But for 1D, we should be able to just input a single array. if len(np.shape(histogram_parameters["bin_edges"])) == 1: histogram_parameters["bin_edges"] = [histogram_parameters["bin_edges"]] self.histogram_parameters = histogram_parameters self.histogram_data = None if self.n_bootstraps > 0: self.histogram_datas = list() else: self.histogram_datas = None def _generate_fes_histogram(self, b, x_n, w_nb, log_w_nb): """ Parameters ---------- b : int the bootstrap sample this is on (for non-bootstrap methods, will be 0) x_n : ndarray, length self.N The data point set used in this bootstrap log_w_n : np.ndarray, float, shape=(self.N) Normalized log weights for each sample for the state in which we want the FES (usually, the unbiased state). Doing it outside the loop to avoid redoing it each time. Returns ------- None Doesn't return, rather adds histogram_data (for b==0) to self, or for b>0, adds histogram_data to the array self.histogram_data for further processing by get_fes. """ # store the data that will be regenerated each time. # We will not try to regenerate the bin locations each time, # as that would make it hard to calculate uncertainties. # We will just recalculate the populations of each bin. histogram_parameters = self.histogram_parameters bins = histogram_parameters["bin_edges"] dims = len(bins) histogram_data = {} histogram_data["dims"] = dims # store the dimensionality for checking later. histogram_data["bins"] = bins # save for other functions. # create the bins from the data. # it's a 1D array, instead of a Nx1 array. Reshape. if len(np.shape(x_n)) == 1: x_n = x_n.reshape(-1, 1) bin_n = np.zeros(x_n.shape, int) bin_length = np.zeros(dims, int) for d in range(dims): bin_length[d] = len(bins[d]) # bins returns 0 as out of bin. We want to use -1 as out # of bin instead. bin_n[:, d] = np.digitize(x_n[:, d], bins[d]) - 1 histogram_data["bin_n"] = bin_n # bin counts in each bin # number each of the bins with samples with an integer # Assign each sample this integer label (in addition to the tuple label) nonzero_bins = list() bin_label = {} sample_label = np.zeros(self.N, int) for n in range(self.N): bin = tuple(bin_n[n]) # which bin (labeled N-D) sample n is in if np.any(bin_n[n] < 0): # this sample is out of grid sample_label[n] = -1 else: # how do we label the bins? if N-dimensional: # bins[0] + bins[1]*bin_length[1]**1 + bins[2]*bin_length[2]**2 sample_label[n] = int( np.sum([bin_n[n][d] * bin_length[d] ** d for d in range(dims)]) ) if bin not in nonzero_bins: nonzero_bins.append(bin) bin_label[bin] = sample_label[n] histogram_data["nonzero_bins"] = nonzero_bins histogram_data["sample_label"] = sample_label # problem with bins above: # # all over bootstraps, not all nonzero bins will occur, as in some bootstraps, # some labels won't appear. # However, if they appear in the list of nonzero bins, they will have the same labels. # that may be OK, since we are only really interested in # uncertainties for the ones that have nonzero samples in the original list. # Need to come up with an order for the labels for all bootstraps, so free energies # are always assigned. if b == 0: bin_order = {} i = 0 for bv in bin_label.values(): if bv not in bin_order: bin_order[bv] = i i += 1 histogram_data["bin_order"] = bin_order histogram_data["bin_label"] = bin_label else: bin_order = self.histogram_data["bin_order"] # Compute the free energies for the histogram bins # with samples. We cannot calculate free energes # for bins w/o samples. f_i = np.zeros(len(bin_label), np.float64) for i, label in enumerate(bin_label.values()): # Get linear n-indices of samples that fall in this bin. indices = np.where(sample_label == label) # Sanity check. if len(indices) == 0: raise DataError( f"WARNING: bin {i} has no samples -- all bins must have at least one sample." ) # Compute dimensionless free energy of occupying state i. f_i[bin_order[label]] = -logsumexp(log_w_nb[indices]) # store the free energies for this bin histogram_data["f"] = f_i if b == 0: self.histogram_data = histogram_data else: self.histogram_datas.append(histogram_data) def _setup_fes_kde(self, kde_parameters): """ Does initial processsing of kde_parameters Parameters ---------- kde_parameters : dict() A options:values dictonary for parameters to create the FES using the kernel density approach. Parameters are passsed on to sklearn KernelDensity. Returns ------- None Internally, creates kde object and kdes list of kde ibjects to store information generated by generate_fes_kde """ try: from sklearn.neighbors import KernelDensity except ImportError: raise ImportError( "Cannot use 'kde' type FES without the scikit-learn module. Could not import sklearn" ) kde = KernelDensity() # get the default params to set them. kde_defaults = kde.get_params() for k in kde_defaults: if k in kde_parameters: kde_defaults[k] = kde_parameters[k] # make sure we didn't pass any arguments that DON'T belong here for k in kde_parameters: if k not in kde_defaults: raise ParameterError("Warning: {:s} is not a parameter in KernelDensity".format(k)) kde.set_params(**kde_defaults) self.kde_parameters = kde_parameters if self.n_bootstraps > 0: self.kdes = list() else: self.kdes = None self.kde = kde def _generate_fes_kde(self, b, x_n, w_n): """ Given an fes object with the kde data set up, determine the information necessary to define a FES using a kernel density approximation Parameters ---------- b : int Which bootstrap this is: b==0 is the initial value with the "untouched" data. x_n : np.ndarray, float, shape=(N,D) x_n[n] is the d-dimensional coordinates of the samples, where D is the reduced dimensional space. w_n : np.ndarray, float, shape=(sself.N) Weights for each sample for the state in which we want the FES (usually, the unbiased state) Returns ------- None Data is stored in ``self.kde`` or ``self.kdes`` (for bootstrap replicates). """ # reshape data if needed. # it's a 1D array, instead of a Nx1 array. Reshape. if len(np.shape(x_n)) == 1: x_n = x_n.reshape(-1, 1) # TODO: figure out if this should be called each bootstrap run or not (shouldn't cost?) # basically, just need to have KernelDensity defined here. try: from sklearn.neighbors import KernelDensity except ImportError: raise ImportError( "Cannot use 'kde' type FES without the scikit-learn module. Could not import sklearn" ) if b > 0: kde = KernelDensity() # need to create a new one so won't get refit # Will take a refactor to get this correct for pylint params = self.kde.get_params() # pylint: disable=access-member-before-definition kde.set_params(**params) else: kde = self.kde kde.fit(x_n, sample_weight=self.w_n) if b > 0: self.kdes.append(kde) def _setup_fes_spline(self, spline_parameters): """ Does initial processsing of spline_parameters Parameters ---------- spline_parameters : dict() A options:values dictonary for parameters to create the FES using the spline approach. Parameters are explained in docstring of Returns ------- None Internally, creates spline_data object and spline_datas list of spline_Data objects to store information generated by generate_fes_spline """ if "objective" not in spline_parameters: spline_parameters["objective"] = "ml" # default objective = spline_parameters["objective"] if objective not in ["ml", "map"]: raise ParameterError( "objective may only be 'ml' or 'map': you have selected {:s}".format(objective) ) if ( objective == "ml" ): # we are doing maximum likelihood minimization, shouldn't be any prior defined if "map_data" in spline_parameters: if spline_parameters["map_data"] is not None: raise ParameterError( "if 'objective' is 'ml' then 'map_data' structure containing priors should not be included" ) # Fill them in with Nones, so the code logic can proceed spline_parameters["map_data"] = {} spline_parameters["map_data"]["logprior"] = None spline_parameters["map_data"]["dlogprior"] = None spline_parameters["map_data"]["ddlogprior"] = None elif objective == "map": if "map_data" not in spline_parameters: raise ParameterError( "if 'objective' is 'map' you must include 'map_data' structure" ) elif spline_parameters["map_data"] == None: raise ParameterError("MAP data must be defined if objective is MAP") else: map_data = spline_parameters["map_data"] if map_data["logprior"] == None: raise ParameterError("log prior must be included if objective is MAP") if map_data["dlogprior"] == None: raise ParameterError("d(log prior) must be included if objective is MAP") if map_data["ddlogprior"] == None: raise ParameterError("d^2(log prior) must be included if objective is MAP") # we need to intialize our function starting point if spline_parameters["optimization_algorithm"] != "Custom-NR": # we are passing it on to scipy if "optimize_options" not in spline_parameters: spline_parameters["optimize_options"] = { "disp": True, "ftol": 10 ** (-7), "xtol": 10 ** (-7), } if "tol" in spline_parameters["optimize_options"]: # scipy doesn't like 'tol' within options spline_parameters["scipy_tol"] = spline_parameters["optimize_options"]["tol"] spline_parameters["optimize_options"].pop("tol", None) else: spline_parameters["scipy_tol"] = None # this is just the default anyway. if spline_parameters["optimization_algorithm"] not in [ "Newton-CG", "CG", "BFGS", "L-BFGS-B", "TNC", "SLSQP", ]: raise ParameterError( "Optimization method {:s} is not supported".format( spline_parameters["optimization_algorithm"] ) ) else: if "optimize_options" not in spline_parameters: spline_parameters["optimize_options"] = dict() if "gtol" not in spline_parameters["optimize_options"]: spline_parameters["optimize_options"]["tol"] = 10 ** (-7) self.spline_parameters = spline_parameters xinit, yinit = self._get_initial_spline_points() self.spline_data = self._get_initial_spline(xinit, yinit) if self.n_bootstraps > 0: self.fes_functions = list() else: self.fes_functions = None def _get_initial_spline_points(self): """ Uses information from spline_parameters to construct initial points to create a spline frmo which to start the minimization. Parameters ---------- None Returns ------- xinit : ndarray, float x-values of spline to be fit for start of minimizaton yinit : ndarray, float, shape y-values of spline to be fit for start of minimizaton """ spline_parameters = self.spline_parameters nspline = spline_parameters["nspline"] kdegree = spline_parameters["kdegree"] xrange = spline_parameters["xrange"] if spline_parameters["spline_initialize"] == "bias_free_energies": initvals = self.mbar.f_k # initialize to the bias free energies if "bias_centers" in spline_parameters: # if we are provided bias center, use them bias_centers = spline_parameters["bias_centers"] sort_indices = np.argsort(bias_centers) K = self.mbar.K if K < 2 * nspline: noverfit = int(np.round(K / 2)) tinit = np.zeros(noverfit + kdegree + 1) tinit[0:kdegree] = xrange[0] tinit[kdegree : noverfit + 1] = np.linspace( xrange[0], xrange[1], num=noverfit + 1 - kdegree, endpoint=True ) tinit[noverfit + 1 : noverfit + kdegree + 1] = xrange[1] # problem: bin centers might not actually be sorted. binit = make_lsq_spline( bias_centers[sort_indices], initvals[sort_indices], tinit, k=kdegree ) xinit = np.linspace(xrange[0], xrange[1], num=2 * nspline) yinit = binit(xinit) else: xinit = bias_centers[sort_indices] yinit = initvals[sort_indices] else: # assume equally spaced bias scenters xinit = np.linspace(xrange[0], xrange[1], self.mbar.K + 1)[1:-1] yinit = initvals elif spline_parameters["spline_initialize"] == "explicit": if "xinit" in spline_parameters: xinit = spline_parameters["xinit"] else: raise ParameterError( "spline_initialize set as explicit, but no xinit array specified" ) if "yinit" in spline_parameters: yinit = spline_parameters["yinit"] else: raise ParameterError( "spline_initialize set as explicit, but no yinit array specified" ) elif spline_parameters["spline_initialize"] == "zeros": # initialize to zero xinit = np.linspace(xrange[0], xrange[1], nspline + kdegree) yinit = np.zeros(len(xinit)) else: spline_initialization = spline_parameters["spline_initialize"] raise ParameterError(f"Initialization type {spline_initialization} not recognized") return xinit, yinit def _get_initial_spline(self, xinit, yinit): """ Uses information from spline_parameters to construct initial points to create a spline frmo which to start the minimization. Parameters ---------- xinit : ndarray, float: x-values of spline to fit yinit : ndarray, float: y-values of spline to fit Returns ------- spline_data : dict() of information used in optimizing the spline parameters """ spline_data = {} spline_parameters = self.spline_parameters kdegree = spline_parameters["kdegree"] # degree of the spline nspline = spline_parameters["nspline"] # number of spline points. xrange = spline_parameters["xrange"] # this is kind of duplicated: figure out how to simplify. # t has to be of size nsplines + kdegree + 1 t = np.zeros(nspline + kdegree + 1) t[0:kdegree] = xrange[0] t[kdegree : nspline + 1] = np.linspace( xrange[0], xrange[1], num=nspline + 1 - kdegree, endpoint=True ) t[nspline + 1 : nspline + kdegree + 1] = xrange[1] # initial fit function # problem: bin centers might not actually be sorted. Should data # getting to here be already sorted? sort_indices = np.argsort(xinit) b = make_lsq_spline(xinit[sort_indices], yinit[sort_indices], t, k=kdegree) # one, since the FES is only determined up to a constant. b.c = b.c - b.c[0] # We zero out the first # the bspline coefficients are the variables we care about. xi = b.c[1:] xold = xi.copy() # The function is \sum_n F(phi(x_n)) + \sum_k ln \int exp(-F(xi) - u_k(xi)) dxi # if we assume bsplines are of the form f(x) = a*b_i(x), then # dF/dtheta is simply the basis function that has support over that # region of space # we now need the derivative of the function WRT the coefficients. Doesn't change knots or degree. # A vector function that is db_c = list() for i in range(nspline): dc = np.zeros(nspline) dc[i] = 1.0 db_c.append(BSpline(b.t, dc, b.k)) # OK, we've defined the derivatives. # same for the next execution. Not sure if best time to save it. # We also construct integration ranges for the derivatives, since no point in integrating when # the function is zero. xrangei = np.zeros([nspline, 2]) for i in range(0, nspline): xrangei[i, 0] = t[i] xrangei[i, 1] = t[i + kdegree + 1] # set integration ranges for derivative products; saves time on # integration. xrangeij = np.zeros([nspline, nspline, 2]) for i in range(0, nspline): for j in range(0, nspline): xrangeij[i, j, 0] = np.max([xrangei[i, 0], xrangei[j, 0]]) xrangeij[i, j, 1] = np.min([xrangei[i, 1], xrangei[j, 1]]) spline_data["initial_coefficients"] = xi spline_data["bspline_derivatives"] = db_c spline_data["bspline"] = b spline_data["xrangei"] = xrangei spline_data["xrangeij"] = xrangeij return spline_data def _generate_fes_spline(self, b, x_n, w_n): """ Given an fes object with the spline set up, determine the information necessary to define a FES. Parameters ---------- b : int Which bootstrap this is: b==0 is the initial value with the "untouched" data. x_n : np.ndarray, float, shape=(N,D) x_n[n] is the d-dimensional coordinates of the samples, where D is the reduced dimensional space. w_n : np.ndarray, float, shape=(sself.N) Weights for each sample for the state in which we want the FES (usually, the unbiased state) Returns ------- None Data is stored in self.fes_function or self.fes_functions (for bootstrap replicates). """ if b == 0: xi = self.spline_data["initial_coefficients"].copy() else: xi = self.spline_data["first_coefficients"].copy() spline_parameters = self.spline_parameters spline_data = self.spline_data func = self._bspline_calculate_f grad = self._bspline_calculate_g hess = self._bspline_calculate_h if spline_parameters["optimization_algorithm"] != "Custom-NR": if spline_parameters["optimization_algorithm"] == "Newton-CG": hessian = hess else: hessian = None spline_args = (x_n, w_n) results = minimize( func, xi, args=spline_args, method=spline_parameters["optimization_algorithm"], jac=grad, tol=spline_parameters["scipy_tol"], hess=hess, options=spline_parameters["optimize_options"], ) bspline = self._val_to_spline(results["x"], form="log") # TODO: where is this saved? savexi = results["x"] else: if "gtol" in spline_parameters["optimize_options"]: tol = spline_parameters["optimize_options"]["gtol"] elif "tol" in spline_parameters["optimize_options"]: tol = spline_parameters["optimize_options"]["tol"] # should come up with better way to make sure it passes the first time. dg = tol * 1e10 firsttime = True while dg > tol: # until we reach the tolerance. f = func(xi, *spline_args) # we need some error handling: if we stepped too far, we should go back # still not great error handling. Requires something # close. if firsttime: firsttime = False else: count = 0 # we went too far! Pull back. # pylint: disable=used-before-assignment,undefined-variable while (f >= fold * (1.1) and count < 5) or (np.isinf(f)): f = fold # let's not step as far: dx = 0.9 * dx # step back 90% of dx xi = xold - dx xold = xi.copy() f = func(xi, *spline_args) count += 1 fold = f xold = xi.copy() g = grad(xi, *spline_args) h = hess(xi, *spline_args) # now find the new point. # x_n+1 = x_n - f''(x_n)^-1 f'(x_n) # which we solve more stably as: # x_n - x_n+1 = f''(x_n)^-1 f'(x_n) # f''(x_n)(x_n-x_n+1) = f'(x_n) # solution is dx = x_n-x_n+1 dx = np.linalg.lstsq(h, g, rcond=None)[0] xi = xold - dx if spline_parameters["optimize_options"]["disp"]: dg = np.sqrt(np.dot(g, g)) logger.info("f = {:.10f}. gradient norm = {:.10f}".format(f, np.sqrt(dg))) bspline = self._val_to_spline(xi, form="log") savexi = xi if b == 0: nparameters = len(savexi) minus_log_likelihood = func(savexi, *spline_args) self.spline_data["first_coefficients"] = savexi # now store BIC and AIC results = self._calculate_information_criteria( nparameters, minus_log_likelihood, self.N ) self.spline_data["aic"] = results["aic"] self.spline_data["bic"] = results["bic"] if b == 0: self.fes_function = bspline else: self.fes_functions.append(bspline) @staticmethod def _calculate_information_criteria(nparameters, minus_log_likelihood, N): """ Calculate and store various informaton criterias Parameters ---------- nparameters : int, number of parameters in the model minus_log_likelihood : float, minus the log likelihood of the model. N : int, number of samples Results ------- results : dictionary Keys are "aic" (Akaike information critera) and "bic" (Bayesian information criteria) """ results = {} # calculate the AIC # formula is: 2(number of parameters) - 2 * loglikelihood results["aic"] = 2 * nparameters + 2 * minus_log_likelihood # calculate the BIC # formula is: ln(number of data points) * (number of parameters) - 2 ln (likelihood at maximum) results["bic"] = 2 * np.log(N) * nparameters + 2 * minus_log_likelihood # potential problems: we don't compute the full log likelihood currently since we # exclude the potential energy part - will have to see what this is in reference to. # this shouldn't be too much of a problem, since we are interested in choosing BETWEEN models. return results def get_information_criteria(self, type="akaike"): """ returns the Akaike or Bayesian Informatiton Criteria for the model if it exists. Parameters ---------- type : string either 'Akaike' (or 'akaike' or 'aic') or 'Bayesian' (or 'bayesian' or 'bic') Returns ------- float : value of information criteria """ if self.fes_type != "spline": raise ParameterError( "Information criteria currently only defined for spline approaches, you are currently using {:s}".format( type ) ) if type in ["akaike", "Akaike", "AIC", "aic"]: return self.spline_data["aic"] elif type in ["bayesian", "Bayesian", "BIC", "bic"]: return self.spline_data["bic"] else: raise ParameterError("Information criteria of type '{:s}' not defined".format(type)) def get_fes( self, x, reference_point="from-lowest", fes_reference=None, uncertainty_method=None ): """ Returns values of the FES at the specified x points. Parameters ---------- x : numpy.ndarray of D dimensions, where D is the dimensionality of the FES defined. reference_point : str, optional Method for reporting values and uncertainties (default : 'from-lowest') * 'from-lowest' - the uncertainties in the free energy difference with lowest point on FES are reported * 'from-specified' - same as from lowest, but from a user specified point * 'from-normalization' - the normalization \\sum_i p_i = 1 is used to determine uncertainties spread out through the FES * 'all-differences' - the nbins x nbins matrix df_ij of uncertainties in free energy differences is returned instead of df_i uncertainty_method : str, optional Method for computing uncertainties (default: None) fes_reference: an N-d point specifying the reference state. Ignored except with uncertainty method ``from_specified`` Returns ------- dict 'f_i' : np.ndarray, float, shape=(K) result_vals['f_i'][i] is the dimensionless free energy of the x_i point, relative to the reference point 'df_i' : np.ndarray, float, shape=(K) result_vals['df_i'][i] is the uncertainty in the difference of x_i with respect to the reference point Only included if uncertainty_method is not None """ # if it's not an array, make it one. x = np.array(x) # it's a 1D array, instead of a Nx1 array. Reshape. if len(np.shape(x)) <= 1: x = x.reshape(-1, 1) if reference_point == "from-specified" and fes_reference is None: logger.info( "No reference state specified for FES, using uncertainty_method = from-specified" ) if self.fes_type == "histogram": result_vals = self._get_fes_histogram( x, reference_point, fes_reference, uncertainty_method ) elif self.fes_type == "kde": # TODO: check dimensionality here result_vals = self._get_fes_kde(x, reference_point, fes_reference, uncertainty_method) elif self.fes_type == "spline": result_vals = self._get_fes_spline( x, reference_point, fes_reference, uncertainty_method ) else: raise ParameterError("fes_type {self.fes_type} is not supported") return result_vals def get_mbar(self): """return the MBAR object being used by the FES Returns ------- MBAR object """ if self.mbar is not None: return self.mbar else: raise DataError("MBAR in the FES object is not initialized, cannot return it.") def get_kde(self): """return the KernelDensity object if it exists. Returns ------- sklearn KernelDensity object """ if self.fes_type == "kde": if self.kde != None: return self.kde else: raise ParameterError( "Can't return the KernelDensity object because kde not yet defined" ) else: raise ParameterError("Can't return the KernelDensity object because fes_type != kde") def _get_fes_histogram( self, x, reference_point="from-lowest", fes_reference=None, uncertainty_method=None ): """ Returns values of the FES at the specified x points for histogram FESs. Parameters ---------- x : numpy.ndarray of D dimensions, where D is the dimensionality of the FES defined. reference_point : str, optional Method for reporting values and uncertainties (default : 'from-lowest') * 'from-lowest' - the uncertainties in the free energy difference with lowest point on FES are reported * 'from-specified' - same as from lowest, but from a user specified point * 'from-normalization' - the normalization \\sum_i p_i = 1 is used to determine uncertainties spread out through the FES * 'all-differences' - the nbins x nbins matrix df_ij of uncertainties in free energy differences is returned instead of df_i uncertainty_method : str, optional Method for computing uncertainties (default: None) fes_reference: an N-d point specifying the reference state. Ignored except with uncertainty method ``from_specified`` Returns ------- dict 'f_i' : np.ndarray, float, shape=(K) result_vals['f_i'][i] is the dimensionless free energy of the x_i point, relative to the reference point 'df_i' : np.ndarray, float, shape=(K) result_vals['df_i'][i] is the uncertainty in the difference of x_i with respect to the reference point Only included if uncertainty_method is not None """ if np.shape(x)[1] != self.histogram_data["dims"]: raise DataError( "query coordinates have inconsistent dimension with the data the FES is fit to." ) if ( uncertainty_method not in ["bootstrap", "analytical"] and uncertainty_method is not None ): raise ParameterError(f"Uncertainty_method {uncertainty_method} is not a valid option") if uncertainty_method == "bootstrap": if self.histogram_datas is None: raise ParameterError( "Can't calculate uncertainties via bootstrap if bootstrapping was not performed when running get_fes" ) else: n_bootstraps = len(self.histogram_datas) # set up structure for return data result_vals = {} # retrive data for each bootstrap histogram_data = self.histogram_data histogram_datas = self.histogram_datas bins = histogram_data["bins"] dims = histogram_data["dims"] bin_order = histogram_data["bin_order"] nbins = len(bin_order) # figure out which bins the values are in. if dims == 1: # what gridpoint does each x fall into? # -1 and nbinsperdim are out of range loc_indices = np.digitize(x, bins[0]) - 1 else: loc_indices = np.zeros([len(x), dims], dtype=int) for d in range(dims): # -1 and nbinsperdim are out of range loc_indices[:, d] = np.digitize(x[:, d], bins[d]) - 1 # figure out which grid point the fes_reference is at if reference_point == "from-specified": if fes_reference is not None: if dims == 1: # make it a list for reduced code duplication. fes_reference = [fes_reference] fes_ref_grid = np.zeros([dims], dtype=int) for d in range(dims): # -1 and nbins_per_dim are out of range fes_ref_grid[d] = np.digitize(fes_reference[d], bins[d]) - 1 if fes_ref_grid[d] == -1 or fes_ref_grid[d] == len(bins[d]): raise ParameterError( "Specified reference point coordinate {:f} in dim {:d} grid point is out of the FES region [{:f},{:f}]".format( fes_ref_grid[d], d, np.min(bins[d]), np.max(bins[d]) ) ) else: raise ParameterError("Specified reference point for FES not given") if reference_point in ["from-lowest", "from-specified", "all-differences"]: if reference_point == "from-lowest": # Determine free energy with lowest free energy to serve as reference point j = histogram_data["f"].argmin() elif reference_point == "from-specified": # find the label of this bin ref_bin_label = histogram_data["bin_label"][tuple(fes_ref_grid)] # then find the invariant free energy index of this bin j = bin_order[ref_bin_label] elif reference_point == "all-differences": raise ParameterError( "reference point method of 'all-differences' is not yet supported for histogram " "FES types (not implemented)" ) f_i = histogram_data["f"] - histogram_data["f"][j] # now calculate uncertainty for these reference_method approaches. df_i = np.zeros(len(histogram_data["f"]), np.float64) if uncertainty_method == "analytical": K = self.mbar.K # Compute uncertainties in free energy at each gridpoint by # forming matrix of W_nk. N_k = np.zeros([K + nbins], np.int64) N_k[0:K] = self.mbar.N_k W_nk = np.zeros([self.mbar.N, K + nbins], np.float64) W_nk[:, 0:K] = np.exp(self.mbar.Log_W_nk) log_w_n = self.mbar._computeUnnormalizedLogWeights(self.u_n) # loop over the nonzero bins for label in histogram_data["bin_label"].values(): # Get indices of samples that fall in this bin. indices = np.where(histogram_data["sample_label"] == label) # Compute normalized weights for this state. flabel = bin_order[label] W_nk[indices, K + flabel] = np.exp( log_w_n[indices] + self.histogram_data["f"][flabel] ) # Compute asymptotic covariance matrix using specified # method. Theta_ij = self.mbar._computeAsymptoticCovarianceMatrix(W_nk, N_k) # Compute uncertainties with respect to difference in free energy # from this state j. for i in range(nbins): df_i[i] = math.sqrt( Theta_ij[K + i, K + i] + Theta_ij[K + j, K + j] - 2.0 * Theta_ij[K + i, K + j] ) elif uncertainty_method == "bootstrap": fall = np.zeros([len(histogram_data["f"]), n_bootstraps]) for b in range(n_bootstraps): h = histogram_datas[b] # just to make this shorter fall[:, b] = h["f"] - h["f"][j] df_i = np.std(fall, axis=1) elif reference_point == "from-normalization": # Determine uncertainties from normalization that \sum_i p_i = 1. # need to reimplement this . . . maybe. # # Compute bin probabilities p_i # p_i = np.exp(-self.fbin - logsumexp(-self.fbin)) # # # todo -- eliminate triple loop over nbins! # # Compute uncertainties in bin probabilities. # d2p_i = np.zeros([nbins], np.float64) # for k in range(nbins): # for i in range(nbins): # for j in range(nbins): # delta_ik = 1.0 * (i == k) # delta_jk = 1.0 * (j == k) # d2p_i[k] += p_i[k] * (p_i[i] - delta_ik) * p_i[k] * ( # p_i[j] - delta_jk) * Theta_ij[K + i, K + j] # # # Transform from d2p_i to df_i # d2f_i = d2p_i / p_i ** 2 # df_i = np.sqrt(d2f_i) raise ParameterError( "uncertainty_method 'from-normalization' is not currently supported for histograms" ) fx_vals = np.zeros(len(x)) dfx_vals = np.zeros(len(x)) # figure out how many grid points in each direction maxp = np.zeros(dims, int) for d in range(dims): maxp[d] = len(bins[d]) for i, l in enumerate(loc_indices): # Must be a way to list comprehend this? if np.any(l < 0): # out of index below fx_vals[i] = np.nan dfx_vals[i] = np.nan continue if np.any(l >= maxp - 1): # out of index above fx_vals[i] = np.nan dfx_vals[i] = np.nan continue bin_label = histogram_data["bin_label"][tuple(l)] if bin_label >= 0: fx_vals[i] = f_i[bin_order[bin_label]] dfx_vals[i] = df_i[bin_order[bin_label]] else: fx_vals[i] = np.nan dfx_vals[i] = np.nan # Return dimensionless free energy and uncertainty. result_vals["f_i"] = fx_vals if uncertainty_method is not None: result_vals["df_i"] = dfx_vals if reference_point == "all-differences": if uncertainty_method == "analytical": # Report uncertainties in all free energy differences as # well. diag = Theta_ij.diagonal() dii = diag[K, K + nbins] # appears broken? Not used? d2f_ij = dii + dii.transpose() - 2 * Theta_ij[K : K + nbins, K : K + nbins] # unsquare uncertainties df_ij = np.sqrt(d2f_ij) dfxij_vals = np.zeros([len(x), len(x)]) bin_is = () for i, l in enumerate(loc_indices): bin_i = histogram_data["bin_order"][self.histogram_data["bin_label"][tuple(l)]] bin_is.append(bin_i) for i, vi in bin_is: for j, vj in bin_i: if vi != -1 and vj != 1: dfxij_vals[i, j] = df_ij[vi, vj] else: dfxij_vals[i, j] = np.nan elif uncertainty_method == "bootstrap": # TODO: check this is working! dfxij_vals = np.zeros([len(histogram_data["f"]), len(histogram_data["f"])]) fall = np.zeros([len(histogram_data["f"]), len(histogram_data["f"]), n_bootstraps]) for b in range(n_bootstraps): h = histogram_datas[b] for i in range(nbins): fall[i, j, b] = ( histogram_datas[b]["f"] - histogram_datas[b]["f"].transpose() ) dfxij_vals = np.std(fall, axis=2) if uncertainty_method is not None: # Return dimensionless free energy and uncertainty. result_vals["df_ij"] = dfxij_vals return result_vals def _get_fes_kde( self, x, reference_point="from-normalization", fes_reference=None, uncertainty_method=None ): """ Returns values of the FES at the specified x points for kde FESs. Parameters ---------- x : numpy.ndarray of D dimensions, where D is the dimensionality of the FES defined. reference_point : str, optional Method for reporting values and uncertainties (default : 'from-lowest') * 'from-lowest' - the uncertainties in the free energy difference with lowest point on FES are reported * 'from-specified' - same as from lowest, but from a user specified point * 'from-normalization' - the normalization \\sum_i p_i = 1 is used to determine uncertainties spread out through the FES * 'all-differences' - the nbins x nbins matrix df_ij of uncertainties in free energy differences is returned instead of df_i uncertainty_method : str, optional Method for computing uncertainties (default : None) fes_reference: an N-d point specifying the reference state. Ignored except with uncertainty method ``from_specified`` Returns ------- dict 'f_i' : np.ndarray, float, shape=(K) result_vals['f_i'][i] is the dimensionless free energy of the x_i point, relative to the reference point 'df_i' : np.ndarray, float, shape=(K) result_vals['df_i'][i] is the uncertainty in the difference of x_i with respect to the reference point Only included if uncertainty_method is not None """ if np.shape(x)[1] != np.shape(self.kde.sample())[1]: raise DataError( "query coordinates have inconsistent dimension with the data the FES is fit to." ) result_vals = {} f_i = -self.kde.score_samples(x) # gives the LOG density, which is what we want. if reference_point == "from-lowest": fmin = np.min(f_i) f_i = f_i - fmin elif reference_point == "from-specified": fmin = -self.kde.score_samples(np.array(fes_reference).reshape(1, -1)) f_i = f_i - fmin elif reference_point == "from-normalization": # uncertainites "from normalization" reference is already applied, since # the density is normalized. pass else: raise ParameterError( f"reference point choice {reference_point} for kde is unavailable" ) result_vals["f_i"] = f_i # now calculate bootstrap uncertainties if uncertainty_method is None: df_i = None elif uncertainty_method == "bootstrap": if self.kdes is None: raise ParameterError( f"Cannot calculate bootstrap error of boostrap KDE's not determined" ) else: n_bootstraps = len(self.kdes) fall = np.zeros([len(x), n_bootstraps]) for b in range(n_bootstraps): fall[:, b] = -self.kdes[b].score_samples(x) - fmin df_i = np.std(fall, axis=1) else: raise ParameterError( f"Uncertainty method {uncertainty_method} for kde is not implemented" ) result_vals["df_i"] = df_i return result_vals def _get_fes_spline( self, x, reference_point="from_lowest", fes_reference=0.0, uncertainty_method=None ): """ Returns values of the FES at the specified x points for spline FESs. Parameters ---------- x : numpy.ndarray of D dimensions, where D is the dimensionality of the FES defined. reference_point : str, optional Method for reporting values and uncertainties (default : 'from-lowest') * 'from-lowest' - the uncertainties in the free energy difference with lowest point on FES are reported * 'from-specified' - same as from lowest, but from a user specified point * 'from-normalization' - the normalization \\sum_i p_i = 1 is used to determine uncertainties spread out through the FES * 'all-differences' - the nbins x nbins matrix df_ij of uncertainties in free energy differences is returned instead of df_i uncertainty_method : str, optional Method for computing uncertainties (default : None) fes_reference: an N-d point specifying the reference state. Ignored except with uncertainty method ``from_specified`` Returns ------- dict 'f_i' : np.ndarray, float, shape=(K) result_vals['f_i'][i] is the dimensionless free energy of the x_i point, relative to the reference point 'df_i' : np.ndarray, float, shape=(K) result_vals['df_i'][i] is the uncertainty in the difference of x_i with respect to the reference point Only included if uncertainty_method is not None """ if np.shape(x)[1] != 1: raise DataError("splines FES only supported in 1D") result_vals = {} # for splines now, should only be 1D. x is passed in as a 2D array, need to covert # back to 1D array before being put in fes_function (which will preserve shape) x = x[:, 0] # before being put here, needs to be converted back f_i = self.fes_function(x) if reference_point == "from-lowest": fmin = np.min(f_i) f_i = f_i - fmin elif reference_point == "from-specified": fmin = -self.fes_function(np.array(fes_reference).reshape(1, -1)) f_i = f_i - fmin else: raise ParameterError( f"reference point {reference_point} not implemented for spline fes" ) if uncertainty_method is None: df_i = None if uncertainty_method == "bootstrap": if self.fes_functions is None: raise ParameterError( "Cannot calculate via uncertainties error if boostrapping was not peformed running get_fes" ) else: n_bootstraps = len(self.fes_functions) dim_breakdown = [d for d in x.shape] + [n_bootstraps] fall = np.zeros(dim_breakdown) for b in range(n_bootstraps): fall[:, b] = self.fes_functions[b](x) - fmin df_i = np.std(fall, axis=-1) # uncertainites "from normalization" reference is applied, since # the density is normalized. result_vals["f_i"] = f_i result_vals["df_i"] = df_i return result_vals def sample_parameter_distribution( self, x_n, mc_parameters=None, decorrelate=True, verbose=True ): """ Samples the valus of the spline parameters with MC. Parameters ---------- x_n : numpy.ndarray of D dimensions D is the dimensionality of the FES defined. mc_parameters: dictionary niteratons : int number of iterations of the Monte Carlo procedure fraction_change : float which fraction of the range of input parameters is used to make new MC moves. sample_every : int the frequency in steps at which the MC timeseries is saved. print_every : int the frequency in steps aat which the MC timeseries is saved to log. logprior : function, the function of parameters, the Function must take a single argument, an array the size of parameters in in the same order as a used by the internal functions. decorrelate : boolean Whether to decorrelate the time series of output points verbose : boolean Whether to print high levels of information to the logger Returns ------- None """ # determine the range of the bspline at the start of the # process: changes are made as fractions of this if self.fes_type != "spline": ParameterError("Sampling of posterior is only supported for spline type") spline_parameters = self.spline_parameters if spline_parameters is None: ParameterError("Must specify spline_parameters to sample the distributions") spline_weights = spline_parameters["spline_weights"] # need the x-range for all methods, since we need to xrange = spline_parameters["xrange"] # numerically integrate over this range if self.fes_function is None: ParameterError( "Need to generate an initial splined FES using generate_fes before performing MCMC sampling" ) if mc_parameters is None: logger.info("Using default MC parameters") mc_parameters = dict() if "niterations" not in mc_parameters: mc_parameters["niterations"] = 5000 if "fraction_change" not in mc_parameters: mc_parameters["fraction_change"] = 0.01 if "sample_every" not in mc_parameters: mc_parameters["sample_every"] = 50 if "print_every" not in mc_parameters: mc_parameters["print_every"] = 1000 if "logprior" not in mc_parameters: mc_parameters["logprior"] = lambda x: 0 niterations = mc_parameters["niterations"] fraction_change = mc_parameters["fraction_change"] sample_every = mc_parameters["sample_every"] print_every = mc_parameters["print_every"] logprior = mc_parameters["logprior"] # create a new copy of the spline for MC sampling. self.mc_data = dict() # we would like to make this below a copy, but BSpline doesn't have copy. self.mc_data["bspline"] = self.fes_function bspline = self.mc_data["bspline"] # ensure normalization of spline def prob(x): return np.exp(-bspline(x)) norm = self._integrate(prob, xrange[0], xrange[1]) bspline.c = bspline.c + np.log(norm) # make a copy of the original spline to preserve it. self.mc_data["original_spline"] = BSpline(bspline.t, bspline.c, bspline.k) # this might not work as well for probability c = bspline.c crange = np.max(c) - np.min(c) dc = fraction_change * crange self.mc_data["naccept"] = 0 csamples = np.zeros([len(c), int(niterations) // int(sample_every)]) logposteriors = np.zeros(int(niterations) // int(sample_every)) self.mc_data["first_step"] = True for n in range(niterations): results = self._MC_step(x_n, self.w_n, dc, xrange, spline_weights, logprior) if n % sample_every == 0: csamples[:, n // sample_every] = results["c"] logposteriors[n // sample_every] = results["logposterior"] if n % print_every == 0 and verbose: logger.info( "MC Step {:d} of {:d}".format(n, niterations), str(results["logposterior"]), str(bspline.c), ) # We now have a distribution of samples of parameters sampled according # to the posterior. # decorrelate the data t_mc = 0 g_mc = None if verbose: logger.info("Done MC sampling") if decorrelate: t_mc, g_mc, Neff = timeseries.detect_equilibration(logposteriors) logger.info( "First equilibration sample is {:d} of {:d}".format(t_mc, len(logposteriors)) ) equil_logp = logposteriors[t_mc:] g_mc = timeseries.statistical_inefficiency(equil_logp) if verbose: logger.info("Statistical inefficiency of log posterior is {:.3g}".format(g_mc)) g_c = np.zeros(len(c)) for nc in range(len(c)): g_c[nc] = timeseries.statistical_inefficiency(csamples[nc, t_mc:]) if verbose: logger.info("Time series for spline parameters are : {:s}".format(str(g_c))) maxgc = np.max(g_c) meangc = np.mean(g_c) guse = g_mc # doesn't affect the distribution that much indices = timeseries.subsample_correlated_data(equil_logp, g=guse) logposteriors = equil_logp[indices] csamples = (csamples[:, t_mc:])[:, indices] if verbose: logger.info("samples after decorrelation : {:d}".format(np.shape(csamples)[1])) self.mc_data["samples"] = csamples self.mc_data["logposteriors"] = logposteriors self.mc_data["mc_parameters"] = mc_parameters self.mc_data["acceptance_ratio"] = self.mc_data["naccept"] / niterations if verbose: logger.info("Acceptance rate : {:5.3f}".format(self.mc_data["acceptance_ratio"])) self.mc_data["nequil"] = t_mc # the start of the "equilibrated" data set self.mc_data["g_logposterior"] = g_mc # statistical efficiency of the log posterior self.mc_data["g_parameters"] = g_c # statistical efficiency of the parametere self.mc_data["g"] = guse # statistical efficiency used for subsampling def get_confidence_intervals(self, xplot, plow, phigh, reference="zero"): """ Parameters ---------- xplot: data points we want to plot at plow: lowest percentile phigh: highest percentile Returns ------- Dictionary of results. Contains: plow : ndarray of float len(xplot) value of the parameter at plow percentile of the distribution at each x in xplot. phigh : ndarray of float value of the parameter at phigh percentile of the distribution at each x in xplot. median : ndarray of float value of the parameter at the median of the distribution at each x in xplot. values : ndarray of float shape [niterations//sample_every, len(xplot)] of the FES saved during the MCMC sampling at each input value of xplot. """ if self.mc_data is None: raise DataError("No MC sampling has been done, cannot construct confidence intervals") # determine confidence intervals nplot = len(xplot) # number of data points to plot. nsamples = len(self.mc_data["logposteriors"]) samplevals = np.zeros([nplot, nsamples]) csamples = self.mc_data["samples"] base_spline = self.mc_data["original_spline"] yvals = base_spline(xplot) for n in range(nsamples): # create a sample spline pcurve = BSpline(base_spline.t, csamples[:, n], base_spline.k) samplevals[:, n] = pcurve(xplot) # now determine the Bayesian confidence intervals. ylows = np.zeros(len(xplot)) yhighs = np.zeros(len(xplot)) ymedians = np.zeros(len(xplot)) for n in range(len(xplot)): ylows[n] = np.percentile(samplevals[n, :], plow) yhighs[n] = np.percentile(samplevals[n, :], phigh) ymedians[n] = np.percentile(samplevals[n, :], 50) return_vals = dict() if reference == "zero": ref = np.min(yvals) elif reference == None: ref = 0 else: raise ParameterError("{:s} is not a valid value for 'reference'") return_vals["plow"] = ylows - ref return_vals["phigh"] = yhighs - ref return_vals["median"] = ymedians - ref return_vals["values"] = yvals - ref return return_vals def get_mc_data(self): """convenience function to retrieve MC data Parameters ---------- None Returns ------- dict samples : samples of the parameters with size [len(parameters) times niterations/sample_every] logposteriors : log posteriors (which might be defined with respect to some reference) as a time series size [# points] mc_parameters : dictionary of parameters that were run with (see definitons in sample_parameter_distrbution) acceptance_ratio : overall acceptance ratio of the MC chain nequil : the start of the "equilibrated" data set (i.e. nequil-1 is the number that werer thrown out) g_logposterior : statistical efficiency of the log posterior g_parameters : statistical efficiency of the parametere g : statistical efficiency used for subsampling """ if self.mc_data is None: raise DataError("No MC sampling has been done, cannot construct confidence intervals") else: return self.mc_data def _get_MC_loglikelihood(self, x_n, w_n, spline_weights, spline, xrange): """ Parameters ---------- x_n : np.ndarray, float, shape=(N,D) x_n[n] is the d-dimensional coordinates of the samples, where D is the reduced dimensional space. w_n : np.ndarray, float, shape=(sself.N) Weights for each sample for the state in which we want the FES (usually, the unbiased state) spline_weights : string which type of fit to the likelihood to use (see `generate_fes` options) spline : function of ndarray argument function current value of the spline for which likelihood is being calculated xrange : float, shape=(2) range over which the FES is defined (defined in spline_parameters ["xrange"] Returns ------- loglikelihood : float log-likelihood of this spline """ N = self.N K = self.K if spline_weights in ["simplesum", "biasedstates"]: loglikelihood = 0 def splinek(x, kf): return spline(x) + self.spline_parameters["fkbias"][kf](x) def expk(x, kf): return np.exp(-splinek(x, kf)) for k in range(K): x_kn = x_n[self.mbar.x_kindices == k] normalize = np.log(self._integrate(expk, xrange[0], xrange[1], args=(k))) if spline_weights == "simplesum": loglikelihood += (N / K) * np.mean(splinek(x_kn, k)) loglikelihood += (N / K) * normalize elif spline_weights == "biasedstates": loglikelihood += np.sum(splinek(x_kn, k)) loglikelihood += self.N_k[k] * normalize elif spline_weights == "unbiasedstate": loglikelihood = N * np.dot(w_n, spline(x_n)) # no need to add normalization, should be normalized. return loglikelihood def _MC_step(self, x_n, w_n, stepsize, xrange, spline_weights, logprior): """sample over the posterior space of the FES as splined. Parameters ---------- x_n: samples from the biased distribution w_n: weights of each sample. stepsize: sigma of the normal distribution used to propose steps xrange: Range the probility distribution is defined o er. spline_weights: Type of weighting used for maximum likelihood for splines. See class definition for description of types. logprior: function describing the prior of the parameters. Default is uniform. Outputs ------- dict * 'c' : the value of the spline constants (len nsplines - we always assume normalized * 'logposterior' : the current value of the logposterior. Notes ----- Modifies several saved variables saved in the structure.x """ mc_data = self.mc_data # keep it shorter bspline = self.mc_data["bspline"] if self.mc_data["first_step"]: c = bspline.c mc_data["previous_logposterior"] = self._get_MC_loglikelihood( x_n, w_n, self.spline_parameters["spline_weights"], bspline, self.spline_parameters["xrange"], ) - logprior(c) cold = bspline.c mc_data["first_step"] = True # create an extra one we can carry around mc_data["newspline"] = BSpline(bspline.t, bspline.c, bspline.k) mc_data["cold"] = bspline.c psize = len(mc_data["cold"]) rchange = stepsize * np.random.normal() cnew = mc_data["cold"].copy() ci = np.random.randint(psize) cnew[ci] += rchange mc_data["newspline"].c = cnew # determine the change in the integral def prob(x): return np.exp(-mc_data["newspline"](x)) new_integral = self._integrate(prob, xrange[0], xrange[1]) cnew = cnew + np.log(new_integral) mc_data["newspline"].c = cnew # this spline should now be normalized. # now calculate the change in log likelihood loglikelihood = self._get_MC_loglikelihood( x_n, w_n, spline_weights, mc_data["newspline"], xrange ) newlogposterior = loglikelihood - logprior(cnew) dlogposterior = newlogposterior - (mc_data["previous_logposterior"]) accept = False if dlogposterior <= 0: accept = True if dlogposterior > 0: if np.random.random() < np.exp(-dlogposterior): accept = True if accept: mc_data["bspline"].c = mc_data["newspline"].c mc_data["cold"] = bspline.c mc_data["previous_logposterior"] = newlogposterior mc_data["naccept"] = mc_data["naccept"] + 1 results = dict() results["c"] = mc_data["bspline"].c results["logposterior"] = mc_data["previous_logposterior"] return results def _bspline_calculate_f(self, xi, x_n, w_n): """Calculate the maximum likelihood / KL divergence of the FES represented using B-splines. Parameters ---------- xi : array of floats size nspline-1 spline coefficients, w_n: weights for each sample. x_n: values of each sample. Output ------ float function value at spline coefficients xi """ mbar = self.mbar K = mbar.K N_k = mbar.N_k N = self.N bloc = self._val_to_spline(xi) # convert the spline coefficients into a spline object spline_weights = self.spline_parameters[ "spline_weights" ] # how to weight the integrated splines in the final likelihood nspline = self.spline_parameters["nspline"] # number of spline points kdegree = self.spline_parameters["kdegree"] # degree of spline xrange = self.spline_parameters["xrange"] # the range FES is defined over fkbias = self.spline_parameters["fkbias"] # K biasing functions if spline_weights in ["simplesum", "biasedstates"]: pF = np.zeros(K) if spline_weights == "simplesum": f = 0 for k in range(K): f += (N / K) * np.mean(bloc(x_n[mbar.x_kindices == k])) elif spline_weights == "biasedstates": # multiply by K to get it in the same order of magnitude f = np.sum(bloc(x_n)) if spline_weights == "simplesum": integral_scaling = (N / K) * np.ones(K) elif spline_weights == "biasedstates": integral_scaling = N_k expf = list() for k in range(K): # what is the biasing function for this state? # define the biasing function # define the exponential of f based on the current parameters # t. def expfk(x, kf=k): return np.exp(-bloc(x) - fkbias[kf](x)) # compute the partition function pF[k] = self._integrate(expfk, xrange[0], xrange[1], args=(k)) expf.append(expfk) # subtract the free energy (add log partition function) f += np.dot(integral_scaling, np.log(pF)) elif spline_weights == "unbiasedstate": # just KL divergence of the unbiased potential # may need to recast w_n f = N * np.dot(w_n, bloc(x_n)) def expf(x): return np.exp(-bloc(x)) # setting limit to try to eliminate errors: hard time because it # goes so small. pF = self._integrate(expf, xrange[0], xrange[1]) # subtract the free energy (add log partition function) f += N * np.log(pF) self.spline_data["bspline_expf"] = expf self.spline_data["bspline_pF"] = pF logprior = self.spline_parameters["map_data"]["logprior"] if logprior != None: # need to add the zero explicitly to the front f -= logprior(np.concatenate([[0], xi], axis=None)) return f def _bspline_calculate_g(self, xi, x_n, w_n): """Calculate the gradient of the maximum likelihood / KL divergence of the FES represented using B-splines. Parameters ----------- xi : array of floats, size nspline-1 spline coefficients, w_n : array of floats weights for each sample. x_n : array of floats values of each sample. Output ------ gradient: array of floats, size (nspline-1) """ ##### COMPUTE THE GRADIENT ####### # The gradient of the function is \sum_n [\sum_k W_k(x_n)] dF(phi(x_n))/dtheta_i - \sum_k _k # # where _k = \int O(xi) exp(-F(xi) - u_k(xi)) dxi / \int exp(-F(xi) # - u_k(xi)) dxi mbar = self.mbar K = mbar.K N_k = mbar.N_k N = self.N bloc = self._val_to_spline(xi) # convert the spline coefficients into a spline object spline_weights = self.spline_parameters[ "spline_weights" ] # how to weight the integrated splines in the final likelihood nspline = self.spline_parameters["nspline"] # number of spline points kdegree = self.spline_parameters["kdegree"] # degree of spline xrange = self.spline_parameters["xrange"] # the range FES is defined over fkbias = self.spline_parameters["fkbias"] # K biasing functions db_c = self.spline_data[ "bspline_derivatives" ] # coefficients of the derivatives of the splines xrangei = self.spline_data[ "xrangei" ] # range the ith basis function of the spline is defined over pF = np.zeros(K) if spline_weights == "simplesum": integral_scaling = (N / K) * np.ones(K) elif spline_weights == "biasedstates": integral_scaling = N_k g = np.zeros(nspline - 1) for i in range(1, nspline): if spline_weights == "simplesum": for k in range(K): g[i - 1] += (N / K) * np.mean(db_c[i](x_n[mbar.x_kindices == k])) elif spline_weights == "biasedstates": g[i - 1] = np.sum(db_c[i](x_n)) elif spline_weights == "unbiasedstate": g[i - 1] = N * np.dot(w_n, db_c[i](x_n)) # now the second part of the gradient. if spline_weights in ["biasedstates", "simplesum"]: gkquad = np.zeros([nspline - 1, K]) def expf(x, k): return np.exp(-bloc(x) - fkbias[k](x)) def dexpf(x, k): return db_c[i + 1](x) * expf(x, k) for k in range(K): # putting this in rather than saving the term so gradient and f # can be called independently pF[k] = self._integrate(expf, xrange[0], xrange[1], args=(k)) for i in range(nspline - 1): # Boltzmann weighted derivative with each biasing function # now compute the expectation of each derivative pE = self._integrate(dexpf, xrangei[i + 1, 0], xrangei[i + 1, 1], args=(k)) # normalize the expectation gkquad[i, k] = pE / pF[k] g -= np.dot(gkquad, integral_scaling) elif spline_weights == "unbiasedstate": gkquad = 0 # not used here, but saved for Hessian calls. def expf(x): return np.exp(-bloc(x)) # 0 is the value of gkquad. Recomputed here to avoid problems pF = self._integrate(expf, xrange[0], xrange[1]) # with other scipy solvers pE = np.zeros(nspline - 1) def dexpf(x, index): return db_c[index + 1](x) * expf(x) for i in range(nspline - 1): # Boltzmann weighted derivative # now compute the expectation of each derivative pE[i] = self._integrate(dexpf, xrangei[i + 1, 0], xrangei[i + 1, 1], args=(i,)) # normalize the expectation. pE[i] /= pF g -= N * pE dlogprior = self.spline_parameters["map_data"]["dlogprior"] if dlogprior != None: # need to add the zero explicitly to the front g -= dlogprior(np.concatenate([[0], xi], axis=None)) self.spline_data["bspline_gkquad"] = gkquad self.spline_data["bspline_pE"] = pE return g def _bspline_calculate_h(self, xi, x_n, w_n): """Calculate the Hessian of the maximum likelihood / KL divergence of the FES represented using B-splines. Parameters ---------- xi : array of floats size nspline-1 spline coefficients w_n : array of floats, size number of points weights for each sample x_n : array of floats, size number of points values of each sample Output ------ Hessian nfloat, size (nspline-1) x (nspline - 1) Notes ----- CURRENTLY assumes that the gradient has already been called at the current value of the parameters. Otherwise, it fails. This means it only works for certain algorithms. """ mbar = self.mbar K = mbar.K N_k = mbar.N_k N = self.N bloc = self._val_to_spline(xi) # convert the spline coefficients into a spline object spline_weights = self.spline_parameters[ "spline_weights" ] # how to weight the integrated splines in the final likelihood nspline = self.spline_parameters["nspline"] # number of spline points kdegree = self.spline_parameters["kdegree"] # degree of spline fkbias = self.spline_parameters["fkbias"] # K biasing functions db_c = self.spline_data[ "bspline_derivatives" ] # coefficients of the derivatives of the splines xrangeij = self.spline_data[ "xrangeij" ] # range in x and y the 2d integration of basis functions i and j are defined over. expf = self.spline_data["bspline_expf"] gkquad = self.spline_data["bspline_gkquad"] pF = self.spline_data["bspline_pF"] pE = self.spline_data["bspline_pE"] if spline_weights == "simplesum": integral_scaling = N / K * np.ones(K) elif spline_weights == "biasedstates": integral_scaling = N_k # now, compute the Hessian. First, the first order components h = np.zeros([nspline - 1, nspline - 1]) if spline_weights in ["simplesum", "biasedstates"]: for k in range(K): h += -integral_scaling[k] * np.outer(gkquad[:, k], gkquad[:, k]) elif spline_weights == "unbiasedstate": h = -N * np.outer(pE, pE) if spline_weights in ["simplesum", "biasedstates"]: for i in range(nspline - 1): for j in range(0, i + 1): if np.abs(i - j) <= kdegree: def ddexpf(x, k_inner): # Disable the PyLint check here because this is the behavior we want # pylint: disable=cell-var-from-loop return db_c[i + 1](x) * db_c[j + 1](x) * expf[k_inner](x) for k in range(K): # now compute the expectation of each derivative pE = integral_scaling[k] * self._integrate( ddexpf, xrangeij[i + 1, j + 1, 0], xrangeij[i + 1, j + 1, 1], args=(k,), ) h[i, j] += pE / pF[k] elif spline_weights == "unbiasedstate": def ddexpf(x, index_i, index_j): return db_c[index_i + 1](x) * db_c[index_j + 1](x) * expf(x) for i in range(nspline - 1): for j in range(0, i + 1): if np.abs(i - j) <= kdegree: # now compute the expectation of each derivative pE = self._integrate( ddexpf, xrangeij[i + 1, j + 1, 0], xrangeij[i + 1, j + 1, 1], args=(i, j), ) h[i, j] += N * pE / pF for i in range(nspline - 1): for j in range(i + 1, nspline - 1): h[i, j] = h[j, i] ddlogprior = self.spline_parameters["map_data"]["ddlogprior"] if ddlogprior != None: # add hessian of prior # need to add the zero explicitly to the front h -= ddlogprior(np.concatenate([[0], xi], axis=None)) return h @staticmethod def _integrate(func, xlow, xhigh, args=(), method="quad"): # could add more ability to specify """ wrapper for integration in case we decide to replace quad with something analytical """ if method == "quad": # just use scipy quadrature results = quad(func, xlow, xhigh, args)[0] else: raise ParameterError("integration method {:s} not yet implemented".format(method)) return results def _val_to_spline(self, x, form=None): """ Convert a set of B-spline coefficients into a BSpline object Parameters ---------- x : float, array the last N-1 coefficients for a bspline; we assume the initial coefficient is set to zero. Returns ------- bspline : A bspline object (or function returning -log (bspline) object if we need it) """ template_bspline = self.spline_data["bspline"] # create new spline with values xnew = np.zeros(len(x) + 1) xnew[0] = (template_bspline).c[0] xnew[1:] = x bspline = BSpline((template_bspline).t, xnew, (template_bspline).k) if form == "exp": return lambda x: -np.log(bspline(x)) elif form == "log": return bspline elif form is None: return bspline ######## # Integration notes: # # If we wanted to integrate the function, can we do it analytically? # Let's focus on the likelihood integral, which is what is slowing down # the MC, which is the slow part. # # for k=0, then B_i,0:t = 1 if t_i < x < t_i+i, 0 otherwise # for k=1, It is a piecewise sum of 2 linear terms, so linear. # f(x) = \\int exp(ax+b)_{t_i}^{t_i+1) = (1/a) e^b (e^a*t2 - e^a t1) # for k=2, it is piecewise sum of 3 quadratic terms, which is quadradic # f(x) = \\int exp(-a(x-b)(x-c))_{t_i)+{t_i+1) = (exp^(1/4 a (b - c)^2) Sqrt[\\pi]] (Erf[1/2 Sqrt[a] (b + c - 2 t1)] - # Erf[1/2 Sqrt[a] (b + c - 2 t2)]))/(2 Sqrt[a]), for a > 0, switch for a<0. # # for k=3, piecewise sum of cubic terms, which appears hard in general. # # Of course, even with linear, we need to be able to integrate with the # bias functions. If it's a Gaussian bias, then linear and quadratic should integrate fine. ########