Examples

Chiron can be used to perform Langevin dynamics simulations of systems, use different MCMC Moves to target the Boltzmann distribution more efficiently and use different enhanced sampling methods to converge faster. In the following examples are given to outline its functionality.

Langevin dynamics of a test system in vacuum

In this example, a Lennard-Jones fluid is simulated, and the trajectory is reported. We start by importing the potential (this is a custom potential using the familiar 12-6 Lennard-Jones potential and the LennardJonesFluid testsystem from openmmtools.

from chiron.potential import LJPotential
from openmmtools.testsystems import LennardJonesFluid

lj_fluid = LennardJonesFluid(reduced_density=0.8, n_particles=100)
lj_potential = LJPotential(lj_fluid.topology)

We need to initialize the pseudo-random number generator. Next, we import the ThermodynamicState and SamplerState, which define the probability density of the MCMC chain and the current state of the sampler.

# set the seed of the pseudo-random number generator
from chiron.utils import PRNG
PRNG.set_seed(1234)

# set up the thermodynamic and sampler state. the thermodynamic state contains 
# all information about the probability density that is targeted, and the sampler state defines 
# the state of the sampler (initially, only the starting positions and a seed to initialize a 
# new random number stream)
  
from chiron.states import SamplerState, ThermodynamicState

thermodynamic_state = ThermodynamicState(
        potential=potential, temperature=300 * unit.kelvin
    )

sampler_state = SamplerState(testsystem.positions, PRNG.get_random_key())

In the next code block, the Langevin Integrator and its reporter are set up (which will automatically write properties of interest). The reporter will automatically write in the directory specified in the BaseReporter.set_directory. The LangevinIntegrator consumes an instance of the SamplerState and ThermodynamicState and returns a new SamplerState.

from chiron.reporters import LangevinDynamicsReporter
from chiron.reporters import BaseReporter

# set up reporter directory
BaseReporter.set_directory(prep_temp_dir.join("langevin_dymamics"))
reporter = LangevinDynamicsReporter()

report_frequency = 100 # report every nth integration step
n_steps = 10_000 # simulation for nth steps (with timestep=1 fs)

integrator = LangevinIntegrator(reporter=reporter, report_frequency=report_frequency)
sampler_state = integrator.run(
        sampler_state,
        thermodynamic_state,
        n_steps=n_steps,
    )

Free energy difference of harmonic oscillators

The free energy difference between 4 harmonic oscillators with different force constants is calculated. We use the MultistateSampler (a simplified re-implementation of the same class in the openmmtools library).

The MultistateSampler uses a MCMCSampler and a list of ThermodynamicStates (defining the different harmonic oscillators) and SamplerStates. The MultistateSampler propagates a markov chain for each of the harmonic oscillators according to the MCMCMoves defined in MCMCSampler (in the following, we use LangevinDynamicsMove). After each iteration, the coordinates are saved, and the energy matrix (n_replica, n_states, n_iterations) is calculated (and used for the free energy estimate using mbar).

# setting up the MCMC moves, sampler and scheduler
from chiron.mcmc import LangevinDynamicsMove
from chiron.mcmc import MCMCSampler, MoveSchedule

lang_move = LangevinDynamicsMove(stepsize=1.0 * unit.femtoseconds, nr_of_steps=100)
move_schedule = MoveSchedule([("LangevinDynamicsMove", lang_move)])
mcmc_sampler = MCMCSampler(
    move_schedule,
)

# setting up the reporter and directory where to save trajectories and log files
from chiron.reporters import MultistateReporter, BaseReporter
BaseReporter.set_directory("multistate_test")
reporter = MultistateReporter()
reporter.reset_reporter_file() # reset log files

# set up the multistate sampler
multistate_sampler = MultiStateSampler(mcmc_sampler=mcmc_sampler, reporter=reporter)

# importing the hamronic oscillator system from openmmtools
from openmmtools.testsystems import HarmonicOscillator
ho = HarmonicOscillator()

# the free energy between 4 harmonic oscillators with different force constants 
# is calculated (and, therefore, different thermodynamic states)
from chiron.states import ThermodynamicState, SamplerState
from chiron.potential import HarmonicOscillatorPotential

# calculate the dimensionless force constant
n_states = 4
T = 300.0 * unit.kelvin  # Minimum temperature.
kT = unit.BOLTZMANN_CONSTANT_kB * T * unit.AVOGADRO_CONSTANT_NA

sigmas = [
    unit.Quantity(2.0 + 0.2 * state_index, unit.angstrom)
    for state_index in range(n_states)
]
Ks = [kT / sigma**2 for sigma in sigmas]

thermodynamic_states = [
    ThermodynamicState(HarmonicOscillatorPotential(ho.topology, k=k), temperature=T)
    for k in Ks
]

print(f"Initialize harmonic oscillator with {n_states} states and ks {Ks}")
# set the pseudo-random number stream
from chiron.utils import PRNG

PRNG.set_seed(1234)
# initialize 4 sampler states, each with their pseudo-random number stream
sampler_state = [
    SamplerState(ho.positions, current_PRNG_key=PRNG.get_random_key())
    for _ in sigmas
]

# The free energy of the harmonic oscillator can be calculated analytically 
import numpy as np

f_i = np.array(
    [
        -np.log(2 * np.pi * (sigma / unit.angstroms) ** 2) * (3.0 / 2.0)
        for sigma in sigmas
    ]
)

multistate_sampler.create(
    thermodynamic_states=thermodynamic_states,
    sampler_states=sampler_state,
    nbr_list=None, # harmonic oscillators don't interact with each other
)
    
print(f'True free energy difference: f_i: {f_i - f_i[:, np.newaxis]}')
# which will return [ 0.        , -0.28593054, -0.54696467, -0.78709279]    

# set up the MCMC sampling (using Langevin dynamics)
n_iteratinos = 100 # le't use 100 MCMC moves
ho_sampler.run(n_iteratinos)

# print the resulting free energies
print(ho_sampler.analytical_f_i)