From c653c8cda8f49d3bbe07190a6477367290ff7f04 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sun, 19 Apr 2020 16:28:10 -0400 Subject: Begin re-writing in Rust --- .gitattributes | 1 + .gitignore | 15 +- COPYING | 13 -- Cargo.lock | 5 + Cargo.toml | 5 + LICENSE | 12 ++ Makefile | 47 ------ color.c | 176 --------------------- color.h | 30 ---- generate.c | 82 ---------- generate.h | 21 --- hilbert.c | 170 -------------------- hilbert.h | 31 ---- kd-forest.c | 327 --------------------------------------- kd-forest.h | 56 ------- main.c | 480 --------------------------------------------------------- options.c | 431 --------------------------------------------------- options.h | 58 ------- src/main.rs | 1 + util.c | 66 -------- util.h | 23 --- 21 files changed, 36 insertions(+), 2014 deletions(-) create mode 100644 .gitattributes delete mode 100644 COPYING create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 LICENSE delete mode 100644 Makefile delete mode 100644 color.c delete mode 100644 color.h delete mode 100644 generate.c delete mode 100644 generate.h delete mode 100644 hilbert.c delete mode 100644 hilbert.h delete mode 100644 kd-forest.c delete mode 100644 kd-forest.h delete mode 100644 main.c delete mode 100644 options.c delete mode 100644 options.h create mode 100644 src/main.rs delete mode 100644 util.c delete mode 100644 util.h diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..c353fc3 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +Cargo.lock binary diff --git a/.gitignore b/.gitignore index aea0ddb..ecd771b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,13 @@ -*.o -/kd-forest +# Generated by Cargo +# will have compiled files and executables +/target/ + +# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries +# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html +# Cargo.lock + +# These are backup files generated by rustfmt +**/*.rs.bk + +# Generated images *.png -*.mkv diff --git a/COPYING b/COPYING deleted file mode 100644 index c6c7def..0000000 --- a/COPYING +++ /dev/null @@ -1,13 +0,0 @@ - DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE - Version 2, December 2004 - - Copyright (C) 2004 Sam Hocevar - - Everyone is permitted to copy and distribute verbatim or modified - copies of this license document, and changing it is allowed as long - as the name is changed. - - DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE - TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION - - 0. You just DO WHAT THE FUCK YOU WANT TO. diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..a63d677 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,5 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +[[package]] +name = "kd-forest" +version = "2.0.0" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..65ffec8 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,5 @@ +[package] +name = "kd-forest" +version = "2.0.0" +authors = ["Tavian Barnes "] +edition = "2018" diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..9a254ab --- /dev/null +++ b/LICENSE @@ -0,0 +1,12 @@ +Copyright (C) 2015-2020 Tavian Barnes + +Permission to use, copy, modify, and/or distribute this software for any +purpose with or without fee is hereby granted. + +THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. diff --git a/Makefile b/Makefile deleted file mode 100644 index 4e035e0..0000000 --- a/Makefile +++ /dev/null @@ -1,47 +0,0 @@ -##################################################################### -# kd-forest # -# Copyright (C) 2014 Tavian Barnes # -# # -# This program is free software. It comes without any warranty, to # -# the extent permitted by applicable law. You can redistribute it # -# and/or modify it under the terms of the Do What The Fuck You Want # -# To Public License, Version 2, as published by Sam Hocevar. See # -# the COPYING file or http://www.wtfpl.net/ for more details. # -##################################################################### - -CC = gcc -CFLAGS = -std=c99 -D_POSIX_C_SOURCE=200809L -pipe -g -O3 -flto -Wall -Wpedantic -Wextra -Wno-sign-compare -Wno-unused-parameter -Wunreachable-code -Wshadow -Wpointer-arith -Wwrite-strings -Wcast-align -Wstrict-prototypes -LDFLAGS = -Wl,-O1,--sort-common,--as-needed,-z,relro -LIBS = -lm -lpng -RM = rm -f - -DEPS = Makefile color.h hilbert.h generate.h kd-forest.h options.h util.h - -IMAGEFLAGS = -b 24 -s -l min -c Lab -ANIMFLAGS = -b 19 -s -l mean -c Lab - -kd-forest: color.o generate.o hilbert.o kd-forest.o main.o options.o util.o - $(CC) $(CFLAGS) $(LDFLAGS) $^ $(LIBS) -o kd-forest - -%.o: %.c $(DEPS) - $(CC) $(CFLAGS) -c $< -o $@ - -image: kd-forest.png - -kd-forest.png: kd-forest - ./kd-forest $(IMAGEFLAGS) -o kd-forest.png - -anim: kd-forest.mkv - -kd-forest.mkv: kd-forest - $(RM) kd-forest.mkv - mkdir /tmp/kd-frames - ./kd-forest $(ANIMFLAGS) -a -o /tmp/kd-frames - ffmpeg -r 60 -i /tmp/kd-frames/%04d.png -c:v libx264 -preset veryslow -qp 0 kd-forest.mkv - $(RM) -r /tmp/kd-frames - -clean: - $(RM) *.o - $(RM) kd-forest - -.PHONY: image anim clean diff --git a/color.c b/color.c deleted file mode 100644 index 9d15034..0000000 --- a/color.c +++ /dev/null @@ -1,176 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2014 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#include "color.h" -#include - -void -color_unpack(uint8_t pixel[3], uint32_t color) -{ - pixel[0] = (color >> 16) & 0xFF; - pixel[1] = (color >> 8) & 0xFF; - pixel[2] = color & 0xFF; -} - -void -color_set_RGB(double coords[3], uint32_t color) -{ - uint8_t pixel[3]; - color_unpack(pixel, color); - for (int i = 0; i < 3; ++i) { - coords[i] = pixel[i]/255.0; - } -} - -// Inverse gamma for sRGB -double -sRGB_C_inv(double t) -{ - if (t <= 0.040449936) { - return t/12.92; - } else { - return pow((t + 0.055)/1.055, 2.4); - } -} - -static void -color_set_XYZ(double XYZ[3], uint32_t color) -{ - double RGB[3]; - color_set_RGB(RGB, color); - - RGB[0] = sRGB_C_inv(RGB[0]); - RGB[1] = sRGB_C_inv(RGB[1]); - RGB[2] = sRGB_C_inv(RGB[2]); - - XYZ[0] = 0.4123808838268995*RGB[0] + 0.3575728355732478*RGB[1] + 0.1804522977447919*RGB[2]; - XYZ[1] = 0.2126198631048975*RGB[0] + 0.7151387878413206*RGB[1] + 0.0721499433963131*RGB[2]; - XYZ[2] = 0.0193434956789248*RGB[0] + 0.1192121694056356*RGB[1] + 0.9505065664127130*RGB[2]; -} - -// CIE L*a*b* and L*u*v* gamma -static double -Lab_f(double t) -{ - if (t > 216.0/24389.0) { - return pow(t, 1.0/3.0); - } else { - return 841.0*t/108.0 + 4.0/29.0; - } -} - -// sRGB white point (CIE D50) in XYZ coordinates -static const double WHITE[] = { - [0] = 0.9504060171449392, - [1] = 0.9999085943425312, - [2] = 1.089062231497274, -}; - -void -color_set_Lab(double coords[3], uint32_t color) -{ - double XYZ[3]; - color_set_XYZ(XYZ, color); - - double fXYZ[] = { - [0] = Lab_f(XYZ[0]/WHITE[0]), - [1] = Lab_f(XYZ[1]/WHITE[1]), - [2] = Lab_f(XYZ[2]/WHITE[2]), - }; - - coords[0] = 116.0*fXYZ[1] - 16.0; - coords[1] = 500.0*(fXYZ[0] - fXYZ[1]); - coords[2] = 200.0*(fXYZ[1] - fXYZ[2]); -} - -void -color_set_Luv(double coords[3], uint32_t color) -{ - double XYZ[3]; - color_set_XYZ(XYZ, color); - - double uv_denom = XYZ[0] + 15.0*XYZ[1] + 3.0*XYZ[2]; - if (uv_denom == 0.0) { - coords[0] = 0.0; - coords[1] = 0.0; - coords[2] = 0.0; - return; - } - - double white_uv_denom = WHITE[0] + 16.0*WHITE[1] + 3.0*WHITE[2]; - - double fY = Lab_f(XYZ[1]/WHITE[1]); - double uprime = 4.0*XYZ[0]/uv_denom; - double unprime = 4.0*WHITE[0]/white_uv_denom; - double vprime = 9.0*XYZ[1]/uv_denom; - double vnprime = 9.0*WHITE[1]/white_uv_denom; - - coords[0] = 116.0*fY - 16.0; - coords[1] = 13.0*coords[0]*(uprime - unprime); - coords[2] = 13.0*coords[0]*(vprime - vnprime); -} - -int -color_comparator(const void *a, const void *b) -{ - uint8_t aRGB[3], bRGB[3]; - color_unpack(aRGB, *(uint32_t *)a); - color_unpack(bRGB, *(uint32_t *)b); - - int anum = aRGB[1] - aRGB[2], adenom = 2*aRGB[0] - aRGB[1] - aRGB[2]; - int bnum = bRGB[1] - bRGB[2], bdenom = 2*bRGB[0] - bRGB[1] - bRGB[2]; - - // The hue angle is defined as atan2(sqrt(3)*n/d) (+ 2*pi if negative). But - // since atan2() is expensive, we compute an equivalent ordering while - // avoiding trig calls. First, handle the quadrants. We have: - // - // hue(n, d) - // | d >= 0 && n == 0 = 0 - // | d >= 0 && n > 0 = atan(n/d) - // | d >= 0 && n < 0 = atan(n/d) + 2*pi - // | d < 0 = atan(n/d) + pi - // - // and since atan(n/d)'s range is [-pi/2, pi/2], each chunk can be strictly - // ordered relative to the other chunks. - if (adenom >= 0) { - if (anum >= 0) { - if (bdenom < 0 || bnum < 0) { - return -1; - } - } else { - if (bdenom < 0 || bnum >= 0) { - return 1; - } - } - } else if (bdenom >= 0) { - if (bnum >= 0) { - return 1; - } else { - return -1; - } - } - - // Special-case zero numerators, because we treat 0/0 as 0, not NaN - if (anum == 0 || bnum == 0) { - int lhs = adenom >= 0 ? anum : -anum; - int rhs = bdenom >= 0 ? bnum : -bnum; - return lhs - rhs; - } - - // The points are in the same/comparable quadrants. We can still avoid - // calculating atan(n/d) though, because it's an increasing function in n/d. - // We can also avoid a division, by noting that an/ad < bn/bd iff - // an*bd*sgn(ad*bd) < bn*ad*sgn(ad*bd). Due to the logic above, both - // denominators must have the same sign, so the sgn()s are redundant. - int lhs = anum*bdenom; - int rhs = bnum*adenom; - return lhs - rhs; -} diff --git a/color.h b/color.h deleted file mode 100644 index 0abafbd..0000000 --- a/color.h +++ /dev/null @@ -1,30 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2014 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#ifndef COLOR_H -#define COLOR_H - -#include - -// Unpack a color into 8-bit RGB values -void color_unpack(uint8_t pixel[3], uint32_t color); - -// Use RGB coordinates -void color_set_RGB(double coords[3], uint32_t color); -// Use CIE L*a*b* coordinates -void color_set_Lab(double coords[3], uint32_t color); -// Use CIE L*u*v* coordinates -void color_set_Luv(double coords[3], uint32_t color); - -// For sorting by hue -int color_comparator(const void *a, const void *b); - -#endif // COLOR_H diff --git a/generate.c b/generate.c deleted file mode 100644 index 9a4eac8..0000000 --- a/generate.c +++ /dev/null @@ -1,82 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2015 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#include "generate.h" -#include "color.h" -#include "hilbert.h" -#include "util.h" -#include - -uint32_t * -generate_colors(const options_t *options) -{ - const unsigned int bit_depth = options->bit_depth; - mode_t mode = options->mode; - - // Allocate bits from most to least perceptually important - unsigned int grb_bits[3]; - for (unsigned int i = 0; i < 3; ++i) { - grb_bits[i] = (bit_depth + 2 - i)/3; - } - - uint32_t *colors = xmalloc(options->ncolors*sizeof(uint32_t)); - for (uint32_t i = 0; i < (1 << bit_depth); ++i) { - uint32_t n = i; - uint32_t grb[3] = {0, 0, 0}; - - switch (mode) { - case MODE_MORTON: - for (unsigned int j = 0; j < bit_depth; ++j) { - grb[j%3] |= (i & (1 << j)) >> (j - j/3); - } - break; - - case MODE_HILBERT: - hilbert_point(3, grb_bits, n, grb); - break; - - default: - for (unsigned int j = 0; j < 3; ++j) { - grb[j] = n & ((1 << grb_bits[j]) - 1); - n >>= grb_bits[j]; - } - break; - } - - // Pad out colors, and put them in RGB order - grb[0] <<= 16U - grb_bits[0]; - grb[1] <<= 24U - grb_bits[1]; - grb[2] <<= 8U - grb_bits[2]; - - colors[i] = grb[1] | grb[0] | grb[2]; - } - - switch (mode) { - case MODE_HUE_SORT: - qsort(colors, options->ncolors, sizeof(uint32_t), color_comparator); - break; - - case MODE_RANDOM: - // Fisher-Yates shuffle - for (unsigned int i = options->ncolors; i-- > 0;) { - unsigned int j = xrand(i + 1); - uint32_t temp = colors[i]; - colors[i] = colors[j]; - colors[j] = temp; - } - break; - - default: - break; - } - - return colors; -} diff --git a/generate.h b/generate.h deleted file mode 100644 index 2b78150..0000000 --- a/generate.h +++ /dev/null @@ -1,21 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2015 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#ifndef GENERATE_H -#define GENERATE_H - -#include "options.h" -#include - -// Generate the colors according to the mode -uint32_t *generate_colors(const options_t *options); - -#endif // GENERATE_H diff --git a/hilbert.c b/hilbert.c deleted file mode 100644 index 98cf247..0000000 --- a/hilbert.c +++ /dev/null @@ -1,170 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2015 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#include "hilbert.h" -#include - -// These algorithms are described in "Compact Hilbert Indices" by Chris Hamilton - -// Right rotation of x by b bits out of n -static uint32_t -ror(uint32_t x, unsigned int b, unsigned int n) -{ - uint32_t l = x & ((1 << b) - 1); - uint32_t r = x >> b; - return (l << (n - b)) | r; -} - -// Left rotation of x by b bits out of n -static uint32_t -rol(uint32_t x, unsigned int b, unsigned int n) -{ - return ror(x, n - b, n); -} - -// Binary reflected Gray code -uint32_t -gray_code(uint32_t i) -{ - return i ^ (i >> 1); -} - -// e(i), the entry point for the ith sub-hypercube -uint32_t -entry_point(uint32_t i) -{ - if (i == 0) { - return 0; - } else { - return gray_code((i - 1) & ~1U); - } -} - -// g(i), the inter sub-hypercube direction -unsigned int -inter_direction(uint32_t i) -{ - // g(i) counts the trailing set bits in i - unsigned int g = 0; - while (i & 1) { - ++g; - i >>= 1; - } - return g; -} - -// d(i), the intra sub-hypercube direction -unsigned int -intra_direction(uint32_t i) -{ - if (i & 1) { - return inter_direction(i); - } else if (i > 0) { - return inter_direction(i - 1); - } else { - return 0; - } -} - -// T transformation inverse -uint32_t -t_inverse(unsigned int dimensions, uint32_t e, unsigned int d, uint32_t a) -{ - return rol(a, d, dimensions) ^ e; -} - -// GrayCodeRankInverse -void -gray_code_rank_inverse(unsigned int dimensions, uint32_t mu, uint32_t pi, unsigned int r, unsigned int free_bits, uint32_t *i, uint32_t *g) -{ - // *i is the inverse rank of r - // *g is gray_code(i) - - *i = 0; - *g = 0; - - for (unsigned int j = free_bits - 1, k = dimensions; k-- > 0;) { - if (mu & (1 << k)) { - *i |= ((r >> j) & 1) << k; - *g |= (*i ^ (*i >> 1)) & (1 << k); - --j; - } else { - *g |= pi & (1 << k); - *i |= (*g ^ (*i >> 1)) & (1 << k); - } - } -} - -// ExtractMask -void -extract_mask(unsigned int dimensions, const unsigned int extents[], unsigned int i, uint32_t *mu, unsigned int *free_bits) -{ - // *mu is the mask - // *free_bits is popcount(*mu) - - *mu = 0; - *free_bits = 0; - - for (unsigned int j = dimensions; j-- > 0;) { - *mu <<= 1; - - if (extents[j] > i) { - *mu |= 1; - *free_bits += 1; - } - } -} - -// CompactHilbertIndexInverse -void -hilbert_point(unsigned int dimensions, const unsigned int extents[], uint32_t index, uint32_t point[]) -{ - unsigned int max_extent = 0, total_extent = 0; - for (unsigned int i = 0; i < dimensions; ++i) { - if (extents[i] > max_extent) { - max_extent = extents[i]; - } - total_extent += extents[i]; - point[i] = 0; - } - - uint32_t e = 0; - unsigned int k = 0; - - // Next direction; we use d instead of d + 1 everywhere - unsigned int d = 1; - - for (unsigned int i = max_extent; i-- > 0;) { - uint32_t mu; - unsigned int free_bits; - extract_mask(dimensions, extents, i, &mu, &free_bits); - mu = ror(mu, d, dimensions); - - uint32_t pi = ror(e, d, dimensions) & ~mu; - - unsigned int r = (index >> (total_extent - k - free_bits)) & ((1 << free_bits) - 1); - - k += free_bits; - - uint32_t w, l; - gray_code_rank_inverse(dimensions, mu, pi, r, free_bits, &w, &l); - - l = t_inverse(dimensions, e, d, l); - - for (unsigned int j = 0; j < 3; ++j) { - point[j] |= (l & 1) << i; - l >>= 1; - } - - e = e ^ ror(entry_point(w), d, dimensions); - d = (d + intra_direction(w) + 1)%dimensions; - } -} diff --git a/hilbert.h b/hilbert.h deleted file mode 100644 index 4628ad6..0000000 --- a/hilbert.h +++ /dev/null @@ -1,31 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2015 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#ifndef HILBERT_H -#define HILBERT_H - -#include - -/** - * Compute the point corresponding to the given (compact) Hilbert index. - * - * @param dimensions - * The number of spatial dimensions. - * @param extents - * The bit depth of each dimension. - * @param index - * The (compact) Hilbert index of the desired point. - * @param[out] point - * Will hold the point on the Hilbert curve at index. - */ -void hilbert_point(unsigned int dimensions, const unsigned int extents[], uint32_t index, uint32_t point[]); - -#endif // HILBERT_H diff --git a/kd-forest.c b/kd-forest.c deleted file mode 100644 index 3840a61..0000000 --- a/kd-forest.c +++ /dev/null @@ -1,327 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2014 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#include "kd-forest.h" -#include "util.h" -#include -#include -#include -#include - -kd_node_t * -new_kd_node(double coords[KD_DIMEN], unsigned int x, unsigned int y) -{ - kd_node_t *node = xmalloc(sizeof(kd_node_t)); - memcpy(node->coords, coords, sizeof(node->coords)); - node->left = node->right = NULL; - node->x = x; - node->y = y; - node->removed = false; - return node; -} - -static void -kd_destroy(kd_node_t *node) -{ - if (node) { - kd_destroy(node->left); - kd_destroy(node->right); - free(node); - } -} - -static size_t -kd_collect_nodes(kd_node_t *root, kd_node_t **buffer, bool include_removed) -{ - size_t count = 0; - - if (include_removed || !root->removed) { - buffer[0] = root; - ++count; - } - - if (root->left) { - count += kd_collect_nodes(root->left, buffer + count, include_removed); - } - if (root->right) { - count += kd_collect_nodes(root->right, buffer + count, include_removed); - } - - if (root->removed && !include_removed) { - free(root); - } - - return count; -} - -typedef int kd_comparator(const void *a, const void* b); - -static int kd_compare0(const void *a, const void *b) { - double aval = (*(const kd_node_t **)a)->coords[0]; - double bval = (*(const kd_node_t **)b)->coords[0]; - return (aval > bval) - (aval < bval); -} - -static int kd_compare1(const void *a, const void *b) { - double aval = (*(const kd_node_t **)a)->coords[1]; - double bval = (*(const kd_node_t **)b)->coords[1]; - return (aval > bval) - (aval < bval); -} - -static int kd_compare2(const void *a, const void *b) { - double aval = (*(const kd_node_t **)a)->coords[2]; - double bval = (*(const kd_node_t **)b)->coords[2]; - return (aval > bval) - (aval < bval); -} - -static kd_comparator *kd_comparators[KD_DIMEN] = { - kd_compare0, - kd_compare1, - kd_compare2, -}; - -// When building k-d trees, we use KD_DIMEN sorted arrays of nodes plus one -// extra array for scratch space -#define KD_BUFSIZE (KD_DIMEN + 1) - -static kd_node_t * -kd_build_tree_recursive(kd_node_t **buffers[KD_BUFSIZE], size_t size, unsigned int coord) -{ - if (size == 0) { - return NULL; - } - - size_t split = size/2; - size_t left_size = split, right_size = size - left_size - 1; - kd_node_t *root = buffers[coord][split]; - for (size_t i = 0; i < size; ++i) { - buffers[coord][i]->is_left = i < left_size; - } - - kd_node_t **right_buffers[KD_BUFSIZE]; - for (int i = 0; i < KD_DIMEN; ++i) { - right_buffers[i] = buffers[i] + left_size + 1; - } - - kd_node_t **scratch = buffers[KD_DIMEN]; - right_buffers[KD_DIMEN] = scratch; - - for (size_t i = 0; i < KD_DIMEN; ++i) { - if (i == coord) { - continue; - } - - kd_node_t **buffer = buffers[i]; - kd_node_t **right_buffer = right_buffers[i]; - for (size_t j = 0, k = 0, skip = 0; j < size; ++j) { - if (buffer[j]->is_left) { - buffer[j - skip] = buffer[j]; - } else { - if (buffer[j] != root) { - scratch[k] = buffer[j]; - ++k; - } - ++skip; - } - } - for (size_t j = 0; j < right_size; ++j) { - right_buffer[j] = scratch[j]; - } - } - - coord = (coord + 1)%KD_DIMEN; - root->left = kd_build_tree_recursive(buffers, left_size, coord); - root->right = kd_build_tree_recursive(right_buffers, right_size, coord); - - return root; -} - -static kd_node_t * -kd_build_tree(kd_node_t **buffers[KD_BUFSIZE], size_t size) -{ - for (int i = 1; i < KD_DIMEN; ++i) { - memcpy(buffers[i], buffers[0], size*sizeof(kd_node_t *)); - } - for (int i = 0; i < KD_DIMEN; ++i) { - qsort(buffers[i], size, sizeof(kd_node_t *), kd_comparators[i]); - } - return kd_build_tree_recursive(buffers, size, 0); -} - -static double -kd_distance_sq(const double a[KD_DIMEN], const double b[KD_DIMEN]) -{ - double result = 0.0; - for (int i = 0; i < KD_DIMEN; ++i) { - double d = a[i] - b[i]; - result += d*d; - } - return result; -} - -static void -kd_find_nearest_recursive(kd_node_t *node, const double target[KD_DIMEN], const double closest[KD_DIMEN], kd_node_t **best, double *limit, unsigned int coord) -{ - if (!node->removed) { - double node_dist_sq = kd_distance_sq(node->coords, target); - if (node_dist_sq < *limit) { - *limit = node_dist_sq; - *best = node; - } - } - - kd_node_t *first; - kd_node_t *second; - if (target[coord] < node->coords[coord]) { - first = node->left; - second = node->right; - } else { - first = node->right; - second = node->left; - } - - unsigned int next = (coord + 1)%KD_DIMEN; - - if (first) { - kd_find_nearest_recursive(first, target, closest, best, limit, next); - } - - if (second) { - double new_closest[KD_DIMEN]; - memcpy(new_closest, closest, sizeof(new_closest)); - new_closest[coord] = node->coords[coord]; - - if (kd_distance_sq(new_closest, target) < *limit) { - kd_find_nearest_recursive(second, target, new_closest, best, limit, next); - } - } -} - -static void -kd_find_nearest(kd_node_t *root, const double target[KD_DIMEN], kd_node_t **best, double *limit) -{ - kd_find_nearest_recursive(root, target, target, best, limit, 0); -} - -void -kdf_init(kd_forest_t *kdf) -{ - kdf->roots = NULL; - kdf->size = kdf->size_est = 0; - kdf->roots_size = kdf->roots_capacity = 0; -} - -void -kdf_destroy(kd_forest_t *kdf) -{ - for (unsigned int i = 0; i < kdf->roots_size; ++i) { - kd_destroy(kdf->roots[i]); - } - - free(kdf->roots); -} - -static size_t -kdf_collect_nodes(kd_forest_t *kdf, kd_node_t **buffer, unsigned int slot, bool include_removed) -{ - size_t count = 0; - for (unsigned int i = 0; i < slot; ++i) { - if (kdf->roots[i]) { - count += kd_collect_nodes(kdf->roots[i], buffer + count, include_removed); - } - } - return count; -} - -static void -kdf_balance(kd_forest_t *kdf, kd_node_t *new_node, bool force) -{ - ++kdf->size; - - size_t slot, buffer_size; - if (force) { - buffer_size = kdf->size_est = kdf->size; - slot = kdf->roots_size; - } else { - ++kdf->size_est; - for (slot = 0; slot < kdf->roots_size; ++slot) { - if (!kdf->roots[slot]) { - break; - } - } - buffer_size = 1 << slot; - } - - kd_node_t **buffer = xmalloc(buffer_size*sizeof(kd_node_t *)); - buffer[0] = new_node; - kdf_collect_nodes(kdf, buffer + 1, slot, !force); - - kd_node_t **buffers[KD_BUFSIZE]; - for (int i = 1; i < KD_BUFSIZE; ++i) { - buffers[i] = xmalloc(buffer_size*sizeof(kd_node_t *)); - } - - if (slot >= kdf->roots_capacity) { - kdf->roots_capacity = slot + 1; - kdf->roots = xrealloc(kdf->roots, kdf->roots_capacity*sizeof(kd_node_t *)); - } - - size_t i, offset; - for (i = 0, offset = 0; offset < buffer_size; ++i) { - size_t chunk_size = 1 << i; - if (buffer_size & chunk_size) { - buffers[0] = buffer + offset; - kdf->roots[i] = kd_build_tree(buffers, chunk_size); - offset |= chunk_size; - } else { - kdf->roots[i] = NULL; - } - } - if (force || i > kdf->roots_size) { - kdf->roots_size = i; - } - - free(buffer); - for (i = 1; i < KD_BUFSIZE; ++i) { - free(buffers[i]); - } -} - -void -kdf_insert(kd_forest_t *kdf, kd_node_t *node) -{ - // If half or more of the nodes are removed, force a complete rebalance - bool force = (kdf->size_est + 1) >= 2*(kdf->size + 1); - kdf_balance(kdf, node, force); -} - -void -kdf_remove(kd_forest_t *kdf, kd_node_t *node) -{ - --kdf->size; - node->removed = true; -} - -kd_node_t * -kdf_find_nearest(const kd_forest_t *kdf, const double target[KD_DIMEN]) -{ - double limit = INFINITY; - kd_node_t *best = NULL; - - for (unsigned int i = 0; i < kdf->roots_size; ++i) { - kd_node_t *root = kdf->roots[i]; - if (root != NULL) { - kd_find_nearest(root, target, &best, &limit); - } - } - - return best; -} diff --git a/kd-forest.h b/kd-forest.h deleted file mode 100644 index 3651bfe..0000000 --- a/kd-forest.h +++ /dev/null @@ -1,56 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2014 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#ifndef KD_FOREST_H -#define KD_FOREST_H - -#include -#include -#include - -#define KD_DIMEN 3 - -// Single node in a k-d tree -typedef struct kd_node_t { - // Node coordinates - double coords[KD_DIMEN]; - // Sub-trees - struct kd_node_t *left, *right; - // Used to keep track of which sub-tree a node is in during construction - bool is_left; - // Weak delete support - bool removed; - - // Corresponding image position for this node - unsigned int x, y; -} kd_node_t; - -kd_node_t *new_kd_node(double coords[KD_DIMEN], unsigned int x, unsigned int y); - -// A forest of k-d trees -typedef struct { - // Array of k-d tree roots - kd_node_t **roots; - // Size and capacity of the roots array - unsigned int roots_size, roots_capacity; - // The actual size of this tree - size_t size; - // The size estimate for this tree, counting removed nodes - size_t size_est; -} kd_forest_t; - -void kdf_init(kd_forest_t *kdf); -void kdf_destroy(kd_forest_t *kdf); -void kdf_insert(kd_forest_t *kdf, kd_node_t *node); -void kdf_remove(kd_forest_t *kdf, kd_node_t *node); -kd_node_t *kdf_find_nearest(const kd_forest_t *kdf, const double target[KD_DIMEN]); - -#endif // KD_FOREST_H diff --git a/main.c b/main.c deleted file mode 100644 index 109f9f6..0000000 --- a/main.c +++ /dev/null @@ -1,480 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2014 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#include "color.h" -#include "generate.h" -#include "kd-forest.h" -#include "options.h" -#include "util.h" -#include -#include -#include -#include -#include -#include -#include -#include -#if __unix__ -# include -#endif - -// A single pixel in all its glory -typedef struct { - double value[KD_DIMEN]; - kd_node_t *node; - unsigned int x, y; - bool filled; -} pixel_t; - -// All-encompasing state struct -typedef struct { - const options_t *options; - uint32_t *colors; - pixel_t *pixels; - png_byte **bitmap; - unsigned int iteration; - unsigned int update_interval; - unsigned int frame; - size_t max_boundary; -} state_t; - -static void init_state(state_t *state, const options_t *options); -static void generate_image(state_t *state); -static void destroy_state(state_t *state); - -// Entry point -int -main(int argc, char *argv[]) -{ - options_t options; - if (!parse_options(&options, argc, argv)) { - fprintf(stderr, "\n"); - print_usage(stderr, argv[0], options.help); - return EXIT_FAILURE; - } - - if (options.help) { - print_usage(stdout, argv[0], true); - return EXIT_SUCCESS; - } - - state_t state; - init_state(&state, &options); - generate_image(&state); - destroy_state(&state); - return EXIT_SUCCESS; -} - -static pixel_t * -create_pixels(const options_t *options) -{ - pixel_t *pixels = xmalloc(options->npixels*sizeof(pixel_t)); - for (unsigned int y = 0, i = 0; y < options->height; ++y) { - for (unsigned int x = 0; x < options->width; ++x, ++i) { - pixel_t *pixel = pixels + i; - pixel->node = NULL; - pixel->x = x; - pixel->y = y; - pixel->filled = false; - } - } - return pixels; -} - -static png_byte ** -create_bitmap(const options_t *options) -{ - png_byte **rows = xmalloc(options->height*sizeof(png_byte *)); - const size_t row_size = 4*options->width*sizeof(png_byte); - for (unsigned int i = 0; i < options->height; ++i) { - rows[i] = xmalloc(row_size); - memset(rows[i], 0, row_size); - } - return rows; -} - -static void -init_state(state_t *state, const options_t *options) -{ - printf("Generating a %u-bit, %ux%u image (%zu pixels)\n", - options->bit_depth, options->width, options->height, options->npixels); - - xsrand(options->seed); - - state->options = options; - state->colors = generate_colors(options); - state->pixels = create_pixels(options); - state->bitmap = create_bitmap(options); - state->iteration = 0; - state->update_interval = 1U << (state->options->bit_depth + 1)/2; - state->frame = 0; - state->max_boundary = 0; -} - -static void generate_bitmap(state_t *state); -static void write_png(const state_t *state, const char *filename); - -static void -generate_image(state_t *state) -{ - generate_bitmap(state); - - if (!state->options->animate) { - write_png(state, state->options->filename); - } -} - -static void -destroy_state(state_t *state) -{ - for (unsigned int i = 0; i < state->options->height; ++i) { - free(state->bitmap[i]); - } - free(state->bitmap); - free(state->pixels); - free(state->colors); -} - -static pixel_t * -get_pixel(const state_t *state, unsigned int x, unsigned int y) -{ - return state->pixels + state->options->width*y + x; -} - -static pixel_t * -try_neighbor(const state_t *state, pixel_t *pixel, int dx, int dy) -{ - if (dx < 0 && pixel->x < -dx) { - return NULL; - } else if (dx > 0 && pixel->x + dx >= state->options->width) { - return NULL; - } else if (dy < 0 && pixel->y < -dy) { - return NULL; - } else if (dy > 0 && pixel->y + dy >= state->options->height) { - return NULL; - } - - return pixel + (int)state->options->width*dy + dx; -} - -static unsigned int -get_all_neighbors(const state_t *state, pixel_t *pixel, pixel_t *neighbors[8]) -{ - unsigned int size = 0; - - for (int dy = -1; dy <= 1; ++dy) { - for (int dx = -1; dx <= 1; ++dx) { - if (dx == 0 && dy == 0) { - continue; - } - - pixel_t *neighbor = try_neighbor(state, pixel, dx, dy); - if (neighbor) { - neighbors[size++] = neighbor; - } - } - } - - return size; -} - -static unsigned int -get_neighbors(const state_t *state, pixel_t *pixel, pixel_t *neighbors[8], bool filled) -{ - unsigned int size = 0; - - for (int dy = -1; dy <= 1; ++dy) { - for (int dx = -1; dx <= 1; ++dx) { - if (dx == 0 && dy == 0) { - continue; - } - - pixel_t *neighbor = try_neighbor(state, pixel, dx, dy); - if (neighbor && neighbor->filled == filled) { - neighbors[size++] = neighbor; - } - } - } - - return size; -} - -static pixel_t * -select_empty_neighbor(const state_t *state, pixel_t *pixel) -{ - pixel_t *neighbors[8]; - unsigned int size = get_neighbors(state, pixel, neighbors, false); - return neighbors[xrand(size)]; -} - -static pixel_t * -find_next_pixel(const state_t *state, const kd_forest_t *kdf, const double target[KD_DIMEN]) -{ - kd_node_t *nearest = kdf_find_nearest(kdf, target); - pixel_t *pixel = get_pixel(state, nearest->x, nearest->y); - - if (state->options->selection == SELECTION_MIN) { - pixel = select_empty_neighbor(state, pixel); - } - - return pixel; -} - -static void -ensure_pixel_removed(kd_forest_t *kdf, pixel_t *pixel) -{ - if (pixel->node) { - kdf_remove(kdf, pixel->node); - pixel->node = NULL; - } -} - -static bool -has_empty_neighbors(const state_t *state, pixel_t *pixel) -{ - for (int dy = -1; dy <= 1; ++dy) { - for (int dx = -1; dx <= 1; ++dx) { - if (dx == 0 && dy == 0) { - continue; - } - - pixel_t *neighbor = try_neighbor(state, pixel, dx, dy); - if (neighbor && !neighbor->filled) { - return true; - } - } - } - - return false; -} - -static void -insert_new_pixel_min(state_t *state, kd_forest_t *kdf, pixel_t *pixel) -{ - pixel->filled = true; - - if (has_empty_neighbors(state, pixel)) { - pixel->node = new_kd_node(pixel->value, pixel->x, pixel->y); - kdf_insert(kdf, pixel->node); - } - - pixel_t *neighbors[8]; - unsigned int size = get_all_neighbors(state, pixel, neighbors); - for (unsigned int i = 0; i < size; ++i) { - pixel_t *neighbor = neighbors[i]; - if (!has_empty_neighbors(state, neighbor)) { - ensure_pixel_removed(kdf, neighbor); - } - } -} - -static void -insert_new_pixel_mean(state_t *state, kd_forest_t *kdf, pixel_t *pixel) -{ - pixel->filled = true; - ensure_pixel_removed(kdf, pixel); - - pixel_t *neighbors[8]; - unsigned int size = get_neighbors(state, pixel, neighbors, false); - for (unsigned int i = 0; i < size; ++i) { - pixel_t *neighbor = neighbors[i]; - - double value[KD_DIMEN] = { 0.0 }; - - pixel_t *filled[8]; - unsigned int nfilled = get_neighbors(state, neighbor, filled, true); - for (unsigned int j = 0; j < nfilled; ++j) { - for (unsigned int k = 0; k < KD_DIMEN; ++k) { - value[k] += filled[j]->value[k]; - } - } - - for (unsigned int j = 0; j < KD_DIMEN; ++j) { - value[j] /= nfilled; - } - - ensure_pixel_removed(kdf, neighbor); - neighbor->node = new_kd_node(value, neighbor->x, neighbor->y); - kdf_insert(kdf, neighbor->node); - } -} - -static void -insert_new_pixel(state_t *state, kd_forest_t *kdf, pixel_t *pixel) -{ - switch (state->options->selection) { - case SELECTION_MIN: - insert_new_pixel_min(state, kdf, pixel); - break; - - case SELECTION_MEAN: - insert_new_pixel_mean(state, kdf, pixel); - break; - } -} - -static void -print_progress(const char *format, ...) -{ -#if __unix__ - static bool tty_checked = false; - static bool tty = false; - if (!tty_checked) { - tty = isatty(STDOUT_FILENO); - tty_checked = true; - } - const char *clear_line = tty ? "\033[2K\r" : ""; - const char *new_line = tty ? "" : "\n"; -#else - const char *clear_line = ""; - const char *new_line = "\n"; -#endif - - printf("%s", clear_line); - - va_list args; - va_start(args, format); - vprintf(format, args); - va_end(args); - - printf("%s", new_line); - fflush(stdout); -} - -static void -place_color(state_t *state, kd_forest_t *kdf, unsigned int i) -{ - if (state->iteration % state->update_interval == 0) { - if (state->options->animate) { - char filename[strlen(state->options->filename) + 10]; - sprintf(filename, "%s/%04u.png", state->options->filename, state->frame); - write_png(state, filename); - ++state->frame; - } - - print_progress("%.2f%%\t| boundary size: %zu\t| max boundary size: %zu", - 100.0*state->iteration/state->options->ncolors, kdf->size, state->max_boundary); - } - - uint32_t color = state->colors[i]; - - double target[KD_DIMEN]; - switch (state->options->color_space) { - case COLOR_SPACE_RGB: - color_set_RGB(target, color); - break; - case COLOR_SPACE_LAB: - color_set_Lab(target, color); - break; - case COLOR_SPACE_LUV: - color_set_Luv(target, color); - break; - } - - pixel_t *pixel; - if (state->iteration == 0) { - pixel = get_pixel(state, state->options->x, state->options->y); - } else { - pixel = find_next_pixel(state, kdf, target); - } - - memcpy(pixel->value, target, sizeof(target)); - insert_new_pixel(state, kdf, pixel); - if (kdf->size > state->max_boundary) { - state->max_boundary = kdf->size; - } - - png_byte *png_pixel = state->bitmap[pixel->y] + 4*pixel->x; - color_unpack(png_pixel, color); - png_pixel[3] = 0xFF; -} - -static void -generate_bitmap(state_t *state) -{ - // Make the forest - kd_forest_t kdf; - kdf_init(&kdf); - - if (state->options->stripe) { - for (unsigned int i = 1; i <= state->options->bit_depth + 1; ++i) { - unsigned int stripe = 1 << i; - for (unsigned int j = stripe/2 - 1; j < state->options->ncolors; j += stripe, ++state->iteration) { - place_color(state, &kdf, j); - } - } - } else { - for (unsigned int i = 0; i < state->options->ncolors; ++i, ++state->iteration) { - place_color(state, &kdf, i); - } - } - - if (state->options->animate) { - char filename[strlen(state->options->filename) + 10]; - -#if __unix__ - sprintf(filename, "%s/last.png", state->options->filename); - write_png(state, filename); - - for (int i = 0; i < 120; ++i) { - sprintf(filename, "%s/%04u.png", state->options->filename, state->frame + i); - if (symlink("last.png", filename) != 0) { - abort(); - } - } -#else - for (int i = 0; i < 120; ++i) { - sprintf(filename, "%s/%04u.png", state->options->filename, state->frame + i); - write_png(state, filename); - } -#endif - } - - print_progress("%.2f%%\t| boundary size: %zu\t| max boundary size: %zu\n", - 100.0, kdf.size, state->max_boundary); - - kdf_destroy(&kdf); -} - -static void -write_png(const state_t *state, const char *filename) -{ - FILE *file = fopen(filename, "wb"); - if (!file) { - abort(); - } - - png_struct *png_ptr = - png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL); - if (!png_ptr) { - abort(); - } - - png_info *info_ptr = png_create_info_struct(png_ptr); - if (!info_ptr) { - abort(); - } - - // libpng will longjmp here if it encounters an error from now on - if (setjmp(png_jmpbuf(png_ptr))) { - abort(); - } - - png_init_io(png_ptr, file); - png_set_IHDR(png_ptr, info_ptr, state->options->width, state->options->height, 8, - PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_ADAM7, - PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT); - png_set_sRGB_gAMA_and_cHRM(png_ptr, info_ptr, PNG_sRGB_INTENT_ABSOLUTE); - png_set_rows(png_ptr, info_ptr, state->bitmap); - png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL); - png_destroy_write_struct(&png_ptr, &info_ptr); - fclose(file); -} diff --git a/options.c b/options.c deleted file mode 100644 index e25930a..0000000 --- a/options.c +++ /dev/null @@ -1,431 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2014 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#include "options.h" -#include -#include -#include -#include -#include -#if __unix__ -# include -#endif - -static bool -parse_arg(int argc, char *argv[], - const char *short_form, const char *long_form, - const char **value, int *i, bool *error) -{ - size_t short_len = short_form ? strlen(short_form) : 0; - size_t long_len = long_form ? strlen(long_form) : 0; - - const char *actual_form; - const char *arg = argv[*i]; - const char *candidate = NULL; - - if (short_form && strncmp(arg, short_form, short_len) == 0) { - actual_form = short_form; - if (strlen(arg) > short_len) { - candidate = arg + short_len; - } - } else if (long_form && strncmp(arg, long_form, long_len) == 0) { - actual_form = long_form; - if (strlen(arg) > long_len) { - if (arg[long_len] == '=') { - candidate = arg + long_len + 1; - } else { - return false; - } - } - } else { - return false; - } - - if (value) { - if (candidate) { - *value = candidate; - } else if (*i < argc - 1) { - ++*i; - *value = argv[*i]; - } else { - fprintf(stderr, "Expected a value for %s\n", arg); - *error = true; - return false; - } - } else if (candidate) { - fprintf(stderr, "Unexpected value for %s: `%s'\n", - actual_form, candidate); - *error = true; - return false; - } - - return true; -} - -static bool -str_to_uint(const char *str, unsigned int *value) -{ - char *endptr; - long result = strtol(str, &endptr, 10); - if (*str == '\0' || *endptr != '\0') { - return false; - } - if (result < 0 || result > UINT_MAX) { - return false; - } - - *value = result; - return true; -} - -static void -strcatinc(char **destp, const char *src) -{ - strcpy(*destp, src); - *destp += strlen(src); -} - -typedef enum { - COLORIZE_NORMAL, - COLORIZE_AT, - COLORIZE_BANG, - COLORIZE_STAR, - COLORIZE_SHORT_OPTION, - COLORIZE_LONG_OPTION, -} colorize_state_t; - -static void -print_colorized(FILE *file, bool tty, const char *format, ...) -{ - const char *bold = tty ? "\033[1m" : ""; - const char *red = tty ? "\033[1;31m" : ""; - const char *green = tty ? "\033[1;32m" : ""; - const char *normal = tty ? "\033[0m" : ""; - - size_t size = strlen(format) + 1; - char colorized[16*size]; - char *builder = colorized; - - colorize_state_t state = COLORIZE_NORMAL; - for (size_t i = 0; i < size; ++i) { - char c = format[i]; - - if (c == '\\') { - *builder++ = format[++i]; - continue; - } - - switch (state) { - case COLORIZE_AT: - if (c == '@') { - strcatinc(&builder, normal); - state = COLORIZE_NORMAL; - } else { - *builder++ = c; - } - break; - - case COLORIZE_BANG: - if (c == '!') { - strcatinc(&builder, normal); - state = COLORIZE_NORMAL; - } else { - *builder++ = c; - } - break; - - case COLORIZE_STAR: - if (c == '*') { - strcatinc(&builder, normal); - state = COLORIZE_NORMAL; - } else { - *builder++ = c; - } - break; - - case COLORIZE_SHORT_OPTION: - *builder++ = c; - strcatinc(&builder, normal); - state = COLORIZE_NORMAL; - break; - - case COLORIZE_LONG_OPTION: - if (!isalpha(c) && c != '-') { - strcatinc(&builder, normal); - state = COLORIZE_NORMAL; - } - *builder++ = c; - break; - - default: - switch (c) { - case '@': - state = COLORIZE_AT; - strcatinc(&builder, green); - break; - - case '!': - state = COLORIZE_BANG; - strcatinc(&builder, bold); - break; - - case '*': - state = COLORIZE_STAR; - strcatinc(&builder, red); - break; - - case '-': - if (c == '-') { - if (format[i + 1] == '-') { - state = COLORIZE_LONG_OPTION; - } else { - state = COLORIZE_SHORT_OPTION; - } - strcatinc(&builder, red); - } - *builder++ = c; - break; - - default: - *builder++ = c; - break; - } - break; - } - } - - va_list args; - va_start(args, format); - vprintf(colorized, args); - va_end(args); -} - -void -print_usage(FILE *file, const char *command, bool verbose) -{ -#if __unix__ - bool tty = isatty(fileno(file)); -#else - bool tty = false; -#endif - - size_t length = strlen(command); - char whitespace[length + 1]; - memset(whitespace, ' ', length); - whitespace[length] = '\0'; - -#define usage(...) print_colorized(file, tty, __VA_ARGS__) - usage("Usage:\n"); - usage(" !$! *%s* [-b|--bit-depth @DEPTH@]\n", command); - usage(" %s [-s|--hue-sort] [-r|--random] [-M|--morton] [-H|--hilbert]\n", whitespace); - usage(" %s [-t|--stripe] [-T|--no-stripe]\n", whitespace); - usage(" %s [-l|--selection @min@|@mean@]\n", whitespace); - usage(" %s [-c|--color-space @RGB@|@Lab@|@Luv@]\n", whitespace); - usage(" %s [-w|--width @WIDTH@] [-h|--height @HEIGHT@]\n", whitespace); - usage(" %s [-x @X@] [-y @Y@]\n", whitespace); - usage(" %s [-a|--animate]\n", whitespace); - usage(" %s [-o|--output @PATH@]\n", whitespace); - usage(" %s [-e|--seed @SEED@]\n", whitespace); - usage(" %s [-?|--help]\n", whitespace); - - if (!verbose) { - return; - } - - usage("\n"); - usage(" -b, --bit-depth @DEPTH@:\n"); - usage(" Use all @DEPTH@\\-bit colors (!default!: @24@)\n\n"); - usage(" -s, --hue-sort:\n"); - usage(" Sort colors by hue (!default!)\n"); - usage(" -r, --random:\n"); - usage(" Randomize colors\n"); - usage(" -M, --morton:\n"); - usage(" Place colors in Morton order (Z\\-order)\n"); - usage(" -H, --hilbert:\n"); - usage(" Place colors in Hilbert curve order\n\n"); - usage(" -t, --stripe:\n"); - usage(" -T, --no-stripe:\n"); - usage(" Whether to reduce artifacts by iterating through the colors in\n"); - usage(" multiple stripes (!default!: --stripe)\n\n"); - usage(" -l, --selection @min@|@mean@:\n"); - usage(" Specify the selection mode (!default!: @min@)\n\n"); - usage(" @min@: Pick the pixel with the closest neighboring pixel\n"); - usage(" @mean@: Pick the pixel with the closest average of all its neighbors\n\n"); - usage(" -c, --color-space @RGB@|@Lab@|@Luv@:\n"); - usage(" Use the given color space (!default!: @Lab@)\n\n"); - usage(" -w, --width @WIDTH@\n"); - usage(" -h, --height @HEIGHT@:\n"); - usage(" Generate a @WIDTH@x@HEIGHT@ image (!default!: @as small as possible@)\n\n"); - usage(" -x @X@\n"); - usage(" -y @Y@:\n"); - usage(" Place the first pixel at (@X@, @Y@) (!default!: @center@)\n\n"); - usage(" -a, --animate:\n"); - usage(" Generate frames of an animation\n\n"); - usage(" -o, --output @PATH@:\n"); - usage(" Output a PNG file at @PATH@ (!default!: @kd\\-forest.png@)\n\n"); - usage(" If -a/--animate is specified, this is treated as a directory which\n"); - usage(" will hold many frames\n\n"); - usage(" -e, --seed @SEED@:\n"); - usage(" Seed the random number generator (!default!: @0@)\n\n"); - usage(" -?, --help:\n"); - usage(" Show this message\n"); -#undef usage -} - -bool -parse_options(options_t *options, int argc, char *argv[]) -{ - // Set defaults - options->bit_depth = 24; - options->mode = MODE_HUE_SORT; - options->stripe = true; - options->selection = SELECTION_MIN; - options->color_space = COLOR_SPACE_LAB; - options->animate = false; - options->filename = NULL; - options->seed = 0; - options->help = false; - - bool width_set = false, height_set = false; - bool x_set = false, y_set = false; - bool result = true; - - for (int i = 1; i < argc; ++i) { - const char *value; - bool error = false; - - if (parse_arg(argc, argv, "-b", "--bit-depth", &value, &i, &error)) { - if (!str_to_uint(value, &options->bit_depth) - || options->bit_depth <= 1 - || options->bit_depth > 24) { - fprintf(stderr, "Invalid bit depth: `%s'\n", value); - error = true; - } - } else if (parse_arg(argc, argv, "-s", "--hue-sort", NULL, &i, &error)) { - options->mode = MODE_HUE_SORT; - } else if (parse_arg(argc, argv, "-r", "--random", NULL, &i, &error)) { - options->mode = MODE_RANDOM; - } else if (parse_arg(argc, argv, "-M", "--morton", NULL, &i, &error)) { - options->mode = MODE_MORTON; - } else if (parse_arg(argc, argv, "-H", "--hilbert", NULL, &i, &error)) { - options->mode = MODE_HILBERT; - } else if (parse_arg(argc, argv, "-t", "--stripe", NULL, &i, &error)) { - options->stripe = true; - } else if (parse_arg(argc, argv, "-T", "--no-stripe", NULL, &i, &error)) { - options->stripe = false; - } else if (parse_arg(argc, argv, "-l", "--selection", &value, &i, &error)) { - if (strcmp(value, "min") == 0) { - options->selection = SELECTION_MIN; - } else if (strcmp(value, "mean") == 0) { - options->selection = SELECTION_MEAN; - } else { - fprintf(stderr, "Invalid selection mode: `%s'\n", value); - error = true; - } - } else if (parse_arg(argc, argv, "-c", "--color-space", &value, &i, &error)) { - if (strcmp(value, "RGB") == 0) { - options->color_space = COLOR_SPACE_RGB; - } else if (strcmp(value, "Lab") == 0) { - options->color_space = COLOR_SPACE_LAB; - } else if (strcmp(value, "Luv") == 0) { - options->color_space = COLOR_SPACE_LUV; - } else { - fprintf(stderr, "Invalid color space: `%s'\n", value); - error = true; - } - } else if (parse_arg(argc, argv, "-w", "--width", &value, &i, &error)) { - if (str_to_uint(value, &options->width)) { - width_set = true; - } else { - fprintf(stderr, "Invalid width: `%s'\n", value); - error = true; - } - } else if (parse_arg(argc, argv, "-h", "--height", &value, &i, &error)) { - if (str_to_uint(value, &options->height)) { - height_set = true; - } else { - fprintf(stderr, "Invalid height: `%s'\n", value); - error = true; - } - } else if (parse_arg(argc, argv, "-x", NULL, &value, &i, &error)) { - if (str_to_uint(value, &options->x)) { - x_set = true; - } else { - fprintf(stderr, "Invalid x coordinate: `%s'\n", value); - error = true; - } - } else if (parse_arg(argc, argv, "-y", NULL, &value, &i, &error)) { - if (str_to_uint(value, &options->y)) { - y_set = true; - } else { - fprintf(stderr, "Invalid y coordinate: `%s'\n", value); - error = true; - } - } else if (parse_arg(argc, argv, "-a", "--animate", NULL, &i, &error)) { - options->animate = true; - } else if (parse_arg(argc, argv, "-o", "--output", &value, &i, &error)) { - options->filename = value; - } else if (parse_arg(argc, argv, "-e", "--seed", &value, &i, &error)) { - if (!str_to_uint(value, &options->seed)) { - fprintf(stderr, "Invalid random seed: `%s'\n", value); - error = true; - } - } else if (parse_arg(argc, argv, "-?", "--help", NULL, &i, &error)) { - options->help = true; - } else if (!error) { - fprintf(stderr, "Unexpected argument `%s'\n", argv[i]); - error = true; - } - - if (error) { - result = false; - } - } - - options->ncolors = (size_t)1 << options->bit_depth; - - if (!width_set && !height_set) { - // Round width up to make widescreen the default - options->width = 1U << (options->bit_depth + 1)/2; - options->height = 1U << options->bit_depth/2; - } else if (width_set && !height_set) { - options->height = (options->ncolors + options->width - 1)/options->width; - } else if (!width_set && height_set) { - options->width = (options->ncolors + options->height - 1)/options->height; - } - - options->npixels = (size_t)options->width*options->height; - if (options->npixels < options->ncolors) { - fprintf(stderr, "Image too small (at least %zu pixels needed)\n", options->ncolors); - result = false; - } - - if (!x_set) { - options->x = options->width/2; - } else if (options->x >= options->width) { - fprintf(stderr, "-x coordinate too big, should be less than %u\n", options->width); - result = false; - } - - if (!y_set) { - options->y = options->height/2; - } else if (options->y >= options->height) { - fprintf(stderr, "-y coordinate too big, should be less than %u\n", options->height); - result = false; - } - - // Default filename depends on -a flag - if (!options->filename) { - options->filename = options->animate ? "frames" : "kd-forest.png"; - } - - return result; -} diff --git a/options.h b/options.h deleted file mode 100644 index 07ecc45..0000000 --- a/options.h +++ /dev/null @@ -1,58 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2014 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#ifndef OPTIONS_H -#define OPTIONS_H - -#include -#include - -// Possible generation modes -typedef enum { - MODE_HUE_SORT, - MODE_RANDOM, - MODE_MORTON, - MODE_HILBERT, -} mode_t; - -// Possible pixel selection modes -typedef enum { - SELECTION_MIN, - SELECTION_MEAN, -} selection_t; - -// Possible color spaces -typedef enum { - COLOR_SPACE_RGB, - COLOR_SPACE_LAB, - COLOR_SPACE_LUV, -} color_space_t; - -// Command-line options -typedef struct { - unsigned int bit_depth; - mode_t mode; - bool stripe; - selection_t selection; - color_space_t color_space; - unsigned int width, height; - size_t ncolors, npixels; - unsigned int x, y; - bool animate; - const char *filename; - unsigned int seed; - bool help; -} options_t; - -bool parse_options(options_t *options, int argc, char *argv[]); -void print_usage(FILE *file, const char *command, bool verbose); - -#endif // OPTIONS_H diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..f328e4d --- /dev/null +++ b/src/main.rs @@ -0,0 +1 @@ +fn main() {} diff --git a/util.c b/util.c deleted file mode 100644 index a2573a4..0000000 --- a/util.c +++ /dev/null @@ -1,66 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2014 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#include "util.h" -#include - -void * -xmalloc(size_t size) -{ - void *ret = malloc(size); - if (!ret) { - abort(); - } - return ret; -} - -void * -xrealloc(void *ptr, size_t size) -{ - void *ret = realloc(ptr, size); - if (!ret) { - abort(); - } - return ret; -} - -// Based on sample rand() implementation from POSIX.1-2001 - -static unsigned long xrand_next = 0; - -void xsrand(unsigned int seed) { - xrand_next = seed; -} - -static unsigned int xrand_simple(void) { - xrand_next = xrand_next*1103515245 + 12345; - return (unsigned int)(xrand_next/65536)%32768; -} - -static unsigned int xrand_full(void) { - unsigned int low = xrand_simple(); - unsigned int high = xrand_simple(); - return low | (high << 15); -} - -#define XRAND_RANGE 1073741824U - -unsigned int -xrand(unsigned int range) -{ - // Compensate for bias if XRAND_RANGE isn't a multiple of range - unsigned int limit = XRAND_RANGE - XRAND_RANGE%range; - unsigned int res; - do { - res = xrand_full(); - } while (res >= limit); - return res%range; -} diff --git a/util.h b/util.h deleted file mode 100644 index 9cb5013..0000000 --- a/util.h +++ /dev/null @@ -1,23 +0,0 @@ -/********************************************************************* - * kd-forest * - * Copyright (C) 2014 Tavian Barnes * - * * - * This program is free software. It comes without any warranty, to * - * the extent permitted by applicable law. You can redistribute it * - * and/or modify it under the terms of the Do What The Fuck You Want * - * To Public License, Version 2, as published by Sam Hocevar. See * - * the COPYING file or http://www.wtfpl.net/ for more details. * - *********************************************************************/ - -#ifndef UTIL_H -#define UTIL_H - -#include - -void *xmalloc(size_t size); -void *xrealloc(void *ptr, size_t size); - -unsigned int xrand(unsigned int range); -void xsrand(unsigned int seed); - -#endif // UTIL_H -- cgit v1.2.3 From 8aa2911b4c978ad7131a430d07a580cedf6f8f65 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sun, 19 Apr 2020 16:37:16 -0400 Subject: metric: Add some general interfaces for metric spaces --- Cargo.lock | 104 ++++++++++++ Cargo.toml | 4 + src/main.rs | 2 + src/metric.rs | 531 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 641 insertions(+) create mode 100644 src/metric.rs diff --git a/Cargo.lock b/Cargo.lock index a63d677..45602e1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,5 +1,109 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. +[[package]] +name = "autocfg" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8aac770f1885fd7e387acedd76065302551364496e46b3dd00860b2f8359b9d" + +[[package]] +name = "cfg-if" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4785bdd1c96b2a846b2bd7cc02e86b6b3dbf14e7e53446c4f54c92a361040822" + +[[package]] +name = "getrandom" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7abc8dd8451921606d809ba32e95b6111925cd2906060d2dcc29c070220503eb" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + [[package]] name = "kd-forest" version = "2.0.0" +dependencies = [ + "ordered-float", + "rand", +] + +[[package]] +name = "libc" +version = "0.2.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "99e85c08494b21a9054e7fe1374a732aeadaff3980b6990b94bfd3a70f690005" + +[[package]] +name = "num-traits" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c62be47e61d1842b9170f0fdeec8eba98e60e90e5446449a0545e5152acd7096" +dependencies = [ + "autocfg", +] + +[[package]] +name = "ordered-float" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18869315e81473c951eb56ad5558bbc56978562d3ecfb87abb7a1e944cea4518" +dependencies = [ + "num-traits", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74490b50b9fbe561ac330df47c08f3f33073d2d00c150f719147d7c54522fa1b" + +[[package]] +name = "rand" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a6b1679d49b24bbfe0c803429aa1874472f50d9b363131f0e89fc356b544d03" +dependencies = [ + "getrandom", + "libc", + "rand_chacha", + "rand_core", + "rand_hc", +] + +[[package]] +name = "rand_chacha" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4c8ed856279c9737206bf725bf36935d8666ead7aa69b52be55af369d193402" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19" +dependencies = [ + "getrandom", +] + +[[package]] +name = "rand_hc" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca3129af7b92a17112d59ad498c6f81eaf463253766b90396d39ea7a39d6613c" +dependencies = [ + "rand_core", +] + +[[package]] +name = "wasi" +version = "0.9.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519" diff --git a/Cargo.toml b/Cargo.toml index 65ffec8..5a93cf5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,3 +3,7 @@ name = "kd-forest" version = "2.0.0" authors = ["Tavian Barnes "] edition = "2018" + +[dependencies] +ordered-float = "1.0.2" +rand = "0.7.3" diff --git a/src/main.rs b/src/main.rs index f328e4d..0d7989b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1 +1,3 @@ +pub mod metric; + fn main() {} diff --git a/src/metric.rs b/src/metric.rs new file mode 100644 index 0000000..e067771 --- /dev/null +++ b/src/metric.rs @@ -0,0 +1,531 @@ +//! [Metric spaces](https://en.wikipedia.org/wiki/Metric_space). + +use ordered_float::OrderedFloat; + +use std::cmp::Ordering; +use std::collections::BinaryHeap; +use std::iter::FromIterator; + +/// An [order embedding](https://en.wikipedia.org/wiki/Order_embedding) for distances. +/// +/// Implementations of this trait must satisfy, for all non-negative distances `x` and `y`: +/// +/// * `x == Self::from(x).into()` +/// * `x <= y` iff `Self::from(x) <= Self::from(y)` +/// +/// This trait exists to optimize the common case where distances can be compared more efficiently +/// than their exact values can be computed. For example, taking the square root can be avoided +/// when comparing Euclidean distances (see [SquaredDistance]). +pub trait Distance: Copy + From + Into + Ord {} + +/// A raw numerical distance. +#[derive(Debug, Clone, Copy, Eq, Ord, PartialEq, PartialOrd)] +pub struct RawDistance(OrderedFloat); + +impl From for RawDistance { + fn from(value: f64) -> Self { + Self(value.into()) + } +} + +impl From for f64 { + fn from(value: RawDistance) -> Self { + value.0.into_inner() + } +} + +impl Distance for RawDistance {} + +/// A squared distance, to avoid computing square roots unless absolutely necessary. +#[derive(Debug, Clone, Copy, Eq, Ord, PartialEq, PartialOrd)] +pub struct SquaredDistance(OrderedFloat); + +impl SquaredDistance { + /// Create a SquaredDistance from an already squared value. + pub fn from_squared(value: f64) -> Self { + Self(value.into()) + } +} + +impl From for SquaredDistance { + fn from(value: f64) -> Self { + Self::from_squared(value * value) + } +} + +impl From for f64 { + fn from(value: SquaredDistance) -> Self { + value.0.into_inner().sqrt() + } +} + +impl Distance for SquaredDistance {} + +/// A [metric space](https://en.wikipedia.org/wiki/Metric_space). +pub trait Metric { + /// The type used to represent distances. Use [RawDistance] to compare the actual values + /// directly, or another type if comparisons can be implemented more efficiently. + type Distance: Distance; + + /// Computes the distance between this point and another point. This function must satisfy + /// three conditions: + /// + /// * `x.distance(y) == 0` iff `x == y` (identity of indiscernibles) + /// * `x.distance(y) == y.distance(x)` (symmetry) + /// * `x.distance(z) <= x.distance(y) + y.distance(z)` (triangle inequality) + fn distance(&self, other: &T) -> Self::Distance; +} + +/// Blanket [Metric] implementation for references. +impl<'a, 'b, T, U: Metric> Metric<&'a T> for &'b U { + type Distance = U::Distance; + + fn distance(&self, other: &&'a T) -> Self::Distance { + (*self).distance(other) + } +} + +/// The standard [Euclidean distance](https://en.wikipedia.org/wiki/Euclidean_distance) metric. +impl Metric for [f64] { + type Distance = SquaredDistance; + + fn distance(&self, other: &Self) -> Self::Distance { + debug_assert!(self.len() == other.len()); + + let mut sum = 0.0; + for i in 0..self.len() { + let diff = self[i] - other[i]; + sum += diff * diff; + } + + Self::Distance::from_squared(sum) + } +} + +/// A nearest neighbor to a target. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Neighbor { + /// The found item. + pub item: T, + /// The distance from the target. + pub distance: f64, +} + +impl Neighbor { + /// Create a new Neighbor. + pub fn new(item: T, distance: f64) -> Self { + Self { item, distance } + } +} + +/// A candidate nearest neighbor found during a search. +#[derive(Debug)] +struct Candidate { + item: T, + distance: D, +} + +impl Candidate { + fn new(target: U, item: T) -> Self + where + U: Metric, + { + let distance = target.distance(&item); + Self { item, distance } + } + + fn into_neighbor(self) -> Neighbor { + Neighbor::new(self.item, self.distance.into()) + } +} + +impl PartialOrd for Candidate { + fn partial_cmp(&self, other: &Self) -> Option { + self.distance.partial_cmp(&other.distance) + } +} + +impl Ord for Candidate { + fn cmp(&self, other: &Self) -> Ordering { + self.distance.cmp(&other.distance) + } +} + +impl PartialEq for Candidate { + fn eq(&self, other: &Self) -> bool { + self.distance.eq(&other.distance) + } +} + +impl Eq for Candidate {} + +/// Accumulates nearest neighbor search results. +pub trait Neighborhood> { + /// Returns the target of the nearest neighbor search. + fn target(&self) -> U; + + /// Check whether a distance is within this neighborhood. + fn contains(&self, distance: f64) -> bool { + distance < 0.0 || self.contains_distance(distance.into()) + } + + /// Check whether a distance is within this neighborhood. + fn contains_distance(&self, distance: U::Distance) -> bool; + + /// Consider a new candidate neighbor. + fn consider(&mut self, item: T) -> U::Distance; +} + +/// A [Neighborhood] with at most one result. +#[derive(Debug)] +struct SingletonNeighborhood> { + /// The target of the nearest neighbor search. + target: U, + /// The current threshold distance to the farthest result. + threshold: Option, + /// The current nearest neighbor, if any. + candidate: Option>, +} + +impl SingletonNeighborhood +where + U: Copy + Metric, +{ + /// Create a new single metric result tracker. + /// + /// * `target`: The target fo the nearest neighbor search. + /// * `threshold`: The maximum allowable distance. + fn new(target: U, threshold: Option) -> Self { + Self { + target, + threshold: threshold.map(U::Distance::from), + candidate: None, + } + } + + /// Consider a candidate. + fn push(&mut self, candidate: Candidate) -> U::Distance { + let distance = candidate.distance; + + if self.contains_distance(distance) { + self.threshold = Some(distance); + self.candidate = Some(candidate); + } + + distance + } + + /// Convert this result into an optional neighbor. + fn into_option(self) -> Option> { + self.candidate.map(Candidate::into_neighbor) + } +} + +impl Neighborhood for SingletonNeighborhood +where + U: Copy + Metric, +{ + fn target(&self) -> U { + self.target + } + + fn contains_distance(&self, distance: U::Distance) -> bool { + self.threshold.map(|t| distance <= t).unwrap_or(true) + } + + fn consider(&mut self, item: T) -> U::Distance { + self.push(Candidate::new(self.target, item)) + } +} + +/// A [Neighborhood] of up to `k` results, using a binary heap. +#[derive(Debug)] +struct HeapNeighborhood> { + /// The target of the nearest neighbor search. + target: U, + /// The number of nearest neighbors to find. + k: usize, + /// The current threshold distance to the farthest result. + threshold: Option, + /// A max-heap of the best candidates found so far. + heap: BinaryHeap>, +} + +impl HeapNeighborhood +where + U: Copy + Metric, +{ + /// Create a new metric result tracker. + /// + /// * `target`: The target fo the nearest neighbor search. + /// * `k`: The number of nearest neighbors to find. + /// * `threshold`: The maximum allowable distance. + fn new(target: U, k: usize, threshold: Option) -> Self { + Self { + target, + k, + threshold: threshold.map(U::Distance::from), + heap: BinaryHeap::with_capacity(k), + } + } + + /// Consider a candidate. + fn push(&mut self, candidate: Candidate) -> U::Distance { + let distance = candidate.distance; + + if self.contains_distance(distance) { + let heap = &mut self.heap; + + if heap.len() == self.k { + heap.pop(); + } + + heap.push(candidate); + + if heap.len() == self.k { + self.threshold = self.heap.peek().map(|c| c.distance) + } + } + + distance + } + + /// Convert these results into a vector of neighbors. + fn into_vec(self) -> Vec> { + self.heap + .into_sorted_vec() + .into_iter() + .map(Candidate::into_neighbor) + .collect() + } +} + +impl Neighborhood for HeapNeighborhood +where + U: Copy + Metric, +{ + fn target(&self) -> U { + self.target + } + + fn contains_distance(&self, distance: U::Distance) -> bool { + self.k > 0 && self.threshold.map(|t| distance <= t).unwrap_or(true) + } + + fn consider(&mut self, item: T) -> U::Distance { + self.push(Candidate::new(self.target, item)) + } +} + +/// A [nearest neighbor search](https://en.wikipedia.org/wiki/Nearest_neighbor_search) index. +/// +/// Type parameters: +/// * `T`: The search result type. +/// * `U`: The query type. +pub trait NearestNeighbors = T> { + /// Returns the nearest neighbor to `target` (or `None` if this index is empty). + fn nearest(&self, target: &U) -> Option> { + self.search(SingletonNeighborhood::new(target, None)) + .into_option() + } + + /// Returns the nearest neighbor to `target` within the distance `threshold`, if one exists. + fn nearest_within(&self, target: &U, threshold: f64) -> Option> { + self.search(SingletonNeighborhood::new(target, Some(threshold))) + .into_option() + } + + /// Returns the up to `k` nearest neighbors to `target`. + fn k_nearest(&self, target: &U, k: usize) -> Vec> { + self.search(HeapNeighborhood::new(target, k, None)) + .into_vec() + } + + /// Returns the up to `k` nearest neighbors to `target` within the distance `threshold`. + fn k_nearest_within(&self, target: &U, k: usize, threshold: f64) -> Vec> { + self.search(HeapNeighborhood::new(target, k, Some(threshold))) + .into_vec() + } + + /// Search for nearest neighbors and add them to a neighborhood. + fn search<'a, 'b, N>(&'a self, neighborhood: N) -> N + where + T: 'a, + U: 'b, + N: Neighborhood<&'a T, &'b U>; +} + +/// A [NearestNeighbors] implementation that does exhaustive search. +#[derive(Debug)] +pub struct ExhaustiveSearch(Vec); + +impl ExhaustiveSearch { + /// Create an empty ExhaustiveSearch index. + pub fn new() -> Self { + Self(Vec::new()) + } + + /// Add a new item to the index. + pub fn push(&mut self, item: T) { + self.0.push(item); + } +} + +impl FromIterator for ExhaustiveSearch { + fn from_iter>(items: I) -> Self { + Self(items.into_iter().collect()) + } +} + +impl IntoIterator for ExhaustiveSearch { + type Item = T; + type IntoIter = std::vec::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +impl Extend for ExhaustiveSearch { + fn extend>(&mut self, iter: I) { + for value in iter { + self.push(value); + } + } +} + +impl> NearestNeighbors for ExhaustiveSearch { + fn search<'a, 'b, N>(&'a self, mut neighborhood: N) -> N + where + T: 'a, + U: 'b, + N: Neighborhood<&'a T, &'b U>, + { + for e in &self.0 { + neighborhood.consider(e); + } + neighborhood + } +} + +#[cfg(test)] +pub mod tests { + use super::*; + + use rand::prelude::*; + + #[derive(Clone, Copy, Debug, PartialEq)] + pub struct Point(pub [f64; 3]); + + impl Metric for Point { + type Distance = SquaredDistance; + + fn distance(&self, other: &Self) -> Self::Distance { + self.0.distance(&other.0) + } + } + + /// Test a [NearestNeighbors] impl. + pub fn test_nearest_neighbors(from_iter: F) + where + T: NearestNeighbors, + F: Fn(Vec) -> T, + { + test_empty(&from_iter); + test_pythagorean(&from_iter); + test_random_points(&from_iter); + } + + fn test_empty(from_iter: &F) + where + T: NearestNeighbors, + F: Fn(Vec) -> T, + { + let points = Vec::new(); + let index = from_iter(points); + let target = Point([0.0, 0.0, 0.0]); + assert_eq!(index.nearest(&target), None); + assert_eq!(index.nearest_within(&target, 1.0), None); + assert!(index.k_nearest(&target, 0).is_empty()); + assert!(index.k_nearest(&target, 3).is_empty()); + assert!(index.k_nearest_within(&target, 0, 1.0).is_empty()); + assert!(index.k_nearest_within(&target, 3, 1.0).is_empty()); + } + + fn test_pythagorean(from_iter: &F) + where + T: NearestNeighbors, + F: Fn(Vec) -> T, + { + let points = vec![ + Point([3.0, 4.0, 0.0]), + Point([5.0, 0.0, 12.0]), + Point([0.0, 8.0, 15.0]), + Point([1.0, 2.0, 2.0]), + Point([2.0, 3.0, 6.0]), + Point([4.0, 4.0, 7.0]), + ]; + let index = from_iter(points); + let target = Point([0.0, 0.0, 0.0]); + + assert_eq!( + index.nearest(&target), + Some(Neighbor::new(&Point([1.0, 2.0, 2.0]), 3.0)) + ); + + assert_eq!(index.nearest_within(&target, 2.0), None); + assert_eq!( + index.nearest_within(&target, 4.0), + Some(Neighbor::new(&Point([1.0, 2.0, 2.0]), 3.0)) + ); + + assert!(index.k_nearest(&target, 0).is_empty()); + assert_eq!( + index.k_nearest(&target, 3), + vec![ + Neighbor::new(&Point([1.0, 2.0, 2.0]), 3.0), + Neighbor::new(&Point([3.0, 4.0, 0.0]), 5.0), + Neighbor::new(&Point([2.0, 3.0, 6.0]), 7.0), + ] + ); + + assert!(index.k_nearest(&target, 0).is_empty()); + assert_eq!( + index.k_nearest_within(&target, 3, 6.0), + vec![ + Neighbor::new(&Point([1.0, 2.0, 2.0]), 3.0), + Neighbor::new(&Point([3.0, 4.0, 0.0]), 5.0), + ] + ); + assert_eq!( + index.k_nearest_within(&target, 3, 8.0), + vec![ + Neighbor::new(&Point([1.0, 2.0, 2.0]), 3.0), + Neighbor::new(&Point([3.0, 4.0, 0.0]), 5.0), + Neighbor::new(&Point([2.0, 3.0, 6.0]), 7.0), + ] + ); + } + + fn test_random_points(from_iter: &F) + where + T: NearestNeighbors, + F: Fn(Vec) -> T, + { + let mut points = Vec::new(); + for _ in 0..255 { + points.push(Point([random(), random(), random()])); + } + let target = Point([random(), random(), random()]); + + let eindex = ExhaustiveSearch::from_iter(points.clone()); + let index = from_iter(points); + + assert_eq!(index.k_nearest(&target, 3), eindex.k_nearest(&target, 3)); + } + + #[test] + fn test_exhaustive_index() { + test_nearest_neighbors(ExhaustiveSearch::from_iter); + } +} -- cgit v1.2.3 From 8d9de0e1028daed981246174182a39dd917b72bc Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sun, 19 Apr 2020 16:38:24 -0400 Subject: metric/vp: Implement vantage-point trees --- src/metric.rs | 2 + src/metric/vp.rs | 168 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 170 insertions(+) create mode 100644 src/metric/vp.rs diff --git a/src/metric.rs b/src/metric.rs index e067771..549db67 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -1,5 +1,7 @@ //! [Metric spaces](https://en.wikipedia.org/wiki/Metric_space). +pub mod vp; + use ordered_float::OrderedFloat; use std::cmp::Ordering; diff --git a/src/metric/vp.rs b/src/metric/vp.rs new file mode 100644 index 0000000..8d5b091 --- /dev/null +++ b/src/metric/vp.rs @@ -0,0 +1,168 @@ +//! [Vantage-point trees](https://en.wikipedia.org/wiki/Vantage-point_tree). + +use super::{Metric, NearestNeighbors, Neighborhood}; + +use std::iter::FromIterator; + +/// A node in a VP tree. +#[derive(Debug)] +struct VpNode { + /// The vantage point itself. + item: T, + /// The radius of this node. + radius: f64, + /// The subtree inside the radius, if any. + inside: Option>, + /// The subtree outside the radius, if any. + outside: Option>, +} + +impl VpNode { + /// Create a new VpNode. + fn new(mut items: Vec) -> Option> { + if items.is_empty() { + return None; + } + + let item = items.pop().unwrap(); + + items.sort_by_cached_key(|a| item.distance(a)); + + let mid = items.len() / 2; + let outside: Vec = items.drain(mid..).collect(); + + let radius = items.last().map(|l| item.distance(l).into()).unwrap_or(0.0); + + Some(Box::new(Self { + item, + radius, + inside: Self::new(items), + outside: Self::new(outside), + })) + } +} + +trait VpSearch<'a, T, U, N> { + /// Recursively search for nearest neighbors. + fn search(&'a self, neighborhood: &mut N); + + /// Search the inside subtree. + fn search_inside(&'a self, distance: f64, neighborhood: &mut N); + + /// Search the outside subtree. + fn search_outside(&'a self, distance: f64, neighborhood: &mut N); +} + +impl<'a, T, U, N> VpSearch<'a, T, U, N> for VpNode +where + T: 'a, + U: Metric<&'a T>, + N: Neighborhood<&'a T, U>, +{ + fn search(&'a self, neighborhood: &mut N) { + let distance = neighborhood.consider(&self.item).into(); + + if distance <= self.radius { + self.search_inside(distance, neighborhood); + self.search_outside(distance, neighborhood); + } else { + self.search_outside(distance, neighborhood); + self.search_inside(distance, neighborhood); + } + } + + fn search_inside(&'a self, distance: f64, neighborhood: &mut N) { + if let Some(inside) = &self.inside { + if neighborhood.contains(distance - self.radius) { + inside.search(neighborhood); + } + } + } + + fn search_outside(&'a self, distance: f64, neighborhood: &mut N) { + if let Some(outside) = &self.outside { + if neighborhood.contains(self.radius - distance) { + outside.search(neighborhood); + } + } + } +} + +/// A [vantage-point tree](https://en.wikipedia.org/wiki/Vantage-point_tree). +#[derive(Debug)] +pub struct VpTree { + root: Option>>, +} + +impl FromIterator for VpTree { + fn from_iter>(items: I) -> Self { + Self { + root: VpNode::new(items.into_iter().collect::>()), + } + } +} + +impl NearestNeighbors for VpTree +where + T: Metric, + U: Metric, +{ + fn search<'a, 'b, N>(&'a self, mut neighborhood: N) -> N + where + T: 'a, + U: 'b, + N: Neighborhood<&'a T, &'b U>, + { + if let Some(root) = &self.root { + root.search(&mut neighborhood); + } + neighborhood + } +} + +/// An iterator that moves values out of a VP tree. +#[derive(Debug)] +pub struct IntoIter { + stack: Vec>>, +} + +impl IntoIter { + fn new(node: Option>>) -> Self { + Self { + stack: node.into_iter().collect(), + } + } +} + +impl Iterator for IntoIter { + type Item = T; + + fn next(&mut self) -> Option { + self.stack.pop().map(|node| { + self.stack.extend(node.inside); + self.stack.extend(node.outside); + node.item + }) + } +} + +impl IntoIterator for VpTree { + type Item = T; + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + IntoIter::new(self.root) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::metric::tests::test_nearest_neighbors; + + #[test] + fn test_vp_tree() { + test_nearest_neighbors(VpTree::from_iter); + } +} -- cgit v1.2.3 From 1c560791902a4ef72efa671106d8f6d97fea50c1 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sun, 19 Apr 2020 16:40:29 -0400 Subject: metric/kd: Implement k-d trees --- src/metric.rs | 1 + src/metric/kd.rs | 217 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 218 insertions(+) create mode 100644 src/metric/kd.rs diff --git a/src/metric.rs b/src/metric.rs index 549db67..95191da 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -1,5 +1,6 @@ //! [Metric spaces](https://en.wikipedia.org/wiki/Metric_space). +pub mod kd; pub mod vp; use ordered_float::OrderedFloat; diff --git a/src/metric/kd.rs b/src/metric/kd.rs new file mode 100644 index 0000000..ab0cd2e --- /dev/null +++ b/src/metric/kd.rs @@ -0,0 +1,217 @@ +//! [k-d trees](https://en.wikipedia.org/wiki/K-d_tree). + +use super::{Metric, NearestNeighbors, Neighborhood}; + +use ordered_float::OrderedFloat; + +use std::iter::FromIterator; + +/// A point in Cartesian space. +pub trait Cartesian { + /// Returns the number of dimensions necessary to describe this point. + fn dimensions(&self) -> usize; + + /// Returns the value of the `i`th coordinate of this point (`i < self.dimensions()`). + fn coordinate(&self, i: usize) -> f64; +} + +/// Blanket [Cartesian] implementation for references. +impl<'a, T: Cartesian> Cartesian for &'a T { + fn dimensions(&self) -> usize { + (*self).dimensions() + } + + fn coordinate(&self, i: usize) -> f64 { + (*self).coordinate(i) + } +} + +/// Standard cartesian space. +impl Cartesian for [f64] { + fn dimensions(&self) -> usize { + self.len() + } + + fn coordinate(&self, i: usize) -> f64 { + self[i] + } +} + +/// A node in a k-d tree. +#[derive(Debug)] +struct KdNode { + /// The value stored in this node. + item: T, + /// The left subtree, if any. + left: Option>, + /// The right subtree, if any. + right: Option>, +} + +trait KdSearch<'a, T, U, N> { + /// Recursively search for nearest neighbors. + fn search(&'a self, i: usize, neighborhood: &mut N); + + /// Search the left subtree. + fn search_left(&'a self, i: usize, distance: f64, neighborhood: &mut N); + + /// Search the right subtree. + fn search_right(&'a self, i: usize, distance: f64, neighborhood: &mut N); +} + +impl<'a, T, U, N> KdSearch<'a, T, U, N> for KdNode +where + T: 'a + Cartesian, + U: Cartesian + Metric<&'a T>, + N: Neighborhood<&'a T, U>, +{ + fn search(&'a self, i: usize, neighborhood: &mut N) { + neighborhood.consider(&self.item); + + let distance = neighborhood.target().coordinate(i) - self.item.coordinate(i); + let j = (i + 1) % self.item.dimensions(); + if distance <= 0.0 { + self.search_left(j, distance, neighborhood); + self.search_right(j, -distance, neighborhood); + } else { + self.search_right(j, -distance, neighborhood); + self.search_left(j, distance, neighborhood); + } + } + + fn search_left(&'a self, i: usize, distance: f64, neighborhood: &mut N) { + if let Some(left) = &self.left { + if neighborhood.contains(distance) { + left.search(i, neighborhood); + } + } + } + + fn search_right(&'a self, i: usize, distance: f64, neighborhood: &mut N) { + if let Some(right) = &self.right { + if neighborhood.contains(distance) { + right.search(i, neighborhood); + } + } + } +} + +impl KdNode { + /// Create a new KdNode. + fn new(i: usize, mut items: Vec) -> Option> { + if items.is_empty() { + return None; + } + + items.sort_unstable_by_key(|x| OrderedFloat::from(x.coordinate(i))); + + let mid = items.len() / 2; + let right: Vec = items.drain((mid + 1)..).collect(); + let item = items.pop().unwrap(); + let j = (i + 1) % item.dimensions(); + Some(Box::new(Self { + item, + left: Self::new(j, items), + right: Self::new(j, right), + })) + } +} + +/// A [k-d tree](https://en.wikipedia.org/wiki/K-d_tree). +#[derive(Debug)] +pub struct KdTree { + root: Option>>, +} + +impl FromIterator for KdTree { + /// Create a new k-d tree from a set of points. + fn from_iter>(items: I) -> Self { + Self { + root: KdNode::new(0, items.into_iter().collect()), + } + } +} + +impl NearestNeighbors for KdTree +where + T: Cartesian, + U: Cartesian + Metric, +{ + fn search<'a, 'b, N>(&'a self, mut neighborhood: N) -> N + where + T: 'a, + U: 'b, + N: Neighborhood<&'a T, &'b U>, + { + if let Some(root) = &self.root { + root.search(0, &mut neighborhood); + } + neighborhood + } +} + +/// An iterator that the moves values out of a k-d tree. +#[derive(Debug)] +pub struct IntoIter { + stack: Vec>>, +} + +impl IntoIter { + fn new(node: Option>>) -> Self { + Self { + stack: node.into_iter().collect(), + } + } +} + +impl Iterator for IntoIter { + type Item = T; + + fn next(&mut self) -> Option { + self.stack.pop().map(|node| { + self.stack.extend(node.left); + self.stack.extend(node.right); + node.item + }) + } +} + +impl IntoIterator for KdTree { + type Item = T; + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + IntoIter::new(self.root) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::metric::tests::{test_nearest_neighbors, Point}; + use crate::metric::SquaredDistance; + + impl Metric<[f64]> for Point { + type Distance = SquaredDistance; + + fn distance(&self, other: &[f64]) -> Self::Distance { + self.0.distance(other) + } + } + + impl Cartesian for Point { + fn dimensions(&self) -> usize { + self.0.dimensions() + } + + fn coordinate(&self, i: usize) -> f64 { + self.0.coordinate(i) + } + } + + #[test] + fn test_kd_tree() { + test_nearest_neighbors(KdTree::from_iter); + } +} -- cgit v1.2.3 From 5de377b2b00a927a4f6463c1c5a5fd18606ad006 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Thu, 30 Apr 2020 22:51:06 -0400 Subject: metric/kd: Prune k-d tree searches more aggressively --- src/metric/kd.rs | 116 +++++++++++++++++++++++++++++++------------------------ 1 file changed, 65 insertions(+), 51 deletions(-) diff --git a/src/metric/kd.rs b/src/metric/kd.rs index ab0cd2e..db1b2bd 100644 --- a/src/metric/kd.rs +++ b/src/metric/kd.rs @@ -7,7 +7,7 @@ use ordered_float::OrderedFloat; use std::iter::FromIterator; /// A point in Cartesian space. -pub trait Cartesian { +pub trait Cartesian: Metric<[f64]> { /// Returns the number of dimensions necessary to describe this point. fn dimensions(&self) -> usize; @@ -26,6 +26,15 @@ impl<'a, T: Cartesian> Cartesian for &'a T { } } +/// Blanket [Metric<[f64]>](Metric) implementation for [Cartesian] references. +impl<'a, T: Cartesian> Metric<[f64]> for &'a T { + type Distance = T::Distance; + + fn distance(&self, other: &[f64]) -> Self::Distance { + (*self).distance(other) + } +} + /// Standard cartesian space. impl Cartesian for [f64] { fn dimensions(&self) -> usize { @@ -37,6 +46,21 @@ impl Cartesian for [f64] { } } +/// Marker trait for cartesian metric spaces. +pub trait CartesianMetric: + Cartesian + Metric>::Distance> +{ +} + +/// Blanket [CartesianMetric] implementation for cartesian spaces with compatible metric distance +/// types. +impl CartesianMetric for U +where + T: ?Sized, + U: ?Sized + Cartesian + Metric>::Distance>, +{ +} + /// A node in a k-d tree. #[derive(Debug)] struct KdNode { @@ -48,54 +72,6 @@ struct KdNode { right: Option>, } -trait KdSearch<'a, T, U, N> { - /// Recursively search for nearest neighbors. - fn search(&'a self, i: usize, neighborhood: &mut N); - - /// Search the left subtree. - fn search_left(&'a self, i: usize, distance: f64, neighborhood: &mut N); - - /// Search the right subtree. - fn search_right(&'a self, i: usize, distance: f64, neighborhood: &mut N); -} - -impl<'a, T, U, N> KdSearch<'a, T, U, N> for KdNode -where - T: 'a + Cartesian, - U: Cartesian + Metric<&'a T>, - N: Neighborhood<&'a T, U>, -{ - fn search(&'a self, i: usize, neighborhood: &mut N) { - neighborhood.consider(&self.item); - - let distance = neighborhood.target().coordinate(i) - self.item.coordinate(i); - let j = (i + 1) % self.item.dimensions(); - if distance <= 0.0 { - self.search_left(j, distance, neighborhood); - self.search_right(j, -distance, neighborhood); - } else { - self.search_right(j, -distance, neighborhood); - self.search_left(j, distance, neighborhood); - } - } - - fn search_left(&'a self, i: usize, distance: f64, neighborhood: &mut N) { - if let Some(left) = &self.left { - if neighborhood.contains(distance) { - left.search(i, neighborhood); - } - } - } - - fn search_right(&'a self, i: usize, distance: f64, neighborhood: &mut N) { - if let Some(right) = &self.right { - if neighborhood.contains(distance) { - right.search(i, neighborhood); - } - } - } -} - impl KdNode { /// Create a new KdNode. fn new(i: usize, mut items: Vec) -> Option> { @@ -115,6 +91,40 @@ impl KdNode { right: Self::new(j, right), })) } + + /// Recursively search for nearest neighbors. + fn search<'a, U, N>(&'a self, i: usize, closest: &mut [f64], neighborhood: &mut N) + where + T: 'a, + U: CartesianMetric<&'a T>, + N: Neighborhood<&'a T, U>, + { + neighborhood.consider(&self.item); + + let target = neighborhood.target(); + let ti = target.coordinate(i); + let si = self.item.coordinate(i); + let j = (i + 1) % self.item.dimensions(); + + let (near, far) = if ti <= si { + (&self.left, &self.right) + } else { + (&self.right, &self.left) + }; + + if let Some(near) = near { + near.search(j, closest, neighborhood); + } + + if let Some(far) = far { + let saved = closest[i]; + closest[i] = si; + if neighborhood.contains_distance(target.distance(closest)) { + far.search(j, closest, neighborhood); + } + closest[i] = saved; + } + } } /// A [k-d tree](https://en.wikipedia.org/wiki/K-d_tree). @@ -135,7 +145,7 @@ impl FromIterator for KdTree { impl NearestNeighbors for KdTree where T: Cartesian, - U: Cartesian + Metric, + U: CartesianMetric, { fn search<'a, 'b, N>(&'a self, mut neighborhood: N) -> N where @@ -143,8 +153,12 @@ where U: 'b, N: Neighborhood<&'a T, &'b U>, { + let target = neighborhood.target(); + let dims = target.dimensions(); + let mut closest: Vec<_> = (0..dims).map(|i| target.coordinate(i)).collect(); + if let Some(root) = &self.root { - root.search(0, &mut neighborhood); + root.search(0, &mut closest, &mut neighborhood); } neighborhood } -- cgit v1.2.3 From 9699f4657ecaaf4361448f249e4f2e210a854af4 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Thu, 23 Apr 2020 09:55:26 -0400 Subject: metric/vp: Flatten the tree representation --- src/metric/vp.rs | 133 +++++++++++++++++++++---------------------------------- 1 file changed, 51 insertions(+), 82 deletions(-) diff --git a/src/metric/vp.rs b/src/metric/vp.rs index 8d5b091..fae62e5 100644 --- a/src/metric/vp.rs +++ b/src/metric/vp.rs @@ -11,78 +11,62 @@ struct VpNode { item: T, /// The radius of this node. radius: f64, - /// The subtree inside the radius, if any. - inside: Option>, - /// The subtree outside the radius, if any. - outside: Option>, + /// The size of the subtree inside the radius. + inside_len: usize, } impl VpNode { /// Create a new VpNode. - fn new(mut items: Vec) -> Option> { - if items.is_empty() { - return None; + fn new(item: T) -> Self { + Self { + item, + radius: 0.0, + inside_len: 0, } + } - let item = items.pop().unwrap(); - - items.sort_by_cached_key(|a| item.distance(a)); - - let mid = items.len() / 2; - let outside: Vec = items.drain(mid..).collect(); + /// Build a VP tree recursively. + fn build(slice: &mut [VpNode]) { + if let Some((node, children)) = slice.split_first_mut() { + let item = &node.item; + children.sort_by_cached_key(|n| item.distance(&n.item)); - let radius = items.last().map(|l| item.distance(l).into()).unwrap_or(0.0); + let (inside, outside) = children.split_at_mut(children.len() / 2); + if let Some(last) = inside.last() { + node.radius = item.distance(&last.item).into(); + } + node.inside_len = inside.len(); - Some(Box::new(Self { - item, - radius, - inside: Self::new(items), - outside: Self::new(outside), - })) + Self::build(inside); + Self::build(outside); + } } -} -trait VpSearch<'a, T, U, N> { /// Recursively search for nearest neighbors. - fn search(&'a self, neighborhood: &mut N); - - /// Search the inside subtree. - fn search_inside(&'a self, distance: f64, neighborhood: &mut N); - - /// Search the outside subtree. - fn search_outside(&'a self, distance: f64, neighborhood: &mut N); -} + fn recurse<'a, U, N>(slice: &'a [VpNode], neighborhood: &mut N) + where + T: 'a, + U: Metric<&'a T>, + N: Neighborhood<&'a T, U>, + { + let (node, children) = slice.split_first().unwrap(); + let (inside, outside) = children.split_at(node.inside_len); -impl<'a, T, U, N> VpSearch<'a, T, U, N> for VpNode -where - T: 'a, - U: Metric<&'a T>, - N: Neighborhood<&'a T, U>, -{ - fn search(&'a self, neighborhood: &mut N) { - let distance = neighborhood.consider(&self.item).into(); + let distance = neighborhood.consider(&node.item).into(); - if distance <= self.radius { - self.search_inside(distance, neighborhood); - self.search_outside(distance, neighborhood); + if distance <= node.radius { + if !inside.is_empty() && neighborhood.contains(distance - node.radius) { + Self::recurse(inside, neighborhood); + } + if !outside.is_empty() && neighborhood.contains(node.radius - distance) { + Self::recurse(outside, neighborhood); + } } else { - self.search_outside(distance, neighborhood); - self.search_inside(distance, neighborhood); - } - } - - fn search_inside(&'a self, distance: f64, neighborhood: &mut N) { - if let Some(inside) = &self.inside { - if neighborhood.contains(distance - self.radius) { - inside.search(neighborhood); + if !outside.is_empty() && neighborhood.contains(node.radius - distance) { + Self::recurse(outside, neighborhood); } - } - } - - fn search_outside(&'a self, distance: f64, neighborhood: &mut N) { - if let Some(outside) = &self.outside { - if neighborhood.contains(self.radius - distance) { - outside.search(neighborhood); + if !inside.is_empty() && neighborhood.contains(distance - node.radius) { + Self::recurse(inside, neighborhood); } } } @@ -90,15 +74,13 @@ where /// A [vantage-point tree](https://en.wikipedia.org/wiki/Vantage-point_tree). #[derive(Debug)] -pub struct VpTree { - root: Option>>, -} +pub struct VpTree(Vec>); impl FromIterator for VpTree { fn from_iter>(items: I) -> Self { - Self { - root: VpNode::new(items.into_iter().collect::>()), - } + let mut nodes: Vec<_> = items.into_iter().map(VpNode::new).collect(); + VpNode::build(nodes.as_mut_slice()); + Self(nodes) } } @@ -113,36 +95,23 @@ where U: 'b, N: Neighborhood<&'a T, &'b U>, { - if let Some(root) = &self.root { - root.search(&mut neighborhood); + if !self.0.is_empty() { + VpNode::recurse(&self.0, &mut neighborhood); } + neighborhood } } /// An iterator that moves values out of a VP tree. #[derive(Debug)] -pub struct IntoIter { - stack: Vec>>, -} - -impl IntoIter { - fn new(node: Option>>) -> Self { - Self { - stack: node.into_iter().collect(), - } - } -} +pub struct IntoIter(std::vec::IntoIter>); impl Iterator for IntoIter { type Item = T; fn next(&mut self) -> Option { - self.stack.pop().map(|node| { - self.stack.extend(node.inside); - self.stack.extend(node.outside); - node.item - }) + self.0.next().map(|n| n.item) } } @@ -151,7 +120,7 @@ impl IntoIterator for VpTree { type IntoIter = IntoIter; fn into_iter(self) -> Self::IntoIter { - IntoIter::new(self.root) + IntoIter(self.0.into_iter()) } } -- cgit v1.2.3 From e9a81a6d0df149252164003975addf175d5c6f4b Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Thu, 23 Apr 2020 09:55:13 -0400 Subject: metric/kd: Flatten the tree representation --- src/metric/kd.rs | 113 ++++++++++++++++++++++++++----------------------------- 1 file changed, 54 insertions(+), 59 deletions(-) diff --git a/src/metric/kd.rs b/src/metric/kd.rs index db1b2bd..2caf4a3 100644 --- a/src/metric/kd.rs +++ b/src/metric/kd.rs @@ -66,61 +66,71 @@ where struct KdNode { /// The value stored in this node. item: T, - /// The left subtree, if any. - left: Option>, - /// The right subtree, if any. - right: Option>, + /// The size of the left subtree. + left_len: usize, } impl KdNode { /// Create a new KdNode. - fn new(i: usize, mut items: Vec) -> Option> { - if items.is_empty() { - return None; + fn new(item: T) -> Self { + Self { item, left_len: 0 } + } + + /// Build a k-d tree recursively. + fn build(slice: &mut [KdNode], i: usize) { + if slice.is_empty() { + return; } - items.sort_unstable_by_key(|x| OrderedFloat::from(x.coordinate(i))); - - let mid = items.len() / 2; - let right: Vec = items.drain((mid + 1)..).collect(); - let item = items.pop().unwrap(); - let j = (i + 1) % item.dimensions(); - Some(Box::new(Self { - item, - left: Self::new(j, items), - right: Self::new(j, right), - })) + slice.sort_unstable_by_key(|n| OrderedFloat::from(n.item.coordinate(i))); + + let mid = slice.len() / 2; + slice.swap(0, mid); + + let (node, children) = slice.split_first_mut().unwrap(); + let (left, right) = children.split_at_mut(mid); + node.left_len = left.len(); + + let j = (i + 1) % node.item.dimensions(); + Self::build(left, j); + Self::build(right, j); } /// Recursively search for nearest neighbors. - fn search<'a, U, N>(&'a self, i: usize, closest: &mut [f64], neighborhood: &mut N) - where + fn recurse<'a, U, N>( + slice: &'a [KdNode], + i: usize, + closest: &mut [f64], + neighborhood: &mut N, + ) where T: 'a, U: CartesianMetric<&'a T>, N: Neighborhood<&'a T, U>, { - neighborhood.consider(&self.item); + let (node, children) = slice.split_first().unwrap(); + neighborhood.consider(&node.item); let target = neighborhood.target(); let ti = target.coordinate(i); - let si = self.item.coordinate(i); - let j = (i + 1) % self.item.dimensions(); + let ni = node.item.coordinate(i); + let j = (i + 1) % node.item.dimensions(); - let (near, far) = if ti <= si { - (&self.left, &self.right) + let (left, right) = children.split_at(node.left_len); + let (near, far) = if ti <= ni { + (left, right) } else { - (&self.right, &self.left) + (right, left) }; - if let Some(near) = near { - near.search(j, closest, neighborhood); + if !near.is_empty() { + Self::recurse(near, j, closest, neighborhood); } - if let Some(far) = far { + if !far.is_empty() { let saved = closest[i]; - closest[i] = si; + closest[i] = ni; if neighborhood.contains_distance(target.distance(closest)) { - far.search(j, closest, neighborhood); + Self::recurse(far, j, closest, neighborhood); } closest[i] = saved; } @@ -129,16 +139,14 @@ impl KdNode { /// A [k-d tree](https://en.wikipedia.org/wiki/K-d_tree). #[derive(Debug)] -pub struct KdTree { - root: Option>>, -} +pub struct KdTree(Vec>); impl FromIterator for KdTree { /// Create a new k-d tree from a set of points. fn from_iter>(items: I) -> Self { - Self { - root: KdNode::new(0, items.into_iter().collect()), - } + let mut nodes: Vec<_> = items.into_iter().map(KdNode::new).collect(); + KdNode::build(nodes.as_mut_slice(), 0); + Self(nodes) } } @@ -153,40 +161,27 @@ where U: 'b, N: Neighborhood<&'a T, &'b U>, { - let target = neighborhood.target(); - let dims = target.dimensions(); - let mut closest: Vec<_> = (0..dims).map(|i| target.coordinate(i)).collect(); + if !self.0.is_empty() { + let target = neighborhood.target(); + let dims = target.dimensions(); + let mut closest: Vec<_> = (0..dims).map(|i| target.coordinate(i)).collect(); - if let Some(root) = &self.root { - root.search(0, &mut closest, &mut neighborhood); + KdNode::recurse(&self.0, 0, &mut closest, &mut neighborhood); } + neighborhood } } /// An iterator that the moves values out of a k-d tree. #[derive(Debug)] -pub struct IntoIter { - stack: Vec>>, -} - -impl IntoIter { - fn new(node: Option>>) -> Self { - Self { - stack: node.into_iter().collect(), - } - } -} +pub struct IntoIter(std::vec::IntoIter>); impl Iterator for IntoIter { type Item = T; fn next(&mut self) -> Option { - self.stack.pop().map(|node| { - self.stack.extend(node.left); - self.stack.extend(node.right); - node.item - }) + self.0.next().map(|n| n.item) } } @@ -195,7 +190,7 @@ impl IntoIterator for KdTree { type IntoIter = IntoIter; fn into_iter(self) -> Self::IntoIter { - IntoIter::new(self.root) + IntoIter(self.0.into_iter()) } } -- cgit v1.2.3 From a4a75059f302de2a00971f1f485fcf4389710628 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sun, 19 Apr 2020 16:55:17 -0400 Subject: metric/forest: Implement dynamized forests --- src/metric.rs | 1 + src/metric/forest.rs | 152 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 153 insertions(+) create mode 100644 src/metric/forest.rs diff --git a/src/metric.rs b/src/metric.rs index 95191da..7a5f5f7 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -1,5 +1,6 @@ //! [Metric spaces](https://en.wikipedia.org/wiki/Metric_space). +pub mod forest; pub mod kd; pub mod vp; diff --git a/src/metric/forest.rs b/src/metric/forest.rs new file mode 100644 index 0000000..f23c451 --- /dev/null +++ b/src/metric/forest.rs @@ -0,0 +1,152 @@ +//! [Dynamization](https://en.wikipedia.org/wiki/Dynamization) for nearest neighbor search. + +use super::kd::KdTree; +use super::vp::VpTree; +use super::{Metric, NearestNeighbors, Neighborhood}; + +use std::iter::{Extend, Flatten, FromIterator}; + +/// A dynamic wrapper for a static nearest neighbor search data structure. +/// +/// This type applies [dynamization](https://en.wikipedia.org/wiki/Dynamization) to an arbitrary +/// nearest neighbor search structure `T`, allowing new items to be added dynamically. +#[derive(Debug)] +pub struct Forest(Vec>); + +impl Forest +where + U: FromIterator + IntoIterator, +{ + /// Create a new empty forest. + pub fn new() -> Self { + Self(Vec::new()) + } + + /// Add a new item to the forest. + pub fn push(&mut self, item: T) { + let mut items = vec![item]; + + for slot in &mut self.0 { + match slot.take() { + // Collect the items from any trees we encounter... + Some(tree) => { + items.extend(tree); + } + // ... and put them all in the first empty slot + None => { + *slot = Some(items.into_iter().collect()); + return; + } + } + } + + self.0.push(Some(items.into_iter().collect())); + } + + /// Get the number of items in the forest. + pub fn len(&self) -> usize { + let mut len = 0; + for (i, slot) in self.0.iter().enumerate() { + if slot.is_some() { + len |= 1 << i; + } + } + len + } +} + +impl Extend for Forest +where + U: FromIterator + IntoIterator, +{ + fn extend>(&mut self, items: I) { + for item in items { + self.push(item); + } + } +} + +impl FromIterator for Forest +where + U: FromIterator + IntoIterator, +{ + fn from_iter>(items: I) -> Self { + let mut forest = Self::new(); + forest.extend(items); + forest + } +} + +type IntoIterImpl = Flatten>>>; + +/// An iterator that moves items out of a forest. +pub struct IntoIter(IntoIterImpl); + +impl Iterator for IntoIter { + type Item = T::Item; + + fn next(&mut self) -> Option { + self.0.next() + } +} + +impl IntoIterator for Forest { + type Item = T::Item; + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + IntoIter(self.0.into_iter().flatten().flatten()) + } +} + +impl NearestNeighbors for Forest +where + U: Metric, + V: NearestNeighbors, +{ + fn search<'a, 'b, N>(&'a self, neighborhood: N) -> N + where + T: 'a, + U: 'b, + N: Neighborhood<&'a T, &'b U>, + { + self.0 + .iter() + .flatten() + .fold(neighborhood, |n, t| t.search(n)) + } +} + +/// A forest of k-d trees. +pub type KdForest = Forest>; + +/// A forest of vantage-point trees. +pub type VpForest = Forest>; + +#[cfg(test)] +mod tests { + use super::*; + + use crate::metric::tests::test_nearest_neighbors; + use crate::metric::ExhaustiveSearch; + + #[test] + fn test_exhaustive_forest() { + test_nearest_neighbors(Forest::>::from_iter); + } + + #[test] + fn test_forest_forest() { + test_nearest_neighbors(Forest::>>::from_iter); + } + + #[test] + fn test_kd_forest() { + test_nearest_neighbors(KdForest::from_iter); + } + + #[test] + fn test_vp_forest() { + test_nearest_neighbors(VpForest::from_iter); + } +} -- cgit v1.2.3 From b4a39a3f22fac361f6a535d281eee5586078281b Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Thu, 23 Apr 2020 14:58:47 -0400 Subject: metric/forest: Optimize bulk insertion --- src/metric/forest.rs | 47 +++++++++++++++++++++++++++-------------------- 1 file changed, 27 insertions(+), 20 deletions(-) diff --git a/src/metric/forest.rs b/src/metric/forest.rs index f23c451..29b6f55 100644 --- a/src/metric/forest.rs +++ b/src/metric/forest.rs @@ -4,7 +4,7 @@ use super::kd::KdTree; use super::vp::VpTree; use super::{Metric, NearestNeighbors, Neighborhood}; -use std::iter::{Extend, Flatten, FromIterator}; +use std::iter::{self, Extend, Flatten, FromIterator}; /// A dynamic wrapper for a static nearest neighbor search data structure. /// @@ -24,23 +24,7 @@ where /// Add a new item to the forest. pub fn push(&mut self, item: T) { - let mut items = vec![item]; - - for slot in &mut self.0 { - match slot.take() { - // Collect the items from any trees we encounter... - Some(tree) => { - items.extend(tree); - } - // ... and put them all in the first empty slot - None => { - *slot = Some(items.into_iter().collect()); - return; - } - } - } - - self.0.push(Some(items.into_iter().collect())); + self.extend(iter::once(item)); } /// Get the number of items in the forest. @@ -60,9 +44,32 @@ where U: FromIterator + IntoIterator, { fn extend>(&mut self, items: I) { - for item in items { - self.push(item); + let mut vec: Vec<_> = items.into_iter().collect(); + let new_len = self.len() + vec.len(); + + for i in 0.. { + let bit = 1 << i; + + if bit > new_len { + break; + } + + if i >= self.0.len() { + self.0.push(None); + } + + if new_len & bit == 0 { + if let Some(tree) = self.0[i].take() { + vec.extend(tree); + } + } else if self.0[i].is_none() { + let offset = vec.len() - bit; + self.0[i] = Some(vec.drain(offset..).collect()); + } } + + debug_assert!(vec.is_empty()); + debug_assert!(self.len() == new_len); } } -- cgit v1.2.3 From e53ace3e69ed6bacedb7de345df10d3e575a291e Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sun, 19 Apr 2020 16:58:52 -0400 Subject: metric/soft: Implement soft deletes --- src/metric.rs | 1 + src/metric/soft.rs | 282 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 283 insertions(+) create mode 100644 src/metric/soft.rs diff --git a/src/metric.rs b/src/metric.rs index 7a5f5f7..b46c8da 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -2,6 +2,7 @@ pub mod forest; pub mod kd; +pub mod soft; pub mod vp; use ordered_float::OrderedFloat; diff --git a/src/metric/soft.rs b/src/metric/soft.rs new file mode 100644 index 0000000..0d7dcdb --- /dev/null +++ b/src/metric/soft.rs @@ -0,0 +1,282 @@ +//! [Soft deletion](https://en.wiktionary.org/wiki/soft_deletion) for nearest neighbor search. + +use super::forest::{KdForest, VpForest}; +use super::kd::KdTree; +use super::vp::VpTree; +use super::{Metric, NearestNeighbors, Neighborhood}; + +use std::iter; +use std::iter::FromIterator; +use std::mem; + +/// A trait for objects that can be soft-deleted. +pub trait SoftDelete { + /// Check whether this item is deleted. + fn is_deleted(&self) -> bool; +} + +/// Blanket [SoftDelete] implementation for references. +impl<'a, T: SoftDelete> SoftDelete for &'a T { + fn is_deleted(&self) -> bool { + (*self).is_deleted() + } +} + +/// [Neighborhood] wrapper that ignores soft-deleted items. +#[derive(Debug)] +struct SoftNeighborhood(N); + +impl Neighborhood for SoftNeighborhood +where + T: SoftDelete, + U: Metric, + N: Neighborhood, +{ + fn target(&self) -> U { + self.0.target() + } + + fn contains(&self, distance: f64) -> bool { + self.0.contains(distance) + } + + fn contains_distance(&self, distance: U::Distance) -> bool { + self.0.contains_distance(distance) + } + + fn consider(&mut self, item: T) -> U::Distance { + if item.is_deleted() { + self.target().distance(&item) + } else { + self.0.consider(item) + } + } +} + +/// A [NearestNeighbors] implementation that supports [soft deletes](https://en.wiktionary.org/wiki/soft_deletion). +#[derive(Debug)] +pub struct SoftSearch(T); + +impl SoftSearch +where + T: SoftDelete, + U: FromIterator + IntoIterator, +{ + /// Create a new empty soft index. + pub fn new() -> Self { + Self(iter::empty().collect()) + } + + /// Push a new item into this index. + pub fn push(&mut self, item: T) + where + U: Extend, + { + self.0.extend(iter::once(item)); + } + + /// Rebuild this index, discarding deleted items. + pub fn rebuild(&mut self) { + let items = mem::replace(&mut self.0, iter::empty().collect()); + self.0 = items.into_iter().filter(|e| !e.is_deleted()).collect(); + } +} + +impl> Extend for SoftSearch { + fn extend>(&mut self, iter: I) { + self.0.extend(iter); + } +} + +impl> FromIterator for SoftSearch { + fn from_iter>(iter: I) -> Self { + Self(U::from_iter(iter)) + } +} + +impl IntoIterator for SoftSearch { + type Item = T::Item; + type IntoIter = T::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +impl NearestNeighbors for SoftSearch +where + T: SoftDelete, + U: Metric, + V: NearestNeighbors, +{ + fn search<'a, 'b, N>(&'a self, neighborhood: N) -> N + where + T: 'a, + U: 'b, + N: Neighborhood<&'a T, &'b U>, + { + self.0.search(SoftNeighborhood(neighborhood)).0 + } +} + +/// A k-d forest that supports soft deletes. +pub type SoftKdForest = SoftSearch>; + +/// A k-d tree that supports soft deletes. +pub type SoftKdTree = SoftSearch>; + +/// A VP forest that supports soft deletes. +pub type SoftVpForest = SoftSearch>; + +/// A VP tree that supports soft deletes. +pub type SoftVpTree = SoftSearch>; + +#[cfg(test)] +mod tests { + use super::*; + + use crate::metric::kd::Cartesian; + use crate::metric::tests::Point; + use crate::metric::Neighbor; + + #[derive(Debug, PartialEq)] + struct SoftPoint { + point: Point, + deleted: bool, + } + + impl SoftPoint { + fn new(x: f64, y: f64, z: f64) -> Self { + Self { + point: Point([x, y, z]), + deleted: false, + } + } + + fn deleted(x: f64, y: f64, z: f64) -> Self { + Self { + point: Point([x, y, z]), + deleted: true, + } + } + } + + impl SoftDelete for SoftPoint { + fn is_deleted(&self) -> bool { + self.deleted + } + } + + impl Metric for SoftPoint { + type Distance = ::Distance; + + fn distance(&self, other: &Self) -> Self::Distance { + self.point.distance(&other.point) + } + } + + impl Metric<[f64]> for SoftPoint { + type Distance = ::Distance; + + fn distance(&self, other: &[f64]) -> Self::Distance { + self.point.distance(other) + } + } + + impl Cartesian for SoftPoint { + fn dimensions(&self) -> usize { + self.point.dimensions() + } + + fn coordinate(&self, i: usize) -> f64 { + self.point.coordinate(i) + } + } + + impl Metric for Point { + type Distance = ::Distance; + + fn distance(&self, other: &SoftPoint) -> Self::Distance { + self.distance(&other.point) + } + } + + fn test_index(index: &T) + where + T: NearestNeighbors, + { + let target = Point([0.0, 0.0, 0.0]); + + assert_eq!( + index.nearest(&target), + Some(Neighbor::new(&SoftPoint::new(1.0, 2.0, 2.0), 3.0)) + ); + + assert_eq!(index.nearest_within(&target, 2.0), None); + assert_eq!( + index.nearest_within(&target, 4.0), + Some(Neighbor::new(&SoftPoint::new(1.0, 2.0, 2.0), 3.0)) + ); + + assert_eq!( + index.k_nearest(&target, 3), + vec![ + Neighbor::new(&SoftPoint::new(1.0, 2.0, 2.0), 3.0), + Neighbor::new(&SoftPoint::new(3.0, 4.0, 0.0), 5.0), + Neighbor::new(&SoftPoint::new(2.0, 3.0, 6.0), 7.0), + ] + ); + + assert_eq!( + index.k_nearest_within(&target, 3, 6.0), + vec![ + Neighbor::new(&SoftPoint::new(1.0, 2.0, 2.0), 3.0), + Neighbor::new(&SoftPoint::new(3.0, 4.0, 0.0), 5.0), + ] + ); + assert_eq!( + index.k_nearest_within(&target, 3, 8.0), + vec![ + Neighbor::new(&SoftPoint::new(1.0, 2.0, 2.0), 3.0), + Neighbor::new(&SoftPoint::new(3.0, 4.0, 0.0), 5.0), + Neighbor::new(&SoftPoint::new(2.0, 3.0, 6.0), 7.0), + ] + ); + } + + fn test_soft_index(index: &mut SoftSearch) + where + T: Extend, + T: FromIterator, + T: IntoIterator, + T: NearestNeighbors, + { + let points = vec![ + SoftPoint::deleted(0.0, 0.0, 0.0), + SoftPoint::new(3.0, 4.0, 0.0), + SoftPoint::new(5.0, 0.0, 12.0), + SoftPoint::new(0.0, 8.0, 15.0), + SoftPoint::new(1.0, 2.0, 2.0), + SoftPoint::new(2.0, 3.0, 6.0), + SoftPoint::new(4.0, 4.0, 7.0), + ]; + + for point in points { + index.push(point); + } + test_index(index); + + index.rebuild(); + test_index(index); + } + + #[test] + fn test_soft_kd_forest() { + test_soft_index(&mut SoftKdForest::new()); + } + + #[test] + fn test_soft_vp_forest() { + test_soft_index(&mut SoftVpForest::new()); + } +} -- cgit v1.2.3 From 462daaffd5ec720ed80a2e7b1f445a73cabf5833 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sun, 19 Apr 2020 16:59:24 -0400 Subject: metric/approx: Implement approximate nearest neighbor search --- src/metric.rs | 1 + src/metric/approx.rs | 131 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+) create mode 100644 src/metric/approx.rs diff --git a/src/metric.rs b/src/metric.rs index b46c8da..268aefd 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -1,5 +1,6 @@ //! [Metric spaces](https://en.wikipedia.org/wiki/Metric_space). +pub mod approx; pub mod forest; pub mod kd; pub mod soft; diff --git a/src/metric/approx.rs b/src/metric/approx.rs new file mode 100644 index 0000000..c23f9c7 --- /dev/null +++ b/src/metric/approx.rs @@ -0,0 +1,131 @@ +//! [Approximate nearest neighbor search](https://en.wikipedia.org/wiki/Nearest_neighbor_search#Approximate_nearest_neighbor). + +use super::{Metric, NearestNeighbors, Neighborhood}; + +/// An approximate [Neighborhood], for approximate nearest neighbor searches. +#[derive(Debug)] +struct ApproximateNeighborhood { + inner: N, + ratio: f64, + limit: usize, +} + +impl ApproximateNeighborhood { + fn new(inner: N, ratio: f64, limit: usize) -> Self { + Self { + inner, + ratio, + limit, + } + } +} + +impl Neighborhood for ApproximateNeighborhood +where + U: Metric, + N: Neighborhood, +{ + fn target(&self) -> U { + self.inner.target() + } + + fn contains(&self, distance: f64) -> bool { + if self.limit > 0 { + self.inner.contains(self.ratio * distance) + } else { + false + } + } + + fn contains_distance(&self, distance: U::Distance) -> bool { + self.contains(self.ratio * distance.into()) + } + + fn consider(&mut self, item: T) -> U::Distance { + self.limit = self.limit.saturating_sub(1); + self.inner.consider(item) + } +} + +/// An [approximate nearest neighbor search](https://en.wikipedia.org/wiki/Nearest_neighbor_search#Approximate_nearest_neighbor) +/// index. +/// +/// This wrapper converts an exact nearest neighbor search algorithm into an approximate one by +/// modifying the behavior of [Neighborhood::contains]. The approximation is controlled by two +/// parameters: +/// +/// * `ratio`: The [nearest neighbor distance ratio](https://en.wikipedia.org/wiki/Nearest_neighbor_search#Nearest_neighbor_distance_ratio), +/// which controls how much closer new candidates must be to be considered. For example, a ratio +/// of 2.0 means that a neighbor must be less than half of the current threshold away to be +/// considered. A ratio of 1.0 means an exact search. +/// +/// * `limit`: A limit on the number of candidates to consider. Typical nearest neighbor algorithms +/// find a close match quickly, so setting a limit bounds the worst-case search time while keeping +/// good accuracy. +#[derive(Debug)] +pub struct ApproximateSearch { + inner: T, + ratio: f64, + limit: usize, +} + +impl ApproximateSearch { + /// Create a new ApproximateSearch index. + /// + /// * `inner`: The [NearestNeighbors] implementation to wrap. + /// * `ratio`: The nearest neighbor distance ratio. + /// * `limit`: The maximum number of results to consider. + pub fn new(inner: T, ratio: f64, limit: usize) -> Self { + Self { + inner, + ratio, + limit, + } + } +} + +impl NearestNeighbors for ApproximateSearch +where + U: Metric, + V: NearestNeighbors, +{ + fn search<'a, 'b, N>(&'a self, neighborhood: N) -> N + where + T: 'a, + U: 'b, + N: Neighborhood<&'a T, &'b U>, + { + self.inner + .search(ApproximateNeighborhood::new( + neighborhood, + self.ratio, + self.limit, + )) + .inner + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::metric::kd::KdTree; + use crate::metric::tests::test_nearest_neighbors; + use crate::metric::vp::VpTree; + + use std::iter::FromIterator; + + #[test] + fn test_approx_kd_tree() { + test_nearest_neighbors(|iter| { + ApproximateSearch::new(KdTree::from_iter(iter), 1.0, std::usize::MAX) + }); + } + + #[test] + fn test_approx_vp_tree() { + test_nearest_neighbors(|iter| { + ApproximateSearch::new(VpTree::from_iter(iter), 1.0, std::usize::MAX) + }); + } +} -- cgit v1.2.3 From 232995aecf809309848b864a77e9d968c6185a29 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sat, 2 May 2020 13:43:16 -0400 Subject: color: Implement color spaces --- Cargo.lock | 296 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 1 + src/color.rs | 282 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 1 + 4 files changed, 580 insertions(+) create mode 100644 src/color.rs diff --git a/Cargo.lock b/Cargo.lock index 45602e1..7f81f15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,17 +1,119 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. +[[package]] +name = "adler32" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d2e7343e7fc9de883d1b0341e0b13970f764c14101234857d2ddafa1cb1cac2" + [[package]] name = "autocfg" version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8aac770f1885fd7e387acedd76065302551364496e46b3dd00860b2f8359b9d" +[[package]] +name = "bitflags" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf1de2fe8c75bc145a2f577add951f8134889b4795d47466a54a5c846d691693" + +[[package]] +name = "bytemuck" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37fa13df2292ecb479ec23aa06f4507928bef07839be9ef15281411076629431" + +[[package]] +name = "byteorder" +version = "1.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c48aae112d48ed9f069b33538ea9e3e90aa263cfa3d1c24309612b1f7472de" + [[package]] name = "cfg-if" version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4785bdd1c96b2a846b2bd7cc02e86b6b3dbf14e7e53446c4f54c92a361040822" +[[package]] +name = "color_quant" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dbbb57365263e881e805dc77d94697c9118fd94d8da011240555aa7b23445bd" + +[[package]] +name = "crc32fast" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba125de2af0df55319f41944744ad91c71113bf74a4646efff39afe1f6842db1" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-deque" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f02af974daeee82218205558e51ec8768b48cf524bd01d550abe5573a608285" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", + "maybe-uninit", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "058ed274caafc1f60c4997b5fc07bf7dc7cca454af7c6e81edffe5f33f70dace" +dependencies = [ + "autocfg", + "cfg-if", + "crossbeam-utils", + "lazy_static", + "maybe-uninit", + "memoffset", + "scopeguard", +] + +[[package]] +name = "crossbeam-queue" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c695eeca1e7173472a32221542ae469b3e9aac3a4fc81f7696bcad82029493db" +dependencies = [ + "cfg-if", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3c7c73a2d1e9fc0886a08b93e98eb643461230d5f1925e4036204d5f2e261a8" +dependencies = [ + "autocfg", + "cfg-if", + "lazy_static", +] + +[[package]] +name = "deflate" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e5d2a2273fed52a7f947ee55b092c4057025d7a3e04e5ecdbd25d6c3fb1bd7" +dependencies = [ + "adler32", + "byteorder", +] + +[[package]] +name = "either" +version = "1.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb1f6b1ce1c140482ea30ddd3335fc0024ac7ee112895426e0a629a6c20adfe3" + [[package]] name = "getrandom" version = "0.1.14" @@ -23,20 +125,145 @@ dependencies = [ "wasi", ] +[[package]] +name = "gif" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "471d90201b3b223f3451cd4ad53e34295f16a1df17b1edf3736d47761c3981af" +dependencies = [ + "color_quant", + "lzw", +] + +[[package]] +name = "hermit-abi" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61565ff7aaace3525556587bd2dc31d4a07071957be715e63ce7b1eccf51a8f4" +dependencies = [ + "libc", +] + +[[package]] +name = "image" +version = "0.23.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9117f4167a8f21fa2bb3f17a652a760acd7572645281c98e3b612a26242c96ee" +dependencies = [ + "bytemuck", + "byteorder", + "gif", + "jpeg-decoder", + "num-iter", + "num-rational", + "num-traits", + "png", + "scoped_threadpool", + "tiff", +] + +[[package]] +name = "inflate" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1cdb29978cc5797bd8dcc8e5bf7de604891df2a8dc576973d71a281e916db2ff" +dependencies = [ + "adler32", +] + +[[package]] +name = "jpeg-decoder" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b47b4c4e017b01abdc5bcc126d2d1002e5a75bbe3ce73f9f4f311a916363704" +dependencies = [ + "byteorder", + "rayon", +] + [[package]] name = "kd-forest" version = "2.0.0" dependencies = [ + "image", "ordered-float", "rand", ] +[[package]] +name = "lazy_static" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" + [[package]] name = "libc" version = "0.2.69" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "99e85c08494b21a9054e7fe1374a732aeadaff3980b6990b94bfd3a70f690005" +[[package]] +name = "lzw" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d947cbb889ed21c2a84be6ffbaebf5b4e0f4340638cba0444907e38b56be084" + +[[package]] +name = "maybe-uninit" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60302e4db3a61da70c0cb7991976248362f30319e88850c487b9b95bbf059e00" + +[[package]] +name = "memoffset" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4fc2c02a7e374099d4ee95a193111f72d2110197fe200272371758f6c3643d8" +dependencies = [ + "autocfg", +] + +[[package]] +name = "miniz_oxide" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa679ff6578b1cddee93d7e82e263b94a575e0bfced07284eb0c037c1d2416a5" +dependencies = [ + "adler32", +] + +[[package]] +name = "num-integer" +version = "0.1.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f6ea62e9d81a77cd3ee9a2a5b9b609447857f3d358704331e4ef39eb247fcba" +dependencies = [ + "autocfg", + "num-traits", +] + +[[package]] +name = "num-iter" +version = "0.1.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfb0800a0291891dd9f4fe7bd9c19384f98f7fbe0cd0f39a2c6b88b9868bbc00" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c000134b5dbf44adc5cb772486d335293351644b801551abe8f75c84cfa4aef" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.11" @@ -46,6 +273,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05499f3756671c15885fee9034446956fff3f243d6077b91e5767df161f766b3" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "ordered-float" version = "1.0.2" @@ -55,6 +292,18 @@ dependencies = [ "num-traits", ] +[[package]] +name = "png" +version = "0.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c68a431ed29933a4eb5709aca9800989758c97759345860fa5db3cfced0b65d" +dependencies = [ + "bitflags", + "crc32fast", + "deflate", + "inflate", +] + [[package]] name = "ppv-lite86" version = "0.2.6" @@ -102,6 +351,53 @@ dependencies = [ "rand_core", ] +[[package]] +name = "rayon" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db6ce3297f9c85e16621bb8cca38a06779ffc31bb8184e1be4bed2be4678a098" +dependencies = [ + "crossbeam-deque", + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08a89b46efaf957e52b18062fb2f4660f8b8a4dde1807ca002690868ef2c85a9" +dependencies = [ + "crossbeam-deque", + "crossbeam-queue", + "crossbeam-utils", + "lazy_static", + "num_cpus", +] + +[[package]] +name = "scoped_threadpool" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d51f5df5af43ab3f1360b429fa5e0152ac5ce8c0bd6485cae490332e96846a8" + +[[package]] +name = "scopeguard" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" + +[[package]] +name = "tiff" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "002351e428db1eb1d8656d4ca61947c3519ac3191e1c804d4600cd32093b77ad" +dependencies = [ + "byteorder", + "lzw", + "miniz_oxide", +] + [[package]] name = "wasi" version = "0.9.0+wasi-snapshot-preview1" diff --git a/Cargo.toml b/Cargo.toml index 5a93cf5..902c3a3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,5 +5,6 @@ authors = ["Tavian Barnes "] edition = "2018" [dependencies] +image = "0.23.4" ordered-float = "1.0.2" rand = "0.7.3" diff --git a/src/color.rs b/src/color.rs new file mode 100644 index 0000000..a062795 --- /dev/null +++ b/src/color.rs @@ -0,0 +1,282 @@ +//! Colors and color spaces. + +use crate::metric::kd::{Cartesian, CartesianMetric}; +use crate::metric::{Metric, SquaredDistance}; + +use image::Rgb; + +use std::ops::Index; + +/// An 8-bit RGB color. +pub type Rgb8 = Rgb; + +/// A [color space](https://en.wikipedia.org/wiki/Color_space). +pub trait ColorSpace: Copy + From + CartesianMetric { + /// Compute the average of the given colors. + fn average>(colors: I) -> Self; +} + +/// [sRGB](https://en.wikipedia.org/wiki/SRGB) space. +#[derive(Clone, Copy, Debug)] +pub struct RgbSpace([f64; 3]); + +impl Index for RgbSpace { + type Output = f64; + + fn index(&self, i: usize) -> &f64 { + &self.0[i] + } +} + +impl From for RgbSpace { + fn from(rgb8: Rgb8) -> Self { + Self([ + (rgb8[0] as f64) / 255.0, + (rgb8[1] as f64) / 255.0, + (rgb8[2] as f64) / 255.0, + ]) + } +} + +impl Metric<[f64]> for RgbSpace { + type Distance = SquaredDistance; + + fn distance(&self, other: &[f64]) -> Self::Distance { + self.0.distance(other) + } +} + +impl Metric for RgbSpace { + type Distance = SquaredDistance; + + fn distance(&self, other: &Self) -> Self::Distance { + self.0.distance(&other.0) + } +} + +impl Cartesian for RgbSpace { + fn dimensions(&self) -> usize { + self.0.dimensions() + } + + fn coordinate(&self, i: usize) -> f64 { + self.0.coordinate(i) + } +} + +impl ColorSpace for RgbSpace { + fn average>(colors: I) -> Self { + let mut sum = [0.0, 0.0, 0.0]; + let mut len: usize = 0; + for color in colors.into_iter() { + for i in 0..3 { + sum[i] += color[i]; + } + len += 1; + } + for i in 0..3 { + sum[i] /= len as f64; + } + Self(sum) + } +} + +/// [CIE XYZ](https://en.wikipedia.org/wiki/CIE_1931_color_space) space. +#[derive(Clone, Copy, Debug)] +struct XyzSpace([f64; 3]); + +impl Index for XyzSpace { + type Output = f64; + + fn index(&self, i: usize) -> &f64 { + &self.0[i] + } +} + +/// The inverse of the sRGB gamma function. +fn srgb_inv_gamma(t: f64) -> f64 { + if t <= 0.040449936 { + t / 12.92 + } else { + ((t + 0.055) / 1.055).powf(2.4) + } +} + +impl From for XyzSpace { + fn from(rgb8: Rgb8) -> Self { + let rgb = RgbSpace::from(rgb8); + + let r = srgb_inv_gamma(rgb[0]); + let g = srgb_inv_gamma(rgb[1]); + let b = srgb_inv_gamma(rgb[2]); + + Self([ + 0.4123808838268995 * r + 0.3575728355732478 * g + 0.1804522977447919 * b, + 0.2126198631048975 * r + 0.7151387878413206 * g + 0.0721499433963131 * b, + 0.0193434956789248 * r + 0.1192121694056356 * g + 0.9505065664127130 * b, + ]) + } +} + +/// CIE D50 [white point](https://en.wikipedia.org/wiki/Standard_illuminant). +const WHITE: XyzSpace = XyzSpace([0.9504060171449392, 0.9999085943425312, 1.089062231497274]); + +/// CIE L\*a\*b\* (and L\*u\*v\*) gamma +fn lab_gamma(t: f64) -> f64 { + if t > 216.0 / 24389.0 { + t.cbrt() + } else { + 841.0 * t / 108.0 + 4.0 / 29.0 + } +} + +/// [CIE L\*a\*b\*](https://en.wikipedia.org/wiki/CIELAB_color_space) space. +#[derive(Clone, Copy, Debug)] +pub struct LabSpace([f64; 3]); + +impl Index for LabSpace { + type Output = f64; + + fn index(&self, i: usize) -> &f64 { + &self.0[i] + } +} + +impl From for LabSpace { + fn from(rgb8: Rgb8) -> Self { + let xyz = XyzSpace::from(rgb8); + + let x = lab_gamma(xyz[0] / WHITE[0]); + let y = lab_gamma(xyz[1] / WHITE[1]); + let z = lab_gamma(xyz[2] / WHITE[2]); + + let l = 116.0 * y - 16.0; + let a = 500.0 * (x - y); + let b = 200.0 * (y - z); + + Self([l, a, b]) + } +} + +impl Metric<[f64]> for LabSpace { + type Distance = SquaredDistance; + + fn distance(&self, other: &[f64]) -> Self::Distance { + self.0.distance(other) + } +} + +impl Metric for LabSpace { + type Distance = SquaredDistance; + + fn distance(&self, other: &Self) -> Self::Distance { + self.0.distance(&other.0) + } +} + +impl Cartesian for LabSpace { + fn dimensions(&self) -> usize { + self.0.dimensions() + } + + fn coordinate(&self, i: usize) -> f64 { + self.0.coordinate(i) + } +} + +impl ColorSpace for LabSpace { + fn average>(colors: I) -> Self { + let mut sum = [0.0, 0.0, 0.0]; + let mut len: usize = 0; + for color in colors.into_iter() { + for i in 0..3 { + sum[i] += color[i]; + } + len += 1; + } + for i in 0..3 { + sum[i] /= len as f64; + } + Self(sum) + } +} + +/// [CIE L\*u\*v\*](https://en.wikipedia.org/wiki/CIELUV) space. +#[derive(Clone, Copy, Debug)] +pub struct LuvSpace([f64; 3]); + +impl Index for LuvSpace { + type Output = f64; + + fn index(&self, i: usize) -> &f64 { + &self.0[i] + } +} + +/// Computes the u' and v' values for L\*u\*v\*. +fn uv_prime(xyz: &XyzSpace) -> (f64, f64) { + let denom = xyz[0] + 15.0 * xyz[1] + 3.0 * xyz[2]; + if denom == 0.0 { + (0.0, 0.0) + } else { + (4.0 * xyz[0] / denom, 9.0 * xyz[1] / denom) + } +} + +impl From for LuvSpace { + fn from(rgb8: Rgb8) -> Self { + let xyz = XyzSpace::from(rgb8); + + let (uprime, vprime) = uv_prime(&xyz); + let (unprime, vnprime) = uv_prime(&WHITE); + + let l = 116.0 * lab_gamma(xyz[1] / WHITE[1]) - 16.0; + let u = 13.0 * l * (uprime - unprime); + let v = 13.0 * l * (vprime - vnprime); + + Self([l, u, v]) + } +} + +impl Metric<[f64]> for LuvSpace { + type Distance = SquaredDistance; + + fn distance(&self, other: &[f64]) -> Self::Distance { + self.0.distance(other) + } +} + +impl Metric for LuvSpace { + type Distance = SquaredDistance; + + fn distance(&self, other: &Self) -> Self::Distance { + self.0.distance(&other.0) + } +} + +impl Cartesian for LuvSpace { + fn dimensions(&self) -> usize { + self.0.dimensions() + } + + fn coordinate(&self, i: usize) -> f64 { + self.0.coordinate(i) + } +} + +impl ColorSpace for LuvSpace { + fn average>(colors: I) -> Self { + let mut sum = [0.0, 0.0, 0.0]; + let mut len: usize = 0; + for color in colors.into_iter() { + for i in 0..3 { + sum[i] += color[i]; + } + len += 1; + } + for i in 0..3 { + sum[i] /= len as f64; + } + Self(sum) + } +} diff --git a/src/main.rs b/src/main.rs index 0d7989b..a7bda67 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +pub mod color; pub mod metric; fn main() {} -- cgit v1.2.3 From 5aea448bd2a7f6157cf33a6a2c044d043bbcd3ec Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sat, 2 May 2020 13:46:25 -0400 Subject: color/source: Implement color sources --- src/color.rs | 2 ++ src/color/source.rs | 76 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+) create mode 100644 src/color/source.rs diff --git a/src/color.rs b/src/color.rs index a062795..e0a3399 100644 --- a/src/color.rs +++ b/src/color.rs @@ -1,5 +1,7 @@ //! Colors and color spaces. +pub mod source; + use crate::metric::kd::{Cartesian, CartesianMetric}; use crate::metric::{Metric, SquaredDistance}; diff --git a/src/color/source.rs b/src/color/source.rs new file mode 100644 index 0000000..bd752d9 --- /dev/null +++ b/src/color/source.rs @@ -0,0 +1,76 @@ +//! Sources of colors. + +use super::Rgb8; + +use image::RgbImage; + +/// A source of colors in multidimensional space. +pub trait ColorSource { + /// Get the size of each dimension in this space. + fn dimensions(&self) -> &[usize]; + + /// Get the color at some particular coordinates. + fn get_color(&self, coords: &[usize]) -> Rgb8; +} + +/// The entire RGB space. +#[derive(Debug)] +pub struct AllColors { + dims: [usize; 3], + shifts: [usize; 3], +} + +impl AllColors { + /// Create an AllColors source with the given bit depth. + pub fn new(bits: usize) -> Self { + // Allocate bits from most to least perceptually important + let gbits = (bits + 2) / 3; + let rbits = (bits + 1) / 3; + let bbits = bits / 3; + + Self { + dims: [1 << rbits, 1 << gbits, 1 << bbits], + shifts: [8 - rbits, 8 - gbits, 8 - bbits], + } + } +} + +impl ColorSource for AllColors { + fn dimensions(&self) -> &[usize] { + &self.dims + } + + fn get_color(&self, coords: &[usize]) -> Rgb8 { + Rgb8::from([ + (coords[0] << self.shifts[0]) as u8, + (coords[1] << self.shifts[1]) as u8, + (coords[2] << self.shifts[2]) as u8, + ]) + } +} + +/// Colors extracted from an image. +#[derive(Debug)] +pub struct ImageColors { + dims: [usize; 2], + image: RgbImage, +} + +impl From for ImageColors { + fn from(image: RgbImage) -> Self { + Self { + dims: [image.width() as usize, image.height() as usize], + image: image, + } + } +} + +impl ColorSource for ImageColors { + fn dimensions(&self) -> &[usize] { + &self.dims + } + + fn get_color(&self, coords: &[usize]) -> Rgb8 { + *self.image.get_pixel(coords[0] as u32, coords[1] as u32) + } +} -- cgit v1.2.3 From 62e0fec044b5727efa1841138f44d9a1d9537bcf Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sat, 2 May 2020 13:48:07 -0400 Subject: color/order: Implement color orderings --- src/color.rs | 1 + src/color/order.rs | 196 +++++++++++++++++++++++++++++++++++++++++++++++++++++ src/hilbert.rs | 136 +++++++++++++++++++++++++++++++++++++ src/main.rs | 1 + 4 files changed, 334 insertions(+) create mode 100644 src/color/order.rs create mode 100644 src/hilbert.rs 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 { + source: S, + coords: Vec, +} + +impl From for ColorSourceIter { + fn from(source: S) -> Self { + let coords = vec![0; source.dimensions().len()]; + + Self { source, coords } + } +} + +impl Iterator for ColorSourceIter { + type Item = Rgb8; + + fn next(&mut self) -> Option { + 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 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 { + Some(self.cmp(other)) + } +} + +/// Iterate over colors sorted by their hue. +pub fn hue_sorted(source: S) -> Vec { + 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(source: S, rng: &mut R) -> Vec { + 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::() as u32; + nbits - (n - 1).leading_zeros() +} + +/// Iterate over colors in Morton order (Z-order). +pub fn morton(source: S) -> Vec { + 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(source: S) -> Vec { + 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) -> Vec { + 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() {} -- cgit v1.2.3 From 5ac571d8d16a307cea2922587185557bc773e8ed Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sat, 2 May 2020 13:52:28 -0400 Subject: frontier: New trait for choosing color locations --- src/frontier.rs | 18 ++++++++++++++++++ src/main.rs | 1 + 2 files changed, 19 insertions(+) create mode 100644 src/frontier.rs diff --git a/src/frontier.rs b/src/frontier.rs new file mode 100644 index 0000000..2c6f43a --- /dev/null +++ b/src/frontier.rs @@ -0,0 +1,18 @@ +//! Frontiers on which to place pixels. + +use crate::color::Rgb8; + +/// A frontier of pixels. +pub trait Frontier { + /// The width of the image. + fn width(&self) -> u32; + + /// The height of the image. + fn height(&self) -> u32; + + /// The number of pixels currently on the frontier. + fn len(&self) -> usize; + + /// Place the given color on the frontier, and return its position. + fn place(&mut self, rgb8: Rgb8) -> Option<(u32, u32)>; +} diff --git a/src/main.rs b/src/main.rs index a59a0cf..07e138a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,5 @@ pub mod color; +pub mod frontier; pub mod hilbert; pub mod metric; -- cgit v1.2.3 From da9b2ad1310e1db0ccb26a53181fa7f9b9033d04 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sat, 2 May 2020 13:56:09 -0400 Subject: frontier/image: Implement image frontiers --- src/frontier.rs | 71 +++++++++++++++++++++++++++++++++++++++++++++++- src/frontier/image.rs | 74 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 src/frontier/image.rs diff --git a/src/frontier.rs b/src/frontier.rs index 2c6f43a..1c5a173 100644 --- a/src/frontier.rs +++ b/src/frontier.rs @@ -1,6 +1,13 @@ //! Frontiers on which to place pixels. -use crate::color::Rgb8; +pub mod image; + +use crate::color::{ColorSpace, Rgb8}; +use crate::metric::kd::Cartesian; +use crate::metric::soft::SoftDelete; +use crate::metric::Metric; + +use std::cell::Cell; /// A frontier of pixels. pub trait Frontier { @@ -16,3 +23,65 @@ pub trait Frontier { /// Place the given color on the frontier, and return its position. fn place(&mut self, rgb8: Rgb8) -> Option<(u32, u32)>; } + +/// A pixel on a frontier. +#[derive(Debug)] +struct Pixel { + pos: (u32, u32), + color: C, + deleted: Cell, +} + +impl Pixel { + fn new(x: u32, y: u32, color: C) -> Self { + Self { + pos: (x, y), + color, + deleted: Cell::new(false), + } + } + + fn delete(&self) { + self.deleted.set(true); + } +} + +impl Metric> for C { + type Distance = C::Distance; + + fn distance(&self, other: &Pixel) -> Self::Distance { + self.distance(&other.color) + } +} + +impl> Metric<[f64]> for Pixel { + type Distance = C::Distance; + + fn distance(&self, other: &[f64]) -> Self::Distance { + self.color.distance(other) + } +} + +impl Metric for Pixel { + type Distance = C::Distance; + + fn distance(&self, other: &Pixel) -> Self::Distance { + self.color.distance(&other.color) + } +} + +impl Cartesian for Pixel { + fn dimensions(&self) -> usize { + self.color.dimensions() + } + + fn coordinate(&self, i: usize) -> f64 { + self.color.coordinate(i) + } +} + +impl SoftDelete for Pixel { + fn is_deleted(&self) -> bool { + self.deleted.get() + } +} diff --git a/src/frontier/image.rs b/src/frontier/image.rs new file mode 100644 index 0000000..3655580 --- /dev/null +++ b/src/frontier/image.rs @@ -0,0 +1,74 @@ +//! Frontier that targets an image. + +use super::{Frontier, Pixel}; + +use crate::color::{ColorSpace, Rgb8}; +use crate::metric::soft::SoftKdTree; +use crate::metric::NearestNeighbors; + +use image::RgbImage; + +/// A [Frontier] that places colors on the closest pixel of a target image. +#[derive(Debug)] +pub struct ImageFrontier { + nodes: SoftKdTree>, + width: u32, + height: u32, + len: usize, + deleted: usize, +} + +impl ImageFrontier { + /// Create an ImageFrontier from an image. + pub fn new(img: &RgbImage) -> Self { + let width = img.width(); + let height = img.height(); + let len = (width as usize) * (height as usize); + + Self { + nodes: img + .enumerate_pixels() + .map(|(x, y, p)| Pixel::new(x, y, C::from(*p))) + .collect(), + width, + height, + len, + deleted: 0, + } + } +} + +impl Frontier for ImageFrontier { + fn width(&self) -> u32 { + self.width + } + + fn height(&self) -> u32 { + self.height + } + + fn len(&self) -> usize { + self.len - self.deleted + } + + fn place(&mut self, rgb8: Rgb8) -> Option<(u32, u32)> { + let color = C::from(rgb8); + + if let Some(node) = self.nodes.nearest(&color).map(|n| n.item) { + let pos = node.pos; + + node.delete(); + self.deleted += 1; + + if 32 * self.deleted >= self.len { + self.nodes.rebuild(); + self.len -= self.deleted; + self.deleted = 0; + } + + Some(pos) + } else { + None + } + } +} -- cgit v1.2.3 From dd2336f76bfaff01ccb5b57f4453629ff62016b3 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sat, 2 May 2020 13:57:30 -0400 Subject: frontier/min: Implement min selection --- src/frontier.rs | 61 +++++++++++++++++++++ src/frontier/min.rs | 150 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 211 insertions(+) create mode 100644 src/frontier/min.rs diff --git a/src/frontier.rs b/src/frontier.rs index 1c5a173..e1b6528 100644 --- a/src/frontier.rs +++ b/src/frontier.rs @@ -1,6 +1,7 @@ //! Frontiers on which to place pixels. pub mod image; +pub mod min; use crate::color::{ColorSpace, Rgb8}; use crate::metric::kd::Cartesian; @@ -8,6 +9,7 @@ use crate::metric::soft::SoftDelete; use crate::metric::Metric; use std::cell::Cell; +use std::rc::Rc; /// A frontier of pixels. pub trait Frontier { @@ -85,3 +87,62 @@ impl SoftDelete for Pixel { self.deleted.get() } } + +impl> Metric<[f64]> for Rc> { + type Distance = C::Distance; + + fn distance(&self, other: &[f64]) -> Self::Distance { + (**self).distance(other) + } +} + +impl Metric>> for C { + type Distance = C::Distance; + + fn distance(&self, other: &Rc>) -> Self::Distance { + self.distance(&other.color) + } +} + +impl Metric for Rc> { + type Distance = C::Distance; + + fn distance(&self, other: &Self) -> Self::Distance { + (**self).distance(&**other) + } +} + +impl Cartesian for Rc> { + fn dimensions(&self) -> usize { + (**self).dimensions() + } + + fn coordinate(&self, i: usize) -> f64 { + (**self).coordinate(i) + } +} + +impl SoftDelete for Rc> { + fn is_deleted(&self) -> bool { + (**self).is_deleted() + } +} + +/// Return all the neighbors of a pixel location. +fn neighbors(x: u32, y: u32) -> [(u32, u32); 8] { + let xm1 = x.wrapping_sub(1); + let ym1 = y.wrapping_sub(1); + let xp1 = x + 1; + let yp1 = y + 1; + + [ + (xm1, ym1), + (xm1, y), + (xm1, yp1), + (x, ym1), + (x, yp1), + (xp1, ym1), + (xp1, y), + (xp1, yp1), + ] +} diff --git a/src/frontier/min.rs b/src/frontier/min.rs new file mode 100644 index 0000000..b22b290 --- /dev/null +++ b/src/frontier/min.rs @@ -0,0 +1,150 @@ +//! Minimum selection frontier. + +use super::{neighbors, Frontier, Pixel}; + +use crate::color::{ColorSpace, Rgb8}; +use crate::metric::soft::SoftKdForest; +use crate::metric::NearestNeighbors; + +use rand::Rng; + +use std::rc::Rc; + +/// A pixel on a min frontier. +#[derive(Debug)] +struct MinPixel { + pixel: Option>>, + filled: bool, +} + +impl MinPixel { + fn new() -> Self { + Self { + pixel: None, + filled: false, + } + } +} + +/// A [Frontier] that places colors on a neighbor of the closest pixel so far. +#[derive(Debug)] +pub struct MinFrontier { + rng: R, + pixels: Vec>, + forest: SoftKdForest>>, + width: u32, + height: u32, + x0: u32, + y0: u32, + len: usize, + deleted: usize, +} + +impl MinFrontier { + /// Create a MinFrontier with the given dimensions and initial pixel location. + pub fn new(rng: R, width: u32, height: u32, x0: u32, y0: u32) -> Self { + let size = (width as usize) * (height as usize); + let mut pixels = Vec::with_capacity(size); + for _ in 0..size { + pixels.push(MinPixel::new()); + } + + Self { + rng, + pixels, + forest: SoftKdForest::new(), + width, + height, + x0, + y0, + len: 0, + deleted: 0, + } + } + + fn pixel_index(&self, x: u32, y: u32) -> usize { + debug_assert!(x < self.width); + debug_assert!(y < self.height); + + (x + y * self.width) as usize + } + + fn free_neighbor(&mut self, x: u32, y: u32) -> Option<(u32, u32)> { + // Pick a pseudo-random neighbor + let offset: usize = self.rng.gen(); + + let neighbors = neighbors(x, y); + for i in 0..8 { + let (x, y) = neighbors[(i + offset) % 8]; + if x < self.width && y < self.height { + let i = self.pixel_index(x, y); + if !self.pixels[i].filled { + return Some((x, y)); + } + } + } + + None + } + + fn fill(&mut self, x: u32, y: u32, color: C) -> Option<(u32, u32)> { + let i = self.pixel_index(x, y); + let pixel = &mut self.pixels[i]; + if pixel.filled { + return None; + } + + let rc = Rc::new(Pixel::new(x, y, color)); + pixel.pixel = Some(rc.clone()); + pixel.filled = true; + + if self.free_neighbor(x, y).is_some() { + self.forest.push(rc); + self.len += 1; + } + + for &(x, y) in &neighbors(x, y) { + if x < self.width && y < self.height && self.free_neighbor(x, y).is_none() { + let i = self.pixel_index(x, y); + if let Some(pixel) = self.pixels[i].pixel.take() { + pixel.delete(); + self.deleted += 1; + } + } + } + + if 2 * self.deleted >= self.len { + self.forest.rebuild(); + self.len -= self.deleted; + self.deleted = 0; + } + + Some((x, y)) + } +} + +impl Frontier for MinFrontier { + fn width(&self) -> u32 { + self.width + } + + fn height(&self) -> u32 { + self.height + } + + fn len(&self) -> usize { + self.len - self.deleted + } + + fn place(&mut self, rgb8: Rgb8) -> Option<(u32, u32)> { + let color = C::from(rgb8); + let (x, y) = self + .forest + .nearest(&color) + .map(|n| n.item.pos) + .map(|(x, y)| self.free_neighbor(x, y).unwrap()) + .unwrap_or((self.x0, self.y0)); + + self.fill(x, y, color) + } +} -- cgit v1.2.3 From b58924b5a2534f845b996a2d58ab618301f1d450 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sat, 2 May 2020 13:58:14 -0400 Subject: frontier/mean: Implement mean selection --- src/frontier.rs | 1 + src/frontier/mean.rs | 140 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 141 insertions(+) create mode 100644 src/frontier/mean.rs diff --git a/src/frontier.rs b/src/frontier.rs index e1b6528..7143cb7 100644 --- a/src/frontier.rs +++ b/src/frontier.rs @@ -1,6 +1,7 @@ //! Frontiers on which to place pixels. pub mod image; +pub mod mean; pub mod min; use crate::color::{ColorSpace, Rgb8}; diff --git a/src/frontier/mean.rs b/src/frontier/mean.rs new file mode 100644 index 0000000..889c5ba --- /dev/null +++ b/src/frontier/mean.rs @@ -0,0 +1,140 @@ +//! Mean selection frontier. + +use super::{neighbors, Frontier, Pixel}; + +use crate::color::{ColorSpace, Rgb8}; +use crate::metric::soft::SoftKdForest; +use crate::metric::NearestNeighbors; + +use std::iter; +use std::rc::Rc; + +/// A pixel on a mean frontier. +#[derive(Debug)] +enum MeanPixel { + Empty, + Fillable(Rc>), + Filled(C), +} + +impl MeanPixel { + fn filled_color(&self) -> Option { + match self { + Self::Filled(color) => Some(*color), + _ => None, + } + } +} + +/// A [Frontier] that looks at the average color of each pixel's neighbors. +pub struct MeanFrontier { + pixels: Vec>, + forest: SoftKdForest>>, + width: u32, + height: u32, + len: usize, + deleted: usize, +} + +impl MeanFrontier { + /// Create a MeanFrontier with the given dimensions and initial pixel location. + pub fn new(width: u32, height: u32, x0: u32, y0: u32) -> Self { + let size = (width as usize) * (height as usize); + let mut pixels = Vec::with_capacity(size); + for _ in 0..size { + pixels.push(MeanPixel::Empty); + } + + let pixel0 = Rc::new(Pixel::new(x0, y0, C::from(Rgb8::from([0, 0, 0])))); + let i = (x0 + y0 * width) as usize; + pixels[i] = MeanPixel::Fillable(pixel0.clone()); + + Self { + pixels, + forest: iter::once(pixel0).collect(), + width, + height, + len: 1, + deleted: 0, + } + } + + fn pixel_index(&self, x: u32, y: u32) -> usize { + debug_assert!(x < self.width); + debug_assert!(y < self.height); + + (x + y * self.width) as usize + } + + fn fill(&mut self, x: u32, y: u32, color: C) { + let i = self.pixel_index(x, y); + match &self.pixels[i] { + MeanPixel::Empty => {} + MeanPixel::Fillable(pixel) => { + pixel.delete(); + self.deleted += 1; + } + _ => unreachable!(), + } + self.pixels[i] = MeanPixel::Filled(color); + + let mut pixels = Vec::new(); + for &(x, y) in &neighbors(x, y) { + if x < self.width && y < self.height { + let i = self.pixel_index(x, y); + match &self.pixels[i] { + MeanPixel::Empty => {} + MeanPixel::Fillable(pixel) => { + pixel.delete(); + self.deleted += 1; + } + MeanPixel::Filled(_) => continue, + } + let color = C::average( + neighbors(x, y) + .iter() + .filter(|(x, y)| *x < self.width && *y < self.height) + .map(|(x, y)| self.pixel_index(*x, *y)) + .map(|i| &self.pixels[i]) + .map(MeanPixel::filled_color) + .flatten(), + ); + let pixel = Rc::new(Pixel::new(x, y, color)); + self.pixels[i] = MeanPixel::Fillable(pixel.clone()); + pixels.push(pixel); + } + } + + self.len += pixels.len(); + self.forest.extend(pixels); + + if 2 * self.deleted >= self.len { + self.forest.rebuild(); + self.len -= self.deleted; + self.deleted = 0; + } + } +} + +impl Frontier for MeanFrontier { + fn width(&self) -> u32 { + self.width + } + + fn height(&self) -> u32 { + self.height + } + + fn len(&self) -> usize { + self.len - self.deleted + } + + fn place(&mut self, rgb8: Rgb8) -> Option<(u32, u32)> { + let color = C::from(rgb8); + let (x, y) = self.forest.nearest(&color).map(|n| n.item.pos)?; + + self.fill(x, y, color); + + Some((x, y)) + } +} -- cgit v1.2.3 From 2984e8f93fe88d0ee7eb3c0561dcd2da44807429 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Sat, 2 May 2020 13:58:45 -0400 Subject: main: Implement the main binary --- Cargo.lock | 192 +++++++++++++++++++++++++++++ Cargo.toml | 3 + src/main.rs | 397 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 591 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 7f81f15..c066dd8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6,18 +6,67 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d2e7343e7fc9de883d1b0341e0b13970f764c14101234857d2ddafa1cb1cac2" +[[package]] +name = "ansi_term" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee49baf6cb617b853aa8d93bf420db2383fab46d314482ca2803b40d5fde979b" +dependencies = [ + "winapi", +] + +[[package]] +name = "arrayref" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4c527152e37cf757a3f78aae5a06fbeefdb07ccc535c980a3208ee3060dd544" + +[[package]] +name = "arrayvec" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cff77d8686867eceff3105329d4698d96c2391c176d5d03adc90c7389162b5b8" + +[[package]] +name = "atty" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" +dependencies = [ + "hermit-abi", + "libc", + "winapi", +] + [[package]] name = "autocfg" version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8aac770f1885fd7e387acedd76065302551364496e46b3dd00860b2f8359b9d" +[[package]] +name = "base64" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b41b7ea54a0c9d92199de89e20e58d49f02f8e699814ef3fdf266f6f748d15c7" + [[package]] name = "bitflags" version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf1de2fe8c75bc145a2f577add951f8134889b4795d47466a54a5c846d691693" +[[package]] +name = "blake2b_simd" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8fb2d74254a3a0b5cac33ac9f8ed0e44aa50378d9dbb2e5d83bd21ed1dc2c8a" +dependencies = [ + "arrayref", + "arrayvec", + "constant_time_eq", +] + [[package]] name = "bytemuck" version = "1.2.0" @@ -36,12 +85,33 @@ version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4785bdd1c96b2a846b2bd7cc02e86b6b3dbf14e7e53446c4f54c92a361040822" +[[package]] +name = "clap" +version = "2.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5067f5bb2d80ef5d68b4c87db81601f0b75bca627bc2ef76b141d7b846a3c6d9" +dependencies = [ + "ansi_term", + "atty", + "bitflags", + "strsim", + "textwrap", + "unicode-width", + "vec_map", +] + [[package]] name = "color_quant" version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0dbbb57365263e881e805dc77d94697c9118fd94d8da011240555aa7b23445bd" +[[package]] +name = "constant_time_eq" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc" + [[package]] name = "crc32fast" version = "1.2.0" @@ -108,6 +178,28 @@ dependencies = [ "byteorder", ] +[[package]] +name = "dirs" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13aea89a5c93364a98e9b37b2fa237effbb694d5cfe01c5b70941f7eb087d5e3" +dependencies = [ + "cfg-if", + "dirs-sys", +] + +[[package]] +name = "dirs-sys" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "afa0b23de8fd801745c471deffa6e12d248f962c9fd4b4c33787b055599bde7b" +dependencies = [ + "cfg-if", + "libc", + "redox_users", + "winapi", +] + [[package]] name = "either" version = "1.5.3" @@ -185,9 +277,12 @@ dependencies = [ name = "kd-forest" version = "2.0.0" dependencies = [ + "clap", "image", "ordered-float", "rand", + "rand_pcg", + "term", ] [[package]] @@ -351,6 +446,15 @@ dependencies = [ "rand_core", ] +[[package]] +name = "rand_pcg" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16abd0c1b639e9eb4d7c50c0b8100b0d0f849be2349829c740fe8e6eb4816429" +dependencies = [ + "rand_core", +] + [[package]] name = "rayon" version = "1.3.0" @@ -375,6 +479,35 @@ dependencies = [ "num_cpus", ] +[[package]] +name = "redox_syscall" +version = "0.1.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2439c63f3f6139d1b57529d16bc3b8bb855230c8efcc5d3a896c8bea7c3b1e84" + +[[package]] +name = "redox_users" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09b23093265f8d200fa7b4c2c76297f47e681c655f6f1285a8780d6a022f7431" +dependencies = [ + "getrandom", + "redox_syscall", + "rust-argon2", +] + +[[package]] +name = "rust-argon2" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2bc8af4bda8e1ff4932523b94d3dd20ee30a87232323eda55903ffd71d2fb017" +dependencies = [ + "base64", + "blake2b_simd", + "constant_time_eq", + "crossbeam-utils", +] + [[package]] name = "scoped_threadpool" version = "0.1.9" @@ -387,6 +520,31 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" +[[package]] +name = "strsim" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ea5119cdb4c55b55d432abb513a0429384878c15dde60cc77b1c99de1a95a6a" + +[[package]] +name = "term" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0863a3345e70f61d613eab32ee046ccd1bcc5f9105fe402c61fcd0c13eeb8b5" +dependencies = [ + "dirs", + "winapi", +] + +[[package]] +name = "textwrap" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d326610f408c7a4eb6f51c37c330e496b08506c9457c9d34287ecc38809fb060" +dependencies = [ + "unicode-width", +] + [[package]] name = "tiff" version = "0.4.0" @@ -398,8 +556,42 @@ dependencies = [ "miniz_oxide", ] +[[package]] +name = "unicode-width" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caaa9d531767d1ff2150b9332433f32a24622147e5ebb1f26409d5da67afd479" + +[[package]] +name = "vec_map" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05c78687fb1a80548ae3250346c3db86a80a7cdd77bda190189f2d0a0987c81a" + [[package]] name = "wasi" version = "0.9.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519" + +[[package]] +name = "winapi" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8093091eeb260906a183e6ae1abdba2ef5ef2257a21801128899c3fc699229c6" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" diff --git a/Cargo.toml b/Cargo.toml index 902c3a3..de65bb7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,9 @@ authors = ["Tavian Barnes "] edition = "2018" [dependencies] +clap = "2.33.0" image = "0.23.4" ordered-float = "1.0.2" rand = "0.7.3" +rand_pcg = "0.2.1" +term = "0.6.1" diff --git a/src/main.rs b/src/main.rs index 07e138a..f016b4c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,4 +3,399 @@ pub mod frontier; pub mod hilbert; pub mod metric; -fn main() {} +use crate::color::source::{AllColors, ColorSource, ImageColors}; +use crate::color::{order, ColorSpace, LabSpace, LuvSpace, Rgb8, RgbSpace}; +use crate::frontier::image::ImageFrontier; +use crate::frontier::mean::MeanFrontier; +use crate::frontier::min::MinFrontier; +use crate::frontier::Frontier; + +use clap::{self, clap_app, crate_authors, crate_name, crate_version, AppSettings}; + +use image::{self, Rgba, RgbaImage}; + +use rand::SeedableRng; +use rand_pcg::Pcg64; + +use term; + +use std::cmp; +use std::error::Error; +use std::fs; +use std::io::{self, Write}; +use std::path::PathBuf; +use std::str::FromStr; +use std::time::Instant; + +/// The color source specified on the command line. +#[derive(Debug, Eq, PartialEq)] +enum SourceArg { + /// All RGB colors of the given bit depth. + AllRgb(u32), + /// Take the colors from an image. + Image(PathBuf), +} + +/// The order to process colors in. +#[derive(Debug, Eq, PartialEq)] +enum OrderArg { + /// Sorted by hue. + HueSort, + /// Shuffled randomly. + Random, + /// Morton/Z-order. + Morton, + /// Hilbert curve order. + Hilbert, +} + +/// The frontier implementation. +#[derive(Debug, Eq, PartialEq)] +enum FrontierArg { + /// Pick a neighbor of the closest pixel so far. + Min, + /// Pick the pixel with the closest mean color of all its neighbors. + Mean, + /// Target the closest pixel on an image. + Image(PathBuf), +} + +/// The color space to operate in. +#[derive(Debug, Eq, PartialEq)] +enum ColorSpaceArg { + /// sRGB space. + Rgb, + /// CIE L*a*b* space. + Lab, + /// CIE L*u*v* space. + Luv, +} + +/// Parse an argument into the appropriate type. +fn parse_arg(arg: Option<&str>) -> clap::Result> +where + F: FromStr, + F::Err: Error, +{ + match arg.map(|s| s.parse()) { + Some(Ok(f)) => Ok(Some(f)), + Some(Err(e)) => Err(clap::Error::with_description( + &e.to_string(), + clap::ErrorKind::InvalidValue, + )), + None => Ok(None), + } +} + +/// The parsed command line arguments. +#[derive(Debug)] +struct Args { + source: SourceArg, + order: OrderArg, + stripe: bool, + frontier: FrontierArg, + space: ColorSpaceArg, + width: Option, + height: Option, + x0: Option, + y0: Option, + animate: bool, + output: PathBuf, + seed: u64, +} + +impl Args { + fn parse() -> clap::Result { + let args = clap_app!((crate_name!()) => + (version: crate_version!()) + (author: crate_authors!()) + (setting: AppSettings::ColoredHelp) + (@group source => + (@arg DEPTH: -b --("bit-depth") +takes_value "Use all DEPTH-bit colors") + (@arg INPUT: -i --input +takes_value "Use colors from the INPUT image") + ) + (@group order => + (@arg HUE: -s --hue-sort "Sort colors by hue [default]") + (@arg RANDOM: -r --random "Randomize colors") + (@arg MORTON: -M --morton "Place colors in Morton order (Z-order)") + (@arg HILBERT: -H --hilbert "Place colors in Hilbert curve order") + ) + (@group stripe => + (@arg STRIPE: -t --stripe "Reduce artifacts by iterating through the colors in multiple stripes [default]") + (@arg NOSTRIPE: -T --("no-stripe") "Don't stripe") + ) + (@group frontier => + (@arg MODE: -l --selection +takes_value possible_value[min mean] "Specify the selection mode") + (@arg TARGET: -g --target +takes_value "Place colors on the closest pixels of the TARGET image") + ) + (@arg SPACE: -c --("color-space") default_value("Lab") possible_value[RGB Lab Luv] "Use the given color space") + (@arg WIDTH: -w --width +takes_value "The width of the generated image") + (@arg HEIGHT: -h --height +takes_value "The height of the generated image") + (@arg X: -x +takes_value "The x coordinate of the first pixel") + (@arg Y: -y +takes_value "The y coordinate of the first pixel") + (@arg ANIMATE: -a --animate "Generate frames of an animation") + (@arg PATH: -o --output default_value("kd-forest.png") "Save the image to PATH") + (@arg SEED: -e --seed default_value("0") "Seed the random number generator") + ) + .get_matches_safe()?; + + let source = if let Some(input) = args.value_of("INPUT") { + SourceArg::Image(PathBuf::from(input)) + } else { + SourceArg::AllRgb(parse_arg(args.value_of("DEPTH"))?.unwrap_or(24)) + }; + + let order = if args.is_present("RANDOM") { + OrderArg::Random + } else if args.is_present("MORTON") { + OrderArg::Morton + } else if args.is_present("HILBERT") { + OrderArg::Hilbert + } else { + OrderArg::HueSort + }; + + let stripe = !args.is_present("NOSTRIPE") && order != OrderArg::Random; + + let frontier = if let Some(target) = args.value_of("TARGET") { + FrontierArg::Image(PathBuf::from(target)) + } else { + match args.value_of("MODE") { + Some("min") | None => FrontierArg::Min, + Some("mean") => FrontierArg::Mean, + _ => unreachable!(), + } + }; + + let space = match args.value_of("SPACE").unwrap() { + "RGB" => ColorSpaceArg::Rgb, + "Lab" => ColorSpaceArg::Lab, + "Luv" => ColorSpaceArg::Luv, + _ => unreachable!(), + }; + + let width = parse_arg(args.value_of("WIDTH"))?; + let height = parse_arg(args.value_of("HEIGHT"))?; + let x0 = parse_arg(args.value_of("X"))?; + let y0 = parse_arg(args.value_of("Y"))?; + + let animate = args.is_present("ANIMATE"); + + let mut path = args.value_of("PATH").unwrap(); + if animate && args.occurrences_of("PATH") == 0 { + path = "kd-frames"; + } + let output = PathBuf::from(path); + + let seed = parse_arg(args.value_of("SEED"))?.unwrap_or(0); + + Ok(Self { + source, + order, + stripe, + frontier, + space, + width, + height, + x0, + y0, + animate, + output, + seed, + }) + } +} + +/// main() return type. +type MainResult = Result<(), Box>; + +/// The kd-forest application itself. +#[derive(Debug)] +struct App { + args: Args, + rng: Pcg64, + width: Option, + height: Option, + start_time: Instant, +} + +impl App { + /// Make the App. + fn new(args: Args) -> Self { + let rng = Pcg64::seed_from_u64(args.seed); + let width = args.width; + let height = args.height; + let start_time = Instant::now(); + + Self { + args, + rng, + width, + height, + start_time, + } + } + + fn run(&mut self) -> MainResult { + let colors = match self.args.source { + SourceArg::AllRgb(depth) => { + self.width.get_or_insert(1u32 << ((depth + 1) / 2)); + self.height.get_or_insert(1u32 << (depth / 2)); + self.get_colors(AllColors::new(depth as usize)) + } + SourceArg::Image(ref path) => { + let img = image::open(path)?.into_rgb(); + self.width.get_or_insert(img.width()); + self.height.get_or_insert(img.height()); + self.get_colors(ImageColors::from(img)) + } + }; + + match self.args.space { + ColorSpaceArg::Rgb => self.paint::(colors), + ColorSpaceArg::Lab => self.paint::(colors), + ColorSpaceArg::Luv => self.paint::(colors), + } + } + + fn get_colors(&mut self, source: S) -> Vec { + let colors = match self.args.order { + OrderArg::HueSort => order::hue_sorted(source), + OrderArg::Random => order::shuffled(source, &mut self.rng), + OrderArg::Morton => order::morton(source), + OrderArg::Hilbert => order::hilbert(source), + }; + + if self.args.stripe { + order::striped(colors) + } else { + colors + } + } + + fn paint(&mut self, colors: Vec) -> MainResult { + let width = self.width.unwrap(); + let height = self.height.unwrap(); + let x0 = self.args.x0.unwrap_or(width / 2); + let y0 = self.args.x0.unwrap_or(height / 2); + + match &self.args.frontier { + FrontierArg::Image(ref path) => { + let img = image::open(path)?.into_rgb(); + self.paint_on(colors, ImageFrontier::::new(&img)) + } + FrontierArg::Min => { + let rng = Pcg64::from_rng(&mut self.rng)?; + self.paint_on(colors, MinFrontier::::new(rng, width, height, x0, y0)) + } + FrontierArg::Mean => { + self.paint_on(colors, MeanFrontier::::new(width, height, x0, y0)) + } + } + } + + fn paint_on(&mut self, colors: Vec, mut frontier: F) -> MainResult { + let width = frontier.width(); + let height = frontier.height(); + let mut output = RgbaImage::new(width, height); + + let size = cmp::min((width * height) as usize, colors.len()); + println!("Generating a {}x{} image ({} pixels)", width, height, size); + + if self.args.animate { + fs::create_dir_all(&self.args.output)?; + output.save(&self.args.output.join("0000.png"))?; + } + + let interval = cmp::max(width, height) as usize; + + let mut max_frontier = frontier.len(); + + for (i, color) in colors.into_iter().enumerate() { + let pos = frontier.place(color); + if pos.is_none() { + break; + } + + let (x, y) = pos.unwrap(); + let rgba = Rgba([color[0], color[1], color[2], 255]); + output.put_pixel(x, y, rgba); + + max_frontier = cmp::max(max_frontier, frontier.len()); + + if (i + 1) % interval == 0 { + if self.args.animate { + let frame = (i + 1) / interval; + output.save(&self.args.output.join(format!("{:04}.png", frame)))?; + } + + if i + 1 < size { + self.print_progress(i + 1, size, frontier.len())?; + } + } + } + + if self.args.animate && size % interval != 0 { + let frame = size / interval; + output.save(&self.args.output.join(format!("{:04}.png", frame)))?; + } + + self.print_progress(size, size, max_frontier)?; + + if !self.args.animate { + output.save(&self.args.output)?; + } + + Ok(()) + } + + fn print_progress(&self, i: usize, size: usize, frontier_len: usize) -> io::Result<()> { + let mut term = match term::stderr() { + Some(term) => term, + None => return Ok(()), + }; + + let progress = 100.0 * (i as f64) / (size as f64); + let mut rate = (i as f64) / self.start_time.elapsed().as_secs_f64(); + let mut unit = "px/s"; + + if rate >= 10_000.0 { + rate /= 1_000.0; + unit = "Kpx/s"; + } + + if rate >= 10_000.0 { + rate /= 1_000.0; + unit = "Mpx/s"; + } + + if rate >= 10_000.0 { + rate /= 1_000.0; + unit = "Gpx/s"; + } + + let (frontier_label, newline) = if i == size { + ("max frontier size", "\n") + } else { + ("frontier size", "") + }; + + term.carriage_return()?; + term.delete_line()?; + + write!( + term, + "{:>6.2}% | {:4.0} {:>5} | {}: {}{}", + progress, rate, unit, frontier_label, frontier_len, newline, + ) + } +} + +fn main() -> MainResult { + let args = match Args::parse() { + Ok(args) => args, + Err(e) => e.exit(), + }; + + App::new(args).run() +} -- cgit v1.2.3