phasesmith_workflows/
rietveld_objective.rs

1//! Reusable matrix-free structural Rietveld objective products.
2
3use std::error::Error;
4use std::fmt::{Display, Formatter};
5
6use phasesmith_core::SupportPolicy;
7use phasesmith_engine::{
8    PreparedStructuralModelInputView, PreparedStructuralMultiphase, StructuralMultiphaseError,
9};
10
11use crate::rietveld::{
12    assemble_rietveld_calculation, prepare_phase_model, resolve_phase_contributions,
13};
14use crate::{
15    DifferentiableBackground, RietveldCalculation, RietveldCalculationOptions, RietveldError,
16    RietveldInput, RietveldParameterError, RietveldStructuralLayout,
17};
18
19/// One reusable parameter-major physical structural Jacobian and its values.
20#[derive(Clone, Debug, PartialEq)]
21pub struct PreparedRietveldLinearization {
22    /// Complete calculated pattern at the linearized state.
23    pub calculation: RietveldCalculation,
24    /// Structural Jacobian in physical-parameter-major order.
25    pub jacobian: Vec<f64>,
26    /// Number of physical structural parameter rows.
27    pub parameter_count: usize,
28}
29
30/// Reusable structural values and derivative products for one accepted state.
31pub struct PreparedRietveldObjective {
32    input: RietveldInput,
33    options: RietveldCalculationOptions,
34    layout: RietveldStructuralLayout,
35    prepared: PreparedStructuralMultiphase,
36}
37
38impl PreparedRietveldObjective {
39    /// Validate and prepare one structural-only matrix-free objective.
40    ///
41    /// # Errors
42    ///
43    /// Returns [`RietveldObjectiveError`] for invalid inputs or preparation.
44    pub fn new(
45        input: RietveldInput,
46        options: RietveldCalculationOptions,
47        layout: RietveldStructuralLayout,
48    ) -> Result<Self, RietveldObjectiveError> {
49        input.validate()?;
50        options.validate()?;
51        let models = input
52            .phases
53            .iter()
54            .map(|phase| {
55                prepare_phase_model(phase, input.fixed_spectrum.as_ref(), &options.execution)
56            })
57            .collect::<Result<Vec<_>, _>>()?;
58        let prepared = PreparedStructuralMultiphase::new(models, options.execution.clone())?;
59        layout.validate_phases(&input.phases)?;
60        Ok(Self {
61            input,
62            options,
63            layout,
64            prepared,
65        })
66    }
67
68    /// Borrow the stable physical structural layout.
69    #[must_use]
70    pub const fn layout(&self) -> &RietveldStructuralLayout {
71        &self.layout
72    }
73
74    /// Return the native dense element count required by [`Self::linearize`].
75    ///
76    /// # Errors
77    ///
78    /// Returns [`RietveldObjectiveError::AllocationOverflow`] when the product
79    /// is not representable as `usize`.
80    pub fn dense_element_count(&self) -> Result<usize, RietveldObjectiveError> {
81        self.prepared
82            .structural_parameter_counts()
83            .into_iter()
84            .try_fold(0_usize, usize::checked_add)
85            .and_then(|rows| rows.checked_mul(self.input.pattern.sample_count()))
86            .ok_or(RietveldObjectiveError::AllocationOverflow)
87    }
88
89    /// Materialize one reusable physical structural Jacobian and calculation.
90    ///
91    /// Values and native analytical derivatives are produced in the same
92    /// fused engine pass. The layout then projects native rows into stable
93    /// physical parameter order.
94    ///
95    /// # Errors
96    ///
97    /// Returns [`RietveldObjectiveError`] for engine, shape, projection, or
98    /// calculation failures.
99    pub fn linearize(&self) -> Result<PreparedRietveldLinearization, RietveldObjectiveError> {
100        let products = self.with_inputs(|inputs| self.prepared.linearize(inputs))?;
101        let sample_count = self.input.pattern.sample_count();
102        let jacobians = products
103            .iter()
104            .map(|product| (product.d_y.as_slice(), product.parameter_count))
105            .collect::<Vec<_>>();
106        let jacobian = self
107            .layout
108            .project_native_jacobians(&jacobians, sample_count)?;
109        let calculation = assemble_rietveld_calculation(
110            &self.input,
111            &self.options,
112            products.into_iter().map(|product| product.result).collect(),
113        )?;
114        Ok(PreparedRietveldLinearization {
115            calculation,
116            jacobian,
117            parameter_count: self.layout.parameters().specs().len(),
118        })
119    }
120
121    /// Calculate the structural profile and its physical directional derivative.
122    ///
123    /// # Errors
124    ///
125    /// Returns [`RietveldObjectiveError`] for direction or engine failures.
126    pub fn jvp(&self, direction: &[f64]) -> Result<(Vec<f64>, Vec<f64>), RietveldObjectiveError> {
127        let tangents = self.layout.native_tangents(direction)?;
128        let tangent_views = tangents.iter().map(Vec::as_slice).collect::<Vec<_>>();
129        let products = self.with_inputs(|inputs| self.prepared.jvp(inputs, &tangent_views))?;
130        let sample_count = self.input.pattern.sample_count();
131        let mut profile = vec![0.0; sample_count];
132        let mut derivative = vec![0.0; sample_count];
133        for product in products {
134            for sample in 0..sample_count {
135                profile[sample] += product.result.accumulation.y[sample];
136                derivative[sample] += product.d_y[sample];
137            }
138        }
139        Ok((profile, derivative))
140    }
141
142    /// Apply the physical structural pattern Jacobian transpose.
143    ///
144    /// # Errors
145    ///
146    /// Returns [`RietveldObjectiveError`] for weight or engine failures.
147    pub fn vjp(&self, sample_weights: &[f64]) -> Result<Vec<f64>, RietveldObjectiveError> {
148        let products = self.with_inputs(|inputs| self.prepared.vjp(inputs, sample_weights))?;
149        let gradients = products
150            .iter()
151            .map(|item| item.gradient.as_slice())
152            .collect::<Vec<_>>();
153        Ok(self.layout.project_native_gradients(&gradients)?)
154    }
155
156    /// Apply `J^T W J + damping I` without a dense sample Jacobian.
157    ///
158    /// # Errors
159    ///
160    /// Returns [`RietveldObjectiveError`] for invalid damping or product state.
161    pub fn normal_product(
162        &self,
163        direction: &[f64],
164        damping: f64,
165    ) -> Result<Vec<f64>, RietveldObjectiveError> {
166        if !damping.is_finite() || damping < 0.0 {
167            return Err(RietveldObjectiveError::InvalidDamping);
168        }
169        let (_, derivative) = self.jvp(direction)?;
170        let weights = self.weight_samples(&derivative);
171        let mut product = self.vjp(&weights)?;
172        for (value, direction) in product.iter_mut().zip(direction) {
173            *value += damping * direction;
174        }
175        Ok(product)
176    }
177
178    /// Calculate `J^T W (Ycalc - Yobs)` and the complete calculated pattern.
179    ///
180    /// # Errors
181    ///
182    /// Returns [`RietveldObjectiveError`] for calculation or reverse-product
183    /// failures.
184    pub fn gradient(&self) -> Result<(Vec<f64>, Vec<f64>), RietveldObjectiveError> {
185        let zero = vec![0.0; self.layout.parameters().specs().len()];
186        let (profile, _) = self.jvp(&zero)?;
187        let mut background = self.input.pattern.background_y.clone();
188        if let Some(model) = &self.input.background {
189            for (target, value) in background.iter_mut().zip(
190                model
191                    .calculate(&self.input.pattern.x_deg)
192                    .map_err(RietveldError::Background)?,
193            ) {
194                *target += value;
195            }
196        }
197        let calculated = profile
198            .iter()
199            .zip(background)
200            .map(|(profile, background)| profile + background)
201            .collect::<Vec<_>>();
202        if calculated.iter().any(|value| !value.is_finite()) {
203            return Err(RietveldError::NonFiniteCalculation.into());
204        }
205        let observed = self
206            .input
207            .pattern
208            .observed_y
209            .as_deref()
210            .ok_or(RietveldError::MissingObservations)?;
211        let residual = calculated
212            .iter()
213            .zip(observed)
214            .map(|(calculated, observed)| calculated - observed)
215            .collect::<Vec<_>>();
216        Ok((calculated, self.vjp(&self.weight_samples(&residual))?))
217    }
218
219    fn weight_samples(&self, values: &[f64]) -> Vec<f64> {
220        let mask = self.input.pattern.mask.as_deref();
221        let uncertainty = self
222            .options
223            .use_uncertainty
224            .then_some(self.input.pattern.uncertainty.as_deref())
225            .flatten();
226        values
227            .iter()
228            .enumerate()
229            .map(|(index, value)| {
230                if mask.is_some_and(|mask| !mask[index]) {
231                    0.0
232                } else if let Some(sigma) = uncertainty {
233                    value / (sigma[index] * sigma[index])
234                } else {
235                    *value
236                }
237            })
238            .collect()
239    }
240
241    fn with_inputs<R>(
242        &self,
243        operation: impl FnOnce(
244            &[PreparedStructuralModelInputView<'_>],
245        ) -> Result<R, StructuralMultiphaseError>,
246    ) -> Result<R, RietveldObjectiveError> {
247        let owned = self
248            .input
249            .phases
250            .iter()
251            .map(|phase| resolve_phase_contributions(phase, &self.input))
252            .collect::<Result<Vec<_>, _>>()?;
253        let views = owned
254            .iter()
255            .map(|phase| phase.iter().map(|item| item.as_view()).collect::<Vec<_>>())
256            .collect::<Vec<_>>();
257        let inputs = views
258            .iter()
259            .map(|contributions| PreparedStructuralModelInputView {
260                x_deg: &self.input.pattern.x_deg,
261                instrument: self.input.instrument,
262                axial_geometry: self.input.axial_geometry,
263                position_correction: self.input.position_correction,
264                contributions,
265                support: SupportPolicy::FwhmMultiple(self.options.support_fwhm),
266            })
267            .collect::<Vec<_>>();
268        operation(&inputs).map_err(Into::into)
269    }
270}
271
272/// Invalid prepared Rietveld objective state.
273#[derive(Debug)]
274pub enum RietveldObjectiveError {
275    /// Owned Rietveld state is invalid.
276    Rietveld(RietveldError),
277    /// Parameter transform state is invalid.
278    Parameter(RietveldParameterError),
279    /// Structural engine product failed.
280    Structural(StructuralMultiphaseError),
281    /// Damping must be finite and non-negative.
282    InvalidDamping,
283    /// A requested reusable derivative allocation exceeds addressable memory.
284    AllocationOverflow,
285}
286
287impl Display for RietveldObjectiveError {
288    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
289        match self {
290            Self::Rietveld(error) => Display::fmt(error, formatter),
291            Self::Parameter(error) => Display::fmt(error, formatter),
292            Self::Structural(error) => Display::fmt(error, formatter),
293            Self::InvalidDamping => {
294                formatter.write_str("Rietveld damping must be finite and non-negative")
295            }
296            Self::AllocationOverflow => {
297                formatter.write_str("Rietveld dense linearization allocation overflow")
298            }
299        }
300    }
301}
302
303impl Error for RietveldObjectiveError {
304    fn source(&self) -> Option<&(dyn Error + 'static)> {
305        match self {
306            Self::Rietveld(error) => Some(error),
307            Self::Parameter(error) => Some(error),
308            Self::Structural(error) => Some(error),
309            Self::InvalidDamping | Self::AllocationOverflow => None,
310        }
311    }
312}
313impl From<RietveldError> for RietveldObjectiveError {
314    fn from(value: RietveldError) -> Self {
315        Self::Rietveld(value)
316    }
317}
318impl From<RietveldParameterError> for RietveldObjectiveError {
319    fn from(value: RietveldParameterError) -> Self {
320        Self::Parameter(value)
321    }
322}
323impl From<StructuralMultiphaseError> for RietveldObjectiveError {
324    fn from(value: StructuralMultiphaseError) -> Self {
325        Self::Structural(value)
326    }
327}
phasesmith_workflows/rietveld_objective.rs

phasesmith_workflows/
rietveld_objective.rs
