Umbrella sampling example in Jupyter

This notebook mirrors the Umbrella Sampling workflow in the OpenMM Cookbook tutorial. It applies harmonic biases to the distance between two alanine dipeptide atoms, records multiple windows, and plots the reconstructed free-energy relation ((F(x) = -k_B T \ln p(x))).

Table of contents

• Step 1
• Step 2
• Step 3
• Step 4
• Step 5
• Step 6
• Step 7

Step 1

from pathlib import Path
import os

import matplotlib.pyplot as plt
import numpy as np
from openmm import app, unit
import openmm as mm

Step 2

COURSE_DIR = Path(os.environ.get('COURSE_DIR', str(Path.home() / 'Concepcion26'))).expanduser()
DATA_DIR = COURSE_DIR / 'data'
PDB_FILE = DATA_DIR / 'alanine-dipeptide.pdb'
ATOM_PAIR = (0, 1)

Step 3

def build_simulation():
    pdb = app.PDBFile(str(PDB_FILE))
    forcefield = app.ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')
    modeller = app.Modeller(pdb.topology, pdb.positions)
    modeller.addHydrogens(forcefield)
    system = forcefield.createSystem(
        modeller.topology,
        nonbondedMethod=app.NoCutoff,
        constraints=app.HBonds,
    )
    bias = mm.CustomBondForce('0.5*k*(r-r0)^2')
    bias.addPerBondParameter('r0')
    bias.addGlobalParameter('k', 5000.0)
    atom1, atom2 = ATOM_PAIR
    bias.addBond(atom1, atom2, [0.2])
    system.addForce(bias)
    integrator = mm.LangevinIntegrator(
        300 * unit.kelvin,
        1.0 / unit.picosecond,
        2.0 * unit.femtoseconds,
    )
    simulation = app.Simulation(modeller.topology, system, integrator)
    simulation.context.setPositions(modeller.positions)
    return simulation, bias

Step 4

def distance_from_state(state):
    positions = state.getPositions(asNumpy=True)
    atom1, atom2 = ATOM_PAIR
    delta = positions[atom1] - positions[atom2]
    return np.linalg.norm(delta.value_in_unit(unit.nanometer))

Step 5

simulation, bias = build_simulation()
steps_per_window = 400
sample_interval = 40
windows = np.linspace(0.2, 0.35, 4)
records = []
for target in windows:
    bias.setBondParameters(0, ATOM_PAIR[0], ATOM_PAIR[1], [target])
    simulation.context.setVelocitiesToTemperature(300 * unit.kelvin)
    simulation.minimizeEnergy()
    for step in range(steps_per_window):
        simulation.step(1)
        if (step + 1) % sample_interval == 0:
            state = simulation.context.getState(getPositions=True, enforcePeriodicBox=False)
            records.append((target, distance_from_state(state)))
records = np.array(records)

Step 6

hist, bins = np.histogram(records[:, 1], bins=25)
bin_centers = 0.5 * (bins[:-1] + bins[1:])
kT = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA * 300 * unit.kelvin
free_energy = -kT * np.log(hist / hist.sum() + 1e-9)
free_energy = free_energy.value_in_unit(unit.kilojoules_per_mole)
plt.plot(bin_centers, free_energy, '-o')
plt.xlabel('Distance (nm)')
plt.ylabel('Free energy (kJ/mol)')
plt.title('Umbrella sampling rough profile')
plt.grid(True)
plt.show()

Step 7

import csv

output_file = DATA_DIR / 'umbrella_samples.csv'
with output_file.open('w', newline='') as handle:
    writer = csv.writer(handle)
    writer.writerow(['target_nm', 'sampled_nm'])
    writer.writerows(records)
print(f'Saved {len(records)} samples to {output_file}')

The histogram above reproduces the cookbook’s plotting step. Export your data to docs/data/umbrella_samples.csv and run the documented WHAM commands from the Bash section to rebuild the final free-energy curve.