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, )