PitchDetector.h

// DSPark - Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi - MIT License
 
#pragma once
 
#include "../Core/DspMath.h"
#include "../Core/FFT.h"
 
#include <algorithm>
#include <atomic>
#include <cmath>
#include <cstddef>
#include <memory>
#include <span>
#include <vector>
 
namespace dspark {
 
template <FloatType T>
class PitchDetector
{
public:
    void prepare(double sampleRate, int windowSize = 0, int hopSize = 0)
    {
        if (!std::isfinite(sampleRate) || sampleRate <= 0.0)
            return;
 
        fft_.reset(); // gate OFF: pushSamples() is a no-op while rebuilding
 
        sampleRate_ = sampleRate;
        if (windowSize <= 0)
        {
            // Automatic: hold the analysis TIME SPAN constant across sample
            // rates (the span 2048 samples cover at 48 kHz). The reachable
            // register is fs/(windowSize/2 - 1), so a constant span pins the
            // lowest detectable fundamental instead of letting it climb with
            // the rate.
            const double target = sampleRate_ * (kAutoSpanRef / kAutoSpanRate);
            int n = kAutoMinWindow;
            while (n < kAutoMaxWindow && static_cast<double>(n) < target) n <<= 1;
            windowSize_ = n;
        }
        else
        {
            windowSize_ = std::clamp(windowSize, 64, 1 << 20);
        }
        halfWindow_ = windowSize_ / 2;
        hopSize_    = (hopSize <= 0) ? std::max(1, windowSize_ / 4)
                                     : std::clamp(hopSize, 1, windowSize_);
 
        // Mirrored buffer technique: size is 2x windowSize.
        // Guarantees continuous memory layout without modulo operations.
        buffer_.assign(static_cast<size_t>(windowSize_) * 2, T(0));
        yinBuffer_.assign(static_cast<size_t>(halfWindow_), T(0));
        rawDiff_.assign(static_cast<size_t>(halfWindow_), T(0));
 
        // YIN-FFT resources: the difference function is computed via one
        // cross-correlation in the frequency domain (3 FFTs) instead of the
        // O(windowSize^2) direct form - ~20x faster at the default window.
        fftSize_ = 1;
        while (fftSize_ < windowSize_ * 2) fftSize_ <<= 1;
        fftTime_.assign(static_cast<size_t>(fftSize_), T(0));
        specHalf_.assign(static_cast<size_t>(fftSize_) + 2, T(0));
        specFull_.assign(static_cast<size_t>(fftSize_) + 2, T(0));
        corrTime_.assign(static_cast<size_t>(fftSize_), T(0));
        prefixSq_.assign(static_cast<size_t>(windowSize_) + 1, T(0));
 
        writePos_ = 0;
        samplesSinceLastDetect_ = 0;
 
        frequency_.store(T(0), std::memory_order_relaxed);
        confidence_.store(T(0), std::memory_order_relaxed);
 
        fft_ = std::make_unique<FFTReal<T>>(static_cast<size_t>(fftSize_)); // gate ON, last
    }
 
    void pushSamples(std::span<const T> samples) noexcept
    {
        if (fft_ == nullptr)
            return;
 
        for (const T sample : samples)
        {
            // Write into mirrored buffer
            buffer_[static_cast<size_t>(writePos_)] = sample;
            buffer_[static_cast<size_t>(writePos_ + windowSize_)] = sample;
 
            writePos_++;
            if (writePos_ >= windowSize_)
            {
                writePos_ = 0;
            }
 
            samplesSinceLastDetect_++;
            if (samplesSinceLastDetect_ >= hopSize_)
            {
                detect();
                samplesSinceLastDetect_ = 0;
            }
        }
    }
 
