[AT3] Introduce energy aware quantization

5 files changed, 167 insertions, 19 deletions

diff --git a/src/atrac/atrac3_bitstream.cpp b/src/atrac/atrac3_bitstream.cpp
index 9362ca0..3be3323 100644
--- a/src/atrac/atrac3_bitstream.cpp
+++ b/src/atrac/atrac3_bitstream.cpp
@@ -120,24 +120,6 @@ uint32_t TAtrac3BitStreamWriter::VLCEnc(const uint32_t selector, const int manti
     return bitsUsed;
 }
 
-static inline int ToInt(double x) {
-#if defined(_MSC_VER) && !defined(_WIN64)
-    int n;
-    __asm {
-        fld x
-        fistp n
-    }
-    return n;
-#else
-    return lrint(x);
-#endif
-}
-
-static inline void CalcMantisas(const TFloat* values, const uint32_t first, const uint32_t last, const TFloat mul, int* mantisas) {
-    for (uint32_t j = 0, f = first; f < last; f++, j++) {
-        mantisas[f] = ToInt(values[j] * mul);
-    }
-}
 
 std::pair<uint8_t, uint32_t> TAtrac3BitStreamWriter::CalcSpecsBitsConsumption(const TSingleChannelElement& sce,
                                                         const vector<uint32_t>& precisionPerEachBlocks, int* mantisas)
