symbios-genetics 0.2.0

//! CVT-MAP-Elites: a Centroidal Voronoi Tessellation variant of MAP-Elites.
//!
//! Where [`crate::algorithms::map_elites::MapElites`] partitions behaviour space into a
//! regular grid, CVT-MAP-Elites partitions it into Voronoi cells around a
//! fixed set of centroids. This gives more uniform archive coverage when the
//! feasible region of behaviour space is irregular or unevenly distributed,
//! and decouples archive size from descriptor dimensionality (a 10-dim grid
//! at resolution 10 is `10^10 = 10 billion` cells; the CVT equivalent is
//! whatever you set `num_centroids` to).
//!
//! # Algorithm
//!
//! 1. Compute centroids offline (typically via Lloyd's algorithm against a
//!    uniform sample of the unit hypercube — see [`CvtMapElites::with_lloyd`](crate::algorithms::cvt_map_elites::CvtMapElites::with_lloyd)).
//! 2. Each step samples a parent from the archive, mutates and evaluates it,
//!    then assigns the offspring to the nearest centroid (squared Euclidean).
//! 3. The archive stores at most one elite per centroid; replacement uses
//!    strict-better fitness, identical to grid MAP-Elites.
//!
//! Descriptors are not clamped: any descriptor vector matching the centroid
//! dimensionality is accepted. The user controls the descriptor range via
//! their evaluator. Centroids generated by
//! [`CvtMapElites::with_lloyd`](crate::algorithms::cvt_map_elites::CvtMapElites::with_lloyd)
//! live in `[0, 1]^d`.
//!
//! # Determinism
//!
//! With a fixed seed,
//! [`CvtMapElites::with_lloyd`](crate::algorithms::cvt_map_elites::CvtMapElites::with_lloyd)
//! produces identical centroids and the algorithm produces bit-identical
//! archives across runs.
//!
//! # References
//!
//! Vassiliades, V., Chatzilygeroudis, K., & Mouret, J.-B. (2017).
//! Using Centroidal Voronoi Tessellations to Scale Up the Multidimensional
//! Archive of Phenotypic Elites Algorithm. IEEE Transactions on Evolutionary
//! Computation.

use crate::algorithms::archive::Archive;
use crate::{Evaluator, Evolver, Genotype, Phenotype};
use rand::Rng;
use rand::prelude::SeedableRng;
use rand_pcg::Pcg64;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

#[cfg(feature = "parallel")]
use rayon::prelude::*;

/// Internal representation for serialization (excludes transient cache).
#[derive(Serialize, Deserialize)]
#[serde(bound = "G: Genotype")]
struct CvtMapElitesData<G: Genotype> {
    archive: BTreeMap<Vec<usize>, Phenotype<G>>,
    archive_keys_vec: Vec<Vec<usize>>,
    centroids: Vec<Vec<f32>>,
    mutation_rate: f32,
    batch_size: usize,
    rng: Pcg64,
}

/// CVT-MAP-Elites: MAP-Elites with a Voronoi tessellation behaviour space.
///
/// Each cell corresponds to one centroid; a descriptor maps to a single
/// cell key `[centroid_idx]`. See module docs for details.
pub struct CvtMapElites<G: Genotype> {
    archive: Archive<G>,
    centroids: Vec<Vec<f32>>,
    mutation_rate: f32,
    batch_size: usize,
    rng: Pcg64,
}

impl<G: Genotype> Serialize for CvtMapElites<G> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        use serde::ser::SerializeStruct;
        let mut state = serializer.serialize_struct("CvtMapElites", 6)?;
        state.serialize_field("archive", self.archive.cells())?;
        state.serialize_field("archive_keys_vec", self.archive.keys_vec())?;
        state.serialize_field("centroids", &self.centroids)?;
        state.serialize_field("mutation_rate", &self.mutation_rate)?;
        state.serialize_field("batch_size", &self.batch_size)?;
        state.serialize_field("rng", &self.rng)?;
        state.end()
    }
}

impl<'de, G: Genotype> Deserialize<'de> for CvtMapElites<G> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        use serde::de::Error;

        let data = CvtMapElitesData::<G>::deserialize(deserializer)?;

        validate_centroids(&data.centroids).map_err(D::Error::custom)?;
        if data.batch_size == 0 {
            return Err(D::Error::custom("batch_size must be greater than 0"));
        }

        let archive = Archive::from_raw(data.archive, data.archive_keys_vec);
        archive.validate().map_err(D::Error::custom)?;

        // Defensive: every archive key must reference a centroid that exists.
        // Each grid CVT key is a single-element vec [centroid_idx].
        for key in archive.keys() {
            if key.len() != 1 || key[0] >= data.centroids.len() {
                return Err(D::Error::custom(
                    "archive key references a centroid index out of range",
                ));
            }
        }

        Ok(Self {
            archive,
            centroids: data.centroids,
            mutation_rate: data.mutation_rate,
            batch_size: data.batch_size,
            rng: data.rng,
        })
    }
}

