phasesmith_workflows/
rietveld_joint.rs

1//! Joint multi-histogram Rietveld parameter and matrix-free objective contracts.
2
3use std::collections::{BTreeMap, BTreeSet};
4use std::error::Error;
5use std::fmt::{Display, Formatter};
6
7use phasesmith_model::RecordId;
8
9use crate::{
10    LatticeBounds, ParameterError, ParameterKey, ParameterSet, ParameterSpec,
11    PreparedGeneralRietveldObjective, RietveldCalculationOptions, RietveldGeneralObjectiveError,
12    RietveldGeneralParameterError, RietveldInput, RietveldParameterLayout,
13    RietveldParameterSelection,
14};
15
16/// One experiment in a joint native Rietveld objective.
17#[derive(Clone, Debug, PartialEq)]
18pub struct JointRietveldHistogram {
19    /// Stable experiment identity used to namespace local parameters.
20    pub histogram_id: RecordId,
21    /// Complete observed pattern, experiment, and phase state.
22    pub input: RietveldInput,
23    /// Selected complete parameter families for this histogram.
24    pub selection: RietveldParameterSelection,
25    /// Per-phase lattice bounds in the input phase order.
26    pub lattice_bounds: Vec<Option<LatticeBounds>>,
27    /// Native calculation controls for this histogram.
28    pub calculation: RietveldCalculationOptions,
29}
30
31/// One histogram's profile and directional derivative.
32#[derive(Clone, Debug, PartialEq)]
33pub struct JointRietveldProduct {
34    /// Stable histogram identity.
35    pub histogram_id: RecordId,
36    /// Accepted-state calculated profile.
37    pub profile: Vec<f64>,
38    /// Profile directional derivative in joint physical coordinates.
39    pub derivative: Vec<f64>,
40}
41
42/// Complete accepted-state value and gradient of a joint objective.
43#[derive(Clone, Debug, PartialEq)]
44pub struct JointRietveldGradient {
45    /// Calculated profiles in histogram order.
46    pub calculated: Vec<Vec<f64>>,
47    /// Half the summed weighted squared residual over every histogram.
48    pub objective: f64,
49    /// Gradient in stable joint physical-parameter order.
50    pub gradient: Vec<f64>,
51}
52
53/// Stable shared/local parameter packing for a joint objective.
54#[derive(Clone, Debug, PartialEq)]
55pub struct JointRietveldLayout {
56    parameters: ParameterSet,
57    histogram_ids: Vec<RecordId>,
58    local_layouts: Vec<RietveldParameterLayout>,
59    local_to_joint: Vec<Vec<usize>>,
60}
61
62impl JointRietveldLayout {
63    /// Build one physical parameter set across two or more histograms.
64    ///
65    /// Lattice, coordinates, occupancies, and atomic displacement parameters
66    /// are shared by stable phase/site identity. Instrument, background,
67    /// sample-physics, and phase-scale parameters are histogram-local.
68    ///
69    /// # Errors
70    ///
71    /// Returns [`JointRietveldError`] for invalid histogram identity, input,
72    /// selection, or incompatible shared structural state.
73    pub fn new(histograms: &[JointRietveldHistogram]) -> Result<Self, JointRietveldError> {
74        validate_histograms(histograms)?;
75        let local_layouts = histograms
76            .iter()
77            .map(|histogram| {
78                RietveldParameterLayout::new(
79                    &histogram.input,
80                    &histogram.selection,
81                    &histogram.lattice_bounds,
82                )
83            })
84            .collect::<Result<Vec<_>, _>>()?;
85
86        let mut specs = Vec::new();
87        let mut shared = BTreeMap::<ParameterKey, usize>::new();
88        let mut local_to_joint = Vec::with_capacity(histograms.len());
89        for (histogram, layout) in histograms.iter().zip(&local_layouts) {
90            let mut mapping = Vec::with_capacity(layout.parameters().specs().len());
91            for spec in layout.parameters().specs() {
92                let joint_index = if is_shared(spec.key()) {
93                    if let Some(index) = shared.get(spec.key()).copied() {
94                        if specs[index] != *spec {
95                            return Err(JointRietveldError::SharedParameterMismatch {
96                                key: spec.key().clone(),
97                            });
98                        }
99                        index
100                    } else {
101                        let index = specs.len();
102                        specs.push(spec.clone());
103                        shared.insert(spec.key().clone(), index);
104                        index
105                    }
106                } else {
107                    let key = ParameterKey::new(
108                        spec.key().module(),
109                        format!("{}/{}", histogram.histogram_id, spec.key().owner_id()),
110                        spec.key().name(),
111                    )?;
112                    let index = specs.len();
113                    specs.push(ParameterSpec::new(
114                        key,
115                        spec.value(),
116                        spec.unit(),
117                        spec.bounds(),
118                        spec.scale(),
119                        spec.refine(),
120                    )?);
121                    index
122                };
123                mapping.push(joint_index);
124            }
125            local_to_joint.push(mapping);
126        }
127        Ok(Self {
128            parameters: ParameterSet::new(specs)?,
129            histogram_ids: histograms
130                .iter()
131                .map(|histogram| histogram.histogram_id.clone())
132                .collect(),
133            local_layouts,
134            local_to_joint,
135        })
136    }
137
138    /// Borrow the stable joint physical parameter set.
139    #[must_use]
140    pub const fn parameters(&self) -> &ParameterSet {
141        &self.parameters
142    }
143
144    /// Borrow stable histogram identities in packing order.
145    #[must_use]
146    pub fn histogram_ids(&self) -> &[RecordId] {
147        &self.histogram_ids
148    }
149
150    /// Install joint physical values into every histogram request.
151    ///
152    /// # Errors
153    ///
154    /// Returns [`JointRietveldError`] for a stale histogram contract or an
155    /// invalid joint/local value.
156    pub fn apply_values(
157        &self,
158        histograms: &[JointRietveldHistogram],
159        values: &[f64],
160    ) -> Result<Vec<JointRietveldHistogram>, JointRietveldError> {
161        self.validate_contract(histograms)?;
162        if values.len() != self.parameters.specs().len() {
163            return Err(JointRietveldError::ValueLengthMismatch);
164        }
165        histograms
166            .iter()
167            .enumerate()
168            .map(|(histogram_index, histogram)| {
169                let local_values = self.local_to_joint[histogram_index]
170                    .iter()
171                    .map(|index| values[*index])
172                    .collect::<Vec<_>>();
173                let mut updated = histogram.clone();
174                updated.input = self.local_layouts[histogram_index]
175                    .apply_values(&histogram.input, &local_values)?;
176                Ok(updated)
177            })
178            .collect()
179    }
180
181    fn validate_contract(
182        &self,
183        histograms: &[JointRietveldHistogram],
184    ) -> Result<(), JointRietveldError> {
185        if histograms.len() != self.histogram_ids.len()
186            || histograms
187                .iter()
188                .zip(&self.histogram_ids)
189                .any(|(histogram, expected)| histogram.histogram_id != *expected)
190        {
191            return Err(JointRietveldError::HistogramContractMismatch);
192        }
193        validate_histograms(histograms)?;
194        for ((histogram, expected_layout), expected_mapping) in histograms
195            .iter()
196            .zip(&self.local_layouts)
197            .zip(&self.local_to_joint)
198        {
199            let current_layout = RietveldParameterLayout::new(
200                &histogram.input,
201                &histogram.selection,
202                &histogram.lattice_bounds,
203            )?;
204            if current_layout != *expected_layout
205                || current_layout.parameters().specs().len() != expected_mapping.len()
206            {
207                return Err(JointRietveldError::HistogramContractMismatch);
208            }
209        }
210        Ok(())
211    }
212
213    fn local_direction(
214        &self,
215        histogram_index: usize,
216        direction: &[f64],
217    ) -> Result<Vec<f64>, JointRietveldError> {
218        if direction.len() != self.parameters.specs().len() {
219            return Err(JointRietveldError::ValueLengthMismatch);
220        }
221        Ok(self.local_to_joint[histogram_index]
222            .iter()
223            .map(|index| direction[*index])
224            .collect())
225    }
226
227    fn scatter_add(
228        &self,
229        histogram_index: usize,
230        local: &[f64],
231        joint: &mut [f64],
232    ) -> Result<(), JointRietveldError> {
233        if local.len() != self.local_to_joint[histogram_index].len() {
234            return Err(JointRietveldError::LocalProductLengthMismatch);
235        }
236        for (value, index) in local.iter().zip(&self.local_to_joint[histogram_index]) {
237            joint[*index] += value;
238        }
239        Ok(())
240    }
241}
242
243/// Prepared matrix-free sum of all histogram objectives.
244pub struct PreparedJointRietveldObjective {
245    histograms: Vec<JointRietveldHistogram>,
246    layout: JointRietveldLayout,
247    objectives: Vec<PreparedGeneralRietveldObjective>,
248}
249
250impl PreparedJointRietveldObjective {
251    /// Prepare all histogram products against one shared/local layout.
252    ///
253    /// # Errors
254    ///
255    /// Returns [`JointRietveldError`] for a stale layout or objective failure.
256    pub fn new(
257        histograms: Vec<JointRietveldHistogram>,
258        layout: JointRietveldLayout,
259    ) -> Result<Self, JointRietveldError> {
260        layout.validate_contract(&histograms)?;
261        let objectives = histograms
262            .iter()
263            .zip(&layout.local_layouts)
264            .map(|(histogram, local_layout)| {
265                PreparedGeneralRietveldObjective::new(
266                    histogram.input.clone(),
267                    histogram.calculation.clone(),
268                    local_layout.clone(),
269                )
270            })
271            .collect::<Result<Vec<_>, _>>()?;
272        Ok(Self {
273            histograms,
274            layout,
275            objectives,
276        })
277    }
278
279    /// Borrow the joint parameter layout.
280    #[must_use]
281    pub const fn layout(&self) -> &JointRietveldLayout {
282        &self.layout
283    }
284
285    /// Return expensive model products consumed while preparing all gradients.
286    #[must_use]
287    pub fn preparation_evaluation_count(&self) -> usize {
288        self.objectives
289            .iter()
290            .map(PreparedGeneralRietveldObjective::preparation_evaluation_count)
291            .sum()
292    }
293
294    /// Return expensive model products consumed by one joint normal product.
295    #[must_use]
296    pub fn normal_product_evaluation_count(&self) -> usize {
297        self.objectives
298            .iter()
299            .map(PreparedGeneralRietveldObjective::normal_product_evaluation_count)
300            .sum()
301    }
302
303    /// Apply every histogram Jacobian to one joint physical direction.
304    ///
305    /// # Errors
306    ///
307    /// Returns [`JointRietveldError`] for an invalid direction or product.
308    pub fn jvp(&self, direction: &[f64]) -> Result<Vec<JointRietveldProduct>, JointRietveldError> {
309        self.objectives
310            .iter()
311            .enumerate()
312            .map(|(index, objective)| {
313                let local = self.layout.local_direction(index, direction)?;
314                let (profile, derivative) = objective.jvp(&local)?;
315                Ok(JointRietveldProduct {
316                    histogram_id: self.histograms[index].histogram_id.clone(),
317                    profile,
318                    derivative,
319                })
320            })
321            .collect()
322    }
323
324    /// Apply the transpose of the complete joint Jacobian.
325    ///
326    /// Shared structural rows receive the sum of all histogram products.
327    ///
328    /// # Errors
329    ///
330    /// Returns [`JointRietveldError`] for histogram/sample shape or product
331    /// failures.
332    pub fn vjp(&self, sample_weights: &[Vec<f64>]) -> Result<Vec<f64>, JointRietveldError> {
333        if sample_weights.len() != self.objectives.len() {
334            return Err(JointRietveldError::HistogramProductCountMismatch);
335        }
336        let mut result = vec![0.0; self.layout.parameters.specs().len()];
337        for (index, (objective, weights)) in self.objectives.iter().zip(sample_weights).enumerate()
338        {
339            self.layout
340                .scatter_add(index, &objective.vjp(weights)?, &mut result)?;
341        }
342        Ok(result)
343    }
344
345    /// Apply the summed `J^T W J + damping I` joint normal operator.
346    ///
347    /// # Errors
348    ///
349    /// Returns [`JointRietveldError`] for invalid damping or product state.
350    pub fn normal_product(
351        &self,
352        direction: &[f64],
353        damping: f64,
354    ) -> Result<Vec<f64>, JointRietveldError> {
355        if !damping.is_finite() || damping < 0.0 {
356            return Err(JointRietveldError::InvalidDamping);
357        }
358        let mut result = vec![0.0; self.layout.parameters.specs().len()];
359        for (index, objective) in self.objectives.iter().enumerate() {
360            let local = self.layout.local_direction(index, direction)?;
361            let local_product = objective.normal_product(&local, 0.0)?;
362            self.layout
363                .scatter_add(index, &local_product, &mut result)?;
364        }
365        for (value, direction) in result.iter_mut().zip(direction) {
366            *value += damping * direction;
367        }
368        Ok(result)
369    }
370
371    /// Evaluate the summed accepted-state objective and physical gradient.
372    ///
373    /// # Errors
374    ///
375    /// Returns [`JointRietveldError`] for residual or reverse-product state.
376    pub fn gradient(&self) -> Result<JointRietveldGradient, JointRietveldError> {
377        let mut calculated = Vec::with_capacity(self.objectives.len());
378        let mut gradient = vec![0.0; self.layout.parameters.specs().len()];
379        let mut value = 0.0;
380        for (index, objective) in self.objectives.iter().enumerate() {
381            let (profile, local_gradient) = objective.gradient()?;
382            value += 0.5 * objective.calculation().metrics.chi_square;
383            calculated.push(profile);
384            self.layout
385                .scatter_add(index, &local_gradient, &mut gradient)?;
386        }
387        Ok(JointRietveldGradient {
388            calculated,
389            objective: value,
390            gradient,
391        })
392    }
393}
394
395fn is_shared(key: &ParameterKey) -> bool {
396    matches!(key.module(), "lattice" | "site")
397}
398
399fn validate_histograms(histograms: &[JointRietveldHistogram]) -> Result<(), JointRietveldError> {
400    if histograms.len() < 2 {
401        return Err(JointRietveldError::TooFewHistograms);
402    }
403    if histograms
404        .iter()
405        .map(|histogram| &histogram.histogram_id)
406        .collect::<BTreeSet<_>>()
407        .len()
408        != histograms.len()
409    {
410        return Err(JointRietveldError::DuplicateHistogramId);
411    }
412    let shared_selection = histograms[0].selection.structural;
413    if histograms.iter().skip(1).any(|histogram| {
414        let selection = histogram.selection.structural;
415        selection.lattice != shared_selection.lattice
416            || selection.coordinates != shared_selection.coordinates
417            || selection.occupancy != shared_selection.occupancy
418            || selection.u_iso != shared_selection.u_iso
419    }) {
420        return Err(JointRietveldError::SharedSelectionMismatch);
421    }
422    let mut phases = BTreeMap::new();
423    for histogram in histograms {
424        histogram.input.validate()?;
425        for phase in &histogram.input.phases {
426            let definition = phase.definition();
427            let contract = (
428                phase.site_ids(),
429                definition.cell,
430                &definition.space_group,
431                &definition.fractional_xyz,
432                &definition.occupancy,
433                &definition.u_iso_angstrom2,
434                &definition.anisotropic_mask,
435                &definition.u_aniso_cif_angstrom2,
436                definition.coordinate_tolerance.to_bits(),
437            );
438            if let Some(previous) = phases.insert(phase.phase_id().clone(), contract) {
439                if previous != contract {
440                    return Err(JointRietveldError::SharedPhaseMismatch {
441                        phase_id: phase.phase_id().clone(),
442                    });
443                }
444            }
445        }
446    }
447    Ok(())
448}
449
450/// Invalid joint native Rietveld parameter or objective state.
451#[derive(Debug)]
452pub enum JointRietveldError {
453    /// A joint objective requires at least two histograms.
454    TooFewHistograms,
455    /// Stable histogram identities must be unique.
456    DuplicateHistogramId,
457    /// Shared structural selection differs between histograms.
458    SharedSelectionMismatch,
459    /// Physical structure differs for one shared stable phase.
460    SharedPhaseMismatch {
461        /// Stable incompatible phase identity.
462        phase_id: RecordId,
463    },
464    /// One shared scalar differs in value or metadata.
465    SharedParameterMismatch {
466        /// Stable incompatible parameter identity.
467        key: ParameterKey,
468    },
469    /// A prepared layout was paired with different histograms.
470    HistogramContractMismatch,
471    /// Joint value/direction length is wrong.
472    ValueLengthMismatch,
473    /// Number of histogram reverse products is wrong.
474    HistogramProductCountMismatch,
475    /// A local reverse product has a stale parameter length.
476    LocalProductLengthMismatch,
477    /// Damping must be finite and non-negative.
478    InvalidDamping,
479    /// Stable scalar parameter state is invalid.
480    Parameter(ParameterError),
481    /// Complete single-histogram parameter state is invalid.
482    GeneralParameter(RietveldGeneralParameterError),
483    /// Complete single-histogram objective state is invalid.
484    GeneralObjective(RietveldGeneralObjectiveError),
485}
486
487impl Display for JointRietveldError {
488    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
489        match self {
490            Self::TooFewHistograms => {
491                formatter.write_str("joint Rietveld objective requires at least two histograms")
492            }
493            Self::DuplicateHistogramId => {
494                formatter.write_str("joint Rietveld histogram IDs must be unique")
495            }
496            Self::SharedSelectionMismatch => {
497                formatter.write_str("joint Rietveld shared structural selections must match")
498            }
499            Self::SharedPhaseMismatch { phase_id } => write!(
500                formatter,
501                "joint Rietveld phase {phase_id:?} has incompatible shared structural state"
502            ),
503            Self::SharedParameterMismatch { key } => {
504                write!(
505                    formatter,
506                    "joint Rietveld shared parameter {key} is incompatible"
507                )
508            }
509            Self::HistogramContractMismatch => {
510                formatter.write_str("joint Rietveld histogram contract changed under the layout")
511            }
512            Self::ValueLengthMismatch => {
513                formatter.write_str("joint Rietveld value/direction length is wrong")
514            }
515            Self::HistogramProductCountMismatch => {
516                formatter.write_str("joint Rietveld histogram reverse-product count is wrong")
517            }
518            Self::LocalProductLengthMismatch => {
519                formatter.write_str("joint Rietveld local product length is wrong")
520            }
521            Self::InvalidDamping => {
522                formatter.write_str("joint Rietveld damping must be finite and non-negative")
523            }
524            Self::Parameter(error) => Display::fmt(error, formatter),
525            Self::GeneralParameter(error) => Display::fmt(error, formatter),
526            Self::GeneralObjective(error) => Display::fmt(error, formatter),
527        }
528    }
529}
530
531impl Error for JointRietveldError {
532    fn source(&self) -> Option<&(dyn Error + 'static)> {
533        match self {
534            Self::Parameter(error) => Some(error),
535            Self::GeneralParameter(error) => Some(error),
536            Self::GeneralObjective(error) => Some(error),
537            Self::TooFewHistograms
538            | Self::DuplicateHistogramId
539            | Self::SharedSelectionMismatch
540            | Self::SharedPhaseMismatch { .. }
541            | Self::SharedParameterMismatch { .. }
542            | Self::HistogramContractMismatch
543            | Self::ValueLengthMismatch
544            | Self::HistogramProductCountMismatch
545            | Self::LocalProductLengthMismatch
546            | Self::InvalidDamping => None,
547        }
548    }
549}
550
551impl From<ParameterError> for JointRietveldError {
552    fn from(value: ParameterError) -> Self {
553        Self::Parameter(value)
554    }
555}
556
557impl From<RietveldGeneralParameterError> for JointRietveldError {
558    fn from(value: RietveldGeneralParameterError) -> Self {
559        Self::GeneralParameter(value)
560    }
561}
562
563impl From<RietveldGeneralObjectiveError> for JointRietveldError {
564    fn from(value: RietveldGeneralObjectiveError) -> Self {
565        Self::GeneralObjective(value)
566    }
567}
568
569impl From<crate::RietveldError> for JointRietveldError {
570    fn from(value: crate::RietveldError) -> Self {
571        Self::GeneralParameter(RietveldGeneralParameterError::Rietveld(value))
572    }
573}
phasesmith_workflows/rietveld_joint.rs

phasesmith_workflows/
rietveld_joint.rs
