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//! let pop = Population::<Flex, Real>::new_real(tensor);
30//! assert_eq!(pop.pop_size(), 4);
31//! assert_eq!(pop.genome_dim(), 3);
32//! ```
33
34use std::marker::PhantomData;
35
36use burn::tensor::{backend::Backend, Int, Tensor};
37
38use crate::genome::{Binary, Integer, Real};
39
40/// Population stored on a Burn backend device.
41///
42/// The concrete tensor type depends on the genome kind `K`. Most
43/// consumers interact with [`Population<B, Real>`] via [`tensor`](Population::tensor),
44/// but strategies parameterized on the kind can keep the `K` generic and
45/// reach for the right tensor flavor through the inherent impls below.
46///
47/// Invariant: for every `Population<B, K>` produced by the public
48/// constructors, exactly one of `tensor_real` / `tensor_int` is `Some`,
49/// determined by `K`. `Real` populates `tensor_real`; `Binary`,
50/// `Integer`, and `Permutation` populate `tensor_int`. The inherent
51/// `tensor(&self)` accessors `.expect()` on the matching field because
52/// the constructor contract pins the invariant — a mismatch would be a
53/// bug in this module.
54#[derive(Debug, Clone)]
55pub struct Population<B: Backend, K> {
56 pop_size: usize,
57 genome_dim: usize,
58 _kind: PhantomData<K>,
59 tensor_real: Option<Tensor<B, 2>>,
60 tensor_int: Option<Tensor<B, 2, Int>>,
61}
62
63impl<B: Backend, K> Population<B, K> {
64 /// Returns the number of individuals (rows) in the population.
65 ///
66 /// This value equals `tensor.dims()[0]` for any population produced by
67 /// the public constructors.
68 #[must_use]
69 pub fn pop_size(&self) -> usize {
70 self.pop_size
71 }
72
73 /// Returns the genome dimensionality (number of genes, i.e. columns).
74 ///
75 /// This value equals `tensor.dims()[1]` for any population produced by
76 /// the public constructors.
77 #[must_use]
78 pub fn genome_dim(&self) -> usize {
79 self.genome_dim
80 }
81}
82
83impl<B: Backend> Population<B, Real> {
84 /// Constructs a real-valued population from a `Tensor<B, 2>`.
85 ///
86 /// Shape is read from `tensor.dims()` at construction time; subsequent
87 /// calls to [`pop_size`](Population::pop_size) and
88 /// [`genome_dim`](Population::genome_dim) reflect those dimensions.
89 ///
90 /// # Examples
91 ///
92 /// ```no_run
93 /// use burn::backend::Flex;
94 /// use burn::tensor::{Tensor, TensorData};
95 /// use rlevo_evolution::genome::Real;
96 /// use rlevo_evolution::population::Population;
97 ///
98 /// let device = Default::default();
99 /// let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0], [2, 2]);
100 /// let pop = Population::<Flex, Real>::new_real(
101 /// Tensor::from_data(data, &device),
102 /// );
103 /// assert_eq!(pop.pop_size(), 2);
104 /// assert_eq!(pop.genome_dim(), 2);
105 /// ```
106 ///
107 /// # Panics
108 ///
109 /// Panics if `tensor` is not rank 2.
110 #[must_use]
111 pub fn new_real(tensor: Tensor<B, 2>) -> Self {
112 let dims = tensor.dims();
113 assert_eq!(dims.len(), 2, "population tensor must be rank 2");
114 Self {
115 pop_size: dims[0],
116 genome_dim: dims[1],
117 _kind: PhantomData,
118 tensor_real: Some(tensor),
119 tensor_int: None,
120 }
121 }
122
123 /// Borrows the backing real-valued tensor.
124 ///
125 /// The returned tensor has shape `[pop_size, genome_dim]`. Use this
126 /// to pass the population to fitness functions or operator kernels
127 /// without giving up ownership.
128 ///
129 /// # Panics
130 ///
131 /// Never panics in practice: a real-valued population always holds a
132 /// real tensor by construction.
133 #[must_use]
134 pub fn tensor(&self) -> &Tensor<B, 2> {
135 self.tensor_real
136 .as_ref()
137 .expect("real population always has a tensor_real")
138 }
139
140 /// Consumes the wrapper and returns the owned tensor.
141 ///
142 /// Prefer this over [`tensor`](Population::tensor) when handing the
143 /// population off to a strategy or operator that needs ownership (e.g.
144 /// to avoid a clone on the hot path).
145 ///
146 /// # Panics
147 ///
148 /// Never panics in practice: a real-valued population always holds a
149 /// real tensor by construction.
150 #[must_use]
151 pub fn into_tensor(self) -> Tensor<B, 2> {
152 self.tensor_real
153 .expect("real population always has a tensor_real")
154 }
155}
156
157impl<B: Backend> Population<B, Binary> {
158 /// Constructs a binary population from a `Tensor<B, 2, Int>`.
159 ///
160 /// Each element is expected to be `0` or `1`; the constructor does not
161 /// validate element values. Shape is read from `tensor.dims()`.
162 ///
163 /// # Examples
164 ///
165 /// ```no_run
166 /// use burn::backend::Flex;
167 /// use burn::tensor::{Int, Tensor, TensorData};
168 /// use rlevo_evolution::genome::Binary;
169 /// use rlevo_evolution::population::Population;
170 ///
171 /// let device = Default::default();
172 /// // 3 individuals, each with a 4-bit binary genome.
173 /// let data = TensorData::new(vec![0i64, 1, 0, 1,
174 /// 1, 0, 1, 0,
175 /// 0, 0, 1, 1], [3, 4]);
176 /// let pop = Population::<Flex, Binary>::new_binary(
177 /// Tensor::from_data(data, &device),
178 /// );
179 /// assert_eq!(pop.pop_size(), 3);
180 /// assert_eq!(pop.genome_dim(), 4);
181 /// ```
182 ///
183 /// # Panics
184 ///
185 /// Panics if `tensor` is not rank 2.
186 #[must_use]
187 pub fn new_binary(tensor: Tensor<B, 2, Int>) -> Self {
188 let dims = tensor.dims();
189 assert_eq!(dims.len(), 2, "population tensor must be rank 2");
190 Self {
191 pop_size: dims[0],
192 genome_dim: dims[1],
193 _kind: PhantomData,
194 tensor_real: None,
195 tensor_int: Some(tensor),
196 }
197 }
198
199 /// Borrows the backing integer tensor holding 0/1 values.
200 ///
201 /// The returned tensor has shape `[pop_size, genome_dim]` and element
202 /// type `Int`. Callers performing crossover or mutation should work
203 /// directly with this tensor.
204 ///
205 /// # Panics
206 ///
207 /// Never panics in practice: a binary population always holds an integer
208 /// tensor by construction.
209 #[must_use]
210 pub fn tensor(&self) -> &Tensor<B, 2, Int> {
211 self.tensor_int
212 .as_ref()
213 .expect("binary population always has a tensor_int")
214 }
215}
216
217impl<B: Backend> Population<B, Integer> {
218 /// Constructs an integer population from a `Tensor<B, 2, Int>`.
219 ///
220 /// Elements represent non-negative integer indices (e.g. node indices in
221 /// CGP, symbol indices in integer-coded GA). The constructor does not
222 /// validate element bounds. Shape is read from `tensor.dims()`.
223 ///
224 /// # Examples
225 ///
226 /// ```no_run
227 /// use burn::backend::Flex;
228 /// use burn::tensor::{Int, Tensor, TensorData};
229 /// use rlevo_evolution::genome::Integer;
230 /// use rlevo_evolution::population::Population;
231 ///
232 /// let device = Default::default();
233 /// // 2 individuals, each with a 5-gene integer-valued genome.
234 /// let data = TensorData::new(vec![0i64, 3, 1, 4, 2,
235 /// 2, 0, 4, 1, 3], [2, 5]);
236 /// let pop = Population::<Flex, Integer>::new_integer(
237 /// Tensor::from_data(data, &device),
238 /// );
239 /// assert_eq!(pop.pop_size(), 2);
240 /// assert_eq!(pop.genome_dim(), 5);
241 /// ```
242 ///
243 /// # Panics
244 ///
245 /// Panics if `tensor` is not rank 2.
246 #[must_use]
247 pub fn new_integer(tensor: Tensor<B, 2, Int>) -> Self {
248 let dims = tensor.dims();
249 assert_eq!(dims.len(), 2, "population tensor must be rank 2");
250 Self {
251 pop_size: dims[0],
252 genome_dim: dims[1],
253 _kind: PhantomData,
254 tensor_real: None,
255 tensor_int: Some(tensor),
256 }
257 }
258
259 /// Borrows the backing integer tensor.
260 ///
261 /// The returned tensor has shape `[pop_size, genome_dim]` and element
262 /// type `Int`. Element values are non-negative indices whose domain is
263 /// determined by the problem (e.g. `0..n_nodes` for CGP).
264 ///
265 /// # Panics
266 ///
267 /// Never panics in practice: an integer population always holds an integer
268 /// tensor by construction.
269 #[must_use]
270 pub fn tensor(&self) -> &Tensor<B, 2, Int> {
271 self.tensor_int
272 .as_ref()
273 .expect("integer population always has a tensor_int")
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);
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);
301 assert_eq!(pop.pop_size(), 2);
302 assert_eq!(pop.genome_dim(), 3);
303 }
304}