rlevo_evolution/
fitness.rs

1//! Fitness evaluation traits and adapters.
2//!
3//! Two traits model the two evaluation shapes strategies expect:
4//!
5//! - [`FitnessFn`] — evaluates a single member. Callers hand the fitness
6//!   function a host-side genome row (typically `Vec<f32>`) and receive a
7//!   scalar. Useful for simple benchmarks and for unit-testing operators.
8//! - [`BatchFitnessFn`] — evaluates an entire population in one call and
9//!   returns a device-resident `Tensor<B, 1>` of shape `(pop_size,)`. This
10//!   is the hot path — strategies call it once per generation.
11//!
12//! Two adapters bridge host-side scalar fitness code into [`BatchFitnessFn`]:
13//!
14//! - [`FromFitnessEvaluable`] — wraps
15//!   `rlevo-core::fitness::FitnessEvaluable<Individual = Vec<f64>, Landscape = L>`.
16//!   Use this when an evaluator and a landscape type are already defined
17//!   separately (e.g. `RastriginEvaluator` + `RastriginLandscape`).
18//! - [`FromLandscape`] — wraps `rlevo-core::fitness::Landscape` directly.
19//!   Use this when the landscape is self-evaluating (Sphere, Ackley, Rastrigin)
20//!   and a separate evaluator shim would add no value.
21//!
22//! Both adapters pull each population row to host as `f32`, widen to `f64`,
23//! evaluate on the CPU, and re-upload the results as a `Tensor<B, 1>`.
24//! Purpose-built batched-on-device landscapes should implement
25//! [`BatchFitnessFn`] directly to avoid that round-trip.
26
27use burn::tensor::{Tensor, TensorData, backend::Backend};
28
29use rlevo_core::fitness::{FitnessEvaluable, Landscape};
30
31/// Single-member fitness evaluation.
32///
33/// Implementors may hold mutable state (e.g. a counter for number of
34/// evaluations) and are therefore `&mut self`.
35pub trait FitnessFn<G>: Send {
36    /// Evaluates one genome and returns its scalar fitness.
37    fn evaluate_one(&mut self, member: &G) -> f32;
38}
39
40/// Batched fitness evaluation over a population genome container `G`.
41///
42/// The returned tensor has shape `(pop_size,)` on the supplied device.
43/// Implementors must preserve row order — `fitness[i]` refers to the
44/// individual at row `i` of `population`.
45pub trait BatchFitnessFn<B: Backend, G>: Send {
46    /// Evaluates every member of `population` and returns a fitness tensor.
47    ///
48    /// The returned `Tensor<B, 1>` has shape `(pop_size,)` and is placed on
49    /// `device`. Row order is preserved: `fitness[i]` corresponds to the
50    /// individual at row `i` of `population`.
51    fn evaluate_batch(&mut self, population: &G, device: &<B as burn::tensor::backend::BackendTypes>::Device) -> Tensor<B, 1>;
52}
53
54/// Adapter from `FitnessEvaluable` to [`BatchFitnessFn<B, Tensor<B, 2>>`].
55///
56/// Each row of the population is pulled to host, converted to `Vec<f64>`,
57/// and passed to the underlying evaluator with the configured landscape.
58/// Fitness is computed on the host and then re-uploaded as a single
59/// `Tensor<B, 1>`.
60///
61/// # Precision
62///
63/// Populations are read as `f32` and widened to `f64` for the evaluator
64/// call; the returned `f64` fitness is narrowed back to `f32` before it
65/// is uploaded as a `Tensor<B, 1>`. Fitness values that exceed `f32`
66/// range (or rely on sub-ulp precision) will lose information at the
67/// narrowing step. Purpose-built batched-on-device landscapes should
68/// implement [`BatchFitnessFn`] directly to avoid the round-trip.
69///
70/// # Type Parameters
71///
72/// - `FE`: Concrete [`FitnessEvaluable`] implementation.
73/// - `L`: Landscape type; must match `FE::Landscape`.
74///
75/// # Panics
76///
77/// `evaluate_batch` panics if the supplied population tensor is not rank
78/// 2, or if its data cannot be read as `f32` (e.g. an integer backend).
79#[derive(Debug)]
80pub struct FromFitnessEvaluable<FE, L> {
81    evaluator: FE,
82    landscape: L,
83}
84
85impl<FE, L> FromFitnessEvaluable<FE, L> {
86    /// Builds the adapter from an evaluator and a landscape.
87    pub fn new(evaluator: FE, landscape: L) -> Self {
88        Self {
89            evaluator,
90            landscape,
91        }
92    }
93
94    /// Returns a reference to the wrapped landscape.
95    pub fn landscape(&self) -> &L {
96        &self.landscape
97    }
98}
99
100impl<FE, L, B> BatchFitnessFn<B, Tensor<B, 2>> for FromFitnessEvaluable<FE, L>
101where
102    B: Backend,
103    FE: FitnessEvaluable<Individual = Vec<f64>, Landscape = L> + Send,
104    L: Send + Sync,
105{
106    fn evaluate_batch(&mut self, population: &Tensor<B, 2>, device: &<B as burn::tensor::backend::BackendTypes>::Device) -> Tensor<B, 1> {
107        let dims = population.dims();
108        assert_eq!(dims.len(), 2, "population tensor must be rank 2");
109        let pop_size = dims[0];
110        let genome_dim = dims[1];
111
112        let flat = population
113            .clone()
114            .into_data()
115            .into_vec::<f32>()
116            .expect("tensor data must be readable as f32");
117        debug_assert_eq!(flat.len(), pop_size * genome_dim);
118
119        let mut fitness = Vec::with_capacity(pop_size);
120        let mut individual = Vec::with_capacity(genome_dim);
121        for row in 0..pop_size {
122            individual.clear();
123            let start = row * genome_dim;
124            individual.extend(
125                flat[start..start + genome_dim]
126                    .iter()
127                    .map(|&v| f64::from(v)),
128            );
129            let f = self.evaluator.evaluate(&individual, &self.landscape);
130            #[allow(clippy::cast_possible_truncation)]
131            fitness.push(f as f32);
132        }
133
134        let data = TensorData::new(fitness, [pop_size]);
135        Tensor::<B, 1>::from_data(data, device)
136    }
137}
138
139/// Adapter from [`Landscape`] to [`BatchFitnessFn<B, Tensor<B, 2>>`].
140///
141/// Use this when the landscape carries its own `evaluate(&[f64]) -> f64`
142/// (Sphere, Ackley, Rastrigin) so the example does not need a separate
143/// `FitnessEvaluable` shim. Each row is pulled to host as `f32`, widened
144/// to `f64`, evaluated, and re-uploaded as a `Tensor<B, 1>` — same
145/// precision caveats as [`FromFitnessEvaluable`] apply.
146///
147/// # Panics
148///
149/// `evaluate_batch` panics if the supplied population tensor is not rank
150/// 2, or if its data cannot be read as `f32` (e.g. an integer backend).
151#[derive(Debug)]
152pub struct FromLandscape<L> {
153    landscape: L,
154}
155
156impl<L> FromLandscape<L> {
157    /// Builds the adapter from a self-evaluating landscape.
158    pub fn new(landscape: L) -> Self {
159        Self { landscape }
160    }
161
162    /// Returns a reference to the wrapped landscape.
163    pub fn landscape(&self) -> &L {
164        &self.landscape
165    }
166}
167
168impl<L, B> BatchFitnessFn<B, Tensor<B, 2>> for FromLandscape<L>
169where
170    B: Backend,
171    L: Landscape,
172{
173    fn evaluate_batch(&mut self, population: &Tensor<B, 2>, device: &<B as burn::tensor::backend::BackendTypes>::Device) -> Tensor<B, 1> {
174        let dims = population.dims();
175        assert_eq!(dims.len(), 2, "population tensor must be rank 2");
176        let pop_size = dims[0];
177        let genome_dim = dims[1];
178
179        let flat = population
180            .clone()
181            .into_data()
182            .into_vec::<f32>()
183            .expect("tensor data must be readable as f32");
184        debug_assert_eq!(flat.len(), pop_size * genome_dim);
185
186        let mut fitness = Vec::with_capacity(pop_size);
187        let mut individual = Vec::with_capacity(genome_dim);
188        for row in 0..pop_size {
189            individual.clear();
190            let start = row * genome_dim;
191            individual.extend(
192                flat[start..start + genome_dim]
193                    .iter()
194                    .map(|&v| f64::from(v)),
195            );
196            let f = self.landscape.evaluate(&individual);
197            #[allow(clippy::cast_possible_truncation)]
198            fitness.push(f as f32);
199        }
200
201        let data = TensorData::new(fitness, [pop_size]);
202        Tensor::<B, 1>::from_data(data, device)
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209    use burn::backend::Flex;
210    type TestBackend = Flex;
211
212    #[derive(Debug, Clone, Copy)]
213    struct Sphere;
214
215    struct SphereFit;
216    impl FitnessEvaluable for SphereFit {
217        type Individual = Vec<f64>;
218        type Landscape = Sphere;
219        fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
220            x.iter().map(|v| v * v).sum()
221        }
222    }
223
224    #[test]
225    fn from_fitness_evaluable_preserves_row_order() {
226        let device = Default::default();
227        let data = TensorData::new(
228            vec![1.0_f32, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0],
229            [3, 3],
230        );
231        let pop = Tensor::<TestBackend, 2>::from_data(data, &device);
232
233        let mut adapter = FromFitnessEvaluable::new(SphereFit, Sphere);
234        let fitness = adapter.evaluate_batch(&pop, &device);
235
236        let values = fitness.into_data().into_vec::<f32>().unwrap();
237        assert_eq!(values.len(), 3);
238        approx::assert_relative_eq!(values[0], 1.0, epsilon = 1e-6);
239        approx::assert_relative_eq!(values[1], 4.0, epsilon = 1e-6);
240        approx::assert_relative_eq!(values[2], 9.0, epsilon = 1e-6);
241    }
242
243    #[test]
244    fn from_landscape_preserves_row_order() {
245        struct SphereLandscape;
246        impl Landscape for SphereLandscape {
247            fn evaluate(&self, x: &[f64]) -> f64 {
248                x.iter().map(|v| v * v).sum()
249            }
250        }
251
252        let device = Default::default();
253        let data = TensorData::new(
254            vec![1.0_f32, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0],
255            [3, 3],
256        );
257        let pop = Tensor::<TestBackend, 2>::from_data(data, &device);
258
259        let mut adapter = FromLandscape::new(SphereLandscape);
260        let fitness = adapter.evaluate_batch(&pop, &device);
261
262        let values = fitness.into_data().into_vec::<f32>().unwrap();
263        assert_eq!(values.len(), 3);
264        approx::assert_relative_eq!(values[0], 1.0, epsilon = 1e-6);
265        approx::assert_relative_eq!(values[1], 4.0, epsilon = 1e-6);
266        approx::assert_relative_eq!(values[2], 9.0, epsilon = 1e-6);
267    }
268}
rlevo_evolution/fitness.rs

rlevo_evolution/
fitness.rs
