1use chrono::{Datelike, NaiveDate};
11
12pub(crate) mod atomic;
13pub mod convergence_reader;
14pub mod dictionary;
15pub mod error;
16pub mod fixed_delivery;
17pub mod generic_constraints_echo;
18pub mod hydro_models;
19pub mod manifest;
20pub mod parquet_config;
21pub mod policy;
22pub mod provenance;
23pub mod results_writer;
24pub mod scaling_report;
25pub(crate) mod schemas;
26pub mod simulation_writer;
27pub mod solver_stats_writer;
28pub mod stochastic;
29pub mod training_writer;
30
31pub use convergence_reader::{
32 ConvergenceSummary, read_convergence_summary, read_initial_gap_percent,
33};
34pub use dictionary::write_dictionaries;
35pub use error::OutputError;
36pub use fixed_delivery::{FixedDeliveryRow, write_fixed_delivery};
37pub use generic_constraints_echo::{GenericConstraintEchoRow, write_generic_constraint_echo};
38pub use hydro_models::{
39 read_hydro_model_summary, write_evaporation_models, write_fpha_deviation_points,
40 write_fpha_hyperplanes, write_hydro_model_summary,
41};
42pub use manifest::{
43 DeviationSummary, DeviationWorstEntry, DistributionInfo, HostLayout, MetadataBounds,
44 MetadataConfiguration, MetadataConvergence, MetadataCost, MetadataIterations,
45 MetadataProblemDimensions, MetadataRowPool, MetadataScenarios, MetadataSimulationSolveStats,
46 MetadataTrainingSolveStats, OutputContext, SetupTimings, SimulationMetadata, TrainingMetadata,
47 default_bounds, get_hostname, now_iso8601, read_simulation_metadata, read_training_metadata,
48 write_simulation_metadata, write_training_metadata,
49};
50pub use parquet_config::ParquetWriterConfig;
51pub use provenance::{read_provenance_report, write_provenance_report};
52pub use results_writer::{write_results, write_simulation_results, write_training_results};
53pub use scaling_report::write_scaling_report;
54pub use simulation_writer::SimulationParquetWriter;
55pub use solver_stats_writer::{SolverStatsRow, write_simulation_solver_stats, write_solver_stats};
56pub use stochastic::{
57 FittingReductionEntry, FittingReport, HydroFittingEntry, write_correlation_json,
58 write_fitting_report, write_inflow_annual_component, write_inflow_ar_coefficients,
59 write_inflow_seasonal_stats, write_load_seasonal_stats, write_noise_openings,
60};
61pub use training_writer::{TrainingParquetWriter, write_row_selection_records};
62
63pub(crate) fn date32_days(date: NaiveDate) -> i32 {
66 let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).map_or(0, |e| e.num_days_from_ce());
67 date.num_days_from_ce() - epoch
68}
69
70#[derive(Debug, Clone)]
77pub struct IterationRecord {
78 pub iteration: u32,
80
81 pub lower_bound: f64,
83
84 pub upper_bound: f64,
87
88 pub upper_bound_std: f64,
91
92 pub gap_percent: Option<f64>,
96
97 pub cuts_added: u32,
99
100 pub cuts_removed: u32,
102
103 pub cuts_active: u32,
105
106 pub time_forward_ms: u64,
108
109 pub time_backward_ms: u64,
111
112 pub time_total_ms: u64,
114
115 pub time_forward_wall_ms: u64,
117
118 pub time_backward_wall_ms: u64,
120
121 pub time_cut_selection_ms: u64,
123
124 pub time_mpi_allreduce_ms: u64,
126
127 pub time_cut_sync_ms: u64,
129
130 pub time_lower_bound_ms: u64,
132
133 pub time_state_exchange_ms: u64,
135
136 pub time_cut_batch_build_ms: u64,
138
139 pub time_bwd_setup_ms: u64,
141
142 pub time_bwd_load_imbalance_ms: u64,
144
145 pub time_bwd_scheduling_overhead_ms: u64,
147
148 pub time_fwd_setup_ms: u64,
150
151 pub time_fwd_load_imbalance_ms: u64,
153
154 pub time_fwd_scheduling_overhead_ms: u64,
156
157 pub time_overhead_ms: u64,
160
161 pub forward_passes: u32,
163
164 pub lp_solves: u32,
166
167 pub solve_time_ms: f64,
169
170 pub mean_rows_in_lp: f64,
176}
177
178#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq)]
182pub struct RowPoolStatistics {
183 pub total_generated: u64,
185
186 pub total_active: u64,
188
189 pub peak_active: u64,
191
192 pub cuts_active: u64,
194
195 pub rows_in_lp_total: u64,
201
202 pub rows_in_lp_solve_count: u64,
205
206 pub rows_in_lp_max: u64,
209
210 pub total_loaded: u64,
213}
214
215#[derive(Debug, Clone)]
220pub struct RowSelectionRecord {
221 pub iteration: u32,
223 pub stage: u32,
225 pub cuts_populated: u32,
227 pub cuts_active_before: u32,
229 pub cuts_deactivated: u32,
231 pub cuts_reactivated: u32,
233 pub cuts_active_after: u32,
235 pub selection_time_ms: f64,
237 pub budget_evicted: Option<u32>,
241 pub active_after_budget: Option<u32>,
245}
246
247#[derive(Debug, Clone)]
260pub struct WorkerTimingRecord {
261 pub iteration: u32,
263 pub rank: i32,
265 pub worker_id: Option<i32>,
267 pub timings: [u64; 16],
272}
273
274#[derive(Debug, Clone)]
280pub struct TrainingOutput {
281 pub convergence_records: Vec<IterationRecord>,
283
284 pub final_lower_bound: f64,
286
287 pub final_upper_bound: Option<f64>,
291
292 pub final_gap_percent: Option<f64>,
296
297 pub final_upper_bound_std: Option<f64>,
303
304 pub final_upper_bound_kind: String,
308
309 pub iterations_completed: u32,
311
312 pub converged: bool,
314
315 pub termination_reason: String,
317
318 pub total_time_ms: u64,
320
321 pub cut_stats: RowPoolStatistics,
323
324 pub cut_selection_records: Vec<RowSelectionRecord>,
329
330 pub worker_timing_records: Vec<WorkerTimingRecord>,
339
340 pub training_solve_stats: MetadataTrainingSolveStats,
346}
347
348#[derive(Debug, Clone)]
355pub struct SimulationOutput {
356 pub n_scenarios: u32,
358
359 pub completed: u32,
361
362 pub failed: u32,
364
365 pub total_time_ms: u64,
367
368 pub partitions_written: Vec<String>,
373
374 pub cost: Option<MetadataCost>,
381
382 pub solve_stats: MetadataSimulationSolveStats,
388}
389
390impl SimulationOutput {
391 #[must_use]
413 pub fn merge(outputs: &[Self]) -> Self {
414 if outputs.is_empty() {
415 return Self {
416 n_scenarios: 0,
417 completed: 0,
418 failed: 0,
419 total_time_ms: 0,
420 partitions_written: Vec::new(),
421 cost: None,
422 solve_stats: MetadataSimulationSolveStats::default(),
423 };
424 }
425
426 let n_scenarios = outputs.iter().map(|o| o.n_scenarios).sum();
427 let completed = outputs.iter().map(|o| o.completed).sum();
428 let failed = outputs.iter().map(|o| o.failed).sum();
429 let total_time_ms = outputs.iter().map(|o| o.total_time_ms).max().unwrap_or(0);
430
431 let mut partitions_written: Vec<String> = outputs
432 .iter()
433 .flat_map(|o| o.partitions_written.iter().cloned())
434 .collect();
435 partitions_written.sort();
436
437 let cost = outputs.iter().find_map(|o| o.cost.clone());
438
439 let solve_stats = merge_simulation_solve_stats(outputs);
440
441 Self {
442 n_scenarios,
443 completed,
444 failed,
445 total_time_ms,
446 partitions_written,
447 cost,
448 solve_stats,
449 }
450 }
451}
452
453fn sum_optional_u64(
459 outputs: &[SimulationOutput],
460 field: impl Fn(&MetadataSimulationSolveStats) -> Option<u64>,
461) -> Option<u64> {
462 let mut any = false;
463 let mut total: u64 = 0;
464 for output in outputs {
465 if let Some(value) = field(&output.solve_stats) {
466 any = true;
467 total = total.saturating_add(value);
468 }
469 }
470 any.then_some(total)
471}
472
473fn merge_simulation_solve_stats(outputs: &[SimulationOutput]) -> MetadataSimulationSolveStats {
479 let mut solve_seconds_any = false;
480 let mut solve_seconds_total: f64 = 0.0;
481 for output in outputs {
482 if let Some(value) = output.solve_stats.solve_seconds {
483 solve_seconds_any = true;
484 solve_seconds_total += value;
485 }
486 }
487
488 let parallelism = outputs
489 .iter()
490 .filter_map(|o| o.solve_stats.parallelism)
491 .max();
492
493 MetadataSimulationSolveStats {
494 total_lp_solves: sum_optional_u64(outputs, |s| s.total_lp_solves),
495 first_try: sum_optional_u64(outputs, |s| s.first_try),
496 retried: sum_optional_u64(outputs, |s| s.retried),
497 failed: sum_optional_u64(outputs, |s| s.failed),
498 solve_seconds: solve_seconds_any.then_some(solve_seconds_total),
499 parallelism,
500 }
501}
502
503#[cfg(test)]
504#[allow(
505 clippy::unwrap_used,
506 clippy::expect_used,
507 clippy::float_cmp,
508 clippy::cast_possible_truncation
509)]
510mod tests {
511 use super::*;
512
513 #[test]
514 fn training_output_construction_and_field_access() {
515 let records: Vec<IterationRecord> = (1..=5)
516 .map(|i| IterationRecord {
517 iteration: i,
518 lower_bound: 1.0,
519 upper_bound: 2.0,
520 upper_bound_std: 0.1,
521 gap_percent: Some(50.0),
522 cuts_added: 10,
523 cuts_removed: 2,
524 cuts_active: 8,
525 time_forward_ms: 100,
526 time_backward_ms: 200,
527 time_total_ms: 300,
528 forward_passes: 4,
529 lp_solves: 40,
530 time_forward_wall_ms: 100,
531 time_backward_wall_ms: 200,
532 time_cut_selection_ms: 0,
533 time_mpi_allreduce_ms: 0,
534 time_cut_sync_ms: 0,
535 time_lower_bound_ms: 0,
536 time_state_exchange_ms: 0,
537 time_cut_batch_build_ms: 0,
538 time_bwd_setup_ms: 0,
539 time_bwd_load_imbalance_ms: 0,
540 time_bwd_scheduling_overhead_ms: 0,
541 time_fwd_setup_ms: 0,
542 time_fwd_load_imbalance_ms: 0,
543 time_fwd_scheduling_overhead_ms: 0,
544 time_overhead_ms: 0,
545 solve_time_ms: 0.0,
546 mean_rows_in_lp: 0.0,
547 })
548 .collect();
549 let output = TrainingOutput {
550 convergence_records: records,
551 final_lower_bound: 50.0,
552 final_upper_bound: Some(52.0),
553 final_gap_percent: Some(3.85),
554 final_upper_bound_std: Some(0.5),
555 final_upper_bound_kind: "statistical".to_string(),
556 iterations_completed: 5,
557 converged: true,
558 termination_reason: "relative gap < 1%".to_string(),
559 total_time_ms: 12_000,
560 cut_stats: RowPoolStatistics {
561 total_generated: 300,
562 total_active: 120,
563 peak_active: 150,
564 cuts_active: 120,
565 rows_in_lp_total: 0,
566 rows_in_lp_solve_count: 0,
567 rows_in_lp_max: 0,
568 total_loaded: 0,
569 },
570 cut_selection_records: vec![],
571 worker_timing_records: vec![],
572 training_solve_stats: MetadataTrainingSolveStats::default(),
573 };
574
575 assert_eq!(output.convergence_records.len(), 5);
576 assert_eq!(output.final_lower_bound, 50.0);
577 assert_eq!(output.final_upper_bound, Some(52.0));
578 assert_eq!(output.final_gap_percent, Some(3.85));
579 assert_eq!(output.final_upper_bound_std, Some(0.5));
580 assert_eq!(output.iterations_completed, 5);
581 assert!(output.converged);
582 assert_eq!(output.termination_reason, "relative gap < 1%");
583 assert_eq!(output.total_time_ms, 12_000);
584 assert_eq!(output.cut_stats.total_generated, 300);
585 assert_eq!(output.cut_stats.total_active, 120);
586 assert_eq!(output.cut_stats.peak_active, 150);
587 }
588
589 #[test]
590 fn iteration_record_construction_and_field_access() {
591 let record = IterationRecord {
592 iteration: 7,
593 lower_bound: 10.5,
594 upper_bound: 11.0,
595 upper_bound_std: 0.25,
596 gap_percent: Some(4.55),
597 cuts_added: 15,
598 cuts_removed: 3,
599 cuts_active: 42,
600 time_forward_ms: 150,
601 time_backward_ms: 250,
602 time_total_ms: 400,
603 forward_passes: 8,
604 lp_solves: 80,
605 time_forward_wall_ms: 150,
606 time_backward_wall_ms: 250,
607 time_cut_selection_ms: 5,
608 time_mpi_allreduce_ms: 3,
609 time_cut_sync_ms: 2,
610 time_lower_bound_ms: 4,
611 time_state_exchange_ms: 0,
612 time_cut_batch_build_ms: 0,
613 time_bwd_setup_ms: 0,
614 time_bwd_load_imbalance_ms: 0,
615 time_bwd_scheduling_overhead_ms: 0,
616 time_fwd_setup_ms: 0,
617 time_fwd_load_imbalance_ms: 0,
618 time_fwd_scheduling_overhead_ms: 0,
619 time_overhead_ms: 400u64.saturating_sub(150 + 250 + 5 + 3 + 4),
620 solve_time_ms: 0.0,
621 mean_rows_in_lp: 0.0,
622 };
623
624 assert_eq!(record.iteration, 7);
625 assert_eq!(record.lower_bound, 10.5);
626 assert_eq!(record.upper_bound, 11.0);
627 assert_eq!(record.upper_bound_std, 0.25);
628 assert_eq!(record.gap_percent, Some(4.55));
629 assert_eq!(record.cuts_added, 15);
630 assert_eq!(record.cuts_removed, 3);
631 assert_eq!(record.cuts_active, 42);
632 assert_eq!(record.time_forward_ms, 150);
633 assert_eq!(record.time_backward_ms, 250);
634 assert_eq!(record.time_total_ms, 400);
635 assert_eq!(record.forward_passes, 8);
636 assert_eq!(record.lp_solves, 80);
637 assert_eq!(record.time_forward_wall_ms, 150);
638 assert_eq!(record.time_backward_wall_ms, 250);
639 assert_eq!(record.time_cut_selection_ms, 5);
640 assert_eq!(record.time_mpi_allreduce_ms, 3);
641 assert_eq!(record.time_cut_sync_ms, 2);
642 assert_eq!(record.time_lower_bound_ms, 4);
643 }
644
645 #[test]
646 fn simulation_output_construction_and_field_access() {
647 let output = SimulationOutput {
648 n_scenarios: 100,
649 completed: 100,
650 failed: 0,
651 total_time_ms: 3_200,
652 partitions_written: vec![
653 "simulation/costs/year=2030/part-00.parquet".to_string(),
654 "simulation/costs/year=2031/part-00.parquet".to_string(),
655 ],
656 cost: None,
657 solve_stats: MetadataSimulationSolveStats::default(),
658 };
659
660 assert_eq!(output.n_scenarios, 100);
661 assert_eq!(output.completed, 100);
662 assert_eq!(output.failed, 0);
663 assert_eq!(output.total_time_ms, 3_200);
664 assert_eq!(output.partitions_written.len(), 2);
665 }
666
667 #[test]
668 fn row_pool_statistics_construction() {
669 let stats = RowPoolStatistics {
670 total_generated: 500,
671 total_active: 200,
672 peak_active: 250,
673 cuts_active: 200,
674 rows_in_lp_total: 0,
675 rows_in_lp_solve_count: 0,
676 rows_in_lp_max: 0,
677 total_loaded: 0,
678 };
679
680 assert_eq!(stats.total_generated, 500);
681 assert_eq!(stats.total_active, 200);
682 assert_eq!(stats.peak_active, 250);
683 assert_eq!(stats.cuts_active, 200);
684 }
685
686 #[test]
687 fn row_pool_statistics_serializes_with_new_fields() {
688 let stats = RowPoolStatistics {
689 total_generated: 10,
690 total_active: 7,
691 peak_active: 9,
692 cuts_active: 7,
693 rows_in_lp_total: 30,
694 rows_in_lp_solve_count: 6,
695 rows_in_lp_max: 8,
696 total_loaded: 3,
697 };
698 let json = serde_json::to_string(&stats).expect("serialization must succeed");
699 assert!(
700 !json.contains("\"cuts_in_lp\""),
701 "JSON must not contain cuts_in_lp key"
702 );
703 assert!(
704 json.contains("\"cuts_active\""),
705 "JSON must contain cuts_active key"
706 );
707 for key in [
708 "\"rows_in_lp_total\"",
709 "\"rows_in_lp_solve_count\"",
710 "\"rows_in_lp_max\"",
711 "\"total_loaded\"",
712 ] {
713 assert!(json.contains(key), "JSON must contain {key}");
714 }
715 }
716
717 #[test]
718 fn test_merge_empty_slice() {
719 let merged = SimulationOutput::merge(&[]);
720 assert_eq!(merged.n_scenarios, 0);
721 assert_eq!(merged.completed, 0);
722 assert_eq!(merged.failed, 0);
723 assert_eq!(merged.total_time_ms, 0);
724 assert!(merged.partitions_written.is_empty());
725 }
726
727 #[test]
728 fn test_merge_single_output() {
729 let output = SimulationOutput {
730 n_scenarios: 5,
731 completed: 4,
732 failed: 1,
733 total_time_ms: 1000,
734 partitions_written: vec!["simulation/costs/scenario_id=0000/data.parquet".to_string()],
735 cost: None,
736 solve_stats: MetadataSimulationSolveStats::default(),
737 };
738 let merged = SimulationOutput::merge(std::slice::from_ref(&output));
739 assert_eq!(merged.n_scenarios, 5);
740 assert_eq!(merged.completed, 4);
741 assert_eq!(merged.failed, 1);
742 assert_eq!(merged.total_time_ms, 1000);
743 assert_eq!(merged.partitions_written, output.partitions_written);
744 }
745
746 #[test]
747 fn test_merge_two_outputs() {
748 let a = SimulationOutput {
749 n_scenarios: 3,
750 completed: 3,
751 failed: 0,
752 total_time_ms: 500,
753 partitions_written: vec![
754 "simulation/costs/scenario_id=0000/data.parquet".to_string(),
755 "simulation/costs/scenario_id=0001/data.parquet".to_string(),
756 ],
757 cost: None,
758 solve_stats: MetadataSimulationSolveStats::default(),
759 };
760 let b = SimulationOutput {
761 n_scenarios: 2,
762 completed: 1,
763 failed: 1,
764 total_time_ms: 800,
765 partitions_written: vec!["simulation/costs/scenario_id=0002/data.parquet".to_string()],
766 cost: None,
767 solve_stats: MetadataSimulationSolveStats::default(),
768 };
769 let merged = SimulationOutput::merge(&[a, b]);
770 assert_eq!(merged.n_scenarios, 5);
771 assert_eq!(merged.completed, 4);
772 assert_eq!(merged.failed, 1);
773 assert_eq!(merged.total_time_ms, 800);
775 assert_eq!(merged.partitions_written.len(), 3);
776 }
777
778 #[test]
779 fn test_merge_partitions_sorted() {
780 let a = SimulationOutput {
781 n_scenarios: 1,
782 completed: 1,
783 failed: 0,
784 total_time_ms: 100,
785 partitions_written: vec![
786 "simulation/hydros/scenario_id=0002/data.parquet".to_string(),
787 "simulation/costs/scenario_id=0002/data.parquet".to_string(),
788 ],
789 cost: None,
790 solve_stats: MetadataSimulationSolveStats::default(),
791 };
792 let b = SimulationOutput {
793 n_scenarios: 1,
794 completed: 1,
795 failed: 0,
796 total_time_ms: 200,
797 partitions_written: vec![
798 "simulation/costs/scenario_id=0001/data.parquet".to_string(),
799 "simulation/hydros/scenario_id=0001/data.parquet".to_string(),
800 ],
801 cost: None,
802 solve_stats: MetadataSimulationSolveStats::default(),
803 };
804 let merged = SimulationOutput::merge(&[a, b]);
805 let expected = vec![
806 "simulation/costs/scenario_id=0001/data.parquet".to_string(),
807 "simulation/costs/scenario_id=0002/data.parquet".to_string(),
808 "simulation/hydros/scenario_id=0001/data.parquet".to_string(),
809 "simulation/hydros/scenario_id=0002/data.parquet".to_string(),
810 ];
811 assert_eq!(merged.partitions_written, expected);
812 }
813
814 #[test]
815 fn simulation_output_merge_combines_solve_stats_order_invariant() {
816 let a = SimulationOutput {
817 n_scenarios: 2,
818 completed: 2,
819 failed: 0,
820 total_time_ms: 500,
821 partitions_written: vec![],
822 cost: Some(MetadataCost {
823 mean_cost: 100.0,
824 std_cost: 10.0,
825 }),
826 solve_stats: MetadataSimulationSolveStats {
827 total_lp_solves: Some(40),
828 first_try: Some(35),
829 retried: Some(5),
830 failed: Some(0),
831 solve_seconds: Some(1.5),
832 parallelism: Some(4),
833 },
834 };
835 let b = SimulationOutput {
836 n_scenarios: 3,
837 completed: 3,
838 failed: 0,
839 total_time_ms: 800,
840 partitions_written: vec![],
841 cost: Some(MetadataCost {
842 mean_cost: 200.0,
843 std_cost: 20.0,
844 }),
845 solve_stats: MetadataSimulationSolveStats {
846 total_lp_solves: Some(60),
847 first_try: Some(50),
848 retried: Some(8),
849 failed: Some(2),
850 solve_seconds: Some(2.5),
851 parallelism: Some(8),
852 },
853 };
854
855 let merged_ab = SimulationOutput::merge(&[a.clone(), b.clone()]);
856 let merged_ba = SimulationOutput::merge(&[b, a]);
857
858 assert_eq!(
859 merged_ab.solve_stats.total_lp_solves,
860 merged_ba.solve_stats.total_lp_solves
861 );
862 assert_eq!(merged_ab.solve_stats.total_lp_solves, Some(100));
863 assert_eq!(merged_ab.solve_stats.first_try, Some(85));
864 assert_eq!(
865 merged_ab.solve_stats.first_try,
866 merged_ba.solve_stats.first_try
867 );
868 assert_eq!(merged_ab.solve_stats.retried, Some(13));
869 assert_eq!(merged_ab.solve_stats.retried, merged_ba.solve_stats.retried);
870 assert_eq!(merged_ab.solve_stats.failed, Some(2));
871 assert_eq!(merged_ab.solve_stats.failed, merged_ba.solve_stats.failed);
872
873 assert_eq!(merged_ab.solve_stats.solve_seconds, Some(4.0));
874 assert_eq!(
875 merged_ab.solve_stats.solve_seconds,
876 merged_ba.solve_stats.solve_seconds
877 );
878
879 assert_eq!(merged_ab.solve_stats.parallelism, Some(8));
880 assert_eq!(
881 merged_ab.solve_stats.parallelism,
882 merged_ba.solve_stats.parallelism
883 );
884
885 assert_eq!(
886 merged_ab.cost.as_ref().map(|c| c.mean_cost),
887 Some(100.0),
888 "first-present cost wins in [a, b] order"
889 );
890 assert_eq!(
891 merged_ba.cost.as_ref().map(|c| c.mean_cost),
892 Some(200.0),
893 "first-present cost wins in [b, a] order"
894 );
895 }
896
897 #[test]
898 fn simulation_output_merge_solve_stats_none_when_no_input_records() {
899 let a = SimulationOutput {
900 n_scenarios: 1,
901 completed: 1,
902 failed: 0,
903 total_time_ms: 100,
904 partitions_written: vec![],
905 cost: None,
906 solve_stats: MetadataSimulationSolveStats::default(),
907 };
908 let merged = SimulationOutput::merge(std::slice::from_ref(&a));
909 assert_eq!(merged.solve_stats.total_lp_solves, None);
910 assert_eq!(merged.solve_stats.solve_seconds, None);
911 assert_eq!(merged.solve_stats.parallelism, None);
912 assert!(merged.cost.is_none());
913 }
914}