[AT3] Attempt to improve quality:

 * Band energy aware quantization for high frequency bands
 * Energy aware bit allocation for low frequency bands
 * Tune bit allocation table and weight coefficients

5 files changed, 82 insertions, 26 deletions

diff --git a/src/atrac/atrac3_bitstream.cpp b/src/atrac/atrac3_bitstream.cpp
index 3be3323..7163076 100644
--- a/src/atrac/atrac3_bitstream.cpp
+++ b/src/atrac/atrac3_bitstream.cpp
@@ -34,12 +34,26 @@ using std::vector;
 using std::memset;
 
 static const uint32_t FixedBitAllocTable[TAtrac3Data::MaxBfus] = {
-  5, 5, 4, 4, 4, 4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
-  3, 3, 3, 3, 3, 3, 3, 3,
-  3, 2, 2, 1,
-  1, 0
+  4, 4, 4, 4, 4, 4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
+  2, 2, 2, 2, 2, 1, 1, 1,
+  1, 1, 1, 1,
+  0, 0
 };
 
+#define EAQ 1
+
+#ifdef EAQ
+
+static constexpr size_t LOSY_NAQ_START = 18;
+static constexpr size_t BOOST_NAQ_END = 10;
+
+#else
+
+static constexpr size_t LOSY_NAQ_START = 31;
+static constexpr size_t BOOST_NAQ_END = 0;
+
+#endif
+
 std::vector<TFloat> TAtrac3BitStreamWriter::ATH;
 TAtrac3BitStreamWriter::TAtrac3BitStreamWriter(ICompressedOutput* container, const TContainerParams& params, uint32_t bfuIdxConst)
     : Container(container)
@@ -122,14 +136,14 @@ uint32_t TAtrac3BitStreamWriter::VLCEnc(const uint32_t selector, const int manti
 
 
 std::pair<uint8_t, uint32_t> TAtrac3BitStreamWriter::CalcSpecsBitsConsumption(const TSingleChannelElement& sce,
-                                                        const vector<uint32_t>& precisionPerEachBlocks, int* mantisas)
+    const vector<uint32_t>& precisionPerEachBlocks, int* mantisas, vector<float>& energyErr)
 {
 
     const vector<TScaledBlock>& scaledBlocks = sce.ScaledBlocks;
     const uint32_t numBlocks = precisionPerEachBlocks.size();
     uint32_t bitsUsed = numBlocks * 3;
 
-    auto lambda = [this, numBlocks, mantisas, &precisionPerEachBlocks, &scaledBlocks](bool clcMode, bool calcMant) {
+    auto lambda = [this, numBlocks, mantisas, &precisionPerEachBlocks, &scaledBlocks, &energyErr](bool clcMode, bool calcMant) {
         uint32_t bits = 0;
         for (uint32_t i = 0; i < numBlocks; ++i) {
             if (precisionPerEachBlocks[i] == 0)
@@ -141,7 +155,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();
-                QuantMantisas(values, first, last, mul, mantisas);
+                energyErr[i] = QuantMantisas(values, first, last, mul, i > LOSY_NAQ_START, mantisas);
             }
             bits += clcMode ? CLCEnc(precisionPerEachBlocks[i], mantisas + first, blockSize, nullptr) :
                 VLCEnc(precisionPerEachBlocks[i], mantisas + first, blockSize, nullptr);
@@ -171,8 +185,27 @@ static inline bool CheckBfus(uint16_t* numBfu, const vector<uint32_t>& precision
 
 static const std::pair<uint8_t, vector<uint32_t>> DUMMY_ALLOC{1, vector<uint32_t>{0}};
 
+bool ConsiderEnergyErr(const vector<float>& err, vector<uint32_t>& bits)
+{
+    if (err.size() < bits.size())
+        abort();
+
+    bool adjusted = false;
+    size_t lim = std::min((size_t)BOOST_NAQ_END, bits.size());
+    for (size_t i = 0; i < lim; i++) {
+        float e = err[i];
+        if (((e > 0 && e < 0.7) || e > 1.2) & (bits[i] < 7)) {
+            //std::cerr << "adjust: " << i << " e: " << e << " b: " << bits[i] << std::endl;
+            bits[i]++;
+            adjusted = true;
+        }
+    }
+
+    return adjusted;
+}
+
 std::pair<uint8_t, vector<uint32_t>> TAtrac3BitStreamWriter::CreateAllocation(const TSingleChannelElement& sce,
-                                                                              const uint16_t targetBits, int mt[MaxSpecs], float laudness)
+    const uint16_t targetBits, int mt[MaxSpecs], float laudness)
 {
     const vector<TScaledBlock>& scaledBlocks = sce.ScaledBlocks;
     if (scaledBlocks.empty()) {
@@ -193,6 +226,7 @@ std::pair<uint8_t, vector<uint32_t>> TAtrac3BitStreamWriter::CreateAllocation(co
     }
 
     vector<uint32_t> precisionPerEachBlocks(numBfu);
+    vector<float> energyErr(numBfu);
     uint8_t mode;
     bool cont = true;
     while (cont) {
@@ -201,11 +235,17 @@ std::pair<uint8_t, vector<uint32_t>> TAtrac3BitStreamWriter::CreateAllocation(co
         TFloat minShift = -8;
         for (;;) {
             TFloat shift = (maxShift + minShift) / 2;
-            const vector<uint32_t>& tmpAlloc = CalcBitsAllocation(scaledBlocks, numBfu, spread, shift, laudness);
-            auto consumption = CalcSpecsBitsConsumption(sce, tmpAlloc, mt);
+            vector<uint32_t> tmpAlloc = CalcBitsAllocation(scaledBlocks, numBfu, spread, shift, laudness);
+            energyErr.clear();
+            energyErr.resize(numBfu);
+            std::pair<uint8_t, uint32_t> consumption;
+
+            do {
+                consumption = CalcSpecsBitsConsumption(sce, tmpAlloc, mt, energyErr);
+            } while (ConsiderEnergyErr(energyErr, tmpAlloc));
 
             auto bitsUsedByTonal = EncodeTonalComponents(sce, tmpAlloc, nullptr);
-//            std::cerr << consumption.second << " |tonal: " << bitsUsedByTonal << " target: " << targetBits << " shift " << shift << " max | min " << maxShift << " " << minShift << " numBfu: " << numBfu << std::endl;
+            //std::cerr << consumption.second << " |tonal: " << bitsUsedByTonal << " target: " << targetBits << " shift " << shift << " max | min " << maxShift << " " << minShift << " numBfu: " << numBfu << std::endl;
             consumption.second += bitsUsedByTonal;
 
             if (consumption.second < targetBits) {
@@ -471,11 +511,25 @@ vector<uint32_t> TAtrac3BitStreamWriter::CalcBitsAllocation(const std::vector<TS
             bitsPerEachBlock[i] = 0;
         } else {
             const uint32_t fix = FixedBitAllocTable[i];
-            int tmp = spread * ( (TFloat)scaledBlocks[i].ScaleFactorIndex/3.2) + (1.0 - spread) * fix - shift;
+            float x = 6;
+            if (i < 3) {
+                x = 2.8;
+            } else if (i < 10) {
+                x = 2.6;
+            } else if (i < 15) {
+                x = 3.3;
+            } else if (i <= 20) {
+                x = 3.6;
+            } else if (i <= 28) {
+                x = 4.2;
+            }
+            int tmp = spread * ( (TFloat)scaledBlocks[i].ScaleFactorIndex / x) + (1.0 - spread) * fix - shift;
             if (tmp > 7) {
                 bitsPerEachBlock[i] = 7;
             } else if (tmp < 0) {
                 bitsPerEachBlock[i] = 0;
+            } else if (tmp == 0) {
+                bitsPerEachBlock[i] = 1;
             } else {
                 bitsPerEachBlock[i] = tmp;
             }
diff --git a/src/atrac/atrac3_bitstream.h b/src/atrac/atrac3_bitstream.h
index e67f726..71ff3e0 100644
--- a/src/atrac/atrac3_bitstream.h
+++ b/src/atrac/atrac3_bitstream.h
@@ -73,7 +73,7 @@ private:
 
     std::pair<uint8_t, uint32_t> CalcSpecsBitsConsumption(const TSingleChannelElement& sce,
                                                           const std::vector<uint32_t>& precisionPerEachBlocks,
-                                                          int* mantisas);
+                                                          int* mantisas, std::vector<float>& energyErr);
 
     void EncodeSpecs(const TSingleChannelElement& sce, NBitStream::TBitStream* bitStream,
                      const std::pair<uint8_t, std::vector<uint32_t>>&, const int mt[MaxSpecs]);
diff --git a/src/atrac/atrac_scale.cpp b/src/atrac/atrac_scale.cpp
index 05d8c66..eae1d45 100644
--- a/src/atrac/atrac_scale.cpp
+++ b/src/atrac/atrac_scale.cpp
@@ -36,7 +36,7 @@ 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)
+TFloat QuantMantisas(const TFloat* in, const uint32_t first, const uint32_t last, const TFloat mul, bool ea, int* const mantisas)
 {
     float e1 = 0.0;
     float e2 = 0.0;
@@ -52,18 +52,19 @@ void QuantMantisas(const TFloat* in, const uint32_t first, const uint32_t last,
         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 (ea) {
+            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;
+        return e1 / e2;
     }
 
     static auto cmp = [](const std::pair<float, int>& a, const std::pair<float, int>& b) {
@@ -91,7 +92,7 @@ void QuantMantisas(const TFloat* in, const uint32_t first, const uint32_t last,
                 }
             }
         }
-        return;
+        return e1 / e2;
     }
 
     if (e2 > e1) {
@@ -114,8 +115,9 @@ void QuantMantisas(const TFloat* in, const uint32_t first, const uint32_t last,
                 }
             }
         }
-        return;
+        return e1 / e2;
     }
+    return e1 / e2;
 }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/src/atrac/atrac_scale.h b/src/atrac/atrac_scale.h
index cd94af8..0289961 100644
--- a/src/atrac/atrac_scale.h
+++ b/src/atrac/atrac_scale.h
@@ -27,7 +27,7 @@
 
 namespace NAtracDEnc {
 
-void QuantMantisas(const TFloat* in, uint32_t first, uint32_t last, TFloat mul, int* mantisas);
+TFloat QuantMantisas(const TFloat* in, uint32_t first, uint32_t last, TFloat mul, bool ea, int* mantisas);
 
 struct TScaledBlock {
     TScaledBlock(uint8_t sfi) : ScaleFactorIndex(sfi) {}
diff --git a/src/atrac/atrac_scale_ut.cpp b/src/atrac/atrac_scale_ut.cpp
index afc627a..3d64231 100644
--- a/src/atrac/atrac_scale_ut.cpp
+++ b/src/atrac/atrac_scale_ut.cpp
@@ -53,7 +53,7 @@ TEST(Quant, SaveEnergyLost) {
         std::vector<int> mantisas;
         mantisas.resize(in.size());
 
-        QuantMantisas(scaled.data(), 0, mantisas.size(), td.Q, mantisas.data());
+        QuantMantisas(scaled.data(), 0, mantisas.size(), td.Q, true, mantisas.data());
 
         float e2 = 0.0;
         for (auto x : mantisas) {