diff --git a/espaloma/data/md.py b/espaloma/data/md.py index b15c33fa..dc68a773 100644 --- a/espaloma/data/md.py +++ b/espaloma/data/md.py @@ -324,6 +324,7 @@ def __init__( self, forcefield="gaff-1.81", n_samples=100, + n_conformers=10, n_steps_per_sample=1000, temperature=TEMPERATURE, collision_rate=COLLISION_RATE, @@ -336,6 +337,7 @@ def __init__( self.collision_rate = collision_rate self.step_size = step_size self.forcefield = forcefield + self.n_conformers = n_conformers def simulation_from_graph(self, g): @@ -375,19 +377,6 @@ def simulation_from_graph(self, g): platform=openmm.Platform.getPlatformByName("Reference"), ) - import openff.toolkit - - # get conformer - g.mol.generate_conformers( - toolkit_registry=openff.toolkit.utils.RDKitToolkitWrapper(), - ) - - # put conformer in simulation - simulation.context.setPositions(g.mol.conformers[0]) - - # set velocities - simulation.context.setVelocitiesToTemperature(self.temperature) - return simulation def run(self, g, in_place=True): @@ -414,28 +403,70 @@ def run(self, g, in_place=True): # build simulation simulation = self.simulation_from_graph(g) - # initialize empty list for samples. + import openff.toolkit + + # get conformer + g.mol.generate_conformers( + toolkit_registry=openff.toolkit.utils.RDKitToolkitWrapper(), + n_conformers=self.n_conformers, + ) + + # get number of actual conformers + true_n_conformers = len(g.mol.conformers) + samples = [] + for idx in range(true_n_conformers): + # put conformer in simulation + simulation.context.setPositions(g.mol.conformers[idx]) - # minimize - simulation.minimizeEnergy() + # set velocities + simulation.context.setVelocitiesToTemperature(self.temperature) - # loop through number of samples - for _ in range(self.n_samples): + # minimize + simulation.minimizeEnergy() - # run MD for `self.n_steps_per_sample` steps - simulation.step(self.n_steps_per_sample) + # loop through number of samples + for _ in range(self.n_samples // self.n_conformers): + + # run MD for `self.n_steps_per_sample` steps + simulation.step(self.n_steps_per_sample) + + # append samples to `samples` + samples.append( + simulation.context.getState(getPositions=True) + .getPositions(asNumpy=True) + .value_in_unit(DISTANCE_UNIT) + ) - # append samples to `samples` - samples.append( - simulation.context.getState(getPositions=True) - .getPositions(asNumpy=True) - .value_in_unit(DISTANCE_UNIT) - ) + # if the `samples` array is not filled, + # pick a random conformer to do it again + if len(samples) < self.n_samples: + len_samples = len(samples) + import random + idx = random.choice(list(range(true_n_conformers))) + simulation.context.setPositions(g.mol.conformers[idx]) + + # set velocities + simulation.context.setVelocitiesToTemperature(self.temperature) + + # minimize + simulation.minimizeEnergy() + + # loop through number of samples + for _ in range(self.n_samples - len_samples): + + # run MD for `self.n_steps_per_sample` steps + simulation.step(self.n_steps_per_sample) + + # append samples to `samples` + samples.append( + simulation.context.getState(getPositions=True) + .getPositions(asNumpy=True) + .value_in_unit(DISTANCE_UNIT) + ) - # assert that energy is below zero - # final_energy = simulation.context.getState(getEnergy=True).getPotentialEnergy( - # ).value_in_unit(ENERGY_UNIT) + + assert len(samples) == self.n_samples # put samples into an array samples = np.array(samples)