1use std::borrow::Cow;
8use std::cell::RefCell;
9use std::ops::{Range, RangeInclusive};
10use std::sync::Arc;
11
12use faer::{Mat, Side};
13use nereids_endf::resonance::{ResonanceCovariance, ResonanceData, RmlRange};
14use nereids_endf::retrieval::EndfLibrary;
15use nereids_fitting::error::FittingError;
16use nereids_fitting::lm::{FitModel, FlatMatrix};
17use nereids_fitting::parameters::{FitParameter, ParameterSet};
18use nereids_fitting::poisson::{
19 NEWTON_DECREMENT_TOL, Prior, Unbounded, half_deviance, information_inverse,
20};
21use nereids_fitting::statistics::{Consistency, consistency};
22use nereids_physics::continuous_doppler::{
23 Channel, QuadratureBudget, SUPPORT_X, TierOneBroadening, broaden_with_budget,
24};
25use nereids_physics::detector_spread::region_shares;
26use nereids_physics::doppler::DopplerParams;
27use nereids_physics::flight_time_grid::{FlightTimeGrid, Rows};
28use nereids_physics::forward_scatter::{ForwardYield, forward_yield, lowest_exit_ratio};
29use nereids_physics::ikeda_carpenter::IkedaCarpenterParams;
30use nereids_physics::lattice_doppler::{PhononSpectrum, broaden_in_lattice};
31use nereids_physics::resolution::TOF_FACTOR;
32use nereids_physics::thickness::{interacting_per_depth, uneven_transmission};
33use nereids_physics::transmission::resonance_center_energies;
34use rayon::prelude::*;
35
36use crate::beam::BeamSpline;
37use crate::error::PipelineError;
38use crate::open_beam::{
39 BOUND, COUNTS_TO_MEASURE_NOISE, Calibration, OpenBeamFit, PULSE_NUMBERS, Pulse, Recorded,
40 combined, counted, fit_on_halved_grids, laws, open_beam_with_intervals, overdispersion,
41 validate_counts, validate_live, weights_of,
42};
43
44pub const NEGLIGIBLE_PREDICTION: f64 = 1e-10;
47
48const SETTLED_OVERDISPERSION: f64 = 0.01;
49const MOST_SPREAD_CELLS: usize = 1 << 20;
51
52pub(crate) const MOST_PASSES: usize = 20;
53
54pub(crate) const TEMPERATURE_BOUNDS_K: (f64, f64) = (1.0, 5000.0);
57
58const SECOND_START_K: f64 = TEMPERATURE_BOUNDS_K.1 / 2.0;
60
61const REJECTING_SDS: f64 = 3.0;
64
65#[derive(Debug, Clone, Copy, PartialEq)]
70pub enum Value {
71 Known(f64),
72 Fitted(f64),
73 Within { start: f64, lower: f64, upper: f64 },
74 Measured { value: f64, sd: f64 },
75}
76
77impl Value {
78 pub(crate) fn parameter(
79 self,
80 name: impl Into<Cow<'static, str>>,
81 range: RangeInclusive<f64>,
82 allowed: &str,
83 ) -> Result<FitParameter, PipelineError> {
84 let name = name.into();
85 let (value, lower, upper) = match self {
86 Self::Known(v) | Self::Fitted(v) | Self::Measured { value: v, .. } => {
87 (v, *range.start(), *range.end())
88 }
89 Self::Within {
90 start,
91 lower,
92 upper,
93 } => (start, lower, upper),
94 };
95 if !(value.is_finite()
96 && range.contains(&lower)
97 && lower < upper
98 && range.contains(&upper)
99 && (lower..=upper).contains(&value)
100 && !matches!(self, Self::Measured { sd, .. } if !(sd.is_finite() && sd > 0.0)))
101 {
102 return Err(PipelineError::InvalidParameter(match self {
103 Self::Within { .. } => format!(
104 "{name} must be {allowed}, with bounds lower < upper and a finite start \
105 between them; got {self:?}"
106 ),
107 Self::Measured { .. } => format!(
108 "{name} must be finite and {allowed}, with a finite positive sd; got {self:?}"
109 ),
110 _ => format!("{name} must be finite and {allowed}; got {self:?}"),
111 }));
112 }
113 Ok(FitParameter {
114 name,
115 value,
116 lower,
117 upper,
118 fixed: matches!(self, Self::Known(_)),
119 })
120 }
121}
122
123#[derive(Debug, Clone)]
126pub struct Measurement {
127 pub time_edges_us: Vec<f64>,
129 pub charge_ratio: f64,
131 pub normalization: Value,
134 pub time_background: [Value; 3],
140 pub room_background: Option<RoomBackground>,
143 pub spread: Option<Spread>,
146 pub regions: Vec<Region>,
148}
149
150pub const NO_TIME_BACKGROUND: [Value; 3] =
152 [Value::Known(0.0), Value::Known(0.0), Value::Known(1.0)];
153
154#[derive(Debug, Clone, Copy, PartialEq)]
160pub struct RoomBackground {
161 pub rate: Value,
163 pub frames: [f64; 2],
166}
167
168#[derive(Debug, Clone, PartialEq)]
177pub struct Spread {
178 pub share: Value,
180 pub kernel: Vec<(f64, f64)>,
183 pub grid: [usize; 2],
185 pub cells: Vec<Vec<[usize; 2]>>,
188}
189
190#[derive(Debug, Clone)]
193pub struct Region {
194 pub open_counts: Vec<f64>,
196 pub sample_counts: Vec<f64>,
198 pub open_live: Option<Vec<f64>>,
201 pub sample_live: Option<Vec<f64>>,
204 pub background: [Value; 5],
209 pub pixels: f64,
211 pub gamma_transmission: Value,
214 pub capture_gammas: Value,
220 pub in_scatter: Value,
227 pub timing_offset_us: Value,
231 pub material: Option<Material>,
234}
235
236#[derive(Debug, Clone)]
238pub struct Material {
239 pub isotopes: Vec<(ResonanceData, Value)>,
242 pub temperature_k: Value,
245 pub thickness_variance: Value,
251 pub library: EndfLibrary,
257 pub covariances: Vec<(usize, ResonanceCovariance)>,
261 pub lattice: Option<PhononSpectrum>,
265}
266
267#[derive(Debug, Clone)]
270pub struct CountsFit {
271 pub regions: Vec<RegionFit>,
273 pub normalization: f64,
275 pub t0_us: f64,
277 pub flight_path_m: f64,
279 pub alpha: [f64; 2],
281 pub beta: [f64; 2],
283 pub r: f64,
285 pub fwhm_squared_us2: f64,
288 pub r_high: Option<f64>,
290 pub detector_tail: Option<[f64; 2]>,
293 pub time_background: [f64; 3],
296 pub room_background: f64,
299 pub spread_share: f64,
301 pub covariance: Option<FlatMatrix>,
333 pub unbounded: Option<Unbounded>,
338 pub on_bound: Vec<bool>,
341 pub deviance: f64,
345 pub converged: bool,
350 pub measured_pulls: Option<Vec<f64>>,
361 pub pulse_consistency: Option<Consistency>,
374 pub step_us: f64,
376 pub points: usize,
378 pub halvings: usize,
383}
384
385#[derive(Debug, Clone)]
387pub struct RegionFit {
388 pub densities: Vec<f64>,
392 pub temperature_k: Option<f64>,
395 pub thickness_variance: Option<f64>,
398 pub background: [f64; 5],
401 pub gamma_transmission: f64,
403 pub capture_gammas: f64,
405 pub in_scatter: f64,
407 pub timing_offset_us: f64,
409 pub beam: BeamSpline,
413 pub beam_at_limit: bool,
416 pub predicted: [Vec<f64>; 2],
420 pub overdispersion: [Option<f64>; 2],
430 pub excluded_bins: Vec<usize>,
436}
437
438pub fn fit_counts(
606 measurement: &Measurement,
607 calibration: &Calibration,
608) -> Result<CountsFit, PipelineError> {
609 if measurement
610 .regions
611 .iter()
612 .all(|region| region.material.is_none())
613 {
614 return Err(PipelineError::InvalidParameter(
615 "no region of the measurement holds a material to fit".into(),
616 ));
617 }
618 let first = fit_settled(measurement, calibration, MOST_PASSES)?;
619 let Some(started_hot) = rejected(measurement, &first)
620 .then(|| started_hot(measurement))
621 .flatten()
622 else {
623 return Ok(first);
624 };
625 Ok(match fit_settled(&started_hot, calibration, MOST_PASSES) {
626 Ok(second) if better(measurement, &second, &first) => second,
627 _ => first,
628 })
629}
630
631fn rejected(measurement: &Measurement, fit: &CountsFit) -> bool {
635 measurement
636 .regions
637 .iter()
638 .zip(&fit.regions)
639 .any(|(region, at)| {
640 let bins = (region.sample_counts.len() - at.excluded_bins.len()) as f64;
641 at.overdispersion[1].is_some_and(|phi| phi > 1.0 + REJECTING_SDS * (2.0 / bins).sqrt())
642 })
643}
644
645fn started_hot(measurement: &Measurement) -> Option<Measurement> {
648 let mut hot = measurement.clone();
649 let mut moved = false;
650 for material in hot.regions.iter_mut().filter_map(|r| r.material.as_mut()) {
651 let start = match material.temperature_k {
652 Value::Fitted(_) => SECOND_START_K,
653 Value::Within { lower, upper, .. } => SECOND_START_K.max(lower).min(upper),
654 Value::Known(_) | Value::Measured { .. } => continue,
655 };
656 if let Value::Fitted(old) | Value::Within { start: old, .. } = &mut material.temperature_k {
657 moved |= *old != start;
658 *old = start;
659 }
660 }
661 moved.then_some(hot)
662}
663
664fn better(measurement: &Measurement, fit: &CountsFit, other: &CountsFit) -> bool {
667 if fit.converged != other.converged {
668 return fit.converged;
669 }
670 counts_deviance(measurement, fit, other) + NEWTON_DECREMENT_TOL
671 < counts_deviance(measurement, other, fit)
672}
673
674fn counts_deviance(measurement: &Measurement, fit: &CountsFit, other: &CountsFit) -> f64 {
677 let runs = measurement
678 .regions
679 .iter()
680 .zip(&fit.regions)
681 .zip(&other.regions);
682 runs.map(|((region, at), beside)| {
683 let kept = |k: &usize| {
684 at.excluded_bins.binary_search(k).is_err()
685 && beside.excluded_bins.binary_search(k).is_err()
686 };
687 let open: f64 = (region.open_counts.iter().zip(&at.predicted[0]))
688 .map(|(&y, &mu)| half_deviance(y, mu))
689 .sum();
690 let sample: f64 = (region
691 .sample_counts
692 .iter()
693 .zip(&at.predicted[1])
694 .enumerate())
695 .filter(|(k, _)| kept(k))
696 .map(|(_, (&y, &mu))| half_deviance(y, mu))
697 .sum();
698 open + sample
699 })
700 .sum()
701}
702
703fn fit_settled(
704 measurement: &Measurement,
705 calibration: &Calibration,
706 most_passes: usize,
707) -> Result<CountsFit, PipelineError> {
708 let bins = starts_and_widths_us(&measurement.time_edges_us);
709 let frames = open_frames(measurement.room_background.as_ref());
710 let mut open = starting_bypass(measurement)?;
711 let mut offsets = (0..measurement.regions.len())
712 .map(|region| starting(measurement, Role::TimingOffset { region }))
713 .collect::<Result<Vec<f64>, PipelineError>>()?;
714 let mut intervals: Option<Vec<usize>> = None;
715 let shares = spread_shares(measurement)?;
716 for passes in 1.. {
717 let (mut fit, _) = fit_counts_answer(
718 measurement,
719 calibration,
720 BeamStarts::OpenFitsAt(open, intervals.clone(), offsets.clone()),
721 shares.clone(),
722 )?;
723 let fitted = Bypass {
724 time: fit.time_background,
725 room: fit.room_background,
726 };
727 let fitted_offsets: Vec<f64> = fit.regions.iter().map(|r| r.timing_offset_us).collect();
728 let mut spread = 0.0;
729 for (q, (r, region)) in fit.regions.iter().zip(&measurement.regions).enumerate() {
730 let used = open.open_counts_per_pixel(&late(&bins, offsets[q]), frames);
731 let now = fitted.open_counts_per_pixel(&late(&bins, fitted_offsets[q]), frames);
732 for (k, &mu) in r.predicted[0]
733 .iter()
734 .enumerate()
735 .filter(|(_, mu)| **mu > 0.0)
736 {
737 let live = region.open_live.as_ref().map_or(1.0, |live| live[k]);
738 spread += (live * region.pixels * (now[k] - used[k])).powi(2) / mu;
739 }
740 }
741 if !fit.converged || spread <= BOUND || passes == most_passes {
742 if spread > BOUND {
743 fit.converged = false;
744 fit.covariance = None;
745 fit.unbounded = None;
746 fit.measured_pulls = None;
747 fit.pulse_consistency = None;
748 }
749 return Ok(fit);
750 }
751 if passes >= 2 {
752 intervals
753 .get_or_insert_with(|| fit.regions.iter().map(|r| r.beam.intervals()).collect());
754 }
755 open = fitted;
756 offsets = fitted_offsets;
757 }
758 unreachable!("the passes end at most_passes")
759}
760
761#[derive(Debug, Clone, Copy, PartialEq)]
764pub(crate) struct Bypass {
765 pub(crate) time: [f64; 3],
766 pub(crate) room: f64,
767}
768
769impl Bypass {
770 pub(crate) fn open_counts_per_pixel(&self, bins: &[(f64, f64)], frames: f64) -> Vec<f64> {
773 let (time, _) = time_counts_per_pixel(bins, self.time);
774 time.iter()
775 .zip(bins)
776 .map(|(g, (_, width))| g + self.room * frames * width)
777 .collect()
778 }
779}
780
781pub(crate) fn open_frames(room: Option<&RoomBackground>) -> f64 {
783 room.map_or(0.0, |room| room.frames[0])
784}
785
786fn starting(measurement: &Measurement, role: Role) -> Result<f64, PipelineError> {
788 let (value, name, range, allowed) = quantity(measurement, role).expect("a quantity");
789 Ok(value.parameter(name, range, &allowed)?.value)
790}
791
792fn late(bins: &[(f64, f64)], late_us: f64) -> Vec<(f64, f64)> {
794 bins.iter()
795 .map(|&(start, width)| (start - late_us, width))
796 .collect()
797}
798
799pub(crate) fn starting_bypass(measurement: &Measurement) -> Result<Bypass, PipelineError> {
800 let start = |role| starting(measurement, role);
801 Ok(Bypass {
802 time: [
803 start(Role::TimeBackground(0))?,
804 start(Role::TimeBackground(1))?,
805 start(Role::TimeBackground(2))?,
806 ],
807 room: start(Role::RoomBackground)?,
808 })
809}
810
811pub(crate) struct Starts {
812 pub(crate) beams: Vec<BeamStart>,
813 pub(crate) origin_us: f64,
814 pub(crate) shared_leverage: Option<Vec<f64>>,
815 pub(crate) bound: f64,
816 pub(crate) excluded: Vec<Vec<usize>>,
817}
818
819#[derive(Debug, Clone)]
820pub(crate) struct BeamStart {
821 pub(crate) beam: BeamSpline,
822 pub(crate) overdispersion: Option<f64>,
823 pub(crate) at_limit: bool,
824}
825
826impl From<OpenBeamFit> for BeamStart {
827 fn from(open: OpenBeamFit) -> Self {
828 Self {
829 beam: open.beam,
830 overdispersion: open.overdispersion,
831 at_limit: open.at_limit,
832 }
833 }
834}
835
836pub(crate) struct Checked<'a> {
837 shape: Vec<Option<usize>>,
838 sample_parameters: Vec<FitParameter>,
839 measured: Vec<(usize, f64, f64)>,
840 lines: Vec<Lines<'a>>,
841 instrument: Vec<FitParameter>,
842 offsets: Vec<f64>,
843 base: FlightTimeGrid,
844 resonances_in_span: Vec<Vec<Vec<f64>>>,
845}
846
847pub(crate) fn checked<'a>(
848 measurement: &'a Measurement,
849 calibration: &Calibration,
850) -> Result<Checked<'a>, PipelineError> {
851 let Measurement {
852 time_edges_us,
853 charge_ratio,
854 regions,
855 ..
856 } = measurement;
857 let invalid = |message: String| Err(PipelineError::InvalidParameter(message));
858 let materials = || {
859 regions
860 .iter()
861 .enumerate()
862 .filter_map(|(r, region)| Some((r, region.material.as_ref()?)))
863 };
864 if !(charge_ratio.is_finite() && *charge_ratio > 0.0) {
865 return invalid(format!(
866 "the charge ratio must be finite and positive, got {charge_ratio}"
867 ));
868 }
869 if let Some(room) = &measurement.room_background
870 && !room.frames.iter().all(|f| f.is_finite() && *f > 0.0)
871 {
872 return invalid(format!(
873 "the room background's frames must be finite and positive, got {:?}",
874 room.frames
875 ));
876 }
877 if let Some(spread) = &measurement.spread {
878 spread_checked(spread, regions.len())?;
879 if spread.share != Value::Known(0.0)
880 && !regions.iter().all(|region| {
881 matches!(region.timing_offset_us, Value::Known(_))
882 && region.timing_offset_us == regions[0].timing_offset_us
883 })
884 {
885 return invalid(
886 "a detector spread needs every region's timing offset known and the same".into(),
887 );
888 }
889 }
890 let free = |value: &Value| matches!(value, Value::Fitted(_) | Value::Within { .. });
891 if free(&calibration.t0_us)
892 && !regions.is_empty()
893 && regions.iter().all(|region| free(®ion.timing_offset_us))
894 {
895 return invalid(
896 "t0 and every region's timing offset are fitted, and they move together: know or \
897 measure t0 or one region's offset"
898 .into(),
899 );
900 }
901 if let Some((r, region)) = regions
902 .iter()
903 .enumerate()
904 .find(|(_, region)| !(region.pixels.is_finite() && region.pixels > 0.0))
905 {
906 return invalid(format!(
907 "the pixels of region {r} must be finite and positive, got {}",
908 region.pixels
909 ));
910 }
911 if let Some(r) = regions.iter().position(|region| {
912 region.material.is_none()
913 && [region.capture_gammas, region.in_scatter] != [Value::Known(0.0); 2]
914 }) {
915 return invalid(format!(
916 "region {r} holds no material to capture or scatter neutrons, so its capture gammas \
917 and in-scatter must be known 0"
918 ));
919 }
920 for (r, material) in materials() {
921 if material.isotopes.is_empty() {
922 return invalid(format!("the material of region {r} has no isotopes"));
923 }
924 for (i, (isotope, _)) in material.isotopes.iter().enumerate() {
925 if material.isotopes[..i]
926 .iter()
927 .any(|(other, _)| other.za == isotope.za)
928 {
929 return invalid(format!(
930 "{} is listed twice in region {r}; its densities cannot be told apart",
931 isotope.isotope
932 ));
933 }
934 if !finite(isotope) {
935 return invalid(format!(
936 "the resonance data of {} in region {r} are not finite",
937 isotope.isotope
938 ));
939 }
940 }
941 for (k, (i, covariance)) in material.covariances.iter().enumerate() {
942 let n = 3 * covariance.resonances.len();
943 let listed = &covariance.resonances;
944 let placed = material.isotopes.get(*i).is_some_and(|(data, _)| {
945 !listed.is_empty()
946 && listed.iter().enumerate().all(|(x, &[a, b, c])| {
947 !listed[..x].contains(&[a, b, c])
948 && data
949 .ranges
950 .get(a)
951 .and_then(|range| range.l_groups.get(b))
952 .is_some_and(|group| c < group.resonances.len())
953 })
954 });
955 let matrix = &covariance.covariance;
956 let sd = |p: usize| matrix[p * n + p].abs().sqrt();
957 let symmetric = matrix.len() == n * n
958 && (0..n).all(|p| {
959 (0..n).all(|q| {
960 let scale = sd(p) * sd(q);
961 matrix[p * n + q].is_finite()
962 && (matrix[p * n + q] - matrix[q * n + p]).abs() <= 1e-12 * scale
963 && (p == q || scale > 0.0 || matrix[p * n + q] == 0.0)
964 })
965 });
966 let semidefinite = placed
967 && symmetric
968 && Mat::<f64>::from_fn(n, n, |p, q| {
969 let scale = sd(p) * sd(q);
970 if scale > 0.0 {
971 matrix[p * n + q] / scale
972 } else {
973 matrix[p * n + q]
974 }
975 })
976 .self_adjoint_eigenvalues(Side::Lower)
977 .is_ok_and(|e| e[0] >= -f64::EPSILON * n as f64 * e[n - 1].max(1.0));
978 if material.covariances[..k].iter().any(|(j, _)| j == i) || !semidefinite {
979 return invalid(format!(
980 "covariance {k} of region {r} must name an isotope of its material once and \
981 distinct resonances it has, with a finite symmetric positive semi-definite \
982 matrix over their energies and widths"
983 ));
984 }
985 }
986 }
987 let carried: Vec<(usize, u32)> = materials()
988 .flat_map(|(r, m)| {
989 m.covariances
990 .iter()
991 .map(move |(i, _)| (r, m.isotopes[*i].0.za))
992 })
993 .collect();
994 if materials().any(|(s, m)| {
995 m.isotopes
996 .iter()
997 .any(|(d, _)| carried.iter().any(|&(r, za)| r != s && za == d.za))
998 }) {
999 return invalid(
1000 "a resonance covariance is carried for an isotope another region also holds, whose \
1001 counts it would leave out"
1002 .into(),
1003 );
1004 }
1005 if calibration.pulse.prior.is_some() && materials().any(|(_, m)| !m.covariances.is_empty()) {
1006 return invalid(
1007 "a pulse calibration's prior records no sensitivity to the resonance parameters, so a \
1008 fit under it cannot carry their covariance"
1009 .into(),
1010 );
1011 }
1012 let shape = shape(measurement);
1013 let quantity_roles = roles(&vec![0; regions.len()], &shape);
1014 let sample_parameters = quantity_roles
1015 .iter()
1016 .filter_map(|&role| {
1017 let (value, name, range, allowed) = quantity(measurement, role)?;
1018 Some(value.parameter(name, range, &allowed))
1019 })
1020 .collect::<Result<Vec<FitParameter>, PipelineError>>()?;
1021 let measured: Vec<(usize, f64, f64)> = quantities(measurement, calibration)
1022 .enumerate()
1023 .filter_map(|(offset, (_, value))| match *value {
1024 Value::Measured { value, sd } => Some((offset, value, sd)),
1025 _ => None,
1026 })
1027 .collect();
1028 let lines: Vec<Lines<'_>> = materials()
1029 .map(|(region, material)| {
1030 let temperature = &sample_parameters[quantity_roles
1031 .iter()
1032 .position(|&role| role == Role::Temperature { region })
1033 .expect("a material has a temperature")];
1034 let reach_k = if temperature.fixed {
1035 temperature.value
1036 } else {
1037 temperature.upper
1038 };
1039 let dopplers = material
1040 .isotopes
1041 .iter()
1042 .map(|(isotope, _)| {
1043 DopplerParams::new(reach_k, isotope.awr).map_err(|e| {
1044 in_region(region, PipelineError::InvalidParameter(e.to_string()))
1045 })
1046 })
1047 .collect::<Result<Vec<_>, _>>()?;
1048 Ok(Lines {
1049 region,
1050 isotopes: &material.isotopes,
1051 dopplers,
1052 reach_k,
1053 lattice: material.lattice.as_ref(),
1054 scatters: measurement.regions[region].in_scatter != Value::Known(0.0),
1055 })
1056 })
1057 .collect::<Result<_, PipelineError>>()?;
1058 let instrument = calibration.instrument()?;
1059 validate_pulse_prior(calibration)?;
1060 let start: Vec<f64> = instrument.iter().map(|parameter| parameter.value).collect();
1061 let offsets: Vec<f64> = (0..regions.len())
1062 .map(|region| {
1063 let at = quantity_roles
1064 .iter()
1065 .position(|&role| role == Role::TimingOffset { region })
1066 .expect("every region has a timing offset");
1067 sample_parameters[at].value
1068 })
1069 .collect();
1070 let starts: Vec<f64> = offsets.iter().map(|offset| start[0] + offset).collect();
1071 let base = FlightTimeGrid::spanning(
1072 time_edges_us,
1073 span(&starts),
1074 start[1],
1075 &calibration.pulse.at(&start[2..])?,
1076 )?;
1077 calibrated_lines(&lines, &calibration.pulse, time_edges_us, &starts, start[1])?;
1078 let resonances_in_span = in_span(&lines, &base)?;
1079 Ok(Checked {
1080 shape,
1081 sample_parameters,
1082 measured,
1083 lines,
1084 instrument,
1085 offsets,
1086 base,
1087 resonances_in_span,
1088 })
1089}
1090
1091pub(crate) fn excluded_bins(
1092 time_edges_us: &[f64],
1093 material: Option<&Material>,
1094 numbers: &[f64],
1095 storage_span: Option<(f64, f64)>,
1096) -> Vec<usize> {
1097 let laws = laws(&numbers[2..], storage_span);
1098 let arrival_us = |energy: f64| {
1099 numbers[0]
1100 + TOF_FACTOR * numbers[1] / energy.sqrt()
1101 + 3.0 / laws.alpha.eval(energy)
1102 + laws.r.eval(energy) / laws.beta.eval(energy)
1103 + laws
1104 .detector_tail
1105 .map_or(0.0, |tail| tail.fraction * tail.decay_us)
1106 };
1107 let spans: Vec<(f64, f64)> = material
1108 .into_iter()
1109 .flat_map(|m| {
1110 m.isotopes
1111 .iter()
1112 .filter(|(_, density)| *density != Value::Known(0.0))
1113 .flat_map(|(data, _)| m.library.excluded_windows_ev(&data.isotope))
1114 })
1115 .map(|&(low, high)| (arrival_us(high), arrival_us(low)))
1116 .collect();
1117 (0..time_edges_us.len() - 1)
1118 .filter(|&k| {
1119 spans
1120 .iter()
1121 .any(|&(first, last)| time_edges_us[k] < last && first < time_edges_us[k + 1])
1122 })
1123 .collect()
1124}
1125
1126type Slopes = (Vec<Vec<f64>>, FlatMatrix);
1127
1128fn nuclear_slopes(
1129 regions: &[Region],
1130 model: &Recorded<'_, RegionsModel>,
1131 dispersion: &[f64],
1132 parameters: &ParameterSet,
1133) -> Result<Option<Slopes>, PipelineError> {
1134 let (params, free) = (parameters.all_values(), parameters.free_indices());
1135 let mu = model.evaluate(¶ms)?;
1136 let jacobian = model
1137 .analytical_jacobian(¶ms, &free, &mu)
1138 .expect("the counts model has an analytical Jacobian");
1139 let mut picked: Vec<(&ResonanceCovariance, usize, Vec<f64>)> = Vec::new();
1140 for (r, material) in regions
1141 .iter()
1142 .enumerate()
1143 .filter_map(|(r, region)| Some((r, region.material.as_ref()?)))
1144 {
1145 for (i, covariance) in &material.covariances {
1146 let data = &material.isotopes[*i].0;
1147 let n = 3 * covariance.resonances.len();
1148 for p in (0..n).filter(|p| covariance.covariance[p * n + p] > 0.0) {
1149 let [a, b, c] = covariance.resonances[p / 3];
1150 let h = 0.1 * covariance.covariance[p * n + p].sqrt();
1151 let shifted = |by: f64| {
1152 let mut copy = data.clone();
1153 let resonance = &mut copy.ranges[a].l_groups[b].resonances[c];
1154 *[&mut resonance.energy, &mut resonance.gn, &mut resonance.gg][p % 3] += by;
1155 copy
1156 };
1157 let slope =
1158 model
1159 .model
1160 .resonance_slope(¶ms, r, *i, [&shifted(h), &shifted(-h)], h)?;
1161 let cross = (0..free.len())
1162 .map(|j| {
1163 slope
1164 .iter()
1165 .enumerate()
1166 .filter(|(k, _)| mu[*k] > 0.0)
1167 .map(|(k, c)| {
1168 jacobian.get(k, j) * c * model.live[k] / (mu[k] * dispersion[k])
1169 })
1170 .sum()
1171 })
1172 .collect();
1173 picked.push((covariance, p, cross));
1174 }
1175 }
1176 }
1177 let width = picked.len();
1178 if width == 0 {
1179 return Ok(None);
1180 }
1181 let mut prior = FlatMatrix::zeros(width, width);
1182 for (a, (first, p, _)) in picked.iter().enumerate() {
1183 for (b, (second, q, _)) in picked.iter().enumerate() {
1184 if std::ptr::eq(*first, *second) {
1185 *prior.get_mut(a, b) = first.covariance[p * 3 * first.resonances.len() + q];
1186 }
1187 }
1188 }
1189 Ok(Some((
1190 picked.into_iter().map(|(_, _, cross)| cross).collect(),
1191 prior,
1192 )))
1193}
1194
1195fn sandwich(conditional: &FlatMatrix, slopes: &Slopes) -> FlatMatrix {
1196 let (crosses, prior) = slopes;
1197 let (free, width) = (conditional.nrows, crosses.len());
1198 let resolved: Vec<bool> = (0..free)
1199 .map(|i| conditional.get(i, i).is_finite())
1200 .collect();
1201 let sensitivity: Vec<Vec<f64>> = (0..free)
1202 .map(|i| {
1203 crosses
1204 .iter()
1205 .map(|cross| {
1206 if !resolved[i] {
1207 return f64::NAN;
1208 }
1209 -(0..free)
1210 .filter(|&j| resolved[j])
1211 .map(|j| conditional.get(i, j) * cross[j])
1212 .sum::<f64>()
1213 })
1214 .collect()
1215 })
1216 .collect();
1217 let weighted: Vec<Vec<f64>> = sensitivity
1218 .iter()
1219 .map(|sx| {
1220 (0..width)
1221 .map(|b| (0..width).map(|a| sx[a] * prior.get(a, b)).sum())
1222 .collect()
1223 })
1224 .collect();
1225 let mut term = FlatMatrix::zeros(free, free);
1226 for (x, wx) in weighted.iter().enumerate() {
1227 for (y, sy) in sensitivity.iter().enumerate() {
1228 *term.get_mut(x, y) = wx.iter().zip(sy).map(|(w, s)| w * s).sum();
1229 }
1230 }
1231 term
1232}
1233
1234fn span(starts: &[f64]) -> (f64, f64) {
1236 starts
1237 .iter()
1238 .fold((f64::INFINITY, f64::NEG_INFINITY), |(low, high), &t| {
1239 (low.min(t), high.max(t))
1240 })
1241}
1242
1243fn calibrated_lines(
1244 lines: &[Lines<'_>],
1245 pulse: &Pulse,
1246 time_edges_us: &[f64],
1247 starts: &[f64],
1248 flight_path_m: f64,
1249) -> Result<(), PipelineError> {
1250 for lines in lines {
1251 if let Some(((low, high), line)) = pulse.line_span_ev.zip(pulse.uncalibrated_line(
1252 lines.isotopes,
1253 time_edges_us,
1254 starts[lines.region],
1255 flight_path_m,
1256 )) {
1257 return Err(PipelineError::InvalidParameter(format!(
1258 "region {} has a resonance at {line} eV, outside the {low}–{high} eV of the \
1259 resonances the pulse was calibrated on",
1260 lines.region
1261 )));
1262 }
1263 }
1264 Ok(())
1265}
1266
1267fn in_span(
1268 lines: &[Lines<'_>],
1269 grid: &FlightTimeGrid,
1270) -> Result<Vec<Vec<Vec<f64>>>, PipelineError> {
1271 let energies = grid.energies_ev();
1272 let span_ev = (energies[energies.len() - 1], energies[0]);
1273 lines
1274 .iter()
1275 .map(|lines| {
1276 if lines.scatters {
1277 let lowest = span_ev.0
1278 * lines
1279 .isotopes
1280 .iter()
1281 .map(|(isotope, _)| lowest_exit_ratio(isotope.awr))
1282 .fold(1.0, f64::min);
1283 if let Some((isotope, e)) = lines.isotopes.iter().find_map(|(isotope, _)| {
1284 resonance_center_energies(&[isotope])
1285 .into_iter()
1286 .find(|e| (lowest..span_ev.0).contains(e))
1287 .map(|e| (isotope, e))
1288 }) {
1289 return Err(PipelineError::InvalidParameter(format!(
1290 "the in-scatter of region {} leaves with energies down to {lowest:.6e} \
1291 eV, below the flight-time grid's lowest {:.6e} eV, and {} has a \
1292 resonance at {e:.6e} eV there",
1293 lines.region, span_ev.0, isotope.isotope
1294 )));
1295 }
1296 }
1297 lines
1298 .isotopes
1299 .iter()
1300 .zip(&lines.dopplers)
1301 .map(|((isotope, _), doppler)| {
1302 let read = match lines.lattice {
1303 Some(spectrum) => {
1304 let reach = |e: f64| spectrum.reach_ev(e, lines.reach_k, isotope.awr);
1305 (span_ev.0 - reach(span_ev.0), span_ev.1 + reach(span_ev.1))
1306 }
1307 None => (
1308 (span_ev.0.sqrt() - SUPPORT_X * doppler.u())
1309 .max(0.0)
1310 .powi(2),
1311 (span_ev.1.sqrt() + SUPPORT_X * doppler.u()).powi(2),
1312 ),
1313 };
1314 if !isotope.ranges.iter().any(|range| {
1315 range.is_evaluable()
1316 && range.energy_low <= read.0
1317 && read.1 <= range.energy_high
1318 }) {
1319 return Err(PipelineError::InvalidParameter(format!(
1320 "the Doppler-broadened cross section of {} in region {} reads \
1321 {:.6e}–{:.6e} eV, which no single one of its evaluated resolved \
1322 ranges holds",
1323 isotope.isotope, lines.region, read.0, read.1
1324 )));
1325 }
1326 Ok(resonance_center_energies(&[isotope])
1327 .into_iter()
1328 .filter(|e| (span_ev.0..=span_ev.1).contains(e))
1329 .collect())
1330 })
1331 .collect()
1332 })
1333 .collect()
1334}
1335
1336pub(crate) enum BeamStarts<'a> {
1337 Given(&'a Starts),
1338 OpenFitsAt(Bypass, Option<Vec<usize>>, Vec<f64>),
1339}
1340
1341pub(crate) fn spread_shares(
1344 measurement: &Measurement,
1345) -> Result<Option<Arc<Vec<Vec<f64>>>>, PipelineError> {
1346 let Some(spread) = &measurement.spread else {
1347 return Ok(None);
1348 };
1349 spread_checked(spread, measurement.regions.len())?;
1350 Ok(Some(Arc::new(region_shares(
1351 &spread.kernel,
1352 spread.grid,
1353 &spread.cells,
1354 ))))
1355}
1356
1357pub(crate) fn fit_counts_answer(
1360 measurement: &Measurement,
1361 calibration: &Calibration,
1362 beam_starts: BeamStarts<'_>,
1363 mixing: Option<Arc<Vec<Vec<f64>>>>,
1364) -> Result<(CountsFit, Answer), PipelineError> {
1365 let starts = match beam_starts {
1366 BeamStarts::Given(starts) => Some(starts),
1367 BeamStarts::OpenFitsAt(..) => None,
1368 };
1369 let Measurement {
1370 time_edges_us,
1371 charge_ratio,
1372 regions,
1373 ..
1374 } = measurement;
1375 let Checked {
1376 shape,
1377 sample_parameters,
1378 measured,
1379 lines,
1380 instrument,
1381 offsets,
1382 mut base,
1383 mut resonances_in_span,
1384 } = checked(measurement, calibration)?;
1385 let pulse = &calibration.pulse;
1386 let start: Vec<f64> = instrument.iter().map(|parameter| parameter.value).collect();
1387 let beam_origin_us = starts.map_or(start[0], |starts| starts.origin_us);
1388 let frames = measurement
1389 .room_background
1390 .map_or([0.0; 2], |room| room.frames);
1391 let spreading = measurement
1392 .spread
1393 .as_ref()
1394 .is_some_and(|spread| spread.share != Value::Known(0.0));
1395 let bins = time_edges_us.len() - 1;
1396 if let Some(starts) = starts {
1397 assert_eq!(starts.beams.len(), regions.len(), "one beam per region");
1398 if let Some(leverage) = &starts.shared_leverage {
1399 assert_eq!(leverage.len(), regions.len(), "one leverage per region");
1400 }
1401 }
1402 let excluded_at = |numbers: &[f64], offsets: &[f64]| -> Vec<Vec<usize>> {
1403 let own: Vec<Vec<usize>> = regions
1404 .iter()
1405 .zip(offsets)
1406 .map(|(region, offset)| {
1407 let mut numbers = numbers.to_vec();
1408 numbers[0] += offset;
1409 excluded_bins(
1410 time_edges_us,
1411 region.material.as_ref(),
1412 &numbers,
1413 calibration.pulse.storage_span(),
1414 )
1415 })
1416 .collect();
1417 match &mixing {
1418 Some(shares) if spreading => shares
1419 .iter()
1420 .enumerate()
1421 .map(|(r, from)| {
1422 let mut out: Vec<usize> = own
1423 .iter()
1424 .enumerate()
1425 .filter(|&(q, _)| q == r || from[q] > 0.0)
1426 .flat_map(|(_, out)| out.iter().copied())
1427 .collect();
1428 out.sort_unstable();
1429 out.dedup();
1430 out
1431 })
1432 .collect(),
1433 _ => own,
1434 }
1435 };
1436 let dispersion_of = |excluded: &[Vec<usize>], weights: &[f64]| -> Vec<f64> {
1437 let mut dispersion: Vec<f64> = (0..2 * regions.len() * bins)
1438 .map(|k| weights[k / bins])
1439 .collect();
1440 for (r, out) in excluded.iter().enumerate() {
1441 for &k in out {
1442 dispersion[(2 * r + 1) * bins + k] = f64::INFINITY;
1443 }
1444 }
1445 dispersion
1446 };
1447 let mut excluded = starts.map_or_else(|| excluded_at(&start, &offsets), |s| s.excluded.clone());
1448 let mut live = Vec::with_capacity(2 * regions.len() * bins);
1449 for (r, region) in regions.iter().enumerate() {
1450 let checked = || -> Result<[Vec<f64>; 2], PipelineError> {
1451 validate_counts("open-beam", ®ion.open_counts, bins)?;
1452 validate_counts("sample", ®ion.sample_counts, bins)?;
1453 Ok([
1454 validate_live("open-beam", region.open_live.as_deref(), bins)?,
1455 validate_live("sample", region.sample_live.as_deref(), bins)?,
1456 ])
1457 };
1458 live.extend(checked().map_err(|e| in_region(r, e))?.concat());
1459 }
1460
1461 let opens: Vec<BeamStart> = match beam_starts {
1462 BeamStarts::Given(starts) => starts.beams.clone(),
1463 BeamStarts::OpenFitsAt(numbers, ref intervals, ref offsets) => regions
1464 .iter()
1465 .enumerate()
1466 .map(|(r, region)| {
1467 let per_pixel = numbers.open_counts_per_pixel(
1468 &late(&starts_and_widths_us(time_edges_us), offsets[r]),
1469 open_frames(measurement.room_background.as_ref()),
1470 );
1471 let known: Vec<f64> = per_pixel.iter().map(|g| region.pixels * g).collect();
1472 let open = open_beam_with_intervals(
1473 time_edges_us,
1474 ®ion.open_counts,
1475 calibration,
1476 Some(&live[2 * r * bins..(2 * r + 1) * bins]),
1477 Some(&known),
1478 intervals.as_ref().map(|chosen| chosen[r]),
1479 )
1480 .map_err(|e| in_region(r, e))?;
1481 Ok(BeamStart::from(open))
1482 })
1483 .collect::<Result<_, PipelineError>>()?,
1484 };
1485 let beams: Vec<BeamSpline> = opens.iter().map(|open| open.beam.clone()).collect();
1486 let sizes: Vec<usize> = beams.iter().map(|beam| beam.coefficients().len()).collect();
1487 let pixels: Vec<f64> = regions.iter().map(|region| region.pixels).collect();
1488 let switches: Vec<[bool; 2]> = regions
1489 .iter()
1490 .map(|region| {
1491 [region.capture_gammas, region.in_scatter].map(|value| value != Value::Known(0.0))
1492 })
1493 .collect();
1494 let mut layout = Layout::new(roles(&sizes, &shape), &pixels, &switches);
1495 layout.storage_span = calibration.pulse.storage_span();
1496 let layout = Arc::new(layout);
1497 let mut parameters = ParameterSet::new(
1498 beams
1499 .iter()
1500 .enumerate()
1501 .flat_map(|(r, beam)| {
1502 beam.coefficients().iter().enumerate().map(move |(i, &c)| {
1503 FitParameter::unbounded(format!("beam {i} of region {r}"), c)
1504 })
1505 })
1506 .chain(sample_parameters)
1507 .chain(instrument)
1508 .collect(),
1509 );
1510 let lattices: Vec<Option<PhononSpectrum>> = regions
1511 .iter()
1512 .map(|region| region.material.as_ref().and_then(|m| m.lattice.clone()))
1513 .collect();
1514 let resonances: Vec<Arc<[ResonanceData]>> = regions
1515 .iter()
1516 .map(|region| {
1517 region
1518 .material
1519 .iter()
1520 .flat_map(|material| material.isotopes.iter().map(|(data, _)| data.clone()))
1521 .collect()
1522 })
1523 .collect();
1524 let observed: Vec<f64> = regions
1525 .iter()
1526 .flat_map(|region| region.open_counts.iter().chain(®ion.sample_counts))
1527 .copied()
1528 .collect();
1529 let calibrated_parameters: Vec<usize> = pulse
1530 .prior
1531 .iter()
1532 .flat_map(|prior| prior.numbers.iter().map(|n| layout.pulse + n))
1533 .collect();
1534 let pulse_prior = pulse
1535 .prior
1536 .as_ref()
1537 .map(|prior| Prior::correlated(&calibrated_parameters, &prior.mean, &prior.covariance))
1538 .transpose()?;
1539 let priors: Vec<Prior> = measured
1540 .iter()
1541 .map(|&(offset, mean, sd)| Prior::measured(layout.quantities + offset, mean, sd))
1542 .chain(pulse_prior.clone())
1543 .collect();
1544 let half_maximum = 2.0 * std::f64::consts::LN_2.sqrt();
1545 let narrowest_us = |grid: &FlightTimeGrid,
1546 resonances_in_span: &[Vec<Vec<f64>>],
1547 params: &[f64]|
1548 -> Result<f64, PipelineError> {
1549 let clock = TOF_FACTOR * grid.flight_path_m();
1550 let mut narrowest = f64::INFINITY;
1551 for (lines, energies) in lines.iter().zip(resonances_in_span) {
1552 let temperature_k = params[layout.regions[lines.region]
1553 .temperature
1554 .expect("a material has a temperature")];
1555 for ((isotope, _), energies) in lines.isotopes.iter().zip(energies) {
1556 let effective_k = lines
1557 .lattice
1558 .map_or(temperature_k, |s| s.effective_temperature_k(temperature_k));
1559 let doppler = DopplerParams::new(effective_k, isotope.awr)
1560 .map_err(|e| PipelineError::InvalidParameter(e.to_string()))?;
1561 for &energy in energies {
1562 let width_ev = half_maximum * doppler.doppler_width(energy);
1563 narrowest = narrowest.min(clock / energy.sqrt() * width_ev / (2.0 * energy));
1564 }
1565 }
1566 }
1567 Ok(narrowest)
1568 };
1569 let first_grid = |grid: &FlightTimeGrid,
1570 resonances_in_span: &[Vec<Vec<f64>>],
1571 params: &[f64]|
1572 -> Result<(Arc<FlightTimeGrid>, usize), PipelineError> {
1573 let rule_us = 0.5 * narrowest_us(grid, resonances_in_span, params)?;
1574 let mut first = grid.clone();
1575 let mut halvings = 0;
1576 while first.step_us() > rule_us {
1577 first = first.halved()?;
1578 halvings += 1;
1579 }
1580 Ok((Arc::new(first), halvings))
1581 };
1582 let covering = |grid: &FlightTimeGrid, params: &[f64]| -> Result<bool, PipelineError> {
1583 let pulse = laws(¶ms[layout.pulse..], layout.storage_span);
1584 let (first, last) = span(&layout.starts(params));
1585 let covers = |t0_us: f64| grid.covers(t0_us, params[layout.flight_path], &pulse);
1586 Ok(covers(first)? && covers(last)?)
1587 };
1588 let starting: Vec<f64> = parameters.params.iter().map(|p| p.value).collect();
1589 let mut first = first_grid(&base, &resonances_in_span, &starting)?;
1590 let mut weights: Vec<f64> = opens
1591 .iter()
1592 .flat_map(|open| [open.overdispersion.unwrap_or(1.0); 2])
1593 .collect();
1594 let mut sample_measured = vec![false; regions.len()];
1595 let mut noise_bins: Option<Vec<Vec<usize>>> = None;
1596 let mut passes = 0;
1597 let (fit, rule_halvings, overdispersions, settled) = loop {
1598 passes += 1;
1599 let dispersion = dispersion_of(&excluded, &weights);
1600 let fit = fit_on_halved_grids(
1601 &first.0,
1602 &mut parameters,
1603 &observed,
1604 &dispersion,
1605 &priors,
1606 |grid| {
1607 Ok(Recorded {
1608 model: RegionsModel::new(
1609 grid,
1610 &beams,
1611 &resonances,
1612 (*charge_ratio, frames),
1613 beam_origin_us,
1614 &layout,
1615 )
1616 .lattices(&lattices)
1617 .spread(mixing.clone()),
1618 live: &live,
1619 })
1620 },
1621 |spread, grid, params| {
1622 Ok(spread <= starts.map_or(BOUND, |starts| starts.bound)
1623 || !covering(grid, params)?)
1624 },
1625 )?;
1626 let fitted = |index: usize| fit.result.params[index];
1627 let noise_bins = noise_bins.get_or_insert_with(|| {
1628 (0..regions.len())
1629 .map(|r| counted(&fit, (2 * r + 1) * bins..(2 * r + 2) * bins))
1630 .collect()
1631 });
1632 let samples: Vec<Option<f64>> = noise_bins
1633 .iter()
1634 .enumerate()
1635 .map(|(r, counted)| {
1636 let weighted: Vec<usize> = counted
1637 .iter()
1638 .copied()
1639 .filter(|&k| dispersion[k].is_finite())
1640 .collect();
1641 overdispersion(
1642 &observed,
1643 &fit,
1644 &weighted,
1645 starts
1646 .and_then(|starts| starts.shared_leverage.as_ref())
1647 .map_or(0.0, |extra| extra[r]),
1648 )
1649 })
1650 .collect();
1651 let next: Vec<f64> = samples
1652 .iter()
1653 .enumerate()
1654 .map(|(r, sample)| sample.unwrap_or(weights[2 * r]))
1655 .collect();
1656 let each_settled: Vec<bool> = next
1657 .iter()
1658 .enumerate()
1659 .map(|(r, next)| (next / weights[2 * r + 1] - 1.0).abs() <= SETTLED_OVERDISPERSION)
1660 .collect();
1661 let at_fit = starts.map_or_else(
1662 || {
1663 excluded_at(
1664 &[fitted(layout.t0), fitted(layout.flight_path)]
1665 .into_iter()
1666 .chain(fit.result.params[layout.pulse..].iter().copied())
1667 .collect::<Vec<f64>>(),
1668 &layout
1669 .regions
1670 .iter()
1671 .map(|place| fitted(place.offset))
1672 .collect::<Vec<f64>>(),
1673 )
1674 },
1675 |s| s.excluded.clone(),
1676 );
1677 let settled = each_settled.iter().all(|&s| s) && at_fit == excluded;
1678 let resolved = 2.0 * fit.step_us
1679 <= 0.5 * narrowest_us(&base, &resonances_in_span, &fit.result.params)?;
1680 let covered = fit.converged && covering(&fit.coarse, &fit.result.params)?;
1681 if !fit.converged || (settled && resolved && covered) || passes == MOST_PASSES {
1682 let weighted: Vec<[Option<f64>; 2]> = (0..regions.len())
1683 .map(|r| {
1684 [
1685 opens[r].overdispersion,
1686 (sample_measured[r] || (each_settled[r] && samples[r].is_some()))
1687 .then_some(weights[2 * r + 1]),
1688 ]
1689 })
1690 .collect();
1691 break (fit, first.1, weighted, settled && resolved && covered);
1692 }
1693 if covered {
1694 let resumed = (Arc::clone(&fit.coarse), first.1 + fit.halvings - 1);
1695 first = [
1696 first_grid(&base, &resonances_in_span, &fit.result.params)?,
1697 resumed,
1698 ]
1699 .into_iter()
1700 .min_by(|a, b| a.0.step_us().total_cmp(&b.0.step_us()))
1701 .expect("two grids");
1702 } else {
1703 base = FlightTimeGrid::spanning(
1704 time_edges_us,
1705 span(&layout.starts(&fit.result.params)),
1706 fitted(layout.flight_path),
1707 &calibration.pulse.at(&fit.result.params[layout.pulse..])?,
1708 )?;
1709 resonances_in_span = in_span(&lines, &base)?;
1710 first = first_grid(&base, &resonances_in_span, &fit.result.params)?;
1711 }
1712 for (r, sample) in samples.iter().enumerate() {
1713 weights[2 * r + 1] = next[r];
1714 sample_measured[r] = sample.is_some();
1715 }
1716 excluded = at_fit;
1717 };
1718 let dispersion = dispersion_of(&excluded, &weights);
1719 let converged = fit.converged && settled;
1720 if converged {
1721 calibrated_lines(
1722 &lines,
1723 pulse,
1724 time_edges_us,
1725 &layout.starts(&fit.result.params),
1726 fit.result.params[layout.flight_path],
1727 )?;
1728 }
1729
1730 if let Some((k, (&counts, &predicted))) =
1731 observed
1732 .iter()
1733 .zip(&fit.predicted)
1734 .enumerate()
1735 .find(|(k, (y, mu))| {
1736 !mu.is_finite()
1737 || **mu < 0.0
1738 || (dispersion[*k].is_finite() && **y > 0.0 && **mu < NEGLIGIBLE_PREDICTION)
1739 })
1740 {
1741 return Err(PipelineError::UnmodelledCounts {
1742 region: k / (2 * bins),
1743 run: ["open-beam", "sample"][(k / bins) % 2],
1744 bin: k % bins,
1745 counts,
1746 predicted,
1747 });
1748 }
1749
1750 let carried = regions
1751 .iter()
1752 .flat_map(|region| ®ion.material)
1753 .any(|m| !m.covariances.is_empty());
1754 let slopes = if converged && carried && fit.result.covariance.is_some() {
1755 let model = Recorded {
1756 model: RegionsModel::new(
1757 &fit.fine,
1758 &beams,
1759 &resonances,
1760 (*charge_ratio, frames),
1761 beam_origin_us,
1762 &layout,
1763 )
1764 .lattices(&lattices)
1765 .spread(mixing.clone()),
1766 live: &live,
1767 };
1768 nuclear_slopes(regions, &model, &dispersion, ¶meters)?
1769 } else {
1770 None
1771 };
1772 let carried_with = |c: Option<&FlatMatrix>| Some(sandwich(c?, slopes.as_ref()?));
1773 let nuclear = carried_with(fit.result.covariance.as_ref());
1774 let unbounded_nuclear = carried_with(fit.result.unbounded.as_ref().map(|u| &u.covariance));
1775 let (t_low, t_high) = TEMPERATURE_BOUNDS_K;
1776 let free = parameters.free_indices();
1777 let on_edge = layout
1778 .regions
1779 .iter()
1780 .filter_map(|region| region.temperature)
1781 .any(|t| free.contains(&t) && [t_low, t_high].contains(&fit.result.params[t]));
1782 let sample_quantities: Vec<usize> = (0..free.len())
1783 .filter(|&p| free[p] >= layout.quantities)
1784 .collect();
1785 let block = |full: &FlatMatrix| {
1786 let size = sample_quantities.len();
1787 let mut block = FlatMatrix::zeros(size, size);
1788 for (a, &p) in sample_quantities.iter().enumerate() {
1789 for (b, &q) in sample_quantities.iter().enumerate() {
1790 *block.get_mut(a, b) = if on_edge { f64::NAN } else { full.get(p, q) };
1791 }
1792 }
1793 block
1794 };
1795 let covariance = fit
1796 .result
1797 .covariance
1798 .as_ref()
1799 .filter(|_| converged)
1800 .map(|conditional| {
1801 let mut full = conditional.clone();
1802 for (entry, added) in full
1803 .data
1804 .iter_mut()
1805 .zip(nuclear.iter().flat_map(|n| &n.data))
1806 {
1807 *entry += added;
1808 }
1809 block(&full)
1810 });
1811 let unbounded = fit
1812 .result
1813 .unbounded
1814 .as_ref()
1815 .filter(|_| converged)
1816 .map(|full| Unbounded {
1817 mean: sample_quantities
1818 .iter()
1819 .map(|&p| if on_edge { f64::NAN } else { full.mean[p] })
1820 .collect(),
1821 covariance: block(&full.covariance),
1822 });
1823 let params = &fit.result.params;
1824 let on_bound: Vec<bool> = sample_quantities
1825 .iter()
1826 .map(|&p| fit.result.on_bound[p])
1827 .collect();
1828 let position = |parameter: usize| {
1829 sample_quantities
1830 .iter()
1831 .position(|&p| free[p] == parameter)
1832 .expect("a measured or calibrated quantity is fitted")
1833 };
1834 let pulse_consistency = match (&pulse_prior, &unbounded) {
1835 (Some(prior), Some(unbounded)) => {
1836 let at: Vec<usize> = calibrated_parameters
1837 .iter()
1838 .map(|¶meter| position(parameter))
1839 .collect();
1840 let mut posterior = FlatMatrix::zeros(at.len(), at.len());
1841 for (a, &i) in at.iter().enumerate() {
1842 for (b, &j) in at.iter().enumerate() {
1843 *posterior.get_mut(a, b) = unbounded.covariance.get(i, j);
1844 }
1845 }
1846 let estimate: Vec<f64> = at.iter().map(|&i| unbounded.mean[i]).collect();
1847 consistency(prior, &estimate, &posterior)?
1848 }
1849 _ => None,
1850 };
1851 let measured_pulls = unbounded.as_ref().map(|unbounded| {
1852 measured
1853 .iter()
1854 .map(|&(offset, mean, sd)| {
1855 let a = position(layout.quantities + offset);
1856 let p = sample_quantities[a];
1857 let nuclear = unbounded_nuclear.as_ref().map_or(0.0, |n| n.get(p, p));
1858 (unbounded.mean[a] - mean)
1859 / (sd * sd - unbounded.covariance.get(a, a) + nuclear).sqrt()
1860 })
1861 .collect()
1862 });
1863 let region_fits = layout
1864 .regions
1865 .iter()
1866 .zip(&opens)
1867 .zip(overdispersions)
1868 .enumerate()
1869 .map(|(r, ((place, open), overdispersion))| RegionFit {
1870 densities: params[place.densities.clone()].to_vec(),
1871 temperature_k: place.temperature.map(|t| params[t]),
1872 thickness_variance: place.thickness.map(|v| params[v]),
1873 background: std::array::from_fn(|i| params[place.background + i]),
1874 gamma_transmission: params[place.gamma],
1875 capture_gammas: params[place.capture],
1876 in_scatter: params[place.scatter],
1877 timing_offset_us: params[place.offset],
1878 beam: open.beam.with_coefficients(¶ms[place.beam.clone()]),
1879 beam_at_limit: open.at_limit,
1880 predicted: [
1881 fit.predicted[2 * r * bins..(2 * r + 1) * bins].to_vec(),
1882 fit.predicted[(2 * r + 1) * bins..(2 * r + 2) * bins].to_vec(),
1883 ],
1884 overdispersion,
1885 excluded_bins: excluded[r].clone(),
1886 })
1887 .collect();
1888 let mut fitted = vec![false; params.len()];
1889 let mut at_bound = vec![false; params.len()];
1890 for (&index, &bound) in free.iter().zip(&fit.result.on_bound) {
1891 fitted[index] = true;
1892 at_bound[index] = bound;
1893 }
1894 let answer = Answer {
1895 grid: Arc::clone(&fit.fine),
1896 dispersion,
1897 params: params.clone(),
1898 fitted,
1899 at_bound,
1900 beams,
1901 resonances,
1902 lattices,
1903 charge_ratio: *charge_ratio,
1904 frames,
1905 beam_origin_us,
1906 layout: Arc::clone(&layout),
1907 live,
1908 observed,
1909 };
1910 let fit = CountsFit {
1911 regions: region_fits,
1912 normalization: params[layout.normalization],
1913 t0_us: params[layout.t0],
1914 flight_path_m: params[layout.flight_path],
1915 alpha: [params[layout.pulse], params[layout.pulse + 1]],
1916 beta: [params[layout.pulse + 2], params[layout.pulse + 3]],
1917 r: params[layout.pulse + 4],
1918 fwhm_squared_us2: params[layout.pulse + 5],
1919 r_high: calibration.pulse.r_high.map(|_| params[layout.pulse + 6]),
1920 detector_tail: calibration
1921 .pulse
1922 .detector_tail
1923 .map(|_| [params[layout.pulse + 7], params[layout.pulse + 8]]),
1924 time_background: [0, 1, 2].map(|i| params[layout.time_background + i]),
1925 room_background: params[layout.room_background],
1926 spread_share: params[layout.spread],
1927 covariance,
1928 unbounded,
1929 on_bound,
1930 deviance: fit.result.deviance,
1931 converged,
1932 measured_pulls,
1933 pulse_consistency,
1934 step_us: fit.step_us,
1935 points: fit.points,
1936 halvings: rule_halvings + fit.halvings,
1937 };
1938 Ok((fit, answer))
1939}
1940
1941pub(crate) struct Answer {
1942 grid: Arc<FlightTimeGrid>,
1943 beams: Vec<BeamSpline>,
1944 resonances: Vec<Arc<[ResonanceData]>>,
1945 lattices: Vec<Option<PhononSpectrum>>,
1946 charge_ratio: f64,
1947 frames: [f64; 2],
1948 beam_origin_us: f64,
1949 layout: Arc<Layout>,
1950 live: Vec<f64>,
1951 observed: Vec<f64>,
1952 dispersion: Vec<f64>,
1953 params: Vec<f64>,
1954 fitted: Vec<bool>,
1955 at_bound: Vec<bool>,
1956}
1957
1958pub(crate) struct RegionTerms {
1959 pub(crate) quantities: Vec<Role>,
1960 pub(crate) gradient: Vec<f64>,
1961 pub(crate) information: FlatMatrix,
1962 pub(crate) diagonal: Vec<f64>,
1963 pub(crate) unbounded: Option<(FlatMatrix, Vec<f64>)>,
1964 pub(crate) sensitivity: FlatMatrix,
1965 pub(crate) covariance: FlatMatrix,
1966 pub(crate) counted_information: FlatMatrix,
1967}
1968
1969impl Answer {
1970 pub(crate) fn region_terms(
1971 &self,
1972 region: usize,
1973 shared: &[Role],
1974 ) -> Result<RegionTerms, PipelineError> {
1975 let roles = &self.layout.roles;
1976 let own: Vec<usize> = (0..roles.len())
1977 .filter(|&i| {
1978 self.fitted[i]
1979 && matches!(
1980 roles[i],
1981 Role::Beam { region: r, .. }
1982 | Role::Density { region: r, .. }
1983 | Role::Temperature { region: r }
1984 | Role::ThicknessVariance { region: r }
1985 | Role::Background { region: r, .. }
1986 | Role::GammaTransmission { region: r }
1987 | Role::CaptureGammas { region: r }
1988 | Role::InScatter { region: r }
1989 | Role::TimingOffset { region: r }
1990 if r == region
1991 )
1992 })
1993 .collect();
1994 let shared: Vec<usize> = shared
1995 .iter()
1996 .map(|&role| {
1997 roles
1998 .iter()
1999 .position(|&r| r == role)
2000 .expect("a shared role is in the layout")
2001 })
2002 .collect();
2003 let columns: Vec<usize> = own.iter().chain(&shared).copied().collect();
2004 let model = Recorded {
2005 model: RegionsModel::new(
2006 &self.grid,
2007 &self.beams,
2008 &self.resonances,
2009 (self.charge_ratio, self.frames),
2010 self.beam_origin_us,
2011 &self.layout,
2012 )
2013 .lattices(&self.lattices),
2014 live: &self.live,
2015 };
2016 let predicted = model.evaluate(&self.params)?;
2017 let jacobian = model
2018 .analytical_jacobian(&self.params, &columns, &predicted)
2019 .expect("the counts model has an analytical Jacobian");
2020 let n = columns.len();
2021 let mut full = Mat::<f64>::zeros(n, n);
2022 let mut gradient = vec![0.0; n];
2023 let bins = self.observed.len() / (2 * self.layout.regions.len());
2024 let rows = 2 * region * bins..2 * (region + 1) * bins;
2025 let mut weights = vec![0.0; rows.len()];
2026 for ((((k, &mu), &dispersion), &observed), weight) in rows
2027 .clone()
2028 .zip(&predicted[rows.clone()])
2029 .zip(&self.dispersion[rows.clone()])
2030 .zip(&self.observed[rows.clone()])
2031 .zip(weights.iter_mut())
2032 {
2033 if mu > 0.0 {
2034 *weight = 1.0 / (mu * dispersion);
2035 for a in 0..n {
2036 let slope = jacobian.get(k, a);
2037 gradient[a] += *weight * slope * (mu - observed);
2038 for b in 0..n {
2039 full[(a, b)] += *weight * slope * jacobian.get(k, b);
2040 }
2041 }
2042 } else {
2043 for (a, g) in gradient.iter_mut().enumerate() {
2044 *g += jacobian.get(k, a) / dispersion;
2045 }
2046 }
2047 }
2048 let (o, s) = (own.len(), shared.len());
2049 let counted = rows.len();
2050 let profile = |set: Vec<usize>| -> Result<Profile, PipelineError> {
2051 let m = set.len();
2052 let mut block = FlatMatrix::zeros(m, m);
2053 for (a, &i) in set.iter().enumerate() {
2054 for (b, &j) in set.iter().enumerate() {
2055 *block.get_mut(a, b) = full[(i, j)];
2056 }
2057 }
2058 let diagonal: Vec<f64> = set.iter().map(|&i| full[(i, i)]).collect();
2059 let inverse = information_inverse(&block, &diagonal, counted)?;
2060 let mut moved = FlatMatrix::zeros(m, s);
2061 for a in 0..m {
2062 for j in 0..s {
2063 *moved.get_mut(a, j) = -(0..m)
2064 .map(|b| inverse.spanned.get(a, b) * full[(set[b], o + j)])
2065 .sum::<f64>();
2066 }
2067 }
2068 let profiled = |i: usize, j: usize| {
2069 full[(o + i, o + j)]
2070 + (0..m)
2071 .map(|a| full[(o + i, set[a])] * moved.get(a, j))
2072 .sum::<f64>()
2073 };
2074 let mut information = FlatMatrix::zeros(s, s);
2075 for i in 0..s {
2076 for j in 0..s {
2077 *information.get_mut(i, j) = 0.5 * (profiled(i, j) + profiled(j, i));
2078 }
2079 }
2080 let slope = (0..s)
2081 .map(|j| {
2082 gradient[o + j]
2083 + (0..m)
2084 .map(|a| moved.get(a, j) * gradient[set[a]])
2085 .sum::<f64>()
2086 })
2087 .collect();
2088 Ok(Profile {
2089 set,
2090 inverse: inverse.determined,
2091 resolved: inverse.resolved,
2092 moved,
2093 information,
2094 gradient: slope,
2095 })
2096 };
2097 let bounded = profile((0..o).filter(|&a| !self.at_bound[own[a]]).collect())?;
2098 let unbounded = profile((0..o).collect())?;
2099 let Profile {
2100 set,
2101 inverse,
2102 resolved,
2103 moved,
2104 information,
2105 gradient,
2106 } = bounded;
2107 let quantities: Vec<usize> = (0..set.len())
2108 .filter(|&a| !matches!(roles[own[set[a]]], Role::Beam { .. }))
2109 .collect();
2110 let mut sensitivity = FlatMatrix::zeros(quantities.len(), s);
2111 let mut covariance = FlatMatrix::zeros(quantities.len(), quantities.len());
2112 for (row, &a) in quantities.iter().enumerate() {
2113 for j in 0..s {
2114 *sensitivity.get_mut(row, j) = if resolved[a] {
2115 moved.get(a, j)
2116 } else {
2117 f64::NAN
2118 };
2119 }
2120 for (col, &b) in quantities.iter().enumerate() {
2121 *covariance.get_mut(row, col) = if resolved[a] && resolved[b] {
2122 inverse.get(a, b)
2123 } else {
2124 f64::NAN
2125 };
2126 }
2127 }
2128 let mut counted_information = FlatMatrix::zeros(s, s);
2129 let sample = rows.start + counted / 2..rows.end;
2130 for (k, &weight) in sample.clone().zip(&weights[counted / 2..]) {
2131 if predicted[k] < COUNTS_TO_MEASURE_NOISE {
2132 continue;
2133 }
2134 let slope: Vec<f64> = (0..s)
2135 .map(|j| {
2136 jacobian.get(k, o + j)
2137 + (0..set.len())
2138 .map(|a| jacobian.get(k, set[a]) * moved.get(a, j))
2139 .sum::<f64>()
2140 })
2141 .collect();
2142 for i in 0..s {
2143 for j in 0..s {
2144 *counted_information.get_mut(i, j) += weight * slope[i] * slope[j];
2145 }
2146 }
2147 }
2148 Ok(RegionTerms {
2149 quantities: quantities.iter().map(|&a| roles[own[set[a]]]).collect(),
2150 gradient,
2151 information,
2152 diagonal: (0..s).map(|j| full[(o + j, o + j)]).collect(),
2153 unbounded: Some((unbounded.information, unbounded.gradient)),
2154 sensitivity,
2155 covariance,
2156 counted_information,
2157 })
2158 }
2159}
2160
2161struct Profile {
2162 set: Vec<usize>,
2163 inverse: FlatMatrix,
2164 resolved: Vec<bool>,
2165 moved: FlatMatrix,
2166 information: FlatMatrix,
2167 gradient: Vec<f64>,
2168}
2169
2170pub(crate) const BACKGROUND_TERMS: [&str; 5] = ["b0", "b1", "b2", "d", "f"];
2171
2172pub(crate) const SHARED: [Role; 16] = [
2173 Role::Normalization,
2174 Role::TimeBackground(0),
2175 Role::TimeBackground(1),
2176 Role::TimeBackground(2),
2177 Role::RoomBackground,
2178 Role::T0,
2179 Role::FlightPath,
2180 Role::Pulse(0),
2181 Role::Pulse(1),
2182 Role::Pulse(2),
2183 Role::Pulse(3),
2184 Role::Pulse(4),
2185 Role::Pulse(5),
2186 Role::Pulse(6),
2187 Role::Pulse(7),
2188 Role::Pulse(8),
2189];
2190
2191pub(crate) fn shared_parameters(
2192 measurement: &Measurement,
2193 calibration: &Calibration,
2194) -> Result<Vec<(Value, FitParameter)>, PipelineError> {
2195 validate_pulse_prior(calibration)?;
2196 let mut parameters = Vec::with_capacity(SHARED.len());
2197 for (value, name, range, allowed) in SHARED
2198 .iter()
2199 .filter_map(|&role| quantity(measurement, role))
2200 {
2201 parameters.push(value.parameter(name, range, &allowed)?);
2202 }
2203 parameters.extend(calibration.instrument()?);
2204 let values = quantities(measurement, calibration)
2205 .filter(|(role, _)| SHARED.contains(role))
2206 .map(|(_, value)| *value);
2207 Ok(values.zip(parameters).collect())
2208}
2209
2210fn validate_pulse_prior(calibration: &Calibration) -> Result<(), PipelineError> {
2211 let pulse = &calibration.pulse;
2212 let numbers = pulse.numbers();
2213 match pulse
2214 .prior
2215 .iter()
2216 .flat_map(|prior| &prior.numbers)
2217 .find(|&&n| !matches!(numbers[n], Value::Fitted(_) | Value::Within { .. }))
2218 {
2219 Some(&number) => Err(PipelineError::InvalidParameter(format!(
2220 "the pulse's calibration covers {}, which must be fitted, without a measurement; \
2221 got {:?}",
2222 PULSE_NUMBERS[number], numbers[number]
2223 ))),
2224 None => Ok(()),
2225 }
2226}
2227
2228pub(crate) fn quantities<'a>(
2229 measurement: &'a Measurement,
2230 calibration: &'a Calibration,
2231) -> impl Iterator<Item = (Role, &'a Value)> {
2232 let shape = shape(measurement);
2233 let numbers = calibration.pulse.numbers();
2234 roles(&vec![0; shape.len()], &shape)
2235 .into_iter()
2236 .filter_map(move |role| {
2237 let value = match role {
2238 Role::T0 => &calibration.t0_us,
2239 Role::FlightPath => &calibration.flight_path_m,
2240 Role::Pulse(number) => numbers[number],
2241 _ => quantity(measurement, role)?.0,
2242 };
2243 Some((role, value))
2244 })
2245}
2246
2247fn shape(measurement: &Measurement) -> Vec<Option<usize>> {
2248 measurement
2249 .regions
2250 .iter()
2251 .map(|region| region.material.as_ref().map(|m| m.isotopes.len()))
2252 .collect()
2253}
2254
2255#[derive(Clone, Copy, Debug, PartialEq)]
2256pub(crate) enum Role {
2257 Beam { region: usize, coefficient: usize },
2258 Density { region: usize, isotope: usize },
2259 Temperature { region: usize },
2260 ThicknessVariance { region: usize },
2261 Normalization,
2262 Background { region: usize, term: usize },
2263 GammaTransmission { region: usize },
2264 CaptureGammas { region: usize },
2265 InScatter { region: usize },
2266 TimingOffset { region: usize },
2267 TimeBackground(usize),
2268 RoomBackground,
2269 SpreadShare,
2270 T0,
2271 FlightPath,
2272 Pulse(usize),
2273}
2274
2275fn roles(beams: &[usize], isotopes: &[Option<usize>]) -> Vec<Role> {
2276 let beam = beams.iter().enumerate().flat_map(|(region, &n)| {
2277 (0..n).map(move |coefficient| Role::Beam {
2278 region,
2279 coefficient,
2280 })
2281 });
2282 let materials = isotopes
2283 .iter()
2284 .enumerate()
2285 .filter_map(|(region, n)| Some((region, (*n)?)))
2286 .flat_map(|(region, n)| {
2287 (0..n)
2288 .map(move |isotope| Role::Density { region, isotope })
2289 .chain([
2290 Role::Temperature { region },
2291 Role::ThicknessVariance { region },
2292 ])
2293 });
2294 let backgrounds = (0..isotopes.len())
2295 .flat_map(|region| (0..5).map(move |term| Role::Background { region, term }));
2296 let gammas = (0..isotopes.len()).map(|region| Role::GammaTransmission { region });
2297 let captures = (0..isotopes.len()).map(|region| Role::CaptureGammas { region });
2298 let scatters = (0..isotopes.len()).map(|region| Role::InScatter { region });
2299 let offsets = (0..isotopes.len()).map(|region| Role::TimingOffset { region });
2300 beam.chain(materials)
2301 .chain([Role::Normalization])
2302 .chain(backgrounds)
2303 .chain(gammas)
2304 .chain(captures)
2305 .chain(scatters)
2306 .chain(offsets)
2307 .chain((0..3).map(Role::TimeBackground))
2308 .chain([
2309 Role::RoomBackground,
2310 Role::SpreadShare,
2311 Role::T0,
2312 Role::FlightPath,
2313 ])
2314 .chain((0..9).map(Role::Pulse))
2315 .collect()
2316}
2317
2318fn material(measurement: &Measurement, region: usize) -> &Material {
2319 measurement.regions[region]
2320 .material
2321 .as_ref()
2322 .expect("a density or temperature names a region with a material")
2323}
2324
2325fn quantity(
2326 measurement: &Measurement,
2327 role: Role,
2328) -> Option<(&Value, String, RangeInclusive<f64>, String)> {
2329 let (t_low, t_high) = TEMPERATURE_BOUNDS_K;
2330 Some(match role {
2331 Role::Density { region, isotope } => {
2332 let (data, density) = &material(measurement, region).isotopes[isotope];
2333 (
2334 density,
2335 format!("density of {} in region {region}", data.isotope),
2336 0.0..=f64::INFINITY,
2337 "0 or more".into(),
2338 )
2339 }
2340 Role::Temperature { region } => (
2341 &material(measurement, region).temperature_k,
2342 format!("temperature of region {region}"),
2343 t_low..=t_high,
2344 format!("within {t_low}–{t_high} K"),
2345 ),
2346 Role::ThicknessVariance { region } => (
2347 &material(measurement, region).thickness_variance,
2348 format!("thickness variance of region {region}"),
2349 0.0..=f64::INFINITY,
2350 "0 or more".into(),
2351 ),
2352 Role::Normalization => (
2353 &measurement.normalization,
2354 "normalization".into(),
2355 f64::MIN_POSITIVE..=f64::INFINITY,
2356 "positive".into(),
2357 ),
2358 Role::Background { region, term } => (
2359 &measurement.regions[region].background[term],
2360 format!("{} of region {region}", BACKGROUND_TERMS[term]),
2361 if term == 4 { 0.0 } else { f64::NEG_INFINITY }..=f64::INFINITY,
2362 if term == 4 {
2363 "0 or more"
2364 } else {
2365 "of any sign"
2366 }
2367 .into(),
2368 ),
2369 Role::CaptureGammas { region } => (
2370 &measurement.regions[region].capture_gammas,
2371 format!("capture gammas of region {region}"),
2372 0.0..=f64::INFINITY,
2373 "0 or more".into(),
2374 ),
2375 Role::InScatter { region } => (
2376 &measurement.regions[region].in_scatter,
2377 format!("in-scatter of region {region}"),
2378 0.0..=f64::INFINITY,
2379 "0 or more".into(),
2380 ),
2381 Role::TimingOffset { region } => (
2382 &measurement.regions[region].timing_offset_us,
2383 format!("timing offset of region {region}"),
2384 f64::NEG_INFINITY..=f64::INFINITY,
2385 "any real number".into(),
2386 ),
2387 Role::GammaTransmission { region } => (
2388 &measurement.regions[region].gamma_transmission,
2389 format!("gamma transmission of region {region}"),
2390 0.0..=1.0,
2391 "within 0–1".into(),
2392 ),
2393 Role::TimeBackground(term) => {
2394 let (name, low, allowed) = [
2395 ("C", 0.0, "0 or more"),
2396 ("A", 0.0, "0 or more"),
2397 ("τ", f64::MIN_POSITIVE, "positive"),
2398 ][term];
2399 (
2400 &measurement.time_background[term],
2401 format!("time background {name}"),
2402 low..=f64::INFINITY,
2403 allowed.into(),
2404 )
2405 }
2406 Role::RoomBackground => (
2407 measurement
2408 .room_background
2409 .as_ref()
2410 .map_or(&Value::Known(0.0), |room| &room.rate),
2411 "room background ρ".into(),
2412 0.0..=f64::INFINITY,
2413 "0 or more".into(),
2414 ),
2415 Role::SpreadShare => (
2416 measurement
2417 .spread
2418 .as_ref()
2419 .map_or(&Value::Known(0.0), |spread| &spread.share),
2420 "detector spread h".into(),
2421 0.0..=1.0,
2422 "within 0–1".into(),
2423 ),
2424 Role::Beam { .. } | Role::T0 | Role::FlightPath | Role::Pulse(_) => return None,
2425 })
2426}
2427
2428fn spread_checked(spread: &Spread, regions: usize) -> Result<(), PipelineError> {
2429 let invalid = |message: String| Err(PipelineError::InvalidParameter(message));
2430 let [rows, columns] = spread.grid;
2431 if !rows
2432 .checked_mul(columns)
2433 .is_some_and(|cells| (1..=MOST_SPREAD_CELLS).contains(&cells))
2434 {
2435 return invalid(format!(
2436 "the spread's {rows}×{columns} grid must hold 1 to {MOST_SPREAD_CELLS} cells"
2437 ));
2438 }
2439 if spread.kernel.is_empty()
2440 || !spread
2441 .kernel
2442 .iter()
2443 .all(|&(w, s)| w.is_finite() && w > 0.0 && s.is_finite() && s > 0.0)
2444 {
2445 return invalid(format!(
2446 "the spread kernel needs terms of finite, positive weight and length, got {:?}",
2447 spread.kernel
2448 ));
2449 }
2450 if spread.cells.len() != regions {
2451 return invalid(format!(
2452 "the spread gives the cells of {} regions for {regions} regions",
2453 spread.cells.len()
2454 ));
2455 }
2456 let mut taken = vec![false; rows * columns];
2457 for (r, cells) in spread.cells.iter().enumerate() {
2458 if cells.is_empty() {
2459 return invalid(format!("the spread gives region {r} no cells"));
2460 }
2461 for &[row, column] in cells {
2462 if row >= rows || column >= columns {
2463 return invalid(format!(
2464 "cell [{row}, {column}] of region {r} is off the {rows}×{columns} grid"
2465 ));
2466 }
2467 if std::mem::replace(&mut taken[row * columns + column], true) {
2468 return invalid(format!("cell [{row}, {column}] is in two regions"));
2469 }
2470 }
2471 }
2472 Ok(())
2473}
2474
2475fn in_region(region: usize, error: PipelineError) -> PipelineError {
2476 labelled(&format!("region {region}"), error)
2477}
2478
2479pub(crate) fn labelled(label: &str, error: PipelineError) -> PipelineError {
2480 match error {
2481 PipelineError::InvalidParameter(message) => {
2482 PipelineError::InvalidParameter(format!("{label}: {message}"))
2483 }
2484 PipelineError::ShapeMismatch(message) => {
2485 PipelineError::ShapeMismatch(format!("{label}: {message}"))
2486 }
2487 other => other,
2488 }
2489}
2490
2491fn finite(isotope: &ResonanceData) -> bool {
2492 isotope.awr.is_finite()
2493 && isotope.ranges.iter().all(|range| {
2494 let radii = range
2495 .ap_table
2496 .iter()
2497 .flat_map(|table| table.points.iter().flat_map(|&(e, r)| [e, r]));
2498 let external = range.r_external.iter().flat_map(|r| {
2499 [
2500 r.j, r.e_low, r.e_up, r.r_con, r.r_lin, r.s_con, r.s_lin, r.r_quad,
2501 ]
2502 });
2503 let groups = range.l_groups.iter().flat_map(|group| {
2504 [group.awr, group.apl, group.qx].into_iter().chain(
2505 group
2506 .resonances
2507 .iter()
2508 .flat_map(|r| [r.energy, r.j, r.gn, r.gg, r.gfa, r.gfb]),
2509 )
2510 });
2511 let rml = range.r_matrix_limited.iter().flat_map(rml_values);
2512 [
2513 range.energy_low,
2514 range.energy_high,
2515 range.target_spin,
2516 range.scattering_radius,
2517 ]
2518 .into_iter()
2519 .chain(radii)
2520 .chain(external)
2521 .chain(groups)
2522 .chain(rml)
2523 .all(f64::is_finite)
2524 })
2525}
2526
2527fn rml_values(rml: &RmlRange) -> impl Iterator<Item = f64> + '_ {
2529 let pairs = rml.pairs.iter().flat_map(|p| {
2530 [
2531 p.mass_a, p.mass_b, p.charge_a, p.charge_b, p.spin_a, p.spin_b, p.q_value,
2532 ]
2533 });
2534 let groups = rml.spin_groups.iter().flat_map(|g| {
2535 let channels = g
2536 .channels
2537 .iter()
2538 .flat_map(|c| [c.channel_spin, c.boundary, c.radius_phase, c.radius_true]);
2539 let levels = g
2540 .resonances
2541 .iter()
2542 .flat_map(|r| std::iter::once(r.energy).chain(r.widths.iter().copied()));
2543 let background = g.background.iter().flat_map(|b| {
2544 [
2545 b.e_down, b.e_up, b.r_con, b.r_lin, b.r_quad, b.s_con, b.s_lin,
2546 ]
2547 });
2548 std::iter::once(g.j)
2549 .chain(channels)
2550 .chain(levels)
2551 .chain(background)
2552 });
2553 pairs.chain(groups)
2554}
2555
2556struct Lines<'a> {
2557 region: usize,
2558 isotopes: &'a [(ResonanceData, Value)],
2559 dopplers: Vec<DopplerParams>,
2560 reach_k: f64,
2561 lattice: Option<&'a PhononSpectrum>,
2562 scatters: bool,
2563}
2564
2565struct Layout {
2566 roles: Vec<Role>,
2567 regions: Vec<Place>,
2568 quantities: usize,
2569 normalization: usize,
2570 time_background: usize,
2571 room_background: usize,
2572 spread: usize,
2573 t0: usize,
2574 flight_path: usize,
2575 pulse: usize,
2576 storage_span: Option<(f64, f64)>,
2577}
2578
2579struct Place {
2580 beam: Range<usize>,
2581 densities: Range<usize>,
2582 temperature: Option<usize>,
2583 thickness: Option<usize>,
2584 background: usize,
2585 gamma: usize,
2586 capture: usize,
2587 capture_on: bool,
2588 scatter: usize,
2589 scatter_on: bool,
2590 offset: usize,
2591 pixels: f64,
2592}
2593
2594impl Layout {
2595 fn starts(&self, params: &[f64]) -> Vec<f64> {
2597 self.regions
2598 .iter()
2599 .map(|place| params[self.t0] + params[place.offset])
2600 .collect()
2601 }
2602
2603 fn new(roles: Vec<Role>, pixels: &[f64], switches: &[[bool; 2]]) -> Self {
2604 let at = |wanted: Role| {
2605 roles
2606 .iter()
2607 .position(|&role| role == wanted)
2608 .expect("every region has a background and the instrument its numbers")
2609 };
2610 let span = |of: &dyn Fn(Role) -> bool| {
2611 let first = roles.iter().position(|&role| of(role));
2612 let last = roles.iter().rposition(|&role| of(role));
2613 first
2614 .zip(last)
2615 .map_or(0..0, |(first, last)| first..last + 1)
2616 };
2617 let places = (0..pixels.len())
2618 .map(|r| Place {
2619 beam: span(&|role| matches!(role, Role::Beam { region, .. } if region == r)),
2620 densities: span(
2621 &|role| matches!(role, Role::Density { region, .. } if region == r),
2622 ),
2623 temperature: roles
2624 .iter()
2625 .position(|&role| role == Role::Temperature { region: r }),
2626 thickness: roles
2627 .iter()
2628 .position(|&role| role == Role::ThicknessVariance { region: r }),
2629 background: at(Role::Background { region: r, term: 0 }),
2630 gamma: at(Role::GammaTransmission { region: r }),
2631 capture: at(Role::CaptureGammas { region: r }),
2632 capture_on: switches[r][0],
2633 scatter: at(Role::InScatter { region: r }),
2634 scatter_on: switches[r][1],
2635 offset: at(Role::TimingOffset { region: r }),
2636 pixels: pixels[r],
2637 })
2638 .collect();
2639 Self {
2640 quantities: roles
2641 .iter()
2642 .position(|role| !matches!(role, Role::Beam { .. }))
2643 .expect("the instrument's numbers follow the beams"),
2644 normalization: at(Role::Normalization),
2645 time_background: at(Role::TimeBackground(0)),
2646 room_background: at(Role::RoomBackground),
2647 spread: at(Role::SpreadShare),
2648 t0: at(Role::T0),
2649 flight_path: at(Role::FlightPath),
2650 pulse: at(Role::Pulse(0)),
2651 storage_span: None,
2652 regions: places,
2653 roles,
2654 }
2655 }
2656}
2657
2658struct RegionsModel {
2659 grid: Arc<FlightTimeGrid>,
2660 bins: Vec<(f64, f64)>,
2661 splines: Vec<BeamSpline>,
2662 beam_origin_us: f64,
2663 absorbers: Vec<Absorber>,
2664 energies: Vec<f64>,
2665 flight_energies: Vec<f64>,
2666 shapes: [Vec<f64>; 3],
2667 charge_ratio: f64,
2668 frames: [f64; 2],
2669 layout: Arc<Layout>,
2670 clocks: RefCell<Option<Clocks>>,
2671 mixing: Option<Arc<Vec<Vec<f64>>>>,
2672}
2673
2674struct Clocks {
2677 scales: Vec<Scale>,
2678 of_region: Vec<usize>,
2679 basis: Vec<Vec<[(usize, f64); 5]>>,
2680 basis_slope: Vec<Vec<[(usize, f64); 5]>>,
2681}
2682
2683impl Clocks {
2684 fn of(&self, region: usize) -> &Scale {
2685 &self.scales[self.of_region[region]]
2686 }
2687}
2688
2689struct Absorber {
2690 isotopes: Arc<[ResonanceData]>,
2691 lattice: Option<PhononSpectrum>,
2692 cross_sections: [RefCell<Option<CrossSections>>; 3],
2693}
2694
2695fn broadened(
2698 energies: &[f64],
2699 isotope: &ResonanceData,
2700 temperature_k: f64,
2701 channel: Channel,
2702 lattice: Option<&PhononSpectrum>,
2703) -> Result<TierOneBroadening, String> {
2704 match lattice {
2705 Some(spectrum) => broaden_in_lattice(energies, isotope, temperature_k, spectrum, channel)
2706 .map_err(|e| e.to_string()),
2707 None => broaden_with_budget(
2708 energies,
2709 isotope,
2710 temperature_k,
2711 channel,
2712 QuadratureBudget::default(),
2713 )
2714 .map_err(|e| e.to_string()),
2715 }
2716}
2717
2718const CHANNELS: [Channel; 3] = [Channel::Total, Channel::Capture, Channel::Elastic];
2719const TOTAL: usize = 0;
2720const CAPTURE: usize = 1;
2721const ELASTIC: usize = 2;
2722
2723struct Beams {
2724 open: Vec<f64>,
2725 normalized: Vec<f64>,
2726 sample: Vec<f64>,
2727 per_depth: Vec<f64>,
2728 per_variance: Vec<f64>,
2729 capture: Option<Captured>,
2730 scatter: Option<Scattered>,
2731}
2732
2733struct Captured {
2737 per_efficiency: Vec<f64>,
2738 per_depth: Vec<f64>,
2739 share: Vec<f64>,
2740 per_variance: Vec<f64>,
2741}
2742
2743struct Scattered {
2746 per_acceptance: Vec<f64>,
2747 shares: Vec<ForwardYield>,
2748}
2749
2750struct CrossSections {
2751 temperature_k: f64,
2752 values: Vec<Vec<f64>>,
2753 slopes: Vec<Vec<f64>>,
2754}
2755
2756struct Scale {
2757 t0_us: f64,
2758 flight_path_m: f64,
2759 pulse: IkedaCarpenterParams,
2760 rows: Rows,
2761}
2762
2763impl Clocks {
2764 fn new(
2767 scales: Vec<Scale>,
2768 of_region: Vec<usize>,
2769 grid: &FlightTimeGrid,
2770 splines: &[BeamSpline],
2771 beam_origin_us: f64,
2772 ) -> Self {
2773 let (basis, basis_slope) = splines
2774 .iter()
2775 .zip(&of_region)
2776 .map(|(spline, &group)| {
2777 let scale = &scales[group];
2778 let (shift, stretch) = (
2779 scale.t0_us - beam_origin_us,
2780 scale.flight_path_m / grid.flight_path_m(),
2781 );
2782 let abscissae: Vec<f64> = grid
2783 .flight_times_us()
2784 .iter()
2785 .map(|u| shift + stretch * u)
2786 .collect();
2787 let basis: Vec<_> = abscissae.iter().map(|&a| spline.basis(a)).collect();
2788 let slope: Vec<_> = abscissae.iter().map(|&a| spline.basis_slope(a)).collect();
2789 (basis, slope)
2790 })
2791 .unzip();
2792 Self {
2793 scales,
2794 of_region,
2795 basis,
2796 basis_slope,
2797 }
2798 }
2799}
2800
2801fn predicted(rows: &Rows, values: &[f64]) -> Result<Vec<f64>, FittingError> {
2802 rows.predict(values)
2803 .map_err(|e| FittingError::EvaluationFailed(e.to_string()))
2804}
2805
2806pub(crate) fn starts_and_widths_us(time_edges_us: &[f64]) -> Vec<(f64, f64)> {
2807 time_edges_us
2808 .windows(2)
2809 .map(|edges| (edges[0] - time_edges_us[0], edges[1] - edges[0]))
2810 .collect()
2811}
2812
2813pub(crate) fn time_counts_per_pixel(
2814 starts_and_widths_us: &[(f64, f64)],
2815 [c, a, tau]: [f64; 3],
2816) -> (Vec<f64>, [Vec<f64>; 3]) {
2817 let bins = starts_and_widths_us.len();
2818 let mut counts = Vec::with_capacity(bins);
2819 let mut slopes: [Vec<f64>; 3] = std::array::from_fn(|_| Vec::with_capacity(bins));
2820 for &(start, width) in starts_and_widths_us {
2821 let decay = (-start / tau).exp();
2822 let after = (-(start + width) / tau).exp();
2823 let fraction = -(-width / tau).exp_m1();
2824 let integral = decay * (tau * fraction);
2825 counts.push(c * width + a * integral);
2826 slopes[0].push(width);
2827 slopes[1].push(integral);
2828 slopes[2].push(a * (decay * fraction + (decay * start * fraction - after * width) / tau));
2829 }
2830 (counts, slopes)
2831}
2832
2833impl RegionsModel {
2834 fn new(
2835 grid: &Arc<FlightTimeGrid>,
2836 beams: &[BeamSpline],
2837 isotopes: &[Arc<[ResonanceData]>],
2838 (charge_ratio, frames): (f64, [f64; 2]),
2839 beam_origin_us: f64,
2840 layout: &Arc<Layout>,
2841 ) -> Self {
2842 let flight_energies = grid.energies_ev();
2843 let mut energies = flight_energies.clone();
2844 let shapes = [
2845 vec![1.0; energies.len()],
2846 energies.iter().map(|e| 1.0 / e.sqrt()).collect(),
2847 energies.iter().map(|e| e.sqrt()).collect(),
2848 ];
2849 energies.reverse();
2850 Self {
2851 grid: Arc::clone(grid),
2852 bins: starts_and_widths_us(grid.time_edges_us()),
2853 splines: beams.to_vec(),
2854 beam_origin_us,
2855 absorbers: isotopes
2856 .iter()
2857 .map(|isotopes| Absorber {
2858 isotopes: Arc::clone(isotopes),
2859 lattice: None,
2860 cross_sections: [RefCell::new(None), RefCell::new(None), RefCell::new(None)],
2861 })
2862 .collect(),
2863 energies,
2864 flight_energies,
2865 shapes,
2866 charge_ratio,
2867 frames,
2868 layout: Arc::clone(layout),
2869 clocks: RefCell::new(Some(Clocks::new(
2870 vec![Scale {
2871 t0_us: grid.t0_us(),
2872 flight_path_m: grid.flight_path_m(),
2873 pulse: grid.pulse().params().clone(),
2874 rows: grid.rows().clone(),
2875 }],
2876 vec![0; layout.regions.len()],
2877 grid,
2878 beams,
2879 beam_origin_us,
2880 ))),
2881 mixing: None,
2882 }
2883 }
2884
2885 fn lattices(mut self, lattices: &[Option<PhononSpectrum>]) -> Self {
2888 for (absorber, lattice) in self.absorbers.iter_mut().zip(lattices) {
2889 absorber.lattice = lattice.clone();
2890 }
2891 self
2892 }
2893
2894 fn spread(mut self, mixing: Option<Arc<Vec<Vec<f64>>>>) -> Self {
2896 self.mixing = mixing;
2897 self
2898 }
2899
2900 fn at(
2901 &self,
2902 region: usize,
2903 temperature_k: f64,
2904 channel: usize,
2905 ) -> Result<std::cell::Ref<'_, CrossSections>, FittingError> {
2906 let absorber = &self.absorbers[region];
2907 let current = absorber.cross_sections[channel]
2908 .borrow()
2909 .as_ref()
2910 .is_some_and(|c| c.temperature_k.to_bits() == temperature_k.to_bits());
2911 if !current {
2912 let energies = &self.energies;
2913 let lattice = absorber.lattice.as_ref();
2914 let (values, slopes) = absorber
2915 .isotopes
2916 .par_iter()
2917 .map(|isotope| {
2918 let broadened =
2919 broadened(energies, isotope, temperature_k, CHANNELS[channel], lattice)
2920 .map_err(FittingError::EvaluationFailed)?;
2921 let (mut sigma, mut slope) = (broadened.values, broadened.derivatives);
2922 sigma.reverse();
2923 slope.reverse();
2924 Ok((sigma, slope))
2925 })
2926 .collect::<Result<Vec<_>, FittingError>>()?
2927 .into_iter()
2928 .unzip();
2929 *absorber.cross_sections[channel].borrow_mut() = Some(CrossSections {
2930 temperature_k,
2931 values,
2932 slopes,
2933 });
2934 }
2935 Ok(std::cell::Ref::map(
2936 absorber.cross_sections[channel].borrow(),
2937 |c| c.as_ref().expect("computed above"),
2938 ))
2939 }
2940
2941 fn clocks(&self, params: &[f64]) -> Result<std::cell::Ref<'_, Clocks>, FittingError> {
2942 let starts = self.layout.starts(params);
2943 let flight_path_m = params[self.layout.flight_path];
2944 let pulse = laws(¶ms[self.layout.pulse..], self.layout.storage_span);
2945 let at = |scale: &Scale, t0_us: f64| {
2946 (scale.t0_us.to_bits(), scale.flight_path_m.to_bits())
2947 == (t0_us.to_bits(), flight_path_m.to_bits())
2948 && scale.pulse == pulse
2949 };
2950 let current = self.clocks.borrow().as_ref().is_some_and(|clocks| {
2951 starts
2952 .iter()
2953 .enumerate()
2954 .all(|(r, &start)| at(clocks.of(r), start))
2955 });
2956 if !current {
2957 let mut scales: Vec<Scale> = Vec::new();
2958 let mut of_region = Vec::with_capacity(starts.len());
2959 for start in starts {
2960 if let Some(group) = scales.iter().position(|scale| at(scale, start)) {
2961 of_region.push(group);
2962 continue;
2963 }
2964 let rows = self
2965 .grid
2966 .rows_at(start, flight_path_m, &pulse)
2967 .map_err(|e| FittingError::EvaluationFailed(e.to_string()))?;
2968 of_region.push(scales.len());
2969 scales.push(Scale {
2970 t0_us: start,
2971 flight_path_m,
2972 pulse: pulse.clone(),
2973 rows,
2974 });
2975 }
2976 *self.clocks.borrow_mut() = Some(Clocks::new(
2977 scales,
2978 of_region,
2979 &self.grid,
2980 &self.splines,
2981 self.beam_origin_us,
2982 ));
2983 }
2984 Ok(std::cell::Ref::map(self.clocks.borrow(), |c| {
2985 c.as_ref().expect("computed above")
2986 }))
2987 }
2988
2989 fn beams(&self, params: &[f64], clocks: &Clocks, region: usize) -> Result<Beams, FittingError> {
2990 let place = &self.layout.regions[region];
2991 let densities = ¶ms[place.densities.clone()];
2992 let sigma = place
2993 .temperature
2994 .map(|t| self.at(region, params[t], TOTAL))
2995 .transpose()?;
2996 let open: Vec<f64> = combined(&clocks.basis[region], ¶ms[place.beam.clone()])
2997 .into_iter()
2998 .map(f64::exp)
2999 .collect();
3000 let normalized: Vec<f64> = open
3001 .iter()
3002 .map(|phi| self.charge_ratio * params[self.layout.normalization] * phi)
3003 .collect();
3004 let terms = ¶ms[place.background..place.background + 5];
3005 let exponential = self.exponential(terms[4]);
3006 let variance = place.thickness.map_or(0.0, |v| params[v]);
3007 let capture_sigma = match place.temperature {
3008 Some(t) if place.capture_on => Some(self.at(region, params[t], CAPTURE)?),
3009 _ => None,
3010 };
3011 let efficiency = params[place.capture];
3012 let n = normalized.len();
3013 let depths: Vec<f64> = (0..n)
3014 .map(|j| {
3015 sigma
3016 .iter()
3017 .flat_map(|sigma| densities.iter().zip(&sigma.values))
3018 .map(|(n, sigma)| n * sigma[j])
3019 .sum()
3020 })
3021 .collect();
3022 let elastic = match place.temperature {
3023 Some(t) if place.scatter_on => Some(self.at(region, params[t], ELASTIC)?),
3024 _ => None,
3025 };
3026 let acceptance = params[place.scatter];
3027 let mut scatter = elastic
3028 .as_ref()
3029 .map(|_| {
3030 Ok::<_, FittingError>(Scattered {
3031 per_acceptance: Vec::with_capacity(n),
3032 shares: self.absorbers[region]
3033 .isotopes
3034 .iter()
3035 .map(|isotope| {
3036 forward_yield(&self.flight_energies, &depths, isotope.awr)
3037 .map_err(|e| FittingError::EvaluationFailed(e.to_string()))
3038 })
3039 .collect::<Result<_, _>>()?,
3040 })
3041 })
3042 .transpose()?;
3043 let (mut sample, mut per_depth, mut per_variance) = (
3044 Vec::with_capacity(n),
3045 Vec::with_capacity(n),
3046 Vec::with_capacity(n),
3047 );
3048 let mut capture = capture_sigma.as_ref().map(|_| Captured {
3049 per_efficiency: Vec::with_capacity(n),
3050 per_depth: Vec::with_capacity(n),
3051 share: Vec::with_capacity(n),
3052 per_variance: Vec::with_capacity(n),
3053 });
3054 for (j, beam) in normalized.iter().enumerate() {
3055 let depth = depths[j];
3056 let mut captured = 0.0;
3057 if let (Some(capture), Some(capture_sigma)) = (capture.as_mut(), &capture_sigma) {
3058 let absorbing: f64 = densities
3059 .iter()
3060 .zip(&capture_sigma.values)
3061 .map(|(n, sigma)| n * sigma[j])
3062 .sum();
3063 let (fraction, along_depth, along_variance) =
3064 interacting_per_depth(depth, variance);
3065 captured = efficiency * absorbing * fraction;
3066 capture.per_efficiency.push(beam * absorbing * fraction);
3067 capture
3068 .per_depth
3069 .push(beam * efficiency * absorbing * along_depth);
3070 capture.share.push(beam * efficiency * fraction);
3071 capture
3072 .per_variance
3073 .push(beam * efficiency * absorbing * along_variance);
3074 }
3075 let background: f64 = terms
3076 .iter()
3077 .zip(&self.shapes)
3078 .map(|(b, g)| b * g[j])
3079 .sum::<f64>()
3080 + terms[3] * exponential[j];
3081 let mut scattered = 0.0;
3082 if let (Some(scatter), Some(elastic)) = (scatter.as_mut(), &elastic) {
3083 let per_acceptance: f64 = densities
3084 .iter()
3085 .zip(&elastic.values)
3086 .zip(&scatter.shares)
3087 .map(|((n, sigma), share)| n * sigma[j] * share.share[j])
3088 .sum();
3089 scattered = acceptance * per_acceptance;
3090 scatter.per_acceptance.push(beam * per_acceptance);
3091 }
3092 let (transmission, along_depth, along_variance) = uneven_transmission(depth, variance);
3093 sample.push(beam * (transmission + background + captured + scattered));
3094 per_depth.push(beam * along_depth);
3095 per_variance.push(beam * along_variance);
3096 }
3097 Ok(Beams {
3098 open,
3099 normalized,
3100 sample,
3101 per_depth,
3102 per_variance,
3103 capture,
3104 scatter,
3105 })
3106 }
3107
3108 fn resonance_slope(
3112 &self,
3113 params: &[f64],
3114 region: usize,
3115 isotope: usize,
3116 [up, down]: [&ResonanceData; 2],
3117 h: f64,
3118 ) -> Result<Vec<f64>, PipelineError> {
3119 let place = &self.layout.regions[region];
3120 let temperature = params[place.temperature.expect("a material")];
3121 let density = params[place.densities.start + isotope];
3122 let clocks = self.clocks(params)?;
3123 let scale = clocks.of(region);
3124 let beams = self.beams(params, &clocks, region)?;
3125 let slope = |channel| -> Result<Vec<f64>, PipelineError> {
3126 let lattice = self.absorbers[region].lattice.as_ref();
3127 let sigma = |data| {
3128 broadened(&self.energies, data, temperature, channel, lattice)
3129 .map(|b| b.values)
3130 .map_err(|e| PipelineError::Fitting(FittingError::EvaluationFailed(e)))
3131 };
3132 let (up, down) = (sigma(up)?, sigma(down)?);
3133 Ok(up
3134 .iter()
3135 .zip(&down)
3136 .rev()
3137 .map(|(u, d)| density * (u - d) / (2.0 * h))
3138 .collect())
3139 };
3140 let total = slope(Channel::Total)?;
3141 let mut column: Vec<f64> = total
3142 .iter()
3143 .zip(&beams.per_depth)
3144 .map(|(s, t)| -t * s)
3145 .collect();
3146 if let Some(capture) = &beams.capture {
3147 let captured = slope(Channel::Capture)?;
3148 for (j, c) in column.iter_mut().enumerate() {
3149 *c += capture.per_depth[j] * total[j] + capture.share[j] * captured[j];
3150 }
3151 }
3152 if let Some(scatter) = &beams.scatter {
3153 let scattered = slope(Channel::Elastic)?;
3154 let elastic = self.at(region, temperature, ELASTIC)?;
3155 let densities = ¶ms[place.densities.clone()];
3156 for (j, c) in column.iter_mut().enumerate() {
3157 let mut through = scattered[j] * scatter.shares[isotope].share[j];
3158 for (m, share) in scatter.shares.iter().enumerate() {
3159 let entry = share.per_entry[j] * total[j]
3160 + share.per_exit[j]
3161 .iter()
3162 .map(|&(l, w)| w * total[l])
3163 .sum::<f64>();
3164 through += densities[m] * elastic.values[m][j] * entry;
3165 }
3166 *c += beams.normalized[j] * params[place.scatter] * through;
3167 }
3168 }
3169 let folded = predicted(&scale.rows, &column)?;
3170 let bins = folded.len();
3171 let mut arriving = vec![0.0; 2 * bins * self.layout.regions.len()];
3172 arriving[(2 * region + 1) * bins..(2 * region + 2) * bins].copy_from_slice(&folded);
3173 Ok(match &self.mixing {
3174 Some(shares) => mixed(shares, params[self.layout.spread], &arriving),
3175 None => arriving,
3176 })
3177 }
3178
3179 fn time_background(&self, params: &[f64], region: usize) -> (Vec<f64>, [Vec<f64>; 3]) {
3182 let at = self.layout.time_background;
3183 let bins = late(&self.bins, params[self.layout.regions[region].offset]);
3184 time_counts_per_pixel(&bins, [params[at], params[at + 1], params[at + 2]])
3185 }
3186
3187 fn exponential(&self, f: f64) -> Vec<f64> {
3188 self.shapes[1]
3189 .iter()
3190 .map(|root| (-f * root).exp())
3191 .collect()
3192 }
3193
3194 fn room_per_pixel(&self) -> [Vec<f64>; 2] {
3195 self.frames
3196 .map(|frames| self.bins.iter().map(|(_, width)| frames * width).collect())
3197 }
3198
3199 fn sample_time_scale(&self, params: &[f64], region: usize) -> f64 {
3200 let place = &self.layout.regions[region];
3201 self.charge_ratio * params[self.layout.normalization] * params[place.gamma] * place.pixels
3202 }
3203}
3204
3205impl RegionsModel {
3206 fn arriving(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
3207 let clocks = self.clocks(params)?;
3208 let [open_room, sample_room] = self.room_per_pixel();
3209 let rate = params[self.layout.room_background];
3210 let mut counts = Vec::new();
3211 for region in 0..self.absorbers.len() {
3212 let scale = clocks.of(region);
3213 let (background, _) = self.time_background(params, region);
3214 let Beams { open, sample, .. } = self.beams(params, &clocks, region)?;
3215 let pixels = self.layout.regions[region].pixels;
3216 let sample_scale = self.sample_time_scale(params, region);
3217 counts.extend(
3218 predicted(&scale.rows, &open)?
3219 .iter()
3220 .zip(&background)
3221 .zip(&open_room)
3222 .map(|((c, g), f)| c + pixels * (g + rate * f)),
3223 );
3224 counts.extend(
3225 predicted(&scale.rows, &sample)?
3226 .iter()
3227 .zip(&background)
3228 .zip(&sample_room)
3229 .map(|((c, g), f)| c + sample_scale * g + pixels * rate * f),
3230 );
3231 }
3232 Ok(counts)
3233 }
3234}
3235
3236fn mixed(shares: &[Vec<f64>], h: f64, arriving: &[f64]) -> Vec<f64> {
3240 let block = arriving.len() / shares.len();
3241 (0..arriving.len())
3242 .map(|i| {
3243 let (r, k) = (i / block, i % block);
3244 (1.0 - h) * arriving[i]
3245 + h * shares[r]
3246 .iter()
3247 .enumerate()
3248 .map(|(q, f)| f * arriving[q * block + k])
3249 .sum::<f64>()
3250 })
3251 .collect()
3252}
3253
3254impl FitModel for RegionsModel {
3255 fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
3258 let arriving = self.arriving(params)?;
3259 let Some(shares) = &self.mixing else {
3260 return Ok(arriving);
3261 };
3262 let mut recorded = mixed(shares, params[self.layout.spread], &arriving);
3263 for (recorded, &arrived) in recorded.iter_mut().zip(&arriving) {
3264 if arrived.is_nan() || arrived < 0.0 {
3265 *recorded = arrived;
3266 }
3267 }
3268 Ok(recorded)
3269 }
3270
3271 fn analytical_jacobian(
3272 &self,
3273 params: &[f64],
3274 free_param_indices: &[usize],
3275 y_current: &[f64],
3276 ) -> Option<FlatMatrix> {
3277 let layout = &*self.layout;
3278 let flight_path_m = params[layout.flight_path];
3279 let clocks = self.clocks(params).ok()?;
3280 let regions = self.absorbers.len();
3281 let beams: Vec<Beams> = (0..regions)
3282 .map(|region| self.beams(params, &clocks, region))
3283 .collect::<Result<_, _>>()
3284 .ok()?;
3285 let bins = y_current.len() / (2 * regions);
3286 let (open_rows, sample_rows) = (|r: usize| 2 * r * bins, |r: usize| (2 * r + 1) * bins);
3287 let counts = |r: usize, values: &[f64]| predicted(&clocks.of(r).rows, values).ok();
3288 let times = |beam: &[f64], slope: &[f64]| -> Vec<f64> {
3289 beam.iter().zip(slope).map(|(b, s)| b * s).collect()
3290 };
3291 let groups = clocks.scales.len();
3292 let mut arrival: Vec<Option<Rows>> = (0..groups).map(|_| None).collect();
3293 let mut pulse_slopes: Vec<Option<[Rows; 4]>> = (0..groups).map(|_| None).collect();
3294 let mut tail_slopes: Vec<Option<[Rows; 2]>> = (0..groups).map(|_| None).collect();
3295 let ones = vec![1.0; self.energies.len()];
3296 let root_energy = &self.shapes[2];
3297 let constant = || [ones.clone(), vec![0.0; ones.len()]];
3298 let [low, high] = self.layout.storage_span.map_or_else(constant, |span| {
3299 let [first, last] = [span.0, span.1].map(f64::sqrt);
3300 let share: Vec<f64> = root_energy
3301 .iter()
3302 .map(|root| (root - first) / (last - first))
3303 .collect();
3304 [share.iter().map(|s| 1.0 - s).collect(), share]
3305 });
3306 let pulse_columns: [(usize, &[f64]); 9] = [
3307 (0, root_energy),
3308 (0, &ones),
3309 (1, root_energy),
3310 (1, &ones),
3311 (2, &low),
3312 (3, &ones),
3313 (2, &high),
3314 (4, &ones),
3315 (5, &ones),
3316 ];
3317 let backgrounds: Vec<(Vec<f64>, [Vec<f64>; 3])> = (0..regions)
3318 .map(|r| self.time_background(params, r))
3319 .collect();
3320 let scaled = |factor: f64, values: &[f64]| -> Vec<f64> {
3321 values.iter().map(|v| factor * v).collect()
3322 };
3323 let mut jacobian = FlatMatrix::zeros(y_current.len(), free_param_indices.len());
3324 for (col, &index) in free_param_indices.iter().enumerate() {
3325 let blocks: Vec<(usize, Vec<f64>)> = match layout.roles[index] {
3326 Role::Beam {
3327 region,
3328 coefficient,
3329 } => {
3330 let slope = weights_of(&clocks.basis[region], coefficient);
3331 let Beams { open, sample, .. } = &beams[region];
3332 vec![
3333 (open_rows(region), counts(region, ×(open, &slope))?),
3334 (sample_rows(region), counts(region, ×(sample, &slope))?),
3335 ]
3336 }
3337 Role::Density { region, isotope } => {
3338 let temperature = params[layout.regions[region].temperature?];
3339 let sigma = self.at(region, temperature, TOTAL).ok()?;
3340 let total = &sigma.values[isotope];
3341 let mut column: Vec<f64> = beams[region]
3342 .per_depth
3343 .iter()
3344 .zip(total)
3345 .map(|(d, s)| -d * s)
3346 .collect();
3347 if let Some(capture) = &beams[region].capture {
3348 let captures = self.at(region, temperature, CAPTURE).ok()?;
3349 for (j, c) in column.iter_mut().enumerate() {
3350 *c += capture.per_depth[j] * total[j]
3351 + capture.share[j] * captures.values[isotope][j];
3352 }
3353 }
3354 if let Some(scatter) = &beams[region].scatter {
3355 let place = &layout.regions[region];
3356 let elastic = self.at(region, temperature, ELASTIC).ok()?;
3357 let densities = ¶ms[place.densities.clone()];
3358 for (j, c) in column.iter_mut().enumerate() {
3359 let mut slope =
3360 elastic.values[isotope][j] * scatter.shares[isotope].share[j];
3361 for (m, share) in scatter.shares.iter().enumerate() {
3362 let through = share.per_entry[j] * total[j]
3363 + share.per_exit[j]
3364 .iter()
3365 .map(|&(l, w)| w * total[l])
3366 .sum::<f64>();
3367 slope += densities[m] * elastic.values[m][j] * through;
3368 }
3369 *c += beams[region].normalized[j] * params[place.scatter] * slope;
3370 }
3371 }
3372 vec![(sample_rows(region), counts(region, &column)?)]
3373 }
3374 Role::Temperature { region } => {
3375 let place = &layout.regions[region];
3376 let sigma = self.at(region, params[index], TOTAL).ok()?;
3377 let densities = ¶ms[place.densities.clone()];
3378 let along = |slopes: &[Vec<f64>], j: usize| -> f64 {
3379 densities
3380 .iter()
3381 .zip(slopes)
3382 .map(|(n, slope)| n * slope[j])
3383 .sum()
3384 };
3385 let mut column: Vec<f64> = (0..self.energies.len())
3386 .map(|j| -beams[region].per_depth[j] * along(&sigma.slopes, j))
3387 .collect();
3388 if let Some(capture) = &beams[region].capture {
3389 let captures = self.at(region, params[index], CAPTURE).ok()?;
3390 for (j, c) in column.iter_mut().enumerate() {
3391 *c += capture.per_depth[j] * along(&sigma.slopes, j)
3392 + capture.share[j] * along(&captures.slopes, j);
3393 }
3394 }
3395 if let Some(scatter) = &beams[region].scatter {
3396 let elastic = self.at(region, params[index], ELASTIC).ok()?;
3397 let depth_slopes: Vec<f64> = (0..self.energies.len())
3398 .map(|j| along(&sigma.slopes, j))
3399 .collect();
3400 for (j, c) in column.iter_mut().enumerate() {
3401 let mut slope = 0.0;
3402 for (m, share) in scatter.shares.iter().enumerate() {
3403 let through = share.per_entry[j] * depth_slopes[j]
3404 + share.per_exit[j]
3405 .iter()
3406 .map(|&(l, w)| w * depth_slopes[l])
3407 .sum::<f64>();
3408 slope += densities[m]
3409 * (elastic.slopes[m][j] * share.share[j]
3410 + elastic.values[m][j] * through);
3411 }
3412 *c += beams[region].normalized[j] * params[place.scatter] * slope;
3413 }
3414 }
3415 vec![(sample_rows(region), counts(region, &column)?)]
3416 }
3417 Role::ThicknessVariance { region } => {
3418 let mut column = beams[region].per_variance.clone();
3419 if let Some(capture) = &beams[region].capture {
3420 for (c, v) in column.iter_mut().zip(&capture.per_variance) {
3421 *c += v;
3422 }
3423 }
3424 vec![(sample_rows(region), counts(region, &column)?)]
3425 }
3426 Role::CaptureGammas { region } => vec![(
3427 sample_rows(region),
3428 match &beams[region].capture {
3429 Some(capture) => counts(region, &capture.per_efficiency)?,
3430 None => vec![0.0; bins],
3431 },
3432 )],
3433 Role::InScatter { region } => vec![(
3434 sample_rows(region),
3435 match &beams[region].scatter {
3436 Some(scatter) => counts(region, &scatter.per_acceptance)?,
3437 None => vec![0.0; bins],
3438 },
3439 )],
3440 Role::Normalization => {
3441 let slope = vec![1.0 / params[index]; self.energies.len()];
3442 (0..regions)
3443 .map(|r| {
3444 let place = &layout.regions[r];
3445 let gamma = self.charge_ratio * params[place.gamma] * place.pixels;
3446 let beam = counts(r, ×(&beams[r].sample, &slope))?;
3447 let column = beam
3448 .iter()
3449 .zip(&backgrounds[r].0)
3450 .map(|(b, g)| b + gamma * g)
3451 .collect();
3452 Some((sample_rows(r), column))
3453 })
3454 .collect::<Option<_>>()?
3455 }
3456 Role::Background { region, term } => {
3457 let at = layout.regions[region].background;
3458 let (d, f) = (params[at + 3], params[at + 4]);
3459 let shape: Vec<f64> = match term {
3460 3 => self.exponential(f),
3461 4 => self
3462 .exponential(f)
3463 .iter()
3464 .zip(&self.shapes[1])
3465 .map(|(x, root)| -d * root * x)
3466 .collect(),
3467 _ => self.shapes[term].clone(),
3468 };
3469 vec![(
3470 sample_rows(region),
3471 counts(region, ×(&beams[region].normalized, &shape))?,
3472 )]
3473 }
3474 Role::GammaTransmission { region } => vec![(
3475 sample_rows(region),
3476 scaled(
3477 self.charge_ratio
3478 * params[layout.normalization]
3479 * layout.regions[region].pixels,
3480 &backgrounds[region].0,
3481 ),
3482 )],
3483 Role::TimeBackground(term) => (0..regions)
3484 .flat_map(|r| {
3485 let slope = &backgrounds[r].1[term];
3486 [
3487 (open_rows(r), scaled(layout.regions[r].pixels, slope)),
3488 (
3489 sample_rows(r),
3490 scaled(self.sample_time_scale(params, r), slope),
3491 ),
3492 ]
3493 })
3494 .collect(),
3495 Role::SpreadShare => vec![],
3496 Role::RoomBackground => {
3497 let [open_room, sample_room] = self.room_per_pixel();
3498 (0..regions)
3499 .flat_map(|r| {
3500 let pixels = layout.regions[r].pixels;
3501 [
3502 (open_rows(r), scaled(pixels, &open_room)),
3503 (sample_rows(r), scaled(pixels, &sample_room)),
3504 ]
3505 })
3506 .collect()
3507 }
3508 Role::Pulse(number) => {
3509 let (law, weight) = pulse_columns[number];
3510 let mut blocks = Vec::with_capacity(2 * regions);
3511 for (r, region) in beams.iter().enumerate() {
3512 let (group, scale) = (clocks.of_region[r], clocks.of(r));
3513 let rows = if law < 4 {
3514 if pulse_slopes[group].is_none() {
3515 pulse_slopes[group] = Some(
3516 self.grid
3517 .pulse_slopes_at(scale.t0_us, flight_path_m, &scale.pulse)
3518 .ok()?,
3519 );
3520 }
3521 &pulse_slopes[group].as_ref()?[law]
3522 } else {
3523 if tail_slopes[group].is_none() {
3524 tail_slopes[group] = Some(
3525 self.grid
3526 .tail_slopes_at(scale.t0_us, flight_path_m, &scale.pulse)
3527 .ok()?,
3528 );
3529 }
3530 &tail_slopes[group].as_ref()?[law - 4]
3531 };
3532 for (row, beam) in [
3533 (open_rows(r), ®ion.open),
3534 (sample_rows(r), ®ion.sample),
3535 ] {
3536 blocks.push((row, predicted(rows, ×(beam, weight)).ok()?));
3537 }
3538 }
3539 blocks
3540 }
3541 Role::T0 | Role::FlightPath | Role::TimingOffset { .. } => {
3542 let only = match layout.roles[index] {
3543 Role::TimingOffset { region } => Some(region),
3544 _ => None,
3545 };
3546 let (reach, per_metre): (Vec<f64>, Option<f64>) = match layout.roles[index] {
3547 Role::FlightPath => (
3548 self.grid
3549 .flight_times_us()
3550 .iter()
3551 .map(|u| u / self.grid.flight_path_m())
3552 .collect(),
3553 Some(flight_path_m),
3554 ),
3555 _ => (ones.clone(), None),
3556 };
3557 let at = layout.time_background;
3558 let per_late = params[at + 1] / params[at + 2];
3559 let mut blocks = Vec::with_capacity(2 * regions);
3560 for r in (0..regions).filter(|&r| only.is_none_or(|region| region == r)) {
3561 let (group, scale) = (clocks.of_region[r], clocks.of(r));
3562 if arrival[group].is_none() {
3563 arrival[group] = Some(
3564 self.grid
3565 .arrival_slopes_at(scale.t0_us, flight_path_m, &scale.pulse)
3566 .ok()?,
3567 );
3568 }
3569 let arrival = arrival[group].as_ref()?;
3570 let log_beam_slope = combined(
3571 &clocks.basis_slope[r],
3572 ¶ms[layout.regions[r].beam.clone()],
3573 );
3574 for (row, beam, background) in [
3575 (open_rows(r), &beams[r].open, layout.regions[r].pixels),
3576 (
3577 sample_rows(r),
3578 &beams[r].sample,
3579 self.sample_time_scale(params, r),
3580 ),
3581 ] {
3582 let reached = times(beam, &reach);
3583 let beam_part = counts(r, ×(&reached, &log_beam_slope))?;
3584 let pulse_part = predicted(arrival, &reached).ok()?;
3585 let mut column: Vec<f64> = beam_part
3586 .iter()
3587 .zip(&pulse_part)
3588 .map(|(b, p)| b + p)
3589 .collect();
3590 if let Some(flight_path_m) = per_metre {
3591 for (c, x) in column.iter_mut().zip(counts(r, beam)?) {
3592 *c += x / flight_path_m;
3593 }
3594 }
3595 if only.is_some() {
3596 for (c, x) in column.iter_mut().zip(&backgrounds[r].1[1]) {
3597 *c += background * per_late * x;
3598 }
3599 }
3600 blocks.push((row, column));
3601 }
3602 }
3603 blocks
3604 }
3605 };
3606 for (row, column) in blocks {
3607 for (offset, value) in column.into_iter().enumerate() {
3608 *jacobian.get_mut(row + offset, col) = value;
3609 }
3610 }
3611 }
3612 let Some(shares) = &self.mixing else {
3613 return Some(jacobian);
3614 };
3615 let h = params[layout.spread];
3616 let arriving = self.arriving(params).ok()?;
3617 let along_share: Vec<f64> = mixed(shares, 1.0, &arriving)
3618 .iter()
3619 .zip(&arriving)
3620 .map(|(m, a)| m - a)
3621 .collect();
3622 let mut recorded = FlatMatrix::zeros(y_current.len(), free_param_indices.len());
3623 for (col, &index) in free_param_indices.iter().enumerate() {
3624 let column = if layout.roles[index] == Role::SpreadShare {
3625 along_share.clone()
3626 } else {
3627 let slope: Vec<f64> = (0..y_current.len()).map(|i| jacobian.get(i, col)).collect();
3628 mixed(shares, h, &slope)
3629 };
3630 for (i, value) in column.into_iter().enumerate() {
3631 *recorded.get_mut(i, col) = value;
3632 }
3633 }
3634 Some(recorded)
3635 }
3636}
3637
3638#[cfg(test)]
3639mod tests {
3640 use nereids_endf::resonance::test_support::synthetic_isotope;
3641
3642 use faer::Side;
3643 use faer::linalg::solvers::DenseSolveCore;
3644
3645 use super::*;
3646 use crate::open_beam::Pulse;
3647 use crate::open_beam::tests::{ALPHA, BETA, EDGES_US, FLIGHT_PATH_M, R, T0_US, grid};
3648
3649 #[test]
3650 fn a_fitted_temperature_starts_its_second_fit_hot_within_its_bounds() {
3651 let region = |temperature_k| Region {
3652 open_counts: Vec::new(),
3653 sample_counts: Vec::new(),
3654 open_live: None,
3655 sample_live: None,
3656 background: [Value::Known(0.0); 5],
3657 pixels: 1.0,
3658 gamma_transmission: Value::Known(1.0),
3659 capture_gammas: Value::Known(0.0),
3660 in_scatter: Value::Known(0.0),
3661 timing_offset_us: Value::Known(0.0),
3662 material: Some(Material {
3663 covariances: Vec::new(),
3664 library: EndfLibrary::EndfB8_0,
3665 isotopes: vec![(
3666 synthetic_isotope(72, 180, 20.0, 0.01, 0.06),
3667 Value::Fitted(1e-4),
3668 )],
3669 temperature_k,
3670 thickness_variance: Value::Known(0.0),
3671 lattice: None,
3672 }),
3673 };
3674 let started = |temperatures: &[Value]| {
3675 let measurement = Measurement {
3676 time_edges_us: Vec::new(),
3677 charge_ratio: 1.0,
3678 normalization: Value::Known(1.0),
3679 time_background: NO_TIME_BACKGROUND,
3680 room_background: None,
3681 spread: None,
3682 regions: temperatures.iter().map(|&t| region(t)).collect(),
3683 };
3684 started_hot(&measurement).map(|hot| {
3685 hot.regions
3686 .iter()
3687 .map(|r| r.material.as_ref().expect("material").temperature_k)
3688 .collect::<Vec<_>>()
3689 })
3690 };
3691 let within = |start| Value::Within {
3692 start,
3693 lower: 250.0,
3694 upper: 400.0,
3695 };
3696 assert_eq!(
3697 started(&[Value::Fitted(300.0), Value::Known(300.0), within(300.0)]),
3698 Some(vec![
3699 Value::Fitted(2500.0),
3700 Value::Known(300.0),
3701 within(400.0)
3702 ])
3703 );
3704 let measured = Value::Measured {
3705 value: 300.0,
3706 sd: 3.0,
3707 };
3708 assert_eq!(
3709 started(&[Value::Fitted(2500.0), within(400.0), measured]),
3710 None
3711 );
3712 }
3713
3714 #[test]
3715 fn the_jacobian_is_the_slope_of_every_region_s_counts() {
3716 for channel_fwhm_us in [None, Some(0.35)] {
3717 let lattice = channel_fwhm_us
3718 .is_none()
3719 .then(|| PhononSpectrum::debye(240.0, 200).expect("Debye"));
3720 jacobian_against_central_differences(&grid(channel_fwhm_us), lattice);
3721 }
3722 let bins = starts_and_widths_us(&[350.0, 351.0, 400.0, 470.0]);
3723 let (_, slopes) = time_counts_per_pixel(&bins, [1.0, 2.0, f64::MIN_POSITIVE]);
3724 assert!(slopes.iter().flatten().all(|s| s.is_finite()), "{slopes:?}");
3725 let (counts, _) = time_counts_per_pixel(&bins, [1.0, 2.0, f64::MAX]);
3726 for (count, (_, width)) in counts.iter().zip(&bins) {
3727 assert!((count / (3.0 * width) - 1.0).abs() <= 1e-12, "{counts:?}");
3728 }
3729 }
3730
3731 #[test]
3732 fn a_region_records_the_time_background_its_timing_offset_late() {
3733 let grid = grid(None);
3734 let (_, u_hi) = grid.range_us();
3735 let beams = [BeamSpline::constant(347.0, u_hi, 1.0e4)];
3736 let layout = Arc::new(Layout::new(
3737 roles(&[beams[0].coefficients().len()], &[None]),
3738 &[1.0],
3739 &[[false; 2]],
3740 ));
3741 let ([c, a, tau], late) = ([2000.0, 30000.0, 10.0], 0.3);
3742 let params: Vec<f64> = layout
3743 .roles
3744 .iter()
3745 .map(|role| match *role {
3746 Role::TimingOffset { .. } => late,
3747 Role::TimeBackground(term) => [c, a, tau][term],
3748 Role::T0 => T0_US,
3749 Role::FlightPath => FLIGHT_PATH_M,
3750 _ => 0.0,
3751 })
3752 .collect();
3753 let model = RegionsModel::new(
3754 &grid,
3755 &beams,
3756 &[Arc::new([])],
3757 (1.2, [0.0; 2]),
3758 T0_US,
3759 &layout,
3760 );
3761 let (recorded, _) = model.time_background(¶ms, 0);
3762 let edges = grid.time_edges_us();
3763 let rate = |t: f64| c + a * (-(t - late - edges[0]) / tau).exp();
3764 for (k, counts) in recorded.iter().enumerate() {
3765 let h = (edges[k + 1] - edges[k]) / 64.0;
3766 let simpson: f64 = (0..64)
3767 .map(|i| edges[k] + h * f64::from(i))
3768 .map(|t| h / 6.0 * (rate(t) + 4.0 * rate(t + 0.5 * h) + rate(t + h)))
3769 .sum();
3770 assert!(
3771 (counts / simpson - 1.0).abs() <= 1e-12,
3772 "{k}: {counts} {simpson}"
3773 );
3774 }
3775 }
3776
3777 #[test]
3778 fn under_a_spread_a_negative_arriving_count_is_recorded_as_it_arrives() {
3779 let grid = grid(None);
3780 let (_, u_hi) = grid.range_us();
3781 let beams = [
3782 BeamSpline::constant(347.0, u_hi, 1.0e5),
3783 BeamSpline::constant(347.0, u_hi, 1.0e4),
3784 ];
3785 let layout = Arc::new(Layout::new(
3786 roles(&[beams[0].coefficients().len(); 2], &[None, None]),
3787 &[400.0, 1.0],
3788 &[[false; 2]; 2],
3789 ));
3790 let params: Vec<f64> = layout
3791 .roles
3792 .iter()
3793 .enumerate()
3794 .map(|(i, role)| match *role {
3795 Role::Beam { region, .. } => [12.0, 9.0][region] + 0.3 * (i as f64).sin(),
3796 Role::Background { region, term } => {
3797 [[0.0; 5], [-1.5, 0.0, 0.0, 0.0, 4.0]][region][term]
3798 }
3799 Role::GammaTransmission { .. } => 1.0,
3800 Role::Normalization => 0.93,
3801 Role::TimeBackground(term) => [50.0, 300.0, 40.0][term],
3802 Role::SpreadShare => 0.5,
3803 Role::T0 => T0_US,
3804 Role::FlightPath => FLIGHT_PATH_M,
3805 Role::Pulse(n) => [0.37, 0.06, 0.01, 0.24, 0.16, 0.3, 0.0, 0.0, 1.0][n],
3806 _ => 0.0,
3807 })
3808 .collect();
3809 let model = RegionsModel::new(
3810 &grid,
3811 &beams,
3812 &[Arc::new([]), Arc::new([])],
3813 (1.2, [3.0e4, 2.5e4]),
3814 T0_US,
3815 &layout,
3816 )
3817 .spread(Some(Arc::new(vec![vec![0.9, 0.1], vec![0.1, 0.9]])));
3818 let arriving = model.arriving(¶ms).expect("arriving");
3819 let recorded = model.evaluate(¶ms).expect("recorded");
3820 let negative: Vec<usize> = (0..arriving.len()).filter(|&k| arriving[k] < 0.0).collect();
3821 assert!(!negative.is_empty());
3822 for k in negative {
3823 assert_eq!(recorded[k], arriving[k], "{k}");
3824 }
3825 }
3826
3827 #[test]
3828 fn a_resonance_slope_is_the_slope_of_the_sample_counts() {
3829 let grid = grid(None);
3830 let (_, u_hi) = grid.range_us();
3831 let beams = [
3832 BeamSpline::constant(347.0, u_hi, 1.0e4),
3833 BeamSpline::constant(347.0, u_hi, 2.0e4),
3834 ];
3835 let layout = Arc::new(Layout::new(
3836 roles(&[beams[0].coefficients().len(); 2], &[Some(1), None]),
3837 &[4.0, 9.0],
3838 &[[true; 2], [false; 2]],
3839 ));
3840 let params: Vec<f64> = layout
3841 .roles
3842 .iter()
3843 .enumerate()
3844 .map(|(i, role)| match *role {
3845 Role::Beam { .. } => 9.0 + 0.3 * (i as f64).sin(),
3846 Role::Density { .. } => 3.0e-4,
3847 Role::Temperature { .. } => 300.0,
3848 Role::ThicknessVariance { .. } => 0.02,
3849 Role::Normalization => 0.93,
3850 Role::Background { term, .. } => [0.05, 0.5, -0.01, 0.03, 4.0][term],
3851 Role::GammaTransmission { region } => [0.8, 1.0][region],
3852 Role::CaptureGammas { region } => [0.05, 0.0][region],
3853 Role::InScatter { region } => [0.3, 0.0][region],
3854 Role::TimingOffset { region } => [0.02, -0.01][region],
3855 Role::TimeBackground(term) => [50.0, 300.0, 40.0][term],
3856 Role::RoomBackground => 2.0e-3,
3857 Role::SpreadShare => 0.15,
3858 Role::T0 => T0_US + 0.05,
3859 Role::FlightPath => FLIGHT_PATH_M + 0.003,
3860 Role::Pulse(n) => [0.37, 0.06, 0.01, 0.24, 0.16, 0.3, 0.0, 0.0, 1.0][n],
3861 })
3862 .collect();
3863 let data = synthetic_isotope(72, 180, 20.0, 0.01, 0.06);
3864 let model = |data: &ResonanceData| {
3865 let isotopes: [Arc<[ResonanceData]>; 2] = [Arc::new([data.clone()]), Arc::new([])];
3866 RegionsModel::new(
3867 &grid,
3868 &beams,
3869 &isotopes,
3870 (1.2, [3.0e4, 2.5e4]),
3871 T0_US,
3872 &layout,
3873 )
3874 .spread(Some(Arc::new(vec![vec![0.9, 0.1], vec![0.2, 0.8]])))
3875 };
3876 let h = 1e-4 * data.ranges[0].l_groups[0].resonances[0].gg;
3877 let moved = |by: f64| {
3878 let mut copy = data.clone();
3879 copy.ranges[0].l_groups[0].resonances[0].gg += by;
3880 copy
3881 };
3882 let (up, down) = (moved(h), moved(-h));
3883 let slope = model(&data)
3884 .resonance_slope(¶ms, 0, 0, [&up, &down], h)
3885 .expect("slope");
3886 let counts = |data| model(data).evaluate(¶ms).expect("counts");
3887 let differences: Vec<f64> = counts(&up)
3888 .iter()
3889 .zip(&counts(&down))
3890 .map(|(u, d)| (u - d) / (2.0 * h))
3891 .collect();
3892 let column = differences.iter().fold(0.0_f64, |m, s| m.max(s.abs()));
3893 for (k, (&analytic, &difference)) in slope.iter().zip(&differences).enumerate() {
3894 assert!(
3895 (analytic - difference).abs() <= 1e-6 * column,
3896 "{k}: {analytic} vs {difference}"
3897 );
3898 }
3899 }
3900
3901 fn jacobian_against_central_differences(
3902 grid: &Arc<FlightTimeGrid>,
3903 lattice: Option<PhononSpectrum>,
3904 ) {
3905 let (_, u_hi) = grid.range_us();
3906 let coarse = BeamSpline::constant(347.0, u_hi, 1.0e4);
3907 let beams = [coarse.refined().refined(), coarse.clone(), coarse.refined()];
3908 let isotopes: [Arc<[ResonanceData]>; 3] = [
3909 Arc::new([
3910 synthetic_isotope(72, 180, 20.0, 0.01, 0.06),
3911 synthetic_isotope(74, 182, 20.3, 0.01, 0.06),
3912 ]),
3913 Arc::new([]),
3914 Arc::new([synthetic_isotope(74, 184, 20.6, 0.01, 0.06)]),
3915 ];
3916 let sizes: Vec<usize> = beams.iter().map(|b| b.coefficients().len()).collect();
3917 let mut layout = Layout::new(
3918 roles(&sizes, &[Some(2), None, Some(1)]),
3919 &[4.0, 9.0, 2.0],
3920 &[[true, true], [false, false], [true, true]],
3921 );
3922 layout.storage_span = Some(grid.pulse().energy_span_ev());
3923 let layout = Arc::new(layout);
3924 let params: Vec<f64> = layout
3925 .roles
3926 .iter()
3927 .enumerate()
3928 .map(|(i, role)| match *role {
3929 Role::Beam { .. } => 9.0 + 0.3 * (i as f64).sin(),
3930 Role::Density { region, isotope } => {
3931 [[3.0e-4, 5.0e-4], [0.0; 2], [4.0e-4, 0.0]][region][isotope]
3932 }
3933 Role::Temperature { region } => [300.0, 0.0, 250.0][region],
3934 Role::ThicknessVariance { region } => [0.02, 0.0, 0.05][region],
3935 Role::Normalization => 0.93,
3936 Role::Background { region, term } => [
3937 [0.05, 0.5, -0.01, 0.03, 4.0],
3938 [0.02, 0.1, 0.003, -0.01, 2.0],
3939 [0.04, -0.2, 0.005, 0.02, 6.0],
3940 ][region][term],
3941 Role::GammaTransmission { region } => [0.8, 1.0, 0.6][region],
3942 Role::CaptureGammas { region } => [0.05, 0.03, 0.02][region],
3943 Role::InScatter { region } => [0.3, 0.2, 0.4][region],
3944 Role::TimingOffset { region } => [0.03, 0.03, -0.02][region],
3945 Role::TimeBackground(term) => [50.0, 300.0, 40.0][term],
3946 Role::RoomBackground => 2.0e-3,
3947 Role::SpreadShare => 0.15,
3948 Role::T0 => T0_US + 0.05,
3949 Role::FlightPath => FLIGHT_PATH_M + 0.003,
3950 Role::Pulse(n) => [0.37, 0.06, 0.01, 0.24, 0.16, 0.3, 0.25, 0.1, 1.5][n],
3951 })
3952 .collect();
3953 let model = RegionsModel::new(
3954 grid,
3955 &beams,
3956 &isotopes,
3957 (1.2, [3.0e4, 2.5e4]),
3958 T0_US,
3959 &layout,
3960 )
3961 .lattices(&[lattice, None, None])
3962 .spread(Some(Arc::new(vec![
3963 vec![0.8, 0.1, 0.05],
3964 vec![0.1, 0.7, 0.1],
3965 vec![0.05, 0.15, 0.8],
3966 ])));
3967 let live: Vec<f64> = (0..model.evaluate(¶ms).expect("counts").len())
3968 .map(|k| 0.9 + 0.1 * (0.3 * k as f64).sin())
3969 .collect();
3970 let model = Recorded { model, live: &live };
3971 let mut elsewhere = params.clone();
3972 for t in layout
3973 .regions
3974 .iter()
3975 .filter_map(|region| region.temperature)
3976 {
3977 elsewhere[t] -= 50.0;
3978 }
3979 elsewhere[layout.t0] = T0_US - 0.05;
3980 elsewhere[layout.flight_path] = FLIGHT_PATH_M - 0.003;
3981 elsewhere[layout.pulse + 4] = 0.2;
3982 model.evaluate(&elsewhere).expect("counts");
3983 let counts = model.evaluate(¶ms).expect("counts");
3984 let indices: Vec<usize> = (0..params.len()).collect();
3985 let jacobian = model
3986 .analytical_jacobian(¶ms, &indices, &counts)
3987 .expect("jacobian");
3988 let middle_us = grid.flight_times_us()[grid.flight_times_us().len() / 2];
3989 for index in indices {
3990 let h = if index == layout.flight_path {
3991 params[index] * 1e-4 * params[layout.t0] / middle_us
3992 } else {
3993 1e-4 * params[index].abs()
3994 };
3995 let shifted = |d: f64| {
3996 let mut p = params.clone();
3997 p[index] += d;
3998 model.evaluate(&p).expect("counts")
3999 };
4000 let (up, down) = (shifted(h), shifted(-h));
4001 let slopes: Vec<f64> = up
4002 .iter()
4003 .zip(&down)
4004 .map(|(u, d)| (u - d) / (2.0 * h))
4005 .collect();
4006 let column = slopes.iter().fold(0.0_f64, |m, s| m.max(s.abs()));
4007 for (row, &slope) in slopes.iter().enumerate() {
4008 let analytic = jacobian.get(row, index);
4009 assert!(
4010 (analytic - slope).abs() <= 1e-6 * column,
4011 "{index} {row}: {analytic} vs {slope}"
4012 );
4013 }
4014 }
4015 }
4016
4017 #[test]
4018 fn each_region_s_terms_rebuild_the_joint_covariance() {
4019 let grid = Arc::new(grid(None).halved().expect("grid").halved().expect("grid"));
4020 let (_, u_hi) = grid.range_us();
4021 let beam = BeamSpline::constant(347.0, u_hi, 1.0e4).refined();
4022 let beams = [beam.clone(), beam.clone(), beam];
4023 let materials = [
4024 Some(synthetic_isotope(72, 180, 20.0, 0.01, 0.06)),
4025 Some(synthetic_isotope(74, 182, 24.0, 0.01, 0.06)),
4026 None,
4027 ];
4028 let isotopes: Vec<Arc<[ResonanceData]>> = materials
4029 .iter()
4030 .map(|m| m.iter().cloned().collect())
4031 .collect();
4032 let shape: Vec<Option<usize>> = materials.iter().map(|m| m.as_ref().map(|_| 1)).collect();
4033 let sizes: Vec<usize> = beams.iter().map(|b| b.coefficients().len()).collect();
4034 let layout = Arc::new(Layout::new(
4035 roles(&sizes, &shape),
4036 &[1.0; 3],
4037 &[[false; 2]; 3],
4038 ));
4039 let truth: Vec<f64> = layout
4040 .roles
4041 .iter()
4042 .map(|role| match *role {
4043 Role::Beam {
4044 region,
4045 coefficient,
4046 } => beams[region].coefficients()[coefficient],
4047 Role::Density { region, .. } => [5.0e-4, 4.0e-4, 0.0][region],
4048 Role::Temperature { .. } => 300.0,
4049 Role::ThicknessVariance { .. } => 0.0,
4050 Role::Normalization => 0.93,
4051 Role::Background { region, term } => [
4052 [0.05, 0.0, 0.0, 0.0, 0.0],
4053 [0.05, 0.0, 0.0, 0.0, 0.0],
4054 [0.0; 5],
4055 ][region][term],
4056 Role::GammaTransmission { .. } => 1.0,
4057 Role::CaptureGammas { .. } => 0.0,
4058 Role::InScatter { .. } => 0.0,
4059 Role::TimingOffset { .. } => 0.0,
4060 Role::SpreadShare => 0.0,
4061 Role::TimeBackground(term) => [0.0, 0.0, 1.0][term],
4062 Role::RoomBackground => 0.0,
4063 Role::T0 => T0_US,
4064 Role::FlightPath => FLIGHT_PATH_M,
4065 Role::Pulse(n) => [0.35, 0.05, 0.0, 0.25, 0.15, 0.0, 0.0, 0.0, 1.0][n],
4066 })
4067 .collect();
4068 let model = RegionsModel::new(&grid, &beams, &isotopes, (1.2, [0.0; 2]), T0_US, &layout);
4069 let bins = model.evaluate(&truth).expect("counts").len() / 6;
4070 let live: Vec<f64> = (0..6 * bins)
4071 .map(|k| 0.9 + 0.1 * (0.3 * k as f64).sin())
4072 .collect();
4073 let counts: Vec<f64> = Recorded { model, live: &live }
4074 .evaluate(&truth)
4075 .expect("counts")
4076 .iter()
4077 .enumerate()
4078 .map(|(k, mu)| (mu + 3.0 * mu.sqrt() * (1.7 * k as f64).sin()).round())
4079 .collect();
4080 let measurement = Measurement {
4081 time_edges_us: EDGES_US.map(f64::from).collect(),
4082 charge_ratio: 1.2,
4083 normalization: Value::Fitted(1.0),
4084 time_background: NO_TIME_BACKGROUND,
4085 room_background: None,
4086 spread: None,
4087 regions: materials
4088 .into_iter()
4089 .enumerate()
4090 .map(|(r, material)| Region {
4091 open_counts: counts[2 * r * bins..(2 * r + 1) * bins].to_vec(),
4092 sample_counts: counts[(2 * r + 1) * bins..(2 * r + 2) * bins].to_vec(),
4093 open_live: Some(live[2 * r * bins..(2 * r + 1) * bins].to_vec()),
4094 sample_live: Some(live[(2 * r + 1) * bins..(2 * r + 2) * bins].to_vec()),
4095 background: match material {
4096 Some(_) => [
4097 Value::Fitted(0.02),
4098 Value::Known(0.0),
4099 Value::Known(0.0),
4100 Value::Known(0.0),
4101 Value::Known(0.0),
4102 ],
4103 None => [Value::Known(0.0); 5],
4104 },
4105 pixels: 1.0,
4106 gamma_transmission: Value::Known(1.0),
4107 capture_gammas: Value::Known(0.0),
4108 in_scatter: Value::Known(0.0),
4109 timing_offset_us: Value::Known(0.0),
4110 material: material.map(|data| Material {
4111 covariances: Vec::new(),
4112 library: EndfLibrary::EndfB8_0,
4113 isotopes: vec![(data, Value::Fitted(3.0e-4))],
4114 temperature_k: Value::Fitted(350.0),
4115 thickness_variance: Value::Known(0.0),
4116 lattice: None,
4117 }),
4118 })
4119 .collect(),
4120 };
4121 let calibration = Calibration {
4122 t0_us: Value::Fitted(T0_US),
4123 flight_path_m: Value::Fitted(FLIGHT_PATH_M),
4124 pulse: Pulse {
4125 alpha: [Value::Known(0.35), Value::Known(0.05)],
4126 beta: [Value::Known(0.0), Value::Known(0.25)],
4127 r: Value::Known(0.15),
4128 fwhm_squared_us2: Value::Known(0.0),
4129 r_high: None,
4130 detector_tail: None,
4131 energy_span_ev: (1.0, 200.0),
4132 n_tau: 256,
4133 line_span_ev: None,
4134 prior: None,
4135 },
4136 };
4137 let open = starting_bypass(&measurement).expect("numbers");
4138 let (fit, answer) = fit_counts_answer(
4139 &measurement,
4140 &calibration,
4141 BeamStarts::OpenFitsAt(open, None, vec![0.0; measurement.regions.len()]),
4142 None,
4143 )
4144 .expect("fit");
4145 assert!(fit.converged);
4146 let covariance = fit.covariance.as_ref().expect("covariance");
4147 let fitted: Vec<Role> = quantities(&measurement, &calibration)
4148 .filter(|(_, value)| !matches!(value, Value::Known(_)))
4149 .map(|(role, _)| role)
4150 .collect();
4151 let at = |role: Role| fitted.iter().position(|&r| r == role).expect("fitted");
4152 let close = |actual: f64, a: Role, b: Role| {
4153 let (i, j) = (at(a), at(b));
4154 let expected = covariance.get(i, j);
4155 let scale = (covariance.get(i, i) * covariance.get(j, j)).sqrt();
4156 assert!(
4157 (actual - expected).abs() <= 1e-8 * scale,
4158 "{a:?} {b:?}: {actual} vs {expected}"
4159 );
4160 };
4161 let shared = [Role::Normalization, Role::T0, Role::FlightPath];
4162 let terms: Vec<RegionTerms> = (0..3)
4163 .map(|r| answer.region_terms(r, &shared).expect("terms"))
4164 .collect();
4165 let k = shared.len();
4166 let joint = Mat::from_fn(k, k, |i, j| {
4167 terms.iter().map(|t| t.information.get(i, j)).sum::<f64>()
4168 })
4169 .llt(Side::Lower)
4170 .expect("positive definite")
4171 .inverse();
4172 for i in 0..k {
4173 for j in 0..k {
4174 close(joint[(i, j)], shared[i], shared[j]);
4175 }
4176 }
4177 let carried = |t: &RegionTerms, a: usize, j: usize| {
4178 (0..k)
4179 .map(|i| t.sensitivity.get(a, i) * joint[(i, j)])
4180 .sum::<f64>()
4181 };
4182 for (p, tp) in terms.iter().enumerate() {
4183 for (a, &ra) in tp.quantities.iter().enumerate() {
4184 for (j, &rs) in shared.iter().enumerate() {
4185 close(carried(tp, a, j), ra, rs);
4186 }
4187 for (q, tq) in terms.iter().enumerate() {
4188 for (b, &rb) in tq.quantities.iter().enumerate() {
4189 let common: f64 = (0..k)
4190 .map(|j| carried(tp, a, j) * tq.sensitivity.get(b, j))
4191 .sum();
4192 let own = if p == q { tp.covariance.get(a, b) } else { 0.0 };
4193 close(own + common, ra, rb);
4194 }
4195 }
4196 }
4197 }
4198 let leverage: f64 = (0..k)
4199 .flat_map(|i| (0..k).map(move |j| (i, j)))
4200 .map(|(i, j)| terms[0].counted_information.get(i, j) * joint[(i, j)])
4201 .sum();
4202 let value = |role: Role| {
4203 answer.params[answer
4204 .layout
4205 .roles
4206 .iter()
4207 .position(|&r| r == role)
4208 .expect("a shared role")]
4209 };
4210 let alone = Measurement {
4211 normalization: Value::Known(value(Role::Normalization)),
4212 regions: vec![measurement.regions[0].clone()],
4213 ..measurement.clone()
4214 };
4215 let held = Calibration {
4216 t0_us: Value::Known(value(Role::T0)),
4217 flight_path_m: Value::Known(value(Role::FlightPath)),
4218 ..calibration.clone()
4219 };
4220 let joined = &fit.regions[0];
4221 let starts = Starts {
4222 excluded: vec![joined.excluded_bins.clone()],
4223 beams: vec![BeamStart {
4224 beam: joined.beam.clone(),
4225 overdispersion: joined.overdispersion[0],
4226 at_limit: joined.beam_at_limit,
4227 }],
4228 origin_us: answer.beam_origin_us,
4229 bound: BOUND,
4230 shared_leverage: Some(vec![leverage]),
4231 };
4232 let (region, _) =
4233 fit_counts_answer(&alone, &held, BeamStarts::Given(&starts), None).expect("fit");
4234 let [Some(phi), Some(expected)] =
4235 [®ion.regions[0], joined].map(|fit| fit.overdispersion[1])
4236 else {
4237 panic!("a sample overdispersion");
4238 };
4239 assert!((phi / expected - 1.0).abs() <= 1e-4, "{phi} vs {expected}");
4240 let gradient: Vec<f64> = (0..k)
4241 .map(|i| terms.iter().map(|t| t.gradient[i]).sum())
4242 .collect();
4243 let decrement: f64 = (0..k)
4244 .flat_map(|i| (0..k).map(move |j| (i, j)))
4245 .map(|(i, j)| 0.5 * gradient[i] * joint[(i, j)] * gradient[j])
4246 .sum();
4247 assert!(decrement <= 1e-6, "{decrement}");
4248 let mut nudged = answer;
4249 let roles = nudged.layout.roles.clone();
4250 for (i, &role) in roles.iter().enumerate() {
4251 if matches!(
4252 role,
4253 Role::Density { region: 0, .. } | Role::Temperature { region: 0 }
4254 ) {
4255 nudged.params[i] += covariance.get(at(role), at(role)).sqrt();
4256 }
4257 }
4258 let nudged_terms: Vec<RegionTerms> = (0..3)
4259 .map(|r| nudged.region_terms(r, &shared).expect("terms"))
4260 .collect();
4261 let gradient: Vec<f64> = (0..k)
4262 .map(|i| nudged_terms.iter().map(|t| t.gradient[i]).sum())
4263 .collect();
4264 let decrement: f64 = (0..k)
4265 .flat_map(|i| (0..k).map(move |j| (i, j)))
4266 .map(|(i, j)| 0.5 * gradient[i] * joint[(i, j)] * gradient[j])
4267 .sum();
4268 assert!(decrement <= 1e-2, "{decrement}");
4269
4270 let drifting = Measurement {
4271 time_background: [Value::Fitted(50.0), Value::Known(0.0), Value::Known(1.0)],
4272 room_background: None,
4273 spread: None,
4274 ..measurement
4275 };
4276 for (most_passes, settled) in [(1, false), (2, true)] {
4277 let fit = fit_settled(&drifting, &calibration, most_passes).expect("fit");
4278 assert_eq!(fit.converged, settled);
4279 assert_eq!(fit.covariance.is_some(), settled);
4280 assert_eq!(fit.unbounded.is_some(), settled);
4281 assert_eq!(fit.measured_pulls.is_some(), settled);
4282 }
4283 }
4284
4285 #[test]
4286 fn a_map_under_a_pulse_calibration_is_the_joint_fit() {
4287 let grid = Arc::new(grid(None).halved().expect("grid").halved().expect("grid"));
4288 let (_, u_hi) = grid.range_us();
4289 let beam = BeamSpline::constant(347.0, u_hi, 1.0e4).refined();
4290 let beams = [beam.clone(), beam];
4291 let hafnium = synthetic_isotope(72, 180, 20.0, 0.01, 0.06);
4292 let isotopes: Vec<Arc<[ResonanceData]>> = vec![Arc::new([hafnium.clone()]), Arc::new([])];
4293 let sizes: Vec<usize> = beams.iter().map(|b| b.coefficients().len()).collect();
4294 let layout = Arc::new(Layout::new(
4295 roles(&sizes, &[Some(1), None]),
4296 &[1.0; 2],
4297 &[[false; 2]; 2],
4298 ));
4299 let truth: Vec<f64> = layout
4300 .roles
4301 .iter()
4302 .map(|role| match *role {
4303 Role::Beam {
4304 region,
4305 coefficient,
4306 } => beams[region].coefficients()[coefficient],
4307 Role::Density { .. } => 5.0e-4,
4308 Role::Temperature { .. } => 300.0,
4309 Role::ThicknessVariance { .. } => 0.0,
4310 Role::Normalization => 0.93,
4311 Role::Background { region, term } => {
4312 [[0.05, 0.0, 0.0, 0.0, 0.0], [0.0; 5]][region][term]
4313 }
4314 Role::GammaTransmission { .. } => 1.0,
4315 Role::CaptureGammas { .. } => 0.0,
4316 Role::InScatter { .. } => 0.0,
4317 Role::TimingOffset { .. } => 0.0,
4318 Role::SpreadShare => 0.0,
4319 Role::TimeBackground(term) => [0.0, 0.0, 1.0][term],
4320 Role::RoomBackground => 0.0,
4321 Role::T0 => T0_US,
4322 Role::FlightPath => FLIGHT_PATH_M,
4323 Role::Pulse(n) => [0.35, 0.05, 0.0, 0.25, 0.15, 0.0, 0.0, 0.0, 1.0][n],
4324 })
4325 .collect();
4326 let counts: Vec<f64> =
4327 RegionsModel::new(&grid, &beams, &isotopes, (1.2, [0.0; 2]), T0_US, &layout)
4328 .evaluate(&truth)
4329 .expect("counts")
4330 .iter()
4331 .map(|c| c.round())
4332 .collect();
4333 let bins = counts.len() / 4;
4334 let run =
4335 |r: usize, run: usize| counts[(2 * r + run) * bins..(2 * r + run + 1) * bins].to_vec();
4336 let cube = |which: usize| {
4337 ndarray::Array3::from_shape_fn((bins, 2, 1), |(k, y, _)| run(y, which)[k])
4338 };
4339 let (open, sample) = (cube(0), cube(1));
4340 let sample_pixel = ndarray::Array2::from_shape_fn((2, 1), |(y, _)| y == 0);
4341 let empty_pixel = ndarray::Array2::from_shape_fn((2, 1), |(y, _)| y == 1);
4342 let none = ndarray::Array2::from_elem((2, 1), false);
4343 let material = Material {
4344 covariances: Vec::new(),
4345 library: EndfLibrary::EndfB8_0,
4346 isotopes: vec![(hafnium, Value::Fitted(3.0e-4))],
4347 temperature_k: Value::Fitted(350.0),
4348 thickness_variance: Value::Known(0.0),
4349 lattice: None,
4350 };
4351 let background = [
4352 Value::Fitted(0.02),
4353 Value::Known(0.0),
4354 Value::Known(0.0),
4355 Value::Known(0.0),
4356 Value::Known(0.0),
4357 ];
4358 let map = crate::counts_map::MapMeasurement {
4359 time_edges_us: EDGES_US.map(f64::from).collect(),
4360 charge_ratio: 1.2,
4361 normalization: Value::Fitted(1.0),
4362 time_background: NO_TIME_BACKGROUND,
4363 room_background: None,
4364 open_counts: open.view(),
4365 sample_counts: sample.view(),
4366 open_live: None,
4367 sample_live: None,
4368 excluded: none.view(),
4369 sample: sample_pixel.view(),
4370 empty: empty_pixel.view(),
4371 binning: 1,
4372 material: material.clone(),
4373 background,
4374 gamma_transmission: Value::Known(1.0),
4375 capture_gammas: Value::Known(0.0),
4376 in_scatter: Value::Known(0.0),
4377 empty_background: [Value::Known(0.0); 5],
4378 };
4379 let calibration = Calibration {
4380 t0_us: Value::Measured {
4381 value: T0_US,
4382 sd: 0.01,
4383 },
4384 flight_path_m: Value::Fitted(FLIGHT_PATH_M),
4385 pulse: Pulse {
4386 alpha: [Value::Fitted(0.35), Value::Fitted(0.05)],
4387 beta: [Value::Known(0.0), Value::Known(0.25)],
4388 r: Value::Known(0.15),
4389 fwhm_squared_us2: Value::Known(0.0),
4390 r_high: None,
4391 detector_tail: None,
4392 energy_span_ev: (1.0, 200.0),
4393 n_tau: 256,
4394 line_span_ev: None,
4395 prior: Some(crate::pulse_calibration::PulsePrior {
4396 numbers: vec![0, 1],
4397 mean: vec![0.36, 0.05],
4398 covariance: FlatMatrix {
4399 data: vec![1.0e-4, -2.0e-5, -2.0e-5, 1.0e-5],
4400 nrows: 2,
4401 ncols: 2,
4402 },
4403 }),
4404 },
4405 };
4406 let result = crate::counts_map::fit_map(&map, &calibration).expect("map");
4407 assert!(result.converged);
4408 let region = |r: usize, material: Option<Material>, background| Region {
4409 open_counts: run(r, 0),
4410 sample_counts: run(r, 1),
4411 open_live: None,
4412 sample_live: None,
4413 background,
4414 pixels: 1.0,
4415 gamma_transmission: Value::Known(1.0),
4416 capture_gammas: Value::Known(0.0),
4417 in_scatter: Value::Known(0.0),
4418 timing_offset_us: Value::Known(0.0),
4419 material,
4420 };
4421 let measurement = Measurement {
4422 time_edges_us: map.time_edges_us.clone(),
4423 charge_ratio: 1.2,
4424 normalization: map.normalization,
4425 time_background: NO_TIME_BACKGROUND,
4426 room_background: None,
4427 spread: None,
4428 regions: vec![
4429 region(0, Some(material), background),
4430 region(1, None, [Value::Known(0.0); 5]),
4431 ],
4432 };
4433 let joint = fit_counts(&measurement, &calibration).expect("joint fit");
4434 assert!(joint.converged);
4435 let mut carrying = measurement.clone();
4436 carrying.regions[0]
4437 .material
4438 .as_mut()
4439 .expect("a material")
4440 .covariances = vec![(
4441 0,
4442 ResonanceCovariance {
4443 resonances: vec![[0, 0, 0]],
4444 covariance: [1e-6, 0.0, 0.0, 0.0, 1e-10, 0.0, 0.0, 0.0, 1e-8].into(),
4445 },
4446 )];
4447 assert!(matches!(
4448 fit_counts(&carrying, &calibration),
4449 Err(PipelineError::InvalidParameter(_))
4450 ));
4451 let covariance = joint.covariance.as_ref().expect("covariance");
4452 let fitted: Vec<Role> = quantities(&measurement, &calibration)
4453 .filter(|(_, value)| !matches!(value, Value::Known(_)))
4454 .map(|(role, _)| role)
4455 .collect();
4456 let at = |role: Role| fitted.iter().position(|&r| r == role).expect("fitted");
4457 let close = |actual: f64, a: Role, b: Role| {
4458 let (i, j) = (at(a), at(b));
4459 let scale = (covariance.get(i, i) * covariance.get(j, j)).sqrt();
4460 assert!(
4461 (actual - covariance.get(i, j)).abs() <= 1e-3 * scale,
4462 "{a:?} {b:?}: {actual} vs {}",
4463 covariance.get(i, j)
4464 );
4465 };
4466 let shared = [
4467 Role::Normalization,
4468 Role::T0,
4469 Role::FlightPath,
4470 Role::Pulse(0),
4471 Role::Pulse(1),
4472 ];
4473 let map_shared = result.shared_covariance.as_ref().expect("covariance");
4474 for (s, &a) in shared.iter().enumerate() {
4475 for (t, &b) in shared.iter().enumerate() {
4476 close(map_shared.get(s, t), a, b);
4477 }
4478 }
4479 let own = [
4480 Role::Density {
4481 region: 0,
4482 isotope: 0,
4483 },
4484 Role::Temperature { region: 0 },
4485 Role::Background { region: 0, term: 0 },
4486 ];
4487 let patch = result.covariances.covariance[[0, 0]]
4488 .as_ref()
4489 .expect("covariance");
4490 for (s, &a) in own.iter().enumerate() {
4491 for (t, &b) in own.iter().enumerate() {
4492 close(patch.get(s, t), a, b);
4493 }
4494 }
4495 let gaps = [
4496 (result.alpha[0] - joint.alpha[0])
4497 / covariance
4498 .get(at(Role::Pulse(0)), at(Role::Pulse(0)))
4499 .sqrt(),
4500 (result.t0_us - joint.t0_us) / covariance.get(at(Role::T0), at(Role::T0)).sqrt(),
4501 (result.densities[0][[0, 0]] - joint.regions[0].densities[0])
4502 / covariance.get(at(own[0]), at(own[0])).sqrt(),
4503 ];
4504 assert!(gaps.iter().all(|g| g.abs() <= 0.01), "{gaps:?}");
4505 let [Some(map_test), Some(joint_test)] =
4506 [result.pulse_consistency, joint.pulse_consistency]
4507 else {
4508 panic!("a pulse consistency");
4509 };
4510 assert_eq!(map_test.dof, joint_test.dof);
4511 assert!(
4512 (map_test.q - joint_test.q).abs() <= 1e-3 * joint_test.q,
4513 "{map_test:?} vs {joint_test:?}"
4514 );
4515 }
4516
4517 #[test]
4518 fn the_background_lags_sammy_s_by_the_pulse_s_mean_delay() {
4519 let grid = grid(None);
4520 let (_, u_hi) = grid.range_us();
4521 let first_edge_us = f64::from(*EDGES_US.start());
4522 let beam = BeamSpline::constant(first_edge_us - T0_US, u_hi, 1.0e4);
4523 let isotopes: Arc<[ResonanceData]> = Arc::new([
4524 synthetic_isotope(72, 180, 20.0, 0.01, 0.06),
4525 synthetic_isotope(74, 182, 20.3, 0.01, 0.06),
4526 ]);
4527 let (charge_ratio, normalization) = (1.2, 0.93);
4528 let layout = Arc::new(Layout::new(
4529 roles(&[beam.coefficients().len()], &[Some(2)]),
4530 &[1.0],
4531 &[[false; 2]],
4532 ));
4533 let model = RegionsModel::new(
4534 &grid,
4535 std::slice::from_ref(&beam),
4536 &[isotopes],
4537 (charge_ratio, [0.0; 2]),
4538 T0_US,
4539 &layout,
4540 );
4541 let counts = |background: [f64; 5]| {
4542 let params: Vec<f64> = beam
4543 .coefficients()
4544 .iter()
4545 .copied()
4546 .chain([3.0e-4, 5.0e-4, 300.0, 0.0, normalization])
4547 .chain(background)
4548 .chain([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0])
4549 .chain([T0_US, FLIGHT_PATH_M])
4550 .chain([0.35, 0.05, 0.0, 0.25, 0.15, 0.0, 0.0, 0.0, 1.0])
4551 .collect();
4552 model.evaluate(¶ms).expect("counts")
4553 };
4554 let [b0, b1, b2] = [0.05, 0.5, -0.01];
4555 let (with, without) = (counts([b0, b1, b2, 0.0, 0.0]), counts([0.0; 5]));
4556 let bins = with.len() / 2;
4557 let clock = TOF_FACTOR * FLIGHT_PATH_M;
4558 for k in 0..bins {
4559 let u = first_edge_us + 0.5 + k as f64 - T0_US;
4560 let root_e = clock / u;
4561 let b = b0 + b1 / root_e + b2 * root_e;
4562 let slope = b1 / clock - b2 * clock / (u * u);
4563 let curvature = 2.0 * b2 * clock / u.powi(3);
4564 let energy = root_e * root_e;
4565 let (alpha, beta, r) = (ALPHA.eval(energy), BETA.eval(energy), R.eval(energy));
4566 let mean = 3.0 / alpha + r / beta;
4567 let square = mean * mean + 3.0 / (alpha * alpha) + r * (2.0 - r) / (beta * beta);
4568 let lagged = (with[bins + k] - without[bins + k]) / (charge_ratio * normalization);
4569 let residual = lagged / with[k] - b + slope * mean;
4570 let tolerance = 0.5 * slope.abs() + 0.5 * curvature.abs() * square;
4571 assert!(
4572 residual.abs() <= tolerance,
4573 "{k}: {residual} vs {tolerance}"
4574 );
4575 }
4576 }
4577
4578 #[test]
4579 fn an_isotope_known_to_be_absent_leaves_no_window() {
4580 let edges: Vec<f64> = (380..=530).map(f64::from).collect();
4581 let numbers = [3.0, 25.0, 0.5, 1.0, 0.08, 0.0, 0.2, 0.0, 0.0, 0.0, 1.0];
4582 let bins = |density| {
4583 let material = Material {
4584 covariances: Vec::new(),
4585 isotopes: vec![(synthetic_isotope(73, 181, 13.9, 0.01, 0.06), density)],
4586 temperature_k: Value::Known(300.0),
4587 library: EndfLibrary::EndfB8_1,
4588 thickness_variance: Value::Known(0.0),
4589 lattice: None,
4590 };
4591 excluded_bins(&edges, Some(&material), &numbers, None)
4592 };
4593 assert!(bins(Value::Known(0.0)).is_empty());
4594 assert!(!bins(Value::Known(1e-6)).is_empty());
4595 }
4596}