chore: add benchmark for Prophet

crates/augurs-clustering/benches/dbscan.rs

@@ -1,3 +1,5 @@
+#![allow(missing_docs)]
+
 use criterion::{criterion_group, criterion_main, Criterion};
 
 use augurs_clustering::DbscanClusterer;

crates/augurs-dtw/benches/dtw.rs

@@ -1,3 +1,5 @@
+#![allow(missing_docs)]
+
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use itertools::Itertools;
 

crates/augurs-ets/benches/air_passengers.rs

@@ -1,3 +1,5 @@
+#![allow(missing_docs)]
+
 use criterion::{criterion_group, criterion_main, Criterion};
 use pprof::criterion::{Output, PProfProfiler};
 

crates/augurs-ets/benches/air_passengers_iai.rs

@@ -1,3 +1,5 @@
+#![allow(missing_docs)]
+
 use iai::{black_box, main};
 
 use augurs_core::Fit;

crates/augurs-mstl/benches/vic_elec.rs

@@ -1,3 +1,5 @@
+#![allow(missing_docs)]
+
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use pprof::criterion::{Output, PProfProfiler};
 

crates/augurs-mstl/benches/vic_elec_iai.rs

@@ -1,3 +1,5 @@
+#![allow(missing_docs)]
+
 use iai::{black_box, main};
 
 use augurs_core::Fit;

crates/augurs-prophet/Cargo.toml

@@ -28,7 +28,9 @@ zip = { version = "2.2.0", optional = true }
 augurs.workspace = true
 augurs-testing.workspace = true
 chrono.workspace = true
+criterion.workspace = true
 pretty_assertions.workspace = true
+tracing-subscriber = { workspace = true, default-features = true }
 
 [build-dependencies]
 tempfile = { version = "3.13.0", optional = true }
@@ -45,6 +47,14 @@ download = ["dep:ureq", "dep:zip"]
 internal-ignore-cmdstan-failure = []
 serde = ["dep:serde"]
 
+[lib]
+bench = false
+
+[[bench]]
+name = "real-life"
+harness = false
+required-features = ["cmdstan", "compile-cmdstan"]
+
 [[bin]]
 name = "download-stan-model"
 path = "src/bin/main.rs"

crates/augurs-prophet/src/prophet.rs

@@ -18,7 +18,7 @@ use crate::{
 };
 
 /// The Prophet time series forecasting model.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Prophet<O> {
     /// Options to be used for fitting.
     opts: ProphetOptions,

crates/augurs-prophet/src/prophet/predict.rs

@@ -167,34 +167,27 @@ impl<O> Prophet<O> {
         let cp_zipped = deltas.iter().zip(changepoints_t);
         let deltas_t = cp_zipped
             .cartesian_product(t)
-            .map(|((delta, cp_t), t)| if cp_t <= t { *delta } else { 0.0 })
-            .collect_vec();
-        // `k_t` is a contiguous array with the rate to apply at each time point.
-        let k_t = deltas_t
+            .map(|((delta, cp_t), t)| if cp_t <= t { *delta } else { 0.0 });
+
+        // Repeat each changepoint effect `n` times so we can zip it up.
+        let changepoints_repeated = changepoints_t
             .iter()
-            .enumerate()
-            .fold(vec![k; t.len()], |mut acc, (i, delta)| {
+            .flat_map(|x| std::iter::repeat(*x).take(t.len()));
+        let indexes = (0..t.len()).cycle();
+        // `k_m_t` is a contiguous array where each element contains the rate and offset to
+        // apply at each time point.
+        let k_m_t = izip!(deltas_t, changepoints_repeated, indexes).fold(
+            vec![(k, m); t.len()],
+            |mut acc, (delta, cp_t, idx)| {
                 // Add the changepoint rate to the initial rate.
-                acc[i % t.len()] += *delta;
-                acc
-            });
-        // `m_t` is a contiguous array with the offset to apply at each time point.
-        let m_t = deltas_t
-            .iter()
-            .zip(
-                // Repeat each changepoint effect `n` times so we can zip it up.
-                changepoints_t
-                    .iter()
-                    .flat_map(|x| std::iter::repeat(*x).take(t.len())),
-            )
-            .enumerate()
-            .fold(vec![m; t.len()], |mut acc, (i, (delta, cp_t))| {
+                acc[idx].0 += delta;
                 // Add the changepoint offset to the initial offset where applicable.
-                acc[i % t.len()] += -cp_t * delta;
+                acc[idx].1 += -cp_t * delta;
                 acc
-            });
+            },
+        );
 
-        izip!(t, k_t, m_t).map(|(t, k, m)| t * k + m)
+        izip!(t, k_m_t).map(|(t, (k, m))| t * k + m)
     }
 
     fn piecewise_logistic<'a>(
