]> git.sesse.net Git - narabu/commitdiff
Add some code for (semi-)optimal assignment of rANS coefficients to streams.
authorSteinar H. Gunderson <sgunderson@bigfoot.com>
Sun, 17 Sep 2017 09:06:32 +0000 (11:06 +0200)
committerSteinar H. Gunderson <sgunderson@bigfoot.com>
Sun, 17 Sep 2017 09:09:36 +0000 (11:09 +0200)
qdc.cpp

diff --git a/qdc.cpp b/qdc.cpp
index a0fe2e7b1a71a36f8f9fe0682cf710218669ca85..b5594b76856918d9b14bee3f55dd5bcec24fdd7a 100644 (file)
--- a/qdc.cpp
+++ b/qdc.cpp
@@ -9,14 +9,24 @@
 #include "ryg_rans/rans_byte.h"
 #include "ryg_rans/renormalize.h"
 
+#include <algorithm>
 #include <memory>
+#include <numeric>
+#include <random>
 #include <vector>
+#include <unordered_map>
 
 #define WIDTH 1280
 #define HEIGHT 720
 #define NUM_SYMS 256
 #define ESCAPE_LIMIT (NUM_SYMS - 1)
 
+// If you set this to 1, the program will try to optimize the placement
+// of coefficients to rANS probability distributions. This is randomized,
+// so you might want to run it a few times.
+#define FIND_OPTIMAL_STREAM_ASSIGNMENT 0
+#define NUM_CLUSTERS 8
+
 static constexpr uint32_t prob_bits = 12;
 static constexpr uint32_t prob_scale = 1 << prob_bits;
 
@@ -137,22 +147,19 @@ void SymbolStats::normalize_freqs(uint32_t target_total)
                 calc_cost, (calc_cost - ideal_cost) / 8.0, total_loss / 8.0, total_loss_with_dp / 8.0);
 }
 
-SymbolStats stats[64];
+SymbolStats stats[128];
+
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+// Distance from one stream to the other, based on a hacked-up K-L divergence.
+float kl_dist[64][64];
+#endif
 
-int pick_stats_for(int y, int x)
+int pick_stats_for(int x, int y, bool is_chroma)
 {
-       //return 0;
-       //return std::min<int>(hypot(x, y), 7);
-       return std::min<int>(x + y, 7);
-       //if (x + y >= 7) return 7;
-       //return x + y;
-       //return y * 8 + x;
-#if 0
-       if (y == 0 && x == 0) {
-               return 0;
-       } else {
-               return 1;
-       }
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+       return y * 8 + x + is_chroma * 64;
+#else
+       return std::min<int>(x + y, 7) + is_chroma * 8;
 #endif
 }
                
