phasesmith_engine/
prepared_structural_phase.rs

1//! Owned structural-phase preparation for application-neutral callers.
2
3use phasesmith_core::{
4    ConstantWavelengthInstrument, CwContributionsView, FcjGeometry, SupportPolicy,
5};
6use phasesmith_crystallography::{
7    IntegratedIntensityCorrectionModel, PreparedNeutronScattering, PreparedXrayScattering,
8    SpaceGroup, UnitCell,
9};
10use phasesmith_execution::ExecutionContext;
11
12use crate::structural_pattern::{
13    BuiltInScatteringModel, MonochromaticPositionCorrection, StructuralPatternDenseResult,
14    StructuralPatternError, StructuralPatternInputView, StructuralPatternJvpResult,
15    StructuralPatternResult, StructuralPatternVjpResult,
16    calculate_structural_pattern_dense_with_context, calculate_structural_pattern_jvp_with_context,
17    calculate_structural_pattern_vjp_with_context, calculate_structural_pattern_with_context,
18};
19
20/// Owned crystallographic and scattering data for one structural phase.
21#[derive(Clone, Debug, PartialEq)]
22pub struct StructuralPhaseDefinition {
23    /// Direct unit cell.
24    pub cell: UnitCell,
25    /// Validated exact symmetry group.
26    pub space_group: SpaceGroup,
27    /// Canonical Miller indices.
28    pub hkl: Vec<[i32; 3]>,
29    /// Powder multiplicity for every reflection.
30    pub multiplicity: Vec<usize>,
31    /// Asymmetric-unit fractional coordinates.
32    pub fractional_xyz: Vec<[f64; 3]>,
33    /// Asymmetric-site occupancies.
34    pub occupancy: Vec<f64>,
35    /// Asymmetric-site isotropic displacement in square ångströms.
36    pub u_iso_angstrom2: Vec<f64>,
37    /// True for asymmetric sites described by fixed CIF U tensors.
38    pub anisotropic_mask: Vec<bool>,
39    /// CIF U tensors in component order `11,22,33,23,13,12`.
40    pub u_aniso_cif_angstrom2: Vec<[f64; 6]>,
41    /// Exact built-in scattering-table key for every asymmetric site.
42    pub scattering_species: Vec<String>,
43    /// Fixed real X-ray dispersion offset for every site, or empty when absent.
44    pub scattering_real_offset: Vec<f64>,
45    /// Fixed imaginary X-ray dispersion offset for every site, or empty when absent.
46    pub scattering_imag_offset: Vec<f64>,
47    /// Structural phase scale.
48    pub scale: f64,
49    /// Fixed symmetry-expansion deduplication tolerance.
50    pub coordinate_tolerance: f64,
51    /// Built-in native scattering selection.
52    pub scattering_model: BuiltInScatteringModel,
53    /// Integrated-intensity correction selection.
54    pub correction_model: IntegratedIntensityCorrectionModel,
55}
56
57impl StructuralPhaseDefinition {
58    /// Validate owned reflection, site, scattering, and offset data.
59    ///
60    /// # Errors
61    ///
62    /// Returns [`StructuralPatternError`] when array shapes disagree, offsets
63    /// are invalid, or a built-in scattering-table key is unknown.
64    pub fn validate(&self) -> Result<(), StructuralPatternError> {
65        validate_definition(self)
66    }
67}
68
69/// Borrowed experiment and sample-physics data for evaluating a prepared phase.
70#[derive(Clone, Copy, Debug)]
71pub struct PreparedStructuralPatternInputView<'a> {
72    /// Sorted pattern grid in degrees `2theta`.
73    pub x_deg: &'a [f64],
74    /// Monochromatic constant-wavelength instrument parameters.
75    pub instrument: ConstantWavelengthInstrument,
76    /// Optional Finger--Cox--Jephcoat axial-divergence geometry.
77    pub axial_geometry: Option<FcjGeometry>,
78    /// Explicit zero/sample-displacement position correction.
79    pub position_correction: MonochromaticPositionCorrection,
80    /// Vectorized sample-physics contribution batch.
81    pub contributions: CwContributionsView<'a>,
82    /// Exact finite profile-support policy.
83    pub support: SupportPolicy,
84}
85
86/// Reusable, application-neutral structural phase with an owned worker budget.
87#[derive(Clone)]
88pub struct PreparedStructuralPhase {
89    definition: StructuralPhaseDefinition,
90    execution: ExecutionContext,
91}
92
93impl PreparedStructuralPhase {
94    /// Validate and take ownership of one structural phase.
95    ///
96    /// # Errors
97    ///
98    /// Returns [`StructuralPatternError`] when reflection/site arrays disagree,
99    /// fixed offsets are invalid, or a scattering-table key is unknown.
100    pub fn new(
101        definition: StructuralPhaseDefinition,
102        execution: ExecutionContext,
103    ) -> Result<Self, StructuralPatternError> {
104        definition.validate()?;
105        Ok(Self {
106            definition,
107            execution,
108        })
109    }
110
111    /// Return the number of prepared reflections.
112    #[must_use]
113    pub fn reflection_count(&self) -> usize {
114        self.definition.hkl.len()
115    }
116
117    /// Return the structural parameter count in the native derivative layout.
118    #[must_use]
119    pub fn structural_parameter_count(&self) -> usize {
120        6 + 5 * self.definition.fractional_xyz.len() + 1
121    }
122
123    /// Return the exact worker budget owned by this prepared phase.
124    #[must_use]
125    pub fn execution_threads(&self) -> usize {
126        self.execution.threads()
127    }
128
129    /// Borrow the validated owned structural definition.
130    #[must_use]
131    pub const fn definition(&self) -> &StructuralPhaseDefinition {
132        &self.definition
133    }
134
135    /// Calculate values for the prepared phase.
136    ///
137    /// # Errors
138    ///
139    /// Returns [`StructuralPatternError`] for invalid dynamic or structural
140    /// inputs.
141    pub fn calculate(
142        &self,
143        input: &PreparedStructuralPatternInputView<'_>,
144    ) -> Result<StructuralPatternResult, StructuralPatternError> {
145        self.with_input(input, calculate_structural_pattern_with_context)
146    }
147
148    /// Calculate values and a dense structural-pattern linearization.
149    ///
150    /// # Errors
151    ///
152    /// Returns [`StructuralPatternError`] for invalid dynamic or structural
153    /// inputs.
154    pub fn linearize(
155        &self,
156        input: &PreparedStructuralPatternInputView<'_>,
157    ) -> Result<StructuralPatternDenseResult, StructuralPatternError> {
158        self.with_input(input, calculate_structural_pattern_dense_with_context)
159    }
160
161    /// Calculate values and a structural forward derivative product.
162    ///
163    /// # Errors
164    ///
165    /// Returns [`StructuralPatternError`] for invalid inputs or tangent shape.
166    pub fn jvp(
167        &self,
168        input: &PreparedStructuralPatternInputView<'_>,
169        tangent: &[f64],
170    ) -> Result<StructuralPatternJvpResult, StructuralPatternError> {
171        self.with_input(input, |cell, group, structural_input, execution| {
172            calculate_structural_pattern_jvp_with_context(
173                cell,
174                group,
175                structural_input,
176                tangent,
177                execution,
178            )
179        })
180    }
181
182    /// Calculate values and a pattern-Jacobian transpose product.
183    ///
184    /// # Errors
185    ///
186    /// Returns [`StructuralPatternError`] for invalid inputs or sample weights.
187    pub fn vjp(
188        &self,
189        input: &PreparedStructuralPatternInputView<'_>,
190        sample_weights: &[f64],
191    ) -> Result<StructuralPatternVjpResult, StructuralPatternError> {
192        self.with_input(input, |cell, group, structural_input, execution| {
193            calculate_structural_pattern_vjp_with_context(
194                cell,
195                group,
196                structural_input,
197                sample_weights,
198                execution,
199            )
200        })
201    }
202
203    fn with_input<R>(
204        &self,
205        input: &PreparedStructuralPatternInputView<'_>,
206        operation: impl FnOnce(
207            UnitCell,
208            &SpaceGroup,
209            &StructuralPatternInputView<'_>,
210            &ExecutionContext,
211        ) -> Result<R, StructuralPatternError>,
212    ) -> Result<R, StructuralPatternError> {
213        let definition = &self.definition;
214        let species = definition
215            .scattering_species
216            .iter()
217            .map(String::as_str)
218            .collect::<Vec<_>>();
219        let structural_input = StructuralPatternInputView {
220            x_deg: input.x_deg,
221            hkl: &definition.hkl,
222            multiplicity: &definition.multiplicity,
223            fractional_xyz: &definition.fractional_xyz,
224            occupancy: &definition.occupancy,
225            u_iso_angstrom2: &definition.u_iso_angstrom2,
226            anisotropic_mask: &definition.anisotropic_mask,
227            u_aniso_cif_angstrom2: &definition.u_aniso_cif_angstrom2,
228            scattering_species: &species,
229            scattering_real_offset: &definition.scattering_real_offset,
230            scattering_imag_offset: &definition.scattering_imag_offset,
231            scale: definition.scale,
232            coordinate_tolerance: definition.coordinate_tolerance,
233            instrument: input.instrument,
234            axial_geometry: input.axial_geometry,
235            position_correction: input.position_correction,
236            correction_model: correction_for_wavelength(
237                definition.correction_model,
238                input.instrument.wavelength_angstrom,
239            ),
240            scattering_model: definition.scattering_model,
241            contributions: input.contributions,
242            support: input.support,
243        };
244        operation(
245            definition.cell,
246            &definition.space_group,
247            &structural_input,
248            &self.execution,
249        )
250    }
251}
252
253fn validate_definition(
254    definition: &StructuralPhaseDefinition,
255) -> Result<(), StructuralPatternError> {
256    definition
257        .cell
258        .geometry()
259        .map(|_| ())
260        .map_err(StructuralPatternError::InvalidCell)?;
261    if definition.hkl.len() != definition.multiplicity.len() {
262        return Err(StructuralPatternError::ReflectionLengthMismatch);
263    }
264    let site_count = definition.fractional_xyz.len();
265    if site_count != definition.occupancy.len()
266        || site_count != definition.u_iso_angstrom2.len()
267        || site_count != definition.anisotropic_mask.len()
268        || site_count != definition.u_aniso_cif_angstrom2.len()
269        || site_count != definition.scattering_species.len()
270        || (!definition.scattering_real_offset.is_empty()
271            && site_count != definition.scattering_real_offset.len())
272        || (!definition.scattering_imag_offset.is_empty()
273            && site_count != definition.scattering_imag_offset.len())
274        || definition.scattering_real_offset.is_empty()
275            != definition.scattering_imag_offset.is_empty()
276    {
277        return Err(StructuralPatternError::SiteLengthMismatch);
278    }
279    if definition
280        .scattering_real_offset
281        .iter()
282        .chain(&definition.scattering_imag_offset)
283        .any(|value| !value.is_finite())
284    {
285        return Err(StructuralPatternError::NonFiniteScatteringOffset);
286    }
287    match definition.scattering_model {
288        BuiltInScatteringModel::XrayNonResonant => {
289            PreparedXrayScattering::new(definition.scattering_species.iter().map(String::as_str))
290                .map(|_| ())
291        }
292        BuiltInScatteringModel::NeutronNuclear => {
293            PreparedNeutronScattering::new(definition.scattering_species.iter().map(String::as_str))
294                .map(|_| ())
295        }
296    }
297    .map_err(StructuralPatternError::Scattering)
298}
299
300fn correction_for_wavelength(
301    correction: IntegratedIntensityCorrectionModel,
302    wavelength_angstrom: f64,
303) -> IntegratedIntensityCorrectionModel {
304    match correction {
305        IntegratedIntensityCorrectionModel::Neutral => IntegratedIntensityCorrectionModel::Neutral,
306        IntegratedIntensityCorrectionModel::BraggBrentanoUnpolarizedLp { .. } => {
307            IntegratedIntensityCorrectionModel::BraggBrentanoUnpolarizedLp {
308                wavelength_angstrom,
309            }
310        }
311        IntegratedIntensityCorrectionModel::BraggBrentanoPolarizedLp { polarization, .. } => {
312            IntegratedIntensityCorrectionModel::BraggBrentanoPolarizedLp {
313                wavelength_angstrom,
314                polarization,
315            }
316        }
317        IntegratedIntensityCorrectionModel::ConstantWavelengthNeutronLorentz { .. } => {
318            IntegratedIntensityCorrectionModel::ConstantWavelengthNeutronLorentz {
319                wavelength_angstrom,
320            }
321        }
322        IntegratedIntensityCorrectionModel::TimeOfFlightNeutronLorentz { two_theta_deg } => {
323            IntegratedIntensityCorrectionModel::TimeOfFlightNeutronLorentz { two_theta_deg }
324        }
325    }
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331    use phasesmith_crystallography::SymmetryOperation;
332
333    fn definition() -> StructuralPhaseDefinition {
334        StructuralPhaseDefinition {
335            cell: UnitCell {
336                a_angstrom: 5.0,
337                b_angstrom: 5.0,
338                c_angstrom: 5.0,
339                alpha_deg: 90.0,
340                beta_deg: 90.0,
341                gamma_deg: 90.0,
342            },
343            space_group: SpaceGroup::new(vec![SymmetryOperation::identity()]).expect("P1"),
344            hkl: vec![[1, 0, 0]],
345            multiplicity: vec![2],
346            fractional_xyz: vec![[0.0, 0.0, 0.0]],
347            occupancy: vec![1.0],
348            u_iso_angstrom2: vec![0.01],
349            anisotropic_mask: vec![false],
350            u_aniso_cif_angstrom2: vec![[0.0; 6]],
351            scattering_species: vec!["Si".to_owned()],
352            scattering_real_offset: Vec::new(),
353            scattering_imag_offset: Vec::new(),
354            scale: 1.0,
355            coordinate_tolerance: 1.0e-10,
356            scattering_model: BuiltInScatteringModel::XrayNonResonant,
357            correction_model: IntegratedIntensityCorrectionModel::Neutral,
358        }
359    }
360
361    #[test]
362    fn prepared_phase_owns_validated_data_and_execution() {
363        let phase = PreparedStructuralPhase::new(
364            definition(),
365            ExecutionContext::new(2).expect("execution context"),
366        )
367        .expect("prepared phase");
368        assert_eq!(phase.reflection_count(), 1);
369        assert_eq!(phase.structural_parameter_count(), 12);
370        assert_eq!(phase.execution_threads(), 2);
371    }
372
373    #[test]
374    fn prepared_phase_rejects_inconsistent_owned_shapes() {
375        let mut invalid_reflections = definition();
376        invalid_reflections.multiplicity.clear();
377        assert!(matches!(
378            PreparedStructuralPhase::new(invalid_reflections, ExecutionContext::serial()),
379            Err(StructuralPatternError::ReflectionLengthMismatch)
380        ));
381
382        let mut invalid_sites = definition();
383        invalid_sites.occupancy.clear();
384        assert!(matches!(
385            PreparedStructuralPhase::new(invalid_sites, ExecutionContext::serial()),
386            Err(StructuralPatternError::SiteLengthMismatch)
387        ));
388    }
389
390    #[test]
391    fn prepared_phase_rejects_unknown_scattering_species() {
392        let mut invalid = definition();
393        invalid.scattering_species[0] = "not-an-element".to_owned();
394        assert!(matches!(
395            PreparedStructuralPhase::new(invalid, ExecutionContext::serial()),
396            Err(StructuralPatternError::Scattering(_))
397        ));
398    }
399
400    #[test]
401    fn dynamic_wavelength_replaces_the_stored_correction_wavelength() {
402        let correction = correction_for_wavelength(
403            IntegratedIntensityCorrectionModel::BraggBrentanoPolarizedLp {
404                wavelength_angstrom: 0.5,
405                polarization: 0.7,
406            },
407            1.5406,
408        );
409        assert_eq!(
410            correction,
411            IntegratedIntensityCorrectionModel::BraggBrentanoPolarizedLp {
412                wavelength_angstrom: 1.5406,
413                polarization: 0.7,
414            }
415        );
416    }
417}
phasesmith_engine/prepared_structural_phase.rs

phasesmith_engine/
prepared_structural_phase.rs
