hoomd_md/modify/
zero_center_momentum.rs1use super::ZeroCenterMomentum;
7use hoomd_microstate::{
8 Body, Microstate, SiteKey, Tagged, Transform,
9 boundary::{GenerateGhosts, Wrap},
10 property::{Momentum, Position},
11};
12use hoomd_spatial::PointUpdate;
13use hoomd_vector::Cartesian;
14
15impl<const N: usize, B, S, X, C> ZeroCenterMomentum<B, S> for Microstate<B, S, X, C>
16where
17 B: Position<Position = Cartesian<N>> + Momentum<Momentum = Cartesian<N>> + Transform<S> + Clone,
18 S: Position<Position = Cartesian<N>> + Default,
19 X: PointUpdate<Cartesian<N>, SiteKey>,
20 C: Wrap<B> + Wrap<S> + GenerateGhosts<S>,
21{
22 #[inline]
23 fn zero_center_momentum_with_filter<F: Fn(&Tagged<Body<B, S>>) -> bool>(
24 &mut self,
25 should_zero_body: F,
26 ) {
27 let (total_momentum, count) = self
28 .bodies()
29 .iter()
30 .filter(|&body| should_zero_body(body))
31 .fold((Cartesian::default(), 0), |(total, count), body| {
32 (total + *body.item.properties.momentum(), count + 1)
33 });
34 let average_momentum = total_momentum / f64::from(count);
35
36 for body_index in 0..self.bodies().len() {
37 let body = &self.bodies()[body_index];
38 if !should_zero_body(body) {
39 continue;
40 }
41
42 let mut body_properties = body.item.properties.clone();
43
44 *body_properties.momentum_mut() -= average_momentum;
45
46 self.update_body_properties(body_index, body_properties)
47 .expect("Bodies and sites should remain in simulation boundary.");
48 }
49 }
50}