Skip to main content

rlevo_evolution/
population.rs

1//! Population containers.
2//!
3//! [`Population<B, K>`] is a thin wrapper around a device tensor plus the
4//! shape metadata strategies need. For real-valued kinds it holds a
5//! `Tensor<B, 2>`; binary and integer kinds use `Tensor<B, 2, Int>`.
6//!
7//! The wrapper exists so operators and strategies have a single shape
8//! contract to validate against (they check `pop_size` and `genome_dim`
9//! rather than repeatedly interrogating `tensor.dims()`).
10//!
11//! # Constructing a population
12//!
13//! Each genome kind has a dedicated constructor that takes the
14//! already-allocated tensor:
15//!
16//! ```no_run
17//! use burn::backend::Flex;
18//! use burn::tensor::{Tensor, TensorData};
19//! use rlevo_evolution::genome::Real;
20//! use rlevo_evolution::population::Population;
21//!
22//! let device = Default::default();
23//! // 4 individuals, each with a 3-gene real-valued genome.
24//! let data = TensorData::new(vec![0.1f32, 0.2, 0.3,
25//!                                 0.4, 0.5, 0.6,
26//!                                 0.7, 0.8, 0.9,
27//!                                 1.0, 1.1, 1.2], [4, 3]);
28//! let tensor = Tensor::<Flex, 2>::from_data(data, &device);
29//! // Construction validates that the tensor is non-empty, so it returns a
30//! // `Result`; a 0×n or m×0 tensor is rejected as a `ConfigError`.
31//! let pop = Population::<Flex, Real>::new_real(tensor).unwrap();
32//! assert_eq!(pop.pop_size(), 4);
33//! assert_eq!(pop.genome_dim(), 3);
34//! ```
35
36use burn::tensor::{Int, Tensor, backend::Backend};
37
38use rlevo_core::config::{self, ConfigError};
39
40use crate::genome::{Binary, Integer, Permutation, Real, TensorGenome};
41
42/// Population stored on a Burn backend device.
43///
44/// The concrete tensor type depends on the genome kind `K`, chosen at compile
45/// time through [`TensorGenome::Tensor`]: `Real` is backed by `Tensor<B, 2>`,
46/// `Binary` and `Integer` by `Tensor<B, 2, Int>`. Because the storage type is a
47/// function of `K`, there is a single tensor field and no run-time tag — a
48/// population can never hold the wrong tensor flavour for its kind, so the
49/// [`tensor`](Population::tensor) accessor is total (it cannot fail).
50///
51/// The `K: TensorGenome` bound is what keeps this honest: kinds without a
52/// rectangular tensor form (e.g. [`Tree`](crate::genome::Tree)) do not implement
53/// `TensorGenome`, so `Population<B, Tree>` does not type-check.
54#[derive(Debug, Clone)]
55pub struct Population<B: Backend, K: TensorGenome> {
56    pop_size: usize,
57    genome_dim: usize,
58    tensor: K::Tensor<B>,
59}
60
61impl<B: Backend, K: TensorGenome> Population<B, K> {
62    /// Returns the number of individuals (rows) in the population.
63    ///
64    /// This value equals `tensor.dims()[0]`.
65    #[must_use]
66    pub fn pop_size(&self) -> usize {
67        self.pop_size
68    }
69
70    /// Returns the genome dimensionality (number of genes, i.e. columns).
71    ///
72    /// This value equals `tensor.dims()[1]`.
73    #[must_use]
74    pub fn genome_dim(&self) -> usize {
75        self.genome_dim
76    }
77
78    /// Borrows the backing tensor for this population's kind.
79    ///
80    /// The concrete type is [`K::Tensor<B>`](TensorGenome::Tensor) — a
81    /// `Tensor<B, 2>` for `Real`, a `Tensor<B, 2, Int>` for `Binary`/`Integer`
82    /// — with shape `[pop_size, genome_dim]`. Use it to pass the population to
83    /// fitness functions or operator kernels without giving up ownership.
84    #[must_use]
85    pub fn tensor(&self) -> &K::Tensor<B> {
86        &self.tensor
87    }
88
89    /// Consumes the wrapper and returns the owned tensor.
90    ///
91    /// Prefer this over [`tensor`](Population::tensor) when handing the
92    /// population off to a strategy or operator that needs ownership (e.g. to
93    /// avoid a clone on the hot path).
94    #[must_use]
95    pub fn into_tensor(self) -> K::Tensor<B> {
96        self.tensor
97    }
98}
99
100impl<B: Backend> Population<B, Real> {
101    /// Constructs a real-valued population from a `Tensor<B, 2>`.
102    ///
103    /// Shape is read from `tensor.dims()` at construction time; subsequent
104    /// calls to [`pop_size`](Population::pop_size) and
105    /// [`genome_dim`](Population::genome_dim) reflect those dimensions.
106    ///
107    /// # Errors
108    ///
109    /// Returns [`ConstraintKind::Zero`](rlevo_core::config::ConstraintKind::Zero)
110    /// (as `field` `"pop_size"` or `"genome_dim"`) if the tensor has zero rows
111    /// or zero columns. Rejecting the empty case here names `Population` as the
112    /// source instead of surfacing later as an opaque operator panic.
113    ///
114    /// # Examples
115    ///
116    /// ```no_run
117    /// use burn::backend::Flex;
118    /// use burn::tensor::{Tensor, TensorData};
119    /// use rlevo_evolution::genome::Real;
120    /// use rlevo_evolution::population::Population;
121    ///
122    /// let device = Default::default();
123    /// let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0], [2, 2]);
124    /// let pop = Population::<Flex, Real>::new_real(
125    ///     Tensor::from_data(data, &device),
126    /// ).unwrap();
127    /// assert_eq!(pop.pop_size(), 2);
128    /// assert_eq!(pop.genome_dim(), 2);
129    /// ```
130    pub fn new_real(tensor: Tensor<B, 2>) -> Result<Self, ConfigError> {
131        let dims = tensor.dims();
132        config::nonzero("Population", "pop_size", dims[0])?;
133        config::nonzero("Population", "genome_dim", dims[1])?;
134        Ok(Self {
135            pop_size: dims[0],
136            genome_dim: dims[1],
137            tensor,
138        })
139    }
140}
141
142impl<B: Backend> Population<B, Binary> {
143    /// Constructs a binary population from a `Tensor<B, 2, Int>`.
144    ///
145    /// Each element is expected to be `0` or `1`; the constructor does not
146    /// validate element values. Shape is read from `tensor.dims()`.
147    ///
148    /// # Errors
149    ///
150    /// Returns [`ConstraintKind::Zero`](rlevo_core::config::ConstraintKind::Zero)
151    /// (as `field` `"pop_size"` or `"genome_dim"`) if the tensor has zero rows
152    /// or zero columns.
153    ///
154    /// # Examples
155    ///
156    /// ```no_run
157    /// use burn::backend::Flex;
158    /// use burn::tensor::{Int, Tensor, TensorData};
159    /// use rlevo_evolution::genome::Binary;
160    /// use rlevo_evolution::population::Population;
161    ///
162    /// let device = Default::default();
163    /// // 3 individuals, each with a 4-bit binary genome.
164    /// let data = TensorData::new(vec![0i64, 1, 0, 1,
165    ///                                 1, 0, 1, 0,
166    ///                                 0, 0, 1, 1], [3, 4]);
167    /// let pop = Population::<Flex, Binary>::new_binary(
168    ///     Tensor::from_data(data, &device),
169    /// ).unwrap();
170    /// assert_eq!(pop.pop_size(), 3);
171    /// assert_eq!(pop.genome_dim(), 4);
172    /// ```
173    pub fn new_binary(tensor: Tensor<B, 2, Int>) -> Result<Self, ConfigError> {
174        let dims = tensor.dims();
175        config::nonzero("Population", "pop_size", dims[0])?;
176        config::nonzero("Population", "genome_dim", dims[1])?;
177        Ok(Self {
178            pop_size: dims[0],
179            genome_dim: dims[1],
180            tensor,
181        })
182    }
183}
184
185impl<B: Backend> Population<B, Integer> {
186    /// Constructs an integer population from a `Tensor<B, 2, Int>`.
187    ///
188    /// Elements represent non-negative integer indices (e.g. node indices in
189    /// CGP, symbol indices in integer-coded GA). The constructor does not
190    /// validate element bounds. Shape is read from `tensor.dims()`.
191    ///
192    /// # Errors
193    ///
194    /// Returns [`ConstraintKind::Zero`](rlevo_core::config::ConstraintKind::Zero)
195    /// (as `field` `"pop_size"` or `"genome_dim"`) if the tensor has zero rows
196    /// or zero columns.
197    ///
198    /// # Examples
199    ///
200    /// ```no_run
201    /// use burn::backend::Flex;
202    /// use burn::tensor::{Int, Tensor, TensorData};
203    /// use rlevo_evolution::genome::Integer;
204    /// use rlevo_evolution::population::Population;
205    ///
206    /// let device = Default::default();
207    /// // 2 individuals, each with a 5-gene integer-valued genome.
208    /// let data = TensorData::new(vec![0i64, 3, 1, 4, 2,
209    ///                                 2, 0, 4, 1, 3], [2, 5]);
210    /// let pop = Population::<Flex, Integer>::new_integer(
211    ///     Tensor::from_data(data, &device),
212    /// ).unwrap();
213    /// assert_eq!(pop.pop_size(), 2);
214    /// assert_eq!(pop.genome_dim(), 5);
215    /// ```
216    pub fn new_integer(tensor: Tensor<B, 2, Int>) -> Result<Self, ConfigError> {
217        let dims = tensor.dims();
218        config::nonzero("Population", "pop_size", dims[0])?;
219        config::nonzero("Population", "genome_dim", dims[1])?;
220        Ok(Self {
221            pop_size: dims[0],
222            genome_dim: dims[1],
223            tensor,
224        })
225    }
226}
227
228impl<B: Backend> Population<B, Permutation> {
229    /// Constructs a permutation population from a `Tensor<B, 2, Int>`.
230    ///
231    /// Each row is *assumed* to be a permutation of `0..genome_dim`, but the
232    /// constructor validates only shape — the per-row bijection invariant is
233    /// **not** checked, mirroring how [`new_binary`](Population::new_binary) and
234    /// [`new_integer`](Population::new_integer) leave element values unchecked.
235    /// Shape is read from `tensor.dims()`.
236    ///
237    /// The permutation operators (Ant Colony Optimization over TSP/QAP) are
238    /// planned for a future release; this constructor exists so downstream code
239    /// can allocate and reference `Population<B, Permutation>` today.
240    ///
241    /// # Errors
242    ///
243    /// Returns [`ConstraintKind::Zero`](rlevo_core::config::ConstraintKind::Zero)
244    /// (as `field` `"pop_size"` or `"genome_dim"`) if the tensor has zero rows
245    /// or zero columns.
246    ///
247    /// # Examples
248    ///
249    /// ```no_run
250    /// use burn::backend::Flex;
251    /// use burn::tensor::{Int, Tensor, TensorData};
252    /// use rlevo_evolution::genome::Permutation;
253    /// use rlevo_evolution::population::Population;
254    ///
255    /// let device = Default::default();
256    /// // 2 ants, each a permutation of a 4-node tour.
257    /// let data = TensorData::new(vec![0i64, 1, 2, 3,
258    ///                                 2, 0, 3, 1], [2, 4]);
259    /// let pop = Population::<Flex, Permutation>::new_permutation(
260    ///     Tensor::from_data(data, &device),
261    /// ).unwrap();
262    /// assert_eq!(pop.pop_size(), 2);
263    /// assert_eq!(pop.genome_dim(), 4);
264    /// ```
265    pub fn new_permutation(tensor: Tensor<B, 2, Int>) -> Result<Self, ConfigError> {
266        let dims = tensor.dims();
267        config::nonzero("Population", "pop_size", dims[0])?;
268        config::nonzero("Population", "genome_dim", dims[1])?;
269        Ok(Self {
270            pop_size: dims[0],
271            genome_dim: dims[1],
272            tensor,
273        })
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use burn::backend::Flex;
281    use burn::tensor::TensorData;
282    type TestBackend = Flex;
283
284    #[test]
285    fn real_population_reports_shape() {
286        let device = Default::default();
287        let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0], [2, 2]);
288        let tensor = Tensor::<TestBackend, 2>::from_data(data, &device);
289        let pop = Population::<TestBackend, Real>::new_real(tensor).unwrap();
290        assert_eq!(pop.pop_size(), 2);
291        assert_eq!(pop.genome_dim(), 2);
292        assert_eq!(pop.tensor().dims(), [2, 2]);
293    }
294
295    #[test]
296    fn binary_population_uses_int_tensor() {
297        let device = Default::default();
298        let data = TensorData::new(vec![0i64, 1, 1, 0, 1, 0], [2, 3]);
299        let tensor = Tensor::<TestBackend, 2, Int>::from_data(data, &device);
300        let pop = Population::<TestBackend, Binary>::new_binary(tensor).unwrap();
301        assert_eq!(pop.pop_size(), 2);
302        assert_eq!(pop.genome_dim(), 3);
303    }
304
305    #[test]
306    fn permutation_population_reports_shape() {
307        let device = Default::default();
308        let data = TensorData::new(vec![0i64, 1, 2, 3, 2, 0, 3, 1], [2, 4]);
309        let tensor = Tensor::<TestBackend, 2, Int>::from_data(data, &device);
310        let pop = Population::<TestBackend, Permutation>::new_permutation(tensor).unwrap();
311        assert_eq!(pop.pop_size(), 2);
312        assert_eq!(pop.genome_dim(), 4);
313    }
314
315    #[test]
316    fn new_real_rejects_zero_rows() {
317        let device = Default::default();
318        let data = TensorData::new(Vec::<f32>::new(), [0, 3]);
319        let tensor = Tensor::<TestBackend, 2>::from_data(data, &device);
320        let err = Population::<TestBackend, Real>::new_real(tensor).unwrap_err();
321        assert_eq!(err.field, "pop_size");
322    }
323
324    #[test]
325    fn new_real_rejects_zero_width() {
326        let device = Default::default();
327        let data = TensorData::new(Vec::<f32>::new(), [3, 0]);
328        let tensor = Tensor::<TestBackend, 2>::from_data(data, &device);
329        let err = Population::<TestBackend, Real>::new_real(tensor).unwrap_err();
330        assert_eq!(err.field, "genome_dim");
331    }
332
333    #[test]
334    fn population_is_send_sync() {
335        fn assert_send_sync<T: Send + Sync>() {}
336        assert_send_sync::<Population<TestBackend, Real>>();
337        assert_send_sync::<Population<TestBackend, Permutation>>();
338    }
339}