laddu_physics/
vectors.rs

1use std::fmt::Display;
2
3use approx::{AbsDiffEq, RelativeEq};
4use auto_ops::{impl_op_ex, impl_op_ex_commutative};
5use laddu_expr::{Expr, P4Component, atan2, event_p4_component, vector};
6use nalgebra::{Vector3, Vector4};
7use serde::{Deserialize, Serialize};
8
9use crate::{LadduPhysicsError, LadduPhysicsResult};
10
11/// A vector with three components.
12///
13/// # Examples
14/// ```rust
15/// use laddu_physics::vectors::RealVec3;
16///
17/// let cross = RealVec3::x().cross(&RealVec3::y());
18/// assert_eq!(cross, RealVec3::z());
19/// ```
20#[derive(Copy, Clone, Debug, PartialEq, Serialize, Deserialize)]
21pub struct RealVec3 {
22    /// The x-component of the vector
23    pub x: f64,
24    /// The y-component of the vector
25    pub y: f64,
26    /// The z-component of the vector
27    pub z: f64,
28}
29
30impl Display for RealVec3 {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        write!(f, "[{:6.3}, {:6.3}, {:6.3}]", self.x, self.y, self.z)
33    }
34}
35
36impl AbsDiffEq for RealVec3 {
37    type Epsilon = <f64 as approx::AbsDiffEq>::Epsilon;
38
39    fn default_epsilon() -> Self::Epsilon {
40        f64::default_epsilon()
41    }
42
43    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
44        f64::abs_diff_eq(&self.x, &other.x, epsilon)
45            && f64::abs_diff_eq(&self.y, &other.y, epsilon)
46            && f64::abs_diff_eq(&self.z, &other.z, epsilon)
47    }
48}
49impl RelativeEq for RealVec3 {
50    fn default_max_relative() -> Self::Epsilon {
51        f64::default_max_relative()
52    }
53
54    fn relative_eq(
55        &self,
56        other: &Self,
57        epsilon: Self::Epsilon,
58        max_relative: Self::Epsilon,
59    ) -> bool {
60        f64::relative_eq(&self.x, &other.x, epsilon, max_relative)
61            && f64::relative_eq(&self.y, &other.y, epsilon, max_relative)
62            && f64::relative_eq(&self.z, &other.z, epsilon, max_relative)
63    }
64}
65
66impl From<RealVec3> for Vector3<f64> {
67    fn from(value: RealVec3) -> Self {
68        Vector3::new(value.x, value.y, value.z)
69    }
70}
71
72impl From<Vector3<f64>> for RealVec3 {
73    fn from(value: Vector3<f64>) -> Self {
74        RealVec3::new(value.x, value.y, value.z)
75    }
76}
77
78impl TryFrom<Vec<f64>> for RealVec3 {
79    type Error = LadduPhysicsError;
80
81    fn try_from(value: Vec<f64>) -> Result<Self, Self::Error> {
82        if value.len() != 3 {
83            return Err(LadduPhysicsError::custom(
84                "Attempted to convert Vec<f64> to RealVec3 for Vec with len != 3",
85            ));
86        }
87        Ok(Self {
88            x: value[0],
89            y: value[1],
90            z: value[2],
91        })
92    }
93}
94
95impl From<RealVec3> for Vec<f64> {
96    fn from(value: RealVec3) -> Self {
97        vec![value.x, value.y, value.z]
98    }
99}
100
101impl From<[f64; 3]> for RealVec3 {
102    fn from(value: [f64; 3]) -> Self {
103        Self {
104            x: value[0],
105            y: value[1],
106            z: value[2],
107        }
108    }
109}
110
111impl From<RealVec3> for [f64; 3] {
112    fn from(value: RealVec3) -> Self {
113        [value.x, value.y, value.z]
114    }
115}
116
117impl Default for RealVec3 {
118    fn default() -> Self {
119        RealVec3::zero()
120    }
121}
122
123impl RealVec3 {
124    /// Create a new 3-vector from its components
125    pub fn new(x: f64, y: f64, z: f64) -> Self {
126        RealVec3 { x, y, z }
127    }
128
129    /// Create a zero vector
130    pub const fn zero() -> Self {
131        RealVec3 {
132            x: 0.0,
133            y: 0.0,
134            z: 0.0,
135        }
136    }
137
138    /// Create a unit vector pointing in the x-direction
139    pub const fn x() -> Self {
140        RealVec3 {
141            x: 1.0,
142            y: 0.0,
143            z: 0.0,
144        }
145    }
146
147    /// Create a unit vector pointing in the y-direction
148    pub const fn y() -> Self {
149        RealVec3 {
150            x: 0.0,
151            y: 1.0,
152            z: 0.0,
153        }
154    }
155
156    /// Create a unit vector pointing in the z-direction
157    pub const fn z() -> Self {
158        RealVec3 {
159            x: 0.0,
160            y: 0.0,
161            z: 1.0,
162        }
163    }
164
165    /// Momentum in the x-direction
166    pub fn px(&self) -> f64 {
167        self.x
168    }
169
170    /// Momentum in the y-direction
171    pub fn py(&self) -> f64 {
172        self.y
173    }
174
175    /// Momentum in the z-direction
176    pub fn pz(&self) -> f64 {
177        self.z
178    }
179
180    /// Create a [`RealVec4`] with this vector as the 3-momentum and the given mass
181    pub fn with_mass(&self, mass: f64) -> RealVec4 {
182        let e = f64::sqrt(mass.powi(2) + self.mag2());
183        RealVec4::new(e, self.px(), self.py(), self.pz())
184    }
185
186    /// Create a [`RealVec4`] with this vector as the 3-momentum and the given energy
187    pub fn with_energy(&self, energy: f64) -> RealVec4 {
188        RealVec4::new(energy, self.px(), self.py(), self.pz())
189    }
190
191    /// Compute the dot product of this [`RealVec3`] and another
192    pub fn dot(&self, other: &RealVec3) -> f64 {
193        self.x * other.x + self.y * other.y + self.z * other.z
194    }
195
196    /// Compute the cross product of this [`RealVec3`] and another
197    pub fn cross(&self, other: &RealVec3) -> RealVec3 {
198        RealVec3::new(
199            self.y * other.z - other.y * self.z,
200            self.z * other.x - other.z * self.x,
201            self.x * other.y - other.x * self.y,
202        )
203    }
204
205    /// The magnitude of the vector
206    pub fn mag(&self) -> f64 {
207        f64::sqrt(self.mag2())
208    }
209
210    /// The squared magnitude of the vector
211    pub fn mag2(&self) -> f64 {
212        self.dot(self)
213    }
214
215    /// The cosine of the polar angle $`\theta`$
216    ///
217    /// # Errors
218    ///
219    /// Returns [`LadduPhysicsError`] when the vector has zero or invalid
220    /// magnitude.
221    pub fn costheta(&self) -> LadduPhysicsResult<f64> {
222        let mag = self.mag();
223        if mag <= 0.0 {
224            return Err(LadduPhysicsError::invalid_value(
225                "vector magnitude",
226                "positive when calculating cos(theta)",
227                mag,
228            ));
229        }
230        Ok(self.z / self.mag())
231    }
232
233    /// The polar angle $`\theta`$
234    ///
235    /// # Errors
236    ///
237    /// Returns [`LadduPhysicsError`] when the vector has zero or invalid
238    /// magnitude.
239    pub fn theta(&self) -> LadduPhysicsResult<f64> {
240        Ok(f64::acos(self.costheta()?))
241    }
242
243    /// The azimuthal angle $`\phi`$
244    pub fn phi(&self) -> f64 {
245        f64::atan2(self.y, self.x)
246    }
247
248    /// Create a unit vector in the same direction as this [`RealVec3`]
249    ///
250    /// # Errors
251    ///
252    /// Returns [`LadduPhysicsError`] when the vector has zero or invalid
253    /// magnitude.
254    pub fn unit(&self) -> LadduPhysicsResult<RealVec3> {
255        let mag = self.mag();
256        if mag <= 0.0 {
257            return Err(LadduPhysicsError::invalid_value(
258                "vector magnitude",
259                "positive when constructing unit vector",
260                mag,
261            ));
262        }
263        Ok(RealVec3::new(self.x / mag, self.y / mag, self.z / mag))
264    }
265}
266
267impl<'a> std::iter::Sum<&'a RealVec3> for RealVec3 {
268    fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
269        iter.fold(Self::zero(), |a, b| a + b)
270    }
271}
272impl std::iter::Sum<RealVec3> for RealVec3 {
273    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
274        iter.fold(Self::zero(), |a, b| a + b)
275    }
276}
277
278impl_op_ex!(+ |a: &RealVec3, b: &RealVec3| -> RealVec3 { RealVec3::new(a.x + b.x, a.y + b.y, a.z + b.z) });
279impl_op_ex!(-|a: &RealVec3, b: &RealVec3| -> RealVec3 {
280    RealVec3::new(a.x - b.x, a.y - b.y, a.z - b.z)
281});
282impl_op_ex!(-|a: &RealVec3| -> RealVec3 { RealVec3::new(-a.x, -a.y, -a.z) });
283impl_op_ex_commutative!(+ |a: &RealVec3, b: &f64| -> RealVec3 { RealVec3::new(a.x + b, a.y + b, a.z + b) });
284impl_op_ex_commutative!(-|a: &RealVec3, b: &f64| -> RealVec3 {
285    RealVec3::new(a.x - b, a.y - b, a.z - b)
286});
287impl_op_ex_commutative!(*|a: &RealVec3, b: &f64| -> RealVec3 {
288    RealVec3::new(a.x * b, a.y * b, a.z * b)
289});
290impl_op_ex!(/ |a: &RealVec3, b: &f64| -> RealVec3 { RealVec3::new(a.x / b, a.y / b, a.z / b) });
291
292/// A four-vector (Lorentz vector) stored in `(E, p_x, p_y, p_z)` order.
293///
294/// # Examples
295/// ```rust
296/// use laddu_physics::vectors::{RealVec3, RealVec4};
297///
298/// let momentum = RealVec3::new(1.0, 0.0, 0.0);
299/// let four_vector = momentum.with_mass(2.0);
300/// assert!((four_vector.m2() - 4.0).abs() < 1e-12);
301/// ```
302#[derive(Copy, Clone, Debug, PartialEq, Serialize, Deserialize)]
303pub struct RealVec4 {
304    /// Energy component.
305    pub e: f64,
306    /// Momentum in the x direction.
307    pub px: f64,
308    /// Momentum in the y direction.
309    pub py: f64,
310    /// Momentum in the z direction.
311    pub pz: f64,
312}
313
314impl Display for RealVec4 {
315    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
316        write!(
317            f,
318            "[{:6.3}; {:6.3}, {:6.3}, {:6.3}]",
319            self.e, self.px, self.py, self.pz
320        )
321    }
322}
323
324impl AbsDiffEq for RealVec4 {
325    type Epsilon = <f64 as approx::AbsDiffEq>::Epsilon;
326
327    fn default_epsilon() -> Self::Epsilon {
328        f64::default_epsilon()
329    }
330
331    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
332        f64::abs_diff_eq(&self.e, &other.e, epsilon)
333            && f64::abs_diff_eq(&self.px, &other.px, epsilon)
334            && f64::abs_diff_eq(&self.py, &other.py, epsilon)
335            && f64::abs_diff_eq(&self.pz, &other.pz, epsilon)
336    }
337}
338impl RelativeEq for RealVec4 {
339    fn default_max_relative() -> Self::Epsilon {
340        f64::default_max_relative()
341    }
342
343    fn relative_eq(
344        &self,
345        other: &Self,
346        epsilon: Self::Epsilon,
347        max_relative: Self::Epsilon,
348    ) -> bool {
349        f64::relative_eq(&self.e, &other.e, epsilon, max_relative)
350            && f64::relative_eq(&self.px, &other.px, epsilon, max_relative)
351            && f64::relative_eq(&self.py, &other.py, epsilon, max_relative)
352            && f64::relative_eq(&self.pz, &other.pz, epsilon, max_relative)
353    }
354}
355
356impl From<RealVec4> for Vector4<f64> {
357    fn from(value: RealVec4) -> Self {
358        Vector4::new(value.e, value.px, value.py, value.pz)
359    }
360}
361
362impl From<Vector4<f64>> for RealVec4 {
363    fn from(value: Vector4<f64>) -> Self {
364        RealVec4::new(value.x, value.y, value.z, value.w)
365    }
366}
367
368impl TryFrom<Vec<f64>> for RealVec4 {
369    type Error = LadduPhysicsError;
370
371    fn try_from(value: Vec<f64>) -> Result<Self, Self::Error> {
372        if value.len() != 4 {
373            return Err(LadduPhysicsError::custom(
374                "Attempted to convert Vec<f64> to RealVec4 for Vec with len != 4",
375            ));
376        }
377        Ok(Self {
378            e: value[0],
379            px: value[1],
380            py: value[2],
381            pz: value[3],
382        })
383    }
384}
385
386impl From<RealVec4> for Vec<f64> {
387    fn from(value: RealVec4) -> Self {
388        vec![value.e, value.px, value.py, value.pz]
389    }
390}
391
392impl From<[f64; 4]> for RealVec4 {
393    fn from(value: [f64; 4]) -> Self {
394        Self {
395            e: value[0],
396            px: value[1],
397            py: value[2],
398            pz: value[3],
399        }
400    }
401}
402
403impl From<RealVec4> for [f64; 4] {
404    fn from(value: RealVec4) -> Self {
405        [value.e, value.px, value.py, value.pz]
406    }
407}
408
409impl RealVec4 {
410    /// Create a four-vector in metric order `(E, p_x, p_y, p_z)`.
411    pub fn new(e: f64, px: f64, py: f64, pz: f64) -> Self {
412        RealVec4 { e, px, py, pz }
413    }
414
415    /// Momentum in the x-direction
416    pub fn px(&self) -> f64 {
417        self.px
418    }
419
420    /// Momentum in the y-direction
421    pub fn py(&self) -> f64 {
422        self.py
423    }
424
425    /// Momentum in the z-direction
426    pub fn pz(&self) -> f64 {
427        self.pz
428    }
429
430    /// The energy of the 4-vector
431    pub fn e(&self) -> f64 {
432        self.e
433    }
434
435    /// The 3-momentum
436    pub fn momentum(&self) -> RealVec3 {
437        self.vec3()
438    }
439
440    /// The $`\gamma`$ factor $`\frac{1}{\sqrt{1 - \beta^2}}`$.
441    ///
442    /// # Errors
443    ///
444    /// Returns [`LadduPhysicsError`] when the energy is not positive or the
445    /// resulting velocity is not subluminal.
446    pub fn gamma(&self) -> LadduPhysicsResult<f64> {
447        let beta = self.beta()?;
448        let b2 = beta.dot(&beta);
449        if b2 >= 1.0 {
450            return Err(LadduPhysicsError::invalid_value("|beta|^2", "< 1", b2));
451        }
452        Ok(1.0 / f64::sqrt(1.0 - b2))
453    }
454
455    /// The $`\vec{\beta}`$ vector $`\frac{\vec{p}}{E}`$.
456    ///
457    /// # Errors
458    ///
459    /// Returns [`LadduPhysicsError`] when the four-momentum energy is not
460    /// positive.
461    pub fn beta(&self) -> LadduPhysicsResult<RealVec3> {
462        let e = self.e();
463        if e <= 0.0 {
464            return Err(LadduPhysicsError::invalid_value(
465                "four-momentum energy",
466                "positive",
467                e,
468            ));
469        }
470        Ok(self.momentum() / e)
471    }
472
473    /// The invariant mass corresponding to this 4-momentum
474    ///
475    /// # Errors
476    ///
477    /// Returns [`LadduPhysicsError`] when the invariant mass squared is
478    /// non-finite or negative.
479    pub fn m(&self) -> LadduPhysicsResult<f64> {
480        self.mag()
481    }
482
483    #[inline(always)]
484    /// Return the invariant mass without checking for a spacelike vector.
485    pub fn m_unchecked(&self) -> f64 {
486        self.m2().sqrt()
487    }
488
489    /// Return the signed invariant mass.
490    ///
491    /// # Errors
492    ///
493    /// Returns [`LadduPhysicsError`] when the invariant mass squared is
494    /// non-finite.
495    pub fn signed_m(&self) -> LadduPhysicsResult<f64> {
496        self.signed_mag()
497    }
498
499    #[inline(always)]
500    /// Return the signed invariant mass without checking for finite components.
501    pub fn signed_m_unchecked(&self) -> f64 {
502        self.signed_mag_unchecked()
503    }
504
505    /// The squared invariant mass corresponding to this 4-momentum
506    pub fn m2(&self) -> f64 {
507        self.mag2()
508    }
509
510    /// Compute the Lorentz inner product with another four-vector.
511    pub fn dot(&self, other: &Self) -> f64 {
512        self.e * other.e - self.px * other.px - self.py * other.py - self.pz * other.pz
513    }
514
515    /// Pretty-prints the four-momentum.
516    pub fn to_p4_string(&self) -> String {
517        let mass = self
518            .m()
519            .map(|m| format!("{m:.5}"))
520            .unwrap_or_else(|_| format!("{:.5}i", (-self.m2()).sqrt()));
521        format!(
522            "[e = {:.5}; p = ({:.5}, {:.5}, {:.5}); m = {}]",
523            self.e(),
524            self.px(),
525            self.py(),
526            self.pz(),
527            mass
528        )
529    }
530
531    /// Alias for [`Self::m`] using the $`+---`$ metric.
532    ///
533    /// # Errors
534    ///
535    /// Returns [`LadduPhysicsError`] for non-finite or spacelike four-vectors,
536    /// where `m2 < 0`.
537    pub fn mag(&self) -> LadduPhysicsResult<f64> {
538        let mag2 = self.mag2();
539
540        if !mag2.is_finite() {
541            return Err(LadduPhysicsError::invalid_value(
542                "magnitude squared",
543                "finite",
544                mag2,
545            ));
546        }
547
548        if mag2 < 0.0 {
549            return Err(LadduPhysicsError::invalid_value(
550                "magnitude squared",
551                "nonnegative",
552                mag2,
553            ));
554        }
555
556        Ok(mag2.sqrt())
557    }
558
559    /// Signed invariant mass useful for diagnostics:
560    ///
561    /// - `sqrt(mag2)` for timelike/null vectors
562    /// - `-sqrt(-mag2)` for spacelike vectors
563    ///
564    /// # Errors
565    ///
566    /// Returns [`LadduPhysicsError`] when the invariant magnitude squared is
567    /// non-finite.
568    pub fn signed_mag(&self) -> LadduPhysicsResult<f64> {
569        let mag2 = self.mag2();
570
571        if !mag2.is_finite() {
572            return Err(LadduPhysicsError::invalid_value(
573                "magnitude squared",
574                "finite",
575                mag2,
576            ));
577        }
578
579        if mag2 >= 0.0 {
580            Ok(mag2.sqrt())
581        } else {
582            Ok(-(-mag2).sqrt())
583        }
584    }
585
586    #[inline(always)]
587    /// Return the signed magnitude without checking for finite components.
588    pub fn signed_mag_unchecked(&self) -> f64 {
589        let mag2 = self.mag2();
590        if mag2 >= 0.0 {
591            mag2.sqrt()
592        } else {
593            -(-mag2).sqrt()
594        }
595    }
596
597    /// Alias for [`Self::m2`], the squared invariant mass in the $`+---`$ metric.
598    pub fn mag2(&self) -> f64 {
599        self.e * self.e - (self.px * self.px + self.py * self.py + self.pz * self.pz)
600    }
601
602    /// Gives the vector boosted along a $`\vec{\beta}`$ vector.
603    pub fn boost(&self, beta: &RealVec3) -> Self {
604        let b2 = beta.dot(beta);
605        if b2 == 0.0 {
606            return *self;
607        }
608        let gamma = 1.0 / f64::sqrt(1.0 - b2);
609        let p3 = self.vec3() + beta * ((gamma - 1.0) * self.vec3().dot(beta) / b2 + gamma * self.e);
610        RealVec4::new(gamma * (self.e + beta.dot(&self.vec3())), p3.x, p3.y, p3.z)
611    }
612
613    /// The 3-vector contained in this 4-vector
614    pub fn vec3(&self) -> RealVec3 {
615        RealVec3 {
616            x: self.px,
617            y: self.py,
618            z: self.pz,
619        }
620    }
621}
622
623impl_op_ex!(+ |a: &RealVec4, b: &RealVec4| -> RealVec4 { RealVec4::new(a.e + b.e, a.px + b.px, a.py + b.py, a.pz + b.pz) });
624impl_op_ex!(-|a: &RealVec4, b: &RealVec4| -> RealVec4 {
625    RealVec4::new(a.e - b.e, a.px - b.px, a.py - b.py, a.pz - b.pz)
626});
627impl_op_ex!(-|a: &RealVec4| -> RealVec4 { RealVec4::new(a.e, -a.px, -a.py, -a.pz) });
628
629impl<'a> std::iter::Sum<&'a RealVec4> for RealVec4 {
630    fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
631        iter.fold(Self::new(0.0, 0.0, 0.0, 0.0), |a, b| a + b)
632    }
633}
634
635impl std::iter::Sum<RealVec4> for RealVec4 {
636    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
637        iter.fold(Self::new(0.0, 0.0, 0.0, 0.0), |a, b| a + b)
638    }
639}
640
641/// Expression-valued vector with three spatial components.
642#[derive(Clone, Debug, Serialize, Deserialize)]
643pub struct Vec3 {
644    /// Symbolic x component.
645    pub x: Expr,
646    /// Symbolic y component.
647    pub y: Expr,
648    /// Symbolic z component.
649    pub z: Expr,
650}
651
652impl Vec3 {
653    /// Construct a symbolic Cartesian three-vector.
654    pub fn new(x: impl Into<Expr>, y: impl Into<Expr>, z: impl Into<Expr>) -> Self {
655        Self {
656            x: x.into(),
657            y: y.into(),
658            z: z.into(),
659        }
660    }
661
662    /// Return the symbolic zero vector.
663    pub fn zero() -> Self {
664        Self::new(0.0, 0.0, 0.0)
665    }
666
667    /// Return the positive x-axis unit vector.
668    pub fn x() -> Self {
669        Self::new(1.0, 0.0, 0.0)
670    }
671
672    /// Return the positive y-axis unit vector.
673    pub fn y() -> Self {
674        Self::new(0.0, 1.0, 0.0)
675    }
676
677    /// Return the positive z-axis unit vector.
678    pub fn z() -> Self {
679        Self::new(0.0, 0.0, 1.0)
680    }
681
682    /// Read the spatial components of a named event four-vector.
683    pub fn event(prefix: &str) -> Self {
684        Self::new(
685            event_p4_component(prefix, P4Component::Px),
686            event_p4_component(prefix, P4Component::Py),
687            event_p4_component(prefix, P4Component::Pz),
688        )
689    }
690
691    /// Return the x momentum component.
692    pub fn px(&self) -> Expr {
693        self.x.clone()
694    }
695
696    /// Return the y momentum component.
697    pub fn py(&self) -> Expr {
698        self.y.clone()
699    }
700
701    /// Return the z momentum component.
702    pub fn pz(&self) -> Expr {
703        self.z.clone()
704    }
705
706    /// Compute the Euclidean inner product.
707    pub fn dot(&self, other: &Self) -> Expr {
708        &self.x * &other.x + &self.y * &other.y + &self.z * &other.z
709    }
710
711    /// Compute the Cartesian cross product.
712    pub fn cross(&self, other: &Self) -> Self {
713        Self::new(
714            &self.y * &other.z - &other.y * &self.z,
715            &self.z * &other.x - &other.z * &self.x,
716            &self.x * &other.y - &other.x * &self.y,
717        )
718    }
719
720    /// Return the squared Euclidean magnitude.
721    pub fn mag2(&self) -> Expr {
722        self.dot(self)
723    }
724
725    /// Return the Euclidean magnitude.
726    pub fn mag(&self) -> Expr {
727        self.mag2().sqrt()
728    }
729
730    /// Return the cosine of the polar angle.
731    pub fn costheta(&self) -> Expr {
732        &self.z / self.mag()
733    }
734
735    /// Return a vector normalized to unit magnitude.
736    pub fn unit(&self) -> Self {
737        self / &self.mag()
738    }
739
740    /// Return the azimuthal angle.
741    pub fn phi(&self) -> Expr {
742        atan2(self.py(), self.px())
743    }
744
745    /// Promote this momentum to a four-vector with the given invariant mass.
746    pub fn with_mass(&self, mass: impl Into<Expr>) -> Vec4 {
747        let mass = mass.into();
748        Vec4::new(
749            (mass.powi(2) + self.mag2()).sqrt(),
750            self.px(),
751            self.py(),
752            self.pz(),
753        )
754    }
755
756    /// Promote this momentum to a four-vector with the given energy.
757    pub fn with_energy(&self, energy: impl Into<Expr>) -> Vec4 {
758        Vec4::new(energy, self.px(), self.py(), self.pz())
759    }
760
761    /// Convert this vector to a vector-valued expression.
762    pub fn as_expr(&self) -> Expr {
763        vector([self.x.clone(), self.y.clone(), self.z.clone()])
764    }
765
766    fn scale(&self, scalar: impl Into<Expr>) -> Self {
767        let scalar = scalar.into();
768        Self::new(&self.x * &scalar, &self.y * &scalar, &self.z * scalar)
769    }
770}
771
772impl From<RealVec3> for Vec3 {
773    fn from(value: RealVec3) -> Self {
774        Self::new(value.x, value.y, value.z)
775    }
776}
777
778impl Default for Vec3 {
779    fn default() -> Self {
780        Self::zero()
781    }
782}
783
784impl_op_ex!(+ |a: &Vec3, b: &Vec3| -> Vec3 { Vec3::new(&a.x + &b.x, &a.y + &b.y, &a.z + &b.z) });
785impl_op_ex!(-|a: &Vec3, b: &Vec3| -> Vec3 { Vec3::new(&a.x - &b.x, &a.y - &b.y, &a.z - &b.z) });
786impl_op_ex!(-|a: &Vec3| -> Vec3 { Vec3::new(-&a.x, -&a.y, -&a.z) });
787impl_op_ex!(*|a: &Vec3, b: &Expr| -> Vec3 { a.scale(b) });
788impl_op_ex!(*|a: &Expr, b: &Vec3| -> Vec3 { b.scale(a) });
789impl_op_ex!(*|a: &Vec3, b: &f64| -> Vec3 { a.scale(b) });
790impl_op_ex!(*|a: &f64, b: &Vec3| -> Vec3 { b.scale(a) });
791impl_op_ex!(/ |a: &Vec3, b: &Expr| -> Vec3 {
792    Vec3::new(&a.x / b, &a.y / b, &a.z / b)
793});
794impl_op_ex!(/ |a: &Vec3, b: &f64| -> Vec3 {
795    Vec3::new(&a.x / b, &a.y / b, &a.z / b)
796});
797
798/// Expression-valued four-vector in `(E, p_x, p_y, p_z)` order with a `+---` metric.
799#[derive(Clone, Debug, Serialize, Deserialize)]
800pub struct Vec4 {
801    /// Energy component.
802    pub e: Expr,
803    /// Momentum in the x direction.
804    pub px: Expr,
805    /// Momentum in the y direction.
806    pub py: Expr,
807    /// Momentum in the z direction.
808    pub pz: Expr,
809}
810
811impl Vec4 {
812    /// Create a symbolic four-vector in metric order `(E, p_x, p_y, p_z)`.
813    pub fn new(
814        e: impl Into<Expr>,
815        px: impl Into<Expr>,
816        py: impl Into<Expr>,
817        pz: impl Into<Expr>,
818    ) -> Self {
819        Self {
820            e: e.into(),
821            px: px.into(),
822            py: py.into(),
823            pz: pz.into(),
824        }
825    }
826
827    /// Read a named event four-vector.
828    pub fn event(prefix: &str) -> Self {
829        Self::new(
830            event_p4_component(prefix, P4Component::E),
831            event_p4_component(prefix, P4Component::Px),
832            event_p4_component(prefix, P4Component::Py),
833            event_p4_component(prefix, P4Component::Pz),
834        )
835    }
836
837    /// Return the x momentum expression.
838    pub fn px(&self) -> Expr {
839        self.px.clone()
840    }
841
842    /// Return the y momentum expression.
843    pub fn py(&self) -> Expr {
844        self.py.clone()
845    }
846
847    /// Return the z momentum expression.
848    pub fn pz(&self) -> Expr {
849        self.pz.clone()
850    }
851
852    /// Return the energy expression.
853    pub fn e(&self) -> Expr {
854        self.e.clone()
855    }
856
857    /// Return the spatial momentum.
858    pub fn momentum(&self) -> Vec3 {
859        self.vec3()
860    }
861
862    /// Return the spatial three-vector.
863    pub fn vec3(&self) -> Vec3 {
864        Vec3::new(self.px(), self.py(), self.pz())
865    }
866
867    /// Return the three-velocity `p / E`.
868    pub fn beta(&self) -> Vec3 {
869        self.momentum() / self.e()
870    }
871
872    /// Return the Lorentz factor.
873    pub fn gamma(&self) -> Expr {
874        1.0 / (1.0 - self.beta().mag2()).sqrt()
875    }
876
877    /// Return the squared invariant mass.
878    pub fn m2(&self) -> Expr {
879        self.mag2()
880    }
881
882    /// Return the invariant mass.
883    pub fn m(&self) -> Expr {
884        self.mag()
885    }
886
887    /// Compute the Lorentz inner product with another symbolic four-vector.
888    pub fn dot(&self, other: &Self) -> Expr {
889        &self.e * &other.e - &self.px * &other.px - &self.py * &other.py - &self.pz * &other.pz
890    }
891
892    /// Alias for [`Self::m2`], the squared invariant mass.
893    pub fn mag2(&self) -> Expr {
894        self.dot(self)
895    }
896
897    /// Alias for [`Self::m`], the invariant mass.
898    pub fn mag(&self) -> Expr {
899        self.mag2().sqrt()
900    }
901
902    /// Apply a Lorentz boost by a three-velocity.
903    pub fn boost(&self, beta: &Vec3) -> Self {
904        let b2 = beta.dot(beta);
905        let gamma = (1.0 - &b2).sqrt();
906        let gamma = 1.0 / gamma;
907        let boost_factor = gamma.powi(2) / (&gamma + 1.0);
908        let p3 = self.vec3() + beta * ((boost_factor * self.vec3().dot(beta)) + &gamma * &self.e);
909        Self::new(gamma * (&self.e + beta.dot(&self.vec3())), p3.x, p3.y, p3.z)
910    }
911
912    /// Convert this four-vector to a vector-valued expression.
913    pub fn as_expr(&self) -> Expr {
914        vector([
915            self.e.clone(),
916            self.px.clone(),
917            self.py.clone(),
918            self.pz.clone(),
919        ])
920    }
921}
922
923impl From<RealVec4> for Vec4 {
924    fn from(value: RealVec4) -> Self {
925        Self::new(value.e, value.px, value.py, value.pz)
926    }
927}
928
929impl_op_ex!(+ |a: &Vec4, b: &Vec4| -> Vec4 {
930    Vec4::new(&a.e + &b.e, &a.px + &b.px, &a.py + &b.py, &a.pz + &b.pz)
931});
932impl_op_ex!(-|a: &Vec4, b: &Vec4| -> Vec4 {
933    Vec4::new(&a.e - &b.e, &a.px - &b.px, &a.py - &b.py, &a.pz - &b.pz)
934});
935impl_op_ex!(-|a: &Vec4| -> Vec4 { Vec4::new(-&a.e, -&a.px, -&a.py, -&a.pz) });
936
937#[cfg(test)]
938mod tests {
939    use approx::{assert_abs_diff_eq, assert_relative_eq};
940    use laddu_compile::CompiledModel;
941    use laddu_runtime::CpuBackend;
942    use nalgebra::{Vector3, Vector4};
943    use num::complex::Complex64;
944
945    use super::*;
946
947    fn evaluate(expr: laddu_expr::Expr) -> Complex64 {
948        let model = CompiledModel::from_expr(&expr).unwrap();
949        let params = model.params().default_values();
950        CpuBackend.prepare(&model).evaluate(&params).unwrap()
951    }
952
953    #[test]
954    fn test_display() {
955        let v3 = RealVec3::new(1.2341, -2.3452, 3.4563);
956        assert_eq!(format!("{}", v3), "[ 1.234, -2.345,  3.456]");
957        let v4 = RealVec4::new(4.5674, 1.2341, -2.3452, 3.4563);
958        assert_eq!(format!("{}", v4), "[ 4.567;  1.234, -2.345,  3.456]");
959    }
960
961    #[test]
962    fn test_vec_vector_conversion() {
963        let v = RealVec3::new(1.0, 2.0, 3.0);
964        let vector3: Vec<f64> = v.into();
965        assert_eq!(vector3[0], 1.0);
966        assert_eq!(vector3[1], 2.0);
967        assert_eq!(vector3[2], 3.0);
968
969        let v_from_vec: RealVec3 = vector3.try_into().unwrap();
970        assert_eq!(v_from_vec, v);
971
972        let v = RealVec4::new(1.0, 2.0, 3.0, 4.0);
973        let vector4: Vec<f64> = v.into();
974        assert_eq!(vector4[0], 1.0);
975        assert_eq!(vector4[1], 2.0);
976        assert_eq!(vector4[2], 3.0);
977        assert_eq!(vector4[3], 4.0);
978
979        let v_from_vec: RealVec4 = vector4.try_into().unwrap();
980        assert_eq!(v_from_vec, v);
981    }
982
983    #[test]
984    fn test_vec_array_conversion() {
985        let arr = [1.0, 2.0, 3.0];
986        let v: RealVec3 = arr.into();
987        assert_eq!(v, RealVec3::new(1.0, 2.0, 3.0));
988
989        let back_to_array: [f64; 3] = v.into();
990        assert_eq!(back_to_array, arr);
991
992        let arr = [1.0, 2.0, 3.0, 4.0];
993        let v: RealVec4 = arr.into();
994        assert_eq!(v, RealVec4::new(1.0, 2.0, 3.0, 4.0));
995
996        let back_to_array: [f64; 4] = v.into();
997        assert_eq!(back_to_array, arr);
998    }
999
1000    #[test]
1001    fn test_vec_nalgebra_conversion() {
1002        let v = RealVec3::new(1.0, 2.0, 3.0);
1003        let vector3: Vector3<f64> = v.into();
1004        assert_eq!(vector3.x, 1.0);
1005        assert_eq!(vector3.y, 2.0);
1006        assert_eq!(vector3.z, 3.0);
1007
1008        let v_from_vec: RealVec3 = vector3.into();
1009        assert_eq!(v_from_vec, v);
1010
1011        let v = RealVec4::new(1.0, 2.0, 3.0, 4.0);
1012        let vector4: Vector4<f64> = v.into();
1013        assert_eq!(vector4.x, 1.0);
1014        assert_eq!(vector4.y, 2.0);
1015        assert_eq!(vector4.z, 3.0);
1016        assert_eq!(vector4.w, 4.0);
1017
1018        let v_from_vec: RealVec4 = vector4.into();
1019        assert_eq!(v_from_vec, v);
1020    }
1021
1022    #[test]
1023    fn test_vec_sums() {
1024        let vectors = [RealVec3::new(1.0, 2.0, 3.0), RealVec3::new(4.0, 5.0, 6.0)];
1025        let sum: RealVec3 = vectors.iter().sum();
1026        assert_eq!(sum, RealVec3::new(5.0, 7.0, 9.0));
1027        let sum: RealVec3 = vectors.into_iter().sum();
1028        assert_eq!(sum, RealVec3::new(5.0, 7.0, 9.0));
1029
1030        let vectors = [
1031            RealVec4::new(1.0, 2.0, 3.0, 4.0),
1032            RealVec4::new(4.0, 5.0, 6.0, 7.0),
1033        ];
1034        let sum: RealVec4 = vectors.iter().sum();
1035        assert_eq!(sum, RealVec4::new(5.0, 7.0, 9.0, 11.0));
1036        let sum: RealVec4 = vectors.into_iter().sum();
1037        assert_eq!(sum, RealVec4::new(5.0, 7.0, 9.0, 11.0));
1038    }
1039
1040    #[test]
1041    fn test_three_to_four_momentum_conversion() {
1042        let p3 = RealVec3::new(1.0, 2.0, 3.0);
1043        let target_p4 = RealVec4::new(10.0, 1.0, 2.0, 3.0);
1044        let p4_from_mass = p3.with_mass(target_p4.m().unwrap());
1045        assert_eq!(target_p4.e(), p4_from_mass.e());
1046        assert_eq!(target_p4.px(), p4_from_mass.px());
1047        assert_eq!(target_p4.py(), p4_from_mass.py());
1048        assert_eq!(target_p4.pz(), p4_from_mass.pz());
1049        let p4_from_energy = p3.with_energy(target_p4.e());
1050        assert_eq!(target_p4.e(), p4_from_energy.e());
1051        assert_eq!(target_p4.px(), p4_from_energy.px());
1052        assert_eq!(target_p4.py(), p4_from_energy.py());
1053        assert_eq!(target_p4.pz(), p4_from_energy.pz());
1054    }
1055
1056    #[test]
1057    fn test_four_momentum_basics() {
1058        let p = RealVec4::new(10.0, 3.0, 4.0, 5.0);
1059        assert_eq!(p.e(), 10.0);
1060        assert_eq!(p.px(), 3.0);
1061        assert_eq!(p.py(), 4.0);
1062        assert_eq!(p.pz(), 5.0);
1063        assert_eq!(p.momentum().px(), 3.0);
1064        assert_eq!(p.momentum().py(), 4.0);
1065        assert_eq!(p.momentum().pz(), 5.0);
1066        assert_relative_eq!(p.beta().unwrap().x, 0.3);
1067        assert_relative_eq!(p.beta().unwrap().y, 0.4);
1068        assert_relative_eq!(p.beta().unwrap().z, 0.5);
1069        assert_relative_eq!(p.m2(), 50.0);
1070        assert_relative_eq!(p.m().unwrap(), f64::sqrt(50.0));
1071        assert_eq!(
1072            p.to_p4_string().to_string(),
1073            "[e = 10.00000; p = (3.00000, 4.00000, 5.00000); m = 7.07107]"
1074        );
1075        assert_relative_eq!(RealVec3::x().x, 1.0);
1076        assert_relative_eq!(RealVec3::x().y, 0.0);
1077        assert_relative_eq!(RealVec3::x().z, 0.0);
1078        assert_relative_eq!(RealVec3::y().x, 0.0);
1079        assert_relative_eq!(RealVec3::y().y, 1.0);
1080        assert_relative_eq!(RealVec3::y().z, 0.0);
1081        assert_relative_eq!(RealVec3::z().x, 0.0);
1082        assert_relative_eq!(RealVec3::z().y, 0.0);
1083        assert_relative_eq!(RealVec3::z().z, 1.0);
1084        assert_relative_eq!(RealVec3::default().x, 0.0);
1085        assert_relative_eq!(RealVec3::default().y, 0.0);
1086        assert_relative_eq!(RealVec3::default().z, 0.0);
1087    }
1088
1089    #[test]
1090    fn test_three_momentum_basics() {
1091        let p = RealVec4::new(10.0, 3.0, 4.0, 5.0);
1092        let q = RealVec4::new(0.0, 1.2, -3.4, 7.6);
1093        let p3_view = p.momentum();
1094        let q3_view = q.momentum();
1095        assert_eq!(p3_view.px(), 3.0);
1096        assert_eq!(p3_view.py(), 4.0);
1097        assert_eq!(p3_view.pz(), 5.0);
1098        assert_relative_eq!(p3_view.mag2(), 50.0);
1099        assert_relative_eq!(p3_view.mag(), f64::sqrt(50.0));
1100        assert_relative_eq!(p3_view.costheta().unwrap(), 5.0 / f64::sqrt(50.0));
1101        assert_relative_eq!(p3_view.theta().unwrap(), f64::acos(5.0 / f64::sqrt(50.0)));
1102        assert_relative_eq!(p3_view.phi(), f64::atan2(4.0, 3.0));
1103        assert_relative_eq!(
1104            p3_view.unit().unwrap(),
1105            RealVec3::new(
1106                3.0 / f64::sqrt(50.0),
1107                4.0 / f64::sqrt(50.0),
1108                5.0 / f64::sqrt(50.0)
1109            )
1110        );
1111        assert_relative_eq!(p3_view.cross(&q3_view), RealVec3::new(47.4, -16.8, -15.0));
1112    }
1113
1114    #[test]
1115    fn test_vec_equality() {
1116        let p = RealVec3::new(1.1, 2.2, 3.3);
1117        let p2 = RealVec3::new(1.1 * 2.0, 2.2 * 2.0, 3.3 * 2.0);
1118        assert_abs_diff_eq!(p * 2.0, p2);
1119        assert_relative_eq!(p * 2.0, p2);
1120    }
1121
1122    #[test]
1123    fn test_boost_com() {
1124        let p = RealVec4::new(10.0, 3.0, 4.0, 5.0);
1125        let zero = p.boost(&-p.beta().unwrap()).momentum();
1126        assert_relative_eq!(zero, RealVec3::zero());
1127    }
1128
1129    #[test]
1130    fn test_boost() {
1131        let p0 = RealVec4::new(1.0, 0.0, 0.0, 0.0);
1132        assert_relative_eq!(p0.gamma().unwrap(), 1.0);
1133        let p0 = RealVec4::new(1.0, f64::sqrt(3.0) / 2.0, 0.0, 0.0);
1134        assert_relative_eq!(p0.gamma().unwrap(), 2.0);
1135        let p1 = RealVec4::new(10.0, 3.0, 4.0, 5.0);
1136        let p2 = RealVec4::new(9.0, 3.4, 2.3, 1.2);
1137        let p1_boosted = p1.boost(&-p2.beta().unwrap());
1138        assert_relative_eq!(p1_boosted.e(), 8.157632144622882);
1139        assert_relative_eq!(p1_boosted.px(), -0.6489200627053444);
1140        assert_relative_eq!(p1_boosted.py(), 1.5316128987581492);
1141        assert_relative_eq!(p1_boosted.pz(), 3.712145860221643);
1142    }
1143
1144    #[test]
1145    fn expression_vectors_build_and_evaluate_scalar_observables() {
1146        let p = Vec4::new(10.0, 3.0, 4.0, 5.0);
1147        assert_eq!(evaluate(p.m2()), Complex64::from(50.0));
1148        assert_eq!(evaluate(p.momentum().mag2()), Complex64::from(50.0));
1149
1150        let a = Vec3::new(1.0, 2.0, 3.0);
1151        let b = Vec3::new(4.0, 5.0, 6.0);
1152        assert_eq!(evaluate(a.dot(&b)), Complex64::from(32.0));
1153        assert_eq!(evaluate(a.cross(&b).z), Complex64::from(-3.0));
1154    }
1155}
laddu_physics/vectors.rs

laddu_physics/
vectors.rs
