Optimizing Lambda Schedules for Alchemical Free Energy Calculations =================================================================== | Harsh Amin\ :sup:`1`, Zeke A. Piskulich\ :sup:`1`, and Darrin M. York\ :sup:`1` | :sup:`1`\ Laboratory for Biomolecular Simulation Research, Institute | for Quantitative Biomedicine and Department of Chemistry and Chemical | Biology, Rutgers University, Piscataway, NJ 08854, USA Learning objectives ------------------- .. start-learning-objectives - Optimize **lambda schedules** for alchemical free energy calculations with **fetkutils-tischedule.py**. - Compare the three optimization criteria — **phase-space overlap**, **acceptance ratio**, and **Kullback-Leibler divergence** — to choose an appropriate method. - Evaluate an existing schedule with the ``--read`` flag and benchmark it against optimized schedules. - Interpret the **schedule summary statistics** and **prediction/overlap plots** to detect bottlenecks and judge schedule quality. .. end-learning-objectives Activities ---------- After completing burn-in simulations for your alchemical free energy calculations, the next critical step is to optimize the lambda schedule. An optimized schedule maximizes replica exchange efficiency and enhances sampling, leading to more accurate and converged free energy estimates. This tutorial demonstrates how to use the FE-ToolKit utilities to analyze simulation data and generate optimized lambda schedules. :footcite:t:`Zhang_JChemTheoryComput_2024_v20_p3935` .. contents:: :local: :depth: 4 .. start-tutorial Introduction to Optimizing Lambda Schedules ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Why the Lambda Schedule Matters ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ In alchemical free energy simulations, the coupling parameter :math:`\lambda` defines an artificial pathway between two thermodynamic states, such as ligand A :math:`\rightarrow` ligand B in an RBFE calculation, or an interacting ligand :math:`\rightarrow` noninteracting dummy state in an ABFE calculation. Although the free energy is a state function, the reliability of alchemical free energy simulations depends on the ability to robustly sample phase space along an artificial coordinate :math:`\lambda`. In methods like Alchemical Enhanced Sampling (ACES), efficiency is driven by Hamiltonian replica exchange, whose success is highly correlated with the degree of phase space overlap between adjacent :math:`\lambda`-windows. Traditional :math:`\lambda`-spacing schedules often create bottlenecks-intervals with very low exchange probabilities-that stall sampling and prevent the necessary traversal of conformational states. Optimizing the :math:`\lambda`-schedule is necessary to equalize these overlaps, ensuring a high throughput of end-to-end "single passes" and "round trips" across the alchemical path. By tailoring the schedule to the specific phase space of a system, optimization eliminates these bottlenecks, leading to significantly more precise and robust free energy estimates. Optimization Methods: PSO, AR, KL ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Optimized lambda scheduling uses short preliminary, or "burn-in," simulations to estimate how well different lambda states overlap. These data are used to build a two-dimensional map, :math:`O(\lambda,\lambda')`, that predicts the overlap or exchange behavior between any pair of lambda values. For a fixed number of lambda windows, the optimized schedule is chosen (with the endpoints fixed at :math:`\lambda=0` and :math:`\lambda=1`) so that the *variance* of the overlap index across adjacent intervals is minimized, making the overlap between neighboring windows as uniform as possible. **Opt-PSO: Optimized Phase-Space Overlap.** Opt-PSO directly optimizes the phase-space overlap between neighboring lambda windows. For each lambda state, short simulations generate samples from that state. Energy differences between pairs of lambda states are computed and approximated as distributions. The overlap between these distributions gives an overlap index :math:`O_{ij}`, where larger values indicate better similarity between ensembles. The discrete overlap values from burn-in simulations are interpolated into a continuous function :math:`O(\lambda,\lambda')`. The first and last windows are fixed at :math:`\lambda=0` and :math:`\lambda=1`, while the internal lambda values are adjusted to minimize the variance in nearest-neighbor overlaps. In simple terms, Opt-PSO tries to make every adjacent lambda pair share comparable statistical overlap; since PSO and the H-REMD acceptance ratio are strongly correlated, this also tends to produce comparably efficient replica exchange between every pair of neighboring windows. **Opt-KL: Optimized Kullback-Leibler Divergence.** Opt-KL uses the Kullback-Leibler divergence as the measure of dissimilarity between neighboring lambda ensembles. KL divergence measures how different one probability distribution is from another. A symmetric KL divergence can be used so that the difference between state :math:`i` and state :math:`j` is treated consistently in both directions. When two distributions are very similar, the KL divergence is small; when they differ strongly, the KL divergence becomes large. To use KL divergence in a scheduling procedure similar to Opt-PSO, the divergence is converted into a KL-based overlap index that decreases as the divergence grows. The resulting continuous map is optimized in the same general way: choose lambda values that make the neighboring intervals more uniform. In practice, the KL method is no better than PSO. **Opt-AR: Optimized Acceptance Ratio.** Opt-AR optimizes the lambda schedule using the replica-exchange acceptance ratio directly. Instead of estimating overlap through distribution similarity, it estimates the probability that replicas at two lambda states would exchange according to the Metropolis criterion. The average exchange probability becomes the optimization index. Because exchange acceptance is the practical quantity that controls how efficiently walkers move through lambda space, Opt-AR is especially intuitive for HREMD-based workflows, and its underlying acceptance-ratio map tends to decay monotonically away from the diagonal, which can make it the most predictable choice when a schedule needs to be selected automatically. The final optimization is analogous to Opt-PSO: construct a continuous acceptance-ratio map and adjust the internal lambda points so that neighboring exchange probabilities are more balanced across the full lambda path. Running fetkutils-tischedule.py ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Input Files """"""""""" The lambda schedule optimization tool reads the ``efep_*_*.dat`` files produced from your burn-in simulation by ``edgembar-amber2dats.py`` (see :doc:`/ModularTutorials/Alchemical/fe-toolkit/analyzing-abfe-fe-toolkit`). These files must follow a specific naming convention and format: **File naming**: - ``dvdl_.dat``: the :math:`\partial V/\partial\lambda` time series produced at ``lambda``. ``fetkutils-tischedule.py`` uses the presence of these files to automatically detect which lambda values were sampled during the burn-in. - ``efep__.dat``: potential energies evaluated at ``lambda2``, using configurations sampled from the trajectory run at ``lambda1``. **File format** (``efep_*_*.dat``): Two-column, whitespace-separated values - Column 1: Simulation time (in step numbers) - Column 2: Potential energy (kcal/mol) **Example file structure**: .. code-block:: text 0 -12543.234 500 -12538.912 1000 -12541.456 1500 -12539.782 ... To obtain the most reliable schedule, you want a **complete matrix** of energy evaluation files. This means: - For every trajectory generated at lambda value :math:`\lambda_i` - You have energy evaluations at **all** lambda values :math:`\lambda_j` (where :math:`j` ranges from 0 to :math:`N_{windows}`) For example, if you have 16 initial lambda windows (0 to 15), you would ideally have 16 × 16 = 256 total ``efep_*_*.dat`` files. .. note:: The diagonal files (``efep_i_i.dat``) represent energies evaluated at the same lambda as the simulation, and are required to compute the PSO/AR/KL indices for that window. Off-diagonal files are the cross-evaluations needed for the two-dimensional overlap map, as well as for MBAR analysis. Optimizing a Schedule """"""""""""""""""""" For this activity, we provide a set of burn-in energy files (the ``efep_*_*.dat`` files for ligand **ejm_31** bound to Tyk2), already generated with ``edgembar-amber2dats.py`` (see :doc:`/ModularTutorials/Alchemical/fe-toolkit/analyzing-abfe-fe-toolkit`). .. start-skip-download The files you will need are located here: - :download:`Download the activity files ` Extract the archive and change into the resulting directory: .. code-block:: bash tar -xvzf tischedule_example.tar.gz cd tischedule_example .. end-skip-download Inside, you will find the ``dats/ejm_31/complex/vdw/t1/`` directory, organized following the same edge/environment/stage/trial layout that ``edgembar-amber2dats.py`` produces from your Amber AFE outputs (mdouts and rem.log). This directory contains the burn-in energy files that ``fetkutils-tischedule.py`` reads. The FE-ToolKit ``fetkutils-tischedule.py`` utility implements the three optimization criteria described above (Opt-PSO, Opt-AR, and Opt-KL), each estimated from a short burn-in simulation. Choose one of the three: .. tab-set:: .. tab-item:: Acceptance Ratio (AR) .. code-block:: bash fetkutils-tischedule.py --opt 32 --ar \ --plot opt_ar_32.png -o opt_ar_32.txt \ dats/ejm_31/complex/vdw/t1/ .. tab-item:: Phase Space Overlap (PSO) .. code-block:: bash fetkutils-tischedule.py --opt 32 --pso \ --plot opt_pso_32.png -o opt_pso_32.txt \ dats/ejm_31/complex/vdw/t1/ .. tab-item:: Kullback-Leibler (KL) .. code-block:: bash fetkutils-tischedule.py --opt 32 --kl \ --plot opt_kl_32.png -o opt_kl_32.txt \ dats/ejm_31/complex/vdw/t1/ - ``--pso``: Optimize based on the phase space overlap index - ``--ar``: Optimize based on the predicted replica exchange acceptance ratio (recommended) - ``--kl``: Optimize based on the symmetrized Kullback-Leibler divergence - ``--opt N``: Optimize to N lambda windows (commonly 8-24; more windows improve exchange statistics and free energy precision at increased computational cost) - ``--sym``: Force the schedule to be symmetric about :math:`\lambda=0.5`, halving the number of free parameters - ``--ssc``: (optional) Restrict the schedule to the analytic SSC(:math:`\alpha`) smoothstep family instead of optimizing each lambda value independently; 1 parameter with ``--sym``, otherwise 2 (``--alpha0``/``--alpha1``) - ``--maxgap VALUE``: Maximum allowed spacing between adjacent lambda values (default ``1.0``) - ``--plot FILE``: Write a two-panel figure: the predicted index vs. lambda, and the full 2D overlap map with the schedule overlaid - ``-o FILE`` or ``--out FILE``: Write the optimized lambda schedule to file Evaluating an Existing Schedule """"""""""""""""""""""""""""""" To predict the PSO/AR/KL profile of a schedule you already have, without optimizing, pass the schedule file to ``--read``. The file should contain one lambda value per line: This is useful whenever you already have a schedule and want to see how it performs. For example, if you have a custom schedule from some other method — you can predict its AR/PSO/KL statistics with the same ``--read`` command and compare them directly against the optimized schedules produced with ``--opt``. To try this, save the schedule you want to evaluate as a text file with one lambda value per line. For example below is a 24-window evenly spaced (linear) schedule. Copy it into a file named ``production_schedule.txt`` in your working directory. Then run with fetkutils-tischedule.py ``--read`` flag above to predict this linear schedule's statistics. .. code-block:: text 0.00000000 0.04347826 0.08695652 0.13043478 0.17391304 0.21739130 0.26086957 0.30434783 0.34782609 0.39130435 0.43478261 0.47826087 0.52173913 0.56521739 0.60869565 0.65217391 0.69565217 0.73913043 0.78260870 0.82608696 0.86956522 0.91304348 0.95652174 1.00000000 .. code-block:: bash fetkutils-tischedule.py --read production_schedule.txt --pso \ --plot production_schedule_pso.png \ dats/ejm_31/complex/vdw/t1/ - ``--read FILE``: Evaluate an existing schedule without optimizing. ``FILE`` should contain one column; the rows are the lambda values. - ``--pso``, ``--ar``, or ``--kl``: Choose the metric to evaluate - ``--plot FILE``: Generate the two-panel plot of the metric profile Interpreting the Results ^^^^^^^^^^^^^^^^^^^^^^^^ The Schedule File and Summary """"""""""""""""""""""""""""" The output file (e.g., ``opt_ar_32.txt``) contains the optimized lambda schedule, one value per line formatted to 8 decimal places. For example, the 32-window AR-optimized schedule for the ejm31 complex looks like: .. code-block:: text 0.00000000 0.16355000 0.21471000 ........ 0.83827000 0.87418000 1.00000000 These values represent the lambda schedule to use in your production simulations. Note: - First value is always 0.0 - Last value is always 1.0 Every run of ``fetkutils-tischedule.py`` also prints a summary of the schedule to the terminal, including the maximum gap between windows and the mean, standard deviation, maximum, and minimum of the chosen index across all intervals: .. code-block:: text Max Gap: 0.1636 Mean Value: 0.4813 Std Value: 0.0001 Max Value: 0.4815 Min Value: 0.4812 Read this summary as follows: - **Std Value** is the key indicator of schedule quality. Here it is ``0.0001``, meaning the predicted acceptance ratio is essentially identical across all 31 intervals-the schedule is very well optimized, with no bottlenecks. - **Mean Value** (``0.4813``) is the average predicted index across the intervals; every adjacent pair of windows is predicted to exchange with roughly 48% probability. - **Min Value** (``0.4812``) is the weakest interval in the schedule. Here it is essentially equal to the mean, so no interval is a bottleneck. A ``Min Value`` that dropped toward zero would flag an interval where the windows are spaced too far apart and should be subdivided by increasing ``--opt``. - **Max Gap** (``0.1636``) is the largest spacing between adjacent lambda values-here the first interval, from :math:`\lambda=0` to :math:`\lambda=0.16355`, where the phase space changes slowly enough that a wide step still maintains high overlap. The Prediction Plot """"""""""""""""""" The ``--plot`` option writes a two-panel figure that summarizes the same information visually. The figure below is the plot for the 32-window AR-optimized schedule above (``opt_ar_32.png``): .. figure:: /_static/files/ModularTutorials/Alchemical/fe-toolkit/ejm31_complex_ar_32_1.png :alt: Predicted acceptance ratio for the optimized 32-window schedule :width: 700px **Left panel — predicted index vs. lambda.** Each black point is the predicted index (here the acceptance ratio, ``AR``) at the midpoint of an interval, and the red dotted vertical lines mark the 32 optimized lambda windows. The black profile is essentially flat which shows every adjacent pair of windows has the same predicted acceptance ratio, so there is no bottleneck. A poor schedule would instead show dips in this profile, marking intervals where exchange stalls. **Right panel — the two-dimensional overlap map.** This is the interpolated map :math:`O(\lambda_1,\lambda_2)`: red indicates high acceptance/overlap (near the diagonal, where two states are similar) and blue indicates low acceptance/overlap (states too far apart to exchange). The black dots are the adjacent-window pairs :math:`(\lambda_i,\lambda_{i+1})` of the optimized schedule. Because the optimization equalizes the index, the dots all lie along the same narrow red band just off the diagonal — every neighboring pair sits at the same, high acceptance. If a schedule were poorly spaced, some black dots would stray outward into the pale or blue region, revealing a bottleneck interval. .. end-tutorial Relevant Literature ------------------- .. footbibliography::