Add Optuna's sampler tests

rustuna_core/src/sampler.rs

@@ -54,6 +54,62 @@ impl RandomSampler {
         }
     }
 }
+
+fn round_to_step(value: f64, low: f64, high: f64, step: f64) -> f64 {
+    let mut stepped = low + ((value - low) / step).round() * step;
+    if stepped < low {
+        stepped = low;
+    }
+    if stepped > high {
+        stepped = high;
+    }
+    stepped
+}
+
+fn sample_float_with_step(
+    rng: &mut StdRng,
+    low: f64,
+    high: f64,
+    step: Option<f64>,
+    log: bool,
+) -> f64 {
+    match (step, log) {
+        (None, false) => rng.gen_range(low..high),
+        (None, true) => rng.gen_range(low.ln()..high.ln()).exp(),
+        (Some(step), false) => {
+            let max_index = ((high - low) / step).floor().max(0.0) as i64;
+            let index = rng.gen_range(0..=max_index);
+            low + (index as f64) * step
+        }
+        (Some(step), true) => {
+            let value = rng.gen_range(low.ln()..high.ln()).exp();
+            round_to_step(value, low, high, step)
+        }
+    }
+}
+
+fn sample_int_with_step(rng: &mut StdRng, low: i64, high: i64, step: i64, log: bool) -> f64 {
+    let low_f = low as f64;
+    let high_f = high as f64;
+    let step_f = step as f64;
+    if log {
+        let value = rng.gen_range(low_f.ln()..high_f.ln()).exp();
+        let max_index = ((high_f - low_f) / step_f).floor().max(0.0) as i64;
+        let mut index = ((value - low_f) / step_f).round() as i64;
+        if index < 0 {
+            index = 0;
+        }
+        if index > max_index {
+            index = max_index;
+        }
+        low_f + (index as f64) * step_f
+    } else {
+        let max_index = ((high - low) / step).max(0);
+        let index = rng.gen_range(0..=max_index);
+        (low + index * step) as f64
+    }
+}
+
 impl Sampler for RandomSampler {
     fn sample_independent(
         &mut self,
@@ -63,24 +119,23 @@ impl Sampler for RandomSampler {
         distribution: &Distribution,
     ) -> Result<f64> {
         match distribution {
-            // TODO(c-bata): Support step and log
             Distribution::Float {
                 low,
                 high,
-                step: _,
-                log: _,
+                step,
+                log,
             } => {
-                let param_value = self.rng.gen_range(*low..*high);
-                Ok(param_value)
+                let value = sample_float_with_step(&mut self.rng, *low, *high, *step, *log);
+                Ok(value)
             }
             Distribution::Int {
                 low,
                 high,
-                step: _,
-                log: _,
+                step,
+                log,
             } => {
-                let param_value = self.rng.gen_range(*low..*high);
-                Ok(param_value as f64)
+                let value = sample_int_with_step(&mut self.rng, *low, *high, *step, *log);
+                Ok(value)
             }
             Distribution::Categorical { cardinality } => {
                 let param_value = self.rng.gen_range(0..*cardinality);

rustuna_importance/src/test_utils.rs

@@ -20,7 +20,7 @@ fn single_objective(mut trial: Trial) -> Result<Vec<f64>> {
     let x3 = trial.suggest(
         "x3",
         &Distribution::Float {
-            low: 0.0,
+            low: 1.0,
             high: 3.0,
             step: Some(1.0),
             log: true,
@@ -30,7 +30,7 @@ fn single_objective(mut trial: Trial) -> Result<Vec<f64>> {
     let x5 = trial.suggest(
         "x5",
         &Distribution::Int {
-            low: -3,
+            low: 1,
             high: 3,
             step: 1,
             log: true,
@@ -54,7 +54,7 @@ fn multi_objective(mut trial: Trial) -> Result<Vec<f64>> {
     let x3 = trial.suggest(
         "x3",
         &Distribution::Float {
-            low: 0.0,
+            low: 0.1,
             high: 3.0,
             step: Some(1.0),
             log: true,
@@ -64,7 +64,7 @@ fn multi_objective(mut trial: Trial) -> Result<Vec<f64>> {
     let x5 = trial.suggest(
         "x5",
         &Distribution::Int {
-            low: -3,
+            low: 1,
             high: 3,
             step: 1,
             log: true,

rustuna_pyo3/rustuna/_rustuna.pyi

@@ -755,11 +755,18 @@ class Sampler:
     """Factory class for creating sampler instances."""
 
     @classmethod
-    def tpe(cls, seed: int | None = None, multivariate: bool = True) -> Sampler:
+    def tpe(
+        cls,
+        seed: int | None = None,
+        n_startup_trials: int = 10,
+        multivariate: bool = True,
+    ) -> Sampler:
         """Create a Tree-structured Parzen Estimator sampler.
 
         Args:
             seed: Random seed. If None, a random seed is used.
+            n_startup_trials: Number of startup trials before using TPE.
+            multivariate: Whether to use multivariate TPE.
 
         Returns:
             A TPE sampler instance.

rustuna_pyo3/src/sampler.rs

@@ -32,14 +32,11 @@ impl PySampler {
         n_startup_trials: usize,
         multivariate: bool,
     ) -> PyResult<Self> {
-        let rs_sampler = match seed {
-            Some(seed) => TpeSampler::from_config(TpeConfig {
-                seed: Some(seed),
-                n_startup_trials,
-                multivariate,
-            }),
-            None => TpeSampler::new(),
-        };
+        let rs_sampler = TpeSampler::from_config(TpeConfig {
+            seed,
+            n_startup_trials,
+            multivariate,
+        });
         Ok(PySampler {
             sampler: Arc::new(Mutex::new(rs_sampler)),
             kind: "tpe",

rustuna_pyo3/tests/test_optuna_samplers.py

@@ -0,0 +1,22 @@
+from typing import Callable
+
+import pytest
+from optuna.samplers import BaseSampler
+from optuna.testing.pytest_samplers import (
+    BasicSamplerTestCase,
+    MultiObjectiveSamplerTestCase,
+    RelativeSamplerTestCase,
+)
+
+import rustuna
+from rustuna.converter import ToOptunaSampler
+
+
+class TestTpeSampler(
+    BasicSamplerTestCase, RelativeSamplerTestCase, MultiObjectiveSamplerTestCase
+):
+    @pytest.fixture
+    def sampler(self) -> Callable[[], BaseSampler]:
+        return lambda: ToOptunaSampler(
+            rustuna.Sampler.tpe(n_startup_trials=0, multivariate=True)
+        )