Add Affine and InvAffine for sparse latticed tensors

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<T: BoxedTensor + 'static>(
+    physical: Rc<T>,
+    basis: Vec<Vec<i64>>,
+    offset: Vec<i64>,
+    shape: Vec<usize>,
+) -> Sum<InvAffine> {
+    let affine = Affine::new(physical.hrect().clone(), basis, offset, HRect::new(shape));
+    Sum::new(physical, Rc::new(InvAffine::new(affine)))
+}
+
 pub fn sum<T: BoxedTensor + 'static>(input: Rc<T>, dim: usize) -> Sum<InvProjection> {
     let codomain = input.hrect().clone();
     Sum::new(input, Rc::new(InvProjection::new(codomain, dim)))

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<Vec<i64>>,
+    h_pivots: Vec<usize>,
+    particular: Vec<Vec<i64>>,
+    kernel: Vec<Vec<i64>>,
+    kernel_pivots: Vec<usize>,
+    extents: Vec<i64>,
+}
+
+impl Preimage {
+    pub(crate) fn new(basis: &[Vec<i64>], 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<Vec<usize>> {
+        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<i64> = 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<Vec<usize>>) {
+        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<Vec<i64>>,
+    cols: usize,
+) -> (Vec<Vec<i64>>, Vec<Vec<i64>>, Vec<usize>) {
+    let mut operations: Vec<Vec<i64>> = (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<i64>], 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)
+}

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::{

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<Vec<i64>>,
+    offset: Vec<i64>,
+}
+
+impl Affine {
+    pub fn new(domain: HRect, basis: Vec<Vec<i64>>, offset: Vec<i64>, 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<usize> {
+        // 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::<i64>();
+                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<i64>,
+    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<i64> = 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
+    }
+}