    [[nodiscard]] T getFrequencyHz() const noexcept
    {
        return frequency_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getConfidence() const noexcept
    {
        return confidence_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] int getMidiNote() const noexcept
    {
        const T freq = getFrequencyHz();
        if (freq <= T(0)) return -1;
        return static_cast<int>(std::round(T(69) + T(12) * std::log2(freq / T(440))));
    }
 
    [[nodiscard]] T getCentsOffset() const noexcept
    {
        const T freq = getFrequencyHz();
        if (freq <= T(0)) return T(0);
        T midiExact = T(69) + T(12) * std::log2(freq / T(440));
        return (midiExact - std::round(midiExact)) * T(100);
    }
 
    void setThreshold(T threshold) noexcept
    {
        if (!std::isfinite(threshold)) return;
        threshold_.store(std::clamp(threshold, T(0.01), T(0.5)), std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getThreshold() const noexcept
    {
        return threshold_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] int getWindowSize() const noexcept { return windowSize_; }
 
    [[nodiscard]] int getHopSize() const noexcept { return hopSize_; }
 
    void reset() noexcept
    {
        std::fill(buffer_.begin(), buffer_.end(), T(0));
        writePos_ = 0;
        samplesSinceLastDetect_ = 0;
        frequency_.store(T(0), std::memory_order_relaxed);
        confidence_.store(T(0), std::memory_order_relaxed);
    }
 
private:
    void detect() noexcept
    {
        // Obtain a perfectly contiguous block of memory representing the current window.
        // writePos_ points to the oldest sample in the mirrored buffer.
        const T* currentWindow = &buffer_[static_cast<size_t>(writePos_)];
 
        // Silence check / Energy calculation on contiguous memory. The sum
        // spans the whole window, so it doubles as the non-finite gate: with
        // any NaN/Inf sample present the old code filled the CMND with zeros
        // ((NaN > 0) is false) and published fs/2 at confidence 1.0 - a fake
        // detection with maximum confidence. Report unvoiced instead; the
        // bad samples flush out of the window on their own.
        T energy = T(0);
        for (int i = 0; i < windowSize_; ++i) {
            energy += currentWindow[i] * currentWindow[i];
        }
 
        if (energy < T(1e-10) || !std::isfinite(energy))
        {
            frequency_.store(T(0), std::memory_order_relaxed);
            confidence_.store(T(0), std::memory_order_relaxed);
            return;
        }
 
        // YIN-FFT difference function:
        //   d(tau) = E1 + E2(tau) - 2*r(tau)
        // with E1 = sum of x[0..W)^2 (constant), E2(tau) the energy of the
        // shifted window (prefix sums), and r(tau) the cross-correlation of
        // the first half against the full window - computed with 3 FFTs.
        const int W = halfWindow_;
 
        // (a) prefix sums of squared samples over the full window
        prefixSq_[0] = T(0);
        for (int i = 0; i < windowSize_; ++i)
            prefixSq_[static_cast<size_t>(i + 1)] =
                prefixSq_[static_cast<size_t>(i)] + currentWindow[i] * currentWindow[i];
        const T e1 = prefixSq_[static_cast<size_t>(W)];
 
        // (b) r(tau) via FFT cross-correlation: IFFT(conj(FFT(first half)) * FFT(window))
        std::fill(fftTime_.begin(), fftTime_.end(), T(0));
        std::copy(currentWindow, currentWindow + W, fftTime_.begin());
        fft_->forward(fftTime_.data(), specHalf_.data());
 
        std::fill(fftTime_.begin(), fftTime_.end(), T(0));
        std::copy(currentWindow, currentWindow + windowSize_, fftTime_.begin());
        fft_->forward(fftTime_.data(), specFull_.data());
 
        const int numBins = fftSize_ / 2 + 1;
        for (int k = 0; k < numBins; ++k)
        {
            const T aRe = specHalf_[static_cast<size_t>(2 * k)];
            const T aIm = specHalf_[static_cast<size_t>(2 * k + 1)];
            const T bRe = specFull_[static_cast<size_t>(2 * k)];
            const T bIm = specFull_[static_cast<size_t>(2 * k + 1)];
            // conj(A) * B
            specFull_[static_cast<size_t>(2 * k)]     = aRe * bRe + aIm * bIm;
            specFull_[static_cast<size_t>(2 * k + 1)] = aRe * bIm - aIm * bRe;
        }
        fft_->inverse(specFull_.data(), corrTime_.data());
 
        // (c) CMND from the closed-form difference function
        const T threshold = threshold_.load(std::memory_order_relaxed);
        yinBuffer_[0] = T(1);
        T runningSum = T(0);
 
        for (int tau = 1; tau < W; ++tau)
        {
            const T e2 = prefixSq_[static_cast<size_t>(tau + W)] - prefixSq_[static_cast<size_t>(tau)];
            T d = e1 + e2 - T(2) * corrTime_[static_cast<size_t>(tau)];
            if (d < T(0)) d = T(0); // guard tiny negative round-off
 
            rawDiff_[static_cast<size_t>(tau)] = d;
            runningSum += d;
            yinBuffer_[static_cast<size_t>(tau)] =
                (runningSum > T(0)) ? d * static_cast<T>(tau) / runningSum : T(0);
        }
 
        // Search for dip below threshold
        int tauEstimate = -1;
        for (int tau = 2; tau < halfWindow_; ++tau)
        {
            if (yinBuffer_[static_cast<size_t>(tau)] < threshold)
            {
                while (tau + 1 < halfWindow_ &&
                       yinBuffer_[static_cast<size_t>(tau + 1)] < yinBuffer_[static_cast<size_t>(tau)])
                {
                    ++tau;
                }
                tauEstimate = tau;
                break;
            }
        }
 
        if (tauEstimate < 0)
        {
            frequency_.store(T(0), std::memory_order_relaxed);
            confidence_.store(T(0), std::memory_order_relaxed);
            return;
        }
 
        // Sub-sample precision from the raw difference function (see
        // parabolicInterp()); the CMND only selects the dip.
        T betterTau = parabolicInterp(tauEstimate);
        T finalConfidence = std::clamp(T(1) - yinBuffer_[static_cast<size_t>(tauEstimate)], T(0), T(1));
 
        frequency_.store(static_cast<T>(sampleRate_) / betterTau, std::memory_order_relaxed);
        confidence_.store(finalConfidence, std::memory_order_relaxed);
    }
 
    [[nodiscard]] T parabolicInterp(int tau) const noexcept
    {
        if (tau < 2 || tau >= halfWindow_ - 1)
            return static_cast<T>(tau);
 
        // Centre on the raw function's own lowest sample next to the CMND dip.
        const int centre = tau;
        if (rawDiff_[static_cast<size_t>(centre - 1)] < rawDiff_[static_cast<size_t>(tau)]) tau = centre - 1;
        if (rawDiff_[static_cast<size_t>(centre + 1)] < rawDiff_[static_cast<size_t>(tau)]) tau = centre + 1;
        if (tau < 2 || tau >= halfWindow_ - 1)
            return static_cast<T>(tau);
 
        const double s0 = static_cast<double>(rawDiff_[static_cast<size_t>(tau - 1)]);
        const double s1 = static_cast<double>(rawDiff_[static_cast<size_t>(tau)]);
        const double s2 = static_cast<double>(rawDiff_[static_cast<size_t>(tau + 1)]);
 
        const double denom = s0 - 2.0 * s1 + s2;
 
        // Flat or inverted curvature (no usable dip on the raw function).
        if (!(denom > 1e-30))
            return static_cast<T>(tau);
 
        const double vertex = std::clamp((s0 - s2) / (2.0 * denom), -0.5, 0.5);
        const double w = 6.283185307179586 / static_cast<double>(tau);
        const double offset = std::atan(2.0 * vertex * std::tan(0.5 * w)) / w;
        return static_cast<T>(static_cast<double>(tau) + offset);
    }
 
    static constexpr double kAutoSpanRef = 2048.0;
    static constexpr double kAutoSpanRate = 48000.0;
    static constexpr int kAutoMinWindow = 512;
    static constexpr int kAutoMaxWindow = 16384;
 
    double sampleRate_ = 44100.0;
    int windowSize_ = 2048;
    int halfWindow_ = 1024;
    int hopSize_ = 512;
    int writePos_ = 0;
    int samplesSinceLastDetect_ = 0;
    std::atomic<T> threshold_{ T(0.10) };
 
    // Thread-safe outputs
    std::atomic<T> frequency_{T(0)};
    std::atomic<T> confidence_{T(0)};
 
    // Mirrored buffer keeps every analysis window contiguous
    std::vector<T> buffer_;     // Size: 2 * windowSize_
    std::vector<T> yinBuffer_;  // Size: halfWindow_
    std::vector<T> rawDiff_;    // Size: halfWindow_ (difference function d(tau))
 
    // YIN-FFT resources (cross-correlation difference function)
    int fftSize_ = 4096;
    std::unique_ptr<FFTReal<T>> fft_; // doubles as the "prepared" gate
    std::vector<T> fftTime_;    // Size: fftSize_
    std::vector<T> specHalf_;   // Size: fftSize_ + 2
    std::vector<T> specFull_;   // Size: fftSize_ + 2
    std::vector<T> corrTime_;   // Size: fftSize_
    std::vector<T> prefixSq_;   // Size: windowSize_ + 1
};
 
} // namespace dspark