<jupyter_start><jupyter_text>Compute PMF of butane dihedral with umbrella sampling and FastMBAR<jupyter_code>## import required packages
import os, math, sys
import numpy as np
import matplotlib.pyplot as plt
import simtk.openmm.app  as omm_app
import simtk.openmm as omm
import simtk.unit as unit
from tqdm import tqdm
import mdtraj
from FastMBAR import *<jupyter_output><empty_output><jupyter_text>**NOTE:** The following commands use relative directory pathees and assume that the current working directory is FastMBAR/examples/butane. Therefore, the current working directory has to be changed into FastMBAR/example/butane before running the following commands.<jupyter_code>## change the working directory into ./examples/butane
os.chdir("../")
print("current working directory: {}".format(os.getcwd()))
assert os.getcwd()[-7:] == "/butane", ( 
       f"Error: your current working directory is {os.getcwd()}!"
        "You need to change working directory to ./examples/butane."
       )<jupyter_output>current working directory: /Users/dingxq/Google Drive/ProjectsOnGitHub/FastMBAR/examples/butane<jupyter_text>1. Construct an OpenMM system of butaneBecause we are using OpenMM as our MD engine, we need to setup theMD molecular system in the format required by OpenMM. The format/objectused by OpenMM for a molecular system happens to be a class called [System](http://docs.openmm.org/latest/api-python/generated/simtk.openmm.openmm.System.htmlsimtk.openmm.openmm.System).Therefore, we will prepare our MD molecular system as an OpenMM System.When we prepare the OpenMM system, we add a [CustomTorsionForce](http://docs.openmm.org/latest/api-python/generated/simtk.openmm.openmm.CustomTorsionForce.htmlsimtk.openmm.openmm.CustomTorsionForce) sothat we can add biasing potentials to the system in the following umbrellasampling. Read **psf** and **pdb** files of butane: butane.psf and butane.pdb.The psf file, butane.psf, contains topolgy of butane and it is the topolgy file format used by CHARMM.The psf file, butane.psf, used here is generated using CHARMM.In your study, you usually alreay have a pdb file of your system. You can generate the topology file of your system using various MD softwares such as CHARMM, Gromacs and Amber among others. Just note that different softwares use different format for topology files and OpenMM has several parser for topology files with different foramt.<jupyter_code>psf = omm_app.CharmmPsfFile('./data/butane.psf')
pdb = omm_app.PDBFile('./data/butane.pdb')<jupyter_output><empty_output><jupyter_text>Read CHARMM force field for butane. The CHARMM force field is downloaded from [here](http://mackerell.umaryland.edu/charmm_ff.shtml)<jupyter_code>params = omm_app.CharmmParameterSet('./data/top_all35_ethers.rtf',
                                    './data/par_all35_ethers.prm')<jupyter_output><empty_output><jupyter_text>Create a OpenMM system based on the psf file of butane and the CHARMM force field.Then a harmonic biasing potential for butane dihedral (3-6-9-13)is added to the system so that we can use the biasing potential in the following umbrella sampling.Adding a biasing potential to a system is very easy in OpenMM. We don't have to change any source code of OpenMM. All we need to do is to tell OpenMM the formula of the biasing potential and degree of freedome we want to add the biasing potential to.<jupyter_code>## creay an OpenMM system
system = psf.createSystem(params, nonbondedMethod=omm_app.NoCutoff)

## add a harmonic biasing potential on butane dihedral to the OpenMM system
bias_torsion = omm.CustomTorsionForce("0.5*K*dtheta^2; dtheta = min(diff, 2*Pi-diff); diff = abs(theta - theta0)")
bias_torsion.addGlobalParameter("Pi", math.pi)
bias_torsion.addGlobalParameter("K", 1.0)
bias_torsion.addGlobalParameter("theta0", 0.0)
## 3, 6, 9, 13 are indices of the four carton atoms in butane, between which
## the dihedral angle is biased.
bias_torsion.addTorsion(3, 6, 9, 13)
system.addForce(bias_torsion)<jupyter_output><empty_output><jupyter_text>After constructing the OpenMM system of butane, we can save it in an XML formated text file, which can be used later for simulations. Therefore, if we want to use the same system in another script, we can just read the text file in an OpenMM system instead of constructing it again.You can even open the XML formated text file using a text editor and see what information about the system is included in the XML file.<jupyter_code>## save the OpenMM system of butane
with open("./output/system.xml", 'w') as file_handle:
    file_handle.write(omm.XmlSerializer.serialize(system))<jupyter_output><empty_output><jupyter_text>2. Run umbrella samplingWe run umbrella sampling for the butane dihedral (atom indices: 3-6-9-13).The dihedral is split into multiple windows and in each window, the dihedralis restrainted around a center using a harmonic biasing potential. In thisscript, we run simulations in each window sequentially, but they can be run inparallel if you have a computer cluster with multiple nodes.<jupyter_code>## read the OpenMM system of butane
with open("./output/system.xml", 'r') as file_handle:
    xml = file_handle.read()
system = omm.XmlSerializer.deserialize(xml)

## read psf and pdb file of butane
psf = omm_app.CharmmPsfFile("./data/butane.psf")
pdb = omm_app.PDBFile('./data/butane.pdb')<jupyter_output><empty_output><jupyter_text>In order to run simulations in OpenMM, we need to construct an OpenMM context, which consists of a system, an integrator and a platform.The system is just the butane system we have constructed above.The integrator specifies what kind integration method we should use. Here, we will use Langevin dynamics for NVP ensemble simulation, which corresponds to the OpenMM.LangevinIntegrator.The platform specifies what kind of hardware we will run simulation on. Here, we choose to use CPUs because the size of the system is quite small.<jupyter_code>#### setup an OpenMM context

## platform
platform = omm.Platform.getPlatformByName('CPU')

## intergrator
T = 298.15 * unit.kelvin  ## temperature
fricCoef = 10/unit.picoseconds ## friction coefficient 
stepsize = 1 * unit.femtoseconds ## integration step size
integrator = omm.LangevinIntegrator(T, fricCoef, stepsize)

## construct an OpenMM context
context = omm.Context(system, integrator, platform)<jupyter_output><empty_output><jupyter_text>Here we set the force constant and centers for the harmonic biasing potentials on butane dihedral.<jupyter_code>## set force constant K for the biasing potential. 
## the unit here is kJ*mol^{-1}*nm^{-2}, which is the default unit used in OpenMM
K = 100
context.setParameter("K", K)

## M centers of harmonic biasing potentials
M = 10
theta0 = np.linspace(-math.pi, math.pi, M, endpoint = False)
np.savetxt("./output/theta0.csv", theta0, delimiter = ",")<jupyter_output><empty_output><jupyter_text>The following loop is used to run umbrella sampling window by window.In each iteration, we first set the center of the harmonic biasing potential.Then the configuration of butane is minimized and equilibrated with the biasing potential.After initial equilibration, configurations are sampled and saved.<jupyter_code>## the main loop to run umbrella sampling window by window
for theta0_index in range(M):
    print(f"sampling at theta0 index: {theta0_index} out of {M}")

    ## set the center of the biasing potential
    context.setParameter("theta0", theta0[theta0_index])

    ## minimize
    context.setPositions(pdb.positions)
    state = context.getState(getEnergy = True)
    energy = state.getPotentialEnergy()
    for i in range(50):
        omm.LocalEnergyMinimizer_minimize(context, 1, 20)
        state = context.getState(getEnergy = True)
        energy = state.getPotentialEnergy()

    ## initial equilibrium
    integrator.step(5000)

    ## sampling production. trajectories are saved in dcd files
    file_handle = open(f"./output/traj/traj_{theta0_index}.dcd", 'bw')
    dcd_file = omm_app.dcdfile.DCDFile(file_handle, psf.topology, dt = stepsize)
    for i in tqdm(range(1000)):
        integrator.step(100)
        state = context.getState(getPositions = True)
        positions = state.getPositions()
        dcd_file.writeModel(positions)
    file_handle.close()<jupyter_output>sampling at theta0 index: 0 out of 10<jupyter_text>3. Compute and collect the butane dihedral valuesFor configurations in trajectories sampled using umbrella sampling, we compute the butane dihedral andsave them in csv files. The dihedral is between four atoms with indices of 3, 6, 9, and 13. Here we use the Python package mdtraj to compute dihedrals.<jupyter_code>topology = mdtraj.load_psf("./data/butane.psf")
for theta0_index in range(M):
    traj = mdtraj.load_dcd(f"./output/traj/traj_{theta0_index}.dcd", topology)
    theta = mdtraj.compute_dihedrals(traj, [[3, 6, 9, 13]])
    np.savetxt(f"./output/dihedral/dihedral_{theta0_index}.csv", theta, fmt = "%.5f", delimiter = ",")<jupyter_output><empty_output><jupyter_text>4. Use FastMBAR to solve MBAR/UWHAM equations and compute the PMF Two steps are required to compute PMF using FastMBAR based on umbrella sampling.Firstly, we need to compute the relative free energies of the biased ensembles used in umbrella sampling, i.e., the NVT ensembles with biased potential energies.Secondly, samples from umbreall sampling are reweighted to compute the PMF. Simulations in umbrealla sampling have different biasing potential energies. They are viewed as different thermodynamic states. Therefore, we have $M$ states and samples from these states.As shown in Usage, we can use FastMBAR to compute the relative free energies of these M states.In order to do it, we need to compute the reduced energy matrix $A_{M,N}$ as shown in Fig. 1, where $U(x)$ is the potential energy function; $B_i(x)$ is the biasing potential added in the $i$th state. In this case, the biasing potential is added to the dihedral (3-6-9-13), so $B_i(x) = 0.5*k*\Delta^2$, where $\Delta = min(|\theta(x) - \theta^0_i|, 2\pi - |\theta(x) - \theta^0_i|)$  and $\theta(x)$ is the dihedral (3-6-9-13) calculated based on Cartesian coordinates $x$. ![Fig_1.png](attachment:Fig_1.png) Compared to general cases, the reduced potential energy matrix $A_{M,N}$ in umbrella sampling has a special property. The energy functions of the $M$ states are $U(x) + B_i(x)$. They all have the common component $U(x)$.Removing the common component $U(x)$ from the energy matrix $A_{M,N}$ does not affect the relative free energies of the $M$ states. Therefore, we can ommitting computing $U(x)$ when compute the energy matrix $A_{M,N}$, as shown in Fig. 2 ![Fig_2.png](attachment:Fig_2.png) As shown in Fig. 2, we can compute the reduced energy matrix $A_{M,N}$ just based on dihedral values from umbrella sampling. In the following script, we read the dihedral values and compute the reduced energy matrix $A_{M,N}$. Based on the reduced energy matrix and the number of conformations sampled from each state, we can compute the relative free enegies of the $M$ states using FastMBAR.<jupyter_code>## read dihedral values from umbrella sampling
thetas = []
num_conf = []
for theta0_index in range(M):
    theta = np.loadtxt(f"./output/dihedral/dihedral_{theta0_index}.csv", delimiter = ",")
    thetas.append(theta)
    num_conf.append(len(theta))
thetas = np.concatenate(thetas)
num_conf = np.array(num_conf).astype(np.float64)
N = len(thetas)

## compute reduced energy matrix A
A = np.zeros((M, N))
K = 100
T = 298.15 * unit.kelvin
kbT = unit.BOLTZMANN_CONSTANT_kB * 298.15 * unit.kelvin * unit.AVOGADRO_CONSTANT_NA
kbT = kbT.value_in_unit(unit.kilojoule_per_mole)

theta0 = np.loadtxt("./output/theta0.csv", delimiter = ",")

for theta0_index in range(M):
    current_theta0 = theta0[theta0_index]
    diff = np.abs(thetas - current_theta0)
    diff = np.minimum(diff, 2*math.pi-diff)
    A[theta0_index, :] = 0.5*K*diff**2/kbT

## solve MBAR equations using FastMBAR
fastmbar = FastMBAR(energy = A, num_conf = num_conf, cuda=False, verbose = True)
print("Relative free energies: ", fastmbar.F)<jupyter_output>solve MBAR equation using CPU
Relative free energies:  [0.         2.69648922 3.96118042 0.94195765 3.07990215 8.28713239
 3.27849569 1.07600058 4.24381136 2.39899694]<jupyter_text>Now we are ready to compute the PMF.Solving MBAR equations yields the relative free energies of the $M$ states, all of which have biasing potential enegies. Knowing the relative free enegies of the $M$ states enables us to compute the PMF using an easy reweighting procesure.In order to do that, we need to compute the energy matrix $B_{L,N}$ as shown in Fig. 1 and Fig. 2. To represent the PMF of the dihedral, we split the dihedral range, $[-\pi, \pi]$ into $L$ windows: $[\theta_{l-1}, \theta_l]$ for $l = 1, ..., L$.Then we can represent the PMF by computing the relative free energies of these $L$ states each of which has a potential energy of $U(x)$.Because the $l$th state is constrainted in the dihedral range $[\theta_{l-1}, \theta_l]$, we need to add a biasing potential $R_l(\theta)$ to enforce the constraint.The value of the biasing potential $R_l(\theta) = R_l(\theta(x))$ is 0 when $\theta \in [\theta_{l-1}, \theta_l]$, infinity otherwise.<jupyter_code>## compute the reduced energy matrix B
L = 25
theta_PMF = np.linspace(-math.pi, math.pi, L, endpoint = False)
width = 2*math.pi / L
B = np.zeros((L, N))

for i in range(L):
    theta_center = theta_PMF[i]
    theta_low = theta_center - 0.5*width
    theta_high = theta_center + 0.5*width

    indicator = ((thetas > theta_low) & (thetas <= theta_high)) | \
                 ((thetas + 2*math.pi > theta_low) & (thetas + 2*math.pi <= theta_high)) | \
                 ((thetas - 2*math.pi > theta_low) & (thetas - 2*math.pi <= theta_high))

    B[i, ~indicator] = np.inf

## compute PMF using the energy matrix B
PMF, _ = fastmbar.calculate_free_energies_of_perturbed_states(B)

## plot the PMF
fig = plt.figure(0)
fig.clf()
plt.plot(theta_PMF*180/math.pi, PMF, '-o')
plt.xlim(-180, 180)
plt.xlabel("dihedral")
plt.ylabel("reduced free energy")
plt.savefig("./output/PMF_fast_mbar.pdf")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09565831274267561, 'fraction_numerical': 0.018355100600070597, 'mean_word_length': 4.74219700040535, 'pattern_counts': {'":': 0, '<': 53, '<?xml version=': 0, '>': 53, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 4}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '13', 'dir': 'jupyter-structured-clean-dedup', 'id': '132696', 'license': 'MIT', 'n_tokens_mistral': 4663, 'n_tokens_neox': 4372, 'n_words': 1617, 'path': 'examples/butane/script/main.ipynb', 'repo_name': 'xqding/FastMBAR', 'size': '14165'}