Skip to main content

nucleide_vr_tools/
sampling.rs

1//! Walker/Vose alias-table source sampling.
2//!
3//! [`AliasTable`] implements Walker's method with Vose's construction
4//! (Wuttke 2013), including the reverted small/large index ordering and the
5//! `prob == 1` drain for numerically degenerate leftovers.
6//! [`MeshSourceSampler`] applies it to [`MeshTallyData`]: ANALOG samples the
7//! (volume-weighted) tally totals, UNIFORM samples space uniformly, and USER
8//! consumes an external density array. Only voxel-level sampling is provided;
9//! subvoxel/cell-fraction modes are out of scope.
10
11use nucleide_mcnp_io::meshtal::MeshTallyData;
12
13use crate::Error;
14
15/// Walker/Vose alias table over a discrete PDF.
16///
17/// Build once with [`AliasTable::new`], then draw with
18/// [`AliasTable::sample`] using two uniform random numbers in `[0, 1)`.
19#[derive(Debug, Clone, PartialEq)]
20pub struct AliasTable {
21    n: usize,
22    /// Normalized input PDF, stored for reference/inspection.
23    pdf: Vec<f64>,
24    /// Per-bin cut probability (`<= 1`) as built by Vose's algorithm.
25    prob: Vec<f64>,
26    /// Per-bin alias index.
27    alias: Vec<usize>,
28}
29
30impl AliasTable {
31    /// Build an alias table from a non-negative PDF.
32    ///
33    /// The PDF is normalized internally so the standalone API is safe with
34    /// unnormalized input (the underlying construction assumes a unit sum).
35    pub fn new(pdf: &[f64]) -> Result<Self, Error> {
36        if pdf.is_empty() {
37            return Err(Error::EmptyPdf);
38        }
39        let mut sum = 0.0;
40        for (i, &p) in pdf.iter().enumerate() {
41            if !p.is_finite() {
42                return Err(Error::NonFinitePdf { index: i });
43            }
44            if p < 0.0 {
45                return Err(Error::NegativePdf { index: i, value: p });
46            }
47            sum += p;
48        }
49        if sum <= 0.0 {
50            return Err(Error::ZeroSumPdf);
51        }
52
53        let n = pdf.len();
54        // Normalized PDF: also what `pdf()` reports, since this is what the
55        // table actually samples.
56        let normalized_pdf: Vec<f64> = pdf.iter().map(|&x| x / sum).collect();
57        let mut p = normalized_pdf.clone();
58
59        // Scale so the mean bin probability is exactly 1.
60        for x in p.iter_mut() {
61            *x *= n as f64;
62        }
63
64        // Separate index lists for small and large probabilities. As in the
65        // Wuttke implementation, indices are visited in reverted order.
66        let mut small = Vec::with_capacity(n);
67        let mut large = Vec::with_capacity(n);
68        let mut i = n;
69        while i > 0 {
70            i -= 1;
71            if p[i] < 1.0 {
72                small.push(i);
73            } else {
74                large.push(i);
75            }
76        }
77
78        let mut prob = vec![0.0; n];
79        let mut alias = vec![0usize; n];
80        while !small.is_empty() && !large.is_empty() {
81            let a = small.pop().expect("small non-empty");
82            let g = large.pop().expect("large non-empty");
83            prob[a] = p[a];
84            alias[a] = g;
85            p[g] += p[a] - 1.0;
86            if p[g] < 1.0 {
87                small.push(g);
88            } else {
89                large.push(g);
90            }
91        }
92        while let Some(g) = large.pop() {
93            prob[g] = 1.0;
94        }
95        // Can only happen through numeric instability.
96        while let Some(a) = small.pop() {
97            prob[a] = 1.0;
98        }
99
100        Ok(Self {
101            n,
102            pdf: normalized_pdf,
103            prob,
104            alias,
105        })
106    }
107
108    /// Draw one index using two uniforms in `[0, 1)`.
109    ///
110    /// Mirrors `sample_pdf(rand1, rand2)`: pick column `n * rand1`, return it
111    /// when `rand2 < prob[column]`, else the column's alias.
112    ///
113    /// `r1` is saturated to `[0, n)` before indexing: negative values map to
114    /// column `0` and values `>= 1.0` map to column `n - 1`. This is
115    /// intentional defensive behavior; callers should still pass uniforms in
116    /// `[0, 1)`.
117    pub fn sample(&self, r1: f64, r2: f64) -> usize {
118        let mut i = (self.n as f64 * r1) as usize;
119        if i >= self.n {
120            i = self.n - 1;
121        }
122        if r2 < self.prob[i] {
123            i
124        } else {
125            self.alias[i]
126        }
127    }
128
129    /// The normalized PDF the table samples from.
130    pub fn pdf(&self) -> &[f64] {
131        &self.pdf
132    }
133
134    /// Number of bins.
135    pub fn len(&self) -> usize {
136        self.n
137    }
138
139    /// Whether the table has no bins.
140    pub fn is_empty(&self) -> bool {
141        self.n == 0
142    }
143
144    /// Exact outcome probabilities implied by the table:
145    /// `P(i) = (prob[i] + sum over j with alias[j] == i of (1 - prob[j])) / n`.
146    #[cfg(test)]
147    fn exact_probabilities(&self) -> Vec<f64> {
148        let mut out = vec![0.0; self.n];
149        for c in 0..self.n {
150            let base = self.prob[c].min(1.0);
151            out[c] += base / self.n as f64;
152            if base < 1.0 {
153                out[self.alias[c]] += (1.0 - base) / self.n as f64;
154            }
155        }
156        out
157    }
158}
159
160/// Source-sampling bias mode (values 0–2 at voxel level).
161#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
162pub enum Mode {
163    /// Sample where particles are born: PDF ∝ total source strength per voxel.
164    #[default]
165    Analog,
166    /// Sample uniformly in space: PDF ∝ voxel volume.
167    Uniform,
168    /// Sample from an external (unnormalized) density array.
169    User,
170}
171
172/// One sampled particle birth site (voxel resolution).
173#[derive(Debug, Clone, Copy, PartialEq)]
174pub struct SampledVoxel {
175    /// Flat volume-element index, x slowest → z fastest.
176    pub index: usize,
177    /// X cell index.
178    pub i: usize,
179    /// Y cell index.
180    pub j: usize,
181    /// Z cell index.
182    pub k: usize,
183    /// Birth weight: 1.0 for analog sampling, otherwise
184    /// `analog_pdf[voxel] / biased_pdf[voxel]`.
185    pub weight: f64,
186}
187
188/// Mesh source sampler over [`MeshTallyData`] energy-integrated totals.
189///
190/// The PDF lives over volume elements only (no energy dimension):
191/// densities are multiplied by cell volumes before
192/// normalization, since structured meshes may have unequal cells.
193#[derive(Debug, Clone, PartialEq)]
194pub struct MeshSourceSampler {
195    dims: [usize; 3],
196    num_ves: usize,
197    mode: Mode,
198    table: AliasTable,
199    /// Birth weight per bin: `pdf / bias_pdf` elementwise (all ones in analog).
200    biased_weights: Vec<f64>,
201}
202
203impl MeshSourceSampler {
204    /// Build a sampler over `tally`'s totals in `mode`.
205    ///
206    /// For [`Mode::User`], `user_pdf` supplies one unnormalized density value
207    /// per volume element (length must equal `tally.num_ves()`); it is
208    /// scaled by cell volumes and normalized, like a bias tag.
209    pub fn new(tally: &MeshTallyData, mode: Mode, user_pdf: Option<&[f64]>) -> Result<Self, Error> {
210        let num_ves = tally.num_ves();
211        if num_ves == 0 {
212            return Err(Error::EmptyTally);
213        }
214        if tally.total_result.len() != num_ves {
215            return Err(Error::LengthMismatch {
216                expected: num_ves,
217                got: tally.total_result.len(),
218            });
219        }
220
221        let volumes = cell_volumes(tally);
222        let analog_pdf: Vec<f64> = tally
223            .total_result
224            .iter()
225            .zip(volumes.iter())
226            .enumerate()
227            .map(|(i, (q, v))| {
228                if !q.is_finite() {
229                    return Err(Error::NonFiniteTally {
230                        field: "total_result",
231                        index: i,
232                    });
233                }
234                if *q < 0.0 {
235                    return Err(Error::NegativeTally {
236                        index: i,
237                        value: *q,
238                    });
239                }
240                Ok(q * v)
241            })
242            .collect::<Result<Vec<_>, _>>()?;
243        if analog_pdf.iter().all(|&x| x == 0.0) {
244            return Err(Error::ZeroSumPdf);
245        }
246
247        let bias_pdf: Vec<f64> = match mode {
248            Mode::Analog => analog_pdf.clone(),
249            Mode::Uniform => volumes.clone(),
250            Mode::User => {
251                let user = user_pdf.ok_or(Error::EmptyPdf)?;
252                if user.len() != num_ves {
253                    return Err(Error::LengthMismatch {
254                        expected: num_ves,
255                        got: user.len(),
256                    });
257                }
258                user.iter()
259                    .zip(volumes.iter())
260                    .enumerate()
261                    .map(|(i, (q, v))| {
262                        if !q.is_finite() {
263                            return Err(Error::NonFiniteTally {
264                                field: "user_pdf",
265                                index: i,
266                            });
267                        }
268                        if *q < 0.0 {
269                            return Err(Error::NegativeTally {
270                                index: i,
271                                value: *q,
272                            });
273                        }
274                        Ok(q * v)
275                    })
276                    .collect::<Result<Vec<_>, _>>()?
277            }
278        };
279
280        let normalized_bias = normalized(&bias_pdf).ok_or(Error::ZeroSumPdf)?;
281        let biased_weights: Vec<f64> = match mode {
282            Mode::Analog => vec![1.0; num_ves],
283            _ => {
284                let normalized_analog = normalized(&analog_pdf).ok_or(Error::ZeroSumPdf)?;
285                normalized_analog
286                    .iter()
287                    .zip(normalized_bias.iter())
288                    .map(|(&a, &b)| if b > 0.0 { a / b } else { 1.0 })
289                    .collect()
290            }
291        };
292
293        Ok(Self {
294            dims: tally.dims(),
295            num_ves,
296            mode,
297            table: AliasTable::new(&normalized_bias)?,
298            biased_weights,
299        })
300    }
301
302    /// Sample a birth voxel with two uniforms in `[0, 1)`.
303    pub fn sample(&self, r1: f64, r2: f64) -> SampledVoxel {
304        let index = self.table.sample(r1, r2);
305        let nz = self.dims[2];
306        let ny = self.dims[1];
307        let k = index % nz;
308        let j = (index / nz) % ny;
309        let i = index / (nz * ny);
310        SampledVoxel {
311            index,
312            i,
313            j,
314            k,
315            weight: self.biased_weights[index],
316        }
317    }
318
319    /// The bias mode this sampler was constructed with.
320    pub fn mode(&self) -> Mode {
321        self.mode
322    }
323
324    /// Number of voxels in the sampling domain.
325    pub fn num_voxels(&self) -> usize {
326        self.num_ves
327    }
328
329    /// The underlying alias table (biased/analog PDF included).
330    pub fn table(&self) -> &AliasTable {
331        &self.table
332    }
333}
334
335fn cell_volumes(tally: &MeshTallyData) -> Vec<f64> {
336    let d = tally.dims();
337    let mut vols = Vec::with_capacity(tally.num_ves());
338    for i in 0..d[0] {
339        let dx = tally.x_bounds[i + 1] - tally.x_bounds[i];
340        for j in 0..d[1] {
341            let dy = tally.y_bounds[j + 1] - tally.y_bounds[j];
342            for k in 0..d[2] {
343                let dz = tally.z_bounds[k + 1] - tally.z_bounds[k];
344                vols.push(dx * dy * dz);
345            }
346        }
347    }
348    vols
349}
350
351fn normalized(pdf: &[f64]) -> Option<Vec<f64>> {
352    let sum: f64 = pdf.iter().sum();
353    if sum.is_nan() || sum <= 0.0 {
354        return None;
355    }
356    Some(pdf.iter().map(|&x| x / sum).collect())
357}
358
359#[cfg(test)]
360mod tests {
361    use super::*;
362    use crate::magic::magic;
363
364    /// Park–Miller minimal standard LCG; deterministic, returns values in
365    /// (0, 1) exclusive.
366    struct Lcg(u64);
367
368    impl Lcg {
369        fn next_f64(&mut self) -> f64 {
370            self.0 = (16807 * self.0) % 2147483647;
371            self.0 as f64 / 2147483647.0
372        }
373    }
374
375    fn approx(actual: f64, want: f64, tol: f64) {
376        assert!(
377            (actual - want).abs() <= tol * want.abs().max(1e-30),
378            "expected {want}, got {actual}"
379        );
380    }
381
382    #[test]
383    fn uniform_two_bin_pdf_builds_degenerate_table() {
384        // [0.5, 0.5]: after scaling both bins equal 1, no aliasing needed.
385        let t = AliasTable::new(&[0.5, 0.5]).unwrap();
386        assert_eq!(t.pdf(), &[0.5, 0.5]);
387        assert_eq!(t.prob, vec![1.0, 1.0]);
388        assert_eq!(t.len(), 2);
389        assert!(!t.is_empty());
390    }
391
392    #[test]
393    fn skewed_pdf_exact_outcome_probabilities_match_input() {
394        // The alias method is exact by construction: reconstructing P(i) from
395        // (prob, alias) must reproduce the normalized PDF to machine precision.
396        for pdf in [
397            vec![0.9, 0.06, 0.03, 0.01],
398            vec![0.5, 0.5],
399            vec![0.01, 0.01, 0.96],
400            vec![1.0, 2.0, 3.0, 4.0, 5.0],
401        ] {
402            let t = AliasTable::new(&pdf).unwrap();
403            let exact = t.exact_probabilities();
404            for (e, w) in exact.iter().zip(t.pdf().iter()) {
405                approx(*e, *w, 1e-12);
406            }
407        }
408    }
409
410    #[test]
411    fn skewed_nine_bin_table_exact_probabilities() {
412        let pdf = [0.9, 0.02, 0.02, 0.02, 0.01, 0.01, 0.008, 0.007, 0.005];
413        let t = AliasTable::new(&pdf).unwrap();
414        assert_eq!(t.len(), 9);
415        for (c, (&p, &a)) in t.prob.iter().zip(t.alias.iter()).enumerate() {
416            assert!((0.0..=1.0).contains(&p), "prob[{c}] = {p} outside [0, 1]");
417            assert!(a < 9, "alias[{c}] = {a} out of range");
418        }
419        let exact = t.exact_probabilities();
420        for (&e, &w) in exact.iter().zip(t.pdf().iter()) {
421            approx(e, w, 1e-12);
422        }
423    }
424
425    #[test]
426    fn sample_saturates_out_of_range_uniforms() {
427        let t = AliasTable::new(&[0.25, 0.75]).unwrap();
428        // r1 = 1.0 would read index n in the C++; here it clamps to n - 1.
429        // Column 1 has prob 1.0, so any finite r2 keeps index 1.
430        assert_eq!(t.sample(1.0, 0.999), 1);
431        // Negative r1 clamps to column 0 (prob 0.5); small r2 keeps it,
432        // large r2 takes its alias.
433        assert_eq!(t.sample(-3.0, 0.25), 0);
434        assert_eq!(t.sample(-3.0, 0.75), t.alias[0]);
435    }
436
437    #[test]
438    fn bad_pdfs_rejected() {
439        assert_eq!(AliasTable::new(&[]), Err(Error::EmptyPdf));
440        assert_eq!(
441            AliasTable::new(&[-0.1, 0.5]),
442            Err(Error::NegativePdf {
443                index: 0,
444                value: -0.1
445            })
446        );
447        assert!(matches!(
448            AliasTable::new(&[f64::NAN, 0.5]),
449            Err(Error::NonFinitePdf { index: 0 })
450        ));
451        assert_eq!(AliasTable::new(&[0.0, 0.0]), Err(Error::ZeroSumPdf));
452    }
453
454    #[test]
455    fn exhaustive_sampling_statistics_match_pdf() {
456        // 100k deterministic draws against the task's skewed [0.9, ...] pdf;
457        // every realized frequency within ~1% absolute of its probability.
458        let pdf = [0.9, 0.06, 0.03, 0.01];
459        let t = AliasTable::new(&pdf).unwrap();
460        let mut rng = Lcg(123_456_789);
461        let n_draws = 100_000;
462        let mut counts = [0usize; 4];
463        for _ in 0..n_draws {
464            let idx = t.sample(rng.next_f64(), rng.next_f64());
465            counts[idx] += 1;
466        }
467        for (i, &c) in counts.iter().enumerate() {
468            let freq = c as f64 / n_draws as f64;
469            assert!(
470                (freq - pdf[i]).abs() < 0.01,
471                "bin {i}: freq {freq} vs pdf {}",
472                pdf[i]
473            );
474        }
475    }
476
477    #[test]
478    fn lcg_coverage_hits_every_column_and_alias_branch() {
479        // With 16 bins, both outcomes of every column get exercised across
480        // 100k draws; all bins receive samples.
481        let pdf = [1.0f64; 16];
482        let t = AliasTable::new(&pdf).unwrap();
483        let mut rng = Lcg(987_654_321);
484        let mut seen = [false; 16];
485        for _ in 0..100_000 {
486            seen[t.sample(rng.next_f64(), rng.next_f64())] = true;
487        }
488        assert!(seen.iter().all(|&s| s));
489    }
490
491    fn fixture_tally(name: &str, num: u32) -> MeshTallyData {
492        let path = format!(
493            "{}/../../fixtures/mcnp/meshtal/{name}",
494            env!("CARGO_MANIFEST_DIR")
495        );
496        let m = nucleide_mcnp_io::meshtal::Meshtal::from_file(path).unwrap();
497        m.tallies[&num].clone()
498    }
499
500    #[test]
501    fn fixture_sampler_analog_pdf_proportional_to_totals_times_volume() {
502        let t = fixture_tally("mcnp_meshtal_single_meshtal.txt", 4);
503        let s = MeshSourceSampler::new(&t, Mode::Analog, None).unwrap();
504        assert_eq!(s.mode(), Mode::Analog);
505        assert_eq!(s.num_voxels(), t.num_ves());
506
507        let d = t.dims();
508        let mut expect = Vec::with_capacity(t.num_ves());
509        for ve in 0..t.num_ves() {
510            let k = ve % d[2];
511            let j = (ve / d[2]) % d[1];
512            let i = ve / (d[2] * d[1]);
513            let vol = (t.x_bounds[i + 1] - t.x_bounds[i])
514                * (t.y_bounds[j + 1] - t.y_bounds[j])
515                * (t.z_bounds[k + 1] - t.z_bounds[k]);
516            expect.push(t.total_result[ve] * vol);
517        }
518        let norm: f64 = expect.iter().sum();
519        let got = s.table().pdf();
520        for (&g, &w) in got.iter().zip(expect.iter()) {
521            approx(g, w / norm, 1e-12);
522        }
523        // Analog births carry unit weight regardless of draw.
524        let mut rng = Lcg(555);
525        for _ in 0..100 {
526            let sv = s.sample(rng.next_f64(), rng.next_f64());
527            approx(sv.weight, 1.0, 1e-15);
528        }
529    }
530
531    #[test]
532    fn magic_bounds_proportional_to_analog_pdf_on_equal_cells() {
533        // On an equal-volume mesh both quantities are proportional to the
534        // cell flux: MAGIC lower bound = flux / (2 * max_flux) (no nulls at
535        // default tolerance here), analog pdf = flux / sum(flux).
536        let t = MeshTallyData {
537            tally_number: 1,
538            particle: nucleide_mcnp_io::meshtal::ParticleKind::Neutron,
539            dose_response: false,
540            x_bounds: vec![0.0, 1.0],
541            y_bounds: vec![0.0, 1.0],
542            z_bounds: vec![0.0, 1.0, 2.0, 3.0],
543            e_bounds: vec![0.0, 1.0],
544            column_idx: Default::default(),
545            result: vec![vec![2.0], vec![1.0], vec![4.0]],
546            rel_error: vec![vec![0.1], vec![0.1], vec![0.1]],
547            total_result: vec![2.0, 1.0, 4.0],
548            total_rel_error: vec![0.1, 0.1, 0.1],
549        };
550        let ww = magic(&t).unwrap();
551        let s = MeshSourceSampler::new(&t, Mode::Analog, None).unwrap();
552        let pdf = s.table().pdf();
553
554        let ratio = ww.lower_bounds_ww[0] / pdf[0];
555        for (&w, &p) in ww.lower_bounds_ww[1..].iter().zip(pdf[1..].iter()) {
556            approx(w / p, ratio, 1e-12);
557        }
558        // Spot-check both normalizations independently.
559        approx(pdf[2], 4.0 / 7.0, 1e-12);
560        approx(ww.lower_bounds_ww[2], 0.5, 1e-12);
561    }
562
563    #[test]
564    fn uniform_mode_weights_cells_by_volume() {
565        // Two stacked x-cells of widths 1 and 3 with densities 10 and 1:
566        // analog PDF ∝ strength (density × volume) = [10, 3],
567        // uniform PDF ∝ volume = [1, 3].
568        let t = MeshTallyData {
569            tally_number: 1,
570            particle: nucleide_mcnp_io::meshtal::ParticleKind::Neutron,
571            dose_response: false,
572            x_bounds: vec![0.0, 1.0, 4.0],
573            y_bounds: vec![0.0, 1.0],
574            z_bounds: vec![0.0, 1.0],
575            e_bounds: vec![0.0, 1.0],
576            column_idx: Default::default(),
577            result: vec![vec![10.0], vec![1.0]],
578            rel_error: vec![vec![0.0], vec![0.0]],
579            total_result: vec![10.0, 1.0],
580            total_rel_error: vec![0.0, 0.0],
581        };
582        let uni = MeshSourceSampler::new(&t, Mode::Uniform, None).unwrap();
583        let ana = MeshSourceSampler::new(&t, Mode::Analog, None).unwrap();
584
585        let mut rng = Lcg(42_424_242);
586        let (mut u0, mut a0) = (0u32, 0u32);
587        for _ in 0..100_000 {
588            if uni.sample(rng.next_f64(), rng.next_f64()).index == 0 {
589                u0 += 1;
590            }
591            if ana.sample(rng.next_f64(), rng.next_f64()).index == 0 {
592                a0 += 1;
593            }
594        }
595        // Analog hits cell 0 w.p. 10/13; uniform only 1/4.
596        approx(u0 as f64 / 100_000.0, 0.25, 0.02);
597        approx(a0 as f64 / 100_000.0, 10.0 / 13.0, 0.02);
598
599        // Biased birth weights = analog pdf / biased pdf:
600        // narrow cell (10/13)/(1/4) = 40/13, wide cell (3/13)/(3/4) = 4/13.
601        approx(uni.biased_weights[0], 40.0 / 13.0, 1e-12);
602        approx(uni.biased_weights[1], 4.0 / 13.0, 1e-12);
603        approx(ana.biased_weights[0], 1.0, 1e-12);
604    }
605
606    #[test]
607    fn user_mode_respects_external_pdf_and_reports_weights() {
608        let t = MeshTallyData {
609            tally_number: 1,
610            particle: nucleide_mcnp_io::meshtal::ParticleKind::Photon,
611            dose_response: false,
612            x_bounds: vec![0.0, 1.0, 2.0],
613            y_bounds: vec![0.0, 1.0],
614            z_bounds: vec![0.0, 1.0],
615            e_bounds: vec![0.0, 1.0],
616            column_idx: Default::default(),
617            result: vec![vec![10.0], vec![0.001]],
618            rel_error: vec![vec![0.0], vec![0.0]],
619            total_result: vec![10.0, 0.001],
620            total_rel_error: vec![0.0, 0.0],
621        };
622        // Bias everything into the weakly-populated second cell.
623        let s = MeshSourceSampler::new(&t, Mode::User, Some(&[0.0, 5.0])).unwrap();
624        assert_eq!(s.mode(), Mode::User);
625        assert_eq!(s.table().pdf(), &[0.0, 1.0]);
626
627        let mut rng = Lcg(777_777);
628        for _ in 0..1_000 {
629            let sv = s.sample(rng.next_f64(), rng.next_f64());
630            assert_eq!(sv.index, 1);
631            assert_eq!((sv.i, sv.j, sv.k), (1, 0, 0));
632            // Birth weight = analog pdf / biased pdf ≈ 9.999e-5: births are
633            // concentrated where particles rarely stream, so each carries a
634            // small weight.
635            approx(sv.weight, 0.001 / 10.001, 1e-9);
636        }
637    }
638
639    #[test]
640    fn sampled_indices_decompose_consistently() {
641        let t = fixture_tally("mcnp_meshtal_single_meshtal.txt", 4);
642        let s = MeshSourceSampler::new(&t, Mode::Analog, None).unwrap();
643        let d = t.dims();
644        let mut rng = Lcg(20_260_824);
645        for _ in 0..10_000 {
646            let sv = s.sample(rng.next_f64(), rng.next_f64());
647            assert!(sv.i < d[0] && sv.j < d[1] && sv.k < d[2], "{sv:?}");
648            assert_eq!((sv.i * d[1] + sv.j) * d[2] + sv.k, sv.index);
649            assert_eq!(t.ve_index(sv.i, sv.j, sv.k), sv.index);
650        }
651    }
652
653    #[test]
654    fn user_pdf_length_mismatch_rejected() {
655        let t = fixture_tally("mcnp_meshtal_single_meshtal.txt", 4);
656        assert_eq!(
657            MeshSourceSampler::new(&t, Mode::User, Some(&[1.0; 3])),
658            Err(Error::LengthMismatch {
659                expected: 45,
660                got: 3
661            })
662        );
663    }
664
665    #[test]
666    fn all_zero_totals_rejected() {
667        let mut t = fixture_tally("mcnp_meshtal_single_meshtal.txt", 4);
668        t.total_result.iter_mut().for_each(|v| *v = 0.0);
669        assert_eq!(
670            MeshSourceSampler::new(&t, Mode::Analog, None),
671            Err(Error::ZeroSumPdf)
672        );
673    }
674
675    #[test]
676    fn negative_total_result_rejected() {
677        let mut t = fixture_tally("mcnp_meshtal_single_meshtal.txt", 4);
678        t.total_result[3] = -1.0;
679        assert_eq!(
680            MeshSourceSampler::new(&t, Mode::Analog, None),
681            Err(Error::NegativeTally {
682                index: 3,
683                value: -1.0
684            })
685        );
686    }
687
688    #[test]
689    fn negative_user_pdf_rejected() {
690        let t = fixture_tally("mcnp_meshtal_single_meshtal.txt", 4);
691        let mut user = vec![1.0; t.num_ves()];
692        user[7] = -0.5;
693        assert_eq!(
694            MeshSourceSampler::new(&t, Mode::User, Some(&user)),
695            Err(Error::NegativeTally {
696                index: 7,
697                value: -0.5
698            })
699        );
700    }
701
702    #[test]
703    fn non_finite_total_result_rejected() {
704        let mut t = fixture_tally("mcnp_meshtal_single_meshtal.txt", 4);
705        t.total_result[5] = f64::NAN;
706        assert!(matches!(
707            MeshSourceSampler::new(&t, Mode::Analog, None),
708            Err(Error::NonFiniteTally {
709                field: "total_result",
710                index: 5
711            })
712        ));
713
714        t.total_result[5] = f64::INFINITY;
715        assert!(matches!(
716            MeshSourceSampler::new(&t, Mode::Analog, None),
717            Err(Error::NonFiniteTally {
718                field: "total_result",
719                index: 5
720            })
721        ));
722    }
723
724    #[test]
725    fn total_result_length_mismatch_rejected() {
726        let mut t = fixture_tally("mcnp_meshtal_single_meshtal.txt", 4);
727        t.total_result.truncate(10);
728        assert_eq!(
729            MeshSourceSampler::new(&t, Mode::Analog, None),
730            Err(Error::LengthMismatch {
731                expected: 45,
732                got: 10
733            })
734        );
735    }
736
737    #[test]
738    fn analog_mode_results_unchanged_for_positive_tally() {
739        // Removing the silent abs() should not alter behavior when all values
740        // are already non-negative. Equal volumes keep the PDF proportional
741        // to the raw totals.
742        let t = MeshTallyData {
743            tally_number: 1,
744            particle: nucleide_mcnp_io::meshtal::ParticleKind::Neutron,
745            dose_response: false,
746            x_bounds: vec![0.0, 1.0, 2.0],
747            y_bounds: vec![0.0, 1.0],
748            z_bounds: vec![0.0, 1.0],
749            e_bounds: vec![0.0, 1.0],
750            column_idx: Default::default(),
751            result: vec![vec![10.0], vec![1.0]],
752            rel_error: vec![vec![0.0], vec![0.0]],
753            total_result: vec![10.0, 1.0],
754            total_rel_error: vec![0.0, 0.0],
755        };
756        let s = MeshSourceSampler::new(&t, Mode::Analog, None).unwrap();
757        assert_eq!(s.mode(), Mode::Analog);
758        approx(s.table().pdf()[0], 10.0 / 11.0, 1e-12);
759        approx(s.table().pdf()[1], 1.0 / 11.0, 1e-12);
760        assert!(s.biased_weights.iter().all(|&w| (w - 1.0).abs() < 1e-15));
761    }
762}