From e0aa4ab33be219cd3e386ec8183673f530c3ac3d Mon Sep 17 00:00:00 2001 From: Charlles Abreu Date: Mon, 4 Mar 2024 16:56:04 -0500 Subject: [PATCH 1/2] Improve velocity handling in contexts and states --- ufedmm/ufedmm.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/ufedmm/ufedmm.py b/ufedmm/ufedmm.py index 1c266de..a27c568 100644 --- a/ufedmm/ufedmm.py +++ b/ufedmm/ufedmm.py @@ -774,7 +774,7 @@ def getPositions(self, asNumpy=False, extended=False, split=True): positions, xvars = self._split(all_positions, asNumpy) return (positions, xvars) if extended else positions - def getVelocities(self, asNumpy=False, extended=False): + def getVelocities(self, asNumpy=False, extended=False, split=True): """Gets the velocities of all physical particles and optionally also gets the velocities of the extra particles associated to the extended- space dynamical variables. @@ -785,6 +785,9 @@ def getVelocities(self, asNumpy=False, extended=False): Whether to return Numpy arrays instead of lists of openmm.Vec3. extended : bool, default=False Whether to include the velocities of the extra particles. + split : bool, default=True + Whether to return the velocities of the physical particles and the + velocities of the extra particles separately. Returns ------- @@ -802,10 +805,11 @@ def getVelocities(self, asNumpy=False, extended=False): Exception If velocities were not requested in the ``context.getState()`` call. """ - velocities = super().getVelocities(asNumpy) - if not extended: - velocities, _ = self._split(velocities, asNumpy) - return velocities + all_velocities = super().getVelocities(asNumpy) + if extended and not split: + return all_velocities + velocities, vvars = self._split(all_velocities, asNumpy) + return (velocities, vvars) if extended else velocities def getDynamicalVariables(self): """Gets the values of the extended-space dynamical variables. @@ -1035,8 +1039,12 @@ def setVelocitiesToTemperature(self, temperature, randomSeed=None): + [v.temperature for v in self.variables] ) sigma = np.sqrt(_standardized(unit.MOLAR_GAS_CONSTANT_R) * T / m) - random_state = np.random.RandomState(randomSeed) - velocities = sigma[:, np.newaxis] * random_state.normal(0, 1, (Ntotal, 3)) + random = np.random.RandomState(randomSeed) + velocities = sigma[:, None] * random.normal(0, 1, (Ntotal, 3)) + vcm = np.sum(m[:Natoms, None] * velocities[:Natoms], axis=0) / m[:Natoms].sum() + velocities[:Natoms] -= vcm + for i in range(Natoms, Ntotal): + velocities[i][1:] = 0 super().setVelocities(velocities) From 5e9b2f386e4522c081c6b06f4af0569394e72f9c Mon Sep 17 00:00:00 2001 From: Charlles Abreu Date: Mon, 4 Mar 2024 18:17:31 -0500 Subject: [PATCH 2/2] Rescale velocities to target temperature --- ufedmm/integrators.py | 2 +- ufedmm/ufedmm.py | 30 ++++++++++++++++++++++-------- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/ufedmm/integrators.py b/ufedmm/integrators.py index ada6493..0c3e4f8 100644 --- a/ufedmm/integrators.py +++ b/ufedmm/integrators.py @@ -213,7 +213,7 @@ def update_temperatures(self, system_temperature, extended_space_temperatures): ) temperatures = [system_temperature] * nparticles + extended_space_temperatures kT = [ - unit.MOLAR_GAS_CONSTANT_R * T * openmm.Vec3(1, 1, 1) for T in temperatures + unit.MOLAR_GAS_CONSTANT_R * T * openmm.Vec3(1, 0, 0) for T in temperatures ] self.setPerDofVariableByName("kT", kT) self._up_to_date = True diff --git a/ufedmm/ufedmm.py b/ufedmm/ufedmm.py index a27c568..8d2d094 100644 --- a/ufedmm/ufedmm.py +++ b/ufedmm/ufedmm.py @@ -1030,21 +1030,35 @@ def setVelocitiesToTemperature(self, temperature, randomSeed=None): randomSeed : int, default=None A seed for the random number generator. """ + temperature = _standardized(temperature) system = self.getSystem() Ntotal = system.getNumParticles() Natoms = Ntotal - len(self.variables) - m = np.array([system.getParticleMass(i) / unit.dalton for i in range(Ntotal)]) - T = np.array( - [_standardized(temperature)] * Natoms - + [v.temperature for v in self.variables] + masses = np.array( + [system.getParticleMass(i) / unit.dalton for i in range(Ntotal)] ) - sigma = np.sqrt(_standardized(unit.MOLAR_GAS_CONSTANT_R) * T / m) + temperatures = np.array( + [temperature] * Natoms + [v.temperature for v in self.variables] + ) + kboltz = _standardized(unit.MOLAR_GAS_CONSTANT_R) + sigmas = np.sqrt(kboltz * temperatures / masses)[:, None] random = np.random.RandomState(randomSeed) - velocities = sigma[:, None] * random.normal(0, 1, (Ntotal, 3)) - vcm = np.sum(m[:Natoms, None] * velocities[:Natoms], axis=0) / m[:Natoms].sum() + velocities = sigmas * random.normal(0, 1, (Ntotal, 3)) + + total_mass = masses.sum() + vcm = np.sum(masses[:Natoms, None] * velocities[:Natoms], axis=0) / total_mass velocities[:Natoms] -= vcm - for i in range(Natoms, Ntotal): + num_constraints = system.getNumConstraints() + two_ke = np.sum(masses[:Natoms] * np.sum(velocities[:Natoms] ** 2, axis=1)) + target_two_ke = (3 * Natoms - num_constraints) * kboltz * temperature + velocities[:Natoms] *= np.sqrt(target_two_ke / two_ke) + + for i, v in enumerate(self.variables, start=Natoms): + two_ke = masses[i] * velocities[i][0] ** 2 + target_two_ke = kboltz * v.temperature + velocities[i][0] *= np.sqrt(target_two_ke / two_ke) velocities[i][1:] = 0 + super().setVelocities(velocities)