Skip to main content

phasesmith_engine/
structural_spectrum.rs

1//! Native fixed-wavelength structural spectrum composition.
2
3use std::error::Error;
4use std::fmt::{Display, Formatter};
5
6use phasesmith_core::{
7    Accumulation, ConstantWavelengthInstrument, CwContributionsView, DenseJacobian, FcjGeometry,
8    PatternDerivatives, SupportJacobian, SupportPolicy,
9};
10use phasesmith_crystallography::StructureFactorValues;
11use phasesmith_execution::ExecutionPolicy;
12
13use crate::structural_pattern::CW_INSTRUMENT_PARAMETER_COUNT;
14use crate::{
15    MonochromaticPositionCorrection, PreparedStructuralPatternInputView, PreparedStructuralPhase,
16    StructuralPatternDenseResult, StructuralPatternError, StructuralPatternJvpResult,
17    StructuralPatternResult, StructuralPatternVjpResult, StructuralPhaseDefinition,
18};
19
20const WAVELENGTH_GLOBAL_PARAMETER_INDEX: usize = CW_INSTRUMENT_PARAMETER_COUNT;
21
22/// Borrowed dynamic inputs for one prepared fixed-wavelength spectrum.
23#[derive(Clone, Copy, Debug)]
24pub struct PreparedStructuralSpectrumInputView<'a> {
25    /// Sorted pattern grid in degrees `2theta`.
26    pub x_deg: &'a [f64],
27    /// Reference instrument; each component replaces only its wavelength.
28    pub instrument: ConstantWavelengthInstrument,
29    /// Optional Finger--Cox--Jephcoat axial-divergence geometry.
30    pub axial_geometry: Option<FcjGeometry>,
31    /// Explicit zero/sample-displacement position correction.
32    pub position_correction: MonochromaticPositionCorrection,
33    /// One sample-physics contribution batch per wavelength component.
34    pub contributions: &'a [CwContributionsView<'a>],
35    /// Exact finite profile-support policy.
36    pub support: SupportPolicy,
37}
38
39/// Invalid fixed-wavelength structural spectrum request.
40#[derive(Debug)]
41pub enum StructuralSpectrumError {
42    /// Wavelength and relative-intensity arrays differ in length.
43    ComponentLengthMismatch,
44    /// No wavelength component was supplied.
45    EmptyComponents,
46    /// A wavelength is non-finite or non-positive.
47    InvalidWavelength,
48    /// A relative intensity is non-finite or negative, or the reference is zero.
49    InvalidRelativeIntensity,
50    /// Dynamic contribution batches do not match the prepared component count.
51    ContributionCountMismatch,
52    /// Component calculations do not have compatible result layouts.
53    IncompatibleComponentResult,
54    /// Component composition allocation arithmetic overflowed.
55    AllocationOverflow,
56    /// A monochromatic structural calculation failed.
57    Structural(StructuralPatternError),
58}
59
60impl Display for StructuralSpectrumError {
61    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
62        match self {
63            Self::ComponentLengthMismatch => formatter.write_str(
64                "wavelength and relative-intensity arrays must have the same component count",
65            ),
66            Self::EmptyComponents => {
67                formatter.write_str("at least one wavelength component is required")
68            }
69            Self::InvalidWavelength => {
70                formatter.write_str("component wavelengths must be positive and finite")
71            }
72            Self::InvalidRelativeIntensity => formatter.write_str(
73                "component relative intensities must be finite and non-negative, with a positive reference component",
74            ),
75            Self::ContributionCountMismatch => formatter.write_str(
76                "sample-physics contributions must match the wavelength component count",
77            ),
78            Self::IncompatibleComponentResult => {
79                formatter.write_str("wavelength component result layouts must match")
80            }
81            Self::AllocationOverflow => {
82                formatter.write_str("wavelength component composition allocation overflow")
83            }
84            Self::Structural(error) => Display::fmt(error, formatter),
85        }
86    }
87}
88
89impl Error for StructuralSpectrumError {
90    fn source(&self) -> Option<&(dyn Error + 'static)> {
91        match self {
92            Self::Structural(error) => Some(error),
93            Self::ComponentLengthMismatch
94            | Self::EmptyComponents
95            | Self::InvalidWavelength
96            | Self::InvalidRelativeIntensity
97            | Self::ContributionCountMismatch
98            | Self::IncompatibleComponentResult
99            | Self::AllocationOverflow => None,
100        }
101    }
102}
103
104/// Reusable native fixed-wavelength spectrum for one structural phase.
105#[derive(Clone)]
106pub struct PreparedStructuralSpectrum {
107    phases: Vec<PreparedStructuralPhase>,
108    wavelengths_angstrom: Vec<f64>,
109    normalized_weights: Vec<f64>,
110    execution: ExecutionPolicy,
111}
112
113impl PreparedStructuralSpectrum {
114    /// Prepare every wavelength component from one base structural phase.
115    ///
116    /// Relative intensities are normalized to unit sum and multiply the base
117    /// phase scale. All prepared components share the policy's persistent pool.
118    ///
119    /// # Errors
120    ///
121    /// Returns [`StructuralSpectrumError`] for invalid component arrays or an
122    /// invalid structural phase.
123    pub fn new(
124        definition: &StructuralPhaseDefinition,
125        wavelengths_angstrom: Vec<f64>,
126        relative_intensities: &[f64],
127        execution: ExecutionPolicy,
128    ) -> Result<Self, StructuralSpectrumError> {
129        if wavelengths_angstrom.len() != relative_intensities.len() {
130            return Err(StructuralSpectrumError::ComponentLengthMismatch);
131        }
132        if wavelengths_angstrom.is_empty() {
133            return Err(StructuralSpectrumError::EmptyComponents);
134        }
135        if wavelengths_angstrom
136            .iter()
137            .any(|value| !value.is_finite() || *value <= 0.0)
138        {
139            return Err(StructuralSpectrumError::InvalidWavelength);
140        }
141        if relative_intensities[0] <= 0.0
142            || relative_intensities
143                .iter()
144                .any(|value| !value.is_finite() || *value < 0.0)
145        {
146            return Err(StructuralSpectrumError::InvalidRelativeIntensity);
147        }
148        let intensity_sum = relative_intensities.iter().sum::<f64>();
149        if !intensity_sum.is_finite() || intensity_sum <= 0.0 {
150            return Err(StructuralSpectrumError::InvalidRelativeIntensity);
151        }
152        let normalized_weights = relative_intensities
153            .iter()
154            .map(|value| value / intensity_sum)
155            .collect::<Vec<_>>();
156        let mut phases = Vec::with_capacity(wavelengths_angstrom.len());
157        for (&weight, &wavelength_angstrom) in normalized_weights.iter().zip(&wavelengths_angstrom)
158        {
159            let mut component = definition.clone();
160            component.scale *= weight;
161            component.correction_model = component
162                .correction_model
163                .with_wavelength(wavelength_angstrom);
164            phases.push(
165                PreparedStructuralPhase::new(component, execution.context().clone())
166                    .map_err(StructuralSpectrumError::Structural)?,
167            );
168        }
169        Ok(Self {
170            phases,
171            wavelengths_angstrom,
172            normalized_weights,
173            execution,
174        })
175    }
176
177    /// Return the number of prepared wavelength components.
178    #[must_use]
179    pub fn component_count(&self) -> usize {
180        self.phases.len()
181    }
182
183    /// Return normalized component weights in input order.
184    #[must_use]
185    pub fn normalized_weights(&self) -> &[f64] {
186        &self.normalized_weights
187    }
188
189    /// Return the structural parameter count shared by all components.
190    #[must_use]
191    pub fn structural_parameter_count(&self) -> usize {
192        self.phases[0].structural_parameter_count()
193    }
194
195    /// Calculate and combine all fixed wavelength components.
196    ///
197    /// # Errors
198    ///
199    /// Returns [`StructuralSpectrumError`] for invalid dynamic inputs or a
200    /// failed component calculation.
201    pub fn calculate(
202        &self,
203        input: &PreparedStructuralSpectrumInputView<'_>,
204    ) -> Result<StructuralPatternResult, StructuralSpectrumError> {
205        let results = self.map_components(input, |phase, component_input| {
206            phase.calculate(&component_input)
207        })?;
208        combine_pattern_results(results)
209    }
210
211    /// Calculate and combine dense structural linearizations.
212    ///
213    /// # Errors
214    ///
215    /// Returns [`StructuralSpectrumError`] for invalid inputs or incompatible
216    /// component layouts.
217    pub fn linearize(
218        &self,
219        input: &PreparedStructuralSpectrumInputView<'_>,
220    ) -> Result<StructuralPatternDenseResult, StructuralSpectrumError> {
221        let products = self.map_components(input, |phase, component_input| {
222            phase.linearize(&component_input)
223        })?;
224        combine_dense_products(products, &self.normalized_weights)
225    }
226
227    /// Calculate and combine structural forward derivative products.
228    ///
229    /// # Errors
230    ///
231    /// Returns [`StructuralSpectrumError`] for invalid inputs or tangent shape.
232    pub fn jvp(
233        &self,
234        input: &PreparedStructuralSpectrumInputView<'_>,
235        tangent: &[f64],
236    ) -> Result<StructuralPatternJvpResult, StructuralSpectrumError> {
237        let products =
238            self.map_components_indexed(input, |component, phase, component_input| {
239                let mut component_tangent = tangent.to_vec();
240                let scale_direction = component_tangent
241                    .last_mut()
242                    .ok_or(StructuralSpectrumError::IncompatibleComponentResult)?;
243                *scale_direction *= self.normalized_weights[component];
244                phase
245                    .jvp(&component_input, &component_tangent)
246                    .map_err(StructuralSpectrumError::Structural)
247            })?;
248        combine_jvp_products(products)
249    }
250
251    /// Calculate and combine structural transpose products.
252    ///
253    /// # Errors
254    ///
255    /// Returns [`StructuralSpectrumError`] for invalid inputs or sample weights.
256    pub fn vjp(
257        &self,
258        input: &PreparedStructuralSpectrumInputView<'_>,
259        sample_weights: &[f64],
260    ) -> Result<StructuralPatternVjpResult, StructuralSpectrumError> {
261        let products = self.map_components(input, |phase, component_input| {
262            phase.vjp(&component_input, sample_weights)
263        })?;
264        combine_vjp_products(products, &self.normalized_weights)
265    }
266
267    fn map_components<R: Send>(
268        &self,
269        input: &PreparedStructuralSpectrumInputView<'_>,
270        operation: impl Fn(
271            &PreparedStructuralPhase,
272            PreparedStructuralPatternInputView<'_>,
273        ) -> Result<R, StructuralPatternError>
274        + Send
275        + Sync,
276    ) -> Result<Vec<R>, StructuralSpectrumError> {
277        self.map_components_indexed(input, |_index, phase, component_input| {
278            operation(phase, component_input).map_err(StructuralSpectrumError::Structural)
279        })
280    }
281
282    fn map_components_indexed<R: Send>(
283        &self,
284        input: &PreparedStructuralSpectrumInputView<'_>,
285        operation: impl Fn(
286            usize,
287            &PreparedStructuralPhase,
288            PreparedStructuralPatternInputView<'_>,
289        ) -> Result<R, StructuralSpectrumError>
290        + Send
291        + Sync,
292    ) -> Result<Vec<R>, StructuralSpectrumError> {
293        if input.contributions.len() != self.component_count() {
294            return Err(StructuralSpectrumError::ContributionCountMismatch);
295        }
296        self.execution
297            .context()
298            .map_ordered(
299                self.component_count(),
300                self.execution.minimum_parallel_tasks(),
301                |component| {
302                    let mut instrument = input.instrument;
303                    instrument.wavelength_angstrom = self.wavelengths_angstrom[component];
304                    operation(
305                        component,
306                        &self.phases[component],
307                        PreparedStructuralPatternInputView {
308                            x_deg: input.x_deg,
309                            instrument,
310                            axial_geometry: input.axial_geometry,
311                            position_correction: input.position_correction,
312                            contributions: input.contributions[component],
313                            support: input.support,
314                        },
315                    )
316                },
317            )
318            .into_iter()
319            .collect()
320    }
321}
322
323fn combine_pattern_results(
324    results: Vec<StructuralPatternResult>,
325) -> Result<StructuralPatternResult, StructuralSpectrumError> {
326    let mut iterator = results.into_iter();
327    let first = iterator
328        .next()
329        .ok_or(StructuralSpectrumError::EmptyComponents)?;
330    let mut components = vec![first];
331    components.extend(iterator);
332    let accumulation = combine_accumulations(
333        components
334            .iter_mut()
335            .map(|result| std::mem::replace(&mut result.accumulation, empty_accumulation()))
336            .collect(),
337    )?;
338    let mut structure_factors = StructureFactorValues {
339        f_real: Vec::new(),
340        f_imag: Vec::new(),
341        f_squared: Vec::new(),
342        intensity: Vec::new(),
343        q_squared_inverse_angstrom2: Vec::new(),
344        s_inverse_angstrom: Vec::new(),
345    };
346    let mut d_spacing_angstrom = Vec::new();
347    let mut two_theta_deg = Vec::new();
348    for result in components {
349        structure_factors
350            .f_real
351            .extend(result.structure_factors.f_real);
352        structure_factors
353            .f_imag
354            .extend(result.structure_factors.f_imag);
355        structure_factors
356            .f_squared
357            .extend(result.structure_factors.f_squared);
358        structure_factors
359            .intensity
360            .extend(result.structure_factors.intensity);
361        structure_factors
362            .q_squared_inverse_angstrom2
363            .extend(result.structure_factors.q_squared_inverse_angstrom2);
364        structure_factors
365            .s_inverse_angstrom
366            .extend(result.structure_factors.s_inverse_angstrom);
367        d_spacing_angstrom.extend(result.d_spacing_angstrom);
368        two_theta_deg.extend(result.two_theta_deg);
369    }
370    Ok(StructuralPatternResult {
371        structure_factors,
372        d_spacing_angstrom,
373        two_theta_deg,
374        accumulation,
375    })
376}
377
378fn combine_accumulations(
379    accumulations: Vec<Accumulation>,
380) -> Result<Accumulation, StructuralSpectrumError> {
381    let first = accumulations
382        .first()
383        .ok_or(StructuralSpectrumError::EmptyComponents)?;
384    let sample_count = first.sample_count;
385    let local_parameter_count = first.derivatives.local.parameter_count;
386    let global_parameter_count = first
387        .derivatives
388        .global
389        .as_ref()
390        .ok_or(StructuralSpectrumError::IncompatibleComponentResult)?
391        .parameter_count;
392    if global_parameter_count <= WAVELENGTH_GLOBAL_PARAMETER_INDEX {
393        return Err(StructuralSpectrumError::IncompatibleComponentResult);
394    }
395    let combined_global_count = global_parameter_count - 1;
396    let global_length = combined_global_count
397        .checked_mul(sample_count)
398        .ok_or(StructuralSpectrumError::AllocationOverflow)?;
399    let mut y = vec![0.0; sample_count];
400    let mut starts = Vec::new();
401    let mut offsets: Vec<usize> = vec![0];
402    let mut local_values = Vec::new();
403    let mut global_values = vec![0.0; global_length];
404    for accumulation in accumulations {
405        let global = accumulation
406            .derivatives
407            .global
408            .ok_or(StructuralSpectrumError::IncompatibleComponentResult)?;
409        if accumulation.sample_count != sample_count
410            || accumulation.y.len() != sample_count
411            || accumulation.derivatives.local.parameter_count != local_parameter_count
412            || global.parameter_count != global_parameter_count
413            || global.sample_count != sample_count
414        {
415            return Err(StructuralSpectrumError::IncompatibleComponentResult);
416        }
417        for (combined, value) in y.iter_mut().zip(accumulation.y) {
418            *combined += value;
419        }
420        let cursor = offsets
421            .last()
422            .copied()
423            .ok_or(StructuralSpectrumError::IncompatibleComponentResult)?;
424        starts.extend(accumulation.derivatives.local.starts);
425        for offset in accumulation.derivatives.local.offsets.into_iter().skip(1) {
426            offsets.push(
427                cursor
428                    .checked_add(offset)
429                    .ok_or(StructuralSpectrumError::AllocationOverflow)?,
430            );
431        }
432        local_values.extend(accumulation.derivatives.local.values);
433        let mut combined_row = 0;
434        for row in 0..global_parameter_count {
435            if row == WAVELENGTH_GLOBAL_PARAMETER_INDEX {
436                continue;
437            }
438            let source = row * sample_count;
439            let target = combined_row * sample_count;
440            for sample in 0..sample_count {
441                global_values[target + sample] += global.values[source + sample];
442            }
443            combined_row += 1;
444        }
445    }
446    Ok(Accumulation {
447        y,
448        derivatives: PatternDerivatives {
449            local: SupportJacobian {
450                starts,
451                offsets,
452                values: local_values,
453                parameter_count: local_parameter_count,
454            },
455            global: Some(DenseJacobian {
456                values: global_values,
457                parameter_count: combined_global_count,
458                sample_count,
459            }),
460        },
461        sample_count,
462    })
463}
464
465fn combine_dense_products(
466    products: Vec<StructuralPatternDenseResult>,
467    weights: &[f64],
468) -> Result<StructuralPatternDenseResult, StructuralSpectrumError> {
469    let first = products
470        .first()
471        .ok_or(StructuralSpectrumError::EmptyComponents)?;
472    let parameter_count = first.parameter_count;
473    let sample_count = first.result.accumulation.sample_count;
474    if parameter_count == 0 || products.len() != weights.len() {
475        return Err(StructuralSpectrumError::IncompatibleComponentResult);
476    }
477    let element_count = parameter_count
478        .checked_mul(sample_count)
479        .ok_or(StructuralSpectrumError::AllocationOverflow)?;
480    let mut d_y = vec![0.0; element_count];
481    let mut results = Vec::with_capacity(products.len());
482    for (product, &weight) in products.into_iter().zip(weights) {
483        if product.parameter_count != parameter_count || product.d_y.len() != element_count {
484            return Err(StructuralSpectrumError::IncompatibleComponentResult);
485        }
486        for parameter in 0..parameter_count {
487            let factor = if parameter + 1 == parameter_count {
488                weight
489            } else {
490                1.0
491            };
492            let row = parameter * sample_count;
493            for sample in 0..sample_count {
494                d_y[row + sample] += factor * product.d_y[row + sample];
495            }
496        }
497        results.push(product.result);
498    }
499    Ok(StructuralPatternDenseResult {
500        result: combine_pattern_results(results)?,
501        d_y,
502        parameter_count,
503    })
504}
505
506fn combine_jvp_products(
507    products: Vec<StructuralPatternJvpResult>,
508) -> Result<StructuralPatternJvpResult, StructuralSpectrumError> {
509    let first = products
510        .first()
511        .ok_or(StructuralSpectrumError::EmptyComponents)?;
512    let sample_count = first.d_y.len();
513    let mut d_y = vec![0.0; sample_count];
514    let mut d_integrated_intensity = Vec::new();
515    let mut d_two_theta_deg = Vec::new();
516    let mut results = Vec::with_capacity(products.len());
517    for product in products {
518        if product.d_y.len() != sample_count {
519            return Err(StructuralSpectrumError::IncompatibleComponentResult);
520        }
521        for (combined, value) in d_y.iter_mut().zip(product.d_y) {
522            *combined += value;
523        }
524        d_integrated_intensity.extend(product.d_integrated_intensity);
525        d_two_theta_deg.extend(product.d_two_theta_deg);
526        results.push(product.result);
527    }
528    Ok(StructuralPatternJvpResult {
529        result: combine_pattern_results(results)?,
530        d_y,
531        d_integrated_intensity,
532        d_two_theta_deg,
533    })
534}
535
536fn combine_vjp_products(
537    products: Vec<StructuralPatternVjpResult>,
538    weights: &[f64],
539) -> Result<StructuralPatternVjpResult, StructuralSpectrumError> {
540    let first = products
541        .first()
542        .ok_or(StructuralSpectrumError::EmptyComponents)?;
543    let parameter_count = first.gradient.len();
544    if parameter_count == 0 || products.len() != weights.len() {
545        return Err(StructuralSpectrumError::IncompatibleComponentResult);
546    }
547    let mut gradient = vec![0.0; parameter_count];
548    let mut results = Vec::with_capacity(products.len());
549    for (product, &weight) in products.into_iter().zip(weights) {
550        if product.gradient.len() != parameter_count {
551            return Err(StructuralSpectrumError::IncompatibleComponentResult);
552        }
553        for (parameter, value) in product.gradient.into_iter().enumerate() {
554            let factor = if parameter + 1 == parameter_count {
555                weight
556            } else {
557                1.0
558            };
559            gradient[parameter] += factor * value;
560        }
561        results.push(product.result);
562    }
563    Ok(StructuralPatternVjpResult {
564        result: combine_pattern_results(results)?,
565        gradient,
566    })
567}
568
569fn empty_accumulation() -> Accumulation {
570    Accumulation {
571        y: Vec::new(),
572        derivatives: PatternDerivatives {
573            local: SupportJacobian {
574                starts: Vec::new(),
575                offsets: vec![0],
576                values: Vec::new(),
577                parameter_count: 0,
578            },
579            global: None,
580        },
581        sample_count: 0,
582    }
583}
584
585#[cfg(test)]
586mod tests {
587    use super::*;
588    use phasesmith_core::CwContributionArrays;
589    use phasesmith_crystallography::{
590        IntegratedIntensityCorrectionModel, Rational, SpaceGroup, SymmetryOperation, UnitCell,
591    };
592    fn definition() -> StructuralPhaseDefinition {
593        StructuralPhaseDefinition {
594            cell: UnitCell {
595                a_angstrom: 4.7,
596                b_angstrom: 5.1,
597                c_angstrom: 6.2,
598                alpha_deg: 82.0,
599                beta_deg: 87.0,
600                gamma_deg: 74.0,
601            },
602            space_group: SpaceGroup::new(vec![
603                SymmetryOperation::identity(),
604                SymmetryOperation::new([[-1, 0, 0], [0, -1, 0], [0, 0, -1]], [Rational::zero(); 3])
605                    .expect("inversion"),
606            ])
607            .expect("P-1"),
608            hkl: vec![[1, 0, 1], [2, 1, 1], [1, 2, 3]],
609            multiplicity: vec![2, 4, 2],
610            fractional_xyz: vec![[0.17, 0.23, 0.31], [0.37, 0.11, 0.19]],
611            occupancy: vec![0.82, 0.55],
612            u_iso_angstrom2: vec![0.012, 0.018],
613            anisotropic_mask: vec![false, false],
614            u_aniso_cif_angstrom2: vec![[0.0; 6]; 2],
615            scattering_species: vec!["Si".to_owned(), "O".to_owned()],
616            scattering_real_offset: Vec::new(),
617            scattering_imag_offset: Vec::new(),
618            scale: 1.4,
619            coordinate_tolerance: 1.0e-10,
620            scattering_model: crate::BuiltInScatteringModel::XrayNonResonant,
621            correction_model: IntegratedIntensityCorrectionModel::Neutral,
622        }
623    }
624
625    fn instrument() -> ConstantWavelengthInstrument {
626        ConstantWavelengthInstrument {
627            wavelength_angstrom: 1.5406,
628            u_deg2: 2.0e-4,
629            v_deg2: -1.0e-4,
630            w_deg2: 1.2e-4,
631            x_deg: 1.5e-3,
632            y_deg: 3.0e-3,
633        }
634    }
635
636    fn contribution<'a>(zeros: &'a [f64], ones: &'a [f64]) -> CwContributionsView<'a> {
637        CwContributionsView::new(
638            zeros.len(),
639            0,
640            CwContributionArrays {
641                gaussian_variance_deg2: zeros,
642                lorentzian_fwhm_deg: zeros,
643                intensity_multiplier: ones,
644                d_gaussian_variance_d_position: zeros,
645                d_lorentzian_fwhm_d_position: zeros,
646                d_intensity_multiplier_d_position: zeros,
647                d_gaussian_variance_d_parameters: &[],
648                d_lorentzian_fwhm_d_parameters: &[],
649                d_intensity_multiplier_d_parameters: &[],
650            },
651        )
652        .expect("contribution")
653    }
654
655    #[test]
656    fn spectrum_values_and_derivatives_follow_fixed_component_chain_rules() {
657        let policy = ExecutionPolicy::new(Some(2), 2).expect("policy");
658        let spectrum = PreparedStructuralSpectrum::new(
659            &definition(),
660            vec![1.5406, 1.54439],
661            &[1.0, 0.5],
662            policy,
663        )
664        .expect("spectrum");
665        assert_eq!(spectrum.component_count(), 2);
666        assert_eq!(spectrum.normalized_weights(), [2.0 / 3.0, 1.0 / 3.0]);
667
668        let x = (0..9_001)
669            .map(|index| 10.0 + f64::from(index) * 0.01)
670            .collect::<Vec<_>>();
671        let zeros = [0.0; 3];
672        let ones = [1.0; 3];
673        let contributions = [contribution(&zeros, &ones), contribution(&zeros, &ones)];
674        let input = PreparedStructuralSpectrumInputView {
675            x_deg: &x,
676            instrument: instrument(),
677            axial_geometry: None,
678            position_correction: MonochromaticPositionCorrection {
679                zero_shift_deg: 0.0,
680                bragg_brentano_mm: None,
681                debye_scherrer_micrometre: None,
682            },
683            contributions: &contributions,
684            support: SupportPolicy::FwhmMultiple(20.0),
685        };
686        let values = spectrum.calculate(&input).expect("values");
687        assert_eq!(values.structure_factors.intensity.len(), 6);
688        assert_eq!(values.d_spacing_angstrom.len(), 6);
689        assert_eq!(values.accumulation.derivatives.local.peak_count(), 6);
690        assert_eq!(
691            values
692                .accumulation
693                .derivatives
694                .global
695                .as_ref()
696                .expect("global")
697                .parameter_count,
698            6
699        );
700
701        let parameter_count = spectrum.phases[0].structural_parameter_count();
702        let tangent = (0..parameter_count)
703            .map(|index| f64::from(u32::try_from(index + 1).expect("small index")) * 2.0e-5)
704            .collect::<Vec<_>>();
705        let dense = spectrum.linearize(&input).expect("dense");
706        let jvp = spectrum.jvp(&input, &tangent).expect("JVP");
707        for sample in 0..x.len() {
708            let product = tangent
709                .iter()
710                .enumerate()
711                .map(|(parameter, direction)| direction * dense.d_y[parameter * x.len() + sample])
712                .sum::<f64>();
713            assert!((product - jvp.d_y[sample]).abs() < 3.0e-11 * product.abs().max(1.0));
714        }
715        let sample_weights = x
716            .iter()
717            .map(|value| (0.03 * value).sin())
718            .collect::<Vec<_>>();
719        let vjp = spectrum.vjp(&input, &sample_weights).expect("VJP");
720        for (parameter, gradient) in vjp.gradient.iter().enumerate() {
721            let expected = dense.d_y[parameter * x.len()..(parameter + 1) * x.len()]
722                .iter()
723                .zip(&sample_weights)
724                .map(|(derivative, weight)| derivative * weight)
725                .sum::<f64>();
726            assert!((expected - gradient).abs() < 3.0e-10 * expected.abs().max(1.0));
727        }
728    }
729
730    #[test]
731    fn spectrum_rebinds_intensity_correction_to_each_component_wavelength() {
732        let mut definition = definition();
733        definition.correction_model =
734            IntegratedIntensityCorrectionModel::BraggBrentanoPolarizedLp {
735                wavelength_angstrom: 1.0,
736                polarization: 0.73,
737            };
738        let spectrum = PreparedStructuralSpectrum::new(
739            &definition,
740            vec![1.5406, 1.54439],
741            &[1.0, 0.5],
742            ExecutionPolicy::bounded_default().expect("policy"),
743        )
744        .expect("spectrum");
745
746        for (phase, wavelength_angstrom) in spectrum.phases.iter().zip([1.5406_f64, 1.54439_f64]) {
747            assert_eq!(
748                phase.definition().correction_model,
749                IntegratedIntensityCorrectionModel::BraggBrentanoPolarizedLp {
750                    wavelength_angstrom,
751                    polarization: 0.73,
752                }
753            );
754        }
755    }
756
757    #[test]
758    fn spectrum_validates_component_and_contribution_counts() {
759        let definition = definition();
760        assert!(matches!(
761            PreparedStructuralSpectrum::new(
762                &definition,
763                vec![1.0],
764                &[],
765                ExecutionPolicy::bounded_default().expect("policy"),
766            ),
767            Err(StructuralSpectrumError::ComponentLengthMismatch)
768        ));
769        let spectrum = PreparedStructuralSpectrum::new(
770            &definition,
771            vec![1.0],
772            &[1.0],
773            ExecutionPolicy::bounded_default().expect("policy"),
774        )
775        .expect("spectrum");
776        let input = PreparedStructuralSpectrumInputView {
777            x_deg: &[10.0, 11.0],
778            instrument: instrument(),
779            axial_geometry: None,
780            position_correction: MonochromaticPositionCorrection {
781                zero_shift_deg: 0.0,
782                bragg_brentano_mm: None,
783                debye_scherrer_micrometre: None,
784            },
785            contributions: &[],
786            support: SupportPolicy::FwhmMultiple(20.0),
787        };
788        assert!(matches!(
789            spectrum.calculate(&input),
790            Err(StructuralSpectrumError::ContributionCountMismatch)
791        ));
792    }
793}