From 5d6ee771101c6fdb661b7aa7923df8f194c48940 Mon Sep 17 00:00:00 2001 From: KuangYu Date: Sun, 5 Nov 2023 09:26:19 +0800 Subject: [PATCH] Update doc --- docs/user_guide/4.6MBAR.md | 50 +++++++++++++++++++++++++++----------- 1 file changed, 36 insertions(+), 14 deletions(-) diff --git a/docs/user_guide/4.6MBAR.md b/docs/user_guide/4.6MBAR.md index 1a0662f19..cf573805f 100644 --- a/docs/user_guide/4.6MBAR.md +++ b/docs/user_guide/4.6MBAR.md @@ -7,45 +7,57 @@ In molecular dynamics (MD) simulations, the deep computational graph spanning th In the MBAR theory, it is assumed that there are K ensembles defined by potential energies $$ -u_{i}(x)(i=1,2,3,……,K) \label{eq-1} +\tag{1} +u_{i}(x)(i=1,2,3,……,K) $$ For each ensemble, the Boltzmann weight, partition function, and probability function are defined as: + $$ +\tag{2} w_i = \exp(-\beta_i u_i(x)) \\ c_i = \int dx \cdot w_i(x) \\ -p_i(x) = \frac{w_i(x)}{c_i} \label{eq-2} +p_i(x) = \frac{w_i(x)}{c_i} $$ For each ensemble $i$, select $N_{i}$ configurations, represented by {${x_{in}}$} (where $n=1,2,3,……,N_{i}$), and the total number of configurations across ensembles is represented by {${x_{n}}$} ($n=1,2,3,……,N$), where N is: $$ -N = \sum_{i=1}^{K} N_i \label{eq-3} +\tag{3} +N = \sum_{i=1}^{K} N_i $$ Within the context of MBAR, for any ensemble K, the weighted average of the observable is defined as: $$ -\hat{c}_i = \sum_{n=1}^{N} \frac{w_{i}(x_n)}{\sum_{k=1}^{K} N_{k} \hat{c}_k^{-1} w_{k}(x_n)} \label{eq-4} +\tag{4} +\hat{c}_i = \sum_{n=1}^{N} \frac{w_{i}(x_n)}{\sum_{k=1}^{K} N_{k} \hat{c}_k^{-1} w_{k}(x_n)} $$ To compute the average of a physical quantity $A$ in ensemble $i$, one can utilize the above values to define a virtual ensemble $j$ , with its corresponding Boltzmann weight and partition function: $$ +\tag{5} w_j = w_i(x)A(x) \\ -c_i = \int dx \cdot w_j(x) \label{eq-5} +c_i = \int dx \cdot w_j(x) $$ + Thus, the ensemble average of A is: + $$ -\langle A \rangle_i = \frac{\hat{c}_j}{\hat{c}_i} = \frac{\int dx \cdot w_i(x)A(x)}{\int dx \cdot w_i(x)} \label{eq-6} +\tag{6} +\langle A \rangle_i = \frac{\hat{c}_j}{\hat{c}_i} = \frac{\int dx \cdot w_i(x)A(x)}{\int dx \cdot w_i(x)} $$ + Thus, the MBAR theory provides a method for estimating the average of physical properties using multiple samples. In the MBAR framework, $\hat{c}_i$ in Eqn $\ref{eq-4}$ needs to be solved iteratively; however, the differentiable reweighting algorithm can simplify this estimation process. During gradient descent training for parameter optimization, the parameters undergo only slight perturbations in each training cycle. This allows for the usage of samples from the previous cycles, such that resampling is not necessary until the optimized ensemble deviates significantly from the sampling ensemble. This reduces the time and computational cost of the optimization considerably. In the reweighted MBAR estimator, we define two types of ensembles: the sampling ensemble, from which all samples are extracted (assuming there are $m$ samples, labeled as $m=1, 2, 3, …, M$), and the target ensemble, which needs optimization (labeled as $p, q$, corresponding to the indices $i, j$ in Eqn $\ref{eq-6}$). The sampling ensemble is updated only when necessary and does not need to be differentiable. Its data can be generated by external samplers like OpenMM. Hence, $\hat{c}_i$ can be transformed into: $$ -\hat{c}_p = \sum_{n=1}^{N} w_{p}(x_n) \left( \sum_{m=1}^{M} N_{m} \hat{c}_m^{-1} w_{m}(x_n) \right)^{-1} \label{eq-7} +\tag{7} +\hat{c}_p = \sum_{n=1}^{N} w_{p}(x_n) \left( \sum_{m=1}^{M} N_{m} \hat{c}_m^{-1} w_{m}(x_n) \right)^{-1} $$ + When resample happens, Eqn. $\ref{eq-4}$ is solved iteratively using standard MBAR to update $\hat{c}_m$, which is stored and used to evaluate $\hat{c}_p$ until the next resampling. Subsequently, during the parameter optimization process, Eqn $\ref{eq-7}$ is employed to compute $\hat{c}_p$, serving as a differentiable estimator. Below, we illustrate the workflow of how to use MBAR Estimator in DMFF through a case study. @@ -53,40 +65,50 @@ Below, we illustrate the workflow of how to use MBAR Estimator in DMFF through a If all sampling ensembles are defined as a single ensemble $w_{0}(x)$, and the target ensemble is defined as $w_{p}(x)$, and for physical quantity A, we have: $$ -w_q(x) = w_p(x) A(x) \label{eq-8} +\tag{8} +w_q(x) = w_p(x) A(x) $$ + and define: $$ -\Delta u_{p_0} = u_p(x) - u_0(x) \label{eq-9} +\tag{9} +\Delta u_{p_0} = u_p(x) - u_0(x) $$ + then: $$ -\langle A \rangle_p = \frac{\hat{c}_q}{\hat{c}_p}= \frac{\sum_{n=1}^{N} A(x_n) \exp(-\beta \Delta u_{p_0}(x_n))}{\sum_{n=1}^{N} \exp(-\beta \Delta u_{p_0}(x_n))} \label{eq-10} +\tag{10} +\langle A \rangle_p = \frac{\hat{c}_q}{\hat{c}_p}= \frac{\sum_{n=1}^{N} A(x_n) \exp(-\beta \Delta u_{p_0}(x_n))}{\sum_{n=1}^{N} \exp(-\beta \Delta u_{p_0}(x_n))} $$ + Refers to equations above, this equation indicates that the trajectory reweighting algorithm is a special case of the reweighted MBAR estimator. In DMFF, when calculating the average of the physical quantity A, the formula is expressed as: $$ -\langle A \rangle_p = \sum_{n=1}^{N} W_n A(x_n) \label{eq-11} +\tag{11} +\langle A \rangle_p = \sum_{n=1}^{N} W_n A(x_n) $$ where $$ -\Delta U_{mp} = U_m(x_n) - U_p(x_n) \label{eq-12} +\tag{12} +\Delta U_{mp} = U_m(x_n) - U_p(x_n) $$ $$ -W_n = \frac{\left[\sum_{m=1}^{M} N_m e^{\hat{f}_m -\beta \Delta U_{mp}(x_n)}\right]^{-1}}{\sum_{n=1}^{N} \left[ \sum_{m=1}^{M} N_m e^{\hat{f}_m -\beta \Delta U_{mp}(x_n)} \right]^{-1}} \label{eq-13} +\tag{13} +W_n = \frac{\left[\sum_{m=1}^{M} N_m e^{\hat{f}_m -\beta \Delta U_{mp}(x_n)}\right]^{-1}}{\sum_{n=1}^{N} \left[ \sum_{m=1}^{M} N_m e^{\hat{f}_m -\beta \Delta U_{mp}(x_n)} \right]^{-1}} $$ $\hat{f}_m$ is the partition function of the sampling state. W is the MBAR weight for each sample. Finally, the effective sample size is given, based on which one can judge the deviation of the sampling ensemble from the target ensemble: $$ -n_{\text{eff}} = \frac{\left(\sum_{n=1}^{N} W_n\right)^2}{\sum_{n=1}^{N} W_n^2} \label{eq-14} +\tag{14} +n_{\text{eff}} = \frac{\left(\sum_{n=1}^{N} W_n\right)^2}{\sum_{n=1}^{N} W_n^2} $$ When $n_{eff}$ is too small, it indicates that the current sampling ensemble deviates too much from the target ensemble and needs to be resampled.