@@ -348,6 +355,114 @@ void convert_ycbcr()
        }
 }
 
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+double find_best_assignment(const int *medoids, int *assignment)
+{
+       double current_score = 0.0;
+       for (int i = 0; i < 64; ++i) {
+               int best_medoid = medoids[0];
+               float best_medoid_score = kl_dist[i][medoids[0]];
+               for (int j = 1; j < NUM_CLUSTERS; ++j) {
+                       if (kl_dist[i][medoids[j]] < best_medoid_score) {
+                               best_medoid = medoids[j];
+                               best_medoid_score = kl_dist[i][medoids[j]];
+                       }
+               }
+               assignment[i] = best_medoid;
+               current_score += best_medoid_score;
+       }
+       return current_score;
+}
+
+void find_optimal_stream_assignment(int base)
+{
+       double inv_sum[64];
+       for (unsigned i = 0; i < 64; ++i) {
+               double s = 0.0;
+               for (unsigned k = 0; k < 256; ++k) {
+                       s += stats[i + base].freqs[k] + 0.5;
+               }
+               inv_sum[i] = 1.0 / s;
+       }
+
+       for (unsigned i = 0; i < 64; ++i) {
+               for (unsigned j = 0; j < 64; ++j) {
+                       double d = 0.0;
+                       for (unsigned k = 0; k < 256; ++k) {
+                               double p1 = (stats[i + base].freqs[k] + 0.5) * inv_sum[i];
+                               double p2 = (stats[j + base].freqs[k] + 0.5) * inv_sum[j];
+
+                               // K-L divergence is asymmetric; this is a hack.
+                               d += p1 * log(p1 / p2);
+                               d += p2 * log(p2 / p1);
+                       }
+                       kl_dist[i][j] = d;
+                       //printf("%.3f ", d);
+               }
+               //printf("\n");
+       }
+
+       // k-medoids init
+       int medoids[64];  // only the first NUM_CLUSTERS matter
+       bool is_medoid[64] = { false };
+       std::iota(medoids, medoids + 64, 0);
+       std::random_device rd;
+       std::mt19937 g(rd());
+       std::shuffle(medoids, medoids + 64, g);
+       for (int i = 0; i < NUM_CLUSTERS; ++i) {
+               printf("%d ", medoids[i]);
+               is_medoid[medoids[i]] = true;
+       }
+       printf("\n");
+
+       // assign each data point to the closest medoid
+       int assignment[64];
+       double current_score = find_best_assignment(medoids, assignment);
+
+       for (int i = 0; i < 1000; ++i) {
+               printf("iter %d\n", i);
+               bool any_changed = false;
+               for (int m = 0; m < NUM_CLUSTERS; ++m) {
+                       for (int o = 0; o < 64; ++o) {
+                               if (is_medoid[o]) continue;
+                               int old_medoid = medoids[m];
+                               medoids[m] = o;
+
+                               int new_assignment[64];
+                               double candidate_score = find_best_assignment(medoids, new_assignment);
+
+                               if (candidate_score < current_score) {
+                                       current_score = candidate_score;
+                                       memcpy(assignment, new_assignment, sizeof(assignment));
+
+                                       is_medoid[old_medoid] = false;
+                                       is_medoid[medoids[m]] = true;
+                                       printf("%f: ", current_score);
+                                       for (int i = 0; i < 64; ++i) {
+                                               printf("%d ", assignment[i]);
+                                       }
+                                       printf("\n");
+                                       any_changed = true;
+                               } else {
+                                       medoids[m] = old_medoid;
+                               }
+                       }
+               }
+               if (!any_changed) break;
+       }
+       printf("\n");
+       std::unordered_map<int, int> rmap;
+       for (int i = 0; i < 64; ++i) {
+               if (i % 8 == 0) printf("\n");
+               if (!rmap.count(assignment[i])) {
+                       rmap.emplace(assignment[i], rmap.size());
+               }
+               printf("%d, ", rmap[assignment[i]]);
+       }
+       printf("\n");
+}
+#endif
+
 int main(int argc, char **argv)
 {
        if (argc >= 2)
@@ -620,9 +735,8 @@ int main(int argc, char **argv)
        }
        for (unsigned y = 0; y < 8; ++y) {
                for (unsigned x = 0; x < 8; ++x) {
-                       SymbolStats &s_luma = stats[pick_stats_for(x, y)];
-                       SymbolStats &s_chroma = stats[pick_stats_for(x, y) + 8];  // HACK
-                       //SymbolStats &s_chroma = stats[pick_stats_for(x, y)];
+                       SymbolStats &s_luma = stats[pick_stats_for(x, y, false)];
+                       SymbolStats &s_chroma = stats[pick_stats_for(x, y, true)];
 
                        // Luma
                        for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
@@ -654,6 +768,15 @@ int main(int argc, char **argv)
                        }
                }
        }
+
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+       printf("Luma:\n");
+       find_optimal_stream_assignment(0);
+       printf("Chroma:\n");
+       find_optimal_stream_assignment(64);
+       exit(0);
+#endif
+
        for (unsigned i = 0; i < 64; ++i) {
                stats[i].freqs[255] /= 2;  // zero, has no sign bits (yes, this is trickery)
                stats[i].normalize_freqs(prob_scale);
@@ -685,7 +808,7 @@ int main(int argc, char **argv)
        // Luma
        for (unsigned y = 0; y < 8; ++y) {
                for (unsigned x = 0; x < 8; ++x) {
-                       SymbolStats &s_luma = stats[pick_stats_for(x, y)];
+                       SymbolStats &s_luma = stats[pick_stats_for(x, y, false)];
                        rans_encoder.init_prob(s_luma);
 
                        // Luma
@@ -730,8 +853,7 @@ int main(int argc, char **argv)
        // Cb
        for (unsigned y = 0; y < 8; ++y) {
                for (unsigned x = 0; x < 8; ++x) {
-                       SymbolStats &s_chroma = stats[pick_stats_for(x, y) + 8];
-                       //SymbolStats &s_chroma = stats[pick_stats_for(x, y)];
+                       SymbolStats &s_chroma = stats[pick_stats_for(x, y, true)];
                        rans_encoder.init_prob(s_chroma);
 
                        rans_encoder.clear();
@@ -756,8 +878,7 @@ int main(int argc, char **argv)
        // Cr
        for (unsigned y = 0; y < 8; ++y) {
                for (unsigned x = 0; x < 8; ++x) {
-                       SymbolStats &s_chroma = stats[pick_stats_for(x, y) + 8];
-                       //SymbolStats &s_chroma = stats[pick_stats_for(x, y)];
+                       SymbolStats &s_chroma = stats[pick_stats_for(x, y, true)];
                        rans_encoder.init_prob(s_chroma);
 
                        rans_encoder.clear();