fn validate_centroids(centroids: &[Vec<f32>]) -> Result<(), &'static str> {
    if centroids.is_empty() {
        return Err("centroids must be non-empty");
    }
    let dim = centroids[0].len();
    if dim == 0 {
        return Err("centroids must have at least one dimension");
    }
    for c in centroids {
        if c.len() != dim {
            return Err("all centroids must have the same dimension");
        }
        if c.iter().any(|v| !v.is_finite()) {
            return Err("centroids must contain only finite values");
        }
    }
    Ok(())
}

impl<G: Genotype> CvtMapElites<G> {
    /// Creates a CVT-MAP-Elites instance with caller-provided centroids.
    ///
    /// Use this when centroids are precomputed (e.g. loaded from a file or
    /// computed once and reused across runs). For the convenience case of
    /// running Lloyd's algorithm inline, see [`with_lloyd`](Self::with_lloyd).
    ///
    /// # Arguments
    ///
    /// * `centroids` - Non-empty list of centroid points. All centroids must
    ///   share the same descriptor dimension and contain only finite values.
    /// * `mutation_rate` - Probability of mutation, typically in `[0.0, 1.0]`.
    /// * `batch_size` - Number of offspring generated per [`step`](Self::step).
    ///   Must be > 0.
    /// * `seed` - RNG seed for deterministic execution.
    ///
    /// # Panics
    ///
    /// Panics if `centroids` is empty, contains differently-shaped vectors,
    /// contains non-finite values, or if `batch_size` is 0.
    pub fn new(centroids: Vec<Vec<f32>>, mutation_rate: f32, batch_size: usize, seed: u64) -> Self {
        validate_centroids(&centroids).expect("invalid centroids");
        assert!(batch_size > 0, "batch_size must be greater than 0");
        Self {
            archive: Archive::new(),
            centroids,
            mutation_rate,
            batch_size,
            rng: Pcg64::seed_from_u64(seed),
        }
    }

    /// Creates a CVT-MAP-Elites instance, computing centroids inline via
    /// Lloyd's algorithm against uniform random samples of `[0, 1]^descriptor_dim`.
    ///
    /// # Arguments
    ///
    /// * `num_centroids` - Number of Voronoi cells in the archive. Must be > 0.
    /// * `descriptor_dim` - Behavioural descriptor dimensionality. Must be > 0.
    /// * `num_samples` - Sample points used per Lloyd iteration. Larger gives
    ///   more uniform tessellation at the cost of one-time setup time.
    ///   Typical values: 10_000 for `descriptor_dim <= 4`, 50_000+ for higher
    ///   dimensions. Must be > 0.
    /// * `lloyd_iters` - Number of Lloyd refinement passes. 25-50 is usually
    ///   sufficient for convergence. Must be > 0.
    /// * `mutation_rate`, `batch_size`, `seed` - As for [`new`](Self::new).
    ///
    /// # Determinism
    ///
    /// Centroid generation is deterministic given `seed`: the same seed
    /// produces the same centroids, allowing reproducible CVT runs.
    ///
    /// # Panics
    ///
    /// Panics if any of `num_centroids`, `descriptor_dim`, `num_samples`,
    /// `lloyd_iters`, or `batch_size` is 0.
    pub fn with_lloyd(
        num_centroids: usize,
        descriptor_dim: usize,
        num_samples: usize,
        lloyd_iters: usize,
        mutation_rate: f32,
        batch_size: usize,
        seed: u64,
    ) -> Self {
        assert!(num_centroids > 0, "num_centroids must be greater than 0");
        assert!(descriptor_dim > 0, "descriptor_dim must be greater than 0");
        assert!(num_samples > 0, "num_samples must be greater than 0");
        assert!(lloyd_iters > 0, "lloyd_iters must be greater than 0");
        assert!(batch_size > 0, "batch_size must be greater than 0");

        let mut rng = Pcg64::seed_from_u64(seed);
        let centroids = lloyd_centroids(
            num_centroids,
            descriptor_dim,
            num_samples,
            lloyd_iters,
            &mut rng,
        );

        Self {
            archive: Archive::new(),
            centroids,
            mutation_rate,
            batch_size,
            rng,
        }
    }

