def __init__(self,
explorer,
n_confinements_per_basin=250,
md_time_before_quench=Quantity(1.0, u.picoseconds),
similarity_threshold=Quantity(0.01, u.angstroms),
quench=L_BFGS,
md=Langevin):
from openexplorer import PES, KTN
self.explorer = explorer
self.quench = quench(self.explorer)
self.md = md(self.explorer)
self.n_confinements_per_basin = n_confinements_per_basin
self.md_time_before_quench = md_time_before_quench
md_timestep = self.md.get_parameters()['timestep']
self.md_steps_before_quench = int(self.md_time_before_quench /
md_timestep)
self.similarity_threshold = similarity_threshold
self.pes = PES(self.explorer.topology,
self.explorer.context.getSystem())
self.ktn = KTN(self.explorer.topology,
self.explorer.context.getSystem())
self.reset()
def getBoxVectors(self, asNumpy=False):
"""Get the periodic box vectors.
Parameters:
- asNumpy (boolean=False) if true, the values are returned as a numpy array instead of a list of Vec3s
"""
if asNumpy:
if self._numpyBoxVectors is None:
self._numpyBoxVectors = []
self._numpyBoxVectors.append(
Quantity(
numpy.array(
self.boxVectors[0].value_in_unit(nanometers)),
nanometers))
self._numpyBoxVectors.append(
Quantity(
numpy.array(
self.boxVectors[1].value_in_unit(nanometers)),
nanometers))
self._numpyBoxVectors.append(
Quantity(
numpy.array(
self.boxVectors[2].value_in_unit(nanometers)),
nanometers))
return self._numpyBoxVectors
return self.boxVectors
def getBoxVectors(self, asNumpy=False):
"""Get the periodic box vectors.
Parameters:
- asNumpy (boolean=False) if true, the values are returned as a numpy array instead of a list of Vec3s
"""
if self.boxVectors is None:
raise AttributeError('Box information not found in %s' % self.file)
if asNumpy:
if self._numpyBoxVectors is None:
self._numpyBoxVectors = []
self._numpyBoxVectors.append(
Quantity(
np.array(self.boxVectors[0].value_in_unit(nanometers)),
nanometers))
self._numpyBoxVectors.append(
Quantity(
np.array(self.boxVectors[1].value_in_unit(nanometers)),
nanometers))
self._numpyBoxVectors.append(
Quantity(
np.array(self.boxVectors[2].value_in_unit(nanometers)),
nanometers))
return self._numpyBoxVectors
return self.boxVectors
def get_graph(collection, record_name):
# get record and trajectory
record = collection.get_record(record_name, specification="default")
entry = collection.get_entry(record_name)
from openff.toolkit.topology import Molecule
mol = Molecule.from_qcschema(entry)
try:
trajectory = record.get_trajectory()
except:
return None
if trajectory is None:
return None
g = esp.Graph(mol)
# energy is already hartree
g.nodes["g"].data["u_ref"] = torch.tensor(
[
Quantity(
snapshot.properties.scf_total_energy,
esp.units.HARTREE_PER_PARTICLE,
).value_in_unit(esp.units.ENERGY_UNIT) for snapshot in trajectory
],
dtype=torch.get_default_dtype(),
)[None, :]
g.nodes["n1"].data["xyz"] = torch.tensor(
np.stack(
[
Quantity(
snapshot.get_molecule().geometry,
unit.bohr,
).value_in_unit(esp.units.DISTANCE_UNIT)
for snapshot in trajectory
],
axis=1,
),
requires_grad=True,
dtype=torch.get_default_dtype(),
)
g.nodes["n1"].data["u_ref_prime"] = torch.stack(
[
torch.tensor(
Quantity(
snapshot.dict()["return_result"],
esp.units.HARTREE_PER_PARTICLE / unit.bohr,
).value_in_unit(esp.units.FORCE_UNIT),
dtype=torch.get_default_dtype(),
) for snapshot in trajectory
],
dim=1,
)
return g
def set_parameters(self, tolerance=Quantity(1.0, u.kilojoules_per_mole),
initial_step_size=Quantity(0.01, u.angstroms)):
if not self._initialized:
self._initialize()
self._tolerance = tolerance.in_units_of(u.kilojoules_per_mole)
self._initial_step_size = initial_step_size.in_units_of(u.nanometers)
self._integrator.setGlobalVariableByName('step_size', self._initial_step_size._value)
def h5_to_dataset(df):
def get_smiles(x):
try:
return x["offmol"].to_smiles()
except:
return np.nan
df["smiles"] = df.apply(get_smiles, axis=1)
df = df.dropna()
groups = df.groupby("smiles")
gs = []
for name, group in groups:
mol_ref = group["offmol"][0]
assert all(mol_ref == entry for entry in group["offmol"])
g = esp.Graph(mol_ref)
u_ref = np.concatenate(group["energies"].values)
u_ref_prime = np.concatenate(group["gradients"].values,
axis=0).transpose(1, 0, 2)
xyz = np.concatenate(group["xyz"].values, axis=0).transpose(1, 0, 2)
assert u_ref_prime.shape[0] == xyz.shape[0] == mol_ref.n_atoms
assert u_ref.shape[0] == u_ref_prime.shape[1] == xyz.shape[1]
# energy is already hartree
g.nodes["g"].data["u_ref"] = torch.tensor(
Quantity(u_ref, esp.units.HARTREE_PER_PARTICLE).value_in_unit(
esp.units.ENERGY_UNIT),
dtype=torch.get_default_dtype(),
)[None, :]
g.nodes["n1"].data["xyz"] = torch.tensor(
Quantity(
xyz,
unit.bohr,
).value_in_unit(esp.units.DISTANCE_UNIT),
requires_grad=True,
dtype=torch.get_default_dtype(),
)
g.nodes["n1"].data["u_ref_prime"] = torch.tensor(
Quantity(
u_ref_prime,
esp.units.HARTREE_PER_PARTICLE / unit.bohr,
).value_in_unit(esp.units.FORCE_UNIT),
dtype=torch.get_default_dtype(),
)
gs.append(g)
return esp.data.dataset.GraphDataset(gs)
def periodic_box_vectors_from_xml(xmlfile):
"""Extracts periodic box vectors from OpenMM XML state file.
Box vectors are returned in the format expected by the box vector setting
function of OpenMM's topology.
"""
# parse XML file:
tree = et.parse(xmlfile)
root = tree.getroot()
# name of box vector field in XML file:
pbv = "PeriodicBoxVectors"
# box vectors need to be tuple of tuples:
box_vectors = tuple([
tuple([float(x) for x in root.find(pbv).find("A").attrib.values()]),
tuple([float(x) for x in root.find(pbv).find("B").attrib.values()]),
tuple([float(x) for x in root.find(pbv).find("C").attrib.values()])
])
# add units:
box_vectors = Quantity(box_vectors, nanometer)
# return dimensions to caller:
return (box_vectors)
def to_python(self):
"""Casts string args to ints, floats, bool..."""
for i in self:
if i.val == '':
i.val = None
elif i.name == "HR_K_PARAM": # Workaround for complex unit
i.val = Quantity(
float(i.val),
simtk.unit.kilojoule_per_mole / simtk.unit.nanometer**2)
elif i.type == str:
continue
elif i.type == int:
i.val = int(i.val)
elif i.type == float:
i.val = float(i.val)
elif i.type == bool:
if i.val.lower() in ['true', '1', 'y', 'yes']:
i.val = True
elif i.val.lower() in ['false', '0', 'n', 'no']:
i.val = False
else:
raise ValueError(f"Can't convert {i.val} into bool type.")
elif i.type == Quantity:
try:
i.val = self.parse_quantity(i.val)
except AttributeError:
raise ValueError(f"Can't parse: {i.name} = {i.val}")
else:
raise ValueError(f"Can't parse: {i.name} = {i.val}")
def get_energy(self, state):
# set the coordinates
coordinates = Quantity(state.positions, angstrom)
# set the box vectors
self._simulation.context.setPositions(coordinates)
if self._options.solvation == "explicit":
box_vector = state.box_vector / 10. # Angstrom to nm
self._simulation.context.setPeriodicBoxVectors(
[box_vector[0], 0., 0.],
[0., box_vector[1], 0.],
[0., 0., box_vector[2]],
)
# get the energy
snapshot = self._simulation.context.getState(
getPositions=True, getVelocities=True, getEnergy=True
)
snapshot = self._simulation.context.getState(getEnergy=True)
e_potential = snapshot.getPotentialEnergy()
e_potential = (
e_potential.value_in_unit(kilojoule / mole)
/ GAS_CONSTANT
/ self._temperature
)
return e_potential
def compute_energy(self, param, offset=None, platform=None):
""" Computes energy for a given structure with a given parameter set
Parameters
----------
param: parmed.charmm.CharmmParameterSet
platform: simtk.openmm.Platform to evaluate energy on (if None, will select automatically)
"""
# Save initial mm energy
save = False
if not self._have_mm_energy:
save = True
# calculate energy
super(QMDataBase, self).compute_energy(param, platform)
# Subtract off minimum of mm_energy and add offset
energy_unit = kilojoules_per_mole
min_energy = self.mm_energy.min()
self.mm_energy -= min_energy
if save:
self.initial_mm = deepcopy(self.mm_energy)
if offset:
offset = Quantity(value=offset.value, unit=energy_unit)
self.mm_energy += offset
self.delta_energy = (self.qm_energy - self.mm_energy)
def compute_energy(self, param, platform=None):
""" Computes energy for a given structure with a given parameter set
Parameters
----------
offset :
param: parmed.charmm.CharmmParameterSet
platform: simtk.openmm.Platform to evaluate energy on (if None, will select automatically)
"""
if self.n_frames == 0:
raise Exception(
"self.n_frames = 0! There are no frames to compute energy for."
)
# Check if context exists.
if not self.context:
self.create_context(param, platform)
else:
# copy new torsion parameters
self.copy_torsions(param, platform)
# Compute potential energies for all snapshots.
self.mm_energy = Quantity(value=np.zeros([self.n_frames], np.float64),
unit=kilojoules_per_mole)
for i in range(self.n_frames):
self.context.setPositions(self.positions[i])
state = self.context.getState(getEnergy=True)
self.mm_energy[i] = state.getPotentialEnergy()
def _prep_sim(self, coords, external_forces=[]):
try:
from simtk.openmm import Platform, LangevinIntegrator, Vec3
from simtk.openmm.app import Modeller, ForceField, \
CutoffNonPeriodic, PME, Simulation, HBonds
from simtk.unit import angstrom, nanometers, picosecond, \
kelvin, Quantity, molar
except ImportError:
raise ImportError(
'Please install PDBFixer and OpenMM in order to use ClustENM.')
positions = Quantity([Vec3(*xyz) for xyz in coords], angstrom)
modeller = Modeller(self._topology, positions)
if self._sol == 'imp':
forcefield = ForceField(*self._force_field)
system = forcefield.createSystem(modeller.topology,
nonbondedMethod=CutoffNonPeriodic,
nonbondedCutoff=1.0 * nanometers,
constraints=HBonds)
if self._sol == 'exp':
forcefield = ForceField(*self._force_field)
modeller.addSolvent(forcefield,
padding=self._padding * nanometers,
ionicStrength=self._ionicStrength * molar)
system = forcefield.createSystem(modeller.topology,
nonbondedMethod=PME,
nonbondedCutoff=1.0 * nanometers,
constraints=HBonds)
for force in external_forces:
system.addForce(force)
integrator = LangevinIntegrator(self._temp * kelvin, 1 / picosecond,
0.002 * picosecond)
# precision could be mixed, but single is okay.
platform = self._platform if self._platform is None else Platform.getPlatformByName(
self._platform)
properties = None
if self._platform is None:
properties = {'Precision': 'single'}
elif self._platform in ['CUDA', 'OpenCL']:
properties = {'Precision': 'single'}
simulation = Simulation(modeller.topology, system, integrator,
platform, properties)
simulation.context.setPositions(modeller.positions)
return simulation
def set_parameters(self,
tolerance=Quantity(1.0, u.kilojoules_per_mole /
u.nanometers),
max_iter=0):
self._tolerance = tolerance.in_units_of(u.kilojoules_per_mole /
u.nanometers)
self._max_iter = max_iter
self._initialize()
def set_parameters(self,
timestep=Quantity(1.0, u.femtoseconds),
tolerance=None,
alpha=0.1,
dt_max=Quantity(10.0, u.femtoseconds),
f_inc=1.1,
f_dec=0.5,
f_alpha=0.99,
N_min=5):
self._timestep = timestep.in_units_of(u.picoseconds)
self._tolerance = tolerance
self._alpha = alpha
self._dt_max = dt_max.in_units_of(u.picoseconds)
self._f_inc = f_inc
self._f_dec = f_dec
self._f_alpha = f_alpha
self._N_min = N_min
self._initialize()
def parse_quantity(self, val: str) -> Union[Quantity, None]:
if val == '':
return None
match_obj = self.quantity_regexp.match(val)
value, unit = match_obj.groups()
try:
unit = getattr(simtk.unit, unit)
except AttributeError:
raise ValueError(
f"I Can't recognise unit {unit} in expresion {val}. Example of valid quantity: 12.3 femtosecond."
)
return Quantity(value=float(value), unit=unit)
def getPositions(self, asNumpy=False):
"""Get the atomic positions.
Parameters:
- asNumpy (boolean=False) if true, the values are returned as a numpy array instead of a list of Vec3s
"""
if asNumpy:
if self._numpyPositions is None:
self._numpyPositions = Quantity(
np.array(self.positions.value_in_unit(nanometers)),
nanometers)
return self._numpyPositions
return self.positions
def set_parameters(self,
temperature=Quantity(value=298.0, unit=u.kelvin),
collision_rate=Quantity(value=1.0,
unit=1.0 / u.picosecond),
timestep=Quantity(value=2.0, unit=u.femtosecond)):
self._timestep = timestep
self._temperature = temperature
self._collision_rate = collision_rate
if self._initialized:
self._integrator.setFriction(
self._collision_rate.value_in_unit(u.picosecond**-1))
self._integrator.setTemperature(
self._temperature.value_in_unit(u.kelvin))
self._integrator.setStepSize(
self._timestep.value_in_unit(u.picoseconds))
else:
self._initialize()
def __init__(self,
positions,
topology,
structure,
torsions,
qm_energies,
angles=None,
steps=None,
directions=None,
optimized=None,
time=None):
"""Create new TorsionScanSet object"""
assert isinstance(topology, object)
super(QMDataBase, self).__init__(positions, topology, structure, time)
self.qm_energy = Quantity(value=qm_energies, unit=kilojoules_per_mole)
self.initial_mm = Quantity()
self.delta_energy = Quantity()
self.torsion_index = torsions
self.direction = directions
self.steps = steps
self.angles = angles
self.optimized = optimized
self.phis = {}
def getPositions(self, asNumpy=False, frame=0):
"""Get the atomic positions.
Parameters:
- asNumpy (boolean=False) if true, the values are returned as a numpy array instead of a list of Vec3s
- frame (int=0) the index of the frame for which to get positions
"""
if asNumpy:
if self._numpyPositions is None:
self._numpyPositions = [None]*len(self._positions)
if self._numpyPositions[frame] is None:
self._numpyPositions[frame] = Quantity(numpy.array(self._positions[frame].value_in_unit(nanometers)), nanometers)
return self._numpyPositions[frame]
return self._positions[frame]
def __init__(self,
explorer,
md_time_before_quench=Quantity(1.0, u.picoseconds),
quench=L_BFGS,
md=Langevin):
from openexplorer import PES
self.explorer = explorer
self.quench = quench(self.explorer)
self.md = md(self.explorer)
self.md.set_parameters(temperature=Quantity(500.0, u.kelvin))
self.md_time_before_quench = md_time_before_quench
md_timestep = self.md.get_parameters()['timestep']
self.md_steps_before_quench = int(self.md_time_before_quench /
md_timestep)
self.pes = PES(self.explorer.topology,
self.explorer.context.getSystem())
self.reset()
def getVelocities(self, asNumpy=False):
"""Get the atomic velocities.
Parameters
----------
asNumpy : bool=False
if true, the vectors are returned as numpy arrays instead of Vec3s
"""
if self.velocities is None:
raise AttributeError('velocities not found in %s' % self.file)
if asNumpy:
if self._numpyVelocities is None:
self._numpyVelocities = Quantity(np.array(self.velocities.value_in_unit(nanometers/picoseconds)), nanometers/picoseconds)
return self._numpyVelocities
return self.velocities
def getVelocities(self, asNumpy=False):
"""Get the atomic velocities.
Parameters:
- asNumpy (boolean=False) if true, the vectors are returned as numpy arrays instead of Vec3s
"""
if asNumpy:
if self._numpyVelocities is None:
self._numpyVelocities = Quantity(
numpy.array(
self.velocities.value_in_unit(nanometers /
picoseconds)),
nanometers / picoseconds)
return self._numpyVelocities
return self.velocities
class L_BFGS():
_explorer = None
_initialized = False
_tolerance = Quantity(1.0, u.kilojoules_per_mole / u.nanometers)
_max_iter = 0
def __init__(self, explorer):
self._explorer = explorer
def _initialize(self):
self._initialized = True
def set_parameters(self,
tolerance=Quantity(1.0, u.kilojoules_per_mole /
u.nanometers),
max_iter=0):
self._tolerance = tolerance.in_units_of(u.kilojoules_per_mole /
u.nanometers)
self._max_iter = max_iter
self._initialize()
def replicate_parameters(self, explorer):
tolerance = explorer.quench.l_bfgs._tolerance
max_iter = explorer.quench.l_bfgs._max_iter
self.set_parameters(tolerance, max_iter)
def run(self):
if not self._initialized:
self._initialize()
LocalEnergyMinimizer.minimize(self._explorer.context, self._tolerance,
self._max_iter)
def __call__(self, *args, **kwargs):
return self.run(*args, **kwargs)
def baseline_energy(self, g, suffix=None):
if suffix is None:
suffix = "_" + self.forcefield
from openmmforcefields.generators import SystemGenerator
# define a system generator
system_generator = SystemGenerator(
small_molecule_forcefield=self.forcefield, )
mol = g.mol
# mol.assign_partial_charges("formal_charge")
# create system
system = system_generator.create_system(
topology=mol.to_topology().to_openmm(),
molecules=mol,
)
# parameterize topology
topology = g.mol.to_topology().to_openmm()
integrator = openmm.LangevinIntegrator(TEMPERATURE, COLLISION_RATE,
STEP_SIZE)
# create simulation
simulation = Simulation(topology=topology,
system=system,
integrator=integrator)
us = []
xs = (Quantity(
g.nodes["n1"].data["xyz"].detach().numpy(),
esp.units.DISTANCE_UNIT,
).value_in_unit(unit.nanometer).transpose((1, 0, 2)))
for x in xs:
simulation.context.setPositions(x)
us.append(
simulation.context.getState(
getEnergy=True).getPotentialEnergy().value_in_unit(
esp.units.ENERGY_UNIT))
g.nodes["g"].data["u%s" % suffix] = torch.tensor(us)[None, :]
return g
def deserialize_box_vectors(xmlInput, is_file=True):
"""
Takes an XML string or file and converts to a 3x3 simtk.unit.Quantity for
representing parmed or OpenMM box vectors.
Parameters
----------
xmlInput : str, Required
The name of the file to read for XML, or, if is_file is False, then
read the xmlInput string itself as the XML
is_file : bool, Optional, default: True
If reading the XML to a file is desired, then enter a valid file path
and name. If is_file is false, then the xmlInput string is read as
the XML itself.
Returns
-------
result : simtk.unit.Quantity
The box vectors in a 3x3 simtk.unit.Quantity object for easy input to
parmed or OpenMM.
"""
if is_file:
tree = ET.parse(xmlInput)
xmlBox_vectors = tree.getroot()
else:
xmlBox_vectors = ET.fromstring(xmlInput)
assert xmlBox_vectors.text is not None
xmlA = xmlBox_vectors.find('A')
xmlAx = float(xmlA.find('x').text)
xmlAy = float(xmlA.find('y').text)
xmlAz = float(xmlA.find('z').text)
xmlB = xmlBox_vectors.find('B')
xmlBx = float(xmlB.find('x').text)
xmlBy = float(xmlB.find('y').text)
xmlBz = float(xmlB.find('z').text)
xmlC = xmlBox_vectors.find('C')
xmlCx = float(xmlC.find('x').text)
xmlCy = float(xmlC.find('y').text)
xmlCz = float(xmlC.find('z').text)
box_vectors = Quantity(
[[xmlAx, xmlAy, xmlAz], [xmlBx, xmlBy, xmlBz], [xmlCx, xmlCy, xmlCz]],
unit=nanometer)
return box_vectors
def __init__(self, positions, topology, structure, time=None):
"""Create new TorsionScanSet object"""
assert isinstance(topology, object)
super(DataBase, self).__init__(positions, topology, time)
self.structure = structure
self.mm_energy = Quantity()
self.positions = positions
self.context = None
self.system = None
self.integrator = mm.VerletIntegrator(0.004 * picoseconds)
# Don't allow an empty TorsionScanSet to be created
if self.n_frames == 0:
msg = 'DataBase has no frames!\n'
msg += '\n'
msg += 'DataBase provided were:\n'
msg += str(positions)
raise Exception(msg)
def collect_minimum(self,
explorer,
coordinates=None,
potential_energy=None,
similarity_criterion='least_rmsd',
similarity_threshold=Quantity(0.01, u.angstroms)):
new_minimum = True
inherent_structure_index = None
if similarity_criterion == 'least_rmsd':
similarity = explorer.distance.least_rmsd
else:
raise NotImplementedError
for minimum_index in reversed(range(self.n_minima)):
if similarity(self.minima[minimum_index]) < similarity_threshold:
new_minimum = False
inherent_structure_index = minimum_index
break
if new_minimum:
if coordinates is None:
coordinates = explorer.get_coordinates()
if potential_energy is None:
potential_energy = explorer.get_potential_energy()
self.minima.append(coordinates)
self.potential_energy_minima.append(potential_energy)
inherent_structure_index = self.n_minima
self.basins_network.add_node(inherent_structure_index)
self.n_minima += 1
if self.global_minimum_potential_energy > potential_energy:
self.global_minimum_index = inherent_structure_index
self.global_minimum_potential_energy = potential_energy
return inherent_structure_index
def _assign_openmm_positions(self, configuration):
from simtk.unit import Quantity
positions = Quantity(value=configuration.reshape(-1, _SPATIAL_DIM),
unit=self._length_scale)
self._openmm_context.setPositions(positions)
def _run(self, state, minimize):
assert abs(state.alpha - self._alpha) < 1e-6 # run Monte Carlo
if minimize:
state = self._run_min_mc(state)
else:
state = self._run_mc(state)
# add units to coordinates and velocities (we store in Angstrom, openmm
# uses nm
coordinates = Quantity(state.positions, angstrom)
velocities = Quantity(state.velocities, angstrom / picosecond)
box_vectors = Quantity(state.box_vector, angstrom)
# set the positions
self._simulation.context.setPositions(coordinates)
# if explicit solvent, then set the box vectors
if self._options.solvation == "explicit":
self._simulation.context.setPeriodicBoxVectors(
[box_vectors[0].value_in_unit(nanometer), 0.0, 0.0],
[0.0, box_vectors[1].value_in_unit(nanometer), 0.0],
[0.0, 0.0, box_vectors[2].value_in_unit(nanometer)],
)
# run energy minimization
if minimize:
self._simulation.minimizeEnergy(maxIterations=self._options.minimize_steps)
# set the velocities
self._simulation.context.setVelocities(velocities)
# run timesteps
self._simulation.step(self._options.timesteps)
# extract coords, vels, energy and strip units
if self._options.solvation == "implicit":
snapshot = self._simulation.context.getState(
getPositions=True, getVelocities=True, getEnergy=True
)
elif self._options.solvation == "explicit":
snapshot = self._simulation.context.getState(
getPositions=True,
getVelocities=True,
getEnergy=True,
enforcePeriodicBox=True,
)
coordinates = snapshot.getPositions(asNumpy=True).value_in_unit(angstrom)
velocities = snapshot.getVelocities(asNumpy=True).value_in_unit(
angstrom / picosecond
)
_check_for_nan(coordinates, velocities, self._rank)
# if explicit solvent, the recover the box vectors
if self._options.solvation == "explicit":
box_vector = snapshot.getPeriodicBoxVectors().value_in_unit(angstrom)
box_vector = np.array(
(box_vector[0][0], box_vector[1][1], box_vector[2][2])
)
# just store zeros for implicit solvent
else:
box_vector = np.zeros(3)
# get the energy
e_potential = (
snapshot.getPotentialEnergy().value_in_unit(kilojoule / mole)
/ GAS_CONSTANT
/ self._temperature
)
# store in state
state.positions = coordinates
state.velocities = velocities
state.energy = e_potential
state.box_vector = box_vector
return state
class Langevin():
_explorer = None
_initialized = False
_context = None
_integrator = None
_timestep = Quantity(value=2.0, unit=u.femtosecond)
_temperature = Quantity(value=298.0, unit=u.kelvin)
_collision_rate = Quantity(value=1.0, unit=1.0 / u.picosecond)
def __init__(self, explorer):
self._explorer = explorer
def _initialize(self):
system = self._explorer.context.getSystem()
platform = self._explorer.context.getPlatform()
properties = {}
if platform.getName() == 'CUDA':
properties['CudaPrecision'] = 'mixed'
self._integrator = LangevinIntegrator(self._temperature,
self._collision_rate,
self._timestep)
self._context = Context(system, self._integrator, platform, properties)
self._initialized = True
def set_parameters(self,
temperature=Quantity(value=298.0, unit=u.kelvin),
collision_rate=Quantity(value=1.0,
unit=1.0 / u.picosecond),
timestep=Quantity(value=2.0, unit=u.femtosecond)):
self._timestep = timestep
self._temperature = temperature
self._collision_rate = collision_rate
if self._initialized:
self._integrator.setFriction(
self._collision_rate.value_in_unit(u.picosecond**-1))
self._integrator.setTemperature(
self._temperature.value_in_unit(u.kelvin))
self._integrator.setStepSize(
self._timestep.value_in_unit(u.picoseconds))
else:
self._initialize()
def get_parameters(self):
parameters = {
'timestep': self._timestep,
'temperature': self._temperature,
'collision_rate': self._collision_rate
}
return parameters
def replicate_parameters(self, explorer):
timestep = explorer.md.langevin._timestep
temperature = explorer.md.langevin._temperature
collision_rate = explorer.md.langevin._collision_rate
self.set_paramters(temperature, collision_rate, timestep)
def _set_coordinates(self, coordinates):
self._context.setPositions(coordinates)
def _get_coordinates(self):
return self._context.getState(getPositions=True).getPositions(
asNumpy=True)
def _set_velocities(self, velocities):
self._context.setVelocities(velocities)
def _get_velocities(self):
return self._context.getState(getVelocities=True).getVelocities(
asNumpy=True)
def _coordinates_to_explorer(self):
self._explorer.set_coordinates(self._get_coordinates())
def _coordinates_from_explorer(self):
self._set_coordinates(self._explorer.get_coordinates())
def _velocities_to_explorer(self):
self._explorer.set_velocities(self._get_velocities())
def _velocities_from_explorer(self):
self._set_velocities(self._explorer.get_velocities())
def get_time(self):
return self._context.getState().getTime()
def run(self, steps=0):
if not self._initialized:
self._initialize()
self._coordinates_from_explorer()
self._velocities_from_explorer()
self._integrator.step(steps)
self._coordinates_to_explorer()
self._velocities_to_explorer()
def __call__(self, *args, **kwargs):
return self.run(*args, **kwargs)
