phasesmith_engine/
structural_multiphase.rs

1//! Native scheduling and ordered composition for multiple structural phases.
2
3use std::error::Error;
4use std::fmt::{Display, Formatter};
5use std::sync::Arc;
6
7use phasesmith_core::{
8    ConstantWavelengthInstrument, CwContributionsView, FcjGeometry, OwnedCwContributions,
9    SupportPolicy,
10};
11use phasesmith_execution::ExecutionPolicy;
12
13use crate::{
14    MonochromaticPositionCorrection, PreparedStructuralPatternInputView, PreparedStructuralPhase,
15    PreparedStructuralSpectrum, PreparedStructuralSpectrumInputView, StructuralPatternDenseResult,
16    StructuralPatternError, StructuralPatternJvpResult, StructuralPatternResult,
17    StructuralPatternVjpResult, StructuralSpectrumError,
18};
19
20/// One native built-in structural model, monochromatic or fixed-spectrum.
21#[derive(Clone)]
22pub enum PreparedStructuralModel {
23    /// One monochromatic structural phase.
24    Monochromatic(Arc<PreparedStructuralPhase>),
25    /// One fixed-wavelength structural spectrum.
26    FixedSpectrum(Arc<PreparedStructuralSpectrum>),
27}
28
29impl PreparedStructuralModel {
30    /// Wrap one prepared monochromatic phase.
31    #[must_use]
32    pub fn monochromatic(phase: PreparedStructuralPhase) -> Self {
33        Self::Monochromatic(Arc::new(phase))
34    }
35
36    /// Wrap one prepared fixed-wavelength spectrum.
37    #[must_use]
38    pub fn fixed_spectrum(spectrum: PreparedStructuralSpectrum) -> Self {
39        Self::FixedSpectrum(Arc::new(spectrum))
40    }
41
42    /// Return the structural parameter count for this phase model.
43    #[must_use]
44    pub fn structural_parameter_count(&self) -> usize {
45        match self {
46            Self::Monochromatic(phase) => phase.structural_parameter_count(),
47            Self::FixedSpectrum(spectrum) => spectrum.structural_parameter_count(),
48        }
49    }
50
51    fn calculate(
52        &self,
53        input: &PreparedStructuralModelInputView<'_>,
54    ) -> Result<StructuralPatternResult, StructuralMultiphaseError> {
55        match self {
56            Self::Monochromatic(phase) => phase
57                .calculate(&input.monochromatic_input()?)
58                .map_err(StructuralMultiphaseError::Structural),
59            Self::FixedSpectrum(spectrum) => spectrum
60                .calculate(&input.spectrum_input())
61                .map_err(StructuralMultiphaseError::Spectrum),
62        }
63    }
64
65    fn linearize(
66        &self,
67        input: &PreparedStructuralModelInputView<'_>,
68    ) -> Result<StructuralPatternDenseResult, StructuralMultiphaseError> {
69        match self {
70            Self::Monochromatic(phase) => phase
71                .linearize(&input.monochromatic_input()?)
72                .map_err(StructuralMultiphaseError::Structural),
73            Self::FixedSpectrum(spectrum) => spectrum
74                .linearize(&input.spectrum_input())
75                .map_err(StructuralMultiphaseError::Spectrum),
76        }
77    }
78
79    fn jvp(
80        &self,
81        input: &PreparedStructuralModelInputView<'_>,
82        tangent: &[f64],
83    ) -> Result<StructuralPatternJvpResult, StructuralMultiphaseError> {
84        match self {
85            Self::Monochromatic(phase) => phase
86                .jvp(&input.monochromatic_input()?, tangent)
87                .map_err(StructuralMultiphaseError::Structural),
88            Self::FixedSpectrum(spectrum) => spectrum
89                .jvp(&input.spectrum_input(), tangent)
90                .map_err(StructuralMultiphaseError::Spectrum),
91        }
92    }
93
94    fn vjp(
95        &self,
96        input: &PreparedStructuralModelInputView<'_>,
97        sample_weights: &[f64],
98    ) -> Result<StructuralPatternVjpResult, StructuralMultiphaseError> {
99        match self {
100            Self::Monochromatic(phase) => phase
101                .vjp(&input.monochromatic_input()?, sample_weights)
102                .map_err(StructuralMultiphaseError::Structural),
103            Self::FixedSpectrum(spectrum) => spectrum
104                .vjp(&input.spectrum_input(), sample_weights)
105                .map_err(StructuralMultiphaseError::Spectrum),
106        }
107    }
108}
109
110/// Borrowed dynamic data for one phase model in a multiphase operation.
111#[derive(Clone, Copy, Debug)]
112pub struct PreparedStructuralModelInputView<'a> {
113    /// Sorted pattern grid in degrees `2theta`.
114    pub x_deg: &'a [f64],
115    /// Reference constant-wavelength instrument.
116    pub instrument: ConstantWavelengthInstrument,
117    /// Optional axial-divergence geometry.
118    pub axial_geometry: Option<FcjGeometry>,
119    /// Explicit position correction.
120    pub position_correction: MonochromaticPositionCorrection,
121    /// One contribution batch for monochromatic models, or one per component.
122    pub contributions: &'a [CwContributionsView<'a>],
123    /// Exact finite profile-support policy.
124    pub support: SupportPolicy,
125}
126
127/// Owned sample-physics inputs for one prepared structural model.
128#[derive(Clone, Debug, PartialEq)]
129pub struct StructuralModelInput {
130    /// One contribution batch for monochromatic models, or one per spectrum component.
131    pub contributions: Vec<OwnedCwContributions>,
132}
133
134/// Owned application-boundary request for one multiphase structural calculation.
135#[derive(Clone, Debug, PartialEq)]
136pub struct StructuralCalculationRequest {
137    /// Sorted pattern grid in degrees `2theta`.
138    pub x_deg: Vec<f64>,
139    /// Reference constant-wavelength instrument.
140    pub instrument: ConstantWavelengthInstrument,
141    /// Optional axial-divergence geometry.
142    pub axial_geometry: Option<FcjGeometry>,
143    /// Explicit position correction.
144    pub position_correction: MonochromaticPositionCorrection,
145    /// Dynamic inputs in prepared-model order.
146    pub phase_inputs: Vec<StructuralModelInput>,
147    /// Exact finite profile-support policy.
148    pub support: SupportPolicy,
149}
150
151impl<'a> PreparedStructuralModelInputView<'a> {
152    fn monochromatic_input(
153        &self,
154    ) -> Result<PreparedStructuralPatternInputView<'a>, StructuralMultiphaseError> {
155        if self.contributions.len() != 1 {
156            return Err(StructuralMultiphaseError::ContributionCountMismatch);
157        }
158        Ok(PreparedStructuralPatternInputView {
159            x_deg: self.x_deg,
160            instrument: self.instrument,
161            axial_geometry: self.axial_geometry,
162            position_correction: self.position_correction,
163            contributions: self.contributions[0],
164            support: self.support,
165        })
166    }
167
168    fn spectrum_input(&self) -> PreparedStructuralSpectrumInputView<'a> {
169        PreparedStructuralSpectrumInputView {
170            x_deg: self.x_deg,
171            instrument: self.instrument,
172            axial_geometry: self.axial_geometry,
173            position_correction: self.position_correction,
174            contributions: self.contributions,
175            support: self.support,
176        }
177    }
178}
179
180/// Ordered multiphase structural values and their combined profile.
181#[derive(Clone, Debug, PartialEq)]
182pub struct StructuralMultiphaseResult {
183    /// Sum of phase profile arrays in phase input order.
184    pub profile_y: Vec<f64>,
185    /// One complete diagnostic result per phase, in input order.
186    pub phases: Vec<StructuralPatternResult>,
187}
188
189/// Display-ready owned result from a structural calculation request.
190#[derive(Clone, Debug, PartialEq)]
191pub struct StructuralCalculationResult {
192    /// Pattern grid copied from the evaluated request.
193    pub x_deg: Vec<f64>,
194    /// Sum of all phase profiles in prepared-model order.
195    pub profile_y: Vec<f64>,
196    /// Complete diagnostic result for each phase in prepared-model order.
197    pub phases: Vec<StructuralPatternResult>,
198}
199
200/// Invalid native multiphase request.
201#[derive(Debug)]
202pub enum StructuralMultiphaseError {
203    /// No structural model was supplied.
204    EmptyPhases,
205    /// Dynamic phase inputs do not match the prepared phase count.
206    PhaseInputCountMismatch,
207    /// Tangent vectors do not match the prepared phase count.
208    TangentCountMismatch,
209    /// A monochromatic model did not receive exactly one contribution batch.
210    ContributionCountMismatch,
211    /// Phase profile sample counts differ.
212    SampleCountMismatch,
213    /// A monochromatic phase failed.
214    Structural(StructuralPatternError),
215    /// A fixed-spectrum phase failed.
216    Spectrum(StructuralSpectrumError),
217}
218
219impl Display for StructuralMultiphaseError {
220    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
221        match self {
222            Self::EmptyPhases => formatter.write_str("at least one structural phase is required"),
223            Self::PhaseInputCountMismatch => {
224                formatter.write_str("dynamic phase inputs must match the prepared phase count")
225            }
226            Self::TangentCountMismatch => {
227                formatter.write_str("structural tangents must match the prepared phase count")
228            }
229            Self::ContributionCountMismatch => {
230                formatter.write_str("a monochromatic phase requires exactly one contribution batch")
231            }
232            Self::SampleCountMismatch => {
233                formatter.write_str("structural phases must share one sample grid")
234            }
235            Self::Structural(error) => Display::fmt(error, formatter),
236            Self::Spectrum(error) => Display::fmt(error, formatter),
237        }
238    }
239}
240
241impl Error for StructuralMultiphaseError {
242    fn source(&self) -> Option<&(dyn Error + 'static)> {
243        match self {
244            Self::Structural(error) => Some(error),
245            Self::Spectrum(error) => Some(error),
246            Self::EmptyPhases
247            | Self::PhaseInputCountMismatch
248            | Self::TangentCountMismatch
249            | Self::ContributionCountMismatch
250            | Self::SampleCountMismatch => None,
251        }
252    }
253}
254
255/// Reusable native scheduler for ordered structural phase models.
256pub struct PreparedStructuralMultiphase {
257    models: Vec<PreparedStructuralModel>,
258    execution: ExecutionPolicy,
259}
260
261impl PreparedStructuralMultiphase {
262    /// Prepare a non-empty ordered phase collection.
263    ///
264    /// # Errors
265    ///
266    /// Returns [`StructuralMultiphaseError::EmptyPhases`] for an empty model list.
267    pub fn new(
268        models: Vec<PreparedStructuralModel>,
269        execution: ExecutionPolicy,
270    ) -> Result<Self, StructuralMultiphaseError> {
271        if models.is_empty() {
272            return Err(StructuralMultiphaseError::EmptyPhases);
273        }
274        Ok(Self { models, execution })
275    }
276
277    /// Return the number of prepared phases.
278    #[must_use]
279    pub fn phase_count(&self) -> usize {
280        self.models.len()
281    }
282
283    /// Return structural parameter counts in phase order.
284    #[must_use]
285    pub fn structural_parameter_counts(&self) -> Vec<usize> {
286        self.models
287            .iter()
288            .map(PreparedStructuralModel::structural_parameter_count)
289            .collect()
290    }
291
292    /// Calculate all phases and their ordered sum.
293    ///
294    /// # Errors
295    ///
296    /// Returns [`StructuralMultiphaseError`] for invalid inputs or phase failures.
297    pub fn calculate(
298        &self,
299        inputs: &[PreparedStructuralModelInputView<'_>],
300    ) -> Result<StructuralMultiphaseResult, StructuralMultiphaseError> {
301        let phases = self.map_models(inputs, PreparedStructuralModel::calculate)?;
302        let profile_y = sum_phase_profiles(phases.iter().map(|phase| &phase.accumulation.y))?;
303        Ok(StructuralMultiphaseResult { profile_y, phases })
304    }
305
306    /// Calculate from an entirely owned application-boundary request.
307    ///
308    /// This is the shared entry point for language and application adapters. The
309    /// request owns all dynamic arrays; kernel views exist only for the duration
310    /// of this call.
311    ///
312    /// # Errors
313    ///
314    /// Returns [`StructuralMultiphaseError`] for invalid inputs or phase failures.
315    pub fn calculate_request(
316        &self,
317        request: StructuralCalculationRequest,
318    ) -> Result<StructuralCalculationResult, StructuralMultiphaseError> {
319        let contribution_views = request
320            .phase_inputs
321            .iter()
322            .map(|model| {
323                model
324                    .contributions
325                    .iter()
326                    .map(OwnedCwContributions::as_view)
327                    .collect::<Vec<_>>()
328            })
329            .collect::<Vec<_>>();
330        let inputs = contribution_views
331            .iter()
332            .map(|contributions| PreparedStructuralModelInputView {
333                x_deg: &request.x_deg,
334                instrument: request.instrument,
335                axial_geometry: request.axial_geometry,
336                position_correction: request.position_correction,
337                contributions,
338                support: request.support,
339            })
340            .collect::<Vec<_>>();
341        let result = self.calculate(&inputs)?;
342        Ok(StructuralCalculationResult {
343            x_deg: request.x_deg,
344            profile_y: result.profile_y,
345            phases: result.phases,
346        })
347    }
348
349    /// Calculate one dense structural linearization per phase.
350    ///
351    /// # Errors
352    ///
353    /// Returns [`StructuralMultiphaseError`] for invalid inputs or phase failures.
354    pub fn linearize(
355        &self,
356        inputs: &[PreparedStructuralModelInputView<'_>],
357    ) -> Result<Vec<StructuralPatternDenseResult>, StructuralMultiphaseError> {
358        self.map_models(inputs, PreparedStructuralModel::linearize)
359    }
360
361    /// Calculate one structural forward product per phase.
362    ///
363    /// # Errors
364    ///
365    /// Returns [`StructuralMultiphaseError`] for invalid inputs or tangents.
366    pub fn jvp(
367        &self,
368        inputs: &[PreparedStructuralModelInputView<'_>],
369        tangents: &[&[f64]],
370    ) -> Result<Vec<StructuralPatternJvpResult>, StructuralMultiphaseError> {
371        if tangents.len() != self.phase_count() {
372            return Err(StructuralMultiphaseError::TangentCountMismatch);
373        }
374        self.map_models_indexed(inputs, |index, model, input| {
375            model.jvp(input, tangents[index])
376        })
377    }
378
379    /// Calculate one structural reverse product per phase.
380    ///
381    /// # Errors
382    ///
383    /// Returns [`StructuralMultiphaseError`] for invalid inputs or sample weights.
384    pub fn vjp(
385        &self,
386        inputs: &[PreparedStructuralModelInputView<'_>],
387        sample_weights: &[f64],
388    ) -> Result<Vec<StructuralPatternVjpResult>, StructuralMultiphaseError> {
389        self.map_models(inputs, |model, input| model.vjp(input, sample_weights))
390    }
391
392    fn map_models<R: Send>(
393        &self,
394        inputs: &[PreparedStructuralModelInputView<'_>],
395        operation: impl Fn(
396            &PreparedStructuralModel,
397            &PreparedStructuralModelInputView<'_>,
398        ) -> Result<R, StructuralMultiphaseError>
399        + Send
400        + Sync,
401    ) -> Result<Vec<R>, StructuralMultiphaseError> {
402        self.map_models_indexed(inputs, |_index, model, input| operation(model, input))
403    }
404
405    fn map_models_indexed<R: Send>(
406        &self,
407        inputs: &[PreparedStructuralModelInputView<'_>],
408        operation: impl Fn(
409            usize,
410            &PreparedStructuralModel,
411            &PreparedStructuralModelInputView<'_>,
412        ) -> Result<R, StructuralMultiphaseError>
413        + Send
414        + Sync,
415    ) -> Result<Vec<R>, StructuralMultiphaseError> {
416        if inputs.len() != self.phase_count() {
417            return Err(StructuralMultiphaseError::PhaseInputCountMismatch);
418        }
419        self.execution
420            .context()
421            .map_ordered(
422                self.phase_count(),
423                self.execution.minimum_parallel_tasks(),
424                |index| operation(index, &self.models[index], &inputs[index]),
425            )
426            .into_iter()
427            .collect()
428    }
429}
430
431fn sum_phase_profiles<'a>(
432    profiles: impl IntoIterator<Item = &'a Vec<f64>>,
433) -> Result<Vec<f64>, StructuralMultiphaseError> {
434    let mut iterator = profiles.into_iter();
435    let first = iterator
436        .next()
437        .ok_or(StructuralMultiphaseError::EmptyPhases)?;
438    let mut combined = vec![0.0; first.len()];
439    for profile in std::iter::once(first).chain(iterator) {
440        if profile.len() != combined.len() {
441            return Err(StructuralMultiphaseError::SampleCountMismatch);
442        }
443        for (total, value) in combined.iter_mut().zip(profile) {
444            *total += value;
445        }
446    }
447    Ok(combined)
448}
449
450#[cfg(test)]
451mod tests {
452    use super::*;
453    use crate::{BuiltInScatteringModel, StructuralPhaseDefinition};
454    use phasesmith_crystallography::{
455        IntegratedIntensityCorrectionModel, SpaceGroup, SymmetryOperation, UnitCell,
456    };
457
458    fn prepared_model() -> PreparedStructuralModel {
459        let definition = StructuralPhaseDefinition {
460            cell: UnitCell {
461                a_angstrom: 5.0,
462                b_angstrom: 5.0,
463                c_angstrom: 5.0,
464                alpha_deg: 90.0,
465                beta_deg: 90.0,
466                gamma_deg: 90.0,
467            },
468            space_group: SpaceGroup::new(vec![SymmetryOperation::identity()]).expect("P1"),
469            hkl: vec![[1, 0, 0]],
470            multiplicity: vec![2],
471            fractional_xyz: vec![[0.0, 0.0, 0.0]],
472            occupancy: vec![1.0],
473            u_iso_angstrom2: vec![0.01],
474            anisotropic_mask: vec![false],
475            u_aniso_cif_angstrom2: vec![[0.0; 6]],
476            scattering_species: vec!["Si".to_owned()],
477            scattering_real_offset: Vec::new(),
478            scattering_imag_offset: Vec::new(),
479            scale: 1.0,
480            coordinate_tolerance: 1.0e-10,
481            scattering_model: BuiltInScatteringModel::XrayNonResonant,
482            correction_model: IntegratedIntensityCorrectionModel::Neutral,
483        };
484        PreparedStructuralModel::monochromatic(
485            PreparedStructuralPhase::new(
486                definition,
487                phasesmith_execution::ExecutionContext::serial(),
488            )
489            .expect("prepared phase"),
490        )
491    }
492
493    #[test]
494    fn multiphase_requires_at_least_one_model() {
495        assert!(matches!(
496            PreparedStructuralMultiphase::new(
497                Vec::new(),
498                ExecutionPolicy::bounded_default().expect("policy"),
499            ),
500            Err(StructuralMultiphaseError::EmptyPhases)
501        ));
502    }
503
504    #[test]
505    fn profile_sum_preserves_phase_order_and_validates_sample_counts() {
506        let first = vec![1.0, 2.0, 3.0];
507        let second = vec![0.5, 0.25, 0.125];
508        assert_eq!(
509            sum_phase_profiles([&first, &second]).expect("sum"),
510            [1.5, 2.25, 3.125]
511        );
512        assert!(matches!(
513            sum_phase_profiles([&first, &vec![1.0]]),
514            Err(StructuralMultiphaseError::SampleCountMismatch)
515        ));
516    }
517
518    #[test]
519    fn owned_request_validates_phase_input_count() {
520        let prepared = PreparedStructuralMultiphase::new(
521            vec![prepared_model()],
522            ExecutionPolicy::bounded_default().expect("policy"),
523        )
524        .expect("multiphase");
525        let request = StructuralCalculationRequest {
526            x_deg: vec![20.0, 21.0],
527            instrument: ConstantWavelengthInstrument {
528                wavelength_angstrom: 1.5406,
529                u_deg2: 0.0,
530                v_deg2: 0.0,
531                w_deg2: 0.01,
532                x_deg: 0.0,
533                y_deg: 0.0,
534            },
535            axial_geometry: None,
536            position_correction: MonochromaticPositionCorrection {
537                zero_shift_deg: 0.0,
538                bragg_brentano_mm: None,
539                debye_scherrer_micrometre: None,
540            },
541            phase_inputs: Vec::new(),
542            support: SupportPolicy::FwhmMultiple(8.0),
543        };
544        assert!(matches!(
545            prepared.calculate_request(request),
546            Err(StructuralMultiphaseError::PhaseInputCountMismatch)
547        ));
548    }
549}
phasesmith_engine/structural_multiphase.rs

phasesmith_engine/
structural_multiphase.rs
