Skip to main content

hoomd_md/modify/
zero_center_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 `ZeroCenterMomentum`
5
6use 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}