1use std::error::Error;
4use std::fmt::{Display, Formatter};
5
6use phasesmith_core::{
7 Accumulation, GridView, TofBankGeometry, TofError, TofInstrument,
8 accumulate_tof_batch_with_context,
9};
10use phasesmith_crystallography::{
11 CELL_PARAMETER_COUNT, CellError, IntegratedIntensityCorrection,
12 IntegratedIntensityCorrectionError, IntegratedIntensityCorrectionModel,
13 PreparedNeutronScattering, ScatteringBatch, ScatteringError, SpaceGroup,
14 StructureFactorBatchError, StructureFactorBatchView, StructureFactorValues, UnitCell,
15 calculate_structure_factor_dense_with_context,
16 calculate_structure_factor_intensity_vjp_with_context,
17 calculate_structure_factor_jvp_with_context, calculate_structure_factor_values_with_context,
18};
19use phasesmith_execution::ExecutionContext;
20
21#[derive(Clone, Copy, Debug)]
23pub struct StructuralTofInputView<'a> {
24 pub tof_us: &'a [f64],
26 pub hkl: &'a [[i32; 3]],
28 pub multiplicity: &'a [usize],
30 pub fractional_xyz: &'a [[f64; 3]],
32 pub occupancy: &'a [f64],
34 pub u_iso_angstrom2: &'a [f64],
36 pub anisotropic_mask: &'a [bool],
38 pub u_aniso_cif_angstrom2: &'a [[f64; 6]],
40 pub scattering_species: &'a [&'a str],
42 pub scale: f64,
44 pub coordinate_tolerance: f64,
46 pub correction_model: IntegratedIntensityCorrectionModel,
48 pub bank_geometry: TofBankGeometry,
50 pub instrument: TofInstrument,
52 pub support_fwhm: f64,
54 pub tail_log: f64,
56}
57
58#[derive(Clone, Debug, PartialEq)]
60pub struct StructuralTofResult {
61 pub structure_factors: StructureFactorValues,
63 pub d_spacing_angstrom: Vec<f64>,
65 pub accumulation: Accumulation,
67}
68
69#[derive(Clone, Debug, PartialEq)]
71pub struct StructuralTofDenseResult {
72 pub result: StructuralTofResult,
74 pub d_y: Vec<f64>,
76 pub parameter_count: usize,
78}
79
80#[derive(Clone, Debug, PartialEq)]
82pub struct StructuralTofJvpResult {
83 pub result: StructuralTofResult,
85 pub d_y: Vec<f64>,
87 pub d_integrated_intensity: Vec<f64>,
89 pub d_spacing_angstrom: Vec<f64>,
91}
92
93#[derive(Clone, Debug, PartialEq)]
95pub struct StructuralTofVjpResult {
96 pub result: StructuralTofResult,
98 pub gradient: Vec<f64>,
100}
101
102#[derive(Debug)]
104pub enum StructuralTofError {
105 Cell(CellError),
107 SpeciesLengthMismatch,
109 Scattering(ScatteringError),
111 Correction(IntegratedIntensityCorrectionError),
113 IncompatibleCorrectionModel,
115 StructureFactor(StructureFactorBatchError),
117 Tof(TofError),
119 AllocationOverflow,
121 PatternWeightLengthMismatch,
123 NonFinitePatternWeight,
125}
126
127impl Display for StructuralTofError {
128 fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
129 match self {
130 Self::Cell(error) => Display::fmt(error, formatter),
131 Self::SpeciesLengthMismatch => formatter
132 .write_str("neutron scattering species must contain one key per asymmetric site"),
133 Self::Scattering(error) => Display::fmt(error, formatter),
134 Self::Correction(error) => Display::fmt(error, formatter),
135 Self::IncompatibleCorrectionModel => formatter.write_str(
136 "structural TOF requires neutral correction or a TOF neutron Lorentz angle exactly matching the bank geometry",
137 ),
138 Self::StructureFactor(error) => Display::fmt(error, formatter),
139 Self::Tof(error) => Display::fmt(error, formatter),
140 Self::AllocationOverflow => formatter.write_str("structural TOF output allocation overflow"),
141 Self::PatternWeightLengthMismatch => {
142 formatter.write_str("pattern reverse weights must match the TOF sample count")
143 }
144 Self::NonFinitePatternWeight => formatter.write_str("pattern reverse weights must be finite"),
145 }
146 }
147}
148
149impl Error for StructuralTofError {}
150
151struct PreparedTofNumerics {
152 scattering: ScatteringBatch,
153 correction: IntegratedIntensityCorrection,
154 d_spacing: Vec<f64>,
155 d_spacing_d_cell: Vec<[f64; CELL_PARAMETER_COUNT]>,
156}
157
158impl PreparedTofNumerics {
159 fn structure_batch<'a>(
160 &'a self,
161 input: &StructuralTofInputView<'a>,
162 ) -> StructureFactorBatchView<'a> {
163 StructureFactorBatchView {
164 hkl: input.hkl,
165 multiplicity: input.multiplicity,
166 fractional_xyz: input.fractional_xyz,
167 occupancy: input.occupancy,
168 u_iso_angstrom2: input.u_iso_angstrom2,
169 anisotropic_mask: input.anisotropic_mask,
170 u_aniso_cif_angstrom2: input.u_aniso_cif_angstrom2,
171 scattering_real: &self.scattering.real,
172 scattering_imag: &self.scattering.imag,
173 d_scattering_real_d_s: &self.scattering.d_real_d_s,
174 d_scattering_imag_d_s: &self.scattering.d_imag_d_s,
175 correction: &self.correction.values,
176 d_correction_d_q_squared: &self.correction.d_values_d_q_squared,
177 scale: input.scale,
178 coordinate_tolerance: input.coordinate_tolerance,
179 }
180 }
181}
182
183pub fn calculate_structural_tof_pattern(
190 cell: UnitCell,
191 space_group: &SpaceGroup,
192 input: &StructuralTofInputView<'_>,
193) -> Result<StructuralTofResult, StructuralTofError> {
194 calculate_structural_tof_pattern_with_context(
195 cell,
196 space_group,
197 input,
198 &ExecutionContext::serial(),
199 )
200}
201
202pub fn calculate_structural_tof_pattern_with_context(
208 cell: UnitCell,
209 space_group: &SpaceGroup,
210 input: &StructuralTofInputView<'_>,
211 execution: &ExecutionContext,
212) -> Result<StructuralTofResult, StructuralTofError> {
213 let prepared = prepare(cell, input)?;
214 let values = calculate_structure_factor_values_with_context(
215 cell,
216 space_group,
217 prepared.structure_batch(input),
218 execution,
219 )
220 .map_err(StructuralTofError::StructureFactor)?;
221 assemble_result(input, prepared, values, execution)
222}
223
224pub fn calculate_structural_tof_pattern_dense(
230 cell: UnitCell,
231 space_group: &SpaceGroup,
232 input: &StructuralTofInputView<'_>,
233) -> Result<StructuralTofDenseResult, StructuralTofError> {
234 calculate_structural_tof_pattern_dense_with_context(
235 cell,
236 space_group,
237 input,
238 &ExecutionContext::serial(),
239 )
240}
241
242pub fn calculate_structural_tof_pattern_dense_with_context(
248 cell: UnitCell,
249 space_group: &SpaceGroup,
250 input: &StructuralTofInputView<'_>,
251 execution: &ExecutionContext,
252) -> Result<StructuralTofDenseResult, StructuralTofError> {
253 let prepared = prepare(cell, input)?;
254 let structural = calculate_structure_factor_dense_with_context(
255 cell,
256 space_group,
257 prepared.structure_batch(input),
258 execution,
259 )
260 .map_err(StructuralTofError::StructureFactor)?;
261 let parameter_count = structural.layout.parameter_count();
262 let sample_count = input.tof_us.len();
263 let element_count = parameter_count
264 .checked_mul(sample_count)
265 .ok_or(StructuralTofError::AllocationOverflow)?;
266 let accumulation = accumulate(
267 input,
268 &prepared.d_spacing,
269 &structural.values.intensity,
270 execution,
271 )?;
272 let mut d_y = vec![0.0; element_count];
273 let reflection_count = input.hkl.len();
274 let local = &accumulation.derivatives.local;
275 for reflection in 0..reflection_count {
276 let begin = local.offsets[reflection];
277 let end = local.offsets[reflection + 1];
278 for active in begin..end {
279 let sample = local.starts[reflection] + active - begin;
280 let local_base = 2 * active;
281 for parameter in 0..parameter_count {
282 let d_spacing = if parameter < CELL_PARAMETER_COUNT {
283 prepared.d_spacing_d_cell[reflection][parameter]
284 } else {
285 0.0
286 };
287 d_y[parameter * sample_count + sample] += local.values[local_base]
288 * structural.d_intensity[parameter * reflection_count + reflection]
289 + local.values[local_base + 1] * d_spacing;
290 }
291 }
292 }
293 Ok(StructuralTofDenseResult {
294 result: StructuralTofResult {
295 structure_factors: structural.values,
296 d_spacing_angstrom: prepared.d_spacing,
297 accumulation,
298 },
299 d_y,
300 parameter_count,
301 })
302}
303
304pub fn calculate_structural_tof_pattern_jvp(
310 cell: UnitCell,
311 space_group: &SpaceGroup,
312 input: &StructuralTofInputView<'_>,
313 tangent: &[f64],
314) -> Result<StructuralTofJvpResult, StructuralTofError> {
315 calculate_structural_tof_pattern_jvp_with_context(
316 cell,
317 space_group,
318 input,
319 tangent,
320 &ExecutionContext::serial(),
321 )
322}
323
324pub fn calculate_structural_tof_pattern_jvp_with_context(
330 cell: UnitCell,
331 space_group: &SpaceGroup,
332 input: &StructuralTofInputView<'_>,
333 tangent: &[f64],
334 execution: &ExecutionContext,
335) -> Result<StructuralTofJvpResult, StructuralTofError> {
336 let prepared = prepare(cell, input)?;
337 let structural = calculate_structure_factor_jvp_with_context(
338 cell,
339 space_group,
340 prepared.structure_batch(input),
341 tangent,
342 execution,
343 )
344 .map_err(StructuralTofError::StructureFactor)?;
345 let d_spacing_angstrom = prepared
346 .d_spacing_d_cell
347 .iter()
348 .map(|derivatives| {
349 derivatives
350 .iter()
351 .zip(&tangent[..CELL_PARAMETER_COUNT])
352 .map(|(derivative, direction)| derivative * direction)
353 .sum()
354 })
355 .collect::<Vec<_>>();
356 let accumulation = accumulate(
357 input,
358 &prepared.d_spacing,
359 &structural.values.intensity,
360 execution,
361 )?;
362 let d_y = chain_jvp(&accumulation, &structural.d_intensity, &d_spacing_angstrom);
363 Ok(StructuralTofJvpResult {
364 result: StructuralTofResult {
365 structure_factors: structural.values,
366 d_spacing_angstrom: prepared.d_spacing,
367 accumulation,
368 },
369 d_y,
370 d_integrated_intensity: structural.d_intensity,
371 d_spacing_angstrom,
372 })
373}
374
375pub fn calculate_structural_tof_pattern_vjp(
381 cell: UnitCell,
382 space_group: &SpaceGroup,
383 input: &StructuralTofInputView<'_>,
384 sample_weights: &[f64],
385) -> Result<StructuralTofVjpResult, StructuralTofError> {
386 calculate_structural_tof_pattern_vjp_with_context(
387 cell,
388 space_group,
389 input,
390 sample_weights,
391 &ExecutionContext::serial(),
392 )
393}
394
395pub fn calculate_structural_tof_pattern_vjp_with_context(
401 cell: UnitCell,
402 space_group: &SpaceGroup,
403 input: &StructuralTofInputView<'_>,
404 sample_weights: &[f64],
405 execution: &ExecutionContext,
406) -> Result<StructuralTofVjpResult, StructuralTofError> {
407 if sample_weights.len() != input.tof_us.len() {
408 return Err(StructuralTofError::PatternWeightLengthMismatch);
409 }
410 if sample_weights.iter().any(|value| !value.is_finite()) {
411 return Err(StructuralTofError::NonFinitePatternWeight);
412 }
413 let prepared = prepare(cell, input)?;
414 let values = calculate_structure_factor_values_with_context(
415 cell,
416 space_group,
417 prepared.structure_batch(input),
418 execution,
419 )
420 .map_err(StructuralTofError::StructureFactor)?;
421 let accumulation = accumulate(input, &prepared.d_spacing, &values.intensity, execution)?;
422 let (intensity_weights, d_spacing_weights) = local_transpose(&accumulation, sample_weights);
423 let mut structural = calculate_structure_factor_intensity_vjp_with_context(
424 cell,
425 space_group,
426 prepared.structure_batch(input),
427 &intensity_weights,
428 execution,
429 )
430 .map_err(StructuralTofError::StructureFactor)?;
431 for (reflection, weight) in d_spacing_weights.into_iter().enumerate() {
432 for parameter in 0..CELL_PARAMETER_COUNT {
433 structural.gradient[parameter] +=
434 weight * prepared.d_spacing_d_cell[reflection][parameter];
435 }
436 }
437 Ok(StructuralTofVjpResult {
438 result: StructuralTofResult {
439 structure_factors: values,
440 d_spacing_angstrom: prepared.d_spacing,
441 accumulation,
442 },
443 gradient: structural.gradient,
444 })
445}
446
447fn prepare(
448 cell: UnitCell,
449 input: &StructuralTofInputView<'_>,
450) -> Result<PreparedTofNumerics, StructuralTofError> {
451 if input.scattering_species.len() != input.fractional_xyz.len() {
452 return Err(StructuralTofError::SpeciesLengthMismatch);
453 }
454 input
455 .bank_geometry
456 .validate()
457 .map_err(StructuralTofError::Tof)?;
458 validate_correction(input.correction_model, input.bank_geometry)?;
459 let geometry = cell.geometry().map_err(StructuralTofError::Cell)?;
460 let mut q_squared = Vec::with_capacity(input.hkl.len());
461 let mut d_spacing = Vec::with_capacity(input.hkl.len());
462 let mut d_spacing_d_cell = Vec::with_capacity(input.hkl.len());
463 for &hkl in input.hkl {
464 let (q, d_q) = geometry.q_squared_and_derivatives(hkl);
465 if !q.is_finite() || q <= 0.0 {
466 return Err(StructuralTofError::StructureFactor(
467 StructureFactorBatchError::ZeroReflection,
468 ));
469 }
470 let d = q.sqrt().recip();
471 q_squared.push(q);
472 d_spacing.push(d);
473 d_spacing_d_cell.push(d_q.map(|derivative| -0.5 * d.powi(3) * derivative));
474 }
475 let s = q_squared
476 .iter()
477 .map(|value| 0.5 * value.sqrt())
478 .collect::<Vec<_>>();
479 let scattering = PreparedNeutronScattering::new(input.scattering_species.iter().copied())
480 .and_then(|model| model.evaluate(&s))
481 .map_err(StructuralTofError::Scattering)?;
482 let correction = input
483 .correction_model
484 .evaluate(&q_squared)
485 .map_err(StructuralTofError::Correction)?;
486 Ok(PreparedTofNumerics {
487 scattering,
488 correction,
489 d_spacing,
490 d_spacing_d_cell,
491 })
492}
493
494fn validate_correction(
495 correction: IntegratedIntensityCorrectionModel,
496 geometry: TofBankGeometry,
497) -> Result<(), StructuralTofError> {
498 match correction {
499 IntegratedIntensityCorrectionModel::Neutral => Ok(()),
500 IntegratedIntensityCorrectionModel::TimeOfFlightNeutronLorentz { two_theta_deg }
501 if two_theta_deg.to_bits() == geometry.two_theta_deg.to_bits() =>
502 {
503 Ok(())
504 }
505 _ => Err(StructuralTofError::IncompatibleCorrectionModel),
506 }
507}
508
509fn accumulate(
510 input: &StructuralTofInputView<'_>,
511 d_spacing: &[f64],
512 intensity: &[f64],
513 execution: &ExecutionContext,
514) -> Result<Accumulation, StructuralTofError> {
515 let grid = GridView::new(input.tof_us)
516 .map_err(|reason| StructuralTofError::Tof(TofError::Profile { reason }))?;
517 accumulate_tof_batch_with_context(
518 grid,
519 d_spacing,
520 intensity,
521 input.instrument,
522 input.support_fwhm,
523 input.tail_log,
524 execution,
525 )
526 .map_err(StructuralTofError::Tof)
527}
528
529fn assemble_result(
530 input: &StructuralTofInputView<'_>,
531 prepared: PreparedTofNumerics,
532 values: StructureFactorValues,
533 execution: &ExecutionContext,
534) -> Result<StructuralTofResult, StructuralTofError> {
535 let accumulation = accumulate(input, &prepared.d_spacing, &values.intensity, execution)?;
536 Ok(StructuralTofResult {
537 structure_factors: values,
538 d_spacing_angstrom: prepared.d_spacing,
539 accumulation,
540 })
541}
542
543fn chain_jvp(accumulation: &Accumulation, d_intensity: &[f64], d_spacing: &[f64]) -> Vec<f64> {
544 let local = &accumulation.derivatives.local;
545 let mut result = vec![0.0; accumulation.sample_count];
546 for reflection in 0..local.peak_count() {
547 let begin = local.offsets[reflection];
548 let end = local.offsets[reflection + 1];
549 for active in begin..end {
550 let sample = local.starts[reflection] + active - begin;
551 let base = 2 * active;
552 result[sample] += local.values[base] * d_intensity[reflection]
553 + local.values[base + 1] * d_spacing[reflection];
554 }
555 }
556 result
557}
558
559fn local_transpose(accumulation: &Accumulation, weights: &[f64]) -> (Vec<f64>, Vec<f64>) {
560 let local = &accumulation.derivatives.local;
561 let mut intensity = vec![0.0; local.peak_count()];
562 let mut d_spacing = vec![0.0; local.peak_count()];
563 for reflection in 0..local.peak_count() {
564 let begin = local.offsets[reflection];
565 let end = local.offsets[reflection + 1];
566 for active in begin..end {
567 let sample = local.starts[reflection] + active - begin;
568 let base = 2 * active;
569 intensity[reflection] += weights[sample] * local.values[base];
570 d_spacing[reflection] += weights[sample] * local.values[base + 1];
571 }
572 }
573 (intensity, d_spacing)
574}
575
576#[cfg(test)]
577mod tests {
578 use super::*;
579 use phasesmith_crystallography::{Rational, SymmetryOperation};
580
581 const HKL: [[i32; 3]; 3] = [[1, 0, 1], [2, 1, 1], [1, 2, 3]];
582 const MULTIPLICITY: [usize; 3] = [2, 4, 2];
583 const XYZ: [[f64; 3]; 2] = [[0.17, 0.23, 0.31], [0.37, 0.11, 0.19]];
584 const OCCUPANCY: [f64; 2] = [0.82, 0.55];
585 const U_ISO: [f64; 2] = [0.012, 0.018];
586 const ANISOTROPIC: [bool; 2] = [false, false];
587 const U_ANISO: [[f64; 6]; 2] = [[0.0; 6]; 2];
588 const SPECIES: [&str; 2] = ["Ni", "O"];
589
590 fn cell() -> UnitCell {
591 UnitCell {
592 a_angstrom: 4.7,
593 b_angstrom: 5.1,
594 c_angstrom: 6.2,
595 alpha_deg: 82.0,
596 beta_deg: 87.0,
597 gamma_deg: 74.0,
598 }
599 }
600
601 fn group() -> SpaceGroup {
602 SpaceGroup::new(vec![
603 SymmetryOperation::new([[1, 0, 0], [0, 1, 0], [0, 0, 1]], [Rational::zero(); 3])
604 .unwrap(),
605 ])
606 .unwrap()
607 }
608
609 fn instrument() -> TofInstrument {
610 TofInstrument {
611 zero_us: 1.2,
612 difc_us_per_angstrom: 5_000.0,
613 difa_us_per_angstrom2: 0.2,
614 difb_us_angstrom: 0.0,
615 alpha_coefficient: 0.2,
616 beta0_per_us: 0.03,
617 beta1_angstrom4_per_us: 0.001,
618 betaq_angstrom2_per_us: 0.0,
619 sigma0_us2: 25.0,
620 sigma1_us2_per_angstrom2: 4.0,
621 sigma2_us2_per_angstrom4: 0.1,
622 sigmaq_us2_per_angstrom: 0.0,
623 x_us_per_angstrom: 1.0,
624 y_us_per_angstrom2: 0.1,
625 z_us: 0.5,
626 }
627 }
628
629 fn grid() -> Vec<f64> {
630 (0..2_401)
631 .map(|index| 1_000.0 + f64::from(index) * 10.0)
632 .collect()
633 }
634
635 fn input<'a>(
636 tof_us: &'a [f64],
637 xyz: &'a [[f64; 3]],
638 occupancy: &'a [f64],
639 u_iso: &'a [f64],
640 scale: f64,
641 ) -> StructuralTofInputView<'a> {
642 StructuralTofInputView {
643 tof_us,
644 hkl: &HKL,
645 multiplicity: &MULTIPLICITY,
646 fractional_xyz: xyz,
647 occupancy,
648 u_iso_angstrom2: u_iso,
649 anisotropic_mask: &ANISOTROPIC,
650 u_aniso_cif_angstrom2: &U_ANISO,
651 scattering_species: &SPECIES,
652 scale,
653 coordinate_tolerance: 1.0e-10,
654 correction_model: IntegratedIntensityCorrectionModel::TimeOfFlightNeutronLorentz {
655 two_theta_deg: 88.05,
656 },
657 bank_geometry: TofBankGeometry {
658 two_theta_deg: 88.05,
659 },
660 instrument: instrument(),
661 support_fwhm: 20.0,
662 tail_log: 20.0,
663 }
664 }
665
666 fn values(
667 selected_cell: UnitCell,
668 tof_us: &[f64],
669 xyz: &[[f64; 3]],
670 occupancy: &[f64],
671 u_iso: &[f64],
672 scale: f64,
673 ) -> Vec<f64> {
674 calculate_structural_tof_pattern(
675 selected_cell,
676 &group(),
677 &input(tof_us, xyz, occupancy, u_iso, scale),
678 )
679 .unwrap()
680 .accumulation
681 .y
682 }
683
684 #[test]
685 #[allow(clippy::too_many_lines)]
686 fn dense_structural_rows_match_centered_pattern_differences() {
687 let tof_us = grid();
688 let dense = calculate_structural_tof_pattern_dense(
689 cell(),
690 &group(),
691 &input(&tof_us, &XYZ, &OCCUPANCY, &U_ISO, 1.3),
692 )
693 .unwrap();
694 assert_eq!(dense.parameter_count, 17);
695 assert_eq!(
696 dense
697 .result
698 .accumulation
699 .derivatives
700 .global
701 .as_ref()
702 .unwrap()
703 .parameter_count,
704 15
705 );
706 let mut support_boundary = vec![false; tof_us.len()];
707 let local = &dense.result.accumulation.derivatives.local;
708 for reflection in 0..local.peak_count() {
709 let count = local.offsets[reflection + 1] - local.offsets[reflection];
710 let start = local.starts[reflection];
711 for offset in 0..3.min(count) {
712 support_boundary[start + offset] = true;
713 support_boundary[start + count - 1 - offset] = true;
714 }
715 }
716 for parameter in 0..dense.parameter_count {
717 let step = 1.0e-6;
718 let mut plus_cell = cell();
719 let mut minus_cell = cell();
720 let mut plus_xyz = XYZ;
721 let mut minus_xyz = XYZ;
722 let mut plus_occupancy = OCCUPANCY;
723 let mut minus_occupancy = OCCUPANCY;
724 let mut plus_u_iso = U_ISO;
725 let mut minus_u_iso = U_ISO;
726 let mut plus_scale = 1.3;
727 let mut minus_scale = 1.3;
728 match parameter {
729 0 => {
730 plus_cell.a_angstrom += step;
731 minus_cell.a_angstrom -= step;
732 }
733 1 => {
734 plus_cell.b_angstrom += step;
735 minus_cell.b_angstrom -= step;
736 }
737 2 => {
738 plus_cell.c_angstrom += step;
739 minus_cell.c_angstrom -= step;
740 }
741 3 => {
742 plus_cell.alpha_deg += step;
743 minus_cell.alpha_deg -= step;
744 }
745 4 => {
746 plus_cell.beta_deg += step;
747 minus_cell.beta_deg -= step;
748 }
749 5 => {
750 plus_cell.gamma_deg += step;
751 minus_cell.gamma_deg -= step;
752 }
753 6..=11 => {
754 let local = parameter - CELL_PARAMETER_COUNT;
755 plus_xyz[local / 3][local % 3] += step;
756 minus_xyz[local / 3][local % 3] -= step;
757 }
758 12..=13 => {
759 plus_occupancy[parameter - 12] += step;
760 minus_occupancy[parameter - 12] -= step;
761 }
762 14..=15 => {
763 plus_u_iso[parameter - 14] += step;
764 minus_u_iso[parameter - 14] -= step;
765 }
766 16 => {
767 plus_scale += step;
768 minus_scale -= step;
769 }
770 _ => unreachable!(),
771 }
772 let plus = values(
773 plus_cell,
774 &tof_us,
775 &plus_xyz,
776 &plus_occupancy,
777 &plus_u_iso,
778 plus_scale,
779 );
780 let minus = values(
781 minus_cell,
782 &tof_us,
783 &minus_xyz,
784 &minus_occupancy,
785 &minus_u_iso,
786 minus_scale,
787 );
788 for sample in 0..tof_us.len() {
789 if support_boundary[sample] {
790 continue;
791 }
792 let finite = (plus[sample] - minus[sample]) / (2.0 * step);
793 let analytical = dense.d_y[parameter * tof_us.len() + sample];
794 let relative_tolerance = if parameter < CELL_PARAMETER_COUNT {
795 2.0e-4
800 } else {
801 3.0e-5
802 };
803 assert!(
804 (analytical - finite).abs() <= relative_tolerance * finite.abs().max(1.0),
805 "parameter={parameter} sample={sample} analytical={analytical} finite={finite} error={}",
806 (analytical - finite).abs(),
807 );
808 }
809 }
810 }
811
812 #[test]
813 fn jvp_vjp_match_dense_and_are_adjoint_consistent() {
814 let tof_us = grid();
815 let request = input(&tof_us, &XYZ, &OCCUPANCY, &U_ISO, 1.3);
816 let dense = calculate_structural_tof_pattern_dense(cell(), &group(), &request).unwrap();
817 let tangent = (0..dense.parameter_count)
818 .map(|index| f64::from(u32::try_from(index + 1).unwrap()) * 1.0e-5)
819 .collect::<Vec<_>>();
820 let jvp =
821 calculate_structural_tof_pattern_jvp(cell(), &group(), &request, &tangent).unwrap();
822 for sample in 0..tof_us.len() {
823 let expected = (0..dense.parameter_count)
824 .map(|parameter| dense.d_y[parameter * tof_us.len() + sample] * tangent[parameter])
825 .sum::<f64>();
826 assert!((jvp.d_y[sample] - expected).abs() < 2.0e-12 * expected.abs().max(1.0));
827 }
828 let weights = tof_us
829 .iter()
830 .map(|value| (value * 1.0e-3).sin())
831 .collect::<Vec<_>>();
832 let vjp =
833 calculate_structural_tof_pattern_vjp(cell(), &group(), &request, &weights).unwrap();
834 for parameter in 0..dense.parameter_count {
835 let expected = (0..tof_us.len())
836 .map(|sample| dense.d_y[parameter * tof_us.len() + sample] * weights[sample])
837 .sum::<f64>();
838 assert!((vjp.gradient[parameter] - expected).abs() < 3.0e-11 * expected.abs().max(1.0));
839 }
840 let forward = jvp
841 .d_y
842 .iter()
843 .zip(&weights)
844 .map(|(left, right)| left * right)
845 .sum::<f64>();
846 let reverse = tangent
847 .iter()
848 .zip(&vjp.gradient)
849 .map(|(left, right)| left * right)
850 .sum::<f64>();
851 assert!((forward - reverse).abs() < 2.0e-11 * forward.abs().max(1.0));
852 }
853
854 #[test]
855 fn correction_angle_and_reverse_weight_boundaries_are_explicit() {
856 let tof_us = grid();
857 let mut request = input(&tof_us, &XYZ, &OCCUPANCY, &U_ISO, 1.3);
858 request.correction_model = IntegratedIntensityCorrectionModel::TimeOfFlightNeutronLorentz {
859 two_theta_deg: 90.0,
860 };
861 assert!(matches!(
862 calculate_structural_tof_pattern(cell(), &group(), &request),
863 Err(StructuralTofError::IncompatibleCorrectionModel)
864 ));
865 request.correction_model = IntegratedIntensityCorrectionModel::Neutral;
866 assert!(matches!(
867 calculate_structural_tof_pattern_vjp(cell(), &group(), &request, &[1.0]),
868 Err(StructuralTofError::PatternWeightLengthMismatch)
869 ));
870 }
871}