Skip to main content

hoomd_md/compute/
translational_kinetic_energy.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 `TranslationalKineticEnergy`
5
6use super::TranslationalKineticEnergy;
7use hoomd_microstate::{
8    Body, Microstate, Tagged,
9    property::{Mass, Momentum},
10};
11use hoomd_vector::InnerProduct;
12
13impl<V, B, S, X, C> TranslationalKineticEnergy<B, S> for Microstate<B, S, X, C>
14where
15    V: InnerProduct,
16    B: Momentum<Momentum = V> + Mass,
17{
18    #[inline]
19    fn translational_kinetic_energy_with_filter<F: Fn(&Tagged<Body<B, S>>) -> bool>(
20        &self,
21        should_sum_body: F,
22    ) -> (f64, usize) {
23        self.bodies()
24            .iter()
25            .filter(|&body| should_sum_body(body))
26            .fold((0.0, 0), |(total, count), body| {
27                let p = body.item.properties.momentum();
28                (
29                    total + p.norm_squared() / (2.0 * body.item.properties.mass()),
30                    count + V::n_dimensions(),
31                )
32            })
33    }
34}
35
36#[cfg(test)]
37mod test {
38    use super::*;
39    use approxim::assert_relative_eq;
40    use assert2::check;
41
42    use hoomd_microstate::{
43        Body,
44        property::{DynamicPoint, Point},
45    };
46    use hoomd_vector::Cartesian;
47
48    #[test]
49    fn kinetic_energy_2d() -> anyhow::Result<()> {
50        let microstate = Microstate::builder()
51            .bodies([
52                Body::single_site(DynamicPoint::default(), Point::default()),
53                Body::single_site(
54                    DynamicPoint {
55                        mass: 2.0,
56                        momentum: Cartesian::<2>::from([2.0, 0.0]),
57                        ..Default::default()
58                    },
59                    Point::default(),
60                ),
61                Body::single_site(
62                    DynamicPoint {
63                        mass: 4.0,
64                        momentum: Cartesian::<2>::from([1.0, 1.0]),
65                        ..Default::default()
66                    },
67                    Point::default(),
68                ),
69                Body::single_site(
70                    DynamicPoint {
71                        mass: 3.0,
72                        momentum: Cartesian::<2>::from([-4.0, -2.0]),
73                        ..Default::default()
74                    },
75                    Point::default(),
76                ),
77            ])
78            .try_build()?;
79
80        let (total_kinetic_energy, total_degrees_of_freedom) =
81            microstate.translational_kinetic_energy();
82        check!(total_degrees_of_freedom == 8);
83        assert_relative_eq!(total_kinetic_energy, 1.0 + 2.0 / 8.0 + 20.0 / 6.0);
84
85        let (filtered_kinetic_energy, filtered_degrees_of_freedom) =
86            microstate.translational_kinetic_energy_with_filter(|b| b.tag <= 1);
87        check!(filtered_degrees_of_freedom == 4);
88        assert_relative_eq!(filtered_kinetic_energy, 1.0);
89
90        Ok(())
91    }
92
93    #[test]
94    fn kinetic_energy_3d() -> anyhow::Result<()> {
95        let microstate = Microstate::builder()
96            .bodies([
97                Body::single_site(DynamicPoint::default(), Point::default()),
98                Body::single_site(
99                    DynamicPoint {
100                        mass: 2.0,
101                        momentum: Cartesian::<3>::from([2.0, 0.0, 0.0]),
102                        ..Default::default()
103                    },
104                    Point::default(),
105                ),
106                Body::single_site(
107                    DynamicPoint {
108                        mass: 4.0,
109                        momentum: Cartesian::<3>::from([1.0, 1.0, 1.0]),
110                        ..Default::default()
111                    },
112                    Point::default(),
113                ),
114                Body::single_site(
115                    DynamicPoint {
116                        mass: 3.0,
117                        momentum: Cartesian::<3>::from([-4.0, -2.0, 1.0]),
118                        ..Default::default()
119                    },
120                    Point::default(),
121                ),
122            ])
123            .try_build()?;
124
125        let (total_kinetic_energy, total_degrees_of_freedom) =
126            microstate.translational_kinetic_energy();
127        check!(total_degrees_of_freedom == 12);
128        assert_relative_eq!(total_kinetic_energy, 1.0 + 3.0 / 8.0 + 21.0 / 6.0);
129
130        let (filtered_kinetic_energy, filtered_degrees_of_freedom) =
131            microstate.translational_kinetic_energy_with_filter(|b| b.tag <= 1);
132        check!(filtered_degrees_of_freedom == 6);
133        assert_relative_eq!(filtered_kinetic_energy, 1.0);
134
135        Ok(())
136    }
137}