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

# 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 oscillator with different force constant is calculated (and, therefore, # different ThermodynamicState
from chiron.states import ThermodynamicState, SamplerState
from chiron.potential import HarmonicOscillatorPotential

# calculate the dimensionless foce 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 own 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'Analytic f_i: {f_i - f_i[:, np.newaxis]}')
# which will return [ 0.        , -0.28593054, -0.54696467, -0.78709279]    

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

    # check that we have the correct number of iterations, replicas and states
    assert ho_sampler.iteration == n_iteratinos
    assert ho_sampler._iteration == n_iteratinos
    assert ho_sampler.n_replicas == 4
    assert ho_sampler.n_states == 4

    u_kn = ho_sampler._reporter.get_property("u_kn")
    assert u_kn.shape == (n_iteratinos, 4, 4)
    # check that the free energies are correct
    print(ho_sampler.analytical_f_i)