diff --git a/src/functions.rs b/src/functions.rs index 0bd5804..93307fb 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -8,8 +8,8 @@ use std::rc::Rc; use crate::display::BoxedTensor; use crate::domain::HRect; use crate::mapping::{ - COO, Concat, Diagonal, EmptyMapping, InvProjection, Permutation, Projection, Squeeze, - Unsqueeze, Where, + Affine, COO, Concat, Diagonal, EmptyMapping, InvAffine, InvProjection, Permutation, Projection, + Squeeze, Unsqueeze, Where, }; use crate::tensor::{ JaggedTensor, LeftComposition, Physical, Product, Reduction, ReductionTensor, RightComposition, @@ -269,6 +269,20 @@ pub fn sparse_coo_tensor( Sum::new(values, Rc::new(COO::new(HRect::new(size), coordinates))) } +/// A sparse latticed tensor of shape `shape`: the value at index `x` of +/// `physical` sits at `basis·x + offset`, and the tensor is zero off that +/// lattice. `basis` has a row per dimension of `shape` and a column per +/// dimension of `physical`. Values the lattice sends to one index add up. +pub fn sparse_latticed_tensor( + physical: Rc, + basis: Vec>, + offset: Vec, + shape: Vec, +) -> Sum { + let affine = Affine::new(physical.hrect().clone(), basis, offset, HRect::new(shape)); + Sum::new(physical, Rc::new(InvAffine::new(affine))) +} + pub fn sum(input: Rc, dim: usize) -> Sum { let codomain = input.hrect().clone(); Sum::new(input, Rc::new(InvProjection::new(codomain, dim))) diff --git a/src/lattice.rs b/src/lattice.rs new file mode 100644 index 0000000..1c7d4d3 --- /dev/null +++ b/src/lattice.rs @@ -0,0 +1,183 @@ +//! Inverting an affine map on a box without scanning the box: finding every +//! index `x` with `basis·x = c` and `0 <= x_j < extents[j]`. +//! +//! Integer column operations of determinant ±1 bring the basis to +//! `basis·[particular | kernel] = [h | 0]`, where `h` is a staircase of full +//! column rank. Every integer solution is then `x = particular·z + kernel·t`, +//! exactly: `z` is unique and found by substitution down the staircase of `h`, +//! and `t` ranges over the integer kernel. The kernel's basis is brought to a +//! staircase too, so that each coordinate of `t` in turn gets an exact range +//! from the extents. + +/// The indices of a box that `x ↦ basis·x` sends to a given point. +pub(crate) struct Preimage { + h: Vec>, + h_pivots: Vec, + particular: Vec>, + kernel: Vec>, + kernel_pivots: Vec, + extents: Vec, +} + +impl Preimage { + pub(crate) fn new(basis: &[Vec], extents: &[usize]) -> Self { + let n = extents.len(); + let (echelon, operations, h_pivots) = column_echelon(basis.to_vec(), n); + let r = h_pivots.len(); + let h = echelon.iter().map(|row| row[..r].to_vec()).collect(); + let particular = operations.iter().map(|row| row[..r].to_vec()).collect(); + let free = operations.iter().map(|row| row[r..].to_vec()).collect(); + let (kernel, _, kernel_pivots) = column_echelon(free, n - r); + Self { + h, + h_pivots, + particular, + kernel, + kernel_pivots, + extents: extents.iter().map(|&extent| extent as i64).collect(), + } + } + + /// Every index `x` of the box with `basis·x = c`. + pub(crate) fn solve(&self, c: &[i64]) -> Vec> { + let mut z = vec![0; self.h_pivots.len()]; + for (i, &row) in self.h_pivots.iter().enumerate() { + let rest = c[row] - dot(&self.h[row][..i], &z[..i]); + let step = self.h[row][i]; + if rest % step != 0 { + return Vec::new(); + } + z[i] = rest / step; + } + // The rows between pivots must hold too. + if self + .h + .iter() + .zip(c) + .any(|(row, &target)| dot(row, &z) != target) + { + return Vec::new(); + } + let mut x: Vec = self.particular.iter().map(|row| dot(row, &z)).collect(); + // Coordinates above the kernel's first pivot are already final. + let fixed = self.kernel_pivots.first().copied().unwrap_or(x.len()); + if !(0..fixed).all(|j| (0..self.extents[j]).contains(&x[j])) { + return Vec::new(); + } + let mut found = Vec::new(); + self.enumerate(0, &mut x, &mut found); + found + } + + /// Tries every value of `t[level]` that keeps the coordinates it settles in + /// range, these being the rows from its pivot to the next one. + fn enumerate(&self, level: usize, x: &mut [i64], found: &mut Vec>) { + if level == self.kernel_pivots.len() { + found.push(x.iter().map(|&coord| coord as usize).collect()); + return; + } + let end = self + .kernel_pivots + .get(level + 1) + .copied() + .unwrap_or(x.len()); + let rows = self.kernel_pivots[level]..end; + let settled = self.kernel[rows.clone()] + .iter() + .zip(&x[rows.clone()]) + .zip(&self.extents[rows]); + let (mut low, mut high) = (i64::MIN, i64::MAX); + for ((row, &coord), &extent) in settled { + let (step, last) = (row[level], extent - 1); + if step > 0 { + low = low.max(ceil_div(-coord, step)); + high = high.min(floor_div(last - coord, step)); + } else if step < 0 { + low = low.max(ceil_div(last - coord, step)); + high = high.min(floor_div(-coord, step)); + } else if !(0..=last).contains(&coord) { + return; + } + } + for t in low..=high { + for (coord, row) in x.iter_mut().zip(&self.kernel) { + *coord += row[level] * t; + } + self.enumerate(level + 1, x, found); + for (coord, row) in x.iter_mut().zip(&self.kernel) { + *coord -= row[level] * t; + } + } + } +} + +/// Brings `matrix` to column echelon form by integer column operations of +/// determinant ±1. Returns the echelon form `e`, the operations `v` with +/// `matrix·v = e`, and the pivot row of each column of `e` that is not zero. +/// Those columns come first: each is zero above its pivot, and every column is +/// zero to the right of a pivot. +fn column_echelon( + mut matrix: Vec>, + cols: usize, +) -> (Vec>, Vec>, Vec) { + let mut operations: Vec> = (0..cols) + .map(|i| (0..cols).map(|j| i64::from(i == j)).collect()) + .collect(); + let mut pivots = Vec::new(); + for row in 0..matrix.len() { + let c = pivots.len(); + if c == cols { + break; + } + // Folds every entry right of column `c` into it, leaving their gcd there. + for j in c + 1..cols { + let (a, b) = (matrix[row][c], matrix[row][j]); + if b != 0 { + let (g, s, t) = ext_gcd(a, b); + let operation = (s, t, -b / g, a / g); + column_op(&mut matrix, c, j, operation); + column_op(&mut operations, c, j, operation); + } + } + if matrix[row][c] != 0 { + pivots.push(row); + } + } + (matrix, operations, pivots) +} + +/// Turns columns `c` and `j` into `s·c + t·j` and `u·c + w·j`. +fn column_op(matrix: &mut [Vec], c: usize, j: usize, (s, t, u, w): (i64, i64, i64, i64)) { + for row in matrix { + let (x, y) = (row[c], row[j]); + row[c] = s * x + t * y; + row[j] = u * x + w * y; + } +} + +/// `(g, s, t)` with `s·a + t·b = g`, where `g` is the gcd of `a` and `b`, and +/// is not negative. +fn ext_gcd(a: i64, b: i64) -> (i64, i64, i64) { + if b == 0 { + if a < 0 { (-a, -1, 0) } else { (a, 1, 0) } + } else { + let (g, s, t) = ext_gcd(b, a % b); + (g, t, s - (a / b) * t) + } +} + +fn dot(a: &[i64], b: &[i64]) -> i64 { + a.iter().zip(b).map(|(x, y)| x * y).sum() +} + +fn floor_div(n: i64, d: i64) -> i64 { + if d < 0 { + (-n).div_euclid(-d) + } else { + n.div_euclid(d) + } +} + +fn ceil_div(n: i64, d: i64) -> i64 { + -floor_div(-n, d) +} diff --git a/src/lib.rs b/src/lib.rs index c13d168..d1e8b96 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,7 @@ pub mod display; pub mod domain; pub mod functions; +mod lattice; pub mod mapping; pub mod signature; pub mod tensor; @@ -12,13 +13,13 @@ pub use display::{BoxedTensor, format_boxed}; pub use domain::{Domain, HRect, JaggedDomain}; pub use functions::{ Dims, clamp, concatenate, diagonalize, eye, eye_rect, mm, movedim, negate, ones, permute, pow, - prod, scale, sparse_coo_tensor, square, squeeze, stack, sum, transpose, unsqueeze, where_, - zeros, + prod, scale, sparse_coo_tensor, sparse_latticed_tensor, square, squeeze, stack, sum, transpose, + unsqueeze, where_, zeros, }; pub use mapping::{ - BoxedMapping, BoxedPowersetMapping, COO, Concat, Diagonal, EmptyMapping, IndexMapping, - InvProjection, Permutation, Powerset, PowersetMapping, Projection, Squeeze, Stride, Unsqueeze, - Where, + Affine, BoxedMapping, BoxedPowersetMapping, COO, Concat, Diagonal, EmptyMapping, IndexMapping, + InvAffine, InvProjection, Permutation, Powerset, PowersetMapping, Projection, Squeeze, Stride, + Unsqueeze, Where, }; pub use signature::{HRectSignature, Signature}; pub use tensor::{ diff --git a/src/mapping.rs b/src/mapping.rs index e05a246..b2642d0 100644 --- a/src/mapping.rs +++ b/src/mapping.rs @@ -3,6 +3,7 @@ use std::rc::Rc; use crate::display::BoxedTensor; use crate::domain::{Domain, HRect, JaggedDomain}; +use crate::lattice::Preimage; use crate::signature::{HRectSignature, Signature}; /// A set of indices, as returned by a `PowersetMapping`. @@ -575,3 +576,130 @@ impl HRectPowersetMapping for InvProjection { &self.signature } } + +/// The affine map `x ↦ basis·x + offset`. `basis` has a row per dimension of the +/// codomain and a column per dimension of the domain. Its entries and `offset` +/// may be negative, as long as every index of the domain lands in the codomain. +/// Read through this, a tensor gives its values on a lattice, e.g. its diagonal +/// with `basis = [[1], [1]]`. +pub struct Affine { + signature: HRectSignature, + basis: Vec>, + offset: Vec, +} + +impl Affine { + pub fn new(domain: HRect, basis: Vec>, offset: Vec, codomain: HRect) -> Self { + assert_eq!( + basis.len(), + codomain.ndim(), + "a basis row per dimension of the codomain is required" + ); + assert!( + basis.iter().all(|row| row.len() == domain.ndim()), + "a basis column per dimension of the domain is required" + ); + assert_eq!( + offset.len(), + codomain.ndim(), + "an offset per dimension of the codomain is required" + ); + // An empty domain has no index to land anywhere. + if !domain.shape().contains(&0) { + for (i, (row, &start)) in basis.iter().zip(&offset).enumerate() { + // Each coordinate varies on its own, so each term is at its + // smallest and largest at one end or the other of its range. + let (lowest, highest) = row.iter().zip(domain.shape()).fold( + (start, start), + |(low, high), (&step, &extent)| { + let end = step * (extent as i64 - 1); + (low + end.min(0), high + end.max(0)) + }, + ); + assert!( + 0 <= lowest && highest < codomain.shape()[i] as i64, + "the lattice reaches {lowest} to {highest} along dimension {i}, which has extent {}", + codomain.shape()[i] + ); + } + } + Self { + signature: HRectSignature::new(domain, codomain), + basis, + offset, + } + } +} + +impl HRectMapping for Affine { + type Codomain = HRect; + + fn map(&self, index: &[usize]) -> Vec { + // Every coordinate is non-negative, as checked when constructing. + self.basis + .iter() + .zip(&self.offset) + .map(|(row, &start)| { + let coord = start + + row + .iter() + .zip(index) + .map(|(&step, &x)| step * x as i64) + .sum::(); + coord as usize + }) + .collect() + } + + fn signature(&self) -> &HRectSignature { + &self.signature + } +} + +/// The inverse of an `Affine`: sends each index of its codomain back to the +/// indices of its domain that land on it. There are none off the lattice, so a +/// sum over this lays its underlying tensor out on the lattice and is zero +/// elsewhere: a sparse latticed tensor. They are found by solving +/// `basis·x = index - offset`, so a read costs about as much as the basis is +/// large, however large the tensors. +pub struct InvAffine { + signature: HRectSignature, + offset: Vec, + preimage: Preimage, +} + +impl InvAffine { + pub fn new(affine: Affine) -> Self { + let Affine { + signature, + basis, + offset, + } = affine; + let preimage = Preimage::new(&basis, signature.domain().shape()); + Self { + signature: HRectSignature::new( + signature.codomain().clone(), + signature.domain().clone(), + ), + offset, + preimage, + } + } +} + +impl HRectPowersetMapping for InvAffine { + type Codomain = HRect; + + fn map(&self, index: &[usize]) -> Powerset { + let target: Vec = index + .iter() + .zip(&self.offset) + .map(|(&coord, &start)| coord as i64 - start) + .collect(); + self.preimage.solve(&target).into_iter().collect() + } + + fn signature(&self) -> &HRectSignature { + &self.signature + } +} diff --git a/tests/functions.rs b/tests/functions.rs index ed89759..c6a7a49 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -3,7 +3,8 @@ use std::rc::Rc; use cten::{ BoxedTensor, COO, Domain, HRect, Physical, RightComposition, Stride, Sum, Tensor, clamp, concatenate, diagonalize, eye, eye_rect, mm, movedim, negate, ones, permute, pow, prod, scale, - sparse_coo_tensor, square, squeeze, stack, sum, transpose, unsqueeze, where_, zeros, + sparse_coo_tensor, sparse_latticed_tensor, square, squeeze, stack, sum, transpose, unsqueeze, + where_, zeros, }; fn physical() -> Rc { @@ -663,6 +664,136 @@ fn sparse_coo_tensor_rejects_a_row_of_the_wrong_length() { sparse_coo_tensor(vec![vec![0, 1]], Rc::new(Physical::new(vec![1.0])), vec![3]); } +// The examples of sparse latticed tensors from albanD/subclass_zoo#54. + +#[test] +fn sparse_latticed_tensor_lays_out_a_diagonal() { + let diagonal = sparse_latticed_tensor( + Rc::new(Physical::new(vec![1.0, 2.0, 3.0])), + vec![vec![1], vec![1]], + vec![0, 0], + vec![3, 3], + ); + + assert_eq!(diagonal.hrect().shape(), &[3, 3]); + assert_eq!( + collect_2d(&diagonal, 3, 3), + vec![ + vec![1.0, 0.0, 0.0], + vec![0.0, 2.0, 0.0], + vec![0.0, 0.0, 3.0] + ] + ); +} + +#[test] +fn sparse_latticed_tensor_lays_out_an_anti_diagonal() { + let anti_diagonal = sparse_latticed_tensor( + Rc::new(Physical::new(vec![1.0, 2.0, 3.0])), + vec![vec![1], vec![-1]], + vec![0, 2], + vec![3, 3], + ); + + assert_eq!( + collect_2d(&anti_diagonal, 3, 3), + vec![ + vec![0.0, 0.0, 1.0], + vec![0.0, 2.0, 0.0], + vec![3.0, 0.0, 0.0] + ] + ); +} + +#[test] +fn sparse_latticed_tensor_lays_out_a_block_diagonal() { + // Three 2 by 2 blocks holding 1 to 12, block after block, row by row. + let values: Vec = (1..=12).map(|value| value as f32).collect(); + let blocks = Rc::new(RightComposition::new( + Rc::new(Physical::new(values)), + Rc::new(Stride::new(HRect::new(vec![3, 2, 2]), vec![4, 2, 1])), + )); + let block_diagonal = sparse_latticed_tensor( + blocks, + vec![vec![2, 1, 0], vec![2, 0, 1]], + vec![0, 0], + vec![6, 6], + ); + + assert_eq!( + collect_2d(&block_diagonal, 6, 6), + vec![ + vec![1.0, 2.0, 0.0, 0.0, 0.0, 0.0], + vec![3.0, 4.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 0.0, 5.0, 6.0, 0.0, 0.0], + vec![0.0, 0.0, 7.0, 8.0, 0.0, 0.0], + vec![0.0, 0.0, 0.0, 0.0, 9.0, 10.0], + vec![0.0, 0.0, 0.0, 0.0, 11.0, 12.0], + ] + ); +} + +#[test] +fn sparse_latticed_tensor_pads_with_its_offset_and_shape() { + let padded = sparse_latticed_tensor( + matrix(2, 2, vec![1.0, 2.0, 3.0, 4.0]), + vec![vec![1, 0], vec![0, 1]], + vec![1, 0], + vec![3, 4], + ); + + assert_eq!( + collect_2d(&padded, 3, 4), + vec![ + vec![0.0, 0.0, 0.0, 0.0], + vec![1.0, 2.0, 0.0, 0.0], + vec![3.0, 4.0, 0.0, 0.0] + ] + ); +} + +// The physical tensor is a kernel repeated along a new first axis by a zero +// stride, as torch's `expand` does, so it still holds only three values. +#[test] +fn sparse_latticed_tensor_lays_out_a_circulant_matrix() { + let kernel = Rc::new(RightComposition::new( + Rc::new(Physical::new(vec![1.0, 2.0, 3.0])), + Rc::new(Stride::new(HRect::new(vec![4, 3]), vec![0, 1])), + )); + let circulant = + sparse_latticed_tensor(kernel, vec![vec![1, 0], vec![1, 1]], vec![0, 0], vec![4, 6]); + + assert_eq!( + collect_2d(&circulant, 4, 6), + vec![ + vec![1.0, 2.0, 3.0, 0.0, 0.0, 0.0], + vec![0.0, 1.0, 2.0, 3.0, 0.0, 0.0], + vec![0.0, 0.0, 1.0, 2.0, 3.0, 0.0], + vec![0.0, 0.0, 0.0, 1.0, 2.0, 3.0], + ] + ); +} + +// torch's `x.repeat(4, 2)` without copying `x`: the physical tensor repeats it +// by zero strides, and the basis puts the two copies side by side on each row. +#[test] +fn sparse_latticed_tensor_repeats_a_vector_without_copying_it() { + let expanded = Rc::new(RightComposition::new( + Rc::new(Physical::new(vec![1.0, 2.0, 3.0])), + Rc::new(Stride::new(HRect::new(vec![4, 2, 3]), vec![0, 0, 1])), + )); + let repeated = sparse_latticed_tensor( + expanded, + vec![vec![1, 0, 0], vec![0, 3, 1]], + vec![0, 0], + vec![4, 6], + ); + + for row in collect_2d(&repeated, 4, 6) { + assert_eq!(row, vec![1.0, 2.0, 3.0, 1.0, 2.0, 3.0]); + } +} + #[test] fn sum_adds_up_along_a_dimension() { let u = matrix(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); diff --git a/tests/logic.rs b/tests/logic.rs index ecd3f01..bab84db 100644 --- a/tests/logic.rs +++ b/tests/logic.rs @@ -1,8 +1,8 @@ use std::rc::Rc; use cten::{ - COO, Diagonal, EmptyMapping, HRect, Inf, LeftComposition, Physical, PowersetMapping, Product, - RightComposition, Stride, Sum, Sup, Tensor, + Affine, COO, Diagonal, EmptyMapping, HRect, Inf, LeftComposition, Physical, PowersetMapping, + Product, RightComposition, Stride, Sum, Sup, Tensor, }; fn physical() -> Rc { @@ -287,3 +287,26 @@ fn a_powerset_mapping_can_be_used_directly() { assert!(selected.contains([4].as_slice())); assert!(diagonal.map(&[4, 5]).is_empty()); } + +// Read through an affine map, a tensor gives its values on a lattice: here the +// diagonal of a matrix, as torch's `diagonal` does. +#[test] +fn an_affine_mapping_reads_a_lattice_out_of_a_tensor() { + let values: Vec = (1..=9).map(|value| value as f32).collect(); + let matrix = Rc::new(RightComposition::new( + Rc::new(Physical::new(values)), + Rc::new(Stride::new(HRect::new(vec![3, 3]), vec![3, 1])), + )); + let diagonal = RightComposition::new( + matrix, + Rc::new(Affine::new( + HRect::new(vec![3]), + vec![vec![1], vec![1]], + vec![0, 0], + HRect::new(vec![3, 3]), + )), + ); + + let read: Vec = (0..3).map(|i| diagonal.get(&[i])).collect(); + assert_eq!(read, vec![1.0, 5.0, 9.0]); +} diff --git a/tests/mapping.rs b/tests/mapping.rs index 27ed4f6..0150e65 100644 --- a/tests/mapping.rs +++ b/tests/mapping.rs @@ -1,9 +1,9 @@ -use std::{assert_eq, rc::Rc}; +use std::{assert_eq, collections::HashMap, rc::Rc}; use cten::{ - COO, Concat, Diagonal, EmptyMapping, HRect, IndexMapping, InvProjection, JaggedDomain, - Permutation, Physical, Powerset, PowersetMapping, Projection, Squeeze, Stride, Unsqueeze, - Where, + Affine, COO, Concat, Diagonal, EmptyMapping, HRect, IndexMapping, InvAffine, InvProjection, + JaggedDomain, Permutation, Physical, Powerset, PowersetMapping, Projection, Squeeze, Stride, + Unsqueeze, Where, }; #[test] @@ -71,6 +71,237 @@ fn an_inv_projection_sends_each_index_to_the_line_along_its_dim() { ); } +// The anti-diagonal of a 3 by 3 matrix: a negative step, made up for by the +// offset. +#[test] +fn an_affine_sends_each_index_to_its_lattice_point() { + let affine = Affine::new( + HRect::new(vec![3]), + vec![vec![1], vec![-1]], + vec![0, 2], + HRect::new(vec![3, 3]), + ); + + assert!(affine.domain().contains(&[2])); + assert!(!affine.domain().contains(&[3])); + assert!(affine.codomain().contains(&[2, 2])); + assert_eq!(affine.map(&[0]), vec![0, 2]); + assert_eq!(affine.map(&[1]), vec![1, 1]); + assert_eq!(affine.map(&[2]), vec![2, 0]); +} + +// Without the offset, the same steps would reach -2 along the second dimension. +#[test] +#[should_panic(expected = "the lattice reaches -2 to 0 along dimension 1")] +fn an_affine_rejects_a_lattice_leaving_its_codomain() { + Affine::new( + HRect::new(vec![3]), + vec![vec![1], vec![-1]], + vec![0, 0], + HRect::new(vec![3, 3]), + ); +} + +#[test] +#[should_panic(expected = "a basis column per dimension of the domain")] +fn an_affine_rejects_a_basis_of_the_wrong_shape() { + Affine::new( + HRect::new(vec![3]), + vec![vec![1, 0]], + vec![0], + HRect::new(vec![3]), + ); +} + +// The flattened diagonal of a 3 by 3 matrix: every fourth position is on the +// lattice. +#[test] +fn an_inv_affine_sends_each_index_to_the_lattice_points_on_it() { + let inv_affine = InvAffine::new(Affine::new( + HRect::new(vec![3]), + vec![vec![4]], + vec![0], + HRect::new(vec![9]), + )); + + assert!(inv_affine.domain().contains(&[8])); + assert!(!inv_affine.domain().contains(&[9])); + assert!(inv_affine.codomain().contains(&[2])); + assert_eq!(inv_affine.map(&[4]), Powerset::from([vec![1]])); + assert_eq!(inv_affine.map(&[5]), Powerset::new()); +} + +// Dropping the middle axis is affine too, and sends several indices onto each +// one, which the inverse finds all of. +#[test] +fn an_inv_affine_of_a_projection_matches_the_inv_projection() { + let inv_affine = InvAffine::new(Affine::new( + HRect::new(vec![2, 3, 4]), + vec![vec![1, 0, 0], vec![0, 0, 1]], + vec![0, 0], + HRect::new(vec![2, 4]), + )); + let inv_proj = InvProjection::new(HRect::new(vec![2, 3, 4]), 1); + + for m in 0..2 { + for p in 0..4 { + assert_eq!(inv_affine.map(&[m, p]), inv_proj.map(&[m, p])); + } + } +} + +/// Every index of a hrect of the given shape. +fn indices(shape: &[usize]) -> Vec> { + shape.iter().fold(vec![vec![]], |prefixes, &extent| { + prefixes + .iter() + .flat_map(|prefix| { + (0..extent).map(move |coord| { + let mut index = prefix.clone(); + index.push(coord); + index + }) + }) + .collect() + }) +} + +/// Checks `InvAffine::map` at every index of `shape` against a scan of every +/// index of `physical`. +fn assert_inv_affine_matches_a_scan( + basis: Vec>, + offset: Vec, + physical: Vec, + shape: Vec, +) { + let mut expected: HashMap, Powerset> = HashMap::new(); + for x in indices(&physical) { + let y = basis + .iter() + .zip(&offset) + .map(|(row, &start)| { + let coord: i64 = row + .iter() + .zip(&x) + .map(|(&step, &coord)| step * coord as i64) + .sum(); + (start + coord) as usize + }) + .collect(); + expected.entry(y).or_default().insert(x); + } + let inv_affine = InvAffine::new(Affine::new( + HRect::new(physical), + basis, + offset, + HRect::new(shape.clone()), + )); + for y in indices(&shape) { + let lands = expected.remove(&y).unwrap_or_default(); + assert_eq!(inv_affine.map(&y), lands, "at {y:?}"); + } +} + +// The lattices of all the sparse latticed tensor examples, one that ends before +// its shape does, and two where several indices land on one. +#[test] +fn an_inv_affine_finds_the_indices_a_scan_finds() { + let lattices = [ + (vec![vec![3]], vec![0], vec![5], vec![18]), + (vec![vec![2], vec![3]], vec![0, 0], vec![3], vec![5, 7]), + (vec![vec![1], vec![1]], vec![0, 0], vec![3], vec![3, 3]), + (vec![vec![1], vec![-1]], vec![0, 2], vec![3], vec![3, 3]), + ( + vec![vec![2, 1, 0], vec![2, 0, 1]], + vec![0, 0], + vec![3, 2, 2], + vec![6, 6], + ), + (vec![vec![4]], vec![0], vec![3], vec![9]), + ( + vec![vec![1], vec![1], vec![1]], + vec![0, 0, 0], + vec![3], + vec![3, 3, 3], + ), + ( + vec![vec![1, 0], vec![1, 1]], + vec![0, 0], + vec![4, 3], + vec![4, 6], + ), + ( + vec![vec![1, 0], vec![0, 1]], + vec![1, 0], + vec![2, 2], + vec![3, 4], + ), + ( + vec![vec![0, 1, 0, 4], vec![1, 0, 4, 0]], + vec![0, 0], + vec![2, 2, 2, 2], + vec![6, 6], + ), + ( + vec![vec![0, 2, 1], vec![2, 0, 1]], + vec![0, 0], + vec![4, 4, 2], + vec![8, 8], + ), + ( + vec![vec![1, 0, 0], vec![0, 3, 1]], + vec![0, 0], + vec![4, 2, 3], + vec![4, 6], + ), + (vec![vec![1, 1]], vec![0], vec![3, 3], vec![5]), + ( + vec![vec![1, 0, 0], vec![0, 0, 1]], + vec![0, 0], + vec![2, 3, 4], + vec![2, 4], + ), + ]; + for (basis, offset, physical, shape) in lattices { + assert_inv_affine_matches_a_scan(basis, offset, physical, shape); + } +} + +// Lattices from a fixed seed: entries from -3 to 3, up to three dimensions to +// land in and four to land from, with the offset and shape chosen to fit. +#[test] +fn an_inv_affine_finds_the_indices_a_scan_finds_on_random_lattices() { + let mut seed: u64 = 0x2545_F491_4F6C_DD1D; + let mut next = |bound: u64| { + seed = seed + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + (seed >> 33) % bound + }; + for _ in 0..1000 { + let rows = 1 + next(3) as usize; + let cols = 1 + next(4) as usize; + let physical: Vec = (0..cols).map(|_| 1 + next(3) as usize).collect(); + let basis: Vec> = (0..rows) + .map(|_| (0..cols).map(|_| next(7) as i64 - 3).collect()) + .collect(); + let mut offset = Vec::new(); + let mut shape = Vec::new(); + for row in &basis { + let (low, high) = + row.iter() + .zip(&physical) + .fold((0, 0), |(low, high), (&step, &extent)| { + let end = step * (extent as i64 - 1); + (low + end.min(0), high + end.max(0)) + }); + offset.push(-low); + shape.push((high - low + 1) as usize + next(2) as usize); + } + assert_inv_affine_matches_a_scan(basis, offset, physical, shape); + } +} + #[test] fn a_permutation_reorders_its_domain_and_keeps_its_codomain() { let permutation = Permutation::new(HRect::new(vec![2, 3, 4]), vec![2, 0, 1]);