diff --git a/Cargo.toml b/Cargo.toml index 16cfcb1..f4a1f35 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,7 @@ tracing = ["dep:tracing"] sobol = ["dep:sobol_burley"] cma-es = ["dep:nalgebra"] visualization = [] +fanova = [] [dev-dependencies] tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } diff --git a/src/fanova.rs b/src/fanova.rs new file mode 100644 index 0000000..bb0dd04 --- /dev/null +++ b/src/fanova.rs @@ -0,0 +1,536 @@ +//! fANOVA (functional ANOVA) parameter importance via random forest. +//! +//! Decomposes the variance of the objective function into contributions +//! from individual parameters (main effects) and parameter interactions. +//! +//! The algorithm: +//! 1. Fits a random forest to `(parameters) -> objective_value` +//! 2. Applies functional ANOVA decomposition to the forest +//! 3. Computes main effects (single-parameter importance) +//! 4. Computes interaction effects (pairwise parameter importance) + +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; + +/// Result of fANOVA analysis. +#[derive(Debug, Clone)] +pub struct FanovaResult { + /// Per-parameter importance (fraction of total variance explained). + /// Sorted by descending importance. + pub main_effects: Vec<(String, f64)>, + /// Pairwise interaction importance (fraction of total variance explained). + /// Sorted by descending importance. + pub interactions: Vec<((String, String), f64)>, +} + +/// Configuration for fANOVA analysis. +#[derive(Debug, Clone)] +pub struct FanovaConfig { + /// Number of trees in the random forest (default: 64). + pub n_trees: usize, + /// Maximum depth of each tree. `None` for unlimited (default: `None`). + pub max_depth: Option, + /// Minimum samples required to split a node (default: 2). + pub min_samples_split: usize, + /// Minimum samples required in a leaf node (default: 1). + pub min_samples_leaf: usize, + /// Random seed for reproducibility (default: `Some(42)`). + pub seed: Option, +} + +impl Default for FanovaConfig { + fn default() -> Self { + Self { + n_trees: 64, + max_depth: None, + min_samples_split: 2, + min_samples_leaf: 1, + seed: Some(42), + } + } +} + +// --- Decision Tree --- + +/// A node in the regression tree (arena-allocated). +#[derive(Debug, Clone)] +enum TreeNode { + Leaf { + value: f64, + n_samples: usize, + }, + Split { + feature: usize, + threshold: f64, + left: usize, + right: usize, + n_samples: usize, + }, +} + +/// A regression decision tree for fANOVA. +#[derive(Debug, Clone)] +struct DecisionTree { + nodes: Vec, +} + +impl DecisionTree { + /// Build a tree from the given data using the specified bootstrap indices. + fn build( + data: &[Vec], + targets: &[f64], + indices: &[usize], + config: &FanovaConfig, + rng: &mut StdRng, + ) -> Self { + let mut tree = Self { nodes: Vec::new() }; + tree.build_node(data, targets, indices, 0, config, rng); + tree + } + + #[allow(clippy::cast_precision_loss)] + fn build_node( + &mut self, + data: &[Vec], + targets: &[f64], + indices: &[usize], + depth: usize, + config: &FanovaConfig, + rng: &mut StdRng, + ) -> usize { + let n = indices.len(); + let mean = indices.iter().map(|&i| targets[i]).sum::() / n as f64; + + // Stopping conditions + if n < config.min_samples_split || config.max_depth.is_some_and(|d| depth >= d) { + let idx = self.nodes.len(); + self.nodes.push(TreeNode::Leaf { + value: mean, + n_samples: n, + }); + return idx; + } + + // Pure node check (all targets identical) + #[allow(clippy::float_cmp)] + if indices.iter().all(|&i| targets[i] == targets[indices[0]]) { + let idx = self.nodes.len(); + self.nodes.push(TreeNode::Leaf { + value: mean, + n_samples: n, + }); + return idx; + } + + let n_features = data[0].len(); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let max_features = ((n_features as f64).sqrt().ceil() as usize) + .max(1) + .min(n_features); + let candidates = partial_shuffle(n_features, max_features, rng); + + // Total variance at this node + let total_var: f64 = indices.iter().map(|&i| (targets[i] - mean).powi(2)).sum(); + if total_var == 0.0 { + let idx = self.nodes.len(); + self.nodes.push(TreeNode::Leaf { + value: mean, + n_samples: n, + }); + return idx; + } + + let mut best_score = f64::NEG_INFINITY; + let mut best_feature = 0; + let mut best_threshold = 0.0; + + for &feat in &candidates { + let mut values: Vec = indices.iter().map(|&i| data[i][feat]).collect(); + values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal)); + values.dedup(); + + if values.len() < 2 { + continue; + } + + for w in values.windows(2) { + let threshold = f64::midpoint(w[0], w[1]); + let (l_sum, l_sq, l_n, r_sum, r_sq, r_n) = + split_stats(data, targets, indices, feat, threshold); + + if l_n < config.min_samples_leaf || r_n < config.min_samples_leaf { + continue; + } + + let l_var = l_sq - l_sum * l_sum / l_n as f64; + let r_var = r_sq - r_sum * r_sum / r_n as f64; + let score = total_var - l_var - r_var; + + if score > best_score { + best_score = score; + best_feature = feat; + best_threshold = threshold; + } + } + } + + if best_score <= 0.0 { + let idx = self.nodes.len(); + self.nodes.push(TreeNode::Leaf { + value: mean, + n_samples: n, + }); + return idx; + } + + let (left_indices, right_indices): (Vec, Vec) = indices + .iter() + .partition(|&&i| data[i][best_feature] <= best_threshold); + + if left_indices.is_empty() || right_indices.is_empty() { + let idx = self.nodes.len(); + self.nodes.push(TreeNode::Leaf { + value: mean, + n_samples: n, + }); + return idx; + } + + // Reserve slot for this split node (placeholder replaced below) + let node_idx = self.nodes.len(); + self.nodes.push(TreeNode::Leaf { + value: 0.0, + n_samples: 0, + }); + + let left = self.build_node(data, targets, &left_indices, depth + 1, config, rng); + let right = self.build_node(data, targets, &right_indices, depth + 1, config, rng); + + self.nodes[node_idx] = TreeNode::Split { + feature: best_feature, + threshold: best_threshold, + left, + right, + n_samples: n, + }; + + node_idx + } + + /// Compute marginal prediction for a given feature subset. + /// + /// Features in `subset` use values from `feature_values`. + /// Features not in `subset` are marginalized by weighting branches + /// proportionally to their training-data fractions. + fn marginal_predict(&self, subset: &[usize], feature_values: &[f64]) -> f64 { + self.marginal_predict_at(0, subset, feature_values) + } + + #[allow(clippy::cast_precision_loss)] + fn marginal_predict_at(&self, idx: usize, subset: &[usize], vals: &[f64]) -> f64 { + match self.nodes[idx] { + TreeNode::Leaf { value, .. } => value, + TreeNode::Split { + feature, + threshold, + left, + right, + n_samples, + } => { + if subset.contains(&feature) { + if vals[feature] <= threshold { + self.marginal_predict_at(left, subset, vals) + } else { + self.marginal_predict_at(right, subset, vals) + } + } else { + let l_n = self.n_samples(left) as f64; + let r_n = self.n_samples(right) as f64; + let total = n_samples as f64; + (l_n / total) * self.marginal_predict_at(left, subset, vals) + + (r_n / total) * self.marginal_predict_at(right, subset, vals) + } + } + } + } + + fn n_samples(&self, idx: usize) -> usize { + match self.nodes[idx] { + TreeNode::Leaf { n_samples, .. } | TreeNode::Split { n_samples, .. } => n_samples, + } + } +} + +// --- Helper Functions --- + +/// Select `k` random indices from `0..n` using partial Fisher-Yates shuffle. +fn partial_shuffle(n: usize, k: usize, rng: &mut StdRng) -> Vec { + let mut indices: Vec = (0..n).collect(); + let k = k.min(n); + for i in 0..k { + let j = rng.random_range(i..n); + indices.swap(i, j); + } + indices.truncate(k); + indices +} + +/// Compute left/right split statistics for variance reduction. +#[allow(clippy::cast_precision_loss)] +fn split_stats( + data: &[Vec], + targets: &[f64], + indices: &[usize], + feature: usize, + threshold: f64, +) -> (f64, f64, usize, f64, f64, usize) { + let (mut l_sum, mut l_sq, mut l_n) = (0.0, 0.0, 0usize); + let (mut r_sum, mut r_sq, mut r_n) = (0.0, 0.0, 0usize); + + for &i in indices { + let y = targets[i]; + if data[i][feature] <= threshold { + l_sum += y; + l_sq += y * y; + l_n += 1; + } else { + r_sum += y; + r_sq += y * y; + r_n += 1; + } + } + + (l_sum, l_sq, l_n, r_sum, r_sq, r_n) +} + +/// Population variance of a slice. +#[allow(clippy::cast_precision_loss)] +fn variance(values: &[f64]) -> f64 { + if values.is_empty() { + return 0.0; + } + let n = values.len() as f64; + let mean = values.iter().sum::() / n; + values.iter().map(|v| (v - mean).powi(2)).sum::() / n +} + +// --- Public API --- + +/// Run fANOVA analysis on pre-processed numerical data. +/// +/// `data` is `n_samples` rows, each with `n_features` columns. +/// `targets` has one entry per sample. +/// `feature_names` maps feature index to human-readable name. +#[allow(clippy::cast_precision_loss)] +pub(crate) fn compute_fanova( + data: &[Vec], + targets: &[f64], + feature_names: &[String], + config: &FanovaConfig, +) -> FanovaResult { + let n_samples = data.len(); + let n_features = data[0].len(); + + let mut rng: StdRng = config + .seed + .map_or_else(rand::make_rng, StdRng::seed_from_u64); + + // Build random forest with bootstrap sampling + let trees: Vec = (0..config.n_trees) + .map(|_| { + let bootstrap: Vec = (0..n_samples) + .map(|_| rng.random_range(0..n_samples)) + .collect(); + DecisionTree::build(data, targets, &bootstrap, config, &mut rng) + }) + .collect(); + + // Compute main effects: V_j = Var[E[f | x_j]] + let main_var: Vec = (0..n_features) + .map(|j| { + let subset = [j]; + let preds: Vec = (0..n_samples) + .map(|i| { + trees + .iter() + .map(|t| t.marginal_predict(&subset, &data[i])) + .sum::() + / trees.len() as f64 + }) + .collect(); + variance(&preds) + }) + .collect(); + + // Compute pairwise interaction effects: V_{j,k} - V_j - V_k + let mut interactions: Vec<((String, String), f64)> = Vec::new(); + for j in 0..n_features { + for k in (j + 1)..n_features { + let subset = [j, k]; + let preds: Vec = (0..n_samples) + .map(|i| { + trees + .iter() + .map(|t| t.marginal_predict(&subset, &data[i])) + .sum::() + / trees.len() as f64 + }) + .collect(); + let joint = variance(&preds); + let interaction = (joint - main_var[j] - main_var[k]).max(0.0); + if interaction > 1e-10 { + interactions.push(( + (feature_names[j].clone(), feature_names[k].clone()), + interaction, + )); + } + } + } + + // Normalize so all importances sum to 1.0 + let total: f64 = + main_var.iter().sum::() + interactions.iter().map(|(_, v)| *v).sum::(); + + let mut main_effects: Vec<(String, f64)> = feature_names + .iter() + .zip(&main_var) + .map(|(name, &v)| (name.clone(), if total > 0.0 { v / total } else { 0.0 })) + .collect(); + main_effects.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(core::cmp::Ordering::Equal)); + + if total > 0.0 { + for entry in &mut interactions { + entry.1 /= total; + } + } + interactions.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(core::cmp::Ordering::Equal)); + + FanovaResult { + main_effects, + interactions, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn single_dominant_parameter() { + // f(x, y) = x — only x matters + let mut rng = StdRng::seed_from_u64(0); + let n = 100; + let data: Vec> = (0..n) + .map(|_| vec![rng.random_range(0.0..10.0), rng.random_range(0.0..10.0)]) + .collect(); + let targets: Vec = data.iter().map(|row| row[0]).collect(); + + let result = compute_fanova( + &data, + &targets, + &["x".into(), "y".into()], + &FanovaConfig::default(), + ); + + assert_eq!(result.main_effects[0].0, "x"); + assert!( + result.main_effects[0].1 > 0.8, + "x importance = {}", + result.main_effects[0].1 + ); + } + + #[test] + fn interaction_detection() { + // f(x, y) = x * y — both matter and interact + let mut rng = StdRng::seed_from_u64(0); + let n = 200; + let data: Vec> = (0..n) + .map(|_| vec![rng.random_range(0.0..10.0), rng.random_range(0.0..10.0)]) + .collect(); + let targets: Vec = data.iter().map(|row| row[0] * row[1]).collect(); + + let config = FanovaConfig { + n_trees: 128, + ..FanovaConfig::default() + }; + let result = compute_fanova(&data, &targets, &["x".into(), "y".into()], &config); + + assert!( + !result.interactions.is_empty(), + "should detect x*y interaction" + ); + assert!( + result.interactions[0].1 > 0.05, + "interaction importance = {}", + result.interactions[0].1 + ); + } + + #[test] + fn variance_computation() { + assert!((variance(&[1.0, 2.0, 3.0, 4.0, 5.0]) - 2.0).abs() < 1e-10); + assert!(variance(&[5.0, 5.0, 5.0]).abs() < 1e-10); + assert!(variance(&[]).abs() < 1e-10); + } + + #[test] + fn three_params_one_dominant() { + // f(x, y, z) = 3*x + 0.1*y + 0*z + let mut rng = StdRng::seed_from_u64(7); + let n = 150; + let data: Vec> = (0..n) + .map(|_| { + vec![ + rng.random_range(0.0..10.0), + rng.random_range(0.0..10.0), + rng.random_range(0.0..10.0), + ] + }) + .collect(); + let targets: Vec = data.iter().map(|r| 3.0 * r[0] + 0.1 * r[1]).collect(); + + let result = compute_fanova( + &data, + &targets, + &["x".into(), "y".into(), "z".into()], + &FanovaConfig::default(), + ); + + // x should be the most important + assert_eq!(result.main_effects[0].0, "x"); + assert!(result.main_effects[0].1 > 0.5); + + // z should have near-zero importance + let z_imp = result + .main_effects + .iter() + .find(|(name, _)| name == "z") + .map_or(0.0, |(_, v)| *v); + assert!(z_imp < 0.1, "z importance = {z_imp}"); + } + + #[test] + fn importances_sum_to_one() { + let mut rng = StdRng::seed_from_u64(3); + let n = 100; + let data: Vec> = (0..n) + .map(|_| vec![rng.random_range(0.0..10.0), rng.random_range(0.0..10.0)]) + .collect(); + let targets: Vec = data.iter().map(|r| r[0] + r[1]).collect(); + + let result = compute_fanova( + &data, + &targets, + &["x".into(), "y".into()], + &FanovaConfig::default(), + ); + + let total: f64 = result.main_effects.iter().map(|(_, v)| *v).sum::() + + result.interactions.iter().map(|(_, v)| *v).sum::(); + assert!( + (total - 1.0).abs() < 1e-10, + "importances should sum to 1.0, got {total}" + ); + } +} diff --git a/src/lib.rs b/src/lib.rs index 41832ba..27e9ac8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -219,6 +219,8 @@ macro_rules! trace_debug { mod distribution; mod error; +#[cfg(feature = "fanova")] +mod fanova; mod importance; mod kde; pub mod multi_objective; @@ -234,6 +236,8 @@ mod types; mod visualization; pub use error::{Error, Result, TrialPruned}; +#[cfg(feature = "fanova")] +pub use fanova::{FanovaConfig, FanovaResult}; pub use multi_objective::{MultiObjectiveSampler, MultiObjectiveStudy, MultiObjectiveTrial}; #[cfg(feature = "derive")] pub use optimizer_derive::Categorical; @@ -274,6 +278,8 @@ pub mod prelude { pub use optimizer_derive::Categorical as DeriveCategory; pub use crate::error::{Error, Result, TrialPruned}; + #[cfg(feature = "fanova")] + pub use crate::fanova::{FanovaConfig, FanovaResult}; pub use crate::multi_objective::{MultiObjectiveStudy, MultiObjectiveTrial}; pub use crate::param::ParamValue; pub use crate::parameter::{ diff --git a/src/study.rs b/src/study.rs index 52bccb9..e2ef3d0 100644 --- a/src/study.rs +++ b/src/study.rs @@ -2027,6 +2027,110 @@ where scores } + + /// Computes parameter importance using fANOVA (functional ANOVA) with + /// default configuration. + /// + /// Fits a random forest to the trial data and decomposes variance into + /// per-parameter main effects and pairwise interaction effects. This is + /// more accurate than correlation-based importance ([`Self::param_importance`]) + /// and can detect non-linear relationships and parameter interactions. + /// + /// # Errors + /// + /// Returns [`Error::NoCompletedTrials`] if fewer than 2 trials have completed. + #[cfg(feature = "fanova")] + pub fn fanova(&self) -> crate::Result { + self.fanova_with_config(&crate::fanova::FanovaConfig::default()) + } + + /// Computes parameter importance using fANOVA with custom configuration. + /// + /// See [`Self::fanova`] for details. The [`FanovaConfig`](crate::fanova::FanovaConfig) + /// allows tuning the number of trees, tree depth, and random seed. + /// + /// # Errors + /// + /// Returns [`Error::NoCompletedTrials`] if fewer than 2 trials have completed. + #[cfg(feature = "fanova")] + #[allow(clippy::cast_precision_loss)] + pub fn fanova_with_config( + &self, + config: &crate::fanova::FanovaConfig, + ) -> crate::Result { + use std::collections::BTreeSet; + + use crate::fanova::compute_fanova; + use crate::param::ParamValue; + use crate::types::TrialState; + + let trials = self.completed_trials.read(); + let complete: Vec<_> = trials + .iter() + .filter(|t| t.state == TrialState::Complete) + .collect(); + + if complete.len() < 2 { + return Err(crate::Error::NoCompletedTrials); + } + + // Collect all parameter IDs in a stable order. + let all_param_ids: Vec<_> = { + let set: BTreeSet<_> = complete.iter().flat_map(|t| t.params.keys()).collect(); + set.into_iter().collect() + }; + + if all_param_ids.is_empty() { + return Ok(crate::fanova::FanovaResult { + main_effects: Vec::new(), + interactions: Vec::new(), + }); + } + + // Build feature matrix (only trials that have all parameters). + let mut data = Vec::new(); + let mut targets = Vec::new(); + + for trial in &complete { + let mut row = Vec::with_capacity(all_param_ids.len()); + let mut has_all = true; + + for &pid in &all_param_ids { + if let Some(pv) = trial.params.get(pid) { + row.push(match *pv { + ParamValue::Float(v) => v, + ParamValue::Int(v) => v as f64, + ParamValue::Categorical(v) => v as f64, + }); + } else { + has_all = false; + break; + } + } + + if has_all { + data.push(row); + targets.push(trial.value.clone().into()); + } + } + + if data.len() < 2 { + return Err(crate::Error::NoCompletedTrials); + } + + // Build feature names from parameter labels. + let feature_names: Vec = all_param_ids + .iter() + .map(|&pid| { + complete + .iter() + .find_map(|t| t.param_labels.get(pid)) + .map_or_else(|| pid.to_string(), Clone::clone) + }) + .collect(); + + Ok(compute_fanova(&data, &targets, &feature_names, config)) + } } impl IntoIterator for &Study diff --git a/tests/fanova_tests.rs b/tests/fanova_tests.rs new file mode 100644 index 0000000..55f5c1b --- /dev/null +++ b/tests/fanova_tests.rs @@ -0,0 +1,88 @@ +//! Integration tests for fANOVA parameter importance. + +#![cfg(feature = "fanova")] + +use optimizer::prelude::*; + +#[test] +fn fanova_dominant_parameter() { + // f(x, y) = x^2 — x should dominate + let x = FloatParam::new(0.0, 10.0).name("x"); + let y = FloatParam::new(0.0, 10.0).name("y"); + + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); + study + .optimize(50, |trial| { + let xv = x.suggest(trial)?; + let _yv = y.suggest(trial)?; + Ok::<_, Error>(xv * xv) + }) + .unwrap(); + + let result = study.fanova().unwrap(); + assert_eq!(result.main_effects[0].0, "x"); + assert!( + result.main_effects[0].1 > 0.7, + "x importance = {}", + result.main_effects[0].1 + ); +} + +#[test] +fn fanova_interaction() { + // f(x, y) = x * y — both matter and interact + let x = FloatParam::new(0.0, 10.0).name("x"); + let y = FloatParam::new(0.0, 10.0).name("y"); + + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(7)); + study + .optimize(100, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, Error>(xv * yv) + }) + .unwrap(); + + let config = FanovaConfig { + n_trees: 128, + ..FanovaConfig::default() + }; + let result = study.fanova_with_config(&config).unwrap(); + + // Should detect interaction + assert!( + !result.interactions.is_empty(), + "should detect x*y interaction" + ); +} + +#[test] +fn fanova_consistent_with_correlation() { + // f(x, y) = 3*x + 0.5*y — x should rank higher in both methods + let x = FloatParam::new(0.0, 10.0).name("x"); + let y = FloatParam::new(0.0, 10.0).name("y"); + + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(99)); + study + .optimize(80, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, Error>(3.0 * xv + 0.5 * yv) + }) + .unwrap(); + + let corr = study.param_importance(); + let fanova = study.fanova().unwrap(); + + // Both methods should rank x above y + assert_eq!(corr[0].0, "x", "correlation should rank x first"); + assert_eq!(fanova.main_effects[0].0, "x", "fanova should rank x first"); +} + +#[test] +fn fanova_too_few_trials() { + let study: Study = Study::new(Direction::Minimize); + + let result = study.fanova(); + assert!(result.is_err(), "should error with no trials"); +}