Skip to main content

hoomd_md/modify/
thermalize_momentum.rs

1// Copyright (c) 2024-2026 The Regents of the University of Michigan.
2// Part of hoomd-rs, released under the BSD 3-Clause License.
3
4//! Implement `ThermalizeMomentum`
5
6use 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}