human
What is a the Python code you requested given the following:
I am working on the file "Annealed importance sampling.ipynb".
The first task is:
5. A more interesting example
Let's sample a biomolecule's configuration space in this way, and maybe estimate its partition function.
For now, let's do alanine dipeptide.
5.1-5.2. Define annealing distributions and transition kernels
Can you write Python code for it?
assistant
#Here is the Python code you requested.

from simtk.openmm.app import *
from simtk.openmm import *
from simtk.unit import *
from openmmtools.integrators import MetropolisMonteCarloIntegrator,HMCIntegrator

# all I want is the alanine dipeptide topology
from msmbuilder.example_datasets import AlanineDipeptide
ala = AlanineDipeptide().get().trajectories

top_md = ala[0][0].topology
topology = top_md.to_openmm()
n_atoms = top_md.n_atoms
dim = n_atoms*3

# create an openmm system
forcefield = ForceField('amber99sb.xml','amber10_obc.xml')
system = forcefield.createSystem(topology,
    nonbondedCutoff=1*nanometer, constraints=HBonds)
integrator = HMCIntegrator(300*kelvin)
simulation = Simulation(topology, system, integrator)

# create a thin wrapper class
class PeptideSystem():
    def __init__(self):
        self.simulation = simulation
        self.positions = self.simulation.context.getState(getPositions=True).getPositions().value_in_unit(nanometer)
        self.n_atoms = len(self.positions)
    
    def evaluate_potential_flat(self,position_vec):
        positions = position_vec.reshape(self.n_atoms,3)
        self.simulation.context.setPositions(positions)
        return self.simulation.context.getState(getEnergy=True).getPotentialEnergy()
    
    def propagate(self,position_vec,n_steps=1000,temp=300):
        integrator = HMCIntegrator(temp*kelvin)
        self.simulation = Simulation(topology,system,integrator)
        positions = position_vec.reshape(self.n_atoms,3)
        self.simulation.context.setPositions(positions)
        simulation.step(n_steps)
        
        return np.array(self.simulation.context.getState(getPositions=True).getPositions().value_in_unit(nanometer)).flatten()
    
    def probability_at_temp(self,position_vec,temp=300.0):
        return np.exp(-self.evaluate_potential_flat(position_vec).value_in_unit(kilojoule/mole)/temp)
    

peptide = PeptideSystem()
temperatures = np.logspace(3,0,1000)*300

# for some reason I can't create a bunch of parametrized anonymous functions in a list comprehension?
# i.e. [lambda...