lindhard/ion/crystal/divergence.rs
1//! Beam divergence: the spread of incidence directions about the nominal
2//! beam direction, sampled on the per-particle random stream.
3//!
4//! A divergence model gives each particle a deflection from the central
5//! direction: a polar angle `ϑ` from it and an azimuth `ψ` about it. The
6//! azimuth is uniform on `[0, 2π)` for both models, and the polar angle has:
7//!
8//! * [`Divergence::Gaussian`] with standard deviation `σ`: the two angular
9//! deviations `(ϑ cos ψ, ϑ sin ψ)` in orthogonal planes through the central
10//! direction are independent normal variables of mean 0 and standard
11//! deviation `σ`. Changing the bivariate normal density
12//! `exp(-(ξ² + η²) / 2σ²) / (2π σ²) dξ dη` to polar coordinates gives
13//! `(1 / 2π) dψ · (ϑ / σ²) exp(-ϑ² / 2σ²) dϑ`: `ψ` uniform and independent
14//! of `ϑ`, and `ϑ` Rayleigh distributed with CDF
15//! `F(ϑ) = 1 - exp(-ϑ² / 2σ²)`. It is drawn by inversion,
16//! `ϑ = σ sqrt(-2 ln(1 - U))`. The angular deviations are treated as plane
17//! angles, which is the small-angle reading of a Gaussian divergence; a
18//! draw with `ϑ > π` has probability `exp(-π² / 2σ²)` and is still
19//! applied as a rotation by `ϑ`.
20//! * [`Divergence::UniformCone`] with half-angle `α`: directions uniform in
21//! solid angle inside the cone `ϑ <= α`. The solid-angle element is
22//! `sin ϑ dϑ dψ`, so `F(ϑ) = (1 - cos ϑ) / (1 - cos α)`. Using
23//! `1 - cos x = 2 sin²(x / 2)`, inversion gives
24//! `ϑ = 2 asin(sqrt(U) sin(α / 2))`, which keeps full relative precision
25//! for small angles.
26//!
27//! # Random numbers
28//!
29//! Draws come from the caller's generator, which for a transport run is the
30//! particle's own stream [`crate::rng::stream`]`(seed, index)`, so the
31//! result depends only on `(seed, index)` and not on the thread count. Each
32//! sample takes exactly two 64-bit draws, in this order: `U` for the polar
33//! angle, then `V` for the azimuth `ψ = 2π V`. Uniforms are the top 53 bits
34//! of a draw scaled to `[0, 1)`, as in [`crate::ion::bca`].
35//! [`Divergence::None`] takes no draws, so a beam without divergence leaves
36//! the stream untouched.
37//!
38//! The deflected direction is the central direction rotated by `ϑ` at
39//! azimuth `ψ` with [`crate::ion::bca::kinematics::rotate`]; its azimuth
40//! zero is a fixed, documented frame, which is immaterial because `ψ` is
41//! uniform.
42
43use std::f64::consts::PI;
44
45use crate::rng::uniform;
46use rand_core::Rng;
47
48use super::CrystalError;
49use crate::ion::bca::kinematics::rotate;
50
51/// Beam divergence model.
52#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize)]
53pub enum Divergence {
54 /// No divergence: every particle gets the central direction.
55 None,
56 /// Gaussian divergence: independent normal angular deviations with
57 /// standard deviation `sigma_rad` in two orthogonal planes.
58 Gaussian {
59 /// Standard deviation per plane, radians (finite, `>= 0`).
60 sigma_rad: f64,
61 },
62 /// Directions uniform in solid angle within a cone.
63 UniformCone {
64 /// Cone half-angle, radians, in `[0, π]`.
65 half_angle_rad: f64,
66 },
67}
68
69impl Divergence {
70 /// Check the parameters.
71 pub fn validate(&self) -> Result<(), CrystalError> {
72 match *self {
73 Self::None => Ok(()),
74 Self::Gaussian { sigma_rad: s } if s.is_finite() && s >= 0.0 => Ok(()),
75 Self::Gaussian { sigma_rad } => Err(CrystalError::InvalidAngle {
76 name: "divergence sigma",
77 value: sigma_rad,
78 why: "must be finite and non-negative",
79 }),
80 Self::UniformCone { half_angle_rad: a } if (0.0..=PI).contains(&a) => Ok(()),
81 Self::UniformCone { half_angle_rad } => Err(CrystalError::InvalidAngle {
82 name: "divergence half-angle",
83 value: half_angle_rad,
84 why: "must be in [0, pi]",
85 }),
86 }
87 }
88
89 /// Target CDF of the polar deflection `ϑ` (see the module docs).
90 /// [`Divergence::None`] is a point mass at 0.
91 pub fn polar_cdf(&self, theta: f64) -> f64 {
92 if theta < 0.0 {
93 return 0.0;
94 }
95 match *self {
96 Self::None => 1.0,
97 Self::Gaussian { sigma_rad: s } => {
98 if s == 0.0 {
99 1.0
100 } else {
101 -(-theta * theta / (2.0 * s * s)).exp_m1()
102 }
103 }
104 Self::UniformCone { half_angle_rad: a } => {
105 if theta >= a {
106 1.0
107 } else {
108 let h = (0.5 * theta).sin() / (0.5 * a).sin();
109 h * h
110 }
111 }
112 }
113 }
114
115 /// Draw a deflection `(ϑ, ψ)`: polar angle from the central direction
116 /// and azimuth about it, radians. Takes two draws (none for
117 /// [`Divergence::None`]); the parameters must be valid
118 /// ([`Self::validate`]).
119 pub fn sample_deflection<R: Rng + ?Sized>(&self, rng: &mut R) -> (f64, f64) {
120 let polar = match *self {
121 Self::None => return (0.0, 0.0),
122 Self::Gaussian { sigma_rad } => {
123 let u = uniform(rng);
124 // 1 - u is in (0, 1], so the logarithm is finite.
125 sigma_rad * (-2.0 * (-u).ln_1p()).sqrt()
126 }
127 Self::UniformCone { half_angle_rad } => {
128 let u = uniform(rng);
129 2.0 * (u.sqrt() * (0.5 * half_angle_rad).sin()).asin()
130 }
131 };
132 let azimuth = 2.0 * PI * uniform(rng);
133 (polar, azimuth)
134 }
135
136 /// Deflect the unit vector `central` by a sampled divergence. For
137 /// [`Divergence::None`] this returns `central` unchanged and draws
138 /// nothing.
139 pub fn sample_direction<R: Rng + ?Sized>(&self, central: [f64; 3], rng: &mut R) -> [f64; 3] {
140 match self {
141 Self::None => central,
142 _ => {
143 let (polar, azimuth) = self.sample_deflection(rng);
144 rotate(central, polar, azimuth)
145 }
146 }
147 }
148}
149
150#[cfg(test)]
151mod tests {
152 use super::*;
153 use crate::rng::stream;
154
155 #[test]
156 fn none_draws_nothing() {
157 let mut a = stream(3, 7);
158 let d = Divergence::None.sample_direction([1.0, 0.0, 0.0], &mut a);
159 assert_eq!(d, [1.0, 0.0, 0.0]);
160 assert_eq!(a.next_u64(), stream(3, 7).next_u64());
161 }
162
163 #[test]
164 fn each_sample_takes_two_draws() {
165 for div in [
166 Divergence::Gaussian { sigma_rad: 0.01 },
167 Divergence::UniformCone {
168 half_angle_rad: 0.02,
169 },
170 ] {
171 let mut a = stream(3, 7);
172 div.sample_deflection(&mut a);
173 let mut b = stream(3, 7);
174 b.next_u64();
175 b.next_u64();
176 assert_eq!(a.next_u64(), b.next_u64());
177 }
178 }
179
180 #[test]
181 fn cone_stays_inside_and_zero_width_is_exact() {
182 let div = Divergence::UniformCone {
183 half_angle_rad: 0.05,
184 };
185 let mut rng = stream(1, 0);
186 for _ in 0..1000 {
187 let (p, a) = div.sample_deflection(&mut rng);
188 assert!((0.0..=0.05).contains(&p));
189 assert!((0.0..2.0 * PI).contains(&a));
190 }
191 let (p, _) = Divergence::Gaussian { sigma_rad: 0.0 }.sample_deflection(&mut rng);
192 assert_eq!(p, 0.0);
193 }
194
195 #[test]
196 fn cdfs_are_normalised() {
197 let g = Divergence::Gaussian { sigma_rad: 0.01 };
198 assert_eq!(g.polar_cdf(0.0), 0.0);
199 assert!((g.polar_cdf(1.0) - 1.0).abs() < 1e-15);
200 let c = Divergence::UniformCone {
201 half_angle_rad: 0.3,
202 };
203 assert_eq!(c.polar_cdf(0.0), 0.0);
204 assert_eq!(c.polar_cdf(0.3), 1.0);
205 // Half the solid angle: 1 - cos(t) = (1 - cos(0.3)) / 2.
206 let t = (1.0 - 0.5 * (1.0 - 0.3f64.cos())).acos();
207 assert!((c.polar_cdf(t) - 0.5).abs() < 1e-12);
208 }
209
210 #[test]
211 fn validation() {
212 assert!(Divergence::None.validate().is_ok());
213 assert!(Divergence::Gaussian { sigma_rad: 0.0 }.validate().is_ok());
214 assert!(Divergence::Gaussian { sigma_rad: -1.0 }.validate().is_err());
215 assert!(Divergence::Gaussian {
216 sigma_rad: f64::NAN
217 }
218 .validate()
219 .is_err());
220 assert!(Divergence::UniformCone { half_angle_rad: PI }
221 .validate()
222 .is_ok());
223 assert!(Divergence::UniformCone {
224 half_angle_rad: 4.0
225 }
226 .validate()
227 .is_err());
228 }
229}