1use std::fmt;
4use std::sync::Arc;
5use std::time::Duration;
6
7use crate::explore::reducer::{ReducerError, ReducerKind, ReducerPlan, ReducerState};
8use crate::explore::spec::{MeasureSettings, RunSettings};
9use crate::explore::stop::{StopCondition, StopError};
10use crate::export::{StatColumns, StatsWriteError};
11use crate::view::StatEntry;
12
13#[derive(Debug, Clone, PartialEq)]
17pub struct MeasurePlan {
18 warmup: u64,
19 total: u64,
20 stats_every: u64,
21 series_every: u64,
22 columns: StatColumns,
23 reducers: ReducerPlan,
24 stop: Option<StopCondition>,
25}
26
27impl MeasurePlan {
28 pub fn check(run: &RunSettings, measure: &MeasureSettings) -> Result<(), MeasureError> {
38 if run.warmup.checked_add(run.steps).is_none() {
39 return Err(MeasureError::TooManyTicks {
40 warmup: run.warmup,
41 steps: run.steps,
42 });
43 }
44 if let Some(timeout) = run.timeout
45 && Duration::try_from_secs_f64(timeout.as_secs_f64()).is_err()
46 {
47 return Err(MeasureError::TimeoutTooLong { timeout });
48 }
49 if measure.stats_every == 0 {
50 return Err(MeasureError::ZeroStatsEvery);
51 }
52 if !measure.series_every.is_multiple_of(measure.stats_every) {
53 return Err(MeasureError::SeriesOffCadence {
54 series_every: measure.series_every,
55 stats_every: measure.stats_every,
56 });
57 }
58 if let Some(stop) = &run.stop {
59 stop.check_threshold().map_err(MeasureError::Stop)?;
60 }
61 for reducer in &measure.reducers {
62 match reducer.kind {
63 ReducerKind::FirstCrossing(comparison) => comparison.check().map_err(|source| {
64 MeasureError::Reducer(ReducerError::Comparison {
65 raw: reducer.kind.to_string(),
66 source,
67 })
68 })?,
69 ReducerKind::WindowMean { start, end } if start > end => {
70 return Err(MeasureError::Reducer(ReducerError::BadWindow {
71 raw: reducer.kind.to_string(),
72 }));
73 }
74 _ => {}
75 }
76 }
77 Ok(())
78 }
79
80 pub fn new(run: &RunSettings, measure: &MeasureSettings, columns: StatColumns) -> Result<Self, MeasureError> {
87 Self::check(run, measure)?;
88 let reducers =
89 ReducerPlan::bind(&columns, &measure.reducers, measure.default_reducers).map_err(MeasureError::Reducer)?;
90 let stop = run
91 .stop
92 .as_ref()
93 .map(|stop| StopCondition::bind(stop, &columns))
94 .transpose()
95 .map_err(MeasureError::Stop)?;
96 Ok(Self {
97 warmup: run.warmup,
98 total: run.warmup + run.steps,
99 stats_every: measure.stats_every,
100 series_every: measure.series_every,
101 columns,
102 reducers,
103 stop,
104 })
105 }
106
107 pub fn warmup(&self) -> u64 {
109 self.warmup
110 }
111
112 pub fn total(&self) -> u64 {
114 self.total
115 }
116
117 pub fn stats_every(&self) -> u64 {
119 self.stats_every
120 }
121
122 pub fn series_every(&self) -> u64 {
124 self.series_every
125 }
126
127 pub fn columns(&self) -> &StatColumns {
129 &self.columns
130 }
131
132 pub fn reducers(&self) -> &ReducerPlan {
134 &self.reducers
135 }
136
137 pub fn stop(&self) -> Option<&StopCondition> {
139 self.stop.as_ref()
140 }
141
142 pub fn first_sample(&self) -> u64 {
144 self.warmup
145 }
146
147 pub fn next_sample(&self, tick: u64) -> Option<u64> {
149 if tick >= self.total {
150 return None;
151 }
152 if tick < self.warmup {
153 return Some(self.warmup);
154 }
155 let intervals = (tick - self.warmup) / self.stats_every + 1;
156 let next = self.warmup.saturating_add(intervals.saturating_mul(self.stats_every));
157 Some(next.min(self.total))
158 }
159
160 pub fn is_sample(&self, tick: u64) -> bool {
162 tick == self.total
163 || (tick >= self.warmup && tick < self.total && (tick - self.warmup).is_multiple_of(self.stats_every))
164 }
165
166 pub fn sample_ticks(&self) -> impl Iterator<Item = u64> + '_ {
168 std::iter::successors(Some(self.first_sample()), |&tick| self.next_sample(tick))
169 }
170
171 pub fn sample_count(&self) -> u64 {
173 count_on_cadence(self.total - self.warmup, self.stats_every)
174 }
175
176 pub fn series_row_count(&self) -> u64 {
178 if self.series_every == 0 {
179 0
180 } else {
181 count_on_cadence(self.total - self.warmup, self.series_every)
182 }
183 }
184
185 fn keeps_in_series(&self, tick: u64) -> bool {
187 self.series_every != 0 && tick >= self.warmup && (tick - self.warmup).is_multiple_of(self.series_every)
188 }
189}
190
191fn count_on_cadence(span: u64, every: u64) -> u64 {
195 (span / every).saturating_add(1 + u64::from(!span.is_multiple_of(every)))
196}
197
198#[derive(Debug, Clone, PartialEq, Eq)]
200pub enum MeasureError {
201 TooManyTicks {
203 warmup: u64,
205 steps: u64,
207 },
208 TimeoutTooLong {
210 timeout: Duration,
212 },
213 ZeroStatsEvery,
215 SeriesOffCadence {
217 series_every: u64,
219 stats_every: u64,
221 },
222 Reducer(ReducerError),
224 Stop(StopError),
226}
227
228impl fmt::Display for MeasureError {
229 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
230 match self {
231 Self::TooManyTicks { warmup, steps } => {
232 write!(
233 f,
234 "a warm-up of {warmup} ticks plus {steps} steps exceeds the largest tick"
235 )
236 }
237 Self::TimeoutTooLong { timeout } => {
238 write!(
239 f,
240 "a timeout of {} seconds is too long to record",
241 timeout.as_secs_f64()
242 )
243 }
244 Self::ZeroStatsEvery => write!(f, "stats_every must be at least 1"),
245 Self::SeriesOffCadence {
246 series_every,
247 stats_every,
248 } => write!(
249 f,
250 "series_every {series_every} is not a multiple of stats_every {stats_every}"
251 ),
252 Self::Reducer(_) => write!(f, "reducer"),
253 Self::Stop(_) => write!(f, "stop condition"),
254 }
255 }
256}
257
258impl std::error::Error for MeasureError {
259 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
260 match self {
261 Self::Reducer(error) => Some(error),
262 Self::Stop(error) => Some(error),
263 _ => None,
264 }
265 }
266}
267
268#[derive(Debug, Clone, Default, PartialEq)]
270pub struct SeriesBuffer {
271 width: usize,
272 ticks: Vec<u64>,
273 values: Vec<f64>,
274}
275
276impl SeriesBuffer {
277 pub fn new(width: usize) -> Self {
279 Self {
280 width,
281 ticks: Vec::new(),
282 values: Vec::new(),
283 }
284 }
285
286 pub fn width(&self) -> usize {
288 self.width
289 }
290
291 pub fn len(&self) -> usize {
293 self.ticks.len()
294 }
295
296 pub fn is_empty(&self) -> bool {
298 self.ticks.is_empty()
299 }
300
301 pub fn push(&mut self, tick: u64, values: &[f64]) {
307 assert_eq!(values.len(), self.width, "a series row holds one value per column");
308 self.ticks.push(tick);
309 self.values.extend_from_slice(values);
310 }
311
312 pub fn shrink_to_fit(&mut self) {
314 self.ticks.shrink_to_fit();
315 self.values.shrink_to_fit();
316 }
317
318 pub fn ticks(&self) -> &[u64] {
320 &self.ticks
321 }
322
323 pub fn row(&self, i: usize) -> &[f64] {
329 &self.values[i * self.width..(i + 1) * self.width]
330 }
331
332 pub fn rows(&self) -> impl Iterator<Item = (u64, &[f64])> + '_ {
334 self.ticks.iter().enumerate().map(|(i, &tick)| (tick, self.row(i)))
335 }
336}
337
338#[derive(Debug, Clone, PartialEq, Eq)]
340pub struct NonFiniteSample {
341 pub column: String,
343 pub tick: u64,
345}
346
347impl fmt::Display for NonFiniteSample {
348 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
349 write!(f, "{} is not finite at tick {}", self.column, self.tick)
350 }
351}
352
353#[derive(Debug, Clone, PartialEq)]
355pub struct Measured {
356 pub reducers: Vec<Option<f64>>,
359 pub series: SeriesBuffer,
361 pub non_finite: Option<NonFiniteSample>,
363}
364
365#[derive(Debug, Clone)]
370pub struct Sampler {
371 plan: Arc<MeasurePlan>,
372 reducers: ReducerState,
373 series: SeriesBuffer,
374 row: Vec<f64>,
376 incoming: Vec<f64>,
378 unkept_tick: Option<u64>,
380 non_finite: Option<NonFiniteSample>,
381}
382
383impl Sampler {
384 pub fn new(plan: Arc<MeasurePlan>) -> Self {
386 Self {
387 reducers: ReducerState::new(&plan.reducers),
388 series: SeriesBuffer::new(plan.columns.len()),
389 row: Vec::with_capacity(plan.columns.len()),
390 incoming: Vec::with_capacity(plan.columns.len()),
391 unkept_tick: None,
392 non_finite: None,
393 plan,
394 }
395 }
396
397 pub fn plan(&self) -> &MeasurePlan {
399 &self.plan
400 }
401
402 pub fn push(&mut self, tick: u64, stats: &[StatEntry]) -> Result<bool, StatsWriteError> {
412 debug_assert!(self.plan.is_sample(tick), "tick {tick} is not a sampled tick");
413 self.plan.columns.extract(tick, stats, &mut self.incoming)?;
414 std::mem::swap(&mut self.row, &mut self.incoming);
415 self.unkept_tick = None;
416 if self.non_finite.is_none()
417 && let Some(column) = self.row.iter().position(|value| !value.is_finite())
418 {
419 self.non_finite = Some(NonFiniteSample {
420 column: self.plan.columns.name(column).to_owned(),
421 tick,
422 });
423 }
424 self.reducers.push(&self.plan.reducers, tick, &self.row);
425 if self.plan.keeps_in_series(tick) {
426 self.series.push(tick, &self.row);
427 } else if self.plan.series_every != 0 {
428 self.unkept_tick = Some(tick);
429 }
430 Ok(self.plan.stop.is_some_and(|stop| stop.holds(tick, &self.row)))
431 }
432
433 pub fn finish(mut self) -> Measured {
435 if let Some(tick) = self.unkept_tick {
436 self.series.push(tick, &self.row);
437 }
438 Measured {
439 reducers: self.reducers.finish(&self.plan.reducers),
440 series: self.series,
441 non_finite: self.non_finite,
442 }
443 }
444}
445
446#[cfg(test)]
447mod tests {
448 use std::sync::Arc;
449 use std::time::Duration;
450
451 use super::{MeasureError, MeasurePlan, NonFiniteSample, Sampler};
452 use crate::explore::reducer::{ReducerError, ReducerKind, ReducerSpec};
453 use crate::explore::spec::{MeasureSettings, RunSettings};
454 use crate::explore::stop::{Comparator, Comparison, ComparisonError, StopError, StopSpec};
455 use crate::export::StatColumns;
456 use crate::helpers::stat;
457 use crate::view::StatEntry;
458
459 const COLOR: [u8; 4] = [0, 0, 0, 255];
460
461 fn sample(infected: f64, recovered: f64) -> Vec<StatEntry> {
462 vec![stat("Infected", infected, COLOR), stat("Recovered", recovered, COLOR)]
463 }
464
465 fn plan(warmup: u64, steps: u64, stats_every: u64, series_every: u64) -> MeasurePlan {
466 let run = RunSettings {
467 steps,
468 warmup,
469 ..RunSettings::default()
470 };
471 let measure = MeasureSettings {
472 stats_every,
473 series_every,
474 ..MeasureSettings::default()
475 };
476 MeasurePlan::new(&run, &measure, StatColumns::plan(&sample(0.0, 0.0))).expect("a valid cadence")
477 }
478
479 #[test]
480 fn sampled_ticks_include_the_final_tick() {
481 let off_cadence = plan(3, 10, 4, 4);
482 assert_eq!(off_cadence.sample_ticks().collect::<Vec<_>>(), [3, 7, 11, 13]);
483 assert_eq!(off_cadence.sample_count(), 4);
484 assert!(off_cadence.is_sample(13) && off_cadence.is_sample(7));
485 assert!(!off_cadence.is_sample(2) && !off_cadence.is_sample(12) && !off_cadence.is_sample(15));
486 assert_eq!(off_cadence.next_sample(0), Some(3));
487 assert_eq!(off_cadence.next_sample(12), Some(13));
488 assert_eq!(off_cadence.next_sample(13), None);
489
490 let on_cadence = plan(0, 8, 4, 4);
491 assert_eq!(on_cadence.sample_ticks().collect::<Vec<_>>(), [0, 4, 8]);
492 assert_eq!(on_cadence.sample_count(), 3);
493
494 let no_steps = plan(5, 0, 4, 4);
495 assert_eq!(no_steps.sample_ticks().collect::<Vec<_>>(), [5]);
496 assert_eq!(no_steps.sample_count(), 1);
497 }
498
499 #[test]
500 fn the_series_is_thinned_and_ends_on_the_final_sample() {
501 let thinned = Arc::new(plan(0, 10, 2, 4));
502 let mut sampler = Sampler::new(Arc::clone(&thinned));
503 for tick in thinned.sample_ticks() {
504 sampler.push(tick, &sample(tick as f64, 0.0)).expect("the layout holds");
505 }
506 let measured = sampler.finish();
507 assert_eq!(measured.series.ticks(), [0, 4, 8, 10]);
508 assert_eq!(thinned.series_row_count(), 4);
509 assert_eq!(measured.series.row(3), [10.0, 0.0]);
510 let rows: Vec<(u64, Vec<f64>)> = measured.series.rows().map(|(tick, row)| (tick, row.to_vec())).collect();
511 assert_eq!(rows[1], (4, vec![4.0, 0.0]));
512
513 let silent = Arc::new(plan(0, 10, 2, 0));
514 let mut sampler = Sampler::new(Arc::clone(&silent));
515 for tick in silent.sample_ticks() {
516 sampler.push(tick, &sample(1.0, 0.0)).expect("the layout holds");
517 }
518 let measured = sampler.finish();
519 assert!(measured.series.is_empty(), "a series_every of 0 keeps no rows");
520 assert_eq!(silent.series_row_count(), 0);
521 assert_eq!(measured.reducers.len(), 8, "reducers still run");
522 }
523
524 #[test]
525 fn non_finite_samples_are_skipped_and_flagged() {
526 let plan = Arc::new(plan(0, 4, 2, 2));
527 let mut sampler = Sampler::new(plan);
528 sampler.push(0, &sample(3.0, 1.0)).expect("the layout holds");
529 sampler.push(2, &sample(f64::NAN, 1.0)).expect("the layout holds");
530 sampler.push(4, &sample(4.0, f64::INFINITY)).expect("the layout holds");
531 let measured = sampler.finish();
532 assert_eq!(
533 measured.non_finite,
534 Some(NonFiniteSample {
535 column: "Infected".to_owned(),
536 tick: 2
537 }),
538 "the first one is reported"
539 );
540 assert_eq!(
541 measured.reducers[..4],
542 [Some(4.0), Some(3.0), Some(4.0), Some(3.5)],
543 "Infected skips its NaN"
544 );
545 assert_eq!(
546 measured.reducers[4..],
547 [Some(1.0), Some(1.0), Some(1.0), Some(1.0)],
548 "Recovered skips its infinity"
549 );
550 assert!(
551 measured.series.row(1)[0].is_nan(),
552 "the series keeps what the model gave"
553 );
554 assert_eq!(
555 measured.non_finite.map(|sample| sample.to_string()).as_deref(),
556 Some("Infected is not finite at tick 2")
557 );
558 }
559
560 #[test]
561 fn a_sample_of_another_shape_is_refused() {
562 let mut sampler = Sampler::new(Arc::new(plan(0, 4, 2, 2)));
563 sampler.push(0, &sample(1.0, 1.0)).expect("the layout holds");
564 assert!(sampler.push(2, &[stat("Infected", 1.0, COLOR)]).is_err());
565 let measured = sampler.finish();
566 assert_eq!(measured.series.ticks(), [0], "the refused sample adds no row");
567
568 let mut sampler = Sampler::new(Arc::new(plan(0, 8, 2, 4)));
569 sampler.push(0, &sample(1.0, 1.0)).expect("the layout holds");
570 sampler.push(2, &sample(2.0, 1.0)).expect("the layout holds");
571 assert!(sampler.push(4, &[stat("Infected", 1.0, COLOR)]).is_err());
572 let measured = sampler.finish();
573 assert_eq!(
574 measured.series.ticks(),
575 [0, 2],
576 "the sample the series had not kept yet outlives the refused one"
577 );
578 assert_eq!(measured.series.row(1), [2.0, 1.0]);
579 }
580
581 #[test]
582 fn a_cadence_that_cannot_sample_is_refused() {
583 let run = RunSettings::default();
584 let measure = |stats_every, series_every| MeasureSettings {
585 stats_every,
586 series_every,
587 ..MeasureSettings::default()
588 };
589 assert_eq!(
590 MeasurePlan::check(&run, &measure(0, 0)),
591 Err(MeasureError::ZeroStatsEvery)
592 );
593 assert_eq!(
594 MeasurePlan::check(&run, &measure(4, 6)),
595 Err(MeasureError::SeriesOffCadence {
596 series_every: 6,
597 stats_every: 4
598 })
599 );
600 assert_eq!(MeasurePlan::check(&run, &measure(4, 0)), Ok(()));
601 assert_eq!(MeasurePlan::check(&run, &measure(4, 8)), Ok(()));
602 let endless = RunSettings {
603 steps: u64::MAX,
604 warmup: 1,
605 ..RunSettings::default()
606 };
607 assert!(MeasurePlan::check(&endless, &measure(1, 1)).is_err());
608 }
609
610 #[test]
613 fn settings_a_spec_file_cannot_record_are_refused() {
614 let reversed = MeasureSettings {
615 reducers: vec![ReducerSpec {
616 column: "Infected".to_owned(),
617 kind: ReducerKind::WindowMean { start: 600, end: 200 },
618 }],
619 ..MeasureSettings::default()
620 };
621 assert_eq!(
622 MeasurePlan::check(&RunSettings::default(), &reversed),
623 Err(MeasureError::Reducer(ReducerError::BadWindow {
624 raw: "mean@600..200".to_owned()
625 }))
626 );
627 let single_tick = MeasureSettings {
628 reducers: vec![ReducerSpec {
629 column: "Infected".to_owned(),
630 kind: ReducerKind::WindowMean { start: 200, end: 200 },
631 }],
632 ..MeasureSettings::default()
633 };
634 assert_eq!(MeasurePlan::check(&RunSettings::default(), &single_tick), Ok(()));
635
636 let endless = RunSettings {
637 timeout: Some(Duration::MAX),
638 ..RunSettings::default()
639 };
640 assert_eq!(
641 MeasurePlan::check(&endless, &MeasureSettings::default()),
642 Err(MeasureError::TimeoutTooLong { timeout: Duration::MAX })
643 );
644 let long = RunSettings {
645 timeout: Some(Duration::from_secs(u64::MAX / 2)),
646 ..RunSettings::default()
647 };
648 assert_eq!(MeasurePlan::check(&long, &MeasureSettings::default()), Ok(()));
649 }
650
651 #[test]
653 fn the_longest_accepted_run_counts_its_samples_without_overflow() {
654 let longest = plan(0, u64::MAX, 1, 1);
655 assert_eq!(longest.sample_count(), u64::MAX, "the count saturates");
656 assert_eq!(longest.series_row_count(), u64::MAX);
657 assert_eq!(plan(0, u64::MAX, 2, 2).series_row_count(), u64::MAX / 2 + 2);
658 }
659
660 #[test]
662 fn a_comparison_against_a_threshold_that_is_not_finite_is_refused() {
663 let infinite = Comparison {
664 comparator: Comparator::LessOrEqual,
665 threshold: f64::INFINITY,
666 };
667 let run = RunSettings {
668 stop: Some(StopSpec {
669 column: "Infected".to_owned(),
670 comparison: infinite,
671 min_tick: 0,
672 }),
673 ..RunSettings::default()
674 };
675 assert_eq!(
676 MeasurePlan::check(&run, &MeasureSettings::default()),
677 Err(MeasureError::Stop(StopError::NonFiniteThreshold {
678 raw: "Infected <= inf".to_owned()
679 }))
680 );
681
682 let measure = MeasureSettings {
683 reducers: vec![ReducerSpec {
684 column: "Infected".to_owned(),
685 kind: ReducerKind::FirstCrossing(Comparison {
686 threshold: f64::NAN,
687 ..infinite
688 }),
689 }],
690 ..MeasureSettings::default()
691 };
692 assert_eq!(
693 MeasurePlan::check(&RunSettings::default(), &measure),
694 Err(MeasureError::Reducer(ReducerError::Comparison {
695 raw: "first<=NaN".to_owned(),
696 source: ComparisonError::BadThreshold { raw: "NaN".to_owned() },
697 }))
698 );
699 assert!(
700 MeasurePlan::new(&run, &MeasureSettings::default(), StatColumns::plan(&sample(0.0, 0.0))).is_err(),
701 "a plan built without a check refuses it too"
702 );
703 }
704
705 #[test]
706 fn the_sampler_reports_the_first_sample_where_the_stop_holds() {
707 let run = RunSettings {
708 steps: 20,
709 stop: Some(StopSpec::parse("Infected <= 0", 6).expect("a well-formed condition")),
710 ..RunSettings::default()
711 };
712 let measure = MeasureSettings {
713 stats_every: 2,
714 series_every: 4,
715 ..MeasureSettings::default()
716 };
717 let plan = MeasurePlan::new(&run, &measure, StatColumns::plan(&sample(0.0, 0.0))).expect("the column exists");
718 let mut sampler = Sampler::new(Arc::new(plan));
719 let mut stops = Vec::new();
720 for (tick, infected) in [(0, 0.0), (2, 3.0), (4, 1.0), (6, f64::NAN), (8, 2.0), (10, 0.0)] {
721 stops.push(sampler.push(tick, &sample(infected, 0.0)).expect("the layout holds"));
722 }
723 assert_eq!(
724 stops,
725 [false, false, false, false, false, true],
726 "tick 0 is before the minimum tick, and NaN never stops"
727 );
728 let measured = sampler.finish();
729 assert_eq!(
730 measured.series.ticks(),
731 [0, 4, 8, 10],
732 "the series ends on the stopping sample"
733 );
734 assert_eq!(measured.reducers[0], Some(0.0), "Infected:final is the stopping value");
735
736 let missing = RunSettings {
737 stop: Some(StopSpec::parse("Susceptible <= 0", 0).expect("a well-formed condition")),
738 ..RunSettings::default()
739 };
740 assert!(matches!(
741 MeasurePlan::new(&missing, &measure, StatColumns::plan(&sample(0.0, 0.0))),
742 Err(MeasureError::Stop(StopError::UnknownColumn { .. }))
743 ));
744 }
745}