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 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 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 _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 calculate_energy(self, x: unit.Quantity): """ Given a coordinate set (x) the energy is calculated in kJ/mol. Parameters ---------- x : array of floats, unit'd (angstroms) initial configuration Returns ------- energy : unit.quantity.Quantity energy in kJ/mol """ assert type(x) in _allowable_quantities coordinates = torch.tensor([x.value_in_unit(unit.angstroms)], requires_grad=True, device=self.device, dtype=torch.float32) energy_in_hartrees = self._reform_as_energy_tensor(coordinates) energy = energy_in_hartrees.item() * self.hartree_to_kJ_per_mole * unit.kilojoule_per_mole return energy
def __init__( self, sigma: unit.Quantity, point: unit.Quantity, atom_idx: list, atoms: str, active_at: int = -1, ): """ Center of mass restraint. Parameters ---------- sigma : in angstrom point : np.array, unit'd atom_idx : list list of atoms idxs atoms: str Str of atoms to retrieve element information """ assert type(sigma) == unit.Quantity assert type(point) == unit.Quantity super().__init__(sigma, point.value_in_unit(unit.angstrom), active_at) cutoff_radius = 0.1 * unit.angstrom self.cutoff_radius = cutoff_radius.value_in_unit(unit.angstrom) # only look at heavy atoms full_mass_list = [mass_dict_in_daltons[atoms[i]] for i in atom_idx] heavy_atoms_idx = [i for i, x in enumerate(full_mass_list) if x != 1.0] self.atom_idx = heavy_atoms_idx self.mass_list = [heavy_atoms_idx[i] for i in heavy_atoms_idx] scaled_masses = np.array(self.mass_list) / np.array( self.mass_list).sum() self.masses = torch.tensor(scaled_masses, dtype=torch.double, device=self.device, requires_grad=True)