From fa4af9a01bfd4e62fb4ea7e4384ba9f2901fb334 Mon Sep 17 00:00:00 2001 From: TOKUNAGA Hiroyuki <96674+tkng@users.noreply.github.com> Date: Wed, 7 Oct 2026 02:03:13 +0900 Subject: [PATCH] perf(tpe): avoid repeated sorting and observation copies in model construction TPE already prepares observation columns in sorted parameter order. Building the good and poor models sorts the same names again and creates temporary observation and distribution maps. Categorical values are also copied into temporary float buffers before conversion to category indices. Pass one sorted view of names, distributions, and observation columns to both model builders. Construct category-index buffers directly from the selected rows and retain the parameter order through mixture construction. This removes repeated sorting, temporary maps, and intermediate copies. --- rustuna_core/src/parzen_estimator/model.rs | 181 +++++++++++++++++- .../probability_distributions.rs | 59 +++--- rustuna_sampler/src/tpe/sampler.rs | 77 ++++---- 3 files changed, 248 insertions(+), 69 deletions(-) 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(