diff --git a/rustuna_core/src/parzen_estimator/model.rs b/rustuna_core/src/parzen_estimator/model.rs index 19595755..84498a7a 100644 --- a/rustuna_core/src/parzen_estimator/model.rs +++ b/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> 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> Iterator for ObservationValues<'_, I> { + type Item = f64; + + fn next(&mut self) -> Option { + 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) { + match self { + Self::Borrowed(values) => (values.len(), Some(values.len())), + Self::Selected(values) => values.size_hint(), + } + } +} + impl ParzenEstimator { pub fn new( observations: &HashMap>, @@ -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::>::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, + V: IntoIterator, + { + 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)>, + V: Iterator, + { + 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, 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, 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, + ); + } } diff --git a/rustuna_core/src/parzen_estimator/probability_distributions.rs b/rustuna_core/src/parzen_estimator/probability_distributions.rs index d92b61a7..f3acd3d3 100644 --- a/rustuna_core/src/parzen_estimator/probability_distributions.rs +++ b/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, 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, // Sorted param names @@ -332,7 +344,7 @@ pub(crate) struct MixtureOfProductDistribution { } impl MixtureOfProductDistribution { - pub fn new(distributions_map: HashMap, weights: Vec) -> Self { + pub fn new(entries: Vec<(String, Distributions)>, weights: Vec) -> Self { let sum_w = weights.iter().sum::(); 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 = - 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; diff --git a/rustuna_sampler/src/tpe/sampler.rs b/rustuna_sampler/src/tpe/sampler.rs index 3c7474b1..054e19d3 100644 --- a/rustuna_sampler/src/tpe/sampler.rs +++ b/rustuna_sampler/src/tpe/sampler.rs @@ -200,10 +200,22 @@ struct TpeObservationsView<'a> { trial_numbers: &'a [u32], values: &'a [f64], n_objectives: usize, - param_columns: Vec]>>, + param_columns: Vec>, feasibles_violations: &'a [(bool, f64)], } +struct TpeObservationColumn<'a> { + name: &'a str, + distribution: &'a Distribution, + values: Option<&'a [Option]>, +} + +impl TpeObservationColumn<'_> { + fn at(&self, row: usize) -> Option { + self.values.and_then(|values| values[row]) + } +} + impl TpeObservationSet { fn insert(&mut self, trial: &PersistedTrial) -> Result<()> { let TrialStateValues::Complete(values) = &trial.state_values else { @@ -263,7 +275,10 @@ impl TpeObservationSet { } } - fn view(&self, search_space: &HashMap) -> TpeObservationsView<'_> { + fn view<'a>( + &'a self, + search_space: &'a HashMap, + ) -> TpeObservationsView<'a> { let mut sorted_keys: Vec<&String> = search_space.keys().collect(); sorted_keys.sort(); TpeObservationsView { @@ -271,8 +286,12 @@ impl TpeObservationSet { values: &self.values, n_objectives: self.n_objectives, param_columns: sorted_keys - .iter() - .map(|name| self.param_columns.get(*name).map(Vec::as_slice)) + .into_iter() + .map(|name| TpeObservationColumn { + name, + distribution: &search_space[name], + values: self.param_columns.get(name).map(Vec::as_slice), + }) .collect(), feasibles_violations: &self.feasibles_violations, } @@ -287,10 +306,6 @@ impl TpeObservationsView<'_> { fn values_row(&self, row: usize) -> &[f64] { &self.values[row * self.n_objectives..(row + 1) * self.n_objectives] } - - fn param_at(&self, row: usize, param_idx: usize) -> Option { - self.param_columns[param_idx].and_then(|col| col[row]) - } } /// Tree-structured Parzen Estimator sampler. @@ -384,7 +399,6 @@ impl TpeSampler { &self, ctx: &Context, observations: &TpeObservationsView<'_>, - search_space: &HashMap, ) -> Result> { let n = observations.len(); let is_multi_objective = ctx.directions.len() > 1; @@ -396,8 +410,8 @@ impl TpeSampler { Self::split_rows_for_single_objective(observations, direction, gamma); // Single-objective: recency ramp for both l(x) and g(x) (Optuna default_weights). ( - Self::build_parzen_estimator(observations, &good_rows, search_space, true), - Self::build_parzen_estimator(observations, &poor_rows, search_space, true), + Self::build_parzen_estimator(observations, &good_rows, true), + Self::build_parzen_estimator(observations, &poor_rows, true), ) } else { let directions: &[Direction] = &ctx.directions; @@ -455,8 +469,8 @@ impl TpeSampler { // Multi-objective: uniform weights for l(x) (below), recency ramp for g(x) (above), // matching Optuna's multi-objective TPE. ( - Self::build_parzen_estimator(observations, &good_rows, search_space, false), - Self::build_parzen_estimator(observations, &poor_rows, search_space, true), + Self::build_parzen_estimator(observations, &good_rows, false), + Self::build_parzen_estimator(observations, &poor_rows, true), ) }; @@ -696,15 +710,10 @@ impl TpeSampler { fn build_parzen_estimator( observations: &TpeObservationsView<'_>, rows: &[usize], - search_space: &HashMap, recency_ramp: bool, ) -> ParzenEstimator { - let mut sorted_keys: Vec<&String> = search_space.keys().collect(); - sorted_keys.sort(); - - let n_params = sorted_keys.len(); + let n_params = observations.param_columns.len(); let n_trials = rows.len(); - debug_assert_eq!(n_params, observations.param_columns.len()); // Process trials in chronological (trial-number ascending) order so that the // recency ramp assigns the highest weights to the most recent trials, matching @@ -713,16 +722,11 @@ impl TpeSampler { let mut order: Vec = rows.to_vec(); order.sort_by_key(|&row| observations.trial_numbers[row]); - let mut observations_vec: Vec> = (0..n_params) - .map(|_| Vec::with_capacity(n_trials)) - .collect(); let mut active_counts: Vec = vec![0; n_trials]; - // Keep each input and output column local while preserving row order. - for (param_idx, column) in observations_vec.iter_mut().enumerate() { + for column in &observations.param_columns { for (trial_idx, &row) in order.iter().enumerate() { - if let Some(v) = observations.param_at(row, param_idx) { - column.push(v); + if column.at(row).is_some() { active_counts[trial_idx] += 1; } } @@ -741,12 +745,17 @@ impl TpeSampler { }; let active_weights: Vec = active_indices.iter().map(|&i| weights[i]).collect(); let prior_weight = 1.0; - let observations: HashMap> = sorted_keys - .iter() - .zip(observations_vec) - .map(|(k, v)| ((*k).clone(), v)) - .collect(); - ParzenEstimator::new(&observations, search_space, &active_weights, prior_weight) + ParzenEstimator::from_ordered_observations( + observations.param_columns.iter().map(|column| { + ( + column.name, + column.distribution, + order.iter().filter_map(|&row| column.at(row)), + ) + }), + &active_weights, + prior_weight, + ) } } @@ -799,7 +808,7 @@ impl Sampler for TpeSampler { .random_sampler .sample_independent(ctx, storage, name, distribution); }; - let params = self.sample(ctx, &observations.view(&search_space), &search_space)?; + let params = self.sample(ctx, &observations.view(&search_space))?; Ok(params[name]) } @@ -834,7 +843,7 @@ impl Sampler for TpeSampler { let Some(observations) = self.observations_for_trial(ctx, &storage)? else { return Ok(HashMap::new()); }; - self.sample(ctx, &observations.view(search_space), search_space) + self.sample(ctx, &observations.view(search_space)) } fn after_trial(