Skip to main content

localsearch/optim/
metropolis.rs

1use ordered_float::NotNan;
2use rayon::prelude::*;
3
4use super::{GenericLocalSearchOptimizer, LocalSearchOptimizer, generic::StepResult};
5use crate::{Duration, OptModel, callback::OptCallbackFn};
6
7pub fn metropolis_transition(beta: f64) -> impl Fn(NotNan<f64>, NotNan<f64>) -> f64 {
8    move |current: NotNan<f64>, trial: NotNan<f64>| {
9        let ds = trial - current;
10        if ds <= NotNan::new(0.0).unwrap() {
11            1.0
12        } else {
13            (-beta * ds.into_inner()).exp()
14        }
15    }
16}
17
18// Calculate target based on target_prob
19// p = exp(-beta * ds ) => beta = -ln(p) / ds
20// Average across all energy differences
21pub(crate) fn calculate_temperature_from_acceptance_prob(
22    energy_diffs: &[f64],
23    target_acceptance_prob: f64,
24) -> f64 {
25    let average_energy_diff = energy_diffs.iter().sum::<f64>() / energy_diffs.len() as f64;
26    let ln_prob = target_acceptance_prob.ln().clamp(-100.0, -0.01);
27    -ln_prob / average_energy_diff.clamp(0.01, 100.0)
28}
29
30pub(crate) fn gather_energy_diffs<M: OptModel<ScoreType = NotNan<f64>>>(
31    model: &M,
32    initial_solution_and_score: Option<(M::SolutionType, M::ScoreType)>,
33    n_warmup: usize,
34) -> Vec<f64> {
35    let mut rng = rand::rng();
36    let (current_solution, current_score) =
37        initial_solution_and_score.unwrap_or(model.generate_random_solution(&mut rng).unwrap());
38
39    let energy_diffs: Vec<f64> = (0..n_warmup)
40        .into_par_iter()
41        .filter_map(|_| {
42            let mut rng = rand::rng();
43            let (_, _, trial_score) =
44                model.generate_trial_solution(current_solution.clone(), current_score, &mut rng);
45            let ds = trial_score - current_score;
46            if ds > NotNan::new(0.0).unwrap() {
47                Some(ds.into_inner())
48            } else {
49                None
50            }
51        })
52        .collect();
53
54    energy_diffs
55}
56
57/// Tune inverse temperature beta based on initial random trials
58pub fn tune_temperature<M: OptModel<ScoreType = NotNan<f64>>>(
59    model: &M,
60    initial_solution_and_score: Option<(M::SolutionType, M::ScoreType)>,
61    n_warmup: usize,
62    target_prob: f64,
63) -> f64 {
64    let energy_diffs = gather_energy_diffs(model, initial_solution_and_score, n_warmup);
65    if energy_diffs.is_empty() {
66        1.0
67    } else {
68        calculate_temperature_from_acceptance_prob(&energy_diffs, target_prob)
69    }
70}
71
72/// Optimizer that implements the Metropolis algorithm with constant beta
73#[derive(Clone, Copy)]
74pub struct MetropolisOptimizer {
75    /// The optimizer will give up if there is no improvement of the score after this number of iterations
76    patience: usize,
77    /// Number of trial solutions to generate and evaluate at each iteration
78    n_trials: usize,
79    /// Returns to the best solution if there is no improvement after this number of iterations
80    return_iter: usize,
81    /// Inverse temperature (beta)
82    beta: f64,
83}
84
85impl MetropolisOptimizer {
86    /// Constructor of MetropolisOptimizer
87    ///
88    /// - `patience` : the optimizer will give up
89    ///   if there is no improvement of the score after this number of iterations
90    /// - `n_trials` : number of trial solutions to generate and evaluate at each iteration
91    /// - `return_iter` : returns to the best solution if there is no improvement after this number of iterations.
92    /// - `beta` : inverse temperature
93    pub fn new(patience: usize, n_trials: usize, return_iter: usize, beta: f64) -> Self {
94        Self {
95            patience,
96            n_trials,
97            return_iter,
98            beta,
99        }
100    }
101
102    /// Perform one optimization step
103    pub fn step<M: OptModel<ScoreType = NotNan<f64>>>(
104        &self,
105        model: &M,
106        initial_solution: M::SolutionType,
107        initial_score: M::ScoreType,
108        n_iter: usize,
109        time_limit: Duration,
110        callback: &mut dyn OptCallbackFn<M::SolutionType, M::ScoreType>,
111    ) -> StepResult<M::SolutionType, M::ScoreType> {
112        let transition = |current: NotNan<f64>, trial: NotNan<f64>| {
113            metropolis_transition(self.beta)(current, trial)
114        };
115        let generic_optimizer = GenericLocalSearchOptimizer::new(
116            self.patience,
117            self.n_trials,
118            self.return_iter,
119            transition,
120        );
121        generic_optimizer.step(
122            model,
123            initial_solution,
124            initial_score,
125            n_iter,
126            time_limit,
127            callback,
128        )
129    }
130}
131
132impl<M: OptModel<ScoreType = NotNan<f64>>> LocalSearchOptimizer<M> for MetropolisOptimizer {
133    /// Start optimization
134    ///
135    /// - `model` : the model to optimize
136    /// - `initial_solution` : the initial solution to start optimization
137    /// - `initial_score` : the initial score of the initial solution
138    /// - `n_iter`: maximum iterations
139    /// - `time_limit`: maximum iteration time
140    /// - `callback` : callback function that will be invoked at the end of each iteration
141    fn optimize(
142        &self,
143        model: &M,
144        initial_solution: <M as OptModel>::SolutionType,
145        initial_score: <M as OptModel>::ScoreType,
146        n_iter: usize,
147        time_limit: Duration,
148        callback: &mut dyn OptCallbackFn<<M as OptModel>::SolutionType, <M as OptModel>::ScoreType>,
149    ) -> (<M as OptModel>::SolutionType, <M as OptModel>::ScoreType) {
150        let step_result = self.step(
151            model,
152            initial_solution,
153            initial_score,
154            n_iter,
155            time_limit,
156            callback,
157        );
158        (step_result.best_solution, step_result.best_score)
159    }
160}