rustitude_core/
four_momentum.rs

1//! # Four-Momentum
2//!
3//! This module contains a helper struct, [`FourMomentum`], which contains useful methods and
4//! manipulations for physics four-vectors representing momentum coordinates. In particular,
5//! this struct has the same layout as a `[Field; 4]` with components identified as
6//! $`(E, p_x, p_y, p_z)`$.
7use crate::Field;
8use nalgebra::{Matrix4, Vector3, Vector4};
9use std::{
10    fmt::Display,
11    ops::{Add, Sub},
12};
13
14/// Struct which holds energy and three-momentum as a four-vector.
15///
16/// A four-momentum structure with helpful methods for boosts.
17///
18/// This is the basic structure of a Lorentz four-vector
19/// of the form $`(E, \vec{p})`$ where $`E`$ is the energy and $`\vec{p}`$ is the
20/// momentum.
21///
22/// # Examples
23/// ```
24/// use rustitude_core::prelude::*;
25///
26/// let vec_a = FourMomentum::new(1.3, 0.2, 0.3, 0.1);
27/// let vec_b = FourMomentum::new(4.2, 0.5, 0.4, 0.5);
28/// ```
29#[derive(Debug, Clone, PartialEq, Copy, Default)]
30pub struct FourMomentum<F: Field + 'static>(Vector4<F>);
31
32impl<F: Field> Eq for FourMomentum<F> {}
33
34impl<F: Field> Display for FourMomentum<F> {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        write!(
37            f,
38            "[{}, ({}, {}, {})]",
39            self.e(),
40            self.px(),
41            self.py(),
42            self.pz(),
43        )
44    }
45}
46
47impl<F: Field> FourMomentum<F> {
48    #[allow(clippy::missing_const_for_fn)]
49    pub fn new(e: F, px: F, py: F, pz: F) -> Self {
50        //! Create a new [`FourMomentum`] from energy and momentum components.
51        //!
52        //! Components are listed in the order $` (E, p_x, p_y, p_z) `$
53        Self(Vector4::new(e, px, py, pz))
54    }
55
56    /// Returns the energy of the given [`FourMomentum`].
57    #[allow(clippy::missing_const_for_fn)]
58    pub fn e(&self) -> F {
59        self.0[0]
60    }
61    /// Returns the momentum along the $`x`$-axis of the given [`FourMomentum`].
62    #[allow(clippy::missing_const_for_fn)]
63    pub fn px(&self) -> F {
64        self.0[1]
65    }
66    /// Returns the momentum along the $`y`$-axis of the given [`FourMomentum`].
67    #[allow(clippy::missing_const_for_fn)]
68    pub fn py(&self) -> F {
69        self.0[2]
70    }
71    /// Returns the momentum along the $`z`$-axis of the given [`FourMomentum`].
72    #[allow(clippy::missing_const_for_fn)]
73    pub fn pz(&self) -> F {
74        self.0[3]
75    }
76
77    /// Sets the energy of the given [`FourMomentum`].
78    pub fn set_e(&mut self, value: F) {
79        self.0[0] = value;
80    }
81    /// Sets the momentum along the $`x`$-axis of the given [`FourMomentum`].
82    pub fn set_px(&mut self, value: F) {
83        self.0[1] = value;
84    }
85    /// Sets the momentum along the $`x`$-axis of the given [`FourMomentum`].
86    pub fn set_py(&mut self, value: F) {
87        self.0[2] = value;
88    }
89    /// Sets the momentum along the $`x`$-axis of the given [`FourMomentum`].
90    pub fn set_pz(&mut self, value: F) {
91        self.0[3] = value;
92    }
93
94    /// Calculate the invariant $`m^2`$ for this [`FourMomentum`] instance.
95    ///
96    /// Calculates $` m^2 = E^2 - \vec{p}^2 `$
97    ///
98    /// # Examples
99    /// ```
100    /// use rustitude_core::prelude::*;
101    ///
102    /// let vec_a = FourMomentum::new(20.0, 1.0, 0.2, -0.1);
103    /// //assert_eq!(vec_a.m2(), 20.0 * 20.0 - (1.0 * 1.0 + 0.0 * 0.2 + (-0.1) * (-0.1)));
104    ///
105    /// ```
106    #[allow(clippy::suboptimal_flops)]
107    pub fn m2(&self) -> F {
108        F::powi(self.e(), 2) - F::powi(self.px(), 2) - F::powi(self.py(), 2) - F::powi(self.pz(), 2)
109    }
110
111    /// Calculate the invariant $`m`$ for this [`FourMomentum`] instance.
112    ///
113    /// Calculates $` m = \sqrt{E^2 - \vec{p}^2} `$
114    ///
115    /// # See Also:
116    ///
117    /// [`FourMomentum::m2`]
118    pub fn m(&self) -> F {
119        F::sqrt(self.m2())
120    }
121
122    /// Boosts an instance of [`FourMomentum`] along the $`\vec{\beta}`$
123    /// vector of another [`FourMomentum`].
124    ///
125    /// Calculates $`\mathbf{\Lambda} \cdot \mathbf{x}`$
126    ///
127    /// # Examples
128    /// ```
129    /// use rustitude_core::prelude::*;
130    ///
131    /// let vec_a = FourMomentum::new(20.0, 1.0, -3.2, 4.0);
132    /// let vec_a_COM = vec_a.boost_along(&vec_a);
133    /// assert!(f64::abs(vec_a_COM.px()) < 1e-7);
134    /// assert!(f64::abs(vec_a_COM.py()) < 1e-7);
135    /// assert!(f64::abs(vec_a_COM.pz()) < 1e-7);
136    /// ```
137    pub fn boost_along(&self, other: &Self) -> Self {
138        let m_boost = other.boost_matrix();
139        (m_boost * self.0).into()
140    }
141    /// Extract the 3-momentum as a [`nalgebra::Vector3<Field>`]
142    ///
143    /// # Examples
144    /// ```
145    /// use rustitude_core::prelude::*;
146    /// use nalgebra::Vector3;
147    ///
148    /// let vec_a = FourMomentum::new(20.0, 1.0, 0.2, -0.1);
149    /// assert_eq!(vec_a.momentum(), Vector3::new(1.0, 0.2, -0.1));
150    /// ```
151    pub fn momentum(&self) -> Vector3<F> {
152        Vector3::new(self.px(), self.py(), self.pz())
153    }
154
155    /// Returns the $`\cos(\theta)`$ of the momentum 3-vector.
156    pub fn costheta(&self) -> F {
157        let v = self.momentum();
158        let r = F::sqrt(v.x * v.x + v.y * v.y + v.z * v.z);
159        v.z / r
160    }
161
162    /// Returns the $`\cos(\theta)`$ of the momentum 3-vector.
163    ///
164    /// Alias for [`FourMomentum::costheta`].
165    pub fn theta_cos(&self) -> F {
166        self.costheta()
167    }
168
169    /// Returns the $`\theta`$ polar angle of the momentum 3-vector.
170    pub fn theta(&self) -> F {
171        F::acos(self.costheta())
172    }
173
174    /// Returns the $`\phi`$ azimuthal angle of the momentum 3-vector.
175    pub fn phi(&self) -> F {
176        let v = self.momentum();
177        F::atan2(v.y, v.x)
178    }
179
180    /// Construct the 3-vector $`\vec{\beta}`$ where
181    ///
182    /// $` \vec{\beta} = \frac{\vec{p}}{E} `$
183    pub fn beta3(&self) -> Vector3<F> {
184        self.momentum() / self.e()
185    }
186
187    /// Constructs the 3-vector normal to the 3-momentum
188    pub fn direction(&self) -> Vector3<F> {
189        let v = self.momentum();
190        let mag = F::sqrt(v.x * v.x + v.y * v.y + v.z * v.z);
191        v / mag
192    }
193
194    /// Construct the Lorentz boost matrix $`\mathbf{\Lambda}`$ where
195    ///
196    /// ```math
197    /// \mathbf{\Lambda} = \begin{pmatrix}
198    /// \gamma & -\gamma \beta_x & -\gamma \beta_y & -\gamma \beta_z \\
199    /// -\gamma \beta_x & 1 + (\gamma - 1) \frac{\beta_x^2}{\vec{\beta}^2} & (\gamma - 1) \frac{\beta_x \beta_y}{\vec{\beta}^2} & (\gamma - 1) \frac{\beta_x \beta_z}{\vec{\beta}^2} \\
200    /// -\gamma \beta_y & (\gamma - 1) \frac{\beta_y \beta_x}{\vec{\beta}^2} & 1 + (\gamma - 1) \frac{\beta_y^2}{\vec{\beta}^2} & (\gamma - 1) \frac{\beta_y \beta_z}{\vec{\beta}^2} \\
201    /// -\gamma \beta_z & (\gamma - 1) \frac{\beta_z \beta_x}{\vec{\beta}^2} & (\gamma - 1) \frac{\beta_z \beta_y}{\vec{\beta}^2} & 1 + (\gamma - 1) \frac{\beta_z^2}{\vec{\beta}^2}
202    /// \end{pmatrix}
203    /// ```
204    /// where
205    /// $`\vec{\beta} = \frac{\vec{p}}{E}`$ and $`\gamma = \frac{1}{\sqrt{1 - \vec{\beta}^2}}`$.
206    pub fn boost_matrix(&self) -> Matrix4<F> {
207        let b = self.beta3();
208        let b2 = b.dot(&b);
209        let g = F::one() / F::sqrt(F::one() - b2);
210        Matrix4::new(
211            g,
212            -g * b[0],
213            -g * b[1],
214            -g * b[2],
215            -g * b[0],
216            F::one() + (g - F::one()) * b[0] * b[0] / b2,
217            (g - F::one()) * b[0] * b[1] / b2,
218            (g - F::one()) * b[0] * b[2] / b2,
219            -g * b[1],
220            (g - F::one()) * b[1] * b[0] / b2,
221            F::one() + (g - F::one()) * b[1] * b[1] / b2,
222            (g - F::one()) * b[1] * b[2] / b2,
223            -g * b[2],
224            (g - F::one()) * b[2] * b[0] / b2,
225            (g - F::one()) * b[2] * b[1] / b2,
226            F::one() + (g - F::one()) * b[2] * b[2] / b2,
227        )
228    }
229}
230
231impl<F: Field> From<FourMomentum<F>> for Vector4<F> {
232    fn from(val: FourMomentum<F>) -> Self {
233        Self::new(val.e(), val.px(), val.py(), val.pz())
234    }
235}
236
237impl<F: Field> From<&FourMomentum<F>> for Vector4<F> {
238    fn from(val: &FourMomentum<F>) -> Self {
239        Self::new(val.e(), val.px(), val.py(), val.pz())
240    }
241}
242
243impl<F: Field> From<Vector4<F>> for FourMomentum<F> {
244    fn from(value: Vector4<F>) -> Self {
245        Self(value)
246    }
247}
248
249impl<F: Field> From<&Vector4<F>> for FourMomentum<F> {
250    fn from(value: &Vector4<F>) -> Self {
251        Self(*value)
252    }
253}
254
255impl<F: Field> From<Vec<F>> for FourMomentum<F> {
256    fn from(value: Vec<F>) -> Self {
257        Self::new(value[0], value[1], value[2], value[3])
258    }
259}
260
261impl<F: Field> From<&Vec<F>> for FourMomentum<F> {
262    fn from(value: &Vec<F>) -> Self {
263        Self::new(value[0], value[1], value[2], value[3])
264    }
265}
266
267impl<F: Field> Add for FourMomentum<F> {
268    type Output = Self;
269    fn add(self, rhs: Self) -> Self::Output {
270        Self(self.0 + rhs.0)
271    }
272}
273
274impl<F: Field> Add for &FourMomentum<F> {
275    type Output = <FourMomentum<F> as Add>::Output;
276    fn add(self, rhs: &FourMomentum<F>) -> Self::Output {
277        FourMomentum::add(*self, *rhs)
278    }
279}
280
281impl<F: Field> Sub for FourMomentum<F> {
282    type Output = Self;
283    fn sub(self, rhs: Self) -> Self::Output {
284        Self(self.0 - rhs.0)
285    }
286}
287
288impl<F: Field> Sub for &FourMomentum<F> {
289    type Output = <FourMomentum<F> as Sub>::Output;
290    fn sub(self, rhs: &FourMomentum<F>) -> Self::Output {
291        FourMomentum::sub(*self, *rhs)
292    }
293}
294
295impl<F: Field> std::iter::Sum<Self> for FourMomentum<F> {
296    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
297        iter.fold(Self::default(), |a, b| a + b)
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304    use crate::assert_is_close;
305    #[test]
306    fn test_set_components() {
307        let mut p = FourMomentum::<f64>::default();
308        p.set_e(1.0);
309        p.set_px(2.0);
310        p.set_py(3.0);
311        p.set_pz(4.0);
312        assert_is_close!(p.e(), 1.0, f64);
313        assert_is_close!(p.px(), 2.0, f64);
314        assert_is_close!(p.py(), 3.0, f64);
315        assert_is_close!(p.pz(), 4.0, f64);
316    }
317
318    #[test]
319    fn test_sum() {
320        let a = FourMomentum::new(0.1, 0.2, 0.3, 0.4);
321        let b = FourMomentum::new(1.0, 2.0, 3.0, 4.0);
322        let c = FourMomentum::new(10.0, 20.0, 30.0, 40.0);
323        let d: FourMomentum<f64> = [a, b, c].into_iter().sum();
324        assert_is_close!(d.e(), 11.1, f64);
325        assert_is_close!(d.px(), 22.2, f64);
326        assert_is_close!(d.py(), 33.3, f64);
327        assert_is_close!(d.pz(), 44.4, f64);
328    }
329
330    #[test]
331    fn test_ops() {
332        let a = FourMomentum::new(0.1, 0.2, 0.3, 0.4);
333        let b = FourMomentum::new(1.0, 2.0, 3.0, 4.0);
334        let c = a + b;
335        let d = b - a;
336        assert_is_close!(c.e(), 1.1, f64);
337        assert_is_close!(c.px(), 2.2, f64);
338        assert_is_close!(c.py(), 3.3, f64);
339        assert_is_close!(c.pz(), 4.4, f64);
340        assert_is_close!(d.e(), 0.9, f64);
341        assert_is_close!(d.px(), 1.8, f64);
342        assert_is_close!(d.py(), 2.7, f64);
343        assert_is_close!(d.pz(), 3.6, f64);
344    }
345}
rustitude_core/four_momentum.rs

rustitude_core/
four_momentum.rs
