1use 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#[derive(Clone, Copy, Debug, PartialEq)]
29pub struct TofInstrumentParameterBound {
30 pub parameter: TofInstrumentParameter,
32 pub lower: f64,
34 pub upper: f64,
36}
37
38impl TofInstrumentParameterBound {
39 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#[derive(Clone, Debug, PartialEq)]
64pub struct TofBankInstrumentModel {
65 bank_id: RecordId,
66 bounds: Vec<TofInstrumentParameterBound>,
67}
68
69impl TofBankInstrumentModel {
70 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 #[must_use]
87 pub const fn bank_id(&self) -> &RecordId {
88 &self.bank_id
89 }
90
91 #[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#[derive(Clone, Debug, PartialEq)]
122pub struct TofMultiBankInstrumentInput {
123 pub multibank: TofMultiBankInput,
125 pub instrument_models: Vec<TofBankInstrumentModel>,
127}
128
129impl TofMultiBankInstrumentInput {
130 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#[derive(Clone, Debug, PartialEq)]
184pub struct TofMultiBankInstrumentOptions {
185 pub lebail: TofLeBailOptions,
187 pub instrument_damping: f64,
189 pub max_scaled_instrument_step: f64,
191 pub max_instrument_backtracks: usize,
193 pub unresolved_correlation: f64,
195}
196
197impl TofMultiBankInstrumentOptions {
198 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 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#[derive(Clone, Debug, PartialEq)]
243pub struct TofBankInstrumentState {
244 pub bank_id: RecordId,
246 pub instrument: TofInstrument,
248}
249
250#[derive(Clone, Debug, PartialEq)]
252pub struct TofInstrumentParameterChange {
253 pub bank_id: RecordId,
255 pub parameter: TofInstrumentParameter,
257 pub before: f64,
259 pub after: f64,
261 pub scaled_change: f64,
263}
264
265#[derive(Clone, Debug, PartialEq)]
267pub struct TofInstrumentCorrelation {
268 pub left_bank_id: RecordId,
270 pub left_parameter: TofInstrumentParameter,
272 pub right_bank_id: RecordId,
274 pub right_parameter: TofInstrumentParameter,
276 pub correlation: f64,
278}
279
280#[derive(Clone, Debug, PartialEq)]
282pub struct TofInstrumentDiagnostics {
283 pub parameter_count: usize,
285 pub jacobian_rank: usize,
287 pub maximum_absolute_correlation: Option<f64>,
289 pub unresolved_correlations: Vec<TofInstrumentCorrelation>,
291}
292
293#[derive(Clone, Debug, PartialEq)]
295pub struct TofMultiBankInstrumentIterationRecord {
296 pub iteration: usize,
298 pub bank_metrics: Vec<ResidualEvaluation>,
300 pub metrics: TofMultiBankMetrics,
302 pub maximum_relative_intensity_change: f64,
304 pub maximum_absolute_background_change: f64,
306 pub scaled_instrument_step_norm: f64,
308 pub instrument_parameter_changes: Vec<TofInstrumentParameterChange>,
310}
311
312#[derive(Clone, Debug, PartialEq)]
314pub struct TofMultiBankInstrumentCheckpointBank {
315 pub bank_id: RecordId,
317 pub instrument: TofInstrument,
319 pub phases: Vec<TofLeBailPhase>,
321 pub background: Option<TofChebyshevBackground>,
323}
324
325#[derive(Clone, Debug, PartialEq)]
327pub struct TofMultiBankInstrumentCheckpoint {
328 pub completed_iterations: usize,
330 pub banks: Vec<TofMultiBankInstrumentCheckpointBank>,
332 pub history: Vec<TofMultiBankInstrumentIterationRecord>,
334}
335
336#[derive(Clone, Debug, PartialEq)]
338pub struct TofMultiBankInstrumentResult {
339 pub banks: Vec<TofMultiBankResultBank>,
341 pub instruments: Vec<TofBankInstrumentState>,
343 pub metrics: TofMultiBankMetrics,
345 pub diagnostics: TofInstrumentDiagnostics,
347 pub history: Vec<TofMultiBankInstrumentIterationRecord>,
349 pub termination_reason: TerminationReason,
351 pub checkpoint: TofMultiBankInstrumentCheckpoint,
353}
354
355impl TofMultiBankInstrumentCheckpoint {
356 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
392pub 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#[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 ¤t,
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#[derive(Debug)]
1198pub enum TofMultiBankInstrumentError {
1199 MultiBank(TofMultiBankError),
1201 Profile(TofError),
1203 InvalidModel(&'static str),
1205 InvalidOptions,
1207 LinearSolve,
1209 InvalidCheckpoint(&'static str),
1211 AllocationOverflow,
1213 Runtime(RuntimeError),
1215 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}