Avoid repeated sorting and observation copies in TPE model construction

rustuna_core/src/parzen_estimator/model.rs

@@ -1,5 +1,5 @@
 use rand::rngs::StdRng;
-use std::{collections::HashMap, vec};
+use std::{borrow::Cow, collections::HashMap, vec};
 
 use super::probability_distributions::CategoricalDistributions;
 use super::probability_distributions::{
@@ -13,6 +13,42 @@ pub struct ParzenEstimator {
     mixuture_distribution: MixtureOfProductDistribution,
 }
 
+enum ObservationValues<'a, I> {
+    Borrowed(&'a [f64]),
+    Selected(I),
+}
+
+impl<'a, I: Iterator<Item = f64>> ObservationValues<'a, I> {
+    fn into_slice(self) -> Cow<'a, [f64]> {
+        match self {
+            Self::Borrowed(values) => Cow::Borrowed(values),
+            Self::Selected(values) => Cow::Owned(values.collect()),
+        }
+    }
+}
+
+impl<I: Iterator<Item = f64>> Iterator for ObservationValues<'_, I> {
+    type Item = f64;
+
+    fn next(&mut self) -> Option<f64> {
+        match self {
+            Self::Borrowed(values) => {
+                let (value, rest) = values.split_first()?;
+                *values = rest;
+                Some(*value)
+            }
+            Self::Selected(values) => values.next(),
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        match self {
+            Self::Borrowed(values) => (values.len(), Some(values.len())),
+            Self::Selected(values) => values.size_hint(),
+        }
+    }
+}
+
 impl ParzenEstimator {
     pub fn new(
         observations: &HashMap<String, Vec<f64>>,
@@ -70,18 +106,89 @@ impl ParzenEstimator {
         let mut keys: Vec<_> = search_space.keys().collect();
         keys.sort();
 
-        let mut distributions = HashMap::with_capacity(keys.len());
-        for key in keys.iter() {
-            let obs_vec = observations.get(*key).map(Vec::as_slice).unwrap_or(&[]);
-            let dist = match &search_space[*key] {
+        Self::with_ordered_builder(
+            keys.into_iter().map(|key| {
+                (
+                    key.as_str(),
+                    &search_space[key],
+                    ObservationValues::<std::iter::Empty<f64>>::Borrowed(
+                        observations.get(key).map(Vec::as_slice).unwrap_or(&[]),
+                    ),
+                )
+            }),
+            weights,
+            prior_weight,
+            num_builder,
+            cat_builder,
+        )
+    }
+
+    /// Builds a model from columns in strictly increasing parameter-name order.
+    /// Each column must provide one value for each observation weight.
+    pub fn from_ordered_observations<'a, I, V>(
+        parameters: I,
+        weights: &[f64],
+        prior_weight: f64,
+    ) -> Self
+    where
+        I: IntoIterator<Item = (&'a str, &'a Distribution, V)>,
+        V: IntoIterator<Item = f64>,
+    {
+        let model = Self::with_ordered_builder(
+            parameters.into_iter().map(|(name, distribution, values)| {
+                (
+                    name,
+                    distribution,
+                    ObservationValues::Selected(values.into_iter()),
+                )
+            }),
+            weights,
+            prior_weight,
+            &DefaultNumericalDistributionBuilder,
+            &DefaultCategoricalDistributionBuilder,
+        );
+        assert!(
+            model
+                .mixuture_distribution
+                .distributions
+                .iter()
+                .all(|distribution| distribution.n_kernels() == weights.len() + 1),
+            "Observation columns and weights differ in length"
+        );
+        model
+    }
+
+    fn with_ordered_builder<'a, I, V>(
+        parameters: I,
+        weights: &[f64],
+        prior_weight: f64,
+        num_builder: &impl NumericalDistributionBuilder,
+        cat_builder: &impl CategoricalDistributionBuilder,
+    ) -> Self
+    where
+        I: IntoIterator<Item = (&'a str, &'a Distribution, ObservationValues<'a, V>)>,
+        V: Iterator<Item = f64>,
+    {
+        let parameters = parameters.into_iter();
+        let mut distributions: Vec<(String, Distributions)> =
+            Vec::with_capacity(parameters.size_hint().0);
+        for (name, search_space, values) in parameters {
+            if let Some((previous, _)) = distributions.last() {
+                assert!(
+                    previous.as_str() < name,
+                    "Parameter names must be sorted and unique"
+                );
+            }
+            let distribution = match search_space {
                 Distribution::Float { .. } | Distribution::Int { .. } => {
-                    num_builder.calculate_numerical_distribution(obs_vec, &search_space[*key])
+                    let values = values.into_slice();
+                    num_builder.calculate_numerical_distribution(&values, search_space)
                 }
                 Distribution::Categorical { .. } => {
-                    cat_builder.calculate_categorical_distribution(obs_vec, &search_space[*key])
+                    cat_builder.calculate_categorical_distribution(values.into_iter(), search_space)
                 }
             };
-            distributions.insert((*key).clone(), dist);
+            distributions.push((name.to_owned(), distribution));
         }
 
         let weights_sum = {
@@ -143,7 +250,7 @@ pub(crate) trait NumericalDistributionBuilder {
 pub(crate) trait CategoricalDistributionBuilder {
     fn calculate_categorical_distribution(
         &self,
-        observations: &[f64],
+        observations: impl Iterator<Item = f64>,
         search_space: &Distribution,
     ) -> Distributions;
 }
@@ -292,7 +399,7 @@ impl NumericalDistributionBuilder for DefaultNumericalDistributionBuilder {
 impl CategoricalDistributionBuilder for DefaultCategoricalDistributionBuilder {
     fn calculate_categorical_distribution(
         &self,
-        observations: &[f64],
+        observations: impl Iterator<Item = f64>,
         search_space: &Distribution,
     ) -> Distributions {
         let cardinality = match search_space {
@@ -367,10 +474,23 @@ mod tests {
 
         let parzen_estimator =
             ParzenEstimator::new(&observations, &search_space, &[0.2, 0.5, 0.3], 1.0);
+        let streamed = ParzenEstimator::from_ordered_observations(
+            ["a", "b", "c", "d", "e"].map(|name| {
+                (
+                    name,
+                    &search_space[name],
+                    observations[name].iter().copied(),
+                )
+            }),
+            &[0.2, 0.5, 0.3],
+            1.0,
+        );
         let mut rng = StdRng::seed_from_u64(42);
         let samples = parzen_estimator.sample(&mut rng, 10);
         let mut ordered_rng = StdRng::seed_from_u64(42);
         let ordered = parzen_estimator.sample_ordered(&mut ordered_rng, 10);
+        let mut streamed_rng = StdRng::seed_from_u64(42);
+        assert_eq!(ordered, streamed.sample_ordered(&mut streamed_rng, 10));
         assert_eq!(
             parzen_estimator.parameter_names(),
             ["a", "b", "c", "d", "e"]
@@ -384,6 +504,10 @@ mod tests {
                 parzen_estimator.log_pdf(sample).to_bits(),
                 parzen_estimator.log_pdf_ordered(values).to_bits()
             );
+            assert_eq!(
+                parzen_estimator.log_pdf_ordered(values).to_bits(),
+                streamed.log_pdf_ordered(values).to_bits()
+            );
         }
         assert_eq!(
             rand::RngCore::next_u64(&mut rng),
@@ -414,4 +538,41 @@ mod tests {
             assert!(*e == 0.0 || *e == 1.0 || *e == 2.0);
         }
     }
+
+    #[test]
+    fn ordered_observations_keep_the_uniform_categorical_prior() {
+        let distribution = Distribution::new_categorical(3);
+        let model = ParzenEstimator::from_ordered_observations(
+            [("x", &distribution, std::iter::empty())],
+            &[],
+            1.0,
+        );
+        for category in 0..3 {
+            assert_eq!(
+                model.log_pdf_ordered(&[category as f64]),
+                (1.0_f64 / 3.0).ln()
+            );
+        }
+    }
+
+    #[test]
+    #[should_panic(expected = "Observation columns and weights differ in length")]
+    fn ordered_observations_reject_inconsistent_weights() {
+        ParzenEstimator::from_ordered_observations(
+            [("x", &Distribution::new_categorical(3), [1.0, 2.0])],
+            &[1.0],
+            1.0,
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "Parameter names must be sorted and unique")]
+    fn ordered_observations_reject_unsorted_parameters() {
+        let distribution = Distribution::new_categorical(3);
+        ParzenEstimator::from_ordered_observations(
+            [("b", &distribution, [0.0]), ("a", &distribution, [1.0])],
+            &[1.0],
+            1.0,
+        );
+    }
 }

rustuna_core/src/parzen_estimator/probability_distributions.rs

@@ -267,17 +267,17 @@ pub(crate) struct CategoricalDistributions {
 }
 
 impl CategoricalDistributions {
-    pub(crate) fn new(observations: &[f64], cardinality: usize) -> Self {
+    pub(crate) fn new(observations: impl IntoIterator<Item = f64>, cardinality: usize) -> Self {
         assert!(cardinality > 0, "Categorical cardinality must be positive");
-        let observed_categories = observations
-            .iter()
-            .map(|&value| {
-                let category = value as usize;
-                assert!(category < cardinality, "Observed category is out of range");
-                category
-            })
-            .collect();
-        let prior_mass = 1.0 / (observations.len() + 1) as f64;
+        let observations = observations.into_iter();
+        let (lower, upper) = observations.size_hint();
+        let mut observed_categories = Vec::with_capacity(upper.unwrap_or(lower));
+        for value in observations {
+            let category = value as usize;
+            assert!(category < cardinality, "Observed category is out of range");
+            observed_categories.push(category);
+        }
+        let prior_mass = 1.0 / (observed_categories.len() + 1) as f64;
         let denominator = 1.0 + cardinality as f64 * prior_mass;
         let matching = (1.0 + prior_mass) / denominator;
         let other = prior_mass / denominator;
@@ -321,6 +321,18 @@ pub(crate) enum Distributions {
     Categorical(CategoricalDistributions),
 }
 
+impl Distributions {
+    pub(crate) fn n_kernels(&self) -> usize {
+        match self {
+            Self::TruncNorm(d) => d.mus.len(),
+            Self::TruncLogNorm(d) => d.mus.len(),
+            Self::DiscreteTruncNorm(d) => d.mus.len(),
+            Self::DiscreteTruncLogNorm(d) => d.mus.len(),
+            Self::Categorical(d) => d.observed_categories.len() + 1,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct MixtureOfProductDistribution {
     pub param_names: Vec<String>,          // Sorted param names
@@ -332,7 +344,7 @@ pub(crate) struct MixtureOfProductDistribution {
 }
 
 impl MixtureOfProductDistribution {
-    pub fn new(distributions_map: HashMap<String, Distributions>, weights: Vec<f64>) -> Self {
+    pub fn new(entries: Vec<(String, Distributions)>, weights: Vec<f64>) -> Self {
         let sum_w = weights.iter().sum::<f64>();
         let log_sum_weights = if sum_w > 0.0 {
             sum_w.ln()
@@ -348,8 +360,6 @@ impl MixtureOfProductDistribution {
         let alias =
             WeightedAliasIndex::new(weights).expect("weights must be non-empty and non-negative");
 
-        let mut entries: Vec<_> = distributions_map.into_iter().collect();
-        entries.sort_unstable_by(|a, b| a.0.cmp(&b.0));
         let (param_names, distributions) = entries.into_iter().unzip();
 
         MixtureOfProductDistribution {
@@ -597,8 +607,8 @@ mod tests {
             5.0,
             1.0,
         );
-        let categorical_dist = CategoricalDistributions::new(&[0.0], 2);
-        let distributions = vec![
+        let categorical_dist = CategoricalDistributions::new([0.0], 2);
+        let mut distributions = vec![
             (
                 "param_truncnorm".to_string(),
                 Distributions::TruncNorm(truncnorm_dist),
@@ -620,10 +630,9 @@ mod tests {
                 Distributions::Categorical(categorical_dist),
             ),
         ];
-        let distributions_map: std::collections::HashMap<String, Distributions> =
-            distributions.into_iter().collect();
+        distributions.sort_by(|a, b| a.0.cmp(&b.0));
         let mixture = MixtureOfProductDistribution::new(
-            distributions_map,
+            distributions,
             vec![0.5, 0.5], // weights.len() == mus.len()
         );
         let mut rng = rand::rngs::StdRng::seed_from_u64(42);
@@ -644,7 +653,7 @@ mod tests {
 
     #[test]
     fn categorical_kernels_keep_smoothed_probabilities_and_uniform_prior() {
-        let distribution = CategoricalDistributions::new(&[1.0, 2.0], 3);
+        let distribution = CategoricalDistributions::new([1.0, 2.0], 3);
         let expected = [
             [1.0 / 6.0, 2.0 / 3.0, 1.0 / 6.0],
             [1.0 / 6.0, 1.0 / 6.0, 2.0 / 3.0],
@@ -659,8 +668,8 @@ mod tests {
             }
         }
         for cardinality in [1, 3, 257] {
-            let prior = CategoricalDistributions::new(&[], cardinality);
-            let singleton = CategoricalDistributions::new(&[0.0; 10], 1);
+            let prior = CategoricalDistributions::new([], cardinality);
+            let singleton = CategoricalDistributions::new([0.0; 10], 1);
             let mut scores = vec![0.0; 1];
             prior.accumulate_log_pdf(0, &mut scores);
             assert_eq!(prior.probability(0, 0), 1.0 / cardinality as f64);
@@ -672,16 +681,16 @@ mod tests {
     #[test]
     fn categorical_joint_density_keeps_observation_correlations_and_weights() {
         let mixture = MixtureOfProductDistribution::new(
-            HashMap::from([
+            vec![
                 (
                     "x".into(),
-                    Distributions::Categorical(CategoricalDistributions::new(&[0.0, 1.0], 2)),
+                    Distributions::Categorical(CategoricalDistributions::new([0.0, 1.0], 2)),
                 ),
                 (
                     "y".into(),
-                    Distributions::Categorical(CategoricalDistributions::new(&[0.0, 1.0], 2)),
+                    Distributions::Categorical(CategoricalDistributions::new([0.0, 1.0], 2)),
                 ),
-            ]),
+            ],
             vec![1.0, 2.0, 0.5],
         );
         let matching = 4.0_f64 / 5.0;