Protonation state sampling

Strategies for constant-pH simulation for atomistic potentials without explicit bonded terms

Theory

Preliminaries

Consider a system at equilibrium in a semigrand ensemble relevant for constant-pH simulation.

The thermodynamic parameters $\theta \equiv \{ \beta, p, \mu_{H+} \}$ denote the external inverse thermal energy $\beta = (k_B T)^{-1}$, external pressure $p$, and external proton chemical potential $\mu_{H+}$. We presume we have fixed the number of all other elements in the system, and exchange with the environment only occurs for protons.

Denote the instantaneous configuration of the system by $x$, where the relevant scalar properties conjugate to the thermodynamic parameters are the potential energy $U(x)$, the instantaneous volume $V(x)$, and the instantaneous number of protons $N_{H+}(x)$.

We can write the reduced potential $u(x; \beta, p, \mu)$ for this ensemble as

$$ u(x; \beta, p, \mu) \equiv \beta U(x) + p V(x) + \mu_{H+} N_{H+}(x) $$

(NOTE: The inverse thermal energy $\beta$ has been absorbed into $p$ and $\mu_{H+}$ in this equation, which differs from how we traditionally write the reduced potential. We may think of $\beta, p, \mu_H$ as derived thermodynamic parameters that express the parameters of interest in a convenient manner. If one of the principal thermodynamic parameters (such as temperature) needs to be changed, it will modify one or more of the derived thermodynamic parameters.)

If we instead fix the number of protons at $N_H$, we can instead write the reduced potential relevant for the isothermal-isobaric ensemble.

$$ u(x; \beta, p, N_{H+}) \equiv \beta U(x) + p V(x) $$

The probability density of interest is given by

$$ \pi(x; \theta) = Z(\theta)^{-1} q(x; \theta) $$

where the unnormalized probability density $q(x; \theta)$ is given by the Boltzmann law

$$ q(x; \theta) \equiv e^{-u(x; \theta)} $$

and the semi-grand canonical partition function $Z(\theta)$ is given by a sum over the number of protons:

$$ Z(\theta) \equiv Q(\beta, p, \mu_{H+}) = \sum_{N_{H+}=1}^{\infty} Q(\beta, p, N_{H+}) e^{-\mu_{H+} N_{H+}} = \sum_{N_{H+}=1}^{\infty} e^{g_{N_{H+}}-\mu N_{H+}} $$

and the isothermal-isobaric partition function $Q(\beta, p, N_{H+})$ is

$$ Q(\beta, p, N_{H+}) \equiv \int dx e^{-u(x; \beta, p, N_{H+})} $$

We have used the shorthand

$$ g_{N_{H+}} \equiv - \ln Q(\beta, p, N_{H+}) $$

in the equation above.

Note: We use the shorthand $\mu_H$ and $N_H$ below, rather than writing $H+$ throughout.

Handling discrepancies between experimental and calculated pKas

Broadly, there are two major strategies that could be adopted to address discrepancies between experimental and calculated solution pKas for water (and other solvents), small molecules, and amino acid residues:

Strategy 1: Calibration

Broadly, this strategy is to compute the effective chemical potential for the ML potential, subtract it off, and add back the desired chemical potential.

First, we consider the case of correcting the water pKa to provide appropriate pH. Roughly, this strategy involves the development of a mapping from the desired pH to an effective chemical potential that corrects for the deviation between experimental and computed chemical potentials for bulk water.

Consider a water box in equilibrium with a chemical reservoir of protons at a given pH.

Presuming we implicitly specify a fixed number $N_O$ (which we will suppress writing for convenience), the probability of a given number of protons, $N_H$, is given by

