def _execute(self, directory, available_resources):
yaml_filename = os.path.join(directory, "yank.yaml")
# Create the yank yaml input file from a dictionary of options.
with open(yaml_filename, "w") as file:
yaml.dump(
self._get_full_input_dictionary(available_resources),
file,
sort_keys=False,
)
setup_only = self.setup_only
# Yank is not safe to be called from anything other than the main thread.
# If the current thread is not detected as the main one, then yank should
# be spun up in a new process which should itself be safe to run yank in.
if threading.current_thread() is threading.main_thread():
logger.info("Launching YANK in the main thread.")
free_energy, free_energy_uncertainty = self._run_yank(
directory, available_resources, setup_only
)
else:
from multiprocessing import Process, Queue
logger.info("Launching YANK in a new process.")
# Create a queue to pass the results back to the main process.
queue = Queue()
# Create the process within which yank will run.
process = Process(
target=BaseYankProtocol._run_yank_as_process,
args=[queue, directory, available_resources, setup_only],
)
# Start the process and gather back the output.
process.start()
free_energy, free_energy_uncertainty, exception = queue.get()
process.join()
if exception is not None:
raise exception
self.estimated_free_energy = openmm_quantity_to_pint(free_energy).plus_minus(
openmm_quantity_to_pint(free_energy_uncertainty)
)
def _parameter_value_from_gradient_key(self, gradient_key):
"""Extracts the value of the parameter in the current
open force field object pointed to by a given
`ParameterGradientKey` object.
Parameters
----------
gradient_key: openff.evaluator.forcefield.ParameterGradientKey
The gradient key which points to the parameter of interest.
Returns
-------
unit.Quantity
The value of the parameter.
bool
Returns True if the parameter is a cosmetic one.
"""
try:
import openmm.unit as simtk_unit
except ImportError:
import simtk.unit as simtk_unit
parameter_handler = self.FF.openff_forcefield.get_parameter_handler(
gradient_key.tag)
parameter = (parameter_handler if gradient_key.smirks is None else
parameter_handler.parameters[gradient_key.smirks])
attribute_split = re.split(r"(\d+)", gradient_key.attribute)
attribute_split = list(filter(None, attribute_split))
parameter_attribute = None
parameter_value = None
if hasattr(parameter, gradient_key.attribute):
parameter_attribute = gradient_key.attribute
parameter_value = getattr(parameter, parameter_attribute)
elif len(attribute_split) == 2:
parameter_attribute = attribute_split[0]
if hasattr(parameter, parameter_attribute):
parameter_index = int(attribute_split[1]) - 1
parameter_value_list = getattr(parameter, parameter_attribute)
parameter_value = parameter_value_list[parameter_index]
is_cosmetic = False
if (parameter_attribute is None
or parameter_attribute in parameter._cosmetic_attribs):
is_cosmetic = True
if not isinstance(parameter_value, simtk_unit.Quantity):
parameter_value = parameter_value * simtk_unit.dimensionless
return openmm_quantity_to_pint(parameter_value), is_cosmetic
def test_openmm_to_pint(openmm_unit, value):
openmm_quantity = value * openmm_unit
openmm_raw_value = openmm_quantity.value_in_unit(openmm_unit)
pint_quantity = openmm_quantity_to_pint(openmm_quantity)
pint_raw_value = pint_quantity.magnitude
assert np.allclose(openmm_raw_value, pint_raw_value)
def test_daltons():
openmm_quantity = random() * simtk_unit.dalton
openmm_raw_value = openmm_quantity.value_in_unit(simtk_unit.gram /
simtk_unit.mole)
pint_quantity = openmm_quantity_to_pint(openmm_quantity)
pint_raw_value = pint_quantity.to(unit.gram / unit.mole).magnitude
assert np.allclose(openmm_raw_value, pint_raw_value)
def _execute(self, directory, available_resources):
force_field_source = ForceFieldSource.from_json(self.force_field_path)
if not isinstance(force_field_source, SmirnoffForceFieldSource):
raise ValueError("Only SMIRNOFF force fields are supported.")
force_field = force_field_source.to_force_field()
parameter_units = {
gradient_key: openmm_quantity_to_pint(
getattr(
force_field.get_parameter_handler(
gradient_key.tag).parameters[gradient_key.smirks],
gradient_key.attribute,
)).units
for gradient_key in self.gradient_parameters
}
self.input_observables.clear_gradients()
if isinstance(self.input_observables, Observable):
self.output_observables = Observable(
value=self.input_observables.value,
gradients=[
ParameterGradient(
key=gradient_key,
value=(0.0 * self.input_observables.value.units /
parameter_units[gradient_key]),
) for gradient_key in self.gradient_parameters
],
)
elif isinstance(self.input_observables, ObservableArray):
self.output_observables = ObservableArray(
value=self.input_observables.value,
gradients=[
ParameterGradient(
key=gradient_key,
value=(
numpy.zeros(self.input_observables.value.shape) *
self.input_observables.value.units /
parameter_units[gradient_key]),
) for gradient_key in self.gradient_parameters
],
)
else:
raise NotImplementedError()
def _compute_charge_derivatives(self, n_atoms: int):
d_charge_d_theta = {key: np.zeros(n_atoms) for key in self.gradient_parameters}
if len(self.gradient_parameters) > 0 and not isinstance(
self.parameterized_system.force_field, SmirnoffForceFieldSource
):
raise ValueError(
"Derivates can only be computed for systems parameterized with "
"SMIRNOFF force fields."
)
force_field = self.parameterized_system.force_field.to_force_field()
topology = self.parameterized_system.topology
for key in self.gradient_parameters:
reverse_system, reverse_value = system_subset(key, force_field, topology)
forward_system, forward_value = system_subset(
key, force_field, topology, 0.1
)
reverse_value = openmm_quantity_to_pint(reverse_value)
forward_value = openmm_quantity_to_pint(forward_value)
reverse_charges = self._extract_charges(reverse_system)
forward_charges = self._extract_charges(forward_system)
if reverse_charges is None and forward_charges is None:
d_charge_d_theta[key] /= forward_value.units
else:
d_charge_d_theta[key] = (forward_charges - reverse_charges) / (
forward_value - reverse_value
)
return d_charge_d_theta
def test_combinatorial_openmm_to_pint():
all_openmm_units = list(_get_all_openmm_units())
for i in range(len(all_openmm_units)):
for j in range(i, len(all_openmm_units)):
openmm_unit = all_openmm_units[i] * all_openmm_units[j]
openmm_quantity = random() * openmm_unit
openmm_raw_value = openmm_quantity.value_in_unit(openmm_unit)
pint_quantity = openmm_quantity_to_pint(openmm_quantity)
pint_raw_value = pint_quantity.magnitude
assert np.isclose(openmm_raw_value, pint_raw_value)
def test_openmm_unit_conversions(openmm_unit, value):
openmm_quantity = value * openmm_unit
openmm_base_quantity = openmm_quantity.in_unit_system(
simtk_unit.md_unit_system)
if not isinstance(openmm_base_quantity, simtk_unit.Quantity):
openmm_base_quantity *= simtk_unit.dimensionless
pint_base_quantity = openmm_quantity_to_pint(openmm_base_quantity)
pint_unit = openmm_unit_to_pint(openmm_unit)
pint_quantity = pint_base_quantity.to(pint_unit)
pint_raw_value = pint_quantity.magnitude
openmm_raw_value = openmm_quantity.value_in_unit(openmm_unit)
assert np.allclose(openmm_raw_value, pint_raw_value)
def _approximate_box_size_by_density(
molecules,
n_copies,
mass_density,
box_aspect_ratio,
box_scaleup_factor=1.1,
):
"""Generate an approximate box size based on the number and molecular
weight of the molecules present, and a target density for the final
solvated mixture.
Parameters
----------
molecules : list of openff.toolkit.topology.Molecule
The molecules in the system.
n_copies : list of int
The number of copies of each molecule.
mass_density : openff.evaluator.unit.Quantity
The target mass density for final system. It should have units
compatible with g / mL.
box_aspect_ratio: List of float
The aspect ratio of the simulation box, used in conjunction with
the `mass_density` parameter.
box_scaleup_factor : float
The factor by which the estimated box size should be
increased.
Returns
-------
openff.evaluator.unit.Quantity
A list of the three box lengths in units compatible with angstroms.
"""
volume = 0.0 * unit.angstrom**3
for (molecule, number) in zip(molecules, n_copies):
molecule_mass = reduce((lambda x, y: x + y),
[atom.mass for atom in molecule.atoms])
molecule_mass = openmm_quantity_to_pint(
molecule_mass) / unit.avogadro_constant
molecule_volume = molecule_mass / mass_density
volume += molecule_volume * number
box_length = volume**(1.0 / 3.0) * box_scaleup_factor
box_length_angstrom = box_length.to(unit.angstrom).magnitude
aspect_ratio_normalizer = (box_aspect_ratio[0] * box_aspect_ratio[1] *
box_aspect_ratio[2])**(1.0 / 3.0)
box_size = [
box_length_angstrom * box_aspect_ratio[0],
box_length_angstrom * box_aspect_ratio[1],
box_length_angstrom * box_aspect_ratio[2],
] * unit.angstrom
box_size /= aspect_ratio_normalizer
return box_size
def _compute_gradients(
gradient_parameters: List[ParameterGradientKey],
observables: ObservableFrame,
force_field: "ForceField",
thermodynamic_state: ThermodynamicState,
topology: "Topology",
trajectory: "Trajectory",
compute_resources: ComputeResources,
enable_pbc: bool = True,
perturbation_amount: float = 0.0001,
):
"""Computes the gradients of the provided observables with respect to
the set of specified force field parameters using the central difference
finite difference method.
Notes
-----
The ``observables`` object will be modified in-place.
Parameters
----------
gradient_parameters
The parameters to differentiate with respect to.
observables
The observables to differentiate.
force_field
The full set force field parameters which contain the parameters to
differentiate.
thermodynamic_state
The state at which the trajectory was sampled
topology
The topology of the system the observables were collected for.
trajectory
The trajectory over which the observables were collected.
compute_resources
The compute resources available for the computations.
enable_pbc
Whether PBC should be enabled when re-evaluating system energies.
perturbation_amount
The amount to perturb for the force field parameter by.
"""
from simtk import openmm
gradients = defaultdict(list)
observables.clear_gradients()
if enable_pbc:
# Make sure the PBC are set on the topology otherwise the cut-off will be
# set incorrectly.
topology.box_vectors = trajectory.openmm_boxes(0)
for parameter_key in gradient_parameters:
# Build the slightly perturbed systems.
reverse_system, reverse_parameter_value = system_subset(
parameter_key, force_field, topology, -perturbation_amount)
forward_system, forward_parameter_value = system_subset(
parameter_key, force_field, topology, perturbation_amount)
# Perform a cheap check to try and catch most cases where the systems energy
# does not depend on this parameter.
reverse_xml = openmm.XmlSerializer.serialize(reverse_system)
forward_xml = openmm.XmlSerializer.serialize(forward_system)
if not enable_pbc:
disable_pbc(reverse_system)
disable_pbc(forward_system)
reverse_parameter_value = openmm_quantity_to_pint(
reverse_parameter_value)
forward_parameter_value = openmm_quantity_to_pint(
forward_parameter_value)
# Evaluate the energies using the reverse and forward sub-systems.
if reverse_xml != forward_xml:
reverse_energies = _evaluate_energies(
thermodynamic_state,
reverse_system,
trajectory,
compute_resources,
enable_pbc,
)
forward_energies = _evaluate_energies(
thermodynamic_state,
forward_system,
trajectory,
compute_resources,
enable_pbc,
)
else:
zeros = np.zeros(len(trajectory))
reverse_energies = forward_energies = ObservableFrame({
ObservableType.PotentialEnergy:
ObservableArray(
zeros * unit.kilojoule / unit.mole,
[
ParameterGradient(
key=parameter_key,
value=(zeros * unit.kilojoule / unit.mole /
reverse_parameter_value.units),
)
],
),
ObservableType.ReducedPotential:
ObservableArray(
zeros * unit.dimensionless,
[
ParameterGradient(
key=parameter_key,
value=(zeros * unit.dimensionless /
reverse_parameter_value.units),
)
],
),
})
potential_gradient = ParameterGradient(
key=parameter_key,
value=(forward_energies[ObservableType.PotentialEnergy].value -
reverse_energies[ObservableType.PotentialEnergy].value) /
(forward_parameter_value - reverse_parameter_value),
)
reduced_potential_gradient = ParameterGradient(
key=parameter_key,
value=(forward_energies[ObservableType.ReducedPotential].value -
reverse_energies[ObservableType.ReducedPotential].value) /
(forward_parameter_value - reverse_parameter_value),
)
gradients[ObservableType.PotentialEnergy].append(potential_gradient)
gradients[ObservableType.TotalEnergy].append(potential_gradient)
gradients[ObservableType.Enthalpy].append(potential_gradient)
gradients[ObservableType.ReducedPotential].append(
reduced_potential_gradient)
if ObservableType.KineticEnergy in observables:
gradients[ObservableType.KineticEnergy].append(
ParameterGradient(
key=parameter_key,
value=(
np.zeros(potential_gradient.value.shape) *
observables[ObservableType.KineticEnergy].value.units /
reverse_parameter_value.units),
))
if ObservableType.Density in observables:
gradients[ObservableType.Density].append(
ParameterGradient(
key=parameter_key,
value=(np.zeros(potential_gradient.value.shape) *
observables[ObservableType.Density].value.units /
reverse_parameter_value.units),
))
if ObservableType.Volume in observables:
gradients[ObservableType.Volume].append(
ParameterGradient(
key=parameter_key,
value=(np.zeros(potential_gradient.value.shape) *
observables[ObservableType.Volume].value.units /
reverse_parameter_value.units),
))
for observable_type in observables:
observables[observable_type] = ObservableArray(
value=observables[observable_type].value,
gradients=gradients[observable_type],
)
def _execute(self, directory, available_resources):
import mdtraj
from openforcefield.topology import Molecule, Topology
with open(self.force_field_path) as file:
force_field_source = ForceFieldSource.parse_json(file.read())
if not isinstance(force_field_source, SmirnoffForceFieldSource):
raise ValueError(
"Only SMIRNOFF force fields are supported by this protocol.", )
# Load in the inputs
force_field = force_field_source.to_force_field()
trajectory = mdtraj.load_dcd(self.trajectory_file_path,
self.coordinate_file_path)
unique_molecules = []
for component in self.substance.components:
molecule = Molecule.from_smiles(smiles=component.smiles)
unique_molecules.append(molecule)
pdb_file = app.PDBFile(self.coordinate_file_path)
topology = Topology.from_openmm(pdb_file.topology,
unique_molecules=unique_molecules)
# Compute the difference between the energies using the reduced force field,
# and the full force field.
energy_corrections = None
if self.use_subset_of_force_field:
target_system, _ = self._build_reduced_system(
force_field, topology)
subset_potentials_path = os.path.join(directory, "subset.csv")
subset_potentials = self._evaluate_reduced_potential(
target_system, trajectory, subset_potentials_path,
available_resources)
full_statistics = StatisticsArray.from_pandas_csv(
self.statistics_path)
energy_corrections = (
full_statistics[ObservableType.PotentialEnergy] -
subset_potentials[ObservableType.PotentialEnergy])
# Build the slightly perturbed system.
reverse_system, reverse_parameter_value = self._build_reduced_system(
force_field, topology, -self.perturbation_scale)
forward_system, forward_parameter_value = self._build_reduced_system(
force_field, topology, self.perturbation_scale)
self.reverse_parameter_value = openmm_quantity_to_pint(
reverse_parameter_value)
self.forward_parameter_value = openmm_quantity_to_pint(
forward_parameter_value)
# Calculate the reduced potentials.
self.reverse_potentials_path = os.path.join(directory, "reverse.csv")
self.forward_potentials_path = os.path.join(directory, "forward.csv")
self._evaluate_reduced_potential(
reverse_system,
trajectory,
self.reverse_potentials_path,
available_resources,
energy_corrections,
)
self._evaluate_reduced_potential(
forward_system,
trajectory,
self.forward_potentials_path,
available_resources,
energy_corrections,
)