Merge pull request #28 from wehs7661/fixed_p_acc

Fix the calculation of acceptance ratio in REXEE simulations

docs/theory.rst

@@ -1,16 +1,20 @@
 .. _doc_basic_idea:
 
+.. note:: This page is still a work in progress. Please refer to our paper for more details before this page is completed.
+
 1. Basic idea
 =============
-Replica exchange of expanded ensemble (REXEE) integrates the core principles of replica exchange 
-molecular dynamics (REMD) and expanded ensemble (EXE).  Specifically, a REXEE simulation includes multiple non-interacting, parallel expanded ensemble simulations 
-that collectively sample a number of alchemical states spanning between the coupled state 
-(:math:`\lambda=0`) and the uncoupled state (:math:`\lambda=1`). Each expanded ensemble 
-samples a subset of these states such that the range of its allowed alchemical states 
-overlaps with that of the adjacent replicas. In REXEE, the exchange of coordinates/alchemical 
-states occurs at a specified frequency, which is beneficial for allowing better mixing 
-in the alchemical space given sufficiently long simulation time, properly specified parameters 
-and highly parallelizable computing architectures. 
+Replica exchange of expanded ensemble (REXEE) integrates the core principles of replica exchange (REX)
+and expanded ensemble (EE) methods.  Specifically, a REXEE simulation includes multiple
+replicas of EE simulations in parallel and periodically exchanges coordinates
+between replicas. Each replica samples a different but overlapping range of alchemical 
+intermediate states to collectively sample the space bwteen the coupled (:math:`\lambda=0`)
+and decoupled states (:math:`\lambda=1`). Naturally, the REXEE method decorrelates
+the number of replicas from the number of states, allowing sampling a large number of intermediate
+states with significantly fewer replicas than those required in the Hamiltonian replica exchange (HREX)
+and other similar methods. By parallelizing replicas, the REXEE method also reduces
+the simulation wall time compared to the EE method. More importantly, such parallelism sets the
+stage for wider applications, such as relative free energy calculations for multi-topology transformations.
 
 Mathematically, we first consider a simulation ensemble that consists of :math:`M` non-interacting replicas 
 of the expanded ensemble simulations all at the same constant temperature :math:`T`. We assume 
@@ -203,7 +207,7 @@ end and there won't be a new pair drawn. In greater detail, this scheme can be d
 -----------------------
 In the current implementation, 3 acceptance schemes are available to be specified 
 in the input YAML file (e.g. :code:`params.yaml`) via the option :code:`acceptance`, including :code:`same-state`/:code:`same_state`, 
-:code:`metropolis`, and :code:`metropolis-eq`/:code:`metropolis_eq`. In our implementation, 
+and :code:`metropolis`. In our implementation, 
 relevant methods include :obj:`.propose_swaps`, :obj:`.calc_prob_acc`, and :obj:`.accept_or_reject`.
 Below we elaborate the details of each of the swapping schemes.
 
@@ -237,20 +241,6 @@ where
 In theory, this swapping scheme should obey the detailed balance condition, as derived 
 in :ref:`doc_basic_idea`.
 
-2.2.3. Equilibrated Metropolis swapping
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-In the limit of inifite simulation length, the alchemical weights of a certain 
-alchemical (e.g. :math:`m` or :math:`n`) should be indepedent of the configuration
-(e.g. :math:`i` or :math:`j`) being sampled, i.e. :math:`g^i_n=g^i_m` and :math:`g^j_n=g^j_m`. 
-At this limit, where the weights are equilibrated, the expression of :math:`\Delta` in the
-standard Metropolis swapping scheme reduces to the following:
-
-.. math::
-  \Delta = \beta[(U^i_n + U^j_m) - (U^i_m+U^j_n)]
-
-Notably, this scheme does not consider the difference in the alchemical weights, which can 
-be non-zero frequently, so this swapping scheme does not strictly obey the detailed balance condition.
-
 2.3. Calculation of Δ in Metropolis-based acceptance schemes
 ------------------------------------------------------------
 The calculation of :math:`\Delta` is important because the acceptance ratio :math:`w(X\rightarrow X')=\min(1, \exp(-\Delta))` is 

ensemble_md/cli/run_REXEE.py

@@ -159,7 +159,7 @@ def main():
                 states = copy.deepcopy(states_)
                 weights = copy.deepcopy(weights_)
                 counts = copy.deepcopy(counts_)
-                swap_pattern, swap_list = REXEE.get_swapping_pattern(dhdl_files, states_, weights_)  # swap_list will only be used for modify_coords  # noqa: E501
+                swap_pattern, swap_list = REXEE.get_swapping_pattern(dhdl_files, states_)  # swap_list will only be used for modify_coords  # noqa: E501
 
                 # 3-3. Perform weight correction/weight combination
                 if wl_delta != [None for i in range(REXEE.n_sim)]:  # weight-updating