    /// Returns the centroid set defining the Voronoi tessellation.
    pub fn centroids(&self) -> &[Vec<f32>] {
        &self.centroids
    }

    /// Returns the number of centroids (= maximum archive size).
    pub fn num_centroids(&self) -> usize {
        self.centroids.len()
    }

    /// Returns the current mutation rate.
    pub fn mutation_rate(&self) -> f32 {
        self.mutation_rate
    }

    /// Sets the mutation rate.
    pub fn set_mutation_rate(&mut self, rate: f32) {
        self.mutation_rate = rate;
    }

    /// Returns the batch size (offspring per step).
    pub fn batch_size(&self) -> usize {
        self.batch_size
    }

    /// Sets the batch size.
    ///
    /// # Panics
    ///
    /// Panics if `size` is 0.
    pub fn set_batch_size(&mut self, size: usize) {
        assert!(size > 0, "batch_size must be greater than 0");
        self.batch_size = size;
    }

    /// Number of occupied cells in the archive.
    pub fn archive_len(&self) -> usize {
        self.archive.len()
    }

    /// Gets the elite assigned to a given centroid index, or `None` if the
    /// cell is empty.
    pub fn archive_get(&self, centroid_idx: usize) -> Option<&Phenotype<G>> {
        self.archive.get(&[centroid_idx])
    }

    /// Iterator over occupied cell keys (each a single-element `Vec<usize>`
    /// containing the centroid index).
    pub fn archive_keys(&self) -> impl Iterator<Item = &Vec<usize>> {
        self.archive.keys()
    }

    /// Iterator over (centroid-key, elite) pairs in deterministic order.
    pub fn archive_iter(&self) -> impl Iterator<Item = (&Vec<usize>, &Phenotype<G>)> {
        self.archive.iter()
    }

    /// Returns the elite with the highest fitness across all cells.
    pub fn best_by_fitness(&self) -> Option<&Phenotype<G>> {
        self.archive.best_by_fitness()
    }

    /// Returns the fraction of centroids occupied: `archive_len / num_centroids`.
    pub fn coverage(&self) -> f64 {
        self.archive.len() as f64 / self.centroids.len() as f64
    }

    /// Returns the QD score (sum of fitness across occupied cells).
    ///
    /// See [`MapElites::qd_score`](crate::algorithms::map_elites::MapElites::qd_score) for
    /// the negative-fitness caveat.
    pub fn qd_score(&self) -> f64 {
        self.archive.qd_score()
    }

    /// Exports the archive as CSV. Available with the `export` feature.
    /// See [`MapElites::export_csv`](crate::algorithms::map_elites::MapElites::export_csv)
    /// for the column layout.
    #[cfg(feature = "export")]
    pub fn export_csv<W: std::io::Write>(&self, writer: W) -> std::io::Result<()> {
        self.archive.export_csv(writer)
    }

    /// Returns the index of the centroid nearest to `descriptor` under
    /// squared Euclidean distance.
    ///
    /// # Panics
    ///
    /// Panics if `descriptor.len() != centroid_dim`.
    pub fn assign_to_centroid(&self, descriptor: &[f32]) -> usize {
        let dim = self.centroids[0].len();
        assert_eq!(
            descriptor.len(),
            dim,
            "descriptor dimension {} does not match centroid dimension {}",
            descriptor.len(),
            dim
        );
        nearest_centroid_index(descriptor, &self.centroids)
    }

    /// Seeds the archive with initial individuals.
    ///
    /// Each individual is evaluated and assigned to its nearest centroid.
    /// Replacement follows the strict-better-fitness rule of grid MAP-Elites.
    pub fn seed_population<E: Evaluator<G>>(&mut self, initial: Vec<G>, evaluator: &E) {
        let dim = self.centroids[0].len();
        for dna in initial {
            let (f, obj, desc) = evaluator.evaluate(&dna);

            if f.is_nan() || desc.iter().any(|v| v.is_nan()) {
                continue;
            }
            // Mismatched descriptor dimensions can't be assigned; skip rather than panic
            // (the archive could otherwise be silently corrupted by wrong-dim seed inputs).
            if desc.len() != dim {
                continue;
            }

            let key = vec![nearest_centroid_index(&desc, &self.centroids)];
            self.archive.insert_if_better(
                key,
                Phenotype {
                    genotype: dna,
                    fitness: f,
                    objectives: obj,
                    descriptor: desc,
                },
            );
        }
    }
}

impl<G: Genotype> Evolver<G> for CvtMapElites<G> {
    fn step<E: Evaluator<G>>(&mut self, evaluator: &E) {
        if self.archive.is_empty() {
            return;
        }

        let mutation_rate = self.mutation_rate;
        let dim = self.centroids[0].len();

        let selections: Vec<(Vec<usize>, u64)> = (0..self.batch_size)
            .map(|_| {
                let key = self
                    .archive
                    .sample_key(&mut self.rng)
                    .expect("archive non-empty checked above")
                    .clone();
                let seed = self.rng.random::<u64>();
                (key, seed)
            })
            .collect();

        let parents: Vec<G> = selections
            .iter()
            .map(|(key, _)| {
                self.archive
                    .get(key)
                    .expect("sampled key exists in archive")
                    .genotype
                    .clone()
            })
            .collect();

        #[cfg(feature = "parallel")]
        let results: Vec<(G, f32, Vec<f32>, Vec<f32>)> = parents
            .into_par_iter()
            .zip(selections.into_par_iter())
            .map(|(mut dna, (_, seed))| {
                let mut rng = Pcg64::seed_from_u64(seed);
                dna.mutate(&mut rng, mutation_rate);
                let (f, obj, desc) = evaluator.evaluate(&dna);
                (dna, f, obj, desc)
            })
            .collect();

        #[cfg(not(feature = "parallel"))]
        let results: Vec<(G, f32, Vec<f32>, Vec<f32>)> = parents
            .into_iter()
            .zip(selections.into_iter())
            .map(|(mut dna, (_, seed))| {
                let mut rng = Pcg64::seed_from_u64(seed);
                dna.mutate(&mut rng, mutation_rate);
                let (f, obj, desc) = evaluator.evaluate(&dna);
                (dna, f, obj, desc)
            })
            .collect();

        for (dna, f, obj, desc) in results {
            if f.is_nan() || desc.iter().any(|v| v.is_nan()) || desc.len() != dim {
                continue;
            }

            let key = vec![nearest_centroid_index(&desc, &self.centroids)];
            self.archive.insert_if_better(
                key,
                Phenotype {
                    genotype: dna,
                    fitness: f,
                    objectives: obj,
                    descriptor: desc,
                },
            );
        }
    }

    fn population(&mut self) -> &[Phenotype<G>] {
        self.archive.population()
    }
}

/// Squared Euclidean distance.
#[inline]
fn squared_distance(a: &[f32], b: &[f32]) -> f32 {
    a.iter().zip(b.iter()).map(|(x, y)| (x - y).powi(2)).sum()
}

/// Index of the centroid nearest `point`. Ties broken by lowest index.
#[inline]
fn nearest_centroid_index(point: &[f32], centroids: &[Vec<f32>]) -> usize {
    let mut best_idx = 0;
    let mut best_dist = squared_distance(point, &centroids[0]);
    for (i, c) in centroids.iter().enumerate().skip(1) {
        let d = squared_distance(point, c);
        if d < best_dist {
            best_dist = d;
            best_idx = i;
        }
    }
    best_idx
}

/// Lloyd's algorithm for Centroidal Voronoi Tessellation in `[0, 1]^dim`.
///
/// Each iteration draws `num_samples` uniform points, assigns them to the
/// nearest centroid, and moves each centroid to the mean of its assigned
/// samples. Centroids that receive no samples in an iteration are left in
/// place.
fn lloyd_centroids<R: Rng>(
    num_centroids: usize,
    dim: usize,
    num_samples: usize,
    iterations: usize,
    rng: &mut R,
) -> Vec<Vec<f32>> {
    // Initial centroids: uniform random in [0, 1]^dim.
    let mut centroids: Vec<Vec<f32>> = (0..num_centroids)
        .map(|_| (0..dim).map(|_| rng.random::<f32>()).collect())
        .collect();

    let mut sums: Vec<Vec<f32>> = vec![vec![0.0; dim]; num_centroids];
    let mut counts: Vec<usize> = vec![0; num_centroids];

    for _ in 0..iterations {
        for s in &mut sums {
            for v in s.iter_mut() {
                *v = 0.0;
            }
        }
        counts.iter_mut().for_each(|c| *c = 0);

        for _ in 0..num_samples {
            let sample: Vec<f32> = (0..dim).map(|_| rng.random::<f32>()).collect();
            let nearest = nearest_centroid_index(&sample, &centroids);
            for (acc, v) in sums[nearest].iter_mut().zip(sample.iter()) {
                *acc += v;
            }
            counts[nearest] += 1;
        }

        for i in 0..num_centroids {
            if counts[i] > 0 {
                let n = counts[i] as f32;
                for d in 0..dim {
                    centroids[i][d] = sums[i][d] / n;
                }
            }
        }
    }

    centroids
}