1use crate::column::ColumnData;
9use crate::duration::parse_duration;
10use crate::error::{PondError, Result};
11use crate::reducer::Reducer;
12use crate::rolling::{Alignment, Frame, Rolling, Span, check_magnitude, roll_column};
13use crate::series::{Column, Key, TimeSeries};
14
15#[derive(Debug, Clone, PartialEq)]
28pub struct Smooth {
29 method: Method,
30 output: Option<String>,
31 warmup: Option<usize>,
32 min_samples: Option<usize>,
33 alignment: Option<Alignment>,
34 skip_missing: bool,
35}
36
37#[derive(Debug, Clone, PartialEq)]
38enum Method {
39 Ema(Rate),
40 MovingAverage(String),
41 Loess(f64),
42}
43
44#[derive(Debug, Clone, Copy, PartialEq)]
45enum Rate {
46 Alpha(f64),
47 Span(f64),
48}
49
50impl Smooth {
51 fn new(method: Method) -> Self {
52 Smooth {
53 method,
54 output: None,
55 warmup: None,
56 min_samples: None,
57 alignment: None,
58 skip_missing: false,
59 }
60 }
61
62 pub fn ema(alpha: f64) -> Self {
67 Self::new(Method::Ema(Rate::Alpha(alpha)))
68 }
69
70 pub fn ema_span(span: f64) -> Self {
74 Self::new(Method::Ema(Rate::Span(span)))
75 }
76
77 pub fn moving_average(window: &str) -> Self {
84 Self::new(Method::MovingAverage(window.to_string()))
85 }
86
87 pub fn loess(span: f64) -> Self {
94 Self::new(Method::Loess(span))
95 }
96
97 pub fn output(mut self, name: &str) -> Self {
101 self.output = Some(name.to_string());
102 self
103 }
104
105 pub fn warmup(mut self, rows: usize) -> Self {
109 self.warmup = Some(rows);
110 self
111 }
112
113 pub fn min_samples(mut self, values: usize) -> Self {
115 self.min_samples = Some(values);
116 self
117 }
118
119 pub fn alignment(mut self, alignment: Alignment) -> Self {
122 self.alignment = Some(alignment);
123 self
124 }
125
126 pub fn skip_missing(mut self) -> Self {
130 self.skip_missing = true;
131 self
132 }
133
134 fn check_options(&self) -> Result<()> {
136 let (method, ema, moving, loess) = match self.method {
137 Method::Ema(_) => ("ema", true, false, false),
138 Method::MovingAverage(_) => ("moving_average", false, true, false),
139 Method::Loess(_) => ("loess", false, false, true),
140 };
141 let misplaced = [
142 ("warmup", self.warmup.is_some() && !ema),
143 ("min_samples", self.min_samples.is_some() && !ema),
144 ("alignment", self.alignment.is_some() && !moving),
145 ("skip_missing", self.skip_missing && !(moving || loess)),
146 ];
147 match misplaced.iter().find(|(_, bad)| *bad) {
148 None => Ok(()),
149 Some((option, _)) => Err(invalid(format!("{method} does not take {option}"))),
150 }
151 }
152}
153
154const SMOOTH: &str = "smooth";
155
156fn invalid(message: String) -> PondError {
157 PondError::InvalidData {
158 call: SMOOTH,
159 message,
160 }
161}
162
163fn anchors(key: &Key) -> Vec<f64> {
167 let begin = key.begin();
168 match key.end() {
169 None => begin.iter().map(|&b| b as f64).collect(),
170 Some(end) => begin
171 .iter()
172 .zip(end)
173 .map(|(&b, &e)| b as f64 + (e - b) as f64 / 2.0)
174 .collect(),
175 }
176}
177
178fn anchor_order(anchors: &[f64]) -> Option<Vec<usize>> {
183 if anchors.windows(2).all(|w| w[0] <= w[1]) {
184 return None;
185 }
186 let mut order: Vec<usize> = (0..anchors.len()).collect();
187 order.sort_by(|&a, &b| anchors[a].total_cmp(&anchors[b]));
188 Some(order)
189}
190
191fn loess_at(x: f64, anchors: &[f64], values: &[f64], span: f64) -> Option<f64> {
194 let n = anchors.len();
195 match n {
196 0 => return None,
197 1 => return Some(values[0]),
198 _ => {}
199 }
200 let neighbours = ((span * n as f64).ceil() as usize).clamp(2, n);
201 let (mut low, mut high) = (0, n - neighbours);
204 while low < high {
205 let mid = (low + high) / 2;
206 if x - anchors[mid] > anchors[mid + neighbours] - x {
207 low = mid + 1;
208 } else {
209 high = mid;
210 }
211 }
212 let (start, end) = (low, low + neighbours);
213 let bandwidth = (x - anchors[start]).abs().max((anchors[end - 1] - x).abs());
214 if bandwidth == 0.0 {
215 let lo = anchors.partition_point(|&a| a < x);
217 let hi = anchors.partition_point(|&a| a <= x);
218 let sum = values[lo..hi].iter().fold(0.0, |acc, v| acc + v);
219 return Some(sum / (hi - lo) as f64);
220 }
221 let (mut weighted, mut sw, mut swx, mut swy, mut swxx, mut swxy) =
222 (0usize, 0.0, 0.0, 0.0, 0.0, 0.0);
223 for i in start..end {
224 let dx = anchors[i] - x;
226 let y = values[i];
227 let ratio = dx.abs() / bandwidth;
228 let cube = ratio * ratio * ratio;
230 let tail = 1.0 - cube;
231 let w = if ratio >= 1.0 {
232 0.0
233 } else {
234 tail * tail * tail
235 };
236 if w == 0.0 {
237 continue;
238 }
239 weighted += 1;
240 sw += w;
241 swx += w * dx;
242 swy += w * y;
243 swxx += w * dx * dx;
244 swxy += w * dx * y;
245 }
246 if weighted == 0 || sw == 0.0 {
247 return None;
248 }
249 let denominator = sw * swxx - swx * swx;
250 if denominator.abs() < f64::EPSILON {
251 return Some(swy / sw);
252 }
253 Some((swy * swxx - swx * swxy) / denominator)
254}
255
256fn skip_gaps(mut values: Vec<f64>, source: &[f64]) -> Vec<f64> {
258 for (v, s) in values.iter_mut().zip(source) {
259 if !s.is_finite() {
260 *v = f64::NAN;
261 }
262 }
263 values
264}
265
266fn finite_or_gap(v: f64) -> f64 {
268 if v.is_finite() { v } else { f64::NAN }
269}
270
271#[derive(Debug, Clone, Copy, PartialEq, Eq)]
273pub enum Sample {
274 Stride(usize),
277 Reservoir {
283 size: usize,
285 seed: u64,
287 },
288}
289
290struct SplitMix64(u64);
292
293impl SplitMix64 {
294 fn next(&mut self) -> u64 {
295 self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
296 let mut z = self.0;
297 z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
298 z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
299 z ^ (z >> 31)
300 }
301
302 fn below(&mut self, bound: u64) -> u64 {
306 ((self.next() as u128 * bound as u128) >> 64) as u64
307 }
308}
309
310impl TimeSeries {
311 pub(crate) fn number_source(
313 &self,
314 call: &'static str,
315 column: &str,
316 ) -> Result<(usize, &[f64])> {
317 let c = self.column_index(call, column)?;
318 let source = &self.columns()[c];
319 match source.data() {
320 ColumnData::Number(v) => Ok((c, v.values())),
321 other => Err(PondError::WrongKind {
322 call,
323 column: source.name().to_string(),
324 kind: other.kind().as_str(),
325 reason: "needs a number column".to_string(),
326 }),
327 }
328 }
329
330 fn check_new_name(&self, call: &'static str, name: &str) -> Result<()> {
332 if name == self.key_name() || self.columns().iter().any(|c| c.name() == name) {
333 return Err(PondError::DuplicateOutput {
334 call,
335 column: name.to_string(),
336 });
337 }
338 Ok(())
339 }
340
341 fn put_numbers(&self, c: usize, output: Option<&str>, values: Vec<f64>) -> TimeSeries {
344 let mut columns = self.columns().to_vec();
345 match output {
346 None => {
347 columns[c] = Column::from_data(columns[c].name().to_string(), values, false);
348 }
349 Some(name) => columns.push(Column::from_data(name.to_string(), values, false)),
350 }
351 self.with_value_columns(columns)
352 }
353
354 pub fn smooth(&self, column: &str, smooth: &Smooth) -> Result<TimeSeries> {
373 smooth.check_options()?;
374 let (c, values) = self.number_source(SMOOTH, column)?;
375 let output = smooth.output.as_deref().filter(|&o| o != column);
376 if let Some(name) = output {
377 self.check_new_name(SMOOTH, name)?;
378 }
379 match &smooth.method {
380 Method::Ema(rate) => {
381 let alpha = match *rate {
382 Rate::Alpha(a) if a.is_finite() && a > 0.0 && a <= 1.0 => a,
383 Rate::Alpha(a) => {
384 return Err(invalid(format!("ema alpha must be in (0, 1], got {a}")));
385 }
386 Rate::Span(s) if s.is_finite() && s >= 1.0 => 2.0 / (s + 1.0),
387 Rate::Span(s) => {
388 return Err(invalid(format!(
389 "ema span must be a finite number >= 1, got {s}"
390 )));
391 }
392 };
393 let min_samples = smooth.min_samples.unwrap_or(0);
394 let mut previous: Option<f64> = None;
395 let mut seen = 0;
396 let out: Vec<f64> = values
397 .iter()
398 .map(|&raw| {
399 if !raw.is_finite() {
400 return f64::NAN;
401 }
402 let v = previous.map_or(raw, |p| alpha * raw + (1.0 - alpha) * p);
403 previous = Some(v);
404 seen += 1;
405 if seen < min_samples {
406 f64::NAN
407 } else {
408 finite_or_gap(v)
409 }
410 })
411 .collect();
412 let result = self.put_numbers(c, output, out);
413 let warmup = smooth.warmup.unwrap_or(0).min(result.len());
414 Ok(if warmup == 0 {
415 result
416 } else {
417 result.row_range(warmup..result.len())
418 })
419 }
420 Method::MovingAverage(window) => {
421 let window_ms = parse_duration(window)?;
422 check_magnitude(SMOOTH, &self.columns()[c], Reducer::Avg)?;
423 let mut times = anchors(self.key());
424 let order = anchor_order(×);
425 let sorted_values: Vec<f64>;
426 let in_order = match &order {
427 None => values,
428 Some(order) => {
429 times = order.iter().map(|&i| times[i]).collect();
430 sorted_values = order.iter().map(|&i| values[i]).collect();
431 &sorted_values
432 }
433 };
434 let frame = Frame {
435 times: ×,
436 span: Span::Time(window_ms),
437 alignment: smooth.alignment.unwrap_or_default(),
438 min_samples: 0,
439 };
440 let mut out = roll_column(&frame, in_order, Reducer::Avg);
441 if let Some(order) = order {
442 let mut by_row = vec![f64::NAN; out.len()];
443 for (k, &row) in order.iter().enumerate() {
444 by_row[row] = out[k];
445 }
446 out = by_row;
447 }
448 if smooth.skip_missing {
449 out = skip_gaps(out, values);
450 }
451 Ok(self.put_numbers(c, output, out))
452 }
453 Method::Loess(span) => {
454 let span = *span;
455 if !(span.is_finite() && span > 0.0 && span <= 1.0) {
456 return Err(invalid(format!("loess span must be in (0, 1], got {span}")));
457 }
458 let times = anchors(self.key());
459 let rows: Vec<usize> = match anchor_order(×) {
460 None => (0..times.len()).collect(),
461 Some(order) => order,
462 };
463 let (fit_at, fit_values): (Vec<f64>, Vec<f64>) = rows
464 .iter()
465 .filter(|&&i| values[i].is_finite())
466 .map(|&i| (times[i], values[i]))
467 .unzip();
468 let rows: Vec<usize> = (0..times.len()).collect();
469 let out: Vec<f64> = crate::par::map(&rows, |&i| {
471 if smooth.skip_missing && !values[i].is_finite() {
472 return f64::NAN;
473 }
474 loess_at(times[i], &fit_at, &fit_values, span).map_or(f64::NAN, finite_or_gap)
475 });
476 Ok(self.put_numbers(c, output, out))
477 }
478 }
479 }
480
481 pub fn baseline(
508 &self,
509 column: &str,
510 window: &str,
511 sigma: f64,
512 options: Rolling,
513 ) -> Result<TimeSeries> {
514 self.baseline_named(
515 column,
516 window,
517 sigma,
518 options,
519 ["avg", "sd", "upper", "lower"],
520 )
521 }
522
523 pub fn baseline_named(
528 &self,
529 column: &str,
530 window: &str,
531 sigma: f64,
532 options: Rolling,
533 names: [&str; 4],
534 ) -> Result<TimeSeries> {
535 const CALL: &str = "baseline";
536 for (i, name) in names.iter().enumerate() {
537 self.check_new_name(CALL, name)?;
538 if names[..i].contains(name) {
539 return Err(PondError::DuplicateOutput {
540 call: CALL,
541 column: name.to_string(),
542 });
543 }
544 }
545 let (avg, sd) = self.band(CALL, column, window, sigma, options)?;
546 let (upper, lower): (Vec<f64>, Vec<f64>) = avg
547 .iter()
548 .zip(&sd)
549 .map(|(&a, &s)| {
550 if s > 0.0 {
552 (a + sigma * s, a - sigma * s)
553 } else {
554 (f64::NAN, f64::NAN)
555 }
556 })
557 .unzip();
558 let mut columns = self.columns().to_vec();
559 for (name, values) in names.iter().zip([avg, sd, upper, lower]) {
560 columns.push(Column::from_data(name.to_string(), values, false));
561 }
562 Ok(self.with_value_columns(columns))
563 }
564
565 pub fn outliers(
581 &self,
582 column: &str,
583 window: &str,
584 sigma: f64,
585 options: Rolling,
586 ) -> Result<TimeSeries> {
587 const CALL: &str = "outliers";
588 let (avg, sd) = self.band(CALL, column, window, sigma, options)?;
589 let (_, values) = self.number_source(CALL, column)?;
590 let rows: Vec<usize> = (0..self.len())
591 .filter(|&i| {
592 let (v, a, s) = (values[i], avg[i], sd[i]);
593 s > 0.0 && (v - a).abs() > sigma * s
595 })
596 .collect();
597 Ok(self.take_rows(&rows))
598 }
599
600 fn band(
602 &self,
603 call: &'static str,
604 column: &str,
605 window: &str,
606 sigma: f64,
607 options: Rolling,
608 ) -> Result<(Vec<f64>, Vec<f64>)> {
609 if !(sigma.is_finite() && sigma > 0.0) {
610 return Err(PondError::InvalidData {
611 call,
612 message: format!("sigma must be a positive finite number, got {sigma}"),
613 });
614 }
615 let (c, values) = self.number_source(call, column)?;
616 let source = &self.columns()[c];
617 check_magnitude(call, source, Reducer::Avg)?;
618 check_magnitude(call, source, Reducer::Stdev)?;
619 let window_ms = parse_duration(window)?;
620 let frame = Frame {
622 times: self.key().begin(),
623 span: Span::Time(window_ms),
624 alignment: options.alignment,
625 min_samples: options.min_samples,
626 };
627 let reducers = [Reducer::Avg, Reducer::Stdev];
628 let mut both = crate::par::map(&reducers, |&r| roll_column(&frame, values, r));
629 let sd = both.pop().expect("two reducers");
630 let avg = both.pop().expect("two reducers");
631 Ok((avg, sd))
632 }
633
634 pub fn scan<A>(
661 &self,
662 column: &str,
663 init: A,
664 step: impl FnMut(A, f64, usize) -> (A, f64),
665 ) -> Result<TimeSeries> {
666 self.scan_into(column, None, init, step)
667 }
668
669 pub fn scan_to<A>(
673 &self,
674 column: &str,
675 output: &str,
676 init: A,
677 step: impl FnMut(A, f64, usize) -> (A, f64),
678 ) -> Result<TimeSeries> {
679 self.scan_into(column, Some(output), init, step)
680 }
681
682 fn scan_into<A>(
683 &self,
684 column: &str,
685 output: Option<&str>,
686 init: A,
687 mut step: impl FnMut(A, f64, usize) -> (A, f64),
688 ) -> Result<TimeSeries> {
689 const CALL: &str = "scan";
690 let (c, values) = self.number_source(CALL, column)?;
691 if let Some(name) = output {
692 self.check_new_name(CALL, name)?;
693 }
694 let mut acc = Some(init);
695 let mut last = f64::NAN;
696 let out: Vec<f64> = values
697 .iter()
698 .enumerate()
699 .map(|(i, &v)| {
700 if v.is_finite() {
701 let (next, o) = step(acc.take().expect("accumulator is put back"), v, i);
702 acc = Some(next);
703 last = finite_or_gap(o);
704 }
705 last
706 })
707 .collect();
708 Ok(self.put_numbers(c, output, out))
709 }
710
711 pub fn sample(&self, how: Sample) -> Result<TimeSeries> {
722 const CALL: &str = "sample";
723 let n = self.len();
724 let rows: Vec<usize> = match how {
725 Sample::Stride(0) | Sample::Reservoir { size: 0, .. } => {
726 return Err(PondError::InvalidData {
727 call: CALL,
728 message: "needs a stride or size of at least 1".into(),
729 });
730 }
731 Sample::Stride(stride) => (stride - 1..n).step_by(stride).collect(),
732 Sample::Reservoir { size, .. } if n <= size => return Ok(self.clone()),
733 Sample::Reservoir { size, seed } => {
734 let mut rng = SplitMix64(seed);
735 let mut picked: Vec<usize> = (0..size).collect();
736 for i in size..n {
737 let j = rng.below(i as u64 + 1) as usize;
738 if j < size {
739 picked[j] = i;
740 }
741 }
742 picked.sort_unstable();
743 picked
744 }
745 };
746 Ok(self.take_rows(&rows))
747 }
748}
749
750#[cfg(test)]
751mod tests {
752 use super::*;
753
754 fn series(times: Vec<i64>, values: Vec<Option<f64>>) -> TimeSeries {
755 TimeSeries::new("s", times, vec![Column::optional("v", values)]).unwrap()
756 }
757
758 fn numbers(s: &TimeSeries, name: &str) -> Vec<Option<f64>> {
759 s.column(name).unwrap().numbers().unwrap().iter().collect()
760 }
761
762 #[test]
763 fn options_must_fit_the_method() {
764 let s = series(vec![0, 1], vec![Some(1.0), Some(2.0)]);
765 for bad in [
766 Smooth::loess(0.5).warmup(1),
767 Smooth::moving_average("1s").min_samples(1),
768 Smooth::ema(0.5).alignment(Alignment::Leading),
769 Smooth::ema(0.5).skip_missing(),
770 ] {
771 let e = s.smooth("v", &bad).unwrap_err();
772 assert!(e.to_string().contains("does not take"), "{e}");
773 }
774 }
775
776 #[test]
777 fn loess_fits_a_line_exactly() {
778 let t: Vec<i64> = (0..20).map(|i| i * 1000).collect();
780 let v: Vec<Option<f64>> = (0..20).map(|i| Some(3.0 * i as f64 - 7.0)).collect();
781 let s = series(t, v.clone());
782 let fit = s.smooth("v", &Smooth::loess(0.4)).unwrap();
783 for (got, want) in numbers(&fit, "v").iter().zip(&v) {
784 assert!((got.unwrap() - want.unwrap()).abs() < 1e-9);
785 }
786 }
787
788 #[test]
789 fn scan_carries_over_gaps_and_drops_non_finite() {
790 let s = series(
791 vec![0, 1, 2, 3, 4],
792 vec![None, Some(1.0), None, Some(0.0), Some(2.0)],
793 );
794 let r = s
796 .scan_to("v", "inv", 0, |n, v, _| (n + 1, 1.0 / v))
797 .unwrap();
798 assert_eq!(
799 numbers(&r, "inv"),
800 [None, Some(1.0), Some(1.0), None, Some(0.5)]
801 );
802 assert!(!r.column("inv").unwrap().is_required());
803 }
804
805 #[test]
806 fn reservoir_picks_rows_evenly() {
807 let s = series((0..10).collect(), (0..10).map(|i| Some(i as f64)).collect());
809 let mut counts = [0usize; 10];
810 let runs = 20_000;
811 for seed in 0..runs {
812 let picked = s.sample(Sample::Reservoir { size: 3, seed }).unwrap();
813 assert_eq!(picked.len(), 3);
814 let begin = picked.key().begin();
815 assert!(begin.windows(2).all(|w| w[0] < w[1]), "in order");
816 for &t in begin {
817 counts[t as usize] += 1;
818 }
819 }
820 let expected = runs as f64 * 0.3;
821 for c in counts {
822 assert!((c as f64 - expected).abs() < 5.0 * 65.0, "{counts:?}");
824 }
825 let a = s.sample(Sample::Reservoir { size: 4, seed: 9 }).unwrap();
827 let b = s.sample(Sample::Reservoir { size: 4, seed: 9 }).unwrap();
828 assert_eq!(a.key().begin(), b.key().begin());
829 }
830}