Simulated annealing (Episode 04 – advanced sampling)

This notebook brings the simulated annealing loop into the Episode 04 notebook collection so you can control the temperature schedule, inspect the potential, and experiment with different cooling rates without leaving the interactive workflow. The same assets we use elsewhere in the project live under COURSE_DIR/data/, so you can copy results into results/04-annealing/ when you want to archive snapshots.

Table of contents

• Step 1
• Paths and data preparation
• System setup
• Simulated annealing schedule
• Run the annealing loop

Step 1

from pathlib import Path
import os
from openmm import app, unit
import openmm as mm
from openmm.app import Modeller, PDBFile, ForceField, Simulation

__all__ = ["Path", "os", "unit", "mm", "Modeller", "PDBFile", "ForceField", "Simulation"]

Paths and data preparation

We reuse the project layout: COURSE_DIR points to the folder that the sync scripts keep populated with the official datasets, so all data/ and results/ paths already live under that tree. The notebook requires alanine-dipeptide.pdb in COURSE_DIR/data/; if it is absent, rerun the standard data sync or contact the maintainer to refresh the Course-MD-Data copy.

COURSE_DIR = Path(os.environ.get("COURSE_DIR", Path.home() / "Concepcion26")).expanduser()
DATA_DIR = COURSE_DIR / "data"
RESULTS_DIR = COURSE_DIR / "results" / "04-annealing"
RESULTS_DIR.mkdir(parents=True, exist_ok=True)

PDB_PATH = DATA_DIR / "alanine-dipeptide.pdb"
if not PDB_PATH.exists():
    raise FileNotFoundError(
        f"Missing alanine-dipeptide.pdb in {DATA_DIR}; rerun the repository's data sync so {DATA_DIR} mirrors Course-MD-Data before continuing."
    )

print("COURSE_DIR:", COURSE_DIR)
print("PDB_path:", PDB_PATH)

COURSE_DIR: /home/jordivilla/Concepcion26
PDB_path: /home/jordivilla/Concepcion26/data/alanine-dipeptide.pdb

System setup

Load the alanine dipeptide structure, add hydrogens for the selected force field, and build the OpenMM system with the usual no-cutoff/HBonds setup plus a Langevin integrator so the annealing schedule has a well-defined Hamiltonian.

forcefield = ForceField("amber14-all.xml", "amber14/tip3pfb.xml")
pdb = PDBFile(str(PDB_PATH))
modeller = Modeller(pdb.topology, pdb.positions)
modeller.addHydrogens(forcefield)
system = forcefield.createSystem(
    modeller.topology,
    nonbondedMethod=app.NoCutoff,
    constraints=app.HBonds,
)
integrator = mm.LangevinMiddleIntegrator(300 * unit.kelvin, 1 / unit.picosecond, 0.004 * unit.picoseconds)
simulation = Simulation(modeller.topology, system, integrator)
simulation.context.setPositions(modeller.positions)
simulation.minimizeEnergy()
print("System initialized for {modeller.topology.getNumAtoms()} atoms.")

System initialized for {modeller.topology.getNumAtoms()} atoms.

Simulated annealing schedule

This helper drives the Langevin integrator through a discrete cooling loop. The temperature decrements by 3 K at each step and can be controlled with steps_per_temp and increments.

def run_annealing(simulation, steps_per_temp=1000, increments=100):
    integrator = simulation.integrator
    for idx in range(increments):
        temperature = 3 * (increments - idx) * unit.kelvin
        integrator.setTemperature(temperature)
        simulation.step(steps_per_temp)
        if idx % max(1, increments // 5) == 0:
            state = simulation.context.getState(getEnergy=True)
            print(f"Step {idx}: T={temperature} E={state.getPotentialEnergy()}")
    state = simulation.context.getState(getEnergy=True, getPositions=True)
    return state

Run the annealing loop

Execute the schedule with a light example (few increments) or increase steps_per_temp once you understand the behavior. The last state is retained for later analysis or saving.

result_state = run_annealing(simulation, steps_per_temp=200, increments=5)
print("Annealing complete:", result_state.getPotentialEnergy())

Step 0: T=15 K E=-89.38961791992188 kJ/mol
Step 1: T=12 K E=-89.14910888671875 kJ/mol
Step 2: T=9 K E=-88.937744140625 kJ/mol
Step 3: T=6 K E=-89.64987182617188 kJ/mol
Step 4: T=3 K E=-90.04116821289062 kJ/mol
Annealing complete: -90.04116821289062 kJ/mol