Skip to main content

lindhard/tally/
exposure.rs

1//! Exposure statistics: the distribution of the energy a small volume
2//! receives from a finite number of incident electrons.
3//!
4//! A tally reports means over many histories. A small volume exposed to a few
5//! electrons receives an energy that fluctuates, both from the number of
6//! electrons and from the wide spread of the energy each one leaves there.
7//!
8//! # The per-electron distribution `X`
9//!
10//! [`ExposureAccumulator`] follows each primary history (the primary and all
11//! its secondaries) and sums the energy it deposits inside every configured
12//! [`ExposureVolume`]. That sum, `X` (eV), is binned into an
13//! [`ExposureHistogram`]: an exact atom at `X = 0` (most histories deposit
14//! nothing in a small volume), `bins` equal-width bins of width
15//! `bin_width_ev` on `(0, bins * bin_width_ev]` (a bin is closed on its
16//! upper edge: bin `j`, `j = 1..=bins`, holds `(j-1) w < X <= j w`), and the
17//! running sums `Σ X^k` for `k = 1..=4`. A history with `X` beyond the last
18//! edge is counted in `overflow_count` and clamped into the last bin; the
19//! sums stay exact.
20//!
21//! # Beam profile and the lateral-uniformity assumption
22//!
23//! The beam is a [`BeamProfile`]: a point, a Gaussian of standard deviation
24//! `sigma` per lateral axis, or a uniform disc. The landing offset `d` of each
25//! electron is drawn from the counter-based stream
26//! `rng::stream(seed, history index)` at the start of its history, and the
27//! deposits are taken relative to the landing point (`pos - start - d`
28//! compared with the volume centre, i.e. the volume is shifted by `-d`).
29//! The stored `X` is therefore already averaged over the beam and the
30//! electrons stay independent and identically distributed.
31//!
32//! **This is valid only while the target is laterally uniform** (a layered
33//! stack, as everywhere in the electron engine): then the history of an
34//! electron landing at `d` is the history of one landing at the axis, shifted
35//! by `d`. Laterally varying targets and beams are out of scope here.
36//!
37//! # Compound Poisson statistics
38//!
39//! With `N ~ Poisson(λ)` electrons incident (`λ` the mean number), the energy
40//! in the volume is `S = X_1 + … + X_N`, `X_i` independent copies of `X`.
41//! Conditioning on `N` (law of total expectation and variance) gives
42//!
43//! ```text
44//! E[S]   = E[N] E[X]                    = λ E[X]
45//! Var[S] = E[N] Var[X] + Var[N] E[X]²   = λ E[X²]
46//! ```
47//!
48//! because `E[N] = Var[N] = λ`. These are the textbook compound-Poisson
49//! results (see W. Feller, *An Introduction to Probability Theory and Its
50//! Applications*, vol. 1, 3rd ed., Wiley 1968, the chapter on compound
51//! distributions; and N. R. Campbell, Proc. Camb. Phil. Soc. 15, 117 (1909),
52//! for the mean and variance of such shot-noise sums). The derivation above
53//! is self-contained; the chapter-level citation was not checked against the
54//! page. The moments use the exact stored sums, with no binning error.
55//!
56//! The full distribution of `S` uses Poisson thinning: the histories that
57//! deposit something number `K ~ Poisson(λ (1 - p0))`, `p0` the mass at
58//! `X = 0`, and `S` is the sum of `K` draws from the conditional nonzero
59//! distribution `f`. So
60//!
61//! ```text
62//! P(S = s) = Σ_k Poisson(k; λ(1-p0)) · f^{*k}(s)
63//! ```
64//!
65//! with `f^{*k}` the `k`-fold convolution of the binned `f`, built
66//! iteratively, and the `k = 0` term the exact atom at `S = 0`. The sum is
67//! truncated where the Poisson tail is below `tail_tolerance`; the dropped
68//! mass is reported.
69//!
70//! **Binning bound.** Representing each bin by its lower edge, centre or
71//! upper edge gives three distributions of `S`; the first is stochastically
72//! smaller and the last larger than the true one (up to the clamped
73//! overflow, reported). Every quantile and threshold probability is reported
74//! as `low`, `central` (bin centres) and `high`: the true value lies
75//! between `low` and `high` for quantiles and for `P(S > b)`, and the same
76//! holds for `P(S < a)` with the roles of the two ends swapped (the field
77//! names refer to the representation, see [`Bracket`]). The width of the
78//! bracket on a quantile is at most `K_max * bin_width_ev`, `K_max` the
79//! largest number of depositing electrons kept.
80//!
81//! Units: energies in eV, lengths in metres.
82
83use crate::electron::transport::{disc_offset, finite_positive, gaussian_offset};
84use crate::rng::{stream, uniform};
85use serde::{Deserialize, Serialize};
86
87/// Largest `bins * k_max * bins` multiply-add count the convolution is
88/// allowed (about a second).
89const MAX_CONVOLUTION_WORK: f64 = 2.0e9;
90
91/// Errors of the exposure configuration.
92#[derive(Debug, Clone, PartialEq, thiserror::Error)]
93pub enum ExposureError {
94    /// A scalar of the configuration is out of range.
95    #[error("exposure: {0}")]
96    Invalid(String),
97}
98
99fn invalid<T>(msg: impl Into<String>) -> Result<T, ExposureError> {
100    Err(ExposureError::Invalid(msg.into()))
101}
102
103/// Lateral profile of the beam.
104#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
105#[serde(tag = "kind", rename_all = "kebab-case")]
106pub enum BeamProfile {
107    /// All electrons land on the axis.
108    Point,
109    /// Landing offset Gaussian with standard deviation `sigma_m` per axis.
110    Gaussian {
111        /// Standard deviation per lateral axis, m.
112        sigma_m: f64,
113    },
114    /// Landing offset uniform on a disc.
115    UniformDisc {
116        /// Disc radius, m.
117        radius_m: f64,
118    },
119}
120
121impl BeamProfile {
122    /// The landing offset `(y, z)` of history `index`, drawn from
123    /// `rng::stream(seed, index)`.
124    pub fn offset(&self, seed: u64, index: u64) -> [f64; 2] {
125        match *self {
126            BeamProfile::Point => [0.0, 0.0],
127            BeamProfile::Gaussian { sigma_m } => {
128                let mut rng = stream(seed, index);
129                let (u1, u2) = (uniform(&mut rng), uniform(&mut rng));
130                gaussian_offset(sigma_m, u1, u2)
131            }
132            BeamProfile::UniformDisc { radius_m } => {
133                let mut rng = stream(seed, index);
134                let (u1, u2) = (uniform(&mut rng), uniform(&mut rng));
135                disc_offset(radius_m, u1, u2)
136            }
137        }
138    }
139
140    fn validate(&self) -> Result<(), ExposureError> {
141        match *self {
142            BeamProfile::Point => Ok(()),
143            BeamProfile::Gaussian { sigma_m } if finite_positive(sigma_m) => Ok(()),
144            BeamProfile::UniformDisc { radius_m } if finite_positive(radius_m) => Ok(()),
145            _ => invalid("beam sigma or radius must be positive and finite"),
146        }
147    }
148}
149
150/// Lateral shape of a target volume, centred on the volume's offset.
151#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
152#[serde(tag = "kind", rename_all = "kebab-case")]
153pub enum VolumeShape {
154    /// A cylinder along the depth axis.
155    Cylinder {
156        /// Radius, m.
157        radius_m: f64,
158    },
159    /// A box with the given half widths along `y` and `z`.
160    Box {
161        /// Half width along `y`, m.
162        half_y_m: f64,
163        /// Half width along `z`, m.
164        half_z_m: f64,
165    },
166}
167
168/// A target volume: a lateral shape over a depth range `depth_lo_m <= x <
169/// depth_hi_m`, centred at lateral offset `(offset_y_m, offset_z_m)` from the
170/// beam axis. Deposits exist only inside the target, so a volume extending
171/// beyond the film collects only the part inside it.
172#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
173pub struct ExposureVolume {
174    /// Label in the outputs.
175    pub name: String,
176    /// Lateral shape.
177    pub shape: VolumeShape,
178    /// Lower depth, m.
179    pub depth_lo_m: f64,
180    /// Upper depth (excluded), m.
181    pub depth_hi_m: f64,
182    /// Lateral centre `y`, m.
183    pub offset_y_m: f64,
184    /// Lateral centre `z`, m.
185    pub offset_z_m: f64,
186}
187
188impl ExposureVolume {
189    fn validate(&self) -> Result<(), ExposureError> {
190        let ok = |v: f64| v.is_finite();
191        let shape_ok = match self.shape {
192            VolumeShape::Cylinder { radius_m } => ok(radius_m) && radius_m > 0.0,
193            VolumeShape::Box { half_y_m, half_z_m } => {
194                ok(half_y_m) && half_y_m > 0.0 && ok(half_z_m) && half_z_m > 0.0
195            }
196        };
197        if !shape_ok {
198            return invalid(format!("volume '{}': size must be positive", self.name));
199        }
200        if !(ok(self.depth_lo_m)
201            && ok(self.depth_hi_m)
202            && self.depth_hi_m > self.depth_lo_m
203            && ok(self.offset_y_m)
204            && ok(self.offset_z_m))
205        {
206            return invalid(format!(
207                "volume '{}': need finite offsets and depth_hi > depth_lo",
208                self.name
209            ));
210        }
211        Ok(())
212    }
213
214    /// Whether `pos` (depth, then the lateral position relative to the
215    /// landing point shifted into the volume's frame) lies inside.
216    fn contains(&self, x: f64, dy: f64, dz: f64) -> bool {
217        if !(x >= self.depth_lo_m && x < self.depth_hi_m) {
218            return false;
219        }
220        let (y, z) = (dy - self.offset_y_m, dz - self.offset_z_m);
221        match self.shape {
222            VolumeShape::Cylinder { radius_m } => y.hypot(z) < radius_m,
223            VolumeShape::Box { half_y_m, half_z_m } => y.abs() < half_y_m && z.abs() < half_z_m,
224        }
225    }
226}
227
228/// Settings of the exposure tally.
229#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
230pub struct ExposureConfig {
231    /// The volumes (a neighbour is simply another volume at its offset).
232    pub volumes: Vec<ExposureVolume>,
233    /// Beam profile.
234    pub beam: BeamProfile,
235    /// Seed of the landing-offset stream (independent of the transport's).
236    pub seed: u64,
237    /// Bins of the histogram of `X`.
238    pub bins: usize,
239    /// Upper edge of the histogram, eV.
240    pub x_max_ev: f64,
241    /// Mean number of incident electrons `λ` of the exposure.
242    pub lambda: f64,
243    /// Probabilities `p` in `(0, 1)` whose quantiles of `S` are reported.
244    pub quantiles: Vec<f64>,
245    /// Thresholds `a`, eV, for `P(S < a)`.
246    pub below_ev: Vec<f64>,
247    /// Thresholds `b`, eV, for `P(S > b)`.
248    pub above_ev: Vec<f64>,
249    /// Poisson tail mass dropped from the convolution.
250    pub tail_tolerance: f64,
251}
252
253impl ExposureConfig {
254    /// Checks the configuration.
255    pub fn validate(&self) -> Result<(), ExposureError> {
256        if self.volumes.is_empty() {
257            return invalid("at least one volume is required");
258        }
259        for v in &self.volumes {
260            v.validate()?;
261        }
262        self.beam.validate()?;
263        if self.bins == 0 || !(self.x_max_ev.is_finite() && self.x_max_ev > 0.0) {
264            return invalid("bins must be at least 1 and x_max_ev positive");
265        }
266        if !(self.lambda.is_finite() && self.lambda > 0.0) {
267            return invalid("lambda must be positive");
268        }
269        if self.quantiles.iter().any(|p| !(*p > 0.0 && *p < 1.0)) {
270            return invalid("quantiles must lie in (0, 1)");
271        }
272        if self
273            .below_ev
274            .iter()
275            .chain(&self.above_ev)
276            .any(|t| !t.is_finite())
277        {
278            return invalid("thresholds must be finite");
279        }
280        if !(self.tail_tolerance > 0.0 && self.tail_tolerance < 1.0) {
281            return invalid("tail_tolerance must lie in (0, 1)");
282        }
283        Ok(())
284    }
285}
286
287/// The histogram of the per-electron deposit `X` in one volume, with exact
288/// moment sums. Mergeable; see [the module docs](self).
289#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
290pub struct ExposureHistogram {
291    /// Primary histories tallied.
292    pub histories: u64,
293    /// Histories with `X = 0` exactly.
294    pub zero_count: u64,
295    /// Bin width, eV.
296    pub bin_width_ev: f64,
297    /// Counts of `(j-1) w < X <= j w`, `j = 1..=bins` (index `j-1`).
298    pub counts: Vec<u64>,
299    /// Histories with `X` beyond the last edge (also counted in the last bin).
300    pub overflow_count: u64,
301    /// `Σ X`, eV.
302    pub sum_x: f64,
303    /// `Σ X²`, eV².
304    pub sum_x2: f64,
305    /// `Σ X³`, eV³.
306    pub sum_x3: f64,
307    /// `Σ X⁴`, eV⁴.
308    pub sum_x4: f64,
309}
310
311impl ExposureHistogram {
312    fn new(bins: usize, x_max_ev: f64) -> Self {
313        Self {
314            histories: 0,
315            zero_count: 0,
316            bin_width_ev: x_max_ev / bins as f64,
317            counts: vec![0; bins],
318            overflow_count: 0,
319            sum_x: 0.0,
320            sum_x2: 0.0,
321            sum_x3: 0.0,
322            sum_x4: 0.0,
323        }
324    }
325
326    fn push(&mut self, x: f64) {
327        self.histories += 1;
328        if x <= 0.0 {
329            self.zero_count += 1;
330            return;
331        }
332        let bins = self.counts.len();
333        let j = (x / self.bin_width_ev).ceil() as usize; // 1-based
334        if j > bins {
335            self.overflow_count += 1;
336        }
337        self.counts[j.clamp(1, bins) - 1] += 1;
338        let x2 = x * x;
339        self.sum_x += x;
340        self.sum_x2 += x2;
341        self.sum_x3 += x2 * x;
342        self.sum_x4 += x2 * x2;
343    }
344
345    fn merge(&mut self, o: &Self) {
346        self.histories += o.histories;
347        self.zero_count += o.zero_count;
348        self.overflow_count += o.overflow_count;
349        for (a, b) in self.counts.iter_mut().zip(&o.counts) {
350            *a += b;
351        }
352        self.sum_x += o.sum_x;
353        self.sum_x2 += o.sum_x2;
354        self.sum_x3 += o.sum_x3;
355        self.sum_x4 += o.sum_x4;
356    }
357
358    /// `E[X]`, eV (0 with no histories).
359    pub fn mean_ev(&self) -> f64 {
360        self.sum_x / self.histories.max(1) as f64
361    }
362
363    /// `E[X²]`, eV².
364    pub fn second_moment_ev2(&self) -> f64 {
365        self.sum_x2 / self.histories.max(1) as f64
366    }
367
368    /// Standard error of `E[X]`: `sqrt((E[X²] - E[X]²)/(n-1))`; 0 for `n < 2`.
369    pub fn mean_std_err_ev(&self) -> f64 {
370        let n = self.histories as f64;
371        if self.histories < 2 {
372            return 0.0;
373        }
374        let var = (self.second_moment_ev2() - self.mean_ev().powi(2)).max(0.0);
375        (var * n / (n - 1.0) / n).sqrt()
376    }
377
378    /// Standard error of `E[X²]` from `E[X⁴] - E[X²]²`.
379    pub fn second_moment_std_err_ev2(&self) -> f64 {
380        let n = self.histories as f64;
381        if self.histories < 2 {
382            return 0.0;
383        }
384        let m2 = self.second_moment_ev2();
385        let var = (self.sum_x4 / n - m2 * m2).max(0.0);
386        (var * n / (n - 1.0) / n).sqrt()
387    }
388
389    /// Fraction of histories with `X = 0`.
390    pub fn zero_fraction(&self) -> f64 {
391        self.zero_count as f64 / self.histories.max(1) as f64
392    }
393
394    /// The compound-Poisson summary of `S = X_1 + … + X_N`, `N ~ Poisson(λ)`,
395    /// for any `λ`, from this histogram alone (no rerun). See
396    /// [the module docs](self#compound-poisson-statistics).
397    pub fn compound_poisson(
398        &self,
399        lambda: f64,
400        request: &CompoundRequest,
401    ) -> Result<CompoundPoisson, ExposureError> {
402        if !(lambda.is_finite() && lambda > 0.0) {
403            return invalid("lambda must be positive");
404        }
405        if self.histories == 0 {
406            return invalid("no histories");
407        }
408        let mean = lambda * self.mean_ev();
409        let variance = lambda * self.second_moment_ev2();
410        let mut out = CompoundPoisson {
411            lambda,
412            mean_ev: mean,
413            mean_std_err_ev: lambda * self.mean_std_err_ev(),
414            variance_ev2: variance,
415            variance_std_err_ev2: lambda * self.second_moment_std_err_ev2(),
416            p_zero: 0.0,
417            k_max: 0,
418            truncated_mass: 0.0,
419            bin_width_ev: self.bin_width_ev,
420            overflow_fraction: self.overflow_count as f64 / self.histories as f64,
421            distribution: None,
422            note: None,
423        };
424        let p0 = self.zero_fraction();
425        let nz_total: u64 = self.counts.iter().sum();
426        // Atom at S = 0: no electron deposits anything.
427        out.p_zero = (-lambda * (1.0 - p0)).exp();
428        if nz_total == 0 {
429            out.k_max = 0;
430            out.distribution = Some(distribution_atom_at_zero(request));
431            return Ok(out);
432        }
433        let mu = lambda * (1.0 - p0); // mean number of depositing electrons
434        let (poisson, k_max, trunc) = poisson_pmf(mu, request.tail_tolerance);
435        out.k_max = k_max;
436        out.truncated_mass = trunc;
437        let bins = self.counts.len();
438        let work = (k_max as f64).powi(2) * (bins as f64).powi(2) / 2.0;
439        if work > MAX_CONVOLUTION_WORK {
440            out.note = Some(format!(
441                "distribution not computed: lambda too large for the convolution \
442                 (k_max {k_max}, {bins} bins); mean and variance are exact"
443            ));
444            return Ok(out);
445        }
446        let f: Vec<f64> = self
447            .counts
448            .iter()
449            .map(|&c| c as f64 / nz_total as f64)
450            .collect();
451        // Mass on the integer grid m = Σ j_i, in half-units of the bin width:
452        // lower edge 2(m-k), centre 2m-k, upper 2m. Index by 2m (so 0..=2 k_max bins).
453        let size = 2 * k_max * bins + 1;
454        let mut lo = vec![0.0; size];
455        let mut mid = vec![0.0; size];
456        let mut hi = vec![0.0; size];
457        lo[0] += poisson[0];
458        mid[0] += poisson[0];
459        hi[0] += poisson[0];
460        let mut g = vec![1.0]; // f^{*0}, support m = 0
461        for k in 1..=k_max {
462            // g <- g * f; support of g is m in [k-1, (k-1) bins] -> [k, k bins].
463            let mut next = vec![0.0; g.len() + bins];
464            for (m, &gm) in g.iter().enumerate() {
465                if gm == 0.0 {
466                    continue;
467                }
468                for (j, &fj) in f.iter().enumerate() {
469                    next[m + j + 1] += gm * fj;
470                }
471            }
472            g = next;
473            // g index is m directly (shifted by the +1 per factor above).
474            let pk = poisson[k];
475            for (m, &gm) in g.iter().enumerate() {
476                if gm == 0.0 {
477                    continue;
478                }
479                let w = pk * gm;
480                // Values in half-units: lower 2(m-k), centre 2m-k, upper 2m.
481                lo[2 * (m - k)] += w;
482                mid[2 * m - k] += w;
483                hi[2 * m] += w;
484            }
485        }
486        out.distribution = Some(summarise(self.bin_width_ev, &lo, &mid, &hi, request));
487        Ok(out)
488    }
489}
490
491/// Which compound-Poisson outputs to compute.
492#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
493pub struct CompoundRequest {
494    /// Probabilities whose quantiles of `S` are wanted.
495    pub quantiles: Vec<f64>,
496    /// Thresholds `a` for `P(S < a)`, eV.
497    pub below_ev: Vec<f64>,
498    /// Thresholds `b` for `P(S > b)`, eV.
499    pub above_ev: Vec<f64>,
500    /// Dropped Poisson tail mass.
501    pub tail_tolerance: f64,
502}
503
504/// A value bracketed by the binning representations of `X` (see
505/// [the module docs](self#compound-poisson-statistics)). `low` is the
506/// result with every bin at its lower edge, `high` at its upper edge,
507/// `central` at the bin centres. For a quantile and for `P(S > b)` the true
508/// value lies in `[low, high]`; for `P(S < a)` it lies in `[high, low]`.
509#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
510pub struct Bracket {
511    /// Lower-edge representation.
512    pub low: f64,
513    /// Bin-centre representation.
514    pub central: f64,
515    /// Upper-edge representation.
516    pub high: f64,
517}
518
519/// A reported quantile of `S`, eV.
520#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
521pub struct QuantileResult {
522    /// The probability.
523    pub p: f64,
524    /// The quantile, eV: the least `s` with `P(S <= s) >= p`.
525    pub energy_ev: Bracket,
526}
527
528/// A reported threshold probability.
529#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
530pub struct ThresholdResult {
531    /// The threshold, eV.
532    pub threshold_ev: f64,
533    /// The probability.
534    pub probability: Bracket,
535}
536
537/// The discrete distribution results of a compound-Poisson sum.
538#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
539pub struct CompoundDistribution {
540    /// Quantiles of `S`.
541    pub quantiles: Vec<QuantileResult>,
542    /// `P(S < a)`.
543    pub below: Vec<ThresholdResult>,
544    /// `P(S > b)`.
545    pub above: Vec<ThresholdResult>,
546}
547
548/// Compound-Poisson summary of the energy `S` in a volume.
549#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
550pub struct CompoundPoisson {
551    /// Mean number of incident electrons.
552    pub lambda: f64,
553    /// `E[S] = λ E[X]`, eV.
554    pub mean_ev: f64,
555    /// Standard error of `mean_ev`, eV.
556    pub mean_std_err_ev: f64,
557    /// `Var[S] = λ E[X²]`, eV².
558    pub variance_ev2: f64,
559    /// Standard error of `variance_ev2`, eV².
560    pub variance_std_err_ev2: f64,
561    /// `P(S = 0)`, exact.
562    pub p_zero: f64,
563    /// Largest number of depositing electrons kept in the convolution.
564    pub k_max: usize,
565    /// Poisson mass beyond `k_max` that was dropped.
566    pub truncated_mass: f64,
567    /// Bin width of `X`, eV (the grid of the representation).
568    pub bin_width_ev: f64,
569    /// Fraction of histories whose `X` was clamped into the last bin.
570    pub overflow_fraction: f64,
571    /// Quantiles and threshold probabilities, if computed.
572    pub distribution: Option<CompoundDistribution>,
573    /// Why `distribution` is absent, if it is.
574    pub note: Option<String>,
575}
576
577fn distribution_atom_at_zero(r: &CompoundRequest) -> CompoundDistribution {
578    let b = |v: f64| Bracket {
579        low: v,
580        central: v,
581        high: v,
582    };
583    CompoundDistribution {
584        quantiles: r
585            .quantiles
586            .iter()
587            .map(|&p| QuantileResult {
588                p,
589                energy_ev: b(0.0),
590            })
591            .collect(),
592        below: r
593            .below_ev
594            .iter()
595            .map(|&t| ThresholdResult {
596                threshold_ev: t,
597                probability: b(if t > 0.0 { 1.0 } else { 0.0 }),
598            })
599            .collect(),
600        above: r
601            .above_ev
602            .iter()
603            .map(|&t| ThresholdResult {
604                threshold_ev: t,
605                probability: b(if t < 0.0 { 1.0 } else { 0.0 }),
606            })
607            .collect(),
608    }
609}
610
611/// Poisson pmf `P(k)` for `k = 0..=k_max`, where `k_max` is the smallest
612/// count whose upper tail is below `tol`; returns also the dropped mass.
613fn poisson_pmf(mu: f64, tol: f64) -> (Vec<f64>, usize, f64) {
614    let mut p = vec![(-mu).exp()];
615    let mut cum = p[0];
616    let mut k = 0usize;
617    // exp(-mu) underflows for very large mu; the work check rejects those.
618    while 1.0 - cum > tol && k < 1_000_000 {
619        k += 1;
620        let next = p[k - 1] * mu / k as f64;
621        p.push(next);
622        cum += next;
623        if next == 0.0 && k as f64 > mu {
624            break;
625        }
626    }
627    (p, k, (1.0 - cum).max(0.0))
628}
629
630/// Quantiles and thresholds from the three half-unit-indexed pmfs.
631fn summarise(
632    w: f64,
633    lo: &[f64],
634    mid: &[f64],
635    hi: &[f64],
636    r: &CompoundRequest,
637) -> CompoundDistribution {
638    // Entry i of each array is the mass at i * (w / 2).
639    let half = w / 2.0;
640    let cdf = |a: &[f64]| -> Vec<f64> {
641        let mut c = 0.0;
642        a.iter()
643            .map(|&m| {
644                c += m;
645                c
646            })
647            .collect()
648    };
649    let (cl, cm, ch) = (cdf(lo), cdf(mid), cdf(hi));
650    let quantile = |c: &[f64], p: f64| -> f64 {
651        let i = c.iter().position(|&v| v >= p).unwrap_or(c.len() - 1);
652        i as f64 * half
653    };
654    // P(S <= t) and P(S < t): mass strictly below t is the cdf at the last
655    // grid point < t.
656    let below = |c: &[f64], t: f64| -> f64 {
657        if t <= 0.0 {
658            return 0.0;
659        }
660        // Grid points i*half < t  <=>  i < t/half.
661        let n = (t / half).ceil() as usize; // count of grid points strictly below t
662        if n == 0 {
663            0.0
664        } else {
665            c[(n - 1).min(c.len() - 1)]
666        }
667    };
668    let le = |c: &[f64], t: f64| -> f64 {
669        if t < 0.0 {
670            return 0.0;
671        }
672        let i = (t / half).floor() as usize;
673        c[i.min(c.len() - 1)]
674    };
675    CompoundDistribution {
676        quantiles: r
677            .quantiles
678            .iter()
679            .map(|&p| QuantileResult {
680                p,
681                energy_ev: Bracket {
682                    low: quantile(&cl, p),
683                    central: quantile(&cm, p),
684                    high: quantile(&ch, p),
685                },
686            })
687            .collect(),
688        below: r
689            .below_ev
690            .iter()
691            .map(|&t| ThresholdResult {
692                threshold_ev: t,
693                probability: Bracket {
694                    low: below(&cl, t),
695                    central: below(&cm, t),
696                    high: below(&ch, t),
697                },
698            })
699            .collect(),
700        above: r
701            .above_ev
702            .iter()
703            .map(|&t| ThresholdResult {
704                threshold_ev: t,
705                probability: Bracket {
706                    low: 1.0 - le(&cl, t),
707                    central: 1.0 - le(&cm, t),
708                    high: 1.0 - le(&ch, t),
709                },
710            })
711            .collect(),
712    }
713}
714
715/// Per-history accumulation of the deposit in each configured volume.
716/// Call [`begin_history`](Self::begin_history), then
717/// [`deposit`](Self::deposit) for every deposit, then
718/// [`end_history`](Self::end_history).
719#[derive(Debug, Clone, PartialEq)]
720pub struct ExposureAccumulator {
721    config: ExposureConfig,
722    hists: Vec<ExposureHistogram>,
723    /// Energy of the current history per volume.
724    scratch: Vec<f64>,
725    /// Landing offset and start of the current history.
726    shift: [f64; 2],
727    origin: [f64; 2],
728}
729
730impl ExposureAccumulator {
731    /// Empty accumulator.
732    pub fn new(config: ExposureConfig) -> Result<Self, ExposureError> {
733        config.validate()?;
734        let n = config.volumes.len();
735        let hists = (0..n)
736            .map(|_| ExposureHistogram::new(config.bins, config.x_max_ev))
737            .collect();
738        Ok(Self {
739            config,
740            hists,
741            scratch: vec![0.0; n],
742            shift: [0.0; 2],
743            origin: [0.0; 2],
744        })
745    }
746
747    /// The configuration.
748    pub fn config(&self) -> &ExposureConfig {
749        &self.config
750    }
751
752    /// Starts history `index` of a primary starting at lateral position
753    /// `start_yz` (m).
754    pub fn begin_history(&mut self, index: u64, start_yz: [f64; 2]) {
755        self.shift = self.config.beam.offset(self.config.seed, index);
756        self.origin = start_yz;
757        self.scratch.iter_mut().for_each(|s| *s = 0.0);
758    }
759
760    /// A deposit of `e` eV at `pos` (m) in the current history.
761    pub fn deposit(&mut self, pos: [f64; 3], e: f64) {
762        if !(e.is_finite() && e > 0.0) {
763            return;
764        }
765        let dy = pos[1] - self.origin[0] + self.shift[0];
766        let dz = pos[2] - self.origin[1] + self.shift[1];
767        for (v, s) in self.config.volumes.iter().zip(&mut self.scratch) {
768            if v.contains(pos[0], dy, dz) {
769                *s += e;
770            }
771        }
772    }
773
774    /// Ends the current history.
775    pub fn end_history(&mut self) {
776        for (h, s) in self.hists.iter_mut().zip(&mut self.scratch) {
777            h.push(*s);
778            *s = 0.0;
779        }
780    }
781
782    /// Field-wise merge of another accumulator's totals.
783    ///
784    /// # Panics
785    /// If the configurations differ.
786    pub fn merge(&mut self, o: &Self) {
787        assert!(
788            self.config == o.config,
789            "merging exposure accumulators with different configurations"
790        );
791        for (a, b) in self.hists.iter_mut().zip(&o.hists) {
792            a.merge(b);
793        }
794    }
795
796    /// The per-volume results.
797    pub fn report(&self) -> ExposureReport {
798        let request = CompoundRequest {
799            quantiles: self.config.quantiles.clone(),
800            below_ev: self.config.below_ev.clone(),
801            above_ev: self.config.above_ev.clone(),
802            tail_tolerance: self.config.tail_tolerance,
803        };
804        let volumes = self
805            .config
806            .volumes
807            .iter()
808            .zip(&self.hists)
809            .map(|(v, h)| {
810                let (compound, error) = match h.compound_poisson(self.config.lambda, &request) {
811                    Ok(c) => (Some(c), None),
812                    Err(e) => (None, Some(e.to_string())),
813                };
814                ExposureVolumeReport {
815                    volume: v.clone(),
816                    mean_ev: h.mean_ev(),
817                    mean_std_err_ev: h.mean_std_err_ev(),
818                    second_moment_ev2: h.second_moment_ev2(),
819                    second_moment_std_err_ev2: h.second_moment_std_err_ev2(),
820                    zero_fraction: h.zero_fraction(),
821                    histogram: h.clone(),
822                    compound,
823                    compound_error: error,
824                }
825            })
826            .collect();
827        ExposureReport {
828            lambda: self.config.lambda,
829            beam: self.config.beam,
830            volumes,
831        }
832    }
833}
834
835/// Results for one volume.
836#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
837pub struct ExposureVolumeReport {
838    /// The volume.
839    pub volume: ExposureVolume,
840    /// `E[X]`, eV.
841    pub mean_ev: f64,
842    /// Its standard error, eV.
843    pub mean_std_err_ev: f64,
844    /// `E[X²]`, eV².
845    pub second_moment_ev2: f64,
846    /// Its standard error, eV².
847    pub second_moment_std_err_ev2: f64,
848    /// `P(X = 0)`.
849    pub zero_fraction: f64,
850    /// The per-electron distribution.
851    pub histogram: ExposureHistogram,
852    /// The compound-Poisson summary at the configured `λ`.
853    pub compound: Option<CompoundPoisson>,
854    /// Why `compound` is absent, if it is.
855    pub compound_error: Option<String>,
856}
857
858/// Exposure results of a run.
859#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
860pub struct ExposureReport {
861    /// Mean number of incident electrons.
862    pub lambda: f64,
863    /// Beam profile the landing offsets were drawn from.
864    pub beam: BeamProfile,
865    /// Per volume, in configuration order.
866    pub volumes: Vec<ExposureVolumeReport>,
867}
868
869impl ExposureReport {
870    /// CSV of the per-electron histograms: `volume,bin,lo_ev,hi_ev,count,
871    /// fraction`, with bin 0 the exact `X = 0` atom (`lo = hi = 0`).
872    pub fn histogram_csv(&self) -> String {
873        let mut s = String::from("volume,bin,lo_ev,hi_ev,count,fraction\n");
874        for v in &self.volumes {
875            let h = &v.histogram;
876            let n = h.histories.max(1) as f64;
877            s += &format!(
878                "{},0,0,0,{},{}\n",
879                v.volume.name,
880                h.zero_count,
881                h.zero_count as f64 / n
882            );
883            for (i, &c) in h.counts.iter().enumerate() {
884                s += &format!(
885                    "{},{},{},{},{},{}\n",
886                    v.volume.name,
887                    i + 1,
888                    i as f64 * h.bin_width_ev,
889                    (i + 1) as f64 * h.bin_width_ev,
890                    c,
891                    c as f64 / n
892                );
893            }
894        }
895        s
896    }
897}
898
899#[cfg(test)]
900mod tests {
901    use super::*;
902
903    fn req() -> CompoundRequest {
904        CompoundRequest {
905            quantiles: vec![0.5],
906            below_ev: vec![],
907            above_ev: vec![],
908            tail_tolerance: 1e-14,
909        }
910    }
911
912    /// X = 1 eV with probability p (a single bin of width 1): then
913    /// S ~ Poisson(λ p), and the three representations are all exact.
914    #[test]
915    fn thinned_poisson_matches_the_closed_form() {
916        let mut h = ExposureHistogram::new(1, 1.0);
917        for i in 0..100 {
918            h.push(if i < 30 { 1.0 } else { 0.0 });
919        }
920        let lambda = 4.0;
921        let r = CompoundRequest {
922            quantiles: vec![0.5, 0.9],
923            below_ev: vec![3.0],
924            above_ev: vec![5.0],
925            tail_tolerance: 1e-15,
926        };
927        let c = h.compound_poisson(lambda, &r).unwrap();
928        let mu: f64 = lambda * 0.3;
929        // Closed form.
930        let pmf = |k: u32| (-mu).exp() * mu.powi(k as i32) / (1..=k).product::<u32>() as f64;
931        let d = c.distribution.unwrap();
932        let below3: f64 = (0..3).map(pmf).sum();
933        let above5: f64 = 1.0 - (0..=5).map(pmf).sum::<f64>();
934        // Upper edge representation is exact for X = 1 (bin (0,1]).
935        assert!((d.below[0].probability.high - below3).abs() < 1e-12);
936        assert!((d.above[0].probability.high - above5).abs() < 1e-12);
937        assert!((c.p_zero - (-mu).exp()).abs() < 1e-15);
938        assert!((c.mean_ev - lambda * 0.3).abs() < 1e-12);
939        assert!((c.variance_ev2 - lambda * 0.3).abs() < 1e-12);
940        // Median of Poisson(1.2) is 1.
941        assert_eq!(d.quantiles[0].energy_ev.high, 1.0);
942        assert_eq!(d.quantiles[0].energy_ev.low, 0.0);
943    }
944
945    #[test]
946    fn empty_volume_is_an_atom_at_zero() {
947        let mut h = ExposureHistogram::new(4, 4.0);
948        for _ in 0..10 {
949            h.push(0.0);
950        }
951        let mut r = req();
952        r.below_ev = vec![1.0];
953        r.above_ev = vec![0.0];
954        let c = h.compound_poisson(7.0, &r).unwrap();
955        assert_eq!(c.mean_ev, 0.0);
956        let d = c.distribution.unwrap();
957        assert_eq!(d.quantiles[0].energy_ev.central, 0.0);
958        assert_eq!(d.below[0].probability.central, 1.0);
959        assert_eq!(d.above[0].probability.central, 0.0);
960    }
961
962    #[test]
963    fn brackets_order_and_moments_are_consistent() {
964        let mut h = ExposureHistogram::new(10, 10.0);
965        for i in 0..1000u32 {
966            h.push(if i % 3 == 0 {
967                0.0
968            } else {
969                (i % 17) as f64 * 0.5
970            });
971        }
972        let mut r = req();
973        r.quantiles = vec![0.1, 0.5, 0.9];
974        r.above_ev = vec![8.0];
975        let c = h.compound_poisson(3.0, &r).unwrap();
976        let d = c.distribution.unwrap();
977        for q in &d.quantiles {
978            let b = q.energy_ev;
979            assert!(b.low <= b.central && b.central <= b.high);
980        }
981        let b = d.above[0].probability;
982        assert!(b.low <= b.central && b.central <= b.high);
983        // Bin-centre mean of S is close to the exact mean to within K_max bins.
984        assert!(c.truncated_mass <= 1e-14);
985    }
986
987    #[test]
988    fn config_validation_rejects_bad_scalars() {
989        let v = ExposureVolume {
990            name: "a".into(),
991            shape: VolumeShape::Cylinder { radius_m: 1e-8 },
992            depth_lo_m: 0.0,
993            depth_hi_m: 1e-8,
994            offset_y_m: 0.0,
995            offset_z_m: 0.0,
996        };
997        let ok = ExposureConfig {
998            volumes: vec![v],
999            beam: BeamProfile::Point,
1000            seed: 1,
1001            bins: 10,
1002            x_max_ev: 100.0,
1003            lambda: 1.0,
1004            quantiles: vec![0.5],
1005            below_ev: vec![],
1006            above_ev: vec![],
1007            tail_tolerance: 1e-12,
1008        };
1009        assert!(ok.validate().is_ok());
1010        let mut b = ok.clone();
1011        b.lambda = 0.0;
1012        assert!(b.validate().is_err());
1013        let mut b = ok.clone();
1014        b.volumes.clear();
1015        assert!(b.validate().is_err());
1016        let mut b = ok;
1017        b.quantiles = vec![1.0];
1018        assert!(b.validate().is_err());
1019    }
1020
1021    #[test]
1022    fn beam_offsets_are_deterministic_and_on_the_disc() {
1023        let b = BeamProfile::UniformDisc { radius_m: 2.0 };
1024        assert_eq!(b.offset(5, 9), b.offset(5, 9));
1025        for i in 0..200 {
1026            let [y, z] = b.offset(5, i);
1027            assert!(y.hypot(z) <= 2.0);
1028        }
1029        assert_eq!(BeamProfile::Point.offset(1, 1), [0.0, 0.0]);
1030    }
1031}