1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
//! Note: We keep most thermostat and barostat code as f64, although we use f32 in most sections.
use na_seq::Element;
use rand::{Rng, distr::Distribution};
use rand_distr::{ChiSquared, StandardNormal};
use crate::{
HydrogenConstraint, MdState, NATIVE_TO_KCAL,
solvent::{H_MASS, MASS_WATER_MOL, O_MASS},
};
// Per-molecule Boltzmann, in kcal/mol/K.
// For assigning velocities from temperature, and other thermostat/barostat use.
pub(crate) const GAS_CONST_R: f64 = 0.001_987_204_1; // kcal mol⁻¹ K⁻¹ (Amber-style units)
// Boltzmann constant in (amu · Ų/ps²) K⁻¹
// We use this for the Langevin and Anderson thermostat, where we need per-particle Gaussian noise or variance.
pub(crate) const KB_A2_PS2_PER_K_PER_AMU: f32 = 0.831_446_26;
// TAU is for the CSVR thermostat. In ps. Lower means more sensitive.
// We set an aggressive thermostat during solvent initialization, then a more relaxed one at runtime.
// This is for the VV/CVSR themostat only.
// Note: These are publically exposed, for use in applications.
pub const TAU_TEMP_DEFAULT: f64 = 0.1; // GROMACS default.
pub const TAU_TEMP_WATER_INIT: f64 = 0.01; // for CSVR
// These are in 1/ps. 1 ps^-1 is a good default for explicit solvent and constrained H bonds.
// Lower is closer to Newtonian dynamics. 2ps (0.5ps^-1) is Gromac's default.
pub const LANGEVIN_GAMMA_DEFAULT: f32 = 0.5;
pub const LANGEVIN_GAMMA_WATER_INIT: f32 = 15.;
impl MdState {
/// Computes total kinetic energy, in native units.
/// Includes all non-static atoms, including solvent.
pub(crate) fn measure_kinetic_energy(&self) -> f64 {
let mut result = 0.0;
for a in &self.atoms {
if !a.static_ {
result += (a.mass * a.vel.magnitude_squared()) as f64;
}
}
// Do not include the M/EP site.
for w in &self.water {
result += (w.o.mass * w.o.vel.magnitude_squared()) as f64;
result += (w.h0.mass * w.h0.vel.magnitude_squared()) as f64;
result += (w.h1.mass * w.h1.vel.magnitude_squared()) as f64;
}
// Add in the 0.5 factor, and convert from amu • (Å/ps)² to kcal/mol.
result * 0.5 * NATIVE_TO_KCAL as f64
}
/// COM-only kinetic energy for solvent + full atomic KE for non-solvent atoms, in kcal/mol.
/// Used for pressure via the molecular virial theorem: rigid solvent molecules contribute
/// only their translational (COM) KE, so no SETTLE constraint virial is needed.
pub(crate) fn measure_kinetic_energy_translational(&self) -> f64 {
let mut result = 0.0;
for a in &self.atoms {
if !a.static_ {
result += (a.mass * a.vel.magnitude_squared()) as f64;
}
}
for w in &self.water {
let v_com = (w.o.vel * O_MASS + w.h0.vel * H_MASS + w.h1.vel * H_MASS) / MASS_WATER_MOL;
result += (MASS_WATER_MOL * v_com.magnitude_squared()) as f64;
}
result * 0.5 * NATIVE_TO_KCAL as f64
}
/// Instantaneous temperature [K]
pub(crate) fn measure_temperature(&self) -> f64 {
(2.0 * self.kinetic_energy) / (self.thermo_dof as f64 * GAS_CONST_R)
}
/// Used in temperature computation. Constraints tracked are Hydrogen if constrained, COM drift removal,
/// and static atoms.
/// We cache this at init. Used for kinetic energy and temperature computations.
pub(crate) fn dof_for_thermo(&self) -> usize {
// 3 positional + 3 rotational for each solvent mol.
let mut result = 6 * self.water.len();
if !self.solvent_only_sim_at_init {
result += 3 * self.atoms.iter().filter(|a| !a.static_).count();
}
let num_constraints = {
let mut c = 0;
if !self.solvent_only_sim_at_init {
for atom in &self.atoms {
if matches!(
self.cfg.hydrogen_constraint,
HydrogenConstraint::Shake { shake_tolerance: _ }
) && atom.element == Element::Hydrogen
&& !atom.static_
{
c += 1;
}
}
}
if self.cfg.zero_com_drift {
c += 6;
}
c
};
result.saturating_sub(num_constraints)
}
/// Canonical Sampling through Velocities Rescaling thermostat. (Also known as Bussi, its
/// primary author)
/// [CSVR thermostat](https://arxiv.org/pdf/0803.4060)
/// A canonical velocity-rescale algorithm.
/// Cheap with gentle coupling, but doesn't imitate solvent drag.
pub(crate) fn apply_thermostat_csvr(&mut self, dt: f64, tau: f64, t_target: f64) {
if tau <= 0.0 {
return;
}
// This value is cached at init.
let dof = self.thermo_dof.max(2) as f64;
// Measure current KE from velocities so we get the post-kick value, not a stale cache.
let ke = self.measure_kinetic_energy(); // In kcal/mol
if ke < 1e-20 {
return;
}
let c = (-dt / tau).exp();
// Draw the two random variates used in the exact CSVR update:
let r: f64 = StandardNormal.sample(&mut self.barostat.rng); // N(0,1)
let chi = ChiSquared::new(dof - 1.0)
.unwrap()
.sample(&mut self.barostat.rng); // χ²_{dof-1}
let ke_target = 0.5 * dof * GAS_CONST_R * t_target;
// Discrete-time exact solution for the OU process in K (from Bussi 2007):
// K' = K*c + ke_bar*(1.0 - c) * [ (chi + r*r)/dof ] + 2.0*r*sqrt(c*(1.0-c)*K*ke_bar/dof)
let k_prime = ke * c
+ ke_target * (1.0 - c) * ((chi + r * r) / dof)
+ 2.0 * r * ((c * (1.0 - c) * ke * ke_target / dof).sqrt());
let k_prime = k_prime.max(1e-20);
let lam = (k_prime / ke).sqrt() as f32;
for a in &mut self.atoms {
if a.static_ {
continue;
}
a.vel *= lam;
}
for w in &mut self.water {
w.o.vel *= lam;
w.h0.vel *= lam;
w.h1.vel *= lam;
}
}
/// A thermostat that integrates the stochastic Langevin equation. Good temperature control
/// and ergodicity, but the friction parameter damps real dynamics as it grows. This applies an OU update.
pub(crate) fn apply_langevin_thermostat(&mut self, dt: f32, gamma: f32, temp_tgt_k: f32) {
let c = (-gamma * dt).exp();
let s2 = (1.0 - c * c).max(0.0); // numerical guard
let sigma_num = KB_A2_PS2_PER_K_PER_AMU * temp_tgt_k * s2;
let sigma_o = (sigma_num / O_MASS).sqrt();
let sigma_h = (sigma_num / H_MASS).sqrt();
for a in &mut self.atoms {
if a.static_ {
continue;
}
// per-component σ for velocity noise
let sigma = (sigma_num / a.mass).sqrt();
let nx: f32 = self.barostat.rng.sample(StandardNormal);
let ny: f32 = self.barostat.rng.sample(StandardNormal);
let nz: f32 = self.barostat.rng.sample(StandardNormal);
a.vel.x = c * a.vel.x + sigma * nx;
a.vel.y = c * a.vel.y + sigma * ny;
a.vel.z = c * a.vel.z + sigma * nz;
}
for w in &mut self.water {
let (ox, oy, oz): (f32, f32, f32) = (
self.barostat.rng.sample(StandardNormal),
self.barostat.rng.sample(StandardNormal),
self.barostat.rng.sample(StandardNormal),
);
let (h0x, h0y, h0z): (f32, f32, f32) = (
self.barostat.rng.sample(StandardNormal),
self.barostat.rng.sample(StandardNormal),
self.barostat.rng.sample(StandardNormal),
);
let (h1x, h1y, h1z): (f32, f32, f32) = (
self.barostat.rng.sample(StandardNormal),
self.barostat.rng.sample(StandardNormal),
self.barostat.rng.sample(StandardNormal),
);
w.o.vel.x = c * w.o.vel.x + sigma_o * ox;
w.o.vel.y = c * w.o.vel.y + sigma_o * oy;
w.o.vel.z = c * w.o.vel.z + sigma_o * oz;
w.h0.vel.x = c * w.h0.vel.x + sigma_h * h0x;
w.h0.vel.y = c * w.h0.vel.y + sigma_h * h0y;
w.h0.vel.z = c * w.h0.vel.z + sigma_h * h0z;
w.h1.vel.x = c * w.h1.vel.x + sigma_h * h1x;
w.h1.vel.y = c * w.h1.vel.y + sigma_h * h1y;
w.h1.vel.z = c * w.h1.vel.z + sigma_h * h1z;
}
}
}