@@ -159,7 +141,7 @@ std::pair<uint8_t, uint32_t> TAtrac3BitStreamWriter::CalcSpecsBitsConsumption(co
             const TFloat mul = MaxQuant[std::min(precisionPerEachBlocks[i], (uint32_t)7)];
             if (calcMant) {
                 const TFloat* values = scaledBlocks[i].Values.data();
-                CalcMantisas(values, first, last, mul, mantisas);
+                QuantMantisas(values, first, last, mul, mantisas);
             }
             bits += clcMode ? CLCEnc(precisionPerEachBlocks[i], mantisas + first, blockSize, nullptr) :
                 VLCEnc(precisionPerEachBlocks[i], mantisas + first, blockSize, nullptr);
diff --git a/src/atrac/atrac_scale.cpp b/src/atrac/atrac_scale.cpp
index edc44a6..05d8c66 100644
--- a/src/atrac/atrac_scale.cpp
+++ b/src/atrac/atrac_scale.cpp
@@ -19,6 +19,7 @@
 #include "atrac_scale.h"
 #include "atrac1.h"
 #include "atrac3.h"
+#include "util.h"
 #include <cmath>
 #include <iostream>
 #include <algorithm>
@@ -35,6 +36,90 @@ using std::abs;
 
 static const TFloat MAX_SCALE = 1.0;
 
+void QuantMantisas(const TFloat* in, const uint32_t first, const uint32_t last, const TFloat mul, int* const mantisas)
+{
+    float e1 = 0.0;
+    float e2 = 0.0;
+
+    std::vector<std::pair<float, int>> candidates;
+    candidates.reserve(last - first);
+
+    const float inv2 = 1.0 / (mul * mul);
+
+    for (uint32_t j = 0, f = first; f < last; f++, j++) {
+        auto t = in[j] * mul;
+        e1 += in[j] * in[j];
+        mantisas[f] = ToInt(t);
+        e2 += mantisas[f] * mantisas[f] * inv2;
+
+        float delta = t - (std::trunc(t) + 0.5);
+
+        // 0 ... 0.25 ... 0.5 ... 0.75 ... 1
+        //        ^----------------^ candidates to be rounded to opposite side
+        // to decrease overall energy error in the band
+        if (std::abs(delta) < 0.25) {
+            candidates.push_back({delta, f});
+        }
+    }
+
+    if (candidates.empty()) {
+        return;
+    }
+
+    static auto cmp = [](const std::pair<float, int>& a, const std::pair<float, int>& b) {
+        return std::abs(a.first) < std::abs(b.first);
+    };
+
+    std::sort(candidates.begin(), candidates.end(), cmp);
+
+    if (e2 < e1) {
+        for (const auto& x : candidates) {
+            auto f = x.second;
+            auto j = f - first;
+            auto t = in[j] * mul;
+            if (static_cast<float>(std::abs(mantisas[f])) < std::abs(t) && static_cast<float>(std::abs(mantisas[f])) < (mul - 1)) {
+                int m = mantisas[f];
+                if (m > 0) m++;
+                if (m < 0) m--;
+                if (m == 0) m = t > 0 ? 1 : -1;
+                auto ex = e2;
+                ex -= mantisas[f] * mantisas[f] * inv2;
+                ex += m * m * inv2;
+                if (std::abs(ex - e1) < std::abs(e2 - e1)) {
+                    mantisas[f] = m;
+                    e2 = ex;
+                }
+            }
+        }
+        return;
+    }
+
+    if (e2 > e1) {
+        for (const auto& x : candidates) {
+            auto f = x.second;
+            auto j = f - first;
+            auto t = in[j] * mul;
+            if (static_cast<float>(std::abs(mantisas[f])) > std::abs(t)) {
+                auto m = mantisas[f];
+                if (m > 0) m--;
+                if (m < 0) m++;
+
+                auto ex = e2;
+                ex -= mantisas[f] * mantisas[f] * inv2;
+                ex += m * m * inv2;
+
+                if (std::abs(ex - e1) < std::abs(e2 - e1)) {
+                    mantisas[f] = m;
+                    e2 = ex;
+                }
+            }
+        }
+        return;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+
 template<class TBaseData>
 TScaledBlock TScaler<TBaseData>::Scale(const TFloat* in, uint16_t len) {
     TFloat maxAbsSpec = 0;
diff --git a/src/atrac/atrac_scale.h b/src/atrac/atrac_scale.h
index ec83059..cd94af8 100644
--- a/src/atrac/atrac_scale.h
+++ b/src/atrac/atrac_scale.h
@@ -27,6 +27,8 @@
 
 namespace NAtracDEnc {
 
+void QuantMantisas(const TFloat* in, uint32_t first, uint32_t last, TFloat mul, int* mantisas);
+
 struct TScaledBlock {
     TScaledBlock(uint8_t sfi) : ScaleFactorIndex(sfi) {}
     /* const */ uint8_t ScaleFactorIndex = 0;
diff --git a/src/atrac/atrac_scale_ut.cpp b/src/atrac/atrac_scale_ut.cpp
new file mode 100644
index 0000000..afc627a
--- /dev/null
+++ b/src/atrac/atrac_scale_ut.cpp
@@ -0,0 +1,67 @@
+/*
+ * This file is part of AtracDEnc.
+ *
+ * AtracDEnc is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * AtracDEnc is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with AtracDEnc; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+#include "atrac_scale.h"
+
+#include <gtest/gtest.h>
+
+using namespace NAtracDEnc;
+
+TEST(Quant, SaveEnergyLost) {
+    struct TTestData {
+        const std::vector<TFloat> In;
+        const TFloat S;
+        const TFloat Q;
+        const TFloat Diff;
+    };
+
+    const std::vector<TTestData> testData {
+        {{-2.35, -0.84,  0.65, -1.39,  1.25, -0.41, -0.85,  0.89}, 2.35001, 2.5, 0.5},
+        {{-1.26,  1.26, -1.26,  1.26, -1.26,  1.26, -1.26,  1.26}, 2.35001, 2.5, 0.4},
+        {{-0.32,  0.13,  0.28,  0.35,  0.63,  0.86,  0.63,  0.04}, 1, 15.5, 0.03}
+    };
+
+    for (const auto& td : testData) {
+        const std::vector<TFloat>& in = td.In;
+        const TFloat scale = td.S;
+
+        float e1 = 0.0;
+
+        std::vector<TFloat> scaled;
+        scaled.reserve(in.size());
+
+        for (auto x : in) {
+            e1 += x * x;
+            scaled.push_back(x / scale);
+        }
+
+        std::vector<int> mantisas;
+        mantisas.resize(in.size());
+
+        QuantMantisas(scaled.data(), 0, mantisas.size(), td.Q, mantisas.data());
+
+        float e2 = 0.0;
+        for (auto x : mantisas) {
+            auto t = x * (scale / td.Q);
+            e2 += t * t;
+        }
+
+        EXPECT_TRUE(std::abs(e2 - e1) < td.Diff);
+        std::cerr << "(e1): " << e1 << " (e2): " << e2 << " (e2-e1) " << e2 - e1 << " (e2/e1) " << e2 / e1 << std::endl;
+    }
+}
diff --git a/src/util.h b/src/util.h
index 0e5d948..7d8d65e 100644
--- a/src/util.h
+++ b/src/util.h
@@ -90,3 +90,15 @@ inline T CalcEnergy(const std::vector<T>& in) {
         });
 }
 
+inline int ToInt(double x) {
+#if defined(_MSC_VER) && !defined(_WIN64)
+    int n;
+    __asm {
+        fld x
+        fistp n
+    }
+    return n;
+#else
+    return lrint(x);
+#endif
+}