Skip to main content

hoomd_md/modify/
mod.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//! Methods for thermalizing and zeroing the momenta.
5
6use hoomd_microstate::{Body, Tagged};
7
8mod thermalize_angular_momentum;
9mod thermalize_momentum;
10mod zero_center_angular_momentum;
11mod zero_center_momentum;
12
13/// Draw random momenta from a thermal distribution.
14///
15/// In the [Maxwell–Boltzmann distribution], each component of the momentum $` p_i `$
16/// is normally distributed with mean 0 and variance $` \sigma^2 = m k T`$:
17/// ```math
18///    f(p_i) = \frac{1}{\sqrt{2 \pi m k T}} \exp{\left( -\frac{p_i^2}{2 m k T} \right)}
19/// ```
20///
21/// where $` f `$ is probability, $` m `$ is mass, $` k `$ is the Boltzmann
22/// constant, and $` T `$ is temperature.
23///
24/// The momenta generated by [`ThermalizeMomentum`] do not collectively sum to
25/// zero. To zero the effective momentum of the system's center of mass, use
26/// [`ZeroCenterMomentum`].
27///
28/// [Maxwell–Boltzmann distribution]: https://en.wikipedia.org/wiki/Maxwell%E2%80%93Boltzmann_distribution
29///
30/// # Example
31///
32/// ```
33/// use hoomd_md::ThermalizeMomentum;
34/// use hoomd_microstate::{
35///     Body, Microstate,
36///     property::{DynamicPoint, Point},
37/// };
38/// use hoomd_vector::Cartesian;
39///
40/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
41/// let mut microstate = Microstate::builder()
42///     .bodies([
43///         Body::single_site(
44///             DynamicPoint {
45///                 position: Cartesian::from([1.0, 2.0]),
46///                 ..Default::default()
47///             },
48///             Point::default(),
49///         ),
50///         Body::single_site(
51///             DynamicPoint {
52///                 position: Cartesian::from([-2.0, 3.0]),
53///                 ..Default::default()
54///             },
55///             Point::default(),
56///         ),
57///     ])
58///     .try_build()?;
59///
60/// microstate.thermalize_momentum(1.5);
61/// # Ok(())
62/// # }
63/// ```
64pub trait ThermalizeMomentum<B, S> {
65    /// Assign thermally distributed random momenta to all bodies in the microstate.
66    #[inline]
67    fn thermalize_momentum(&mut self, temperature: f64) {
68        self.thermalize_momentum_with_filter(temperature, |_| true);
69    }
70
71    /// Assign thermally distributed random momenta to selected bodies in the microstate.
72    fn thermalize_momentum_with_filter<F: Fn(&Tagged<Body<B, S>>) -> bool>(
73        &mut self,
74        temperature: f64,
75        should_thermalize_body: F,
76    );
77}
78
79/// Remove translational motion from the system's center of mass.
80///
81/// [`ZeroCenterMomentum`] subtracts the average momentum from every body's momentum:
82/// ```math
83/// \vec{p}_{i,\mathrm{new}} = \vec{p}_{i,\mathrm{old}} - \langle \vec{p}_\mathrm{old} \rangle
84/// ```
85///
86/// # Example
87///
88/// ```
89/// use hoomd_md::ZeroCenterMomentum;
90/// use hoomd_microstate::{
91///     Body, Microstate,
92///     property::{DynamicPoint, Point},
93/// };
94/// use hoomd_vector::Cartesian;
95///
96/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
97/// let mut microstate = Microstate::builder()
98///     .bodies([
99///         Body::single_site(
100///             DynamicPoint {
101///                 position: Cartesian::from([1.0, 2.0]),
102///                 momentum: Cartesian::from([-2.0, 4.0]),
103///                 ..Default::default()
104///             },
105///             Point::default(),
106///         ),
107///         Body::single_site(
108///             DynamicPoint {
109///                 position: Cartesian::from([-2.0, 3.0]),
110///                 momentum: Cartesian::from([3.0, -6.0]),
111///                 ..Default::default()
112///             },
113///             Point::default(),
114///         ),
115///     ])
116///     .try_build()?;
117///
118/// microstate.zero_center_momentum();
119/// # Ok(())
120/// # }
121/// ```
122pub trait ZeroCenterMomentum<B, S> {
123    /// Subtract the average momentum from each body's momentum.
124    #[inline]
125    fn zero_center_momentum(&mut self) {
126        self.zero_center_momentum_with_filter(|_| true);
127    }
128
129    /// Subtract the average momentum from each selected body's momentum.
130    fn zero_center_momentum_with_filter<F: Fn(&Tagged<Body<B, S>>) -> bool>(
131        &mut self,
132        should_zero_body: F,
133    );
134}
135
136/// Remove angular motion about the system's center of mass.
137///
138/// [`ZeroCenterAngularMomentum`] adjusts the translational momentum of every body in order to zero
139/// out the total angular momentum of the system about the center of mass (ignoring
140/// periodic boundary conditions).
141///
142/// # 2D
143///
144/// In 2D, [`ZeroCenterAngularMomentum`] applies:
145/// ```math
146/// \vec{p}_{i,\mathrm{new}} = \vec{p}_{i,\mathrm{old}} - \left( [-r_{ci}^{y}, r_{ci}^{x}] \right) \frac{L_c}{I_c} m_i
147/// ```
148/// where $`i`$ is the index of each body in a system, $`L_c`$ is the
149/// angular momentum about the center of mass, $`I_c`$ is the moment of
150/// inertia about the center of mass, and $`\vec{r}_{ci}`$ is the position of body *i*
151/// relative to the center of mass.
152///
153/// # 3D
154///
155/// In 3D, [`ZeroCenterAngularMomentum`] applies:
156/// ```math
157/// \vec{p}_{i,\mathrm{new}} = \vec{p}_{i,\mathrm{old}} - \left( \vec{\omega}_c \times \vec{r}_{ci} \right) m_i
158/// ```
159/// where $`i`$ is the index of each body in a system,
160/// $`\vec{\omega}_c`$ is angular velocity about the center of mass, and
161/// $`\vec{r}_{ci}`$ is the position of body *i* relative to the center of mass.
162///
163/// $`\vec{\omega}_c`$ is obtained by solving the following linear system:
164/// ```math
165/// \mathbf{I}_c \vec{\omega}_c = \vec{L}_c
166/// ```
167/// where $`\mathbf{I}_c`$ is the moment of inertia about the center of mass,
168/// and $`\vec{L}_c`$ is the angular momentum about the center of mass.
169///
170/// # Example
171///
172/// ```
173/// use hoomd_md::ZeroCenterAngularMomentum;
174/// use hoomd_microstate::{
175///     Body, Microstate,
176///     property::{DynamicPoint, Point},
177/// };
178/// use hoomd_vector::Cartesian;
179///
180/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
181/// let mut microstate = Microstate::builder()
182///     .bodies([
183///         Body::single_site(
184///             DynamicPoint {
185///                 position: Cartesian::from([1.0, 2.0]),
186///                 momentum: Cartesian::from([-2.0, 4.0]),
187///                 ..Default::default()
188///             },
189///             Point::default(),
190///         ),
191///         Body::single_site(
192///             DynamicPoint {
193///                 position: Cartesian::from([-2.0, 3.0]),
194///                 momentum: Cartesian::from([3.0, -6.0]),
195///                 ..Default::default()
196///             },
197///             Point::default(),
198///         ),
199///     ])
200///     .try_build()?;
201///
202/// microstate.zero_center_angular_momentum();
203/// # Ok(())
204/// # }
205/// ```
206pub trait ZeroCenterAngularMomentum<B, S> {
207    /// Adjust each body's angular translational momentum in order to zero the system's overall
208    /// angular momentum about the center of mass.
209    #[inline]
210    fn zero_center_angular_momentum(&mut self) {
211        self.zero_center_angular_momentum_with_filter(|_| true);
212    }
213
214    /// Adjust each selected body's translational momentum in order to zero the system's overall
215    /// angular momentum about the center of mass..
216    fn zero_center_angular_momentum_with_filter<F: Fn(&Tagged<Body<B, S>>) -> bool>(
217        &mut self,
218        should_zero_body: F,
219    );
220}
221
222/// Draw random angular momenta from a thermal distribution.
223///
224/// In the [Maxwell–Boltzmann distribution], each component of the angular momentum $` L_i `$
225/// (aligned to the principal axes) is normally distributed with mean 0 and variance
226/// $` \sigma^2 = I_i k T`$:
227/// ```math
228///    f(L_i) = \frac{1}{\sqrt{2 \pi I_i k T}} \exp{\left( -\frac{L_i^2}{2 I_i k T} \right)}
229/// ```
230///
231/// where $` f `$ is probability, $` I_i `$ is i<sup>th</sup> component of the diagonalized
232/// moment of inertia, $` k `$ is the Boltzmann constant, and $` T `$ is temperature.
233///
234/// [Maxwell–Boltzmann distribution]: https://en.wikipedia.org/wiki/Maxwell%E2%80%93Boltzmann_distribution
235///
236/// # Example
237///
238/// ```
239/// use hoomd_md::ThermalizeAngularMomentum;
240/// use hoomd_microstate::{
241///     Body, Microstate,
242///     property::{DynamicOrientedPoint, Point},
243/// };
244/// use hoomd_vector::{Angle, Cartesian};
245///
246/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
247/// let mut microstate = Microstate::builder()
248///     .bodies([
249///         Body::single_site(
250///             DynamicOrientedPoint {
251///                 position: Cartesian::from([1.0, 2.0]),
252///                 ..Default::default()
253///             },
254///             Point::default(),
255///         ),
256///         Body::single_site(
257///             DynamicOrientedPoint {
258///                 position: Cartesian::from([-2.0, 3.0]),
259///                 ..Default::default()
260///             },
261///             Point::default(),
262///         ),
263///     ])
264///     .try_build()?;
265///
266/// microstate.thermalize_angular_momentum(1.5);
267/// # Ok(())
268/// # }
269/// ```
270pub trait ThermalizeAngularMomentum<B, S> {
271    /// Assign thermally distributed random angular momenta to all bodies in the microstate.
272    #[inline]
273    fn thermalize_angular_momentum(&mut self, temperature: f64) {
274        self.thermalize_angular_momentum_with_filter(temperature, |_| true);
275    }
276
277    /// Assign thermally distributed random angular momenta to selected bodies in the microstate.
278    fn thermalize_angular_momentum_with_filter<F: Fn(&Tagged<Body<B, S>>) -> bool>(
279        &mut self,
280        temperature: f64,
281        should_thermalize_body: F,
282    );
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288    use approxim::{assert_abs_diff_eq, assert_relative_eq};
289    use hoomd_microstate::{
290        Body, Microstate,
291        property::{AngularMomentum, DynamicOrientedPoint, DynamicPoint, Momentum, Point},
292    };
293    use hoomd_vector::{Cartesian, Versor, Wedge};
294    use rstest::*;
295
296    // When draw n samples from gaussian of N(mu, sigma)
297    // the error in the sample mean is ~ sigma * sqrt(1/n)
298    // the error in the sample variance is ~ sigma^2 * sqrt(2/(n-1))
299    // here I use 4 * sigma * sqrt(1/n) and 4 * sigma^2 * sqrt(2/(n-1)) as the testing tolerance
300    // which should cover 99.99% cases.
301    // sqrt(2/9999) ~ 0.01414284
302    const N_BODIES: usize = 10000;
303    const EPSILON_MEAN_SCALE: f64 = 4.0 * 0.01;
304    const EPSILON_VARIANCE_SCALE: f64 = 4.0 * 0.014_142_840;
305
306    mod momentum {
307        use super::*;
308
309        fn create_point_body_3d(
310            position: Cartesian<3>,
311            mass: f64,
312            momentum: Cartesian<3>,
313        ) -> Body<DynamicPoint<Cartesian<3>>, Point<Cartesian<3>>> {
314            Body {
315                properties: DynamicPoint {
316                    position,
317                    momentum,
318                    mass,
319                    ..Default::default()
320                },
321                sites: vec![Point::new(Cartesian::from([0.0, 0.0, 0.0]))],
322            }
323        }
324
325        #[rstest]
326        fn test_distribution(
327            #[values(0.5, 1.0, 2.0)] mass: f64,
328            #[values(0.5, 1.5)] temperature: f64,
329            #[values(42, 123, 999)] seed: u32,
330        ) -> anyhow::Result<()> {
331            let mut microstate = Microstate::builder().seed(seed).try_build()?;
332            let expected_variance = temperature * mass;
333
334            for _ in 0..N_BODIES {
335                microstate.add_body(create_point_body_3d(
336                    Cartesian::default(),
337                    mass,
338                    Cartesian::default(),
339                ))?;
340            }
341
342            microstate.thermalize_momentum(temperature);
343
344            let momenta: Vec<[f64; 3]> = microstate
345                .bodies()
346                .iter()
347                .map(|b| b.item.properties.momentum().coordinates)
348                .collect();
349
350            for dim in 0..3 {
351                let components: Vec<f64> = momenta.iter().map(|m| m[dim]).collect();
352                let mean = components.iter().sum::<f64>() / N_BODIES as f64;
353                let variance = components.iter().map(|&v| (v - mean).powi(2)).sum::<f64>()
354                    / (N_BODIES - 1) as f64;
355
356                assert_abs_diff_eq!(
357                    mean,
358                    0.0,
359                    epsilon = expected_variance.sqrt() * EPSILON_MEAN_SCALE
360                );
361                assert_abs_diff_eq!(
362                    variance,
363                    expected_variance,
364                    epsilon = expected_variance * EPSILON_VARIANCE_SCALE
365                );
366            }
367
368            Ok(())
369        }
370
371        #[rstest]
372        fn zero_center_momentum(
373            #[values(1.0, 2.0)] mass_a: f64,
374            #[values([1.0, -1.0, 0.0], [-1.0, 1.0, 1.0])] momentum_a: [f64; 3],
375            #[values(1.0, 0.5)] mass_b: f64,
376            #[values([-1.0, 0.0, 1.0], [0.5, 1.5, 3.0])] momentum_b: [f64; 3],
377        ) -> anyhow::Result<()> {
378            let momentum_1 = Cartesian::from(momentum_a);
379            let momentum_2 = Cartesian::from(momentum_b);
380            let center_momentum = momentum_1 + momentum_2;
381
382            let mut microstate = Microstate::builder().try_build()?;
383            microstate.add_body(create_point_body_3d(
384                Cartesian::default(),
385                mass_a,
386                momentum_1,
387            ))?;
388            microstate.add_body(create_point_body_3d(
389                Cartesian::default(),
390                mass_b,
391                momentum_2,
392            ))?;
393
394            microstate.zero_center_momentum();
395
396            let modified_momentum_1 = microstate.bodies()[0].item.properties.momentum;
397            let modified_momentum_2 = microstate.bodies()[1].item.properties.momentum;
398
399            let expected_momentum_1 = momentum_1 - center_momentum / 2.0;
400            let expected_momentum_2 = momentum_2 - center_momentum / 2.0;
401
402            assert_abs_diff_eq!(
403                modified_momentum_1 + modified_momentum_2,
404                Cartesian::default(),
405                epsilon = 1e-15
406            );
407            assert_relative_eq!(modified_momentum_1, expected_momentum_1);
408            assert_relative_eq!(modified_momentum_2, expected_momentum_2);
409            Ok(())
410        }
411
412        #[rstest]
413        fn two_particles_stop_rotation(
414            #[values([1.0, 0.0, 0.0])] position_a: [f64; 3],
415            #[values(1.0, 2.0)] mass_a: f64,
416            #[values([1.0, -1.0, 0.0], [-1.0, 1.0, 1.0])] momentum_a: [f64; 3],
417            #[values([-1.0, 0.0, 0.0])] position_b: [f64; 3],
418            #[values(1.0, 0.5)] mass_b: f64,
419            #[values([-1.0, 0.0, 1.0], [0.5, 1.5, 3.0])] momentum_b: [f64; 3],
420        ) -> anyhow::Result<()> {
421            let position_a = Cartesian::from(position_a);
422            let momentum_a = Cartesian::from(momentum_a);
423            let position_b = Cartesian::from(position_b);
424            let momentum_b = Cartesian::from(momentum_b);
425            let position_center = (position_a * mass_a + position_b * mass_b) / (mass_a + mass_b);
426
427            let mut microstate = Microstate::builder().try_build()?;
428            microstate.add_body(create_point_body_3d(position_a, mass_a, momentum_a))?;
429            microstate.add_body(create_point_body_3d(position_b, mass_b, momentum_b))?;
430
431            microstate.zero_center_angular_momentum();
432
433            let modified_momentum_a = microstate.bodies()[0].item.properties.momentum;
434            let modified_momentum_b = microstate.bodies()[1].item.properties.momentum;
435
436            let modified_angular_momentum_a =
437                (position_a - position_center).wedge(&modified_momentum_a);
438            let modified_angular_momentum_b =
439                (position_b - position_center).wedge(&modified_momentum_b);
440
441            assert_abs_diff_eq!(
442                modified_angular_momentum_a + modified_angular_momentum_b,
443                Cartesian::default(),
444                epsilon = 1e-15
445            );
446            Ok(())
447        }
448    }
449
450    mod angular_momentum {
451        use super::*;
452
453        fn create_body_3d(
454            moment_of_inertia: [f64; 3],
455            angular_momentum: Cartesian<3>,
456        ) -> Body<DynamicOrientedPoint<Cartesian<3>, Versor>, Point<Cartesian<3>>> {
457            Body {
458                properties: DynamicOrientedPoint {
459                    moment_of_inertia,
460                    angular_momentum,
461                    ..Default::default()
462                },
463                sites: vec![Point::new(Cartesian::from([0.0, 0.0, 0.0]))],
464            }
465        }
466
467        #[rstest]
468        fn test_distribution(
469            #[values([1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [1.0, 1.0, 1.0], [4.0, 2.0, 0.5])]
470            inertia: [f64; 3],
471            #[values(0.5, 1.5)] temperature: f64,
472            #[values(42, 123, 999)] seed: u32,
473        ) -> anyhow::Result<()> {
474            let mut microstate = Microstate::builder().seed(seed).try_build()?;
475            let expected_variance = Cartesian::from(inertia) * temperature;
476
477            for _ in 0..N_BODIES {
478                microstate.add_body(create_body_3d(inertia, Cartesian::default()))?;
479            }
480
481            microstate.thermalize_angular_momentum(temperature);
482
483            let angular_momenta: Vec<[f64; 3]> = microstate
484                .bodies()
485                .iter()
486                .map(|b| b.item.properties.angular_momentum().coordinates)
487                .collect();
488
489            for dim in 0..3 {
490                let components: Vec<f64> = angular_momenta.iter().map(|m| m[dim]).collect();
491
492                let mean = components.iter().sum::<f64>() / (N_BODIES as f64);
493                let variance = components.iter().map(|&v| (v - mean).powi(2)).sum::<f64>()
494                    / (N_BODIES - 1) as f64;
495
496                assert_abs_diff_eq!(
497                    mean,
498                    0.0,
499                    epsilon = expected_variance[dim].sqrt() * EPSILON_MEAN_SCALE
500                );
501                assert_abs_diff_eq!(
502                    variance,
503                    expected_variance[dim],
504                    epsilon = expected_variance[dim] * EPSILON_VARIANCE_SCALE
505                );
506            }
507
508            Ok(())
509        }
510    }
511}