#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;
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
}
}
}
+#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;
+}
+
+double find_inv_sum(const SymbolStats &stats)
+{
+ double s = 0.0;
+ for (unsigned j = 0; j < NUM_SYMS; ++j) {
+ s += stats.freqs[j] + 0.5;
+ }
+ return 1.0 / s;
+}
+
+void find_optimal_stream_assignment(int base)
+{
+ // k-means init; make random assignments
+ std::random_device rd;
+ std::mt19937 g(rd());
+ std::uniform_int_distribution<> u(0, NUM_CLUSTERS - 1);
+ int assignment[64];
+ for (unsigned i = 0; i < 64; ++i) {
+ assignment[i] = u(g);
+ }
+ double inv_sum_coeffs[64];
+ for (unsigned i = 0; i < 64; ++i) {
+ inv_sum_coeffs[i] = find_inv_sum(stats[i + base]);
+ }
+
+ for (unsigned iter = 0; iter < 1000; ++iter) {
+ // make new clusters based on the current assignments
+ SymbolStats clusters[NUM_CLUSTERS];
+ for (unsigned i = 0; i < NUM_CLUSTERS; ++i) {
+ clusters[i].clear();
+ }
+ for (unsigned i = 0; i < 64; ++i) {
+ for (unsigned j = 0; j < NUM_SYMS; ++j) {
+ clusters[assignment[i]].freqs[j] += stats[i + base].freqs[j];
+ }
+ }
+
+ double inv_sum_clusters[NUM_CLUSTERS];
+ for (unsigned i = 0; i < NUM_CLUSTERS; ++i) {
+ inv_sum_clusters[i] = find_inv_sum(clusters[i]);
+ }
+
+ // find new assignments based on distance to the clusters
+ bool any_changed = false;
+ double total_d = 0.0;
+ for (unsigned i = 0; i < 64; ++i) {
+ int best_assignment = -1;
+ double best_distance = HUGE_VAL;
+ for (unsigned j = 0; j < NUM_CLUSTERS; ++j) {
+ double d = 0.0;
+ for (unsigned k = 0; k < NUM_SYMS; ++k) {
+ double p1 = (clusters[j].freqs[k] + 0.5) * inv_sum_clusters[j];
+ double p2 = (stats[i + base].freqs[k] + 0.5) * inv_sum_coeffs[i];
+
+ // K-L divergence is asymmetric; this is a hack.
+ d += p1 * log(p1 / p2);
+ d += p2 * log(p2 / p1);
+ }
+ if (d < best_distance) {
+ best_assignment = j;
+ best_distance = d;
+ }
+ }
+ if (assignment[i] != best_assignment) {
+ any_changed = true;
+ }
+ assignment[i] = best_assignment;
+ total_d += best_distance;
+ }
+ printf("iter %u: %.3f\n", iter, total_d);
+ 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)
}
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) {
}
}
}
+
+#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);
// 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
// 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();
// 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();
projects / narabu / blobdiff