@@ -416,9 +409,15 @@ impl<O> Prophet<O> {
                 .take(n_timestamps)
                 .collect_vec(),
         };
+        // Use temporary buffers to avoid allocating a new Vec for each
+        // call to `sample_model`.
+        let (mut yhat, mut trend) = (
+            Vec::with_capacity(n_timestamps),
+            Vec::with_capacity(n_timestamps),
+        );
         for i in 0..n_iterations {
             for _ in 0..samples_per_iter {
-                let (yhat, trend) = self.sample_model(
+                self.sample_model(
                     df,
                     features,
                     params,
@@ -427,11 +426,13 @@ impl<O> Prophet<O> {
                     &component_columns.multiplicative,
                     y_scale,
                     i,
+                    &mut yhat,
+                    &mut trend,
                 )?;
                 // We have to transpose things, unfortunately.
-                for ((i, yhat), trend) in yhat.into_iter().enumerate().zip(trend) {
-                    sim_values.yhat[i].push(yhat);
-                    sim_values.trend[i].push(trend);
+                for ((i, yhat), trend) in yhat.iter().enumerate().zip(&trend) {
+                    sim_values.yhat[i].push(*yhat);
+                    sim_values.trend[i].push(*trend);
                 }
             }
         }
@@ -452,9 +453,14 @@ impl<O> Prophet<O> {
         multiplicative: &[i32],
         y_scale: f64,
         iteration: usize,
-    ) -> Result<(Vec<f64>, Vec<f64>), Error> {
+        yhat_tmp: &mut Vec<f64>,
+        trend_tmp: &mut Vec<f64>,
+    ) -> Result<(), Error> {
+        yhat_tmp.clear();
+        trend_tmp.clear();
         let n = df.ds.len();
-        let trend = self.sample_predictive_trend(df, params, changepoints_t, y_scale, iteration)?;
+        *trend_tmp =
+            self.sample_predictive_trend(df, params, changepoints_t, y_scale, iteration)?;
         let beta = &params.beta;
         let mut xb_a = vec![0.0; n];
         for (feature, b, a) in izip!(&features.data, beta, additive) {
@@ -475,11 +481,12 @@ impl<O> Prophet<O> {
         let mut rng = thread_rng();
         let noise = (&mut rng).sample_iter(dist).take(n).map(|x| x * y_scale);
 
-        let yhat = izip!(&trend, &xb_a, &xb_m, noise)
-            .map(|(t, a, m, n)| t * (1.0 + m) + a + n)
-            .collect();
+        for yhat in izip!(trend_tmp, &xb_a, &xb_m, noise).map(|(t, a, m, n)| *t * (1.0 + m) + a + n)
+        {
+            yhat_tmp.push(yhat);
+        }
 
-        Ok((yhat, trend))
+        Ok(())
     }
 
     fn sample_predictive_trend(

crates/augurs-prophet/src/prophet/prep.rs

@@ -28,7 +28,7 @@ pub(super) struct Scales {
     pub(super) regressors: HashMap<String, RegressorScale>,
 }
 
-#[derive(Debug, Default, PartialEq, Eq)]
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub(super) struct Modes {
     pub(super) additive: HashSet<ComponentName>,
     pub(super) multiplicative: HashSet<ComponentName>,

crates/augurs-seasons/benches/periodogram.rs

@@ -1,3 +1,5 @@
+#![allow(missing_docs)]
+
 use criterion::{criterion_group, criterion_main, Criterion};
 use pprof::criterion::{Output, PProfProfiler};
 

crates/augurs/tests/integration.rs

@@ -1,3 +1,5 @@
+//! Integration tests for the augurs wrapper crate.
+
 #[cfg(feature = "changepoint")]
 #[test]
 fn test_changepoint() {