SpectralDenoiser.h

// DSPark - Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi - MIT License
 
#pragma once
 
#include "../Core/AudioBuffer.h"
#include "../Core/AudioSpec.h"
#include "../Core/DspMath.h"
#include "../Core/SpectralProcessor.h"
#include "../Core/StateBlob.h"
 
#include <algorithm>
#include <atomic>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <vector>
 
namespace dspark {
 
template <FloatType T>
class SpectralDenoiser
{
public:
    // -- Lifecycle ---------------------------------------------------------------
 
    void prepare(const AudioSpec& spec, int fftSize = 2048)
    {
        if (!spec.isValid()) return;
        prepared_.store(false, std::memory_order_relaxed);
        numChannels_ = spec.numChannels;
        stft_.prepare(spec, fftSize, fftSize / 4);
        numBins_ = stft_.getNumBins();
 
        profile_.assign(static_cast<size_t>(numBins_), 0.0f);
        learnFrames_ = 0;
        cleanPower_.assign(static_cast<size_t>(numChannels_),
                           std::vector<float>(static_cast<size_t>(numBins_), 0.0f));
 
        prepared_.store(true, std::memory_order_relaxed);
        reset();
    }
 
    void reset() noexcept
    {
        if (!prepared_.load(std::memory_order_relaxed)) return;
        stft_.reset();
        for (auto& c : cleanPower_)
            std::fill(c.begin(), c.end(), 0.0f);
        callCounter_ = 0;
    }
 
    void clearProfile() noexcept
    {
        std::fill(profile_.begin(), profile_.end(), 0.0f);
        learnFrames_ = 0;
    }
 
    // -- Parameters (thread-safe) ---------------------------------------------------
 
    void setLearning(bool learning) noexcept
    {
        learning_.store(learning, std::memory_order_relaxed);
    }
 
    void setReduction(T db) noexcept
    {
        if (!std::isfinite(db)) return;
        reduction_.store(std::clamp(db, T(0), T(40)), std::memory_order_relaxed);
    }
 
    void setThreshold(T factor) noexcept
    {
        if (!std::isfinite(factor)) return;
        threshold_.store(std::clamp(factor, T(1), T(8)), std::memory_order_relaxed);
    }
 
    [[nodiscard]] bool getLearning() const noexcept { return learning_.load(std::memory_order_relaxed); }
    [[nodiscard]] T getReduction() const noexcept { return reduction_.load(std::memory_order_relaxed); }
    [[nodiscard]] T getThreshold() const noexcept { return threshold_.load(std::memory_order_relaxed); }
 
    [[nodiscard]] int getLatency() const noexcept { return stft_.getLatency(); }
 
    [[nodiscard]] std::vector<uint8_t> getState() const
    {
        StateWriter w(stateId("DNSE"), 1);
        // Explicit float casts: the blob stores float, and with T = double the
        // unqualified write(key, double) would be ambiguous (float/int32/bool).
        w.write("reduction", static_cast<float>(reduction_.load(std::memory_order_relaxed)));
        w.write("threshold", static_cast<float>(threshold_.load(std::memory_order_relaxed)));
        return w.blob();
    }
 
    bool setState(const uint8_t* data, size_t size)
    {
        StateReader r(data, size);
        if (!r.isValid() || r.processorId() != stateId("DNSE")) return false;
        setReduction(static_cast<T>(r.read("reduction", 18.0f)));
        setThreshold(static_cast<T>(r.read("threshold", 2.0f)));
        return true;
    }
 
    // -- Processing -------------------------------------------------------------------
 
    void processBlock(AudioBufferView<T> buffer) noexcept
    {
        if (!prepared_.load(std::memory_order_relaxed)) return;
 
        const bool learning = learning_.load(std::memory_order_relaxed);
        const float floorGain = std::pow(
            10.0f, -static_cast<float>(reduction_.load(std::memory_order_relaxed)) / 20.0f);
        const float thresh = static_cast<float>(threshold_.load(std::memory_order_relaxed));
 
        // The STFT invokes the callback once per PROCESSED channel per hop
        // (in channel order) - and it processes min(buffer, prepared)
        // channels. Modulo by that same effective count, reset per block, or
        // a narrow buffer over a wider spec would rotate the per-channel
        // gain memories between hops (channel 0 alternating onto channel 1's
        // release state - measured 0.004 divergence versus a mono-prepared
        // twin before this fix).
        const int nChEff = std::max(1, std::min(buffer.getNumChannels(), numChannels_));
        callCounter_ = 0;
 
        // Over-subtraction: the learned noise power is scaled by threshold^2.
        const float overSub = thresh * thresh;
 
        stft_.processBlock(buffer, [this, learning, floorGain, overSub, nChEff](T* bins, int numBins)
        {
            auto& clean = cleanPower_[static_cast<size_t>(callCounter_ % nChEff)];
            ++callCounter_;
 
            // The profile is the running MEAN noise power per bin (the
            // cumulative average over every learned frame, then an
            // exponential average once kMaxLearnFrames are in).
            float learnRate = 0.0f;
            if (learning)
            {
                learnFrames_ = std::min(learnFrames_ + 1, kMaxLearnFrames);
                learnRate = 1.0f / static_cast<float>(learnFrames_);
            }
 
            for (int k = 0; k < numBins; ++k)
            {
                const float re = static_cast<float>(bins[2 * k]);
                const float im = static_cast<float>(bins[2 * k + 1]);
                const float power = re * re + im * im;
 
                auto& noise = profile_[static_cast<size_t>(k)];
                if (learning)
                    noise += learnRate * (power - noise);
 
                // Decision-directed Wiener gain (Ephraim-Malah a priori SNR):
                // the a priori SNR mixes the previous frame's clean estimate
                // with this frame's excess power, so it follows speech and
                // music onsets at once while random noise flicker cannot open
                // a bin (the hard per-bin gate did exactly that: musical noise).
                float gain = 1.0f;
                const float lambda = overSub * noise;
                auto& prevClean = clean[static_cast<size_t>(k)];
                if (lambda > 0.0f)
                {
                    const float post = power / lambda;
                    const float prio = std::max(kDDAlpha * prevClean / lambda
                                                + (1.0f - kDDAlpha) * std::max(post - 1.0f, 0.0f),
                                                kMinPrioriSnr);
                    gain = std::max(prio / (1.0f + prio), floorGain);
                }
                prevClean = gain * gain * power;
 
                bins[2 * k] = static_cast<T>(re * gain);
                bins[2 * k + 1] = static_cast<T>(im * gain);
            }
        });
    }
 
private:
    SpectralProcessor<T> stft_;
    int numChannels_ = 0;
    int numBins_ = 0;
    std::atomic<bool> prepared_ { false };
 
    static constexpr float kDDAlpha = 0.98f;           
    static constexpr float kMinPrioriSnr = 0.003162f;  
    static constexpr int kMaxLearnFrames = 4096;       
 
    std::vector<float> profile_;                  
    int learnFrames_ = 0;                         
    std::vector<std::vector<float>> cleanPower_;  
    int callCounter_ = 0;
 
    std::atomic<bool> learning_ { false };
    std::atomic<T> reduction_ { T(18) };
    std::atomic<T> threshold_ { T(2) };
};
 
} // namespace dspark