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 __init__(
self,
sigma: unit.Quantity,
point: unit.Quantity,
radius: unit.Quantity,
atom_idx: int,
active_at: int = -1,
):
"""
Flat well restraint that becomes active when atom moves outside of radius.
Parameters
----------
sigma : float, unit'd
point : np.array, unit'd
radius : float, unit'd
atom_idx : list
list of atoms idxs
active_at : int
Integer to indicccate at which state the restraint is fully active. Either 0 (for
lambda 0), or 1 (for lambda 1) or -1 (always active)
"""
assert type(sigma) == unit.Quantity
assert type(point) == unit.Quantity
super().__init__(sigma, point.value_in_unit(unit.angstrom), active_at)
self.atom_idx = atom_idx
self.cutoff_radius = (radius.value_in_unit(unit.angstrom)
) # slightly decrease radius
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 __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 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 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 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 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 compute_energy(self, param, offset, platform=None):
""" Computes energy for a given structure with a given parameter set
Parameters
----------
param: chemistry.charmm.CharmmParameterSet
platform: simtk.openmm.Platform to evaluate energy on (if None, will select automatically)
"""
# Create Context.
integrator = mm.VerletIntegrator(0.004 * u.picoseconds)
system = self.structure.createSystem(param)
if platform != None:
context = mm.Context(system, integrator, platform)
else:
context = mm.Context(system, integrator)
# 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):
context.setPositions(self.openmm_positions(i))
state = context.getState(getEnergy=True)
self.mm_energy[i] = state.getPotentialEnergy()
# Subtract off minimum of mm_energy
self.mm_energy -= self.mm_energy.min() + Quantity(value=float(offset.value), unit=kilojoules_per_mole)
self.delta_energy = self.qm_energy - self.mm_energy
# Compute deviation between MM and QM energies with offset
# self.delta_energy = mm_energy - self.qm_energy + Quantity(value=offset, unit=kilojoule_per_mole)
# Clean up.
del context
del system
del integrator
def compute_energy(self, param, offset, 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)
"""
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()
# 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()
# Subtract off minimum of mm_energy and add offset
min_energy = self.mm_energy.min()
self.mm_energy -= min_energy
self.mm_energy += offset
self.delta_energy = (self.qm_energy - self.mm_energy)
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 calculate_energy(
self,
coordinate_list: unit.Quantity,
lambda_value: float = 0.0,
original_neural_network: bool = True,
requires_grad_wrt_coordinates: bool = True,
requires_grad_wrt_parameters: bool = True,
include_restraint_energy_contribution: bool = True,
):
"""
Given a coordinate set (x) the energy is calculated in kJ/mol.
Parameters
----------
x : list, [N][K][3] unit'd (distance unit)
initial configuration
lambda_value : float
between 0.0 and 1.0 - at zero contributions of alchemical atoms are zero
Returns
-------
NamedTuple
"""
assert type(coordinate_list) == unit.Quantity
assert 0.0 <= float(lambda_value) <= 1.0
logger.debug(
f"Including restraints: {include_restraint_energy_contribution}")
logger.debug(f"Batch-size: {len(coordinate_list)}")
coordinates = torch.tensor(
coordinate_list.value_in_unit(unit.nanometer),
requires_grad=requires_grad_wrt_coordinates,
device=self.device,
dtype=torch.float32,
)
logger.debug(f"coordinates tensor: {coordinates.size()}")
energy_in_kT, restraint_energy_contribution_in_kT = self._calculate_energy(
coordinates,
lambda_value,
original_neural_network,
include_restraint_energy_contribution,
)
energy = np.array([e.item() for e in energy_in_kT]) * kT
restraint_energy_contribution = (
np.array([e.item()
for e in restraint_energy_contribution_in_kT]) * kT)
if requires_grad_wrt_parameters:
return DecomposedEnergy(energy, restraint_energy_contribution,
energy_in_kT)
else:
return DecomposedEnergy(energy, restraint_energy_contribution,
energy_in_kT.detach())
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 __init__(self, positions, topology, structure, torsions, directions, steps, qm_energies):
"""Create new TorsionScanSet object"""
assert isinstance(topology, object)
super(TorsionScanSet, self).__init__(positions, topology)
self.structure = structure
self.qm_energy = Quantity(value=qm_energies, unit=kilojoules_per_mole)
self.mm_energy = Quantity()
self.delta_energy = Quantity()
self.torsion_index = torsions
self.direction = directions
self.steps = steps
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 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 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 __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 sample_velocities(masses: Quantity, temperature: Quantity) -> np.array:
"""Sample Maxwell-Boltzmann velocities ~ N(0, sqrt(kB T / m)"""
n_particles = len(masses)
spatial_dim = 3
v_unscaled = np.random.randn(n_particles, spatial_dim)
# intended to be consistent with timemachine.integrator:langevin_coefficients
sigma = np.sqrt(BOLTZ * temperature.value_in_unit(kelvin)) * np.sqrt(
1 / masses)
v_scaled = v_unscaled * np.expand_dims(sigma, axis=1)
assert v_scaled.shape == (n_particles, spatial_dim)
return v_scaled
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
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 _from_omm_quantity(val: simtk_unit.Quantity):
"""Helper function to convert float or array quantities tagged with SimTK/OpenMM units to
a Pint-compatible quantity"""
unit_ = val.unit
val_ = val.value_in_unit(unit_)
if type(val_) in {float, int}:
unit_ = val.unit
return val_ * unit.Unit(str(unit_))
elif type(val_) in {tuple, list, np.ndarray}:
array = np.asarray(val_)
return array * unit.Unit(str(unit_))
else:
raise UnitValidationError(
"Found a simtk.unit.Unit wrapped around something other than a float-like "
f"or np.ndarray-like. Found a unit wrapped around type {type(val_)}."
)
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
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 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 __init__(self, positions, topology, structure, torsions, directions, steps, qm_energies):
"""Create new TorsionScanSet object"""
assert isinstance(topology, object)
super(TorsionScanSet, self).__init__(positions, topology)
self.structure = structure
self.qm_energy = Quantity(value=qm_energies, unit=kilojoules_per_mole)
self.mm_energy = Quantity()
self.delta_energy = Quantity()
self.torsion_index = torsions
self.direction = directions
self.steps = steps
self.positions = positions
self.context = None
self.system = None
self.integrator = mm.VerletIntegrator(0.004*u.picoseconds)
self.energy = np.array
# Don't allow an empty TorsionScanSet to be created
if self.n_frames == 0:
msg = 'TorsionScanSet has no frames!\n'
msg += '\n'
msg += 'positions provided were:\n'
msg += str(positions)
raise Exception(msg)
class TorsionScanSet(Trajectory):
"""container object for torsion scan
A TorsionScanSet should be constructed by loading Gaussian 09 torsion scan log files from disk
with an mdtraj.Topology object
Examples
--------
>>> torsion_set = read_scan_logfile('../examples/data/pyrrole/torsion-scan/PRL.scan2.neg.log', '../examples/data/pyrrole/pyrrol.psf')
>>> print(torsion_set)
<torsions.TorsionScanSet with 40 frames, 22 atoms, 1 residues, without MM Energy>
Attributes
----------
structure: ParmEd.Structure
qm_energy: simtk.unit.Quantity((n_frames), unit=kilojoule/mole)
mm_energy: simtk.unit.Quantity((n_frames), unit=kilojoule/mole)
delta_energy: simtk.unit.Quantity((n_frames), unit=kilojoule/mole)
torsion_index: {np.ndarray, shape(n_frames, 4)}
step: {np.ndarray, shape(n_frame, 3)}
direction: {np.ndarray, shape(n_frame)}. 0 = negative, 1 = positive
"""
def __init__(self, positions, topology, structure, torsions, directions, steps, qm_energies):
"""Create new TorsionScanSet object"""
assert isinstance(topology, object)
super(TorsionScanSet, self).__init__(positions, topology)
self.structure = structure
self.qm_energy = Quantity(value=qm_energies, unit=kilojoules_per_mole)
self.mm_energy = Quantity()
self.delta_energy = Quantity()
self.torsion_index = torsions
self.direction = directions
self.steps = steps
self.positions = positions
self.context = None
self.system = None
self.integrator = mm.VerletIntegrator(0.004*u.picoseconds)
self.energy = np.array
# Don't allow an empty TorsionScanSet to be created
if self.n_frames == 0:
msg = 'TorsionScanSet has no frames!\n'
msg += '\n'
msg += 'positions provided were:\n'
msg += str(positions)
raise Exception(msg)
def create_context(self, param, platform=None):
self.structure.load_parameters(param, copy_parameters=False)
self.system = self.structure.createSystem()
if platform != None:
self.context = mm.Context(self.system, self.integrator, platform)
else:
self.context = mm.Context(self.system, self.integrator)
def copy_torsions(self):
forces = {self.system.getForce(i).__class__.__name__: self.system.getForce(i)
for i in range(self.system.getNumForces())}
torsion_force = forces['PeriodicTorsionForce']
# create new force
new_torsion_force = self.structure.omm_dihedral_force()
# copy parameters
for i in range(new_torsion_force.getNumTorsions()):
torsion = new_torsion_force.getTorsionParameters(i)
torsion_force.setTorsionParameters(i, *torsion)
# update parameters in context
torsion_force.updateParametersInContext(self.context)
#clean up
del new_torsion_force
def to_dataframe(self):
""" convert TorsionScanSet to pandas dataframe """
data = []
for i in range(self.n_frames):
if len(self.mm_energy) == self.n_frames and len(self.delta_energy) == self.n_frames:
data.append((self.torsion_index[i], self.direction[i], self.steps[i], self.qm_energy[i], self.mm_energy[i],
self.delta_energy[i]))
else:
data.append((self.torsion_index[i], self.direction[i], self.steps[i], self.qm_energy[i], float('nan'), float('nan')))
torsion_set = pd.DataFrame(data, columns=[ "torsion", "scan_direction", "step_point_total", "QM_energy KJ/mol",
"MM_energy KJ/mole", "delta KJ/mole"])
return torsion_set
def _string_summary_basic(self):
"""Basic summary of TorsionScanSet in string form."""
energy_str = 'with MM Energy' if self._have_mm_energy else 'without MM Energy'
value = "torsions.TorsionScanSet with %d frames, %d atoms, %d residues, %s" % (
self.n_frames, self.n_atoms, self.n_residues, energy_str)
return value
def extract_geom_opt(self):
key = []
for i, step in enumerate(self.steps):
try:
if step[1] != self.steps[i+1][1]:
key.append(i)
except IndexError:
key.append(i)
new_torsionScanSet = self.slice(key)
return new_torsionScanSet
def compute_energy(self, param, offset, 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)
"""
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()
# 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()
# Subtract off minimum of mm_energy and add offset
min_energy = self.mm_energy.min()
self.mm_energy -= min_energy
self.mm_energy += offset
self.delta_energy = (self.qm_energy - self.mm_energy)
# Compute deviation between MM and QM energies with offset
#self.delta_energy = mm_energy - self.qm_energy + Quantity(value=offset, unit=kilojoule_per_mole)
@property
def _have_mm_energy(self):
return len(self.mm_energy) is not 0
# @property
# def _unique_torsions(self):
# Not returning the right amount. debug
# torsions = []
# for i in range(len(self.torsion_index)):
# try:
# if (self.torsion_index[i] != self.torsion_index[i+1]).all():
# torsions.append(self.torsion_index[i]), torsions.append(self.torsion_index[i+1])
# except:
# pass
# return len(torsions), torsions
def __getitem__(self, key):
"Get a slice of this trajectory"
return self.slice(key)
def slice(self, key, copy=True):
"""Slice trajectory, by extracting one or more frames into a separate object
This method can also be called using index bracket notation, i.e
`traj[1] == traj.slice(1)`
Parameters
----------
key : {int, np.ndarray, slice}
The slice to take. Can be either an int, a list of ints, or a slice
object.
copy : bool, default=True
Copy the arrays after slicing. If you set this to false, then if
you modify a slice, you'll modify the original array since they
point to the same data.
"""
xyz = self.xyz[key]
time = self.time[key]
torsions = self.torsion_index[key]
direction = self.direction[key]
steps = self.steps[key]
qm_energy = self.qm_energy[key]
unitcell_lengths, unitcell_angles = None, None
if self.unitcell_angles is not None:
unitcell_angles = self.unitcell_angles[key]
if self.unitcell_lengths is not None:
unitcell_lengths = self.unitcell_lengths[key]
if copy:
xyz = xyz.copy()
time = time.copy()
topology = deepcopy(self._topology)
structure = deepcopy(self.structure)
torsions = torsions.copy()
direction = direction.copy()
steps = steps.copy()
qm_energy = qm_energy.copy()
if self.unitcell_angles is not None:
unitcell_angles = unitcell_angles.copy()
if self.unitcell_lengths is not None:
unitcell_lengths = unitcell_lengths.copy()
newtraj = self.__class__(
xyz, topology, structure, torsions, direction, steps, qm_energy)
if self._rmsd_traces is not None:
newtraj._rmsd_traces = np.array(self._rmsd_traces[key],
ndmin=1, copy=True)
return newtraj
class TorsionScanSet(Trajectory):
"""container object for torsion scan
A TorsionScanSet should be constructed by loading Gaussian 09 torsion scan log files from disk
with an mdtraj.Topology object
Examples
--------
>>> torsion_set = read_scan_logfile('FRG.scanN.dir.log')
>>> print torsion_set
<torsions.TorsionScanSet with 346 frames, 22 atoms, 1 residues, 4 unique torsions without MM Energy at 0x10b099b10>
Attributes
----------
structure: chemistry.Structure
qm_energy: simtk.unit.Quantity((n_frames), unit=kilojoule/mole)
mm_energy: simtk.unit.Quantity((n_frames), unit=kilojoule/mole)
delta_energy: simtk.unit.Quantity((n_frames), unit=kilojoule/mole)
torsion_index: {np.ndarray, shape(n_frames, 4)}
step: {np.ndarray, shape(n_frame, 3)}
direction: {np.ndarray, shape(n_frame)}. 0 = negative, 1 = positive
"""
def __init__(self, positions, topology, structure, torsions, directions, steps, qm_energies):
"""Create new TorsionScanSet object"""
assert isinstance(topology, object)
super(TorsionScanSet, self).__init__(positions, topology)
self.structure = structure
self.qm_energy = Quantity(value=qm_energies, unit=kilojoules_per_mole)
self.mm_energy = Quantity()
self.delta_energy = Quantity()
self.torsion_index = torsions
self.direction = directions
self.steps = steps
def to_dataframe(self):
""" convert TorsionScanSet to pandas dataframe """
data = []
for i in range(self.n_frames):
if len(self.mm_energy) == self.n_frames and len(self.delta_energy) == self.n_frames:
data.append(
(
self.torsion_index[i],
self.direction[i],
self.steps[i],
self.qm_energy[i],
self.mm_energy[i],
self.delta_energy[i],
)
)
else:
data.append(
(
self.torsion_index[i],
self.direction[i],
self.steps[i],
self.qm_energy[i],
float("nan"),
float("nan"),
)
)
torsion_set = pd.DataFrame(
data,
columns=[
"torsion",
"scan_direction",
"step_point_total",
"QM_energy KJ/mol",
"MM_energy KJ/mole",
"delta KJ/mole",
],
)
return torsion_set
def _string_summary_basic(self):
"""Basic summary of TorsionScanSet in string form."""
energy_str = "with MM Energy" if self._have_mm_energy else "without MM Energy"
value = "torsions.TorsionScanSet with %d frames, %d atoms, %d residues, %s" % (
self.n_frames,
self.n_atoms,
self.n_residues,
energy_str,
)
return value
def extract_geom_opt(self):
key = []
for i, step in enumerate(self.steps):
try:
if step[1] != self.steps[i + 1][1]:
key.append(i)
except IndexError:
key.append(i)
new_torsionScanSet = self.slice(key)
return new_torsionScanSet
def compute_energy(self, param, offset, platform=None):
""" Computes energy for a given structure with a given parameter set
Parameters
----------
param: chemistry.charmm.CharmmParameterSet
platform: simtk.openmm.Platform to evaluate energy on (if None, will select automatically)
"""
# Create Context.
integrator = mm.VerletIntegrator(0.004 * u.picoseconds)
system = self.structure.createSystem(param)
if platform != None:
context = mm.Context(system, integrator, platform)
else:
context = mm.Context(system, integrator)
# 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):
context.setPositions(self.openmm_positions(i))
state = context.getState(getEnergy=True)
self.mm_energy[i] = state.getPotentialEnergy()
# Subtract off minimum of mm_energy
self.mm_energy -= self.mm_energy.min() + Quantity(value=float(offset.value), unit=kilojoules_per_mole)
self.delta_energy = self.qm_energy - self.mm_energy
# Compute deviation between MM and QM energies with offset
# self.delta_energy = mm_energy - self.qm_energy + Quantity(value=offset, unit=kilojoule_per_mole)
# Clean up.
del context
del system
del integrator
# print('Heap at end of compute_energy'), hp.heeap()
@property
def _have_mm_energy(self):
return len(self.mm_energy) is not 0
# @property
# def _unique_torsions(self):
# Not returning the right amount. debug
# torsions = []
# for i in range(len(self.torsion_index)):
# try:
# if (self.torsion_index[i] != self.torsion_index[i+1]).all():
# torsions.append(self.torsion_index[i]), torsions.append(self.torsion_index[i+1])
# except:
# pass
# return len(torsions), torsions
def __getitem__(self, key):
"Get a slice of this trajectory"
return self.slice(key)
def slice(self, key, copy=True):
"""Slice trajectory, by extracting one or more frames into a separate object
This method can also be called using index bracket notation, i.e
`traj[1] == traj.slice(1)`
Parameters
----------
key : {int, np.ndarray, slice}
The slice to take. Can be either an int, a list of ints, or a slice
object.
copy : bool, default=True
Copy the arrays after slicing. If you set this to false, then if
you modify a slice, you'll modify the original array since they
point to the same data.
"""
xyz = self.xyz[key]
time = self.time[key]
torsions = self.torsion_index[key]
direction = self.direction[key]
steps = self.steps[key]
qm_energy = self.qm_energy[key]
unitcell_lengths, unitcell_angles = None, None
if self.unitcell_angles is not None:
unitcell_angles = self.unitcell_angles[key]
if self.unitcell_lengths is not None:
unitcell_lengths = self.unitcell_lengths[key]
if copy:
xyz = xyz.copy()
time = time.copy()
topology = deepcopy(self._topology)
structure = deepcopy(self.structure)
torsions = torsions.copy()
direction = direction.copy()
steps = steps.copy()
qm_energy = qm_energy.copy()
if self.unitcell_angles is not None:
unitcell_angles = unitcell_angles.copy()
if self.unitcell_lengths is not None:
unitcell_lengths = unitcell_lengths.copy()
newtraj = self.__class__(xyz, topology, structure, torsions, direction, steps, qm_energy)
if self._rmsd_traces is not None:
newtraj._rmsd_traces = np.array(self._rmsd_traces[key], ndmin=1, copy=True)
return newtraj
