hoomd_md/modify/
thermalize_momentum.rs1use std::array;
7
8use super::ThermalizeMomentum;
9use hoomd_microstate::{
10 Body, Microstate, SiteKey, Tagged, Transform,
11 boundary::{GenerateGhosts, Wrap},
12 property::{Mass, Momentum, Position},
13};
14use hoomd_spatial::PointUpdate;
15use hoomd_vector::Cartesian;
16use rand_distr::{Distribution, Normal};
17
18impl<const N: usize, B, S, X, C> ThermalizeMomentum<B, S> for Microstate<B, S, X, C>
19where
20 B: Position<Position = Cartesian<N>>
21 + Momentum<Momentum = Cartesian<N>>
22 + Mass
23 + Transform<S>
24 + Clone,
25 S: Position<Position = Cartesian<N>> + Default,
26 X: PointUpdate<Cartesian<N>, SiteKey>,
27 C: Wrap<B> + Wrap<S> + GenerateGhosts<S>,
28{
29 #[inline]
30 fn thermalize_momentum_with_filter<F: Fn(&Tagged<Body<B, S>>) -> bool>(
31 &mut self,
32 temperature: f64,
33 should_thermalize_body: F,
34 ) {
35 let mut rng = self.counter().make_rng();
36
37 for body_index in 0..self.bodies().len() {
38 let body = &self.bodies()[body_index];
39 if !should_thermalize_body(body) {
40 continue;
41 }
42
43 let mut body_properties = body.item.properties.clone();
44 let mass = body_properties.mass();
45 let sigma = (temperature * mass).sqrt();
46 let normal = Normal::new(0.0, sigma).expect("Normal distribution should be valid");
47
48 if mass > 0.0 {
49 let random_momentum: Cartesian<N> =
50 array::from_fn(|_| normal.sample(&mut rng)).into();
51 *body_properties.momentum_mut() = random_momentum;
52
53 self.update_body_properties(body_index, body_properties)
54 .expect("Bodies and sites should remain in simulation boundary.");
55 }
56 }
57
58 self.increment_substep();
59 }
60}