Skip to main content

phasesmith_engine/
structural_pattern.rs

1//! Fused built-in scattering, structural intensity, and CW profile composition.
2
3use std::error::Error;
4use std::fmt::{Display, Formatter};
5
6use phasesmith_core::{
7    Accumulation, ConstantWavelengthInstrument, CwContributionsError, CwContributionsView, CwError,
8    FcjGeometry, GridView, ProfileError, SupportPolicy,
9    accumulate_cw_contributions_batch_with_context,
10    accumulate_cw_fcj_contributions_batch_with_context,
11};
12use phasesmith_crystallography::{
13    CellError, IntegratedIntensityCorrection, IntegratedIntensityCorrectionError,
14    IntegratedIntensityCorrectionModel, PreparedNeutronScattering, PreparedXrayScattering,
15    ScatteringBatch, ScatteringError, SpaceGroup, StructureFactorBatchError,
16    StructureFactorBatchView, StructureFactorValues, UnitCell,
17    calculate_structure_factor_dense_with_context,
18    calculate_structure_factor_intensity_vjp_with_context,
19    calculate_structure_factor_jvp_with_context, calculate_structure_factor_values_with_context,
20};
21use phasesmith_execution::ExecutionContext;
22
23const CELL_PARAMETER_COUNT: usize = 6;
24pub(crate) const CW_INSTRUMENT_PARAMETER_COUNT: usize = 5;
25const DEGREES_PER_RADIAN: f64 = 180.0 / std::f64::consts::PI;
26
27/// Monochromatic peak-position corrections evaluated with structural geometry.
28#[derive(Clone, Copy, Debug, PartialEq)]
29pub struct MonochromaticPositionCorrection {
30    /// Constant additive shift in degrees `2theta`.
31    pub zero_shift_deg: f64,
32    /// Optional Bragg--Brentano `(sample displacement, goniometer radius)` in mm.
33    pub bragg_brentano_mm: Option<(f64, f64)>,
34    /// Optional Debye--Scherrer `(X, Y, radius)` with displacements in micrometres
35    /// and the goniometer radius in millimetres.
36    pub debye_scherrer_micrometre: Option<(f64, f64, f64)>,
37}
38
39/// Reflection geometry needed by built-in sample-physics providers.
40#[derive(Clone, Debug, PartialEq)]
41pub struct MonochromaticReflectionGeometry {
42    /// Reflection d-spacings in ångströms.
43    pub d_spacing_angstrom: Vec<f64>,
44    /// Corrected reflection positions in degrees `2theta`.
45    pub two_theta_deg: Vec<f64>,
46}
47
48/// Calculate corrected monochromatic reflection positions without profiles.
49///
50/// # Errors
51///
52/// Returns [`StructuralPatternError`] for invalid cell, wavelength, position
53/// correction, or inaccessible reflections.
54pub fn calculate_monochromatic_reflection_geometry(
55    cell: UnitCell,
56    hkl: &[[i32; 3]],
57    instrument: ConstantWavelengthInstrument,
58    position_correction: MonochromaticPositionCorrection,
59) -> Result<MonochromaticReflectionGeometry, StructuralPatternError> {
60    instrument
61        .validate()
62        .map_err(StructuralPatternError::InvalidInstrument)?;
63    validate_position_correction(position_correction)?;
64    let geometry = cell
65        .geometry()
66        .map_err(StructureFactorBatchError::Cell)
67        .map_err(StructuralPatternError::StructureFactor)?;
68    let mut d_spacing_angstrom = Vec::with_capacity(hkl.len());
69    let mut two_theta_deg = Vec::with_capacity(hkl.len());
70    for &reflection in hkl {
71        let q_value = geometry.q_squared(reflection);
72        if !q_value.is_finite() || q_value <= 0.0 {
73            return Err(StructuralPatternError::ReflectionOutsideAngularDomain);
74        }
75        let root_q = q_value.sqrt();
76        let sin_theta = 0.5 * instrument.wavelength_angstrom * root_q;
77        if !(0.0..1.0).contains(&sin_theta) {
78            return Err(StructuralPatternError::ReflectionOutsideAngularDomain);
79        }
80        d_spacing_angstrom.push(root_q.recip());
81        two_theta_deg.push(
82            corrected_monochromatic_position(2.0 * sin_theta.asin(), position_correction)
83                .position_deg,
84        );
85    }
86    Ok(MonochromaticReflectionGeometry {
87        d_spacing_angstrom,
88        two_theta_deg,
89    })
90}
91
92/// Built-in native scattering model selected without a Python callback.
93#[derive(Clone, Copy, Debug, PartialEq, Eq)]
94pub enum BuiltInScatteringModel {
95    /// Non-resonant Waasmaier--Kirfel X-ray form factors.
96    XrayNonResonant,
97    /// Constant bound coherent nuclear-neutron scattering lengths.
98    NeutronNuclear,
99}
100
101/// Borrowed structure, reflection, profile, and sample-physics inputs.
102#[derive(Clone, Copy, Debug)]
103pub struct StructuralPatternInputView<'a> {
104    /// Sorted pattern grid in degrees `2theta`.
105    pub x_deg: &'a [f64],
106    /// Canonical Miller indices.
107    pub hkl: &'a [[i32; 3]],
108    /// Powder multiplicity for each reflection.
109    pub multiplicity: &'a [usize],
110    /// Asymmetric-unit fractional coordinates.
111    pub fractional_xyz: &'a [[f64; 3]],
112    /// Asymmetric-site occupancies.
113    pub occupancy: &'a [f64],
114    /// Asymmetric-site isotropic displacement in square ångströms.
115    pub u_iso_angstrom2: &'a [f64],
116    /// True for asymmetric sites described by fixed CIF U tensors.
117    pub anisotropic_mask: &'a [bool],
118    /// CIF U tensors in component order `11,22,33,23,13,12`.
119    pub u_aniso_cif_angstrom2: &'a [[f64; 6]],
120    /// Exact built-in table key for every asymmetric site.
121    pub scattering_species: &'a [&'a str],
122    /// Fixed real X-ray dispersion offset for every site, or empty when absent.
123    pub scattering_real_offset: &'a [f64],
124    /// Fixed imaginary X-ray dispersion offset for every site, or empty when absent.
125    pub scattering_imag_offset: &'a [f64],
126    /// Structural phase scale.
127    pub scale: f64,
128    /// Fixed symmetry-expansion deduplication tolerance.
129    pub coordinate_tolerance: f64,
130    /// Monochromatic CW instrument/profile parameters.
131    pub instrument: ConstantWavelengthInstrument,
132    /// Optional Finger--Cox--Jephcoat axial-divergence geometry.
133    pub axial_geometry: Option<FcjGeometry>,
134    /// Explicit zero/sample-displacement position correction.
135    pub position_correction: MonochromaticPositionCorrection,
136    /// Explicit integrated-intensity correction model.
137    pub correction_model: IntegratedIntensityCorrectionModel,
138    /// Built-in native scattering selection.
139    pub scattering_model: BuiltInScatteringModel,
140    /// Vectorized sample-physics contribution batch.
141    pub contributions: CwContributionsView<'a>,
142    /// Exact finite profile-support policy.
143    pub support: SupportPolicy,
144}
145
146/// Structural reflection intermediates and fused profile result.
147#[derive(Clone, Debug, PartialEq)]
148pub struct StructuralPatternResult {
149    /// Structure factors and integrated intensities before sample physics.
150    pub structure_factors: StructureFactorValues,
151    /// Reflection d-spacings in ångströms.
152    pub d_spacing_angstrom: Vec<f64>,
153    /// Monochromatic peak positions in degrees `2theta`.
154    pub two_theta_deg: Vec<f64>,
155    /// Support-limited profile values and local/global derivatives.
156    pub accumulation: Accumulation,
157}
158
159/// Fused values and one structural forward derivative product.
160#[derive(Clone, Debug, PartialEq)]
161pub struct StructuralPatternJvpResult {
162    /// Calculated structural pattern.
163    pub result: StructuralPatternResult,
164    /// Structural directional derivative of the pattern samples.
165    pub d_y: Vec<f64>,
166    /// Directional derivative of integrated reflection intensities.
167    pub d_integrated_intensity: Vec<f64>,
168    /// Directional derivative of reflection positions in degrees.
169    pub d_two_theta_deg: Vec<f64>,
170}
171
172/// Fused values and a reusable parameter-major structural pattern Jacobian.
173#[derive(Clone, Debug, PartialEq)]
174pub struct StructuralPatternDenseResult {
175    /// Calculated structural pattern.
176    pub result: StructuralPatternResult,
177    /// Pattern Jacobian with shape parameter count by sample count.
178    pub d_y: Vec<f64>,
179    /// Number of rows in the pattern Jacobian.
180    pub parameter_count: usize,
181}
182
183/// Fused values and one reverse product from pattern sample weights.
184#[derive(Clone, Debug, PartialEq)]
185pub struct StructuralPatternVjpResult {
186    /// Calculated structural pattern.
187    pub result: StructuralPatternResult,
188    /// Pattern-Jacobian transpose product in structural parameter order.
189    pub gradient: Vec<f64>,
190}
191
192/// Invalid fused structural-pattern request.
193#[derive(Debug)]
194pub enum StructuralPatternError {
195    /// The owned unit cell is invalid.
196    InvalidCell(CellError),
197    /// Reflection indices and multiplicities have different lengths.
198    ReflectionLengthMismatch,
199    /// Owned asymmetric-site arrays have different lengths.
200    SiteLengthMismatch,
201    /// Scattering species count does not match the asymmetric-site count.
202    SpeciesLengthMismatch,
203    /// Offset vectors are neither both empty nor matched to the asymmetric sites.
204    ScatteringOffsetLengthMismatch,
205    /// A fixed scattering offset is non-finite.
206    NonFiniteScatteringOffset,
207    /// Fixed dispersion offsets were supplied to a non-X-ray model.
208    UnsupportedScatteringOffset,
209    /// Built-in scattering preparation or evaluation failed.
210    Scattering(ScatteringError),
211    /// Integrated-intensity correction evaluation failed.
212    Correction(IntegratedIntensityCorrectionError),
213    /// A TOF-only intensity correction was supplied to the CW pattern engine.
214    TimeOfFlightCorrectionInConstantWavelengthPattern,
215    /// General-symmetry structural intensity failed.
216    StructureFactor(StructureFactorBatchError),
217    /// Grid validation failed.
218    Profile(ProfileError),
219    /// CW/sample-physics accumulation failed.
220    Contributions(CwContributionsError),
221    /// A reflection is inaccessible for the monochromatic wavelength.
222    ReflectionOutsideAngularDomain,
223    /// The CW instrument model is invalid.
224    InvalidInstrument(CwError),
225    /// A position-correction parameter is invalid.
226    InvalidPositionCorrection,
227    /// Pattern reverse weights do not match the sample count.
228    PatternWeightLengthMismatch,
229    /// A pattern reverse weight is non-finite.
230    NonFinitePatternWeight,
231}
232
233impl Display for StructuralPatternError {
234    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
235        match self {
236            Self::InvalidCell(error) => Display::fmt(error, formatter),
237            Self::ReflectionLengthMismatch => {
238                formatter.write_str("hkl and multiplicity must have the same reflection count")
239            }
240            Self::SiteLengthMismatch => {
241                formatter.write_str("all structural site arrays must have the same site count")
242            }
243            Self::SpeciesLengthMismatch => {
244                formatter.write_str("scattering species must contain one key per asymmetric site")
245            }
246            Self::ScatteringOffsetLengthMismatch => formatter
247                .write_str("scattering offset vectors must both be empty or match the site count"),
248            Self::NonFiniteScatteringOffset => {
249                formatter.write_str("scattering offsets must be finite")
250            }
251            Self::UnsupportedScatteringOffset => formatter
252                .write_str("fixed scattering offsets are supported only for X-ray scattering"),
253            Self::Scattering(error) => Display::fmt(error, formatter),
254            Self::Correction(error) => Display::fmt(error, formatter),
255            Self::TimeOfFlightCorrectionInConstantWavelengthPattern => formatter.write_str(
256                "TOF neutron Lorentz correction is not valid for a constant-wavelength pattern",
257            ),
258            Self::StructureFactor(error) => Display::fmt(error, formatter),
259            Self::Profile(error) => Display::fmt(error, formatter),
260            Self::Contributions(error) => Display::fmt(error, formatter),
261            Self::ReflectionOutsideAngularDomain => formatter.write_str(
262                "structural CW reflections must lie strictly within 0 < 2theta < 180 degrees",
263            ),
264            Self::InvalidInstrument(error) => Display::fmt(error, formatter),
265            Self::InvalidPositionCorrection => formatter.write_str(
266                "position corrections must be finite and goniometer radius must be positive",
267            ),
268            Self::PatternWeightLengthMismatch => {
269                formatter.write_str("pattern reverse weights must match the sample count")
270            }
271            Self::NonFinitePatternWeight => {
272                formatter.write_str("pattern reverse weights must be finite")
273            }
274        }
275    }
276}
277
278impl Error for StructuralPatternError {}
279
280struct PreparedNumerics {
281    scattering: ScatteringBatch,
282    correction: IntegratedIntensityCorrection,
283    d_spacing: Vec<f64>,
284    two_theta_deg: Vec<f64>,
285    d_two_theta_d_cell: Vec<[f64; CELL_PARAMETER_COUNT]>,
286    d_two_theta_d_wavelength: Vec<f64>,
287    d_two_theta_d_sample_displacement: Option<Vec<f64>>,
288    d_two_theta_d_displace_x: Option<Vec<f64>>,
289    d_two_theta_d_displace_y: Option<Vec<f64>>,
290}
291
292#[derive(Clone, Copy, Debug)]
293struct CorrectedPosition {
294    position_deg: f64,
295    d_position_d_base: f64,
296    d_position_d_sample_displacement: Option<f64>,
297    d_position_d_displace_x: Option<f64>,
298    d_position_d_displace_y: Option<f64>,
299}
300
301fn corrected_monochromatic_position(
302    base_position_radians: f64,
303    correction: MonochromaticPositionCorrection,
304) -> CorrectedPosition {
305    let mut result = CorrectedPosition {
306        position_deg: base_position_radians.to_degrees() + correction.zero_shift_deg,
307        d_position_d_base: 1.0,
308        d_position_d_sample_displacement: None,
309        d_position_d_displace_x: None,
310        d_position_d_displace_y: None,
311    };
312    if let Some((displacement, radius)) = correction.bragg_brentano_mm {
313        let theta = 0.5 * base_position_radians;
314        result.position_deg -= 2.0 * displacement / radius * theta.cos() * DEGREES_PER_RADIAN;
315        result.d_position_d_base += displacement / radius * theta.sin();
316        result.d_position_d_sample_displacement =
317            Some(-2.0 / radius * theta.cos() * DEGREES_PER_RADIAN);
318    }
319    if let Some((displace_x, displace_y, radius)) = correction.debye_scherrer_micrometre {
320        let (sin_position, cos_position) = base_position_radians.sin_cos();
321        let displacement_scale = 0.18 / (std::f64::consts::PI * radius);
322        result.position_deg -=
323            displacement_scale * (displace_x * cos_position + displace_y * sin_position);
324        result.d_position_d_base += displacement_scale.to_radians()
325            * (displace_x * sin_position - displace_y * cos_position);
326        result.d_position_d_displace_x = Some(-displacement_scale * cos_position);
327        result.d_position_d_displace_y = Some(-displacement_scale * sin_position);
328    }
329    result
330}
331
332fn validate_position_correction(
333    correction: MonochromaticPositionCorrection,
334) -> Result<(), StructuralPatternError> {
335    let invalid = !correction.zero_shift_deg.is_finite()
336        || (correction.bragg_brentano_mm.is_some()
337            && correction.debye_scherrer_micrometre.is_some())
338        || correction
339            .bragg_brentano_mm
340            .is_some_and(|(displacement, radius)| {
341                !displacement.is_finite() || !radius.is_finite() || radius <= 0.0
342            })
343        || correction
344            .debye_scherrer_micrometre
345            .is_some_and(|(displace_x, displace_y, radius)| {
346                !displace_x.is_finite()
347                    || !displace_y.is_finite()
348                    || !radius.is_finite()
349                    || radius <= 0.0
350            });
351    if invalid {
352        return Err(StructuralPatternError::InvalidPositionCorrection);
353    }
354    Ok(())
355}
356
357fn validate_scattering_offsets(
358    input: &StructuralPatternInputView<'_>,
359) -> Result<(), StructuralPatternError> {
360    if input.scattering_real_offset.is_empty() && input.scattering_imag_offset.is_empty() {
361        return Ok(());
362    }
363    if input.scattering_real_offset.len() != input.fractional_xyz.len()
364        || input.scattering_imag_offset.len() != input.fractional_xyz.len()
365    {
366        return Err(StructuralPatternError::ScatteringOffsetLengthMismatch);
367    }
368    if input
369        .scattering_real_offset
370        .iter()
371        .chain(input.scattering_imag_offset)
372        .any(|value| !value.is_finite())
373    {
374        return Err(StructuralPatternError::NonFiniteScatteringOffset);
375    }
376    if input.scattering_model == BuiltInScatteringModel::NeutronNuclear
377        && input
378            .scattering_real_offset
379            .iter()
380            .chain(input.scattering_imag_offset)
381            .any(|value| *value != 0.0)
382    {
383        return Err(StructuralPatternError::UnsupportedScatteringOffset);
384    }
385    Ok(())
386}
387
388fn apply_scattering_offsets(
389    scattering: &mut ScatteringBatch,
390    input: &StructuralPatternInputView<'_>,
391) {
392    if input.scattering_real_offset.is_empty() {
393        return;
394    }
395    for reflection in 0..scattering.reflection_count {
396        for site in 0..scattering.site_count {
397            let index = reflection * scattering.site_count + site;
398            scattering.real[index] += input.scattering_real_offset[site];
399            scattering.imag[index] += input.scattering_imag_offset[site];
400        }
401    }
402}
403
404impl PreparedNumerics {
405    fn structure_batch<'a>(
406        &'a self,
407        input: &StructuralPatternInputView<'a>,
408    ) -> StructureFactorBatchView<'a> {
409        StructureFactorBatchView {
410            hkl: input.hkl,
411            multiplicity: input.multiplicity,
412            fractional_xyz: input.fractional_xyz,
413            occupancy: input.occupancy,
414            u_iso_angstrom2: input.u_iso_angstrom2,
415            anisotropic_mask: input.anisotropic_mask,
416            u_aniso_cif_angstrom2: input.u_aniso_cif_angstrom2,
417            scattering_real: &self.scattering.real,
418            scattering_imag: &self.scattering.imag,
419            d_scattering_real_d_s: &self.scattering.d_real_d_s,
420            d_scattering_imag_d_s: &self.scattering.d_imag_d_s,
421            correction: &self.correction.values,
422            d_correction_d_q_squared: &self.correction.d_values_d_q_squared,
423            scale: input.scale,
424            coordinate_tolerance: input.coordinate_tolerance,
425        }
426    }
427}
428
429/// Calculate built-in scattering, structural intensities, and one CW profile.
430///
431/// # Errors
432///
433/// Returns [`StructuralPatternError`] for invalid structural, scattering,
434/// correction, instrument, contribution, grid, or support inputs.
435pub fn calculate_structural_pattern(
436    cell: UnitCell,
437    space_group: &SpaceGroup,
438    input: &StructuralPatternInputView<'_>,
439) -> Result<StructuralPatternResult, StructuralPatternError> {
440    calculate_structural_pattern_with_context(cell, space_group, input, &ExecutionContext::serial())
441}
442
443/// Calculate a structural pattern with an explicit bounded execution context.
444///
445/// # Errors
446///
447/// Returns [`StructuralPatternError`] for invalid inputs.
448pub fn calculate_structural_pattern_with_context(
449    cell: UnitCell,
450    space_group: &SpaceGroup,
451    input: &StructuralPatternInputView<'_>,
452    execution: &ExecutionContext,
453) -> Result<StructuralPatternResult, StructuralPatternError> {
454    let prepared = prepare(cell, input)?;
455    calculate_values(cell, space_group, input, &prepared, execution)
456}
457
458/// Calculate values and a reusable dense structural pattern linearization.
459///
460/// # Errors
461///
462/// Returns an error for invalid inputs or allocation overflow.
463pub fn calculate_structural_pattern_dense(
464    cell: UnitCell,
465    space_group: &SpaceGroup,
466    input: &StructuralPatternInputView<'_>,
467) -> Result<StructuralPatternDenseResult, StructuralPatternError> {
468    calculate_structural_pattern_dense_with_context(
469        cell,
470        space_group,
471        input,
472        &ExecutionContext::serial(),
473    )
474}
475
476/// Calculate a dense structural linearization with a bounded context.
477///
478/// # Errors
479///
480/// Returns an error for invalid inputs or allocation overflow.
481pub fn calculate_structural_pattern_dense_with_context(
482    cell: UnitCell,
483    space_group: &SpaceGroup,
484    input: &StructuralPatternInputView<'_>,
485    execution: &ExecutionContext,
486) -> Result<StructuralPatternDenseResult, StructuralPatternError> {
487    let prepared = prepare(cell, input)?;
488    let structural = calculate_structure_factor_dense_with_context(
489        cell,
490        space_group,
491        prepared.structure_batch(input),
492        execution,
493    )
494    .map_err(StructuralPatternError::StructureFactor)?;
495    let parameter_count = structural.layout.parameter_count();
496    let sample_count = input.x_deg.len();
497    let element_count =
498        parameter_count
499            .checked_mul(sample_count)
500            .ok_or(StructuralPatternError::Contributions(
501                CwContributionsError::AllocationOverflow,
502            ))?;
503    let mut accumulation = accumulate(
504        input,
505        &prepared.two_theta_deg,
506        &structural.values.intensity,
507        execution,
508    )?;
509    append_instrument_derivatives(&mut accumulation, &structural.values, input, &prepared)?;
510    let reflection_count = input.hkl.len();
511    let local = &accumulation.derivatives.local;
512    let d_y = if execution.threads() == 1 || parameter_count < 2 {
513        let mut values = zeroed_values(element_count)?;
514        for reflection in 0..reflection_count {
515            let begin = local.offsets[reflection];
516            let end = local.offsets[reflection + 1];
517            for active in begin..end {
518                let sample = local.starts[reflection] + active - begin;
519                let local_base = 2 * active;
520                for parameter in 0..parameter_count {
521                    let structural_index = parameter * reflection_count + reflection;
522                    let position_derivative = if parameter < CELL_PARAMETER_COUNT {
523                        prepared.d_two_theta_d_cell[reflection][parameter]
524                    } else {
525                        0.0
526                    };
527                    values[parameter * sample_count + sample] += local.values[local_base]
528                        * structural.d_intensity[structural_index]
529                        + local.values[local_base + 1] * position_derivative;
530                }
531            }
532        }
533        values
534    } else {
535        let rows = execution.map_ordered(parameter_count, 2, |parameter| {
536            let mut row = zeroed_values(sample_count)?;
537            for reflection in 0..reflection_count {
538                let begin = local.offsets[reflection];
539                let end = local.offsets[reflection + 1];
540                for active in begin..end {
541                    let sample = local.starts[reflection] + active - begin;
542                    let local_base = 2 * active;
543                    let structural_index = parameter * reflection_count + reflection;
544                    let position_derivative = if parameter < CELL_PARAMETER_COUNT {
545                        prepared.d_two_theta_d_cell[reflection][parameter]
546                    } else {
547                        0.0
548                    };
549                    row[sample] += local.values[local_base]
550                        * structural.d_intensity[structural_index]
551                        + local.values[local_base + 1] * position_derivative;
552                }
553            }
554            Ok::<_, StructuralPatternError>(row)
555        });
556        let mut values = Vec::new();
557        values.try_reserve_exact(element_count).map_err(|_| {
558            StructuralPatternError::Contributions(CwContributionsError::AllocationOverflow)
559        })?;
560        for row in rows {
561            values.extend(row?);
562        }
563        values
564    };
565    Ok(StructuralPatternDenseResult {
566        result: StructuralPatternResult {
567            structure_factors: structural.values,
568            d_spacing_angstrom: prepared.d_spacing,
569            two_theta_deg: prepared.two_theta_deg,
570            accumulation,
571        },
572        d_y,
573        parameter_count,
574    })
575}
576
577fn zeroed_values(count: usize) -> Result<Vec<f64>, StructuralPatternError> {
578    let mut values = Vec::new();
579    values.try_reserve_exact(count).map_err(|_| {
580        StructuralPatternError::Contributions(CwContributionsError::AllocationOverflow)
581    })?;
582    values.resize(count, 0.0);
583    Ok(values)
584}
585
586/// Calculate a full structural-pattern JVP without a dense pattern Jacobian.
587///
588/// # Errors
589///
590/// Returns [`StructuralPatternError`] for invalid inputs or structural tangent.
591pub fn calculate_structural_pattern_jvp(
592    cell: UnitCell,
593    space_group: &SpaceGroup,
594    input: &StructuralPatternInputView<'_>,
595    tangent: &[f64],
596) -> Result<StructuralPatternJvpResult, StructuralPatternError> {
597    calculate_structural_pattern_jvp_with_context(
598        cell,
599        space_group,
600        input,
601        tangent,
602        &ExecutionContext::serial(),
603    )
604}
605
606/// Calculate a structural JVP with a bounded execution context.
607///
608/// # Errors
609///
610/// Returns [`StructuralPatternError`] for invalid inputs or tangent shape.
611pub fn calculate_structural_pattern_jvp_with_context(
612    cell: UnitCell,
613    space_group: &SpaceGroup,
614    input: &StructuralPatternInputView<'_>,
615    tangent: &[f64],
616    execution: &ExecutionContext,
617) -> Result<StructuralPatternJvpResult, StructuralPatternError> {
618    let prepared = prepare(cell, input)?;
619    let structural = calculate_structure_factor_jvp_with_context(
620        cell,
621        space_group,
622        prepared.structure_batch(input),
623        tangent,
624        execution,
625    )
626    .map_err(StructuralPatternError::StructureFactor)?;
627    let d_two_theta_deg = prepared
628        .d_two_theta_d_cell
629        .iter()
630        .map(|derivatives| {
631            derivatives
632                .iter()
633                .zip(&tangent[..CELL_PARAMETER_COUNT])
634                .map(|(derivative, direction)| derivative * direction)
635                .sum::<f64>()
636        })
637        .collect::<Vec<_>>();
638    let mut accumulation = accumulate(
639        input,
640        &prepared.two_theta_deg,
641        &structural.values.intensity,
642        execution,
643    )?;
644    append_instrument_derivatives(&mut accumulation, &structural.values, input, &prepared)?;
645    let d_y = chain_pattern_jvp(&accumulation, &structural.d_intensity, &d_two_theta_deg);
646    Ok(StructuralPatternJvpResult {
647        result: StructuralPatternResult {
648            structure_factors: structural.values,
649            d_spacing_angstrom: prepared.d_spacing,
650            two_theta_deg: prepared.two_theta_deg,
651            accumulation,
652        },
653        d_y,
654        d_integrated_intensity: structural.d_intensity,
655        d_two_theta_deg,
656    })
657}
658
659/// Calculate a full structural-pattern transpose product from sample weights.
660///
661/// # Errors
662///
663/// Returns [`StructuralPatternError`] for invalid inputs or sample weights.
664pub fn calculate_structural_pattern_vjp(
665    cell: UnitCell,
666    space_group: &SpaceGroup,
667    input: &StructuralPatternInputView<'_>,
668    sample_weights: &[f64],
669) -> Result<StructuralPatternVjpResult, StructuralPatternError> {
670    calculate_structural_pattern_vjp_with_context(
671        cell,
672        space_group,
673        input,
674        sample_weights,
675        &ExecutionContext::serial(),
676    )
677}
678
679/// Calculate a structural transpose product with a bounded context.
680///
681/// # Errors
682///
683/// Returns [`StructuralPatternError`] for invalid inputs or sample weights.
684pub fn calculate_structural_pattern_vjp_with_context(
685    cell: UnitCell,
686    space_group: &SpaceGroup,
687    input: &StructuralPatternInputView<'_>,
688    sample_weights: &[f64],
689    execution: &ExecutionContext,
690) -> Result<StructuralPatternVjpResult, StructuralPatternError> {
691    if sample_weights.len() != input.x_deg.len() {
692        return Err(StructuralPatternError::PatternWeightLengthMismatch);
693    }
694    if sample_weights.iter().any(|value| !value.is_finite()) {
695        return Err(StructuralPatternError::NonFinitePatternWeight);
696    }
697    let prepared = prepare(cell, input)?;
698    let values = calculate_structure_factor_values_with_context(
699        cell,
700        space_group,
701        prepared.structure_batch(input),
702        execution,
703    )
704    .map_err(StructuralPatternError::StructureFactor)?;
705    let mut accumulation =
706        accumulate(input, &prepared.two_theta_deg, &values.intensity, execution)?;
707    append_instrument_derivatives(&mut accumulation, &values, input, &prepared)?;
708    let (intensity_weights, position_weights) =
709        local_transpose_weights(&accumulation, sample_weights);
710    let mut structural = calculate_structure_factor_intensity_vjp_with_context(
711        cell,
712        space_group,
713        prepared.structure_batch(input),
714        &intensity_weights,
715        execution,
716    )
717    .map_err(StructuralPatternError::StructureFactor)?;
718    for (reflection, weight) in position_weights.into_iter().enumerate() {
719        for parameter in 0..CELL_PARAMETER_COUNT {
720            structural.gradient[parameter] +=
721                weight * prepared.d_two_theta_d_cell[reflection][parameter];
722        }
723    }
724    Ok(StructuralPatternVjpResult {
725        result: StructuralPatternResult {
726            structure_factors: values,
727            d_spacing_angstrom: prepared.d_spacing,
728            two_theta_deg: prepared.two_theta_deg,
729            accumulation,
730        },
731        gradient: structural.gradient,
732    })
733}
734
735fn prepare(
736    cell: UnitCell,
737    input: &StructuralPatternInputView<'_>,
738) -> Result<PreparedNumerics, StructuralPatternError> {
739    validate_constant_wavelength_correction(input.correction_model)?;
740    if input.scattering_species.len() != input.fractional_xyz.len() {
741        return Err(StructuralPatternError::SpeciesLengthMismatch);
742    }
743    validate_scattering_offsets(input)?;
744    input
745        .instrument
746        .validate()
747        .map_err(StructuralPatternError::InvalidInstrument)?;
748    validate_position_correction(input.position_correction)?;
749    let geometry = cell
750        .geometry()
751        .map_err(StructureFactorBatchError::Cell)
752        .map_err(StructuralPatternError::StructureFactor)?;
753    let mut q_squared = Vec::with_capacity(input.hkl.len());
754    let mut d_spacing = Vec::with_capacity(input.hkl.len());
755    let mut two_theta_deg = Vec::with_capacity(input.hkl.len());
756    let mut d_two_theta_d_cell = Vec::with_capacity(input.hkl.len());
757    let mut d_two_theta_d_wavelength = Vec::with_capacity(input.hkl.len());
758    let mut d_two_theta_d_sample_displacement = input
759        .position_correction
760        .bragg_brentano_mm
761        .map(|_| Vec::with_capacity(input.hkl.len()));
762    let mut d_two_theta_d_displace_x = input
763        .position_correction
764        .debye_scherrer_micrometre
765        .map(|_| Vec::with_capacity(input.hkl.len()));
766    let mut d_two_theta_d_displace_y = input
767        .position_correction
768        .debye_scherrer_micrometre
769        .map(|_| Vec::with_capacity(input.hkl.len()));
770    for &hkl in input.hkl {
771        let (q_value, d_q) = geometry.q_squared_and_derivatives(hkl);
772        if !q_value.is_finite() || q_value <= 0.0 {
773            return Err(StructuralPatternError::ReflectionOutsideAngularDomain);
774        }
775        let root_q = q_value.sqrt();
776        let sin_theta = 0.5 * input.instrument.wavelength_angstrom * root_q;
777        if !(0.0..1.0).contains(&sin_theta) {
778            return Err(StructuralPatternError::ReflectionOutsideAngularDomain);
779        }
780        let theta = sin_theta.asin();
781        let corrected = corrected_monochromatic_position(2.0 * theta, input.position_correction);
782        let d_position_factor =
783            corrected.d_position_d_base * input.instrument.wavelength_angstrom * DEGREES_PER_RADIAN
784                / (2.0 * root_q * theta.cos());
785        let d_position_d_wavelength =
786            corrected.d_position_d_base * DEGREES_PER_RADIAN * root_q / theta.cos();
787        q_squared.push(q_value);
788        d_spacing.push(root_q.recip());
789        two_theta_deg.push(corrected.position_deg);
790        d_two_theta_d_cell.push(d_q.map(|derivative| d_position_factor * derivative));
791        d_two_theta_d_wavelength.push(d_position_d_wavelength);
792        if let (Some(values), Some(derivative)) = (
793            d_two_theta_d_sample_displacement.as_mut(),
794            corrected.d_position_d_sample_displacement,
795        ) {
796            values.push(derivative);
797        }
798        if let (Some(values), Some(derivative)) = (
799            d_two_theta_d_displace_x.as_mut(),
800            corrected.d_position_d_displace_x,
801        ) {
802            values.push(derivative);
803        }
804        if let (Some(values), Some(derivative)) = (
805            d_two_theta_d_displace_y.as_mut(),
806            corrected.d_position_d_displace_y,
807        ) {
808            values.push(derivative);
809        }
810    }
811    let s: Vec<f64> = q_squared.iter().map(|value| 0.5 * value.sqrt()).collect();
812    let mut scattering = match input.scattering_model {
813        BuiltInScatteringModel::XrayNonResonant => {
814            PreparedXrayScattering::new(input.scattering_species.iter().copied())
815                .and_then(|model| model.evaluate(&s))
816        }
817        BuiltInScatteringModel::NeutronNuclear => {
818            PreparedNeutronScattering::new(input.scattering_species.iter().copied())
819                .and_then(|model| model.evaluate(&s))
820        }
821    }
822    .map_err(StructuralPatternError::Scattering)?;
823    apply_scattering_offsets(&mut scattering, input);
824    let correction = input
825        .correction_model
826        .evaluate(&q_squared)
827        .map_err(StructuralPatternError::Correction)?;
828    Ok(PreparedNumerics {
829        scattering,
830        correction,
831        d_spacing,
832        two_theta_deg,
833        d_two_theta_d_cell,
834        d_two_theta_d_wavelength,
835        d_two_theta_d_sample_displacement,
836        d_two_theta_d_displace_x,
837        d_two_theta_d_displace_y,
838    })
839}
840
841fn validate_constant_wavelength_correction(
842    correction: IntegratedIntensityCorrectionModel,
843) -> Result<(), StructuralPatternError> {
844    if matches!(
845        correction,
846        IntegratedIntensityCorrectionModel::TimeOfFlightNeutronLorentz { .. }
847    ) {
848        Err(StructuralPatternError::TimeOfFlightCorrectionInConstantWavelengthPattern)
849    } else {
850        Ok(())
851    }
852}
853
854fn calculate_values(
855    cell: UnitCell,
856    space_group: &SpaceGroup,
857    input: &StructuralPatternInputView<'_>,
858    prepared: &PreparedNumerics,
859    execution: &ExecutionContext,
860) -> Result<StructuralPatternResult, StructuralPatternError> {
861    let structure_factors = calculate_structure_factor_values_with_context(
862        cell,
863        space_group,
864        prepared.structure_batch(input),
865        execution,
866    )
867    .map_err(StructuralPatternError::StructureFactor)?;
868    let mut accumulation = accumulate(
869        input,
870        &prepared.two_theta_deg,
871        &structure_factors.intensity,
872        execution,
873    )?;
874    append_instrument_derivatives(&mut accumulation, &structure_factors, input, prepared)?;
875    Ok(StructuralPatternResult {
876        structure_factors,
877        d_spacing_angstrom: prepared.d_spacing.clone(),
878        two_theta_deg: prepared.two_theta_deg.clone(),
879        accumulation,
880    })
881}
882
883fn append_instrument_derivatives(
884    accumulation: &mut Accumulation,
885    structure_factors: &StructureFactorValues,
886    input: &StructuralPatternInputView<'_>,
887    prepared: &PreparedNumerics,
888) -> Result<(), StructuralPatternError> {
889    let sample_count = accumulation.sample_count;
890    let extra_count = 2
891        + usize::from(prepared.d_two_theta_d_sample_displacement.is_some())
892        + usize::from(prepared.d_two_theta_d_displace_x.is_some())
893        + usize::from(prepared.d_two_theta_d_displace_y.is_some());
894    let global = accumulation
895        .derivatives
896        .global
897        .as_mut()
898        .expect("CW contribution accumulation always has global derivatives");
899    let old_values = std::mem::take(&mut global.values);
900    let mut combined = Vec::new();
901    combined
902        .try_reserve(old_values.len() + extra_count * sample_count)
903        .map_err(|_| {
904            StructuralPatternError::Contributions(CwContributionsError::AllocationOverflow)
905        })?;
906    let mut wavelength = vec![0.0; sample_count];
907    let mut zero_shift = vec![0.0; sample_count];
908    let mut sample_displacement = prepared
909        .d_two_theta_d_sample_displacement
910        .as_ref()
911        .map(|_| vec![0.0; sample_count]);
912    let mut displace_x = prepared
913        .d_two_theta_d_displace_x
914        .as_ref()
915        .map(|_| vec![0.0; sample_count]);
916    let mut displace_y = prepared
917        .d_two_theta_d_displace_y
918        .as_ref()
919        .map(|_| vec![0.0; sample_count]);
920    let local = &accumulation.derivatives.local;
921    for reflection in 0..local.peak_count() {
922        #[allow(clippy::cast_precision_loss)]
923        let multiplicity = input.multiplicity[reflection] as f64;
924        let d_intensity_d_wavelength = input.scale
925            * multiplicity
926            * structure_factors.f_squared[reflection]
927            * prepared.correction.d_values_d_wavelength[reflection];
928        let begin = local.offsets[reflection];
929        let end = local.offsets[reflection + 1];
930        for active in begin..end {
931            let sample = local.starts[reflection] + active - begin;
932            let base = 2 * active;
933            let d_intensity = local.values[base];
934            let d_position = local.values[base + 1];
935            wavelength[sample] += d_intensity * d_intensity_d_wavelength
936                + d_position * prepared.d_two_theta_d_wavelength[reflection];
937            zero_shift[sample] += d_position;
938            if let (Some(values), Some(derivatives)) = (
939                sample_displacement.as_mut(),
940                prepared.d_two_theta_d_sample_displacement.as_ref(),
941            ) {
942                values[sample] += d_position * derivatives[reflection];
943            }
944            if let (Some(values), Some(derivatives)) = (
945                displace_x.as_mut(),
946                prepared.d_two_theta_d_displace_x.as_ref(),
947            ) {
948                values[sample] += d_position * derivatives[reflection];
949            }
950            if let (Some(values), Some(derivatives)) = (
951                displace_y.as_mut(),
952                prepared.d_two_theta_d_displace_y.as_ref(),
953            ) {
954                values[sample] += d_position * derivatives[reflection];
955            }
956        }
957    }
958    let instrument_end = CW_INSTRUMENT_PARAMETER_COUNT * sample_count;
959    combined.extend_from_slice(&old_values[..instrument_end]);
960    combined.extend(wavelength);
961    combined.extend(zero_shift);
962    if let Some(values) = sample_displacement {
963        combined.extend(values);
964    }
965    if let Some(values) = displace_x {
966        combined.extend(values);
967    }
968    if let Some(values) = displace_y {
969        combined.extend(values);
970    }
971    combined.extend_from_slice(&old_values[instrument_end..]);
972    global.values = combined;
973    global.parameter_count += extra_count;
974    Ok(())
975}
976
977fn accumulate(
978    input: &StructuralPatternInputView<'_>,
979    two_theta_deg: &[f64],
980    intensities: &[f64],
981    execution: &ExecutionContext,
982) -> Result<Accumulation, StructuralPatternError> {
983    let grid = GridView::new(input.x_deg).map_err(StructuralPatternError::Profile)?;
984    let result = match input.axial_geometry {
985        Some(geometry) => accumulate_cw_fcj_contributions_batch_with_context(
986            grid,
987            two_theta_deg,
988            intensities,
989            input.instrument,
990            input.contributions,
991            geometry,
992            input.support,
993            execution,
994        ),
995        None => accumulate_cw_contributions_batch_with_context(
996            grid,
997            two_theta_deg,
998            intensities,
999            input.instrument,
1000            input.contributions,
1001            input.support,
1002            execution,
1003        ),
1004    };
1005    result.map_err(StructuralPatternError::Contributions)
1006}
1007
1008fn chain_pattern_jvp(
1009    accumulation: &Accumulation,
1010    d_intensity: &[f64],
1011    d_position: &[f64],
1012) -> Vec<f64> {
1013    let local = &accumulation.derivatives.local;
1014    let mut result = vec![0.0; accumulation.sample_count];
1015    for reflection in 0..local.peak_count() {
1016        let begin = local.offsets[reflection];
1017        let end = local.offsets[reflection + 1];
1018        for active in begin..end {
1019            let sample = local.starts[reflection] + active - begin;
1020            let base = 2 * active;
1021            result[sample] += local.values[base] * d_intensity[reflection]
1022                + local.values[base + 1] * d_position[reflection];
1023        }
1024    }
1025    result
1026}
1027
1028fn local_transpose_weights(
1029    accumulation: &Accumulation,
1030    sample_weights: &[f64],
1031) -> (Vec<f64>, Vec<f64>) {
1032    let local = &accumulation.derivatives.local;
1033    let mut intensity = vec![0.0; local.peak_count()];
1034    let mut position = vec![0.0; local.peak_count()];
1035    for reflection in 0..local.peak_count() {
1036        let begin = local.offsets[reflection];
1037        let end = local.offsets[reflection + 1];
1038        for active in begin..end {
1039            let sample = local.starts[reflection] + active - begin;
1040            let base = 2 * active;
1041            intensity[reflection] += local.values[base] * sample_weights[sample];
1042            position[reflection] += local.values[base + 1] * sample_weights[sample];
1043        }
1044    }
1045    (intensity, position)
1046}
1047
1048#[cfg(test)]
1049mod tests {
1050    use super::*;
1051    use phasesmith_core::CwContributionArrays;
1052    use phasesmith_crystallography::{P1ParameterLayout, Rational, SymmetryOperation};
1053
1054    fn cell() -> UnitCell {
1055        UnitCell {
1056            a_angstrom: 4.7,
1057            b_angstrom: 5.1,
1058            c_angstrom: 6.2,
1059            alpha_deg: 82.0,
1060            beta_deg: 87.0,
1061            gamma_deg: 74.0,
1062        }
1063    }
1064
1065    fn group() -> SpaceGroup {
1066        SpaceGroup::new(vec![
1067            SymmetryOperation::identity(),
1068            SymmetryOperation::new([[-1, 0, 0], [0, -1, 0], [0, 0, -1]], [Rational::zero(); 3])
1069                .expect("inversion"),
1070        ])
1071        .expect("P-1")
1072    }
1073
1074    fn instrument() -> ConstantWavelengthInstrument {
1075        ConstantWavelengthInstrument {
1076            wavelength_angstrom: 1.5406,
1077            u_deg2: 2.0e-4,
1078            v_deg2: -1.0e-4,
1079            w_deg2: 1.2e-4,
1080            x_deg: 1.5e-3,
1081            y_deg: 3.0e-3,
1082        }
1083    }
1084
1085    #[allow(clippy::too_many_arguments)]
1086    fn calculate_case(
1087        selected_cell: UnitCell,
1088        x: &[f64],
1089        hkl: &[[i32; 3]],
1090        multiplicity: &[usize],
1091        xyz: &[[f64; 3]],
1092        occupancy: &[f64],
1093        u_iso: &[f64],
1094        scale: f64,
1095        multiplier: &[f64],
1096    ) -> StructuralPatternResult {
1097        calculate_case_with_correction(
1098            selected_cell,
1099            x,
1100            hkl,
1101            multiplicity,
1102            xyz,
1103            occupancy,
1104            u_iso,
1105            scale,
1106            multiplier,
1107            IntegratedIntensityCorrectionModel::Neutral,
1108        )
1109        .expect("structural pattern")
1110    }
1111
1112    #[allow(clippy::too_many_arguments)]
1113    fn calculate_case_with_correction(
1114        selected_cell: UnitCell,
1115        x: &[f64],
1116        hkl: &[[i32; 3]],
1117        multiplicity: &[usize],
1118        xyz: &[[f64; 3]],
1119        occupancy: &[f64],
1120        u_iso: &[f64],
1121        scale: f64,
1122        multiplier: &[f64],
1123        correction_model: IntegratedIntensityCorrectionModel,
1124    ) -> Result<StructuralPatternResult, StructuralPatternError> {
1125        let zeros = vec![0.0; hkl.len()];
1126        let contributions = CwContributionsView::new(
1127            hkl.len(),
1128            0,
1129            CwContributionArrays {
1130                gaussian_variance_deg2: &zeros,
1131                lorentzian_fwhm_deg: &zeros,
1132                intensity_multiplier: multiplier,
1133                d_gaussian_variance_d_position: &zeros,
1134                d_lorentzian_fwhm_d_position: &zeros,
1135                d_intensity_multiplier_d_position: &zeros,
1136                d_gaussian_variance_d_parameters: &[],
1137                d_lorentzian_fwhm_d_parameters: &[],
1138                d_intensity_multiplier_d_parameters: &[],
1139            },
1140        )
1141        .expect("contributions");
1142        calculate_structural_pattern(
1143            selected_cell,
1144            &group(),
1145            &StructuralPatternInputView {
1146                x_deg: x,
1147                hkl,
1148                multiplicity,
1149                fractional_xyz: xyz,
1150                occupancy,
1151                u_iso_angstrom2: u_iso,
1152                anisotropic_mask: &[false, false],
1153                u_aniso_cif_angstrom2: &[[0.0; 6]; 2],
1154                scattering_species: &["Si", "O"],
1155                scattering_real_offset: &[],
1156                scattering_imag_offset: &[],
1157                scale,
1158                coordinate_tolerance: 1.0e-10,
1159                instrument: instrument(),
1160                axial_geometry: None,
1161                position_correction: MonochromaticPositionCorrection {
1162                    zero_shift_deg: 0.0,
1163                    bragg_brentano_mm: None,
1164                    debye_scherrer_micrometre: None,
1165                },
1166                correction_model,
1167                scattering_model: BuiltInScatteringModel::XrayNonResonant,
1168                contributions,
1169                support: SupportPolicy::FwhmMultiple(20.0),
1170            },
1171        )
1172    }
1173
1174    #[test]
1175    fn constant_wavelength_pattern_rejects_tof_only_correction() {
1176        let error = calculate_case_with_correction(
1177            cell(),
1178            &[10.0, 20.0],
1179            &[[1, 0, 1]],
1180            &[2],
1181            &[[0.17, 0.23, 0.31], [0.37, 0.11, 0.19]],
1182            &[0.82, 0.55],
1183            &[0.012, 0.018],
1184            1.0,
1185            &[1.0],
1186            IntegratedIntensityCorrectionModel::TimeOfFlightNeutronLorentz {
1187                two_theta_deg: 90.0,
1188            },
1189        )
1190        .expect_err("TOF correction must not enter CW accumulation");
1191        assert!(matches!(
1192            error,
1193            StructuralPatternError::TimeOfFlightCorrectionInConstantWavelengthPattern
1194        ));
1195    }
1196
1197    #[test]
1198    fn fused_values_apply_sample_intensity_multiplier_exactly_once() {
1199        let x: Vec<f64> = (0..9_001)
1200            .map(|index| 10.0 + f64::from(index) * 0.01)
1201            .collect();
1202        let hkl = [[1, 0, 1], [2, 1, 1], [1, 2, 3]];
1203        let multiplicity = [2, 4, 2];
1204        let xyz = [[0.17, 0.23, 0.31], [0.37, 0.11, 0.19]];
1205        let occupancy = [0.82, 0.55];
1206        let u_iso = [0.012, 0.018];
1207        let neutral = calculate_case(
1208            cell(),
1209            &x,
1210            &hkl,
1211            &multiplicity,
1212            &xyz,
1213            &occupancy,
1214            &u_iso,
1215            1.4,
1216            &[1.0; 3],
1217        );
1218        let doubled = calculate_case(
1219            cell(),
1220            &x,
1221            &hkl,
1222            &multiplicity,
1223            &xyz,
1224            &occupancy,
1225            &u_iso,
1226            1.4,
1227            &[2.0; 3],
1228        );
1229        assert_eq!(
1230            neutral.structure_factors.intensity,
1231            doubled.structure_factors.intensity
1232        );
1233        for (left, right) in neutral.accumulation.y.iter().zip(&doubled.accumulation.y) {
1234            assert!((2.0 * left - right).abs() < 2.0e-15 * right.abs().max(1.0));
1235        }
1236    }
1237
1238    #[test]
1239    fn reflection_geometry_helper_matches_the_fused_structural_positions() {
1240        let x = (0..9_001)
1241            .map(|index| 10.0 + f64::from(index) * 0.01)
1242            .collect::<Vec<_>>();
1243        let hkl = [[1, 0, 1], [2, 1, 1], [1, 2, 3]];
1244        let correction = MonochromaticPositionCorrection {
1245            zero_shift_deg: 0.0,
1246            bragg_brentano_mm: None,
1247            debye_scherrer_micrometre: None,
1248        };
1249        let fused = calculate_case(
1250            cell(),
1251            &x,
1252            &hkl,
1253            &[2, 4, 2],
1254            &[[0.17, 0.23, 0.31], [0.37, 0.11, 0.19]],
1255            &[0.82, 0.55],
1256            &[0.012, 0.018],
1257            1.4,
1258            &[1.0; 3],
1259        );
1260        let geometry =
1261            calculate_monochromatic_reflection_geometry(cell(), &hkl, instrument(), correction)
1262                .unwrap();
1263        assert_eq!(geometry.d_spacing_angstrom, fused.d_spacing_angstrom);
1264        assert_eq!(geometry.two_theta_deg, fused.two_theta_deg);
1265    }
1266
1267    #[test]
1268    #[allow(clippy::too_many_lines)]
1269    fn fused_jvp_vjp_match_pattern_finite_differences_and_adjoint_identity() {
1270        let x: Vec<f64> = (0..9_001)
1271            .map(|index| 10.0 + f64::from(index) * 0.01)
1272            .collect();
1273        let hkl = [[1, 0, 1], [2, 1, 1], [1, 2, 3]];
1274        let multiplicity = [2, 4, 2];
1275        let xyz = [[0.17, 0.23, 0.31], [0.37, 0.11, 0.19]];
1276        let occupancy = [0.82, 0.55];
1277        let u_iso = [0.012, 0.018];
1278        let scale = 1.4;
1279        let zeros = [0.0; 3];
1280        let ones = [1.0; 3];
1281        let contributions = CwContributionsView::new(
1282            hkl.len(),
1283            0,
1284            CwContributionArrays {
1285                gaussian_variance_deg2: &zeros,
1286                lorentzian_fwhm_deg: &zeros,
1287                intensity_multiplier: &ones,
1288                d_gaussian_variance_d_position: &zeros,
1289                d_lorentzian_fwhm_d_position: &zeros,
1290                d_intensity_multiplier_d_position: &zeros,
1291                d_gaussian_variance_d_parameters: &[],
1292                d_lorentzian_fwhm_d_parameters: &[],
1293                d_intensity_multiplier_d_parameters: &[],
1294            },
1295        )
1296        .expect("contributions");
1297        let input = StructuralPatternInputView {
1298            x_deg: &x,
1299            hkl: &hkl,
1300            multiplicity: &multiplicity,
1301            fractional_xyz: &xyz,
1302            occupancy: &occupancy,
1303            u_iso_angstrom2: &u_iso,
1304            anisotropic_mask: &[false, false],
1305            u_aniso_cif_angstrom2: &[[0.0; 6]; 2],
1306            scattering_species: &["Si", "O"],
1307            scattering_real_offset: &[],
1308            scattering_imag_offset: &[],
1309            scale,
1310            coordinate_tolerance: 1.0e-10,
1311            instrument: instrument(),
1312            axial_geometry: None,
1313            position_correction: MonochromaticPositionCorrection {
1314                zero_shift_deg: 0.0,
1315                bragg_brentano_mm: None,
1316                debye_scherrer_micrometre: None,
1317            },
1318            correction_model: IntegratedIntensityCorrectionModel::Neutral,
1319            scattering_model: BuiltInScatteringModel::XrayNonResonant,
1320            contributions,
1321            support: SupportPolicy::FwhmMultiple(20.0),
1322        };
1323        let layout = P1ParameterLayout { site_count: 2 };
1324        let tangent: Vec<f64> = (0..layout.parameter_count())
1325            .map(|index| f64::from(u32::try_from(index + 1).expect("small index")) * 2.0e-5)
1326            .collect();
1327        let jvp = calculate_structural_pattern_jvp(cell(), &group(), &input, &tangent)
1328            .expect("structural JVP");
1329        let dense = calculate_structural_pattern_dense(cell(), &group(), &input)
1330            .expect("structural dense linearization");
1331        assert_eq!(dense.parameter_count, tangent.len());
1332        for sample in 0..x.len() {
1333            let product = tangent
1334                .iter()
1335                .enumerate()
1336                .map(|(parameter, direction)| direction * dense.d_y[parameter * x.len() + sample])
1337                .sum::<f64>();
1338            assert!((product - jvp.d_y[sample]).abs() < 2.0e-11 * product.abs().max(1.0));
1339        }
1340        let step = 1.0e-5;
1341        let mut plus_cell = cell();
1342        let mut minus_cell = cell();
1343        for (parameter, direction) in tangent
1344            .iter()
1345            .copied()
1346            .take(CELL_PARAMETER_COUNT)
1347            .enumerate()
1348        {
1349            perturb_cell(&mut plus_cell, parameter, step * direction);
1350            perturb_cell(&mut minus_cell, parameter, -step * direction);
1351        }
1352        let mut plus_xyz = xyz;
1353        let mut minus_xyz = xyz;
1354        for site in 0..2 {
1355            for component in 0..3 {
1356                let direction = tangent[layout.coordinate(site, component)];
1357                plus_xyz[site][component] += step * direction;
1358                minus_xyz[site][component] -= step * direction;
1359            }
1360        }
1361        let mut plus_occupancy = occupancy;
1362        let mut minus_occupancy = occupancy;
1363        let mut plus_u = u_iso;
1364        let mut minus_u = u_iso;
1365        for site in 0..2 {
1366            plus_occupancy[site] += step * tangent[layout.occupancy(site)];
1367            minus_occupancy[site] -= step * tangent[layout.occupancy(site)];
1368            plus_u[site] += step * tangent[layout.u_iso(site)];
1369            minus_u[site] -= step * tangent[layout.u_iso(site)];
1370        }
1371        let plus = calculate_case(
1372            plus_cell,
1373            &x,
1374            &hkl,
1375            &multiplicity,
1376            &plus_xyz,
1377            &plus_occupancy,
1378            &plus_u,
1379            scale + step * tangent[layout.scale()],
1380            &ones,
1381        );
1382        let minus = calculate_case(
1383            minus_cell,
1384            &x,
1385            &hkl,
1386            &multiplicity,
1387            &minus_xyz,
1388            &minus_occupancy,
1389            &minus_u,
1390            scale - step * tangent[layout.scale()],
1391            &ones,
1392        );
1393        for ((actual, plus_value), minus_value) in jvp
1394            .d_y
1395            .iter()
1396            .zip(&plus.accumulation.y)
1397            .zip(&minus.accumulation.y)
1398        {
1399            let finite_difference = (plus_value - minus_value) / (2.0 * step);
1400            assert!((actual - finite_difference).abs() < 3.0e-5 * finite_difference.abs().max(1.0));
1401        }
1402        let sample_weights: Vec<f64> = x.iter().map(|value| (0.17 * value).sin()).collect();
1403        let vjp = calculate_structural_pattern_vjp(cell(), &group(), &input, &sample_weights)
1404            .expect("structural VJP");
1405        let forward = jvp
1406            .d_y
1407            .iter()
1408            .zip(&sample_weights)
1409            .map(|(derivative, weight)| derivative * weight)
1410            .sum::<f64>();
1411        let reverse = tangent
1412            .iter()
1413            .zip(&vjp.gradient)
1414            .map(|(direction, gradient)| direction * gradient)
1415            .sum::<f64>();
1416        assert!((forward - reverse).abs() < 2.0e-10 * forward.abs().max(1.0));
1417        for (parameter, actual) in vjp.gradient.iter().copied().enumerate() {
1418            let expected = dense.d_y[parameter * x.len()..(parameter + 1) * x.len()]
1419                .iter()
1420                .zip(&sample_weights)
1421                .map(|(derivative, weight)| derivative * weight)
1422                .sum::<f64>();
1423            assert!((actual - expected).abs() < 2.0e-10 * expected.abs().max(1.0));
1424        }
1425    }
1426
1427    fn perturb_cell(cell: &mut UnitCell, parameter: usize, change: f64) {
1428        let value = match parameter {
1429            0 => &mut cell.a_angstrom,
1430            1 => &mut cell.b_angstrom,
1431            2 => &mut cell.c_angstrom,
1432            3 => &mut cell.alpha_deg,
1433            4 => &mut cell.beta_deg,
1434            5 => &mut cell.gamma_deg,
1435            _ => panic!("invalid cell parameter"),
1436        };
1437        *value += change;
1438    }
1439}