@@ -206,7 +206,7 @@ def main():
                 elif REXEE.N_cutoff != -1 and REXEE.w_combine is None:
                     # only perform weight correction
                     print('Note: No weight combination will be performed.')
-                    weights = REXEE.histogram_correction(weights_input, counts)
+                    weights = REXEE.weights_correction(weights_input, counts)
                     _ = REXEE.combine_weights(counts_, weights, print_values=False)[1]  # just to print the combined weights  # noqa: E501
                 else:
                     print('Note: No weight correction will be performed.')

ensemble_md/replica_exchange_EE.py

@@ -196,8 +196,8 @@ def set_params(self, analysis):
         if self.proposal not in [None, 'single', 'neighboring', 'exhaustive', 'multiple']:
             raise ParameterError("The specified proposal scheme is not available. Available options include 'single', 'neighboring', 'exhaustive', and 'multiple'.")  # noqa: E501
 
-        if self.acceptance not in [None, 'same-state', 'same_state', 'metropolis', 'metropolis-eq', 'metropolis_eq']:
-            raise ParameterError("The specified acceptance scheme is not available. Available options include 'same-state', 'metropolis', and 'metropolis-eq'.")  # noqa: E501
+        if self.acceptance not in [None, 'same-state', 'same_state', 'metropolis']:
+            raise ParameterError("The specified acceptance scheme is not available. Available options include 'same-state' and 'metropolis'.")  # noqa: E501
 
         if self.w_combine not in [None, 'final', 'avg']:
             raise ParameterError("The specified type of weight to be combined is not available. Available options include 'final' and 'avg'.")  # noqa: E501
@@ -831,7 +831,7 @@ def propose_swap(swappables):
 
         return swap
 
-    def get_swapping_pattern(self, dhdl_files, states, weights):
+    def get_swapping_pattern(self, dhdl_files, states):
         """
         Generates a list (:code:`swap_pattern`) that represents how the configurations should be swapped in the
         next iteration. The indices of the output list correspond to the simulation/replica indices, and the
@@ -855,9 +855,6 @@ def get_swapping_pattern(self, dhdl_files, states, weights):
             A list of last sampled states (in global indices) of ALL simulaitons. :code:`states[i]=j` means that
             the configuration in replica :code:`i` is at state :code:`j` at the time when the exchange is performed.
             This list can be generated :obj:`.extract_final_dhdl_info`.
-        weights : list
-            A list of lists of final weights of ALL simulations. :code:`weights[i]` corresponds to the list of weights
-            used in replica :code:`i`. The list can be generated by :obj:`.extract_final_log_info`.
 
         Returns
         -------
@@ -926,7 +923,7 @@ def get_swapping_pattern(self, dhdl_files, states, weights):
                         swap_bool = True  # always accept the move
                     else:
                         # Calculate the acceptance ratio and decide whether to accept the swap.
-                        prob_acc = self.calc_prob_acc(swap, dhdl_files, states, shifts, weights)
+                        prob_acc = self.calc_prob_acc(swap, dhdl_files, states, shifts)
                         swap_bool = self.accept_or_reject(prob_acc)
 
                     # Theoretically, in an REXEE simulation, we could either choose to swap configurations (via
@@ -943,8 +940,8 @@ def get_swapping_pattern(self, dhdl_files, states, weights):
                         swap_list.append(swap)
                         # The assignments need to be done at the same time in just one line.
                         # states[swap[0]], states[swap[1]] = states[swap[1]], states[swap[0]]
+                        # weights[swap[0]], weights[swap[1]] = weights[swap[1]], weights[swap[0]]
                         shifts[swap[0]], shifts[swap[1]] = shifts[swap[1]], shifts[swap[0]]
-                        weights[swap[0]], weights[swap[1]] = weights[swap[1]], weights[swap[0]]
                         dhdl_files[swap[0]], dhdl_files[swap[1]] = dhdl_files[swap[1]], dhdl_files[swap[0]]
                         swap_pattern[swap[0]], swap_pattern[swap[1]] = swap_pattern[swap[1]], swap_pattern[swap[0]]
                         state_ranges[swap[0]], state_ranges[swap[1]] = state_ranges[swap[1]], state_ranges[swap[0]]
@@ -972,7 +969,7 @@ def get_swapping_pattern(self, dhdl_files, states, weights):
 
         return swap_pattern, swap_list
 
-    def calc_prob_acc(self, swap, dhdl_files, states, shifts, weights):
+    def calc_prob_acc(self, swap, dhdl_files, states, shifts):
         """
         Calculates the acceptance ratio given the Monte Carlo scheme for swapping the simulations.
 
@@ -996,10 +993,6 @@ def calc_prob_acc(self, swap, dhdl_files, states, shifts, weights):
             in :code:`dhdl_files[0]`, :code:`dhdl_files[1]`, ... (which importantly, are not necessarily
             :code:`dhdl_0.xvg`, :code:`dhdl_1.xvg`, ...). If multiple swaps are not used, then
             this should always be :code:`list(REXEE.s * np.arange(REXEE.n_sim))`.
-        weights : list
-            A list of lists of final weights of ALL simulations. Typiecally generated by
-            :obj:`.extract_final_log_info`. :code:`weights[i]` and :code:`dhdl_files[i]` should correspond to
-            the same configuration.
 
         Returns
         -------
@@ -1012,7 +1005,7 @@ def calc_prob_acc(self, swap, dhdl_files, states, shifts, weights):
             else:
                 prob_acc = 0  # This must lead to an output of swap_bool = False from the function accept_or_reject
 
-        else:  # i.e. metropolis-eq or metropolis, which both require the calculation of dU
+        else:  # i.e. must be metropolis, which requires the calculation of dU
             # Now we calculate dU
             if self.verbose is True:
                 print("  Proposing a move from (x^i_m, x^j_n) to (x^i_n, x^j_m) ...")
@@ -1044,24 +1037,8 @@ def calc_prob_acc(self, swap, dhdl_files, states, shifts, weights):
                     f"  U^i_n - U^i_m = {dU_0:.2f} kT, U^j_m - U^j_n = {dU_1:.2f} kT, Total dU: {dU:.2f} kT"
                 )
 
-            if self.acceptance == "metropolis_eq" or self.acceptance == "metropolis-eq":
-                prob_acc = min(1, np.exp(-dU))
-            else:  # must be 'metropolis', which consider lambda weights as well
-                g0 = weights[swap[0]]
-                g1 = weights[swap[1]]
-                dg_0 = g0[new_state_0] - g0[old_state_0]  # g^{i}_{n} - g^{i}_{m}
-                dg_1 = g1[new_state_1] - g1[old_state_1]  # g^{j}_{m} - g^{j}_{n}
-                dg = dg_0 + dg_1  # kT
-
-                # Note that simulations with different lambda ranges would have different references
-                # so g^{i}_{n} - g^{j}_{n} or g^{j}_{m} - g^{i}_{m} wouldn't not make sense.
-                # We therefore print g^i_n - g^i_m and g^j_m - g^j_n instead even if they are less interesting.
-                if self.verbose is True:
-                    print(
-                        f"  g^i_n - g^i_m = {dg_0:.2f} kT, g^j_m - g^j_n = {dg_1:.2f} kT, Total dg: {dg:.2f} kT"
-                    )
+            prob_acc = min(1, np.exp(-dU))
 
-                prob_acc = min(1, np.exp(-dU + dg))
         return prob_acc
 
     def accept_or_reject(self, prob_acc):
@@ -1223,7 +1200,7 @@ def prepare_weights(self, weights_avg, weights_final):
     def combine_weights(self, hist, weights, weights_err=None, print_values=True):
         """
         Combine alchemical weights across multiple replicas and adjusts the histogram counts
-        corerspondingly. Note that if :code:`weights_err` is provided, inverse-variance weighting will be used.
+        correspondingly. Note that if :code:`weights_err` is provided, inverse-variance weighting will be used.
         Care must be taken since inverse-variance weighting can lead to slower
         convergence if the provided errors are not accurate. (See :ref:`doc_w_schemes` for mor details.)
 

ensemble_md/tests/data/dhdl/dhdl_1.xvg

@@ -82,4 +82,4 @@
 3.9400    3 -27256.027 -8.4877577 -19.770096 6.3658122 4.2438757 2.1219381 0.0000000 -0.43802848 4.1566980
 3.9600    4 -27277.781 0.62781328 23.077715 -4.9004017 -4.7434649 -4.5865131 -4.4295602 0.0000000 6.5618884
 3.9800    1 -27256.410 9.1837435 20.394253 -2.2959371 0.0000000 2.2959374 4.5918728 10.720802 18.117647
-4.0000    1 -27335.402 5.6179891 21.115717 -1.4045356 0.0000000 1.4045001 2.8089978 9.0623788 16.527695
+4.0000    1 -27335.402 5.6179891 21.115717 -1.4045356 0.0000000 1.4045001 2.8089978 11.0623788 16.527695