color/order: Implement color orderings

4 files changed, 334 insertions, 0 deletions

diff --git a/src/color.rs b/src/color.rs
index e0a3399..64fd82b 100644
--- a/src/color.rs
+++ b/src/color.rs
@@ -1,5 +1,6 @@
 //! Colors and color spaces.
 
+pub mod order;
 pub mod source;
 
 use crate::metric::kd::{Cartesian, CartesianMetric};
diff --git a/src/color/order.rs b/src/color/order.rs
new file mode 100644
index 0000000..300a556
--- /dev/null
+++ b/src/color/order.rs
@@ -0,0 +1,196 @@
+//! Linear orders for colors.
+
+use super::source::ColorSource;
+use super::Rgb8;
+
+use crate::hilbert::hilbert_point;
+
+use rand::seq::SliceRandom;
+use rand::Rng;
+
+use std::cmp::Ordering;
+
+/// An iterator over all colors from a source.
+#[derive(Debug)]
+struct ColorSourceIter<S> {
+    source: S,
+    coords: Vec<usize>,
+}
+
+impl<S: ColorSource> From<S> for ColorSourceIter<S> {
+    fn from(source: S) -> Self {
+        let coords = vec![0; source.dimensions().len()];
+
+        Self { source, coords }
+    }
+}
+
+impl<S: ColorSource> Iterator for ColorSourceIter<S> {
+    type Item = Rgb8;
+
+    fn next(&mut self) -> Option<Rgb8> {
+        if self.coords.is_empty() {
+            return None;
+        }
+
+        let color = self.source.get_color(&self.coords);
+
+        let dims = self.source.dimensions();
+        for i in 0..dims.len() {
+            self.coords[i] += 1;
+            if self.coords[i] < dims[i] {
+                break;
+            } else if i == dims.len() - 1 {
+                self.coords.clear();
+            } else {
+                self.coords[i] = 0;
+            }
+        }
+
+        Some(color)
+    }
+}
+
+/// Wrapper for sorting colors by hue.
+#[derive(Debug, Eq, PartialEq)]
+struct Hue {
+    /// The quadrant of the hue angle.
+    quad: i32,
+    /// The numerator of the hue calculation.
+    num: i32,
+    /// The denominator of the hue calculation.
+    denom: i32,
+}
+
+impl From<Rgb8> for Hue {
+    fn from(rgb8: Rgb8) -> Self {
+        // The hue angle is atan2(sqrt(3) * (G - B), 2 * R - G - B).  We avoid actually computing
+        // the atan2() as an optimization.
+        let r = rgb8[0] as i32;
+        let g = rgb8[1] as i32;
+        let b = rgb8[2] as i32;
+
+        let num = g - b;
+        let mut denom = 2 * r - g - b;
+        if num == 0 && denom == 0 {
+            denom = 1;
+        }
+
+        let quad = match (num >= 0, denom >= 0) {
+            (true, true) => 0,
+            (true, false) => 1,
+            (false, false) => 2,
+            (false, true) => 3,
+        };
+
+        Self { quad, num, denom }
+    }
+}
+
+impl Ord for Hue {
+    fn cmp(&self, other: &Self) -> Ordering {
+        // Within the same quadrant,
+        //
+        //     atan2(n1, d1) < atan2(n2, d2)  iff
+        //           n1 / d1 < n2 / d2        iff
+        //           n1 * d2 < n2 * d1
+        self.quad
+            .cmp(&other.quad)
+            .then_with(|| (self.num * other.denom).cmp(&(other.num * self.denom)))
+    }
+}
+
+impl PartialOrd for Hue {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+/// Iterate over colors sorted by their hue.
+pub fn hue_sorted<S: ColorSource>(source: S) -> Vec<Rgb8> {
+    let mut colors: Vec<_> = ColorSourceIter::from(source).collect();
+    colors.sort_by_key(|c| Hue::from(*c));
+    colors
+}
+
+/// Iterate over colors in random order.
+pub fn shuffled<S: ColorSource, R: Rng>(source: S, rng: &mut R) -> Vec<Rgb8> {
+    let mut colors: Vec<_> = ColorSourceIter::from(source).collect();
+    colors.shuffle(rng);
+    colors
+}
+
+/// ceil(log_2(n)). for rounding up to powers of 2.
+fn log2(n: usize) -> u32 {
+    let nbits = 8 * std::mem::size_of::<usize>() as u32;
+    nbits - (n - 1).leading_zeros()
+}
+
+/// Iterate over colors in Morton order (Z-order).
+pub fn morton<S: ColorSource>(source: S) -> Vec<Rgb8> {
+    let mut colors = Vec::new();
+
+    let dims = source.dimensions();
+    let ndims = dims.len();
+
+    let nbits = ndims * dims.iter().map(|n| log2(*n) as usize).max().unwrap();
+
+    let size = 1usize << nbits;
+    let mut coords = vec![0; ndims];
+    for i in 0..size {
+        for x in &mut coords {
+            *x = 0;
+        }
+        for j in 0..nbits {
+            let bit = (i >> j) & 1;
+            coords[j % ndims] |= bit << (j / ndims);
+        }
+        if coords.iter().zip(dims.iter()).all(|(x, n)| x < n) {
+            colors.push(source.get_color(&coords));
+        }
+    }
+
+    colors
+}
+
+/// Iterate over colors in Hilbert curve order.
+pub fn hilbert<S: ColorSource>(source: S) -> Vec<Rgb8> {
+    let mut colors = Vec::new();
+
+    let dims = source.dimensions();
+    let ndims = dims.len();
+
+    let bits: Vec<_> = dims.iter().map(|n| log2(*n)).collect();
+    let nbits: u32 = bits.iter().sum();
+    let size = 1usize << nbits;
+
+    let mut coords = vec![0; ndims];
+
+    for i in 0..size {
+        hilbert_point(i, &bits, &mut coords);
+        if coords.iter().zip(dims.iter()).all(|(x, n)| x < n) {
+            colors.push(source.get_color(&coords));
+        }
+    }
+
+    colors
+}
+
+/// Stripe an ordered list of colors, to reduce artifacts in the generated image.
+///
+/// The striped ordering gives every other item first, then every other item from the remaining
+/// items, etc. For example, the striped form of `0..16` is
+/// `[0, 2, 4, 6, 8, 10, 12, 14, 1, 5, 9, 13, 3, 11, 7, 15]`.
+pub fn striped(colors: Vec<Rgb8>) -> Vec<Rgb8> {
+    let len = colors.len();
+    let mut result = Vec::with_capacity(len);
+    let mut stripe = 1;
+    while stripe <= len {
+        for i in ((stripe - 1)..len).step_by(2 * stripe) {
+            result.push(colors[i]);
+        }
+        stripe *= 2;
+    }
+
+    result
+}
diff --git a/src/hilbert.rs b/src/hilbert.rs
new file mode 100644
index 0000000..c0982d4
--- /dev/null
+++ b/src/hilbert.rs
@@ -0,0 +1,136 @@
+//! Implementation of [Compact Hilbert Indices](https://dl.acm.org/doi/10.1109/CISIS.2007.16) by
+//! Chris Hamilton.
+
+/// Right rotation of x by b bits out of n.
+fn rotate_right(x: usize, b: u32, n: u32) -> usize {
+    let l = x & ((1 << b) - 1);
+    let r = x >> b;
+    (l << (n - b)) | r
+}
+
+/// Left rotation of x by b bits out of n.
+fn rotate_left(x: usize, b: u32, n: u32) -> usize {
+    rotate_right(x, n - b, n)
+}
+
+/// Binary reflected Gray code.
+fn gray_code(i: usize) -> usize {
+    i ^ (i >> 1)
+}
+
+/// e(i), the entry point for the ith sub-hypercube.
+fn entry_point(i: usize) -> usize {
+    if i == 0 {
+        0
+    } else {
+        gray_code((i - 1) & !1)
+    }
+}
+
+/// g(i), the inter sub-hypercube direction.
+fn inter_direction(i: usize) -> u32 {
+    // g(i) counts the trailing set bits in i
+    (!i).trailing_zeros()
+}
+
+/// d(i), the intra sub-hypercube direction.
+fn intra_direction(i: usize) -> u32 {
+    if i & 1 != 0 {
+        inter_direction(i)
+    } else if i > 0 {
+        inter_direction(i - 1)
+    } else {
+        0
+    }
+}
+
+/// T transformation inverse
+fn t_inverse(dims: u32, e: usize, d: u32, a: usize) -> usize {
+    rotate_left(a, d, dims) ^ e
+}
+
+/// GrayCodeRankInverse
+fn gray_code_rank_inverse(
+    dims: u32,
+    mu: usize,
+    pi: usize,
+    r: usize,
+    free_bits: u32,
+) -> (usize, usize) {
+    // The inverse rank of r
+    let mut i = 0;
+    // gray_code(i)
+    let mut g = 0;
+
+    let mut j = free_bits - 1;
+    for k in (0..dims).rev() {
+        if mu & (1 << k) == 0 {
+            g |= pi & (1 << k);
+            i |= (g ^ (i >> 1)) & (1 << k);
+        } else {
+            i |= ((r >> j) & 1) << k;
+            g |= (i ^ (i >> 1)) & (1 << k);
+            j = j.wrapping_sub(1);
+        }
+    }
+
+    (i, g)
+}
+
+/// ExtractMask.
+fn extract_mask(bits: &[u32], i: u32) -> (usize, u32) {
+    // The mask
+    let mut mu = 0;
+    // popcount(mu)
+    let mut free_bits = 0;
+
+    let dims = bits.len();
+    for j in (0..dims).rev() {
+        mu <<= 1;
+        if bits[j] > i {
+            mu |= 1;
+            free_bits += 1;
+        }
+    }
+
+    (mu, free_bits)
+}
+
+/// Compute the corresponding point for a Hilbert index (CompactHilbertIndexInverse).
+pub fn hilbert_point(index: usize, bits: &[u32], point: &mut [usize]) {
+    let dims = bits.len() as u32;
+    let max = *bits.iter().max().unwrap();
+    let sum: u32 = bits.iter().sum();
+
+    let mut e = 0;
+    let mut k = 0;
+
+    // Next direction; we use d instead of d + 1 everywhere
+    let mut d = 1;
+
+    for x in point.iter_mut() {
+        *x = 0;
+    }
+
+    for i in (0..max).rev() {
+        let (mut mu, free_bits) = extract_mask(bits, i);
+        mu = rotate_right(mu, d, dims);
+
+        let pi = rotate_right(e, d, dims) & !mu;
+
+        let r = (index >> (sum - k - free_bits)) & ((1 << free_bits) - 1);
+
+        k += free_bits;
+
+        let (w, mut l) = gray_code_rank_inverse(dims, mu, pi, r, free_bits);
+        l = t_inverse(dims, e, d, l);
+
+        for x in point.iter_mut() {
+            *x |= (l & 1) << i;
+            l >>= 1;
+        }
+
+        e = e ^ rotate_right(entry_point(w), d, dims);
+        d = (d + intra_direction(w) + 1) % dims;
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index a7bda67..a59a0cf 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,5 @@
 pub mod color;
+pub mod hilbert;
 pub mod metric;
 
 fn main() {}