Skip to main content

phasesmith_workflows/
tof_multibank_instrument.rs

1//! Bounded bank-local instrument refinement over atomic multi-bank TOF Le Bail cycles.
2
3use std::collections::BTreeSet;
4use std::error::Error;
5use std::fmt::{Display, Formatter};
6
7use nalgebra::{DMatrix, DVector};
8use phasesmith_core::{
9    TOF_GLOBAL_PARAMETER_COUNT, TofError, TofInstrument, TofInstrumentParameter,
10    TofProfileParameters,
11};
12use phasesmith_model::RecordId;
13
14use crate::tof_lebail::{initialize_intensities, normal_tof_stop};
15use crate::tof_multibank::{
16    AcceptedBankState, MultiBankCycleCandidate, MultiBankCycleOutcome, aggregate_metrics,
17    calculate_states, evaluate_bank_metrics, fitted_parameter_count, prepare_multibank_cycle,
18    reflection_intensities,
19};
20use crate::{
21    DiagnosticValue, RefinementEventKind, RefinementLimits, RefinementRuntime, ResidualEvaluation,
22    RuntimeError, TerminationReason, TofChebyshevBackground, TofLeBailError, TofLeBailOptions,
23    TofLeBailPhase, TofMultiBankError, TofMultiBankInput, TofMultiBankMetrics,
24    TofMultiBankResultBank,
25};
26
27/// Closed physical bounds for one selected bank-local instrument coefficient.
28#[derive(Clone, Copy, Debug, PartialEq)]
29pub struct TofInstrumentParameterBound {
30    /// Coefficient matching one fused dense global derivative row.
31    pub parameter: TofInstrumentParameter,
32    /// Inclusive physical lower bound.
33    pub lower: f64,
34    /// Inclusive physical upper bound.
35    pub upper: f64,
36}
37
38impl TofInstrumentParameterBound {
39    /// Construct one finite non-degenerate physical interval.
40    ///
41    /// # Errors
42    ///
43    /// Returns [`TofMultiBankInstrumentError::InvalidModel`] for invalid bounds.
44    pub fn new(
45        parameter: TofInstrumentParameter,
46        lower: f64,
47        upper: f64,
48    ) -> Result<Self, TofMultiBankInstrumentError> {
49        if !lower.is_finite() || !upper.is_finite() || lower >= upper {
50            return Err(TofMultiBankInstrumentError::InvalidModel(
51                "TOF instrument bounds must be finite and increasing",
52            ));
53        }
54        Ok(Self {
55            parameter,
56            lower,
57            upper,
58        })
59    }
60}
61
62/// Selected bank-local coefficients in deterministic solver order.
63#[derive(Clone, Debug, PartialEq)]
64pub struct TofBankInstrumentModel {
65    bank_id: RecordId,
66    bounds: Vec<TofInstrumentParameterBound>,
67}
68
69impl TofBankInstrumentModel {
70    /// Construct one non-empty bank-local coefficient selection.
71    ///
72    /// # Errors
73    ///
74    /// Returns [`TofMultiBankInstrumentError::InvalidModel`] for empty or
75    /// duplicate selections.
76    pub fn new(
77        bank_id: RecordId,
78        bounds: Vec<TofInstrumentParameterBound>,
79    ) -> Result<Self, TofMultiBankInstrumentError> {
80        let result = Self { bank_id, bounds };
81        result.validate_selection()?;
82        Ok(result)
83    }
84
85    /// Stable bank identity.
86    #[must_use]
87    pub const fn bank_id(&self) -> &RecordId {
88        &self.bank_id
89    }
90
91    /// Selected coefficients and physical bounds in solver order.
92    #[must_use]
93    pub fn bounds(&self) -> &[TofInstrumentParameterBound] {
94        &self.bounds
95    }
96
97    fn validate_selection(&self) -> Result<(), TofMultiBankInstrumentError> {
98        if self.bounds.is_empty() {
99            return Err(TofMultiBankInstrumentError::InvalidModel(
100                "bank-local TOF instrument selection must not be empty",
101            ));
102        }
103        let mut selected = BTreeSet::new();
104        for bound in &self.bounds {
105            if !bound.lower.is_finite() || !bound.upper.is_finite() || bound.lower >= bound.upper {
106                return Err(TofMultiBankInstrumentError::InvalidModel(
107                    "TOF instrument bounds must be finite and increasing",
108                ));
109            }
110            if !selected.insert(bound.parameter) {
111                return Err(TofMultiBankInstrumentError::InvalidModel(
112                    "bank-local TOF instrument parameters must be unique",
113                ));
114            }
115        }
116        Ok(())
117    }
118}
119
120/// Multi-bank observations plus selected local instrument coefficients.
121#[derive(Clone, Debug, PartialEq)]
122pub struct TofMultiBankInstrumentInput {
123    /// Atomic fixed-cell bank-local TOF request.
124    pub multibank: TofMultiBankInput,
125    /// Selected banks in deterministic parameter packing order.
126    pub instrument_models: Vec<TofBankInstrumentModel>,
127}
128
129impl TofMultiBankInstrumentInput {
130    /// Validate bank identities, selections, bounds, and initial values.
131    ///
132    /// # Errors
133    ///
134    /// Returns [`TofMultiBankInstrumentError`] for invalid bank-local state.
135    pub fn validate(&self) -> Result<(), TofMultiBankInstrumentError> {
136        self.multibank.validate()?;
137        if self.instrument_models.is_empty() {
138            return Err(TofMultiBankInstrumentError::InvalidModel(
139                "multi-bank TOF instrument refinement requires at least one selected bank",
140            ));
141        }
142        let mut bank_ids = BTreeSet::new();
143        for model in &self.instrument_models {
144            model.validate_selection()?;
145            if !bank_ids.insert(model.bank_id.clone()) {
146                return Err(TofMultiBankInstrumentError::InvalidModel(
147                    "bank-local TOF instrument model IDs must be unique",
148                ));
149            }
150            let bank = self
151                .multibank
152                .banks
153                .iter()
154                .find(|bank| &bank.bank_id == model.bank_id())
155                .ok_or(TofMultiBankInstrumentError::InvalidModel(
156                    "selected TOF instrument bank is absent from the request",
157                ))?;
158            let values = bank.input.instrument.values();
159            if model.bounds.iter().any(|bound| {
160                let value = values[bound.parameter.index()];
161                value < bound.lower || value > bound.upper
162            }) {
163                return Err(TofMultiBankInstrumentError::InvalidModel(
164                    "initial TOF instrument coefficient lies outside its declared bounds",
165                ));
166            }
167        }
168        Ok(())
169    }
170
171    pub(crate) fn parameter_count(&self) -> Result<usize, TofMultiBankInstrumentError> {
172        self.instrument_models
173            .iter()
174            .try_fold(0_usize, |count, model| {
175                count
176                    .checked_add(model.bounds.len())
177                    .ok_or(TofMultiBankInstrumentError::AllocationOverflow)
178            })
179    }
180}
181
182/// Controls for alternating local Le Bail updates and one instrument step.
183#[derive(Clone, Debug, PartialEq)]
184pub struct TofMultiBankInstrumentOptions {
185    /// Existing TOF redistribution, support, weighting, and execution controls.
186    pub lebail: TofLeBailOptions,
187    /// Non-negative diagonal regularization in scaled instrument coordinates.
188    pub instrument_damping: f64,
189    /// Maximum absolute coefficient step in scaled coordinates.
190    pub max_scaled_instrument_step: f64,
191    /// Number of objective backtracking halvings after the initial trial.
192    pub max_instrument_backtracks: usize,
193    /// Absolute weighted-column correlation reported as unresolved.
194    pub unresolved_correlation: f64,
195}
196
197impl TofMultiBankInstrumentOptions {
198    /// Construct validated bank-local instrument controls.
199    ///
200    /// # Errors
201    ///
202    /// Returns [`TofMultiBankInstrumentError::InvalidOptions`] for invalid controls.
203    pub fn new(
204        lebail: TofLeBailOptions,
205        instrument_damping: f64,
206        max_scaled_instrument_step: f64,
207        max_instrument_backtracks: usize,
208        unresolved_correlation: f64,
209    ) -> Result<Self, TofMultiBankInstrumentError> {
210        let result = Self {
211            lebail,
212            instrument_damping,
213            max_scaled_instrument_step,
214            max_instrument_backtracks,
215            unresolved_correlation,
216        };
217        result.validate()?;
218        Ok(result)
219    }
220
221    /// Revalidate numerical controls.
222    ///
223    /// # Errors
224    ///
225    /// Returns [`TofMultiBankInstrumentError::InvalidOptions`] for invalid fields.
226    pub fn validate(&self) -> Result<(), TofMultiBankInstrumentError> {
227        self.lebail.validate()?;
228        if !self.instrument_damping.is_finite()
229            || self.instrument_damping < 0.0
230            || !self.max_scaled_instrument_step.is_finite()
231            || self.max_scaled_instrument_step <= 0.0
232            || !self.unresolved_correlation.is_finite()
233            || !(0.0..=1.0).contains(&self.unresolved_correlation)
234        {
235            return Err(TofMultiBankInstrumentError::InvalidOptions);
236        }
237        Ok(())
238    }
239}
240
241/// One accepted bank-local instrument.
242#[derive(Clone, Debug, PartialEq)]
243pub struct TofBankInstrumentState {
244    /// Stable bank identity.
245    pub bank_id: RecordId,
246    /// Accepted calibration/profile coefficients.
247    pub instrument: TofInstrument,
248}
249
250/// One accepted physical bank-local coefficient change.
251#[derive(Clone, Debug, PartialEq)]
252pub struct TofInstrumentParameterChange {
253    /// Stable bank identity.
254    pub bank_id: RecordId,
255    /// Selected coefficient.
256    pub parameter: TofInstrumentParameter,
257    /// Physical value before the step.
258    pub before: f64,
259    /// Physical value after the step.
260    pub after: f64,
261    /// Change divided by the parameter scale.
262    pub scaled_change: f64,
263}
264
265/// One unresolved weighted-Jacobian column pair.
266#[derive(Clone, Debug, PartialEq)]
267pub struct TofInstrumentCorrelation {
268    /// Bank and coefficient for the first column.
269    pub left_bank_id: RecordId,
270    /// First selected coefficient.
271    pub left_parameter: TofInstrumentParameter,
272    /// Bank and coefficient for the second column.
273    pub right_bank_id: RecordId,
274    /// Second selected coefficient.
275    pub right_parameter: TofInstrumentParameter,
276    /// Normalized weighted-column dot product.
277    pub correlation: f64,
278}
279
280/// Identifiability diagnostics for the final selected instrument system.
281#[derive(Clone, Debug, PartialEq)]
282pub struct TofInstrumentDiagnostics {
283    /// Number of selected physical coefficients.
284    pub parameter_count: usize,
285    /// Numerical rank of the weighted Jacobian.
286    pub jacobian_rank: usize,
287    /// Largest finite absolute correlation between nonzero columns.
288    pub maximum_absolute_correlation: Option<f64>,
289    /// Pairs at or above the configured unresolved threshold.
290    pub unresolved_correlations: Vec<TofInstrumentCorrelation>,
291}
292
293/// One atomically accepted extraction plus bank-local instrument cycle.
294#[derive(Clone, Debug, PartialEq)]
295pub struct TofMultiBankInstrumentIterationRecord {
296    /// One-based accepted cycle index.
297    pub iteration: usize,
298    /// Bank-local residual records in input order.
299    pub bank_metrics: Vec<ResidualEvaluation>,
300    /// Aggregate metrics with selected instrument coefficients counted once.
301    pub metrics: TofMultiBankMetrics,
302    /// Largest relative intensity change.
303    pub maximum_relative_intensity_change: f64,
304    /// Largest absolute background-coefficient change.
305    pub maximum_absolute_background_change: f64,
306    /// Norm of the accepted instrument step in scaled coordinates.
307    pub scaled_instrument_step_norm: f64,
308    /// Physical changes in deterministic bank/selection order.
309    pub instrument_parameter_changes: Vec<TofInstrumentParameterChange>,
310}
311
312/// One bank's accepted checkpoint state.
313#[derive(Clone, Debug, PartialEq)]
314pub struct TofMultiBankInstrumentCheckpointBank {
315    /// Stable bank identity.
316    pub bank_id: RecordId,
317    /// Accepted instrument.
318    pub instrument: TofInstrument,
319    /// Accepted extracted phase intensities.
320    pub phases: Vec<TofLeBailPhase>,
321    /// Accepted refinable background, when configured.
322    pub background: Option<TofChebyshevBackground>,
323}
324
325/// Complete last-accepted state for exact instrument continuation.
326#[derive(Clone, Debug, PartialEq)]
327pub struct TofMultiBankInstrumentCheckpoint {
328    /// Number of atomically accepted cycles.
329    pub completed_iterations: usize,
330    /// Accepted bank-local instrument/intensity/background state.
331    pub banks: Vec<TofMultiBankInstrumentCheckpointBank>,
332    /// Complete deterministic accepted history.
333    pub history: Vec<TofMultiBankInstrumentIterationRecord>,
334}
335
336/// Complete result from analytical bank-local TOF instrument refinement.
337#[derive(Clone, Debug, PartialEq)]
338pub struct TofMultiBankInstrumentResult {
339    /// Final bank-local display states and fused derivative calculations.
340    pub banks: Vec<TofMultiBankResultBank>,
341    /// Accepted instruments in bank input order.
342    pub instruments: Vec<TofBankInstrumentState>,
343    /// Final aggregate residual metrics.
344    pub metrics: TofMultiBankMetrics,
345    /// Final selected-system identifiability diagnostics.
346    pub diagnostics: TofInstrumentDiagnostics,
347    /// Complete deterministic accepted history.
348    pub history: Vec<TofMultiBankInstrumentIterationRecord>,
349    /// Stable bounded-runtime termination category.
350    pub termination_reason: TerminationReason,
351    /// Complete last accepted state.
352    pub checkpoint: TofMultiBankInstrumentCheckpoint,
353}
354
355impl TofMultiBankInstrumentCheckpoint {
356    /// Revalidate this continuation against the exact local/instrument contract.
357    ///
358    /// # Errors
359    ///
360    /// Returns [`TofMultiBankInstrumentError`] for stale identities, topology,
361    /// bounds, backgrounds, instruments, or history.
362    pub fn validate_for(
363        &self,
364        input: &TofMultiBankInstrumentInput,
365        options: &TofMultiBankInstrumentOptions,
366    ) -> Result<(), TofMultiBankInstrumentError> {
367        input.validate()?;
368        options.validate()?;
369        if self.completed_iterations != self.history.len()
370            || self.completed_iterations > options.lebail.cycles
371            || self.banks.len() != input.multibank.banks.len()
372        {
373            return Err(TofMultiBankInstrumentError::InvalidCheckpoint(
374                "TOF instrument checkpoint counts differ from the request",
375            ));
376        }
377        for (saved, original) in self.banks.iter().zip(&input.multibank.banks) {
378            validate_checkpoint_bank(saved, original, input)?;
379        }
380        if self.history.iter().enumerate().any(|(index, record)| {
381            record.iteration != index + 1
382                || record.bank_metrics.len() != input.multibank.banks.len()
383        }) {
384            return Err(TofMultiBankInstrumentError::InvalidCheckpoint(
385                "TOF instrument checkpoint history is not contiguous and bank-aligned",
386            ));
387        }
388        Ok(())
389    }
390}
391
392/// Refine selected bank-local instrument coefficients against all TOF banks.
393///
394/// # Errors
395///
396/// Returns [`TofMultiBankInstrumentError`] for invalid contracts, calculation,
397/// linear solve, or runtime state.
398pub fn refine_tof_multibank_instrument(
399    input: &TofMultiBankInstrumentInput,
400    options: &TofMultiBankInstrumentOptions,
401) -> Result<TofMultiBankInstrumentResult, TofMultiBankInstrumentError> {
402    input.validate()?;
403    options.validate()?;
404    let evaluations_per_cycle = options
405        .max_instrument_backtracks
406        .checked_add(4)
407        .ok_or(TofMultiBankInstrumentError::AllocationOverflow)?;
408    let max_evaluations = options
409        .lebail
410        .cycles
411        .checked_mul(evaluations_per_cycle)
412        .ok_or(TofMultiBankInstrumentError::AllocationOverflow)?;
413    let limits = RefinementLimits::new(options.lebail.cycles, max_evaluations, None, 1)?;
414    let mut runtime = RefinementRuntime::new(limits, None)?;
415    refine_tof_multibank_instrument_with_runtime(input, options, None, &mut runtime)
416}
417
418/// Refine selected bank-local instruments with cancellation and continuation.
419///
420/// Candidate instruments, intensities, and backgrounds publish only after the
421/// complete joint cycle succeeds. A stop during backtracking discards the
422/// entire in-progress cycle.
423///
424/// # Errors
425///
426/// Returns [`TofMultiBankInstrumentError`] for invalid contracts, calculation,
427/// linear solve, or non-normal runtime failures.
428#[allow(clippy::too_many_lines)]
429pub fn refine_tof_multibank_instrument_with_runtime(
430    input: &TofMultiBankInstrumentInput,
431    options: &TofMultiBankInstrumentOptions,
432    checkpoint: Option<&TofMultiBankInstrumentCheckpoint>,
433    runtime: &mut RefinementRuntime<TofMultiBankInstrumentCheckpoint>,
434) -> Result<TofMultiBankInstrumentResult, TofMultiBankInstrumentError> {
435    input.validate()?;
436    options.validate()?;
437    let parameter_count = input.parameter_count()?;
438    let restored = restore_state(input, options, checkpoint)?;
439    let mut states = restored.states;
440    let mut instruments = restored.instruments;
441    let mut history = restored.history;
442    if let Some(checkpoint) = checkpoint {
443        runtime.resume_accepted(checkpoint.completed_iterations)?;
444    }
445    runtime.emit(
446        RefinementEventKind::Start,
447        "tof_multibank_instrument",
448        "bank-local multi-bank TOF instrument refinement started",
449        vec![
450            (
451                "bank_count".to_owned(),
452                DiagnosticValue::Unsigned(input.multibank.banks.len() as u64),
453            ),
454            (
455                "instrument_parameter_count".to_owned(),
456                DiagnosticValue::Unsigned(parameter_count as u64),
457            ),
458        ],
459    )?;
460    let mut accepted_calculations = None;
461    let mut termination = TerminationReason::MaxIterations;
462    for iteration in restored.first_iteration..=options.lebail.cycles {
463        if let Err(error) = runtime.begin_iteration(iteration) {
464            termination = normal_tof_stop(error)?;
465            break;
466        }
467        let live = live_input(input, &instruments)?;
468        let candidate = match prepare_multibank_cycle(
469            &live,
470            &states,
471            &options.lebail,
472            parameter_count,
473            runtime,
474        )? {
475            MultiBankCycleOutcome::Candidate(candidate) => candidate,
476            MultiBankCycleOutcome::Stopped(reason) => {
477                termination = reason;
478                break;
479            }
480        };
481        let updated = match instrument_update(input, options, &instruments, candidate, runtime)? {
482            InstrumentUpdateOutcome::Candidate(candidate) => candidate,
483            InstrumentUpdateOutcome::Stopped(reason) => {
484                termination = reason;
485                break;
486            }
487        };
488        states = updated.cycle.states;
489        instruments = updated.instruments;
490        accepted_calculations = Some(updated.cycle.calculations.clone());
491        history.push(TofMultiBankInstrumentIterationRecord {
492            iteration,
493            bank_metrics: updated.cycle.bank_metrics.clone(),
494            metrics: updated.cycle.metrics,
495            maximum_relative_intensity_change: updated.cycle.maximum_relative_intensity_change,
496            maximum_absolute_background_change: updated.cycle.maximum_absolute_background_change,
497            scaled_instrument_step_norm: updated.step_norm,
498            instrument_parameter_changes: updated.changes,
499        });
500        let accepted = checkpoint_from_state(input, &states, &instruments, &history);
501        runtime.accept_step(Some(&accepted))?;
502        runtime.emit(
503            RefinementEventKind::Iteration,
504            "tof_multibank_instrument_iteration",
505            "bank-local multi-bank TOF instrument cycle accepted",
506            vec![
507                (
508                    "rwp".to_owned(),
509                    DiagnosticValue::Float(updated.cycle.metrics.rwp),
510                ),
511                (
512                    "scaled_instrument_step_norm".to_owned(),
513                    DiagnosticValue::Float(updated.step_norm),
514                ),
515            ],
516        )?;
517    }
518    let live = live_input(input, &instruments)?;
519    let calculations = match accepted_calculations {
520        Some(calculations) => calculations,
521        None => calculate_states(&live, &states, &options.lebail)?,
522    };
523    let bank_metrics = evaluate_bank_metrics(&live, &states, &calculations, &options.lebail)?;
524    let fitted = fitted_parameter_count(&states)?
525        .checked_add(parameter_count)
526        .ok_or(TofMultiBankInstrumentError::AllocationOverflow)?;
527    let metrics = aggregate_metrics(&live, &calculations, &options.lebail, fitted)?;
528    let packed = pack_instruments(input, &instruments)?;
529    let (jacobian, _) =
530        instrument_system(input, &live, &calculations, &packed.scales, &options.lebail)?;
531    let diagnostics = instrument_diagnostics(input, &jacobian, options.unresolved_correlation)?;
532    let checkpoint = checkpoint_from_state(input, &states, &instruments, &history);
533    checkpoint.validate_for(input, options)?;
534    let banks = live
535        .banks
536        .iter()
537        .zip(states)
538        .zip(calculations)
539        .zip(bank_metrics)
540        .map(
541            |(((bank, state), calculation), metrics)| TofMultiBankResultBank {
542                bank_id: bank.bank_id.clone(),
543                calculation,
544                metrics,
545                intensities: reflection_intensities(&state.phases),
546                phases: state.phases,
547                background: state.background,
548            },
549        )
550        .collect();
551    let instruments = instrument_states(input, &instruments);
552    runtime.emit(
553        RefinementEventKind::Termination,
554        "tof_multibank_instrument",
555        "bank-local multi-bank TOF instrument refinement terminated",
556        vec![
557            (
558                "termination_reason".to_owned(),
559                DiagnosticValue::String(termination.as_str().to_owned()),
560            ),
561            (
562                "jacobian_rank".to_owned(),
563                DiagnosticValue::Unsigned(diagnostics.jacobian_rank as u64),
564            ),
565        ],
566    )?;
567    Ok(TofMultiBankInstrumentResult {
568        banks,
569        instruments,
570        metrics,
571        diagnostics,
572        history,
573        termination_reason: termination,
574        checkpoint,
575    })
576}
577
578struct RestoredState {
579    states: Vec<AcceptedBankState>,
580    instruments: Vec<TofInstrument>,
581    history: Vec<TofMultiBankInstrumentIterationRecord>,
582    first_iteration: usize,
583}
584
585fn restore_state(
586    input: &TofMultiBankInstrumentInput,
587    options: &TofMultiBankInstrumentOptions,
588    checkpoint: Option<&TofMultiBankInstrumentCheckpoint>,
589) -> Result<RestoredState, TofMultiBankInstrumentError> {
590    let Some(checkpoint) = checkpoint else {
591        let states = input
592            .multibank
593            .banks
594            .iter()
595            .map(|bank| {
596                Ok(AcceptedBankState {
597                    phases: initialize_intensities(&bank.input, &options.lebail)?,
598                    background: bank.input.background.clone(),
599                })
600            })
601            .collect::<Result<Vec<_>, TofMultiBankInstrumentError>>()?;
602        return Ok(RestoredState {
603            states,
604            instruments: input
605                .multibank
606                .banks
607                .iter()
608                .map(|bank| bank.input.instrument)
609                .collect(),
610            history: Vec::new(),
611            first_iteration: 1,
612        });
613    };
614    checkpoint.validate_for(input, options)?;
615    Ok(RestoredState {
616        states: checkpoint
617            .banks
618            .iter()
619            .map(|bank| AcceptedBankState {
620                phases: bank.phases.clone(),
621                background: bank.background.clone(),
622            })
623            .collect(),
624        instruments: checkpoint
625            .banks
626            .iter()
627            .map(|bank| bank.instrument)
628            .collect(),
629        history: checkpoint.history.clone(),
630        first_iteration: checkpoint
631            .completed_iterations
632            .checked_add(1)
633            .ok_or(TofMultiBankInstrumentError::AllocationOverflow)?,
634    })
635}
636
637struct InstrumentUpdate {
638    cycle: MultiBankCycleCandidate,
639    instruments: Vec<TofInstrument>,
640    step_norm: f64,
641    changes: Vec<TofInstrumentParameterChange>,
642}
643
644enum InstrumentUpdateOutcome {
645    Candidate(InstrumentUpdate),
646    Stopped(TerminationReason),
647}
648
649#[allow(clippy::too_many_lines)]
650fn instrument_update<C>(
651    input: &TofMultiBankInstrumentInput,
652    options: &TofMultiBankInstrumentOptions,
653    instruments: &[TofInstrument],
654    mut cycle: MultiBankCycleCandidate,
655    runtime: &mut RefinementRuntime<C>,
656) -> Result<InstrumentUpdateOutcome, TofMultiBankInstrumentError> {
657    let packed = pack_instruments(input, instruments)?;
658    let current = live_input(input, instruments)?;
659    let (jacobian, residual) = instrument_system(
660        input,
661        &current,
662        &cycle.calculations,
663        &packed.scales,
664        &options.lebail,
665    )?;
666    let normal = jacobian.transpose() * &jacobian;
667    let rhs = jacobian.transpose() * &residual;
668    let mut regularized = normal;
669    for parameter in 0..regularized.nrows() {
670        regularized[(parameter, parameter)] += options.instrument_damping;
671    }
672    let mut step = if let Some(solution) = regularized.lu().solve(&rhs) {
673        solution
674    } else {
675        jacobian
676            .svd(true, true)
677            .solve(&residual, f64::EPSILON)
678            .map_err(|_| TofMultiBankInstrumentError::LinearSolve)?
679    };
680    if step.iter().any(|value| !value.is_finite()) {
681        return Err(TofMultiBankInstrumentError::LinearSolve);
682    }
683    for parameter in 0..step.len() {
684        let current = packed.values[parameter] / packed.scales[parameter];
685        let lower = (packed.lower[parameter] / packed.scales[parameter] - current)
686            .max(-options.max_scaled_instrument_step);
687        let upper = (packed.upper[parameter] / packed.scales[parameter] - current)
688            .min(options.max_scaled_instrument_step);
689        step[parameter] = step[parameter].clamp(lower, upper);
690    }
691    if step.norm() == 0.0 {
692        return Ok(InstrumentUpdateOutcome::Candidate(InstrumentUpdate {
693            cycle,
694            instruments: instruments.to_vec(),
695            step_norm: 0.0,
696            changes: Vec::new(),
697        }));
698    }
699    let parameter_count = input.parameter_count()?;
700    let local_parameter_count = fitted_parameter_count(&cycle.states)?;
701    let total_parameter_count = local_parameter_count
702        .checked_add(parameter_count)
703        .ok_or(TofMultiBankInstrumentError::AllocationOverflow)?;
704    let mut factor = 1.0;
705    for _ in 0..=options.max_instrument_backtracks {
706        let trial_values = packed
707            .values
708            .iter()
709            .zip(&packed.scales)
710            .zip(step.iter())
711            .map(|((value, scale), step)| value + factor * scale * step)
712            .collect::<Vec<_>>();
713        let trial_instruments = match unpack_instruments(input, instruments, &trial_values) {
714            Ok(instruments) => instruments,
715            Err(error) if recoverable_trial_error(&error) => {
716                factor *= 0.5;
717                continue;
718            }
719            Err(error) => return Err(error),
720        };
721        let trial_input = match live_input(input, &trial_instruments) {
722            Ok(input) => input,
723            Err(error) if recoverable_trial_error(&error) => {
724                factor *= 0.5;
725                continue;
726            }
727            Err(error) => return Err(error),
728        };
729        if let Err(error) = runtime.begin_evaluation() {
730            return Ok(InstrumentUpdateOutcome::Stopped(normal_tof_stop(error)?));
731        }
732        let trial_calculations =
733            match calculate_states(&trial_input, &cycle.states, &options.lebail) {
734                Ok(calculations) => calculations,
735                Err(error) if recoverable_multibank_trial_error(&error) => {
736                    factor *= 0.5;
737                    continue;
738                }
739                Err(error) => return Err(error.into()),
740            };
741        let trial_metrics = aggregate_metrics(
742            &trial_input,
743            &trial_calculations,
744            &options.lebail,
745            total_parameter_count,
746        )?;
747        if trial_metrics.chi_square < cycle.metrics.chi_square {
748            let trial_bank_metrics = evaluate_bank_metrics(
749                &trial_input,
750                &cycle.states,
751                &trial_calculations,
752                &options.lebail,
753            )?;
754            let changes = instrument_changes(input, &packed, &trial_values)?;
755            cycle.calculations = trial_calculations;
756            cycle.bank_metrics = trial_bank_metrics;
757            cycle.metrics = trial_metrics;
758            return Ok(InstrumentUpdateOutcome::Candidate(InstrumentUpdate {
759                cycle,
760                instruments: trial_instruments,
761                step_norm: factor * step.norm(),
762                changes,
763            }));
764        }
765        factor *= 0.5;
766    }
767    Ok(InstrumentUpdateOutcome::Candidate(InstrumentUpdate {
768        cycle,
769        instruments: instruments.to_vec(),
770        step_norm: 0.0,
771        changes: Vec::new(),
772    }))
773}
774
775fn recoverable_trial_error(error: &TofMultiBankInstrumentError) -> bool {
776    matches!(
777        error,
778        TofMultiBankInstrumentError::Profile(_)
779            | TofMultiBankInstrumentError::MultiBank(TofMultiBankError::LeBail(
780                TofLeBailError::Profile(_) | TofLeBailError::InvalidPhase(_)
781            ))
782    )
783}
784
785fn recoverable_multibank_trial_error(error: &TofMultiBankError) -> bool {
786    matches!(
787        error,
788        TofMultiBankError::LeBail(TofLeBailError::Profile(_) | TofLeBailError::InvalidPhase(_))
789    )
790}
791
792pub(crate) struct PackedInstruments {
793    pub(crate) values: Vec<f64>,
794    pub(crate) scales: Vec<f64>,
795    pub(crate) lower: Vec<f64>,
796    pub(crate) upper: Vec<f64>,
797}
798
799pub(crate) fn pack_instruments(
800    input: &TofMultiBankInstrumentInput,
801    instruments: &[TofInstrument],
802) -> Result<PackedInstruments, TofMultiBankInstrumentError> {
803    if instruments.len() != input.multibank.banks.len() {
804        return Err(TofMultiBankInstrumentError::InternalInvariant);
805    }
806    let mut values = Vec::new();
807    let mut scales = Vec::new();
808    let mut lower = Vec::new();
809    let mut upper = Vec::new();
810    for model in &input.instrument_models {
811        let bank_index = bank_index(input, model.bank_id())?;
812        let instrument_values = instruments[bank_index].values();
813        for bound in &model.bounds {
814            let value = instrument_values[bound.parameter.index()];
815            let half_span = 0.5 * (bound.upper - bound.lower);
816            values.push(value);
817            scales.push(value.abs().max(half_span).max(f64::EPSILON.sqrt()));
818            lower.push(bound.lower);
819            upper.push(bound.upper);
820        }
821    }
822    Ok(PackedInstruments {
823        values,
824        scales,
825        lower,
826        upper,
827    })
828}
829
830pub(crate) fn unpack_instruments(
831    input: &TofMultiBankInstrumentInput,
832    current: &[TofInstrument],
833    values: &[f64],
834) -> Result<Vec<TofInstrument>, TofMultiBankInstrumentError> {
835    if current.len() != input.multibank.banks.len() || values.len() != input.parameter_count()? {
836        return Err(TofMultiBankInstrumentError::InternalInvariant);
837    }
838    let mut result = current.to_vec();
839    let mut offset = 0;
840    for model in &input.instrument_models {
841        let index = bank_index(input, model.bank_id())?;
842        let mut instrument_values = result[index].values();
843        for bound in &model.bounds {
844            instrument_values[bound.parameter.index()] = values[offset];
845            offset += 1;
846        }
847        result[index] = TofInstrument::from_values(instrument_values)?;
848    }
849    Ok(result)
850}
851
852pub(crate) fn live_input(
853    input: &TofMultiBankInstrumentInput,
854    instruments: &[TofInstrument],
855) -> Result<TofMultiBankInput, TofMultiBankInstrumentError> {
856    if instruments.len() != input.multibank.banks.len() {
857        return Err(TofMultiBankInstrumentError::InternalInvariant);
858    }
859    let mut result = input.multibank.clone();
860    for (bank, instrument) in result.banks.iter_mut().zip(instruments) {
861        bank.input.instrument = *instrument;
862    }
863    result.validate()?;
864    Ok(result)
865}
866
867fn bank_index(
868    input: &TofMultiBankInstrumentInput,
869    bank_id: &RecordId,
870) -> Result<usize, TofMultiBankInstrumentError> {
871    input
872        .multibank
873        .banks
874        .iter()
875        .position(|bank| &bank.bank_id == bank_id)
876        .ok_or(TofMultiBankInstrumentError::InternalInvariant)
877}
878
879pub(crate) fn instrument_system(
880    input: &TofMultiBankInstrumentInput,
881    live: &TofMultiBankInput,
882    calculations: &[crate::TofLeBailCalculation],
883    scales: &[f64],
884    options: &TofLeBailOptions,
885) -> Result<(DMatrix<f64>, DVector<f64>), TofMultiBankInstrumentError> {
886    let columns = input.parameter_count()?;
887    if columns != scales.len()
888        || live.banks.len() != input.multibank.banks.len()
889        || calculations.len() != live.banks.len()
890    {
891        return Err(TofMultiBankInstrumentError::InternalInvariant);
892    }
893    let rows = live.banks.iter().try_fold(0_usize, |count, bank| {
894        let included = bank.input.pattern.mask.as_ref().map_or_else(
895            || bank.input.pattern.sample_count(),
896            |mask| mask.iter().filter(|value| **value).count(),
897        );
898        count
899            .checked_add(included)
900            .ok_or(TofMultiBankInstrumentError::AllocationOverflow)
901    })?;
902    if rows == 0 {
903        return Err(TofMultiBankInstrumentError::LinearSolve);
904    }
905    let mut selected_jacobian = DMatrix::zeros(rows, columns);
906    let mut selected_residual = DVector::zeros(rows);
907    let mut model_offsets = Vec::with_capacity(input.instrument_models.len());
908    let mut packed_offset = 0_usize;
909    for model in &input.instrument_models {
910        model_offsets.push(packed_offset);
911        packed_offset = packed_offset
912            .checked_add(model.bounds.len())
913            .ok_or(TofMultiBankInstrumentError::AllocationOverflow)?;
914    }
915    if packed_offset != columns {
916        return Err(TofMultiBankInstrumentError::InternalInvariant);
917    }
918    let mut row_offset = 0;
919    for (bank, calculation) in live.banks.iter().zip(calculations) {
920        let samples = bank.input.pattern.sample_count();
921        let global = calculation
922            .accumulation
923            .derivatives
924            .global
925            .as_ref()
926            .ok_or(TofMultiBankInstrumentError::InternalInvariant)?;
927        if global.parameter_count != TOF_GLOBAL_PARAMETER_COUNT
928            || global.sample_count != samples
929            || global.values.len() != TOF_GLOBAL_PARAMETER_COUNT * samples
930        {
931            return Err(TofMultiBankInstrumentError::InternalInvariant);
932        }
933        let selected = input
934            .instrument_models
935            .iter()
936            .position(|model| model.bank_id() == &bank.bank_id);
937        let observed = bank
938            .input
939            .pattern
940            .observed_y
941            .as_ref()
942            .ok_or(TofLeBailError::MissingObservations)?;
943        for sample in 0..samples {
944            if bank
945                .input
946                .pattern
947                .mask
948                .as_ref()
949                .is_some_and(|mask| !mask[sample])
950            {
951                continue;
952            }
953            let weight = if options.use_uncertainty {
954                bank.input
955                    .pattern
956                    .uncertainty
957                    .as_ref()
958                    .map_or(1.0, |sigma| sigma[sample].recip())
959            } else {
960                1.0
961            };
962            selected_residual[row_offset] = (observed[sample] - calculation.y[sample]) * weight;
963            if let Some(model_index) = selected {
964                let model = &input.instrument_models[model_index];
965                let offset = model_offsets[model_index];
966                for (local, bound) in model.bounds.iter().enumerate() {
967                    selected_jacobian[(row_offset, offset + local)] = global.values
968                        [bound.parameter.index() * samples + sample]
969                        * scales[offset + local]
970                        * weight;
971                }
972            }
973            row_offset += 1;
974        }
975    }
976    if row_offset != rows {
977        return Err(TofMultiBankInstrumentError::InternalInvariant);
978    }
979    Ok((selected_jacobian, selected_residual))
980}
981
982fn instrument_changes(
983    input: &TofMultiBankInstrumentInput,
984    packed: &PackedInstruments,
985    after: &[f64],
986) -> Result<Vec<TofInstrumentParameterChange>, TofMultiBankInstrumentError> {
987    if packed.values.len() != after.len() || packed.values.len() != packed.scales.len() {
988        return Err(TofMultiBankInstrumentError::InternalInvariant);
989    }
990    let mut changes = Vec::new();
991    let mut offset = 0;
992    for model in &input.instrument_models {
993        for bound in &model.bounds {
994            if packed.values[offset].to_bits() != after[offset].to_bits() {
995                changes.push(TofInstrumentParameterChange {
996                    bank_id: model.bank_id.clone(),
997                    parameter: bound.parameter,
998                    before: packed.values[offset],
999                    after: after[offset],
1000                    scaled_change: (after[offset] - packed.values[offset]) / packed.scales[offset],
1001                });
1002            }
1003            offset += 1;
1004        }
1005    }
1006    Ok(changes)
1007}
1008
1009fn instrument_diagnostics(
1010    input: &TofMultiBankInstrumentInput,
1011    jacobian: &DMatrix<f64>,
1012    threshold: f64,
1013) -> Result<TofInstrumentDiagnostics, TofMultiBankInstrumentError> {
1014    let parameter_count = input.parameter_count()?;
1015    if jacobian.ncols() != parameter_count {
1016        return Err(TofMultiBankInstrumentError::InternalInvariant);
1017    }
1018    let singular = jacobian.clone().svd(false, false).singular_values;
1019    let maximum = singular.iter().copied().fold(0.0_f64, f64::max);
1020    let dimension = u32::try_from(jacobian.nrows().max(jacobian.ncols())).unwrap_or(u32::MAX);
1021    let tolerance = f64::from(dimension) * f64::EPSILON * maximum;
1022    let jacobian_rank = singular.iter().filter(|value| **value > tolerance).count();
1023    let keys = instrument_keys(input);
1024    let norms = (0..parameter_count)
1025        .map(|column| jacobian.column(column).norm())
1026        .collect::<Vec<_>>();
1027    let mut maximum_absolute_correlation = None::<f64>;
1028    let mut unresolved_correlations = Vec::new();
1029    for left in 0..parameter_count {
1030        if norms[left] == 0.0 {
1031            continue;
1032        }
1033        for right in left + 1..parameter_count {
1034            if norms[right] == 0.0 {
1035                continue;
1036            }
1037            let correlation = (jacobian.column(left).dot(&jacobian.column(right))
1038                / (norms[left] * norms[right]))
1039                .clamp(-1.0, 1.0);
1040            let absolute = correlation.abs();
1041            maximum_absolute_correlation =
1042                Some(maximum_absolute_correlation.map_or(absolute, |value| value.max(absolute)));
1043            if absolute >= threshold {
1044                unresolved_correlations.push(TofInstrumentCorrelation {
1045                    left_bank_id: keys[left].0.clone(),
1046                    left_parameter: keys[left].1,
1047                    right_bank_id: keys[right].0.clone(),
1048                    right_parameter: keys[right].1,
1049                    correlation,
1050                });
1051            }
1052        }
1053    }
1054    Ok(TofInstrumentDiagnostics {
1055        parameter_count,
1056        jacobian_rank,
1057        maximum_absolute_correlation,
1058        unresolved_correlations,
1059    })
1060}
1061
1062fn instrument_keys(input: &TofMultiBankInstrumentInput) -> Vec<(RecordId, TofInstrumentParameter)> {
1063    input
1064        .instrument_models
1065        .iter()
1066        .flat_map(|model| {
1067            model
1068                .bounds
1069                .iter()
1070                .map(|bound| (model.bank_id.clone(), bound.parameter))
1071        })
1072        .collect()
1073}
1074
1075fn checkpoint_from_state(
1076    input: &TofMultiBankInstrumentInput,
1077    states: &[AcceptedBankState],
1078    instruments: &[TofInstrument],
1079    history: &[TofMultiBankInstrumentIterationRecord],
1080) -> TofMultiBankInstrumentCheckpoint {
1081    TofMultiBankInstrumentCheckpoint {
1082        completed_iterations: history.len(),
1083        banks: input
1084            .multibank
1085            .banks
1086            .iter()
1087            .zip(states)
1088            .zip(instruments)
1089            .map(
1090                |((bank, state), instrument)| TofMultiBankInstrumentCheckpointBank {
1091                    bank_id: bank.bank_id.clone(),
1092                    instrument: *instrument,
1093                    phases: state.phases.clone(),
1094                    background: state.background.clone(),
1095                },
1096            )
1097            .collect(),
1098        history: history.to_vec(),
1099    }
1100}
1101
1102fn instrument_states(
1103    input: &TofMultiBankInstrumentInput,
1104    instruments: &[TofInstrument],
1105) -> Vec<TofBankInstrumentState> {
1106    input
1107        .multibank
1108        .banks
1109        .iter()
1110        .zip(instruments)
1111        .map(|(bank, instrument)| TofBankInstrumentState {
1112            bank_id: bank.bank_id.clone(),
1113            instrument: *instrument,
1114        })
1115        .collect()
1116}
1117
1118fn validate_checkpoint_bank(
1119    saved: &TofMultiBankInstrumentCheckpointBank,
1120    original: &crate::TofLeBailBank,
1121    input: &TofMultiBankInstrumentInput,
1122) -> Result<(), TofMultiBankInstrumentError> {
1123    saved.instrument.validate()?;
1124    if saved.bank_id != original.bank_id || saved.phases.len() != original.input.phases.len() {
1125        return Err(TofMultiBankInstrumentError::InvalidCheckpoint(
1126            "TOF instrument checkpoint bank identity or phase count changed",
1127        ));
1128    }
1129    for (saved_phase, original_phase) in saved.phases.iter().zip(&original.input.phases) {
1130        saved_phase.validate()?;
1131        if saved_phase.phase_id() != original_phase.phase_id()
1132            || saved_phase.name() != original_phase.name()
1133            || saved_phase.reflection_ids() != original_phase.reflection_ids()
1134            || saved_phase.hkl() != original_phase.hkl()
1135            || saved_phase.d_spacing_angstrom() != original_phase.d_spacing_angstrom()
1136            || saved_phase.scale().to_bits() != original_phase.scale().to_bits()
1137        {
1138            return Err(TofMultiBankInstrumentError::InvalidCheckpoint(
1139                "TOF instrument checkpoint bank topology changed",
1140            ));
1141        }
1142        for d_spacing in saved_phase.d_spacing_angstrom() {
1143            TofProfileParameters::from_instrument(*d_spacing, saved.instrument)?;
1144        }
1145    }
1146    let selected = input
1147        .instrument_models
1148        .iter()
1149        .find(|model| model.bank_id() == &saved.bank_id);
1150    let saved_values = saved.instrument.values();
1151    let original_values = original.input.instrument.values();
1152    for parameter in TofInstrumentParameter::ALL {
1153        let bound = selected.and_then(|model| {
1154            model
1155                .bounds
1156                .iter()
1157                .find(|bound| bound.parameter == parameter)
1158        });
1159        match bound {
1160            Some(bound)
1161                if saved_values[parameter.index()] >= bound.lower
1162                    && saved_values[parameter.index()] <= bound.upper => {}
1163            Some(_) => {
1164                return Err(TofMultiBankInstrumentError::InvalidCheckpoint(
1165                    "TOF instrument checkpoint coefficient lies outside its bounds",
1166                ));
1167            }
1168            None if saved_values[parameter.index()].to_bits()
1169                == original_values[parameter.index()].to_bits() => {}
1170            None => {
1171                return Err(TofMultiBankInstrumentError::InvalidCheckpoint(
1172                    "unselected TOF instrument coefficient changed in a checkpoint",
1173                ));
1174            }
1175        }
1176    }
1177    match (&saved.background, &original.input.background) {
1178        (None, None) => Ok(()),
1179        (Some(saved), Some(original))
1180            if saved.background_id() == original.background_id()
1181                && saved
1182                    .domain_us()
1183                    .iter()
1184                    .zip(original.domain_us())
1185                    .all(|(saved, original)| saved.to_bits() == original.to_bits())
1186                && saved.coefficients().len() == original.coefficients().len() =>
1187        {
1188            Ok(())
1189        }
1190        _ => Err(TofMultiBankInstrumentError::InvalidCheckpoint(
1191            "TOF instrument checkpoint background contract changed",
1192        )),
1193    }
1194}
1195
1196/// Invalid bank-local TOF instrument request or numerical state.
1197#[derive(Debug)]
1198pub enum TofMultiBankInstrumentError {
1199    /// Bank-local or fixed-cell multi-bank contract failed.
1200    MultiBank(TofMultiBankError),
1201    /// TOF instrument/profile validation failed.
1202    Profile(TofError),
1203    /// Bank/parameter selection is inconsistent.
1204    InvalidModel(&'static str),
1205    /// Solver controls are invalid.
1206    InvalidOptions,
1207    /// Weighted instrument normal equations could not be solved finitely.
1208    LinearSolve,
1209    /// A continuation state disagrees with the immutable request contract.
1210    InvalidCheckpoint(&'static str),
1211    /// Checked allocation arithmetic overflowed.
1212    AllocationOverflow,
1213    /// Bounded runtime or event/checkpoint delivery failed.
1214    Runtime(RuntimeError),
1215    /// Internal array alignment failed after validated construction.
1216    InternalInvariant,
1217}
1218
1219impl Display for TofMultiBankInstrumentError {
1220    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
1221        match self {
1222            Self::MultiBank(error) => Display::fmt(error, formatter),
1223            Self::Profile(error) => Display::fmt(error, formatter),
1224            Self::InvalidModel(message) | Self::InvalidCheckpoint(message) => {
1225                formatter.write_str(message)
1226            }
1227            Self::InvalidOptions => {
1228                formatter.write_str("invalid multi-bank TOF instrument options")
1229            }
1230            Self::LinearSolve => {
1231                formatter.write_str("multi-bank TOF instrument linear solve failed")
1232            }
1233            Self::AllocationOverflow => {
1234                formatter.write_str("multi-bank TOF instrument allocation overflow")
1235            }
1236            Self::Runtime(error) => Display::fmt(error, formatter),
1237            Self::InternalInvariant => {
1238                formatter.write_str("multi-bank TOF instrument internal array invariant failed")
1239            }
1240        }
1241    }
1242}
1243
1244impl Error for TofMultiBankInstrumentError {
1245    fn source(&self) -> Option<&(dyn Error + 'static)> {
1246        match self {
1247            Self::MultiBank(error) => Some(error),
1248            Self::Profile(error) => Some(error),
1249            Self::Runtime(error) => Some(error),
1250            _ => None,
1251        }
1252    }
1253}
1254
1255impl From<TofMultiBankError> for TofMultiBankInstrumentError {
1256    fn from(value: TofMultiBankError) -> Self {
1257        Self::MultiBank(value)
1258    }
1259}
1260
1261impl From<TofLeBailError> for TofMultiBankInstrumentError {
1262    fn from(value: TofLeBailError) -> Self {
1263        Self::MultiBank(TofMultiBankError::LeBail(value))
1264    }
1265}
1266
1267impl From<TofError> for TofMultiBankInstrumentError {
1268    fn from(value: TofError) -> Self {
1269        Self::Profile(value)
1270    }
1271}
1272
1273impl From<RuntimeError> for TofMultiBankInstrumentError {
1274    fn from(value: RuntimeError) -> Self {
1275        Self::Runtime(value)
1276    }
1277}