$$ p(N_H | \beta, p, \mu) = \frac{\exp[g_{N_H} - \mu N_H]}{\sum_{N_H'=0}^\infty \exp[g_{N_H'} - \mu_H N_H']} $$

If we knew the relative free energies $g_{N_H}$ for $N_H = 0, 1, \ldots$, we could determine $\mu_H$ from a quantity such as $< N_H >$ by solving the nonlinear relationship for $\mu$ given a pH:

$$ < N_H >_\mu = \sum_{N_H=0}^{\infty} N_H p(N_H | \beta, p, \mu) = \frac{N_H \exp[g_{N_H} - \mu N_H]}{\sum_{N_H'=0}^{\infty} \exp[g_{N_H'} - \mu_H N_H']} $$

The definition of pH is in terms of the activity $a_H$ of protons:

$$ pH \equiv - \log_{10} a_H \approx - \log_{10} ([H+]/c_0) = - \log_{10} \frac{< N_H >_\mu}{< V >_\mu} $$

where $[H+]$ is the concentration of protons and $c_0$ is the standard state concentration of 1 Molar.

(TODO: Revise this section for a better handling of the activity.)

This suggests that, should we need to calibrate $\mu_H$ to produce a specified pH from known $\{ g_{N_H} \}$, we could do so by solving for $\mu(pH)$ via the relationship

$$ 10^{-pH} c_0 < V >_\mu = < N_H >_\mu = \frac{N_H \exp[g_{N_H} - \mu N_H]}{\sum_{N_H'=0}^\infty \exp[g_{N_H'} - \mu_H N_H']} $$

This approach was used, for example, in calibrating the chemical potential of neutral salt pairs in TIP3P water in SaltSwap.

However, it is likely we can exploit a more direct relationship between $\mu$ and pH via the activity $a_{H+}$ if we have an accurate potential $U(x)$ that does not require any kind of explicit calibration step. We shall explore this later.

Strategy 2: Training or fine-tuning potentials to correct chemical potentials

Much the same way we expect to include loss terms in training or fine-tuning that penalize deviations from experimental binding free energies, we can penalize deviations between computed and experimental chemical potentials.

Note that there may be finite-size effects that require us to extrapolate unless we use sufficiently large solvent boxes, but this can be rapidly established by computing the chemical potential of proton insertion/deletion for neutral bulk water as a function of $N_O$.

Because experimental pKas may not reliably provide direct access to chemical potentials---see this paper for a more detailed explanation of experimental pKa measurements and their challenges---it may be desirable to fit directly to experimental electrochemical (or UV-metric) titration data that can be directly computed from knowledge of the proton-dependent free energies $\{ g_{N_H} \}$ of the species of interest.

Similar to binding free energies, the free energies of relevant protonation states could be computed at a reference parameter set $\theta_0$ and, as parameters are perturbed to $\theta_1, \theta_2, \ldots$, equilibrium simulations of solvated protonation species could be used to estimate the free energies at the new parameter sets $\theta_n$ via reweighting techniques such as MBAR.

Because protons may migrate in solution, we might reasonably wonder whether the free energy differences we compute are for the species of interest. We don't know! We can use heuristics to assign molecularity and compute chemical species, but this is only approximate. In reality, we can simply compute the ladder of free energies $\{ g_{N_H} \}$ for a box of $N$ $H_2 O$ molecules, and the perturbed $\{ g_{N_H}^{(M)} \}$ for the same box of $N$ $H_2 O$ molecules with a molecule $M$ present. While this is rigorous (assuming convergence) and should provide us with the ability to fine-tune a potential to electrochemical titrations, it can be computationally expensive to do so.

We might therefore explore whether there are shortcuts, such as:

• Using gas-phase proton affinities or related energies instead
• Using restraints to constrain the chemical species being sampled (e.g. preventing proton dissociation or migration)
• Computing the solution free energy difference between protonation states of the molecule of interest only for neutral water, rather than with a significant $\Delta N_H$.

Algorithms

This section discusses several algorithms that may be suitable for sampling protonation states and tautomers in the context of Markov chain Monte Carlo molecular simulations.

Tautomerization

For simplicity, we first consider algorithms for dealign with tautomerization. These algorithms can be easily adapted to handle protonation states in a straightforward manner.

Consider the tautomerization of imidazole

(TODO: Replace this with a better figure)

Instantaneous Monte Carlo with heavy atom centered Gaussians

The simplest approach we will consider is an instantaneous Metropolis Monte Carlo proposal where we propose to "hop" the proton from one titratable heavy atom site to another.

• Step 1: Propose $x_{new} \sim p(x_{new} | x_{old})$ where the proton at position 1 is moved to a position that is a Gaussian random variable about the N at position 3.
• Step 2: Accept or reject the move with the Metropolis-Hastings criteria

$$ P_{accept} = \min \{ 1, \frac{p(x_{old} | x_{new})}{p(x_{new} | x_{old})} \frac{\pi(x_{new})}{\pi(x_{old})} \} $$

where we can use the difference in reduced potentials $\Delta u \equiv u(x_{new}) - u_(x_{old})$ to write

$$ P_{accept} = \min \{ 1, \frac{p(x_{old} | x_{new})}{p(x_{new} | x_{old})} e^{-\Delta u} \} $$

Instantaneous Monte Carlo with improved proposal densities

We expect the average acceptance rate of the instantaneous heavy atom centered Monte Carlo proposal $&lt; P_{accept} &gt;$ to be low in both vacuum and solvent. In particular, the proposal ratio $r \equiv \frac{p(x_{old} | x_{new})}{p(x_{new} | x_{old})}$ will generally be low because the probability of generating the old proton placement $p(x_{old} | x_{new})$ will be low for this poor choice of proposal probability.

Instead, we can significantly improve the choice of proposal probability by learning the parameters of the Gaussian distribution in terms of the heavy atom and its connected neighbors. Suppose we have an arrangement $X-NH-Y$ within a molecule and we have learned deterministic functions for the displacement $\Delta \mu(x)$ and standard deviation $\sigma(x)$ of the position of $H$ given $X-N-Y$ from equilibrium simulations through some manner of equivariant neural network. We can use these in a Gaussian proposal density to sample the new proton position

$$ x_H \sim \mathcal{N}(x_N ; \Delta \mu(x), \sigma(x)) $$

where $x_H$ is the coordinates of the proton H being placed and $x_N$ is the coordinates of the heavy atom $N$ it is being placed near.

We expect this should significantly increase the acceptance rate $&lt; P_{accept} &gt;$ by improving the proposal probability ratio $r$ in vacuum.

Parallel instantaneous Monte Carlo proposals

Even with improved proton proposal distributions, in solvent, instantaneous placement of protons without allowing relaxation of the surrounding solvent will likely lead to low acceptance rates $&lt; P_{accept} &gt;$ due to the probability density ratio $s = \frac{\pi(x_{new})}{\pi(x_{old})}$ appearing in the Metropolis-Hastings acceptance probability.

A very simple way to improve this is to perform parallel proposals: If we propose $M$ moves in parallel, $x_{new}^{(1)}, \ldots, x_{new}^{(M)}$, we can modify the acceptance criteria to select either the old position $x_{old}$ or one of the new proposed positions via parallel Monte Carlo. Here, we select among the list of potential configurations $y \in \{ x_{old}, x_{new}^{(1)}, \ldots, x_{new}^{(M)} \}$ using the symmetric rule:

$$ P_{accept}(y) \propto p(y | x_{old}) \pi(y) $$

(TODO: Check that we have correctly included the proposal probabilities $p(y | x_{old})$.)

In the limit of $&lt; P_{accept} &gt; &lt;&lt; 1/M$, this approach will give near-linear improvements in acceptance probability in $M$. This comes at the cost of $M$ energy evaluations, but it is possible that a very wide GPU can enable $M$ energies to be computed in parallel at a cost of $L &lt; M$.

This idea is succinctly described in Waste recycling Monte Carlo.

(TODO: Check if there are issues with this acceptance probability. Multiple-try Metropolis is an alternative.)

Nonequilibrium candidate Monte Carlo (NCMC) proposals : transient augmentation

It is possible to achieve superlinear improvements in acceptance probability $&lt; P_{accept} &gt;$ through the use of Nonequilibrium Candidate Monte Carlo (NCMC).

There are numerous ways to use NCMC for this scenario. First, we consider building on the previous algorithms that draw a new proton position using one of the hydrogen position proposal probabilities $p(x_H | x)$ above in a scheme that temporarily augments the system with a new proton in addition to the old proton, switching the old proton interactions off as the new proton is switched on.

First, partition the atoms into the proton under consideration for transport---$x_{H, old}$---and all other protons---$x_{core}. We draw the position of a new additional proton $x_{H, new}$:

$$ x_{H, new}^{(0)} \sim p(x_{H} | x_{old}) $$

We then integrate $t = 1, \ldots, T$ steps of molecular dynamics using the alchemical potential energy function $U_t(x)$ given by an interpolative alchemical energy function such as

$$ U_t(x) = (1-\frac{t}{T}) U(x_{H, old}, x_{core}) + \frac{t}{T} U(x_{H, new}, x_{core}) $$

where $x_{core}$ denotes all atoms except the proton being hopped, and integrate the dimensionless work $w$, taking care to properly compute the work appropriate for the integrator being used (see NCMC). For example, if velocity Verlet integration is used, the dimensionless work $w$ is equivalent to the protocol work

$$ w = \sum_{t=1}^T [ u_{t}(x_t) - u_{t-1}(x_t) ] $$

The acceptance probability for $x_T$ is then given by replacing the reduced potential difference $\Delta u$ in the Metropolis-Hastings criteria with the work $w$:

$$ P_{accept} = \min \{ 1, \frac{p(x_{H, old, T} | x_T)}{p(x_{H, new, 0} | x_0)} e^{-\Delta w} \} $$

If rejected, however, the entire trajectory $x_{1}, \ldots, x_{T}$ must be discarded.

The velocity will also need to be reversed on acceptance or rejection to ensure the correct distribution is maintained---see NCMC for more details. If an MCMC scheme is used where protonation state moves and MD moves that begin with a new velocity draw from the Maxwell-Boltzmann distribution, this can be ignored since the velocity will be refreshed immediately afterwards anyway.

Note that, we can explore various alchemical potentials $U_t(x)$ that may not require the computation of the potential energy twice. For example, by introducing a per-particle alchemical parameter, we may be able to avoid this additional expense. However, the integration of per-particle alchemical parameters is likely to be highly architecture-dependent, and may only be feasible for certain architecture classes for neural network potentials. This is an entire field of research, so we leave this to later.

Nonequilibrium candidate Monte Carlo (NCMC) proposals : displacement

As an alternative to transiently augmenting the number of particles, we can instead propose a displacement $\Delta x$ that intends to transport the proton from one likely proton location to another, using an alchemical protocol similar to the Alchemical Transfer Method (ATM):

$$ \Delta x \sim p(\Delta x | x_0) $$

where the alchemical energy function becomes

$$ U_t(x) = (1-\frac{t}{T}) U(x) + \frac{t}{T} U(x + \Delta x) $$

The acceptance probability replaces the proposal ratio $r$ with one involving $\Delta x$:

$$ P_{accept} = \min \{ 1, \frac{p(-\Delta x | x_T)}{p(\Delta x | x_0)} e^{-\Delta w} \} $$

Nonequilibrium candidate Monte Carlo (NCMC) proposals : proton wires

Excess protons are surprisingly mobile in bulk water. Water wires form spontaneously and enable the rapid nonlocal translocation of excess protons. In confined geometries, proton transport over large distances can occur on the sub-picosecond timescale.

As a result, we may not need to be overly concerned about where a proton is initially placed if we use NCMC to transport or insert/delete a proton---spontaneous water wires may effectively shuttle protons to/from relevant locations on their own. Whether this timescale is practical, of course, is a question for the Experiments section below.

Suppose we select a position for placing the proton that is uniform within the simulation box:

$$ p(x_H | x) = 1/V(x) $$

If we again integrate an NCMC switching protocol over $T$ steps as the old proton is switched off and the new proton switched on, the acceptance probability simply becomes

$$ P_{accept} = \min \{ 1, \frac{1/V(x_T)}{1/V(x_0)} e^{-\Delta w} \} = \min \{ 1, \frac{V(x_0)}{V(x_T)} e^{-\Delta w} \} $$

The acceptance probability will increase superlinearly in $T$ until some timescale commensurate with proton transport within the system, when it will begin to increase sub-linearly with $T$.

Protonation state sampling

Nonequilibrium candidate Monte Carlo (NCMC) proposals : proton wires

We can build on this idea by simply attempting to insert/delete protons within the bulk.

Consider the incredibly simple algorithm:

Insertion:

• Propose proton N_H+1 position x_H uniformly within the box
• Perform NCMC insertion under NpT conditions where the protocol work for just the potential is integrated:

$$ U_t(x) = (1 - \frac{t}{T}) U(x; N_H) + \frac{t}{T} U(x, x_H; N_H+1) $$

• Accept or reject with the probability

$$ P_{accept} = min \{ 1, \frac{1}{N_H+1} \frac{V(x_0)}{V(x_T)} e^{-w} e^{-\mu} \} $$

Deletion:

• Select a proton uniformly for deletion
• Perform NCMC deletion under NpT conditions where the protocol work for just the potential is integrated:

$$ U_t(x) = (1 - \frac{t}{T}) U(x; N_H) + \frac{t}{T} U(x, x_H; N_H+1) $$

• Accept or reject with the probability

$$ P_{accept} = min \{ 1, \frac{1}{N_H} \frac{V(x_0)}{V(x_T)} e^{-w} e^{+\mu} \} $$

TODO: This elides a few steps that simplify the resulting algorithm---make the derivation from super-detailed balance with the reduced potential more explicit. Also check accept/reject proposals are correct.

It is possible to improve the acceptance rates further by biasing the selection of protons via effective pKa estimates and correcting for these proposal probabilities in the Metropolis-Hastings acceptance criteria.

Useful experiments

Computing the acceptance rate for instantaneous Monte Carlo in vacuum and solvent

To set a baseline, it is sensible to attempt to compute $\log &lt; P_{accept} &gt;$ for instantaneous Monte Carlo with Gaussian proposal probabilities in vacuum and solvent to within an order of magnitude or so. This will provide us with a clear picture of whether instantaneous Monte Carlo could be a viable strategy alone or with parallel Monte Carlo.

The procedure is roughly:

1. Sample equilibrium configurations $x \sim \pi(x)$ for imidazole in vacuum or in water at 300 K and 1 atm to generate a collection of equilibrium snapshots $x_n$.
2. Collect statistics on the log proposal ratio $\log r$ and log acceptance ratio $\log s = - \Delta u$ by drawing an equilibrium snapshot $x_n$, proposing a new proton position, and computing $\log r$ and $\log s$
3. Analyze the statistics to estimate $\log &lt; P_{accept} &gt;$, bootstrapping to estimate its posterior uncertainty

We can then perform the same exercise using learned optimal equivariant functions $\Delta \mu(x), \sigma(x)$ which should significantly improve $\log r$, though likely still have poor $\log s$.

Nonequilibrium candidate Monte Carlo (NCMC)

We can then implement various NCMC protocols to explore the distribution of the work $w$ as a function of the nonequilibrium switching time $T$. In particular, the variance $var(w)$ should diminish very rapidly with increasing $T$, then reaching a crossover point at some relevant timescale $T^*$ where doubling $T^*$ halves $var(w)$. This timescale $T^*$ should be near the peak of efficiency in terms of the acceptance probability per energy evaluation, $&lt; P_{accept} &gt; / T$.

What is the chemical potential of a proton?

We can use a Wang-Landau-like scheme (e.g. SAMS, as we used in SaltSwap) to estimate the free energy $g_{N_H} \equiv - \ln Q(\beta, p, N_H)$ as a function of $N_H$ given fixed $N_O$ and thermodynamic parameters $\beta, p$. This involves simulating a water box and using the NCMC moves above to insert/delete protons, adjusting the estimates of $g_{N_H}$ on the fly to ensure equal sampling of all $N_H$ over a specified min and max range.

Note that the total charge of the system must be updated every time a proton is inserted or deleted.

When do finite-size effects vanish?

Following the same experimental strategy above to compute the free energy $g_{\Delta N_H, N_O}$, now as the number of excess protons $\Delta N_H \equiv N_H - N_O$, we compute this for different numbers of oxygens $N_O$ to assess how large a system must be until finite-size effects vanish and differences $g_{\Delta N_H, N_O+1} - g_{\Delta N_H, N_O}$ are much smaller than $k_B T$ (which is unity in these reduced free energy units).

How does adding a small molecule to a box of water perturb the chemical potential of a proton?

Following the same experimental strategy above to compute the free energy $g_{\Delta N_H}$ of a box of water as the number of excess protons is varied, we can compute this with and without a titratable molecule $M$ present in the box to ascertain how much simulation time is required to accurately compute the deviation is between neat water and water with the molecule of interest. This deviation can be used to rigorously compute the effective pKa of the small molecule, but the small differences in $g_{\Delta N_H}$ that give rise to this may suggest more efficient perturbative strategies for computing effective pKas (or electrochemical titrations from which the pKa can be estimated in a manner similar to that used experimentally).