From 974011796e26f2005d71b7e5e8227d6014789eb6 Mon Sep 17 00:00:00 2001
From: Stephen Waits <steve@waits.net>
Date: Tue, 5 May 2026 08:02:49 -0600
Subject: [PATCH] feat(algorithms): add AntColonyTsp ant colony optimization
 for TSP-style permutations
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Dorigo-style Ant System for permutation problems on a complete graph:
each generation, every ant constructs a tour by probabilistically
picking the next node from those it has not yet visited, weighted by
`τ_ij^α · η_ij^β` where τ is the pheromone level on edge (i, j) and
η is the heuristic desirability (1 / distance, here). After all ants
finish, pheromone evaporates by a factor `(1 - ρ)` and is reinforced
on each ant's tour proportional to that tour's quality.

Decision type is `Vec<usize>` (a permutation of 0..n_cities). The user
supplies a distance matrix and the n_cities is inferred. Single-objective
only (the cost is total tour length, which the Problem evaluates).

Tests build a 5-city ring and verify ACO finds a near-optimal tour,
plus deterministic reruns and panic on multi-objective.
---
 src/algorithms/ant_colony_tsp.rs | 385 +++++++++++++++++++++++++++++++
 src/algorithms/mod.rs            |   2 +
 src/prelude.rs                   |   3 +-
 3 files changed, 389 insertions(+), 1 deletion(-)
 create mode 100644 src/algorithms/ant_colony_tsp.rs

diff --git a/src/algorithms/ant_colony_tsp.rs b/src/algorithms/ant_colony_tsp.rs
new file mode 100644
index 0000000..ccadabd
--- /dev/null
+++ b/src/algorithms/ant_colony_tsp.rs
@@ -0,0 +1,385 @@
+//! `AntColonyTsp` — Dorigo-style Ant System for permutation problems on a
+//! complete graph (TSP-style).
+
+use rand::Rng as _;
+
+use crate::core::candidate::Candidate;
+use crate::core::objective::Direction;
+use crate::core::population::Population;
+use crate::core::problem::Problem;
+use crate::core::result::OptimizationResult;
+use crate::core::rng::rng_from_seed;
+use crate::traits::Optimizer;
+
+/// Configuration for [`AntColonyTsp`].
+#[derive(Debug, Clone)]
+pub struct AntColonyTspConfig {
+    /// Number of ants per generation.
+    pub ants: usize,
+    /// Number of generations.
+    pub generations: usize,
+    /// Pheromone weight `α`.
+    pub alpha: f64,
+    /// Heuristic weight `β`.
+    pub beta: f64,
+    /// Pheromone evaporation rate `ρ` ∈ [0, 1].
+    pub evaporation: f64,
+    /// Pheromone deposit constant `Q`. Reinforcement on edge (i, j) is
+    /// `Q / tour_length` for every ant whose tour uses (i, j).
+    pub deposit: f64,
+    /// Initial pheromone level on every edge.
+    pub initial_pheromone: f64,
+    /// Seed for the deterministic RNG.
+    pub seed: u64,
+}
+
+impl Default for AntColonyTspConfig {
+    fn default() -> Self {
+        Self {
+            ants: 30,
+            generations: 100,
+            alpha: 1.0,
+            beta: 2.0,
+            evaporation: 0.5,
+            deposit: 1.0,
+            initial_pheromone: 1.0,
+            seed: 42,
+        }
+    }
+}
+
+/// Ant Colony Optimization for permutation-style problems on a complete graph.
+///
+/// `Vec<usize>` decisions only (the permutation `[0, 1, …, n_cities - 1]`).
+/// Single-objective only — typically minimizing total tour length, but the
+/// algorithm is direction-aware for completeness.
+///
+/// Each ant builds a tour by repeatedly choosing the next node with
+/// probability `∝ τ_ij^α · η_ij^β` over the unvisited cities, where
+/// `η_ij = 1 / distance_ij` is the heuristic desirability.
+pub struct AntColonyTsp {
+    /// Algorithm configuration.
+    pub config: AntColonyTspConfig,
+    /// Symmetric distance matrix; size `n_cities × n_cities`. Diagonal must
+    /// be zero.
+    pub distances: Vec<Vec<f64>>,
+}
+
+impl AntColonyTsp {
+    /// Construct an `AntColonyTsp`. Validates that `distances` is square
+    /// and has a zero diagonal.
+    pub fn new(config: AntColonyTspConfig, distances: Vec<Vec<f64>>) -> Self {
+        let n = distances.len();
+        assert!(n >= 2, "AntColonyTsp distances matrix must have >= 2 cities");
+        for (i, row) in distances.iter().enumerate() {
+            assert_eq!(row.len(), n, "AntColonyTsp distances matrix must be square");
+            assert_eq!(row[i], 0.0, "AntColonyTsp distance from city to itself must be 0");
+        }
+        Self { config, distances }
+    }
+}
+
+impl<P> Optimizer<P> for AntColonyTsp
+where
+    P: Problem<Decision = Vec<usize>> + Sync,
+{
+    fn run(&mut self, problem: &P) -> OptimizationResult<P::Decision> {
+        assert!(self.config.ants >= 1, "AntColonyTsp ants must be >= 1");
+        let objectives = problem.objectives();
+        assert!(
+            objectives.is_single_objective(),
+            "AntColonyTsp requires exactly one objective",
+        );
+        let direction = objectives.objectives[0].direction;
+        let n = self.distances.len();
+        let mut rng = rng_from_seed(self.config.seed);
+
+        // Heuristic desirability: 1 / distance (with a small floor to avoid
+        // division by zero for very-close cities).
+        let eta: Vec<Vec<f64>> = self
+            .distances
+            .iter()
+            .map(|row| row.iter().map(|&d| if d > 0.0 { 1.0 / d } else { 0.0 }).collect())
+            .collect();
+
+        // Pheromone matrix.
+        let mut pheromone: Vec<Vec<f64>> =
+            vec![vec![self.config.initial_pheromone; n]; n];
+
+        let mut best_decision: Option<Vec<usize>> = None;
+        let mut best_eval: Option<crate::core::evaluation::Evaluation> = None;
+        let mut evaluations = 0usize;
+
+        for _ in 0..self.config.generations {
+            let mut tours: Vec<Vec<usize>> = Vec::with_capacity(self.config.ants);
+            let mut tour_evals: Vec<crate::core::evaluation::Evaluation> =
+                Vec::with_capacity(self.config.ants);
+
+            for _ in 0..self.config.ants {
+                let start = rng.random_range(0..n);
+                let tour = build_tour(
+                    n,
+                    start,
+                    &pheromone,
+                    &eta,
+                    self.config.alpha,
+                    self.config.beta,
+                    &mut rng,
+                );
+                let eval = problem.evaluate(&tour);
+                evaluations += 1;
+                tours.push(tour);
+                tour_evals.push(eval);
+            }
+
+            // Update best.
+            for (tour, eval) in tours.iter().zip(tour_evals.iter()) {
+                let beats = match &best_eval {
+                    None => true,
+                    Some(b) => better_than_so(eval, b, direction),
+                };
+                if beats {
+                    best_decision = Some(tour.clone());
+                    best_eval = Some(eval.clone());
+                }
+            }
+
+            // Pheromone evaporation.
+            for row in pheromone.iter_mut() {
+                for v in row.iter_mut() {
+                    *v *= 1.0 - self.config.evaporation;
+                }
+            }
+
+            // Pheromone deposit on each ant's tour.
+            for (tour, eval) in tours.iter().zip(tour_evals.iter()) {
+                let length = eval.objectives.first().copied().unwrap_or(f64::INFINITY).max(1e-12);
+                let deposit = self.config.deposit / length;
+                for w in tour.windows(2) {
+                    let (i, j) = (w[0], w[1]);
+                    pheromone[i][j] += deposit;
+                    pheromone[j][i] += deposit;
+                }
+                // Close the loop.
+                let (i, j) = (*tour.last().unwrap(), tour[0]);
+                pheromone[i][j] += deposit;
+                pheromone[j][i] += deposit;
+            }
+        }
+
+        let best = Candidate::new(best_decision.unwrap(), best_eval.unwrap());
+        let population = Population::new(vec![best.clone()]);
+        let front = vec![best.clone()];
+        OptimizationResult::new(
+            population,
+            front,
+            Some(best),
+            evaluations,
+            self.config.generations,
+        )
+    }
+}
+
+fn build_tour(
+    n: usize,
+    start: usize,
+    pheromone: &[Vec<f64>],
+    eta: &[Vec<f64>],
+    alpha: f64,
+    beta: f64,
+    rng: &mut crate::core::rng::Rng,
+) -> Vec<usize> {
+    let mut tour = Vec::with_capacity(n);
+    let mut visited = vec![false; n];
+    tour.push(start);
+    visited[start] = true;
+
+    for _ in 1..n {
+        let current = *tour.last().unwrap();
+        // Build a probability vector over the unvisited candidates.
+        let mut probs: Vec<(usize, f64)> = (0..n)
+            .filter(|&j| !visited[j])
+            .map(|j| {
+                let p = pheromone[current][j].max(0.0).powf(alpha) * eta[current][j].powf(beta);
+                (j, p)
+            })
+            .collect();
+        let total: f64 = probs.iter().map(|(_, p)| *p).sum();
+        let next = if total > 0.0 {
+            let r: f64 = rng.random::<f64>() * total;
+            let mut acc = 0.0;
+            let mut chosen = probs.last().unwrap().0;
+            for (j, p) in &probs {
+                acc += *p;
+                if r <= acc {
+                    chosen = *j;
+                    break;
+                }
+            }
+            chosen
+        } else {
+            // Degenerate case: pheromone × heuristic is 0 for every
+            // unvisited city. Fall back to uniform random.
+            let &(j, _) = probs.choose_uniform(rng);
+            j
+        };
+        let _ = probs;
+        tour.push(next);
+        visited[next] = true;
+    }
+    tour
+}
+
+trait ChooseUniform<T> {
+    fn choose_uniform(&self, rng: &mut crate::core::rng::Rng) -> &T;
+}
+impl<T> ChooseUniform<T> for [T] {
+    fn choose_uniform(&self, rng: &mut crate::core::rng::Rng) -> &T {
+        &self[rng.random_range(0..self.len())]
+    }
+}
+
+fn better_than_so(
+    a: &crate::core::evaluation::Evaluation,
+    b: &crate::core::evaluation::Evaluation,
+    direction: Direction,
+) -> bool {
+    match (a.is_feasible(), b.is_feasible()) {
+        (true, false) => true,
+        (false, true) => false,
+        (false, false) => a.constraint_violation < b.constraint_violation,
+        (true, true) => match direction {
+            Direction::Minimize => a.objectives[0] < b.objectives[0],
+            Direction::Maximize => a.objectives[0] > b.objectives[0],
+        },
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::evaluation::Evaluation;
+    use crate::core::objective::{Objective, ObjectiveSpace};
+
+    /// A 5-city ring problem: cities placed at `(cos(2πi/5), sin(2πi/5))`.
+    /// Optimal tour length: 2·5·sin(π/5) ≈ 5.878 (a regular pentagon).
+    struct RingTsp {
+        distances: Vec<Vec<f64>>,
+    }
+    impl RingTsp {
+        fn new(n: usize) -> Self {
+            use std::f64::consts::PI;
+            let pts: Vec<(f64, f64)> = (0..n)
+                .map(|i| {
+                    let a = 2.0 * PI * (i as f64) / (n as f64);
+                    (a.cos(), a.sin())
+                })
+                .collect();
+            let distances = (0..n)
+                .map(|i| {
+                    (0..n)
+                        .map(|j| {
+                            let (xi, yi) = pts[i];
+                            let (xj, yj) = pts[j];
+                            ((xi - xj).powi(2) + (yi - yj).powi(2)).sqrt()
+                        })
+                        .collect()
+                })
+                .collect();
+            Self { distances }
+        }
+    }
+    impl Problem for RingTsp {
+        type Decision = Vec<usize>;
+
+        fn objectives(&self) -> ObjectiveSpace {
+            ObjectiveSpace::new(vec![Objective::minimize("tour_length")])
+        }
+
+        fn evaluate(&self, tour: &Vec<usize>) -> Evaluation {
+            let n = tour.len();
+            let mut total = 0.0;
+            for w in tour.windows(2) {
+                total += self.distances[w[0]][w[1]];
+            }
+            total += self.distances[tour[n - 1]][tour[0]];
+            Evaluation::new(vec![total])
+        }
+    }
+
+    /// Trivial single-objective problem to test the multi-objective panic.
+    struct DummyMo;
+    impl Problem for DummyMo {
+        type Decision = Vec<usize>;
+
+        fn objectives(&self) -> ObjectiveSpace {
+            ObjectiveSpace::new(vec![
+                Objective::minimize("a"),
+                Objective::minimize("b"),
+            ])
+        }
+
+        fn evaluate(&self, _tour: &Vec<usize>) -> Evaluation {
+            Evaluation::new(vec![0.0, 0.0])
+        }
+    }
+
+    #[test]
+    fn finds_near_optimum_on_5_city_ring() {
+        let problem = RingTsp::new(5);
+        let mut opt = AntColonyTsp::new(
+            AntColonyTspConfig {
+                ants: 10,
+                generations: 30,
+                alpha: 1.0,
+                beta: 3.0,
+                evaporation: 0.5,
+                deposit: 1.0,
+                initial_pheromone: 1.0,
+                seed: 1,
+            },
+            problem.distances.clone(),
+        );
+        let r = opt.run(&problem);
+        let best = r.best.unwrap();
+        // Optimal pentagon perimeter ≈ 5.878. ACO should hit close.
+        assert!(
+            best.evaluation.objectives[0] < 5.95,
+            "got tour length = {}",
+            best.evaluation.objectives[0],
+        );
+    }
+
+    #[test]
+    fn deterministic_with_same_seed() {
+        let problem = RingTsp::new(5);
+        let cfg = AntColonyTspConfig {
+            ants: 8,
+            generations: 10,
+            alpha: 1.0,
+            beta: 2.0,
+            evaporation: 0.5,
+            deposit: 1.0,
+            initial_pheromone: 1.0,
+            seed: 99,
+        };
+        let mut a = AntColonyTsp::new(cfg.clone(), problem.distances.clone());
+        let mut b = AntColonyTsp::new(cfg, problem.distances.clone());
+        let ra = a.run(&problem);
+        let rb = b.run(&problem);
+        assert_eq!(
+            ra.best.unwrap().evaluation.objectives,
+            rb.best.unwrap().evaluation.objectives,
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "exactly one objective")]
+    fn multi_objective_panics() {
+        let mut opt = AntColonyTsp::new(
+            AntColonyTspConfig::default(),
+            vec![vec![0.0, 1.0], vec![1.0, 0.0]],
+        );
+        let _ = opt.run(&DummyMo);
+    }
+}
diff --git a/src/algorithms/mod.rs b/src/algorithms/mod.rs
index bc594ab..ee50aac 100644
--- a/src/algorithms/mod.rs
+++ b/src/algorithms/mod.rs
@@ -1,5 +1,6 @@
 //! Built-in reference optimizers.
 
+pub mod ant_colony_tsp;
 pub mod cma_es;
 pub mod differential_evolution;
 pub mod genetic_algorithm;
@@ -17,6 +18,7 @@ pub mod simulated_annealing;
 pub mod spea2;
 pub mod tabu_search;
 
+pub use ant_colony_tsp::*;
 pub use cma_es::*;
 pub use differential_evolution::*;
 pub use genetic_algorithm::*;
diff --git a/src/prelude.rs b/src/prelude.rs
index c073d09..23bdddc 100644
--- a/src/prelude.rs
+++ b/src/prelude.rs
@@ -22,7 +22,8 @@ pub use crate::operators::{
 };
 
 pub use crate::algorithms::{
-    CmaEs, CmaEsConfig, DifferentialEvolution, DifferentialEvolutionConfig,
+    AntColonyTsp, AntColonyTspConfig, CmaEs, CmaEsConfig, DifferentialEvolution,
+    DifferentialEvolutionConfig,
     GeneticAlgorithm, GeneticAlgorithmConfig, HillClimber, HillClimberConfig, Ibea,
     IbeaConfig, Moead, MoeadConfig, Mopso, MopsoConfig, Nsga2, Nsga2Config, Nsga3, Nsga3Config, Paes, PaesConfig, ParticleSwarm,
     ParticleSwarmConfig, RandomSearch, RandomSearchConfig, SimulatedAnnealing,