stft.hpp

Go to the documentation of this file.

#pragma once

 
#include <pulp/signal/fft.hpp>
#include <pulp/signal/mirrored_history_buffer.hpp>
#include <pulp/signal/windowing.hpp>
#include <vector>
#include <complex>
#include <cmath>
#include <algorithm>
#include <cassert>
#include <cstring>
 
namespace pulp::signal {

struct StftConfig {
    int fft_size = 1024;                              // Must be power of 2, 256–8192
    int hop_size = 256;                               // Samples between frames
    WindowFunction::Type window = WindowFunction::Type::hann;
    float window_param = 0.0f;                        // For Kaiser alpha
};

template <typename SampleType = float>
struct StftFrameT {
    std::vector<SampleType> magnitude; // num_bins values (fft_size/2 + 1)
    std::vector<SampleType> phase;     // num_bins values
    int num_bins = 0;
};
 
using StftFrame = StftFrameT<float>;
using StftFrame64 = StftFrameT<double>;

template <typename SampleType = float>
class StftT {
public:
    StftT() = default;
 
    explicit StftT(const StftConfig& config) { configure(config); }

    void configure(const StftConfig& config) {
        assert(config.fft_size >= 256 && config.fft_size <= 8192);
        assert((config.fft_size & (config.fft_size - 1)) == 0); // power of 2
 
        config_ = config;
        fft_ = FftT<SampleType>(config.fft_size);
        window_ = WindowFunction::generate<SampleType>(
            config.fft_size, config.window, static_cast<SampleType>(config.window_param));
        window_sum_ = SampleType{0};
        for (const auto coefficient : window_) window_sum_ += coefficient;
        num_bins_ = config.fft_size / 2 + 1;
 
        // The complete oldest-to-newest analysis history remains contiguous.
        history_.prepare(static_cast<std::size_t>(config.fft_size));
        hop_counter_ = 0;
        samples_fed_ = 0;
 
        // Working buffers
        windowed_.resize(config.fft_size, SampleType{0.0f});
        freq_.resize(config.fft_size);
 
        frame_.magnitude.resize(num_bins_, SampleType{0.0f});
        frame_.phase.resize(num_bins_, SampleType{0.0f});
        frame_.num_bins = num_bins_;
        frame_ready_ = false;
    }

    bool push_samples(const SampleType* samples, int count) {
        bool computed = false;
        for (int i = 0; i < count; ++i) {
            history_.push(samples[i]);
            ++samples_fed_;
            ++hop_counter_;
 
            if (hop_counter_ >= config_.hop_size && samples_fed_ >= config_.fft_size) {
                compute_frame();
                hop_counter_ = 0;
                computed = true;
            }
        }
        return computed;
    }

    const StftFrameT<SampleType>& latest_frame() const { return frame_; }

    bool frame_ready() const { return frame_ready_; }

    int num_bins() const { return num_bins_; }

    SampleType window_sum() const { return window_sum_; }

    int fft_size() const { return config_.fft_size; }

    int hop_size() const { return config_.hop_size; }

    static void to_db(SampleType* magnitudes,
                      int count,
                      SampleType floor_db = SampleType{-120.0f}) {
        SampleType floor_linear =
            std::pow(SampleType{10.0f}, floor_db / SampleType{20.0f});
        for (int i = 0; i < count; ++i) {
            magnitudes[i] = SampleType{20.0f} *
                std::log10(std::max(magnitudes[i], floor_linear));
        }
    }

    std::vector<SampleType> latest_magnitude_db(
        SampleType floor_db = SampleType{-120.0f}) const {
        auto db = frame_.magnitude;
        to_db(db.data(), static_cast<int>(db.size()), floor_db);
        return db;
    }
 
    void reset() {
        history_.reset();
        hop_counter_ = 0;
        samples_fed_ = 0;
        frame_ready_ = false;
        std::fill(frame_.magnitude.begin(), frame_.magnitude.end(), SampleType{0.0f});
        std::fill(frame_.phase.begin(), frame_.phase.end(), SampleType{0.0f});
    }
 
private:
    void compute_frame() {
        const auto history = history_.window();
        for (int i = 0; i < config_.fft_size; ++i) {
            windowed_[i] = history[static_cast<std::size_t>(i)];
        }
 
        // Apply window
        WindowFunction::apply(windowed_.data(), window_);
 
        // Forward FFT (DC..Nyquist is all the frame reads)
        fft_.forward_real_half(windowed_.data(), freq_.data());
 
        // Extract magnitude and phase for positive frequencies
        for (int i = 0; i < num_bins_; ++i) {
            frame_.magnitude[i] = std::abs(freq_[i]);
            frame_.phase[i] = std::arg(freq_[i]);
        }
 
        frame_ready_ = true;
    }
 
    StftConfig config_;
    FftT<SampleType> fft_{1024};
    std::vector<SampleType> window_;
    SampleType window_sum_ = SampleType{1};
    int num_bins_ = 0;
 
    // Single-thread sample history; not a producer/consumer queue.
    MirroredHistoryBuffer<SampleType> history_;
    int hop_counter_ = 0;
    int samples_fed_ = 0;
 
    // Working buffers (reused to avoid allocation)
    std::vector<SampleType> windowed_;
    std::vector<std::complex<SampleType>> freq_;
 
    // Latest computed frame
    StftFrameT<SampleType> frame_;
    bool frame_ready_ = false;
};
 
using Stft = StftT<float>;
using Stft64 = StftT<double>;
 
} // namespace pulp::signal