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}