ShortTimeFourierTransform.cpp

Go to the documentation of this file.

namespace krotos
{
    ShortTimeFourierTransform::ShortTimeFourierTransform(int winSizeInSamples, int fftSize, int hopSize,
                                                         WindowType winMethod, fftMode mode, int sampleRate)
        : dftParams{winSizeInSamples, fftSize, hopSize, winSizeInSamples - hopSize, sampleRate}, m_mode(mode),
          m_fft(static_cast<int>(log2f(static_cast<float>(fftSize)))), sigFrame(2 * fftSize),
          // TODO: Remove after stl removal
          m_sigFrame(2 * fftSize)
    {
        // Validate input parameters
        jassert(dftParams.winSize > 0);
        jassert(dftParams.fftSize > 0 &&
                (dftParams.fftSize & (dftParams.fftSize - 1)) == 0); // Check if fftSize is a power of 2
        jassert(dftParams.hopSize > 0 && dftParams.hopSize < dftParams.winSize);
        jassert(dftParams.sampleRate > 0);
        jassert(dftParams.winSize <= dftParams.fftSize); // Ensure window size does not exceed fft size
 
        // Set here the fft binSize
        binParams.binSize = static_cast<float>(dftParams.sampleRate) / static_cast<float>(dftParams.fftSize);
 
        // generate window function
        m_window = WindowFunctions::generateWindow(winSizeInSamples + 1, winMethod);
        dftParams.window = WindowFunctions::generateWindowEigen(winSizeInSamples + 1, winMethod);
    }
 
    std::vector<std::vector<float>> ShortTimeFourierTransform::stft(const AudioSampleBuffer& buffer)
    {
        m_inputSize = buffer.getNumSamples();
        auto inputData = buffer.getReadPointer(0);
 
        std::vector<std::vector<float>> outputData;
        // Pre allocate memory for outputData
        outputData.reserve(m_inputSize / dftParams.hopSize + 1);
 
        int posin = 0;
 
        for (; posin < m_inputSize; posin += dftParams.hopSize)
        {
            // Initialize sigFrame with zeros for each frame
            std::fill(m_sigFrame.begin(), m_sigFrame.end(), 0.0f);
 
            // Extract a signal frame and apply window function
            int frameEnd = std::min(posin + dftParams.winSize, m_inputSize);
            for (int i = 0; i < frameEnd - posin; ++i)
            {
                m_sigFrame[i] = inputData[posin + i] * m_window[i];
            }
 
            // No explicit zero-padding needed here since sigFrame was initialized with zeros
            // and it's already twice the size of FFT points as required.
 
            // check which FFT we want to process the frame with
            std::vector<float> frameOutput;
            switch (m_mode)
            {
            case krotos::ShortTimeFourierTransform::freqMagOnly:
                // apply FFT on the frame - get back frequency magnitude
                m_fft.performFrequencyOnlyForwardTransform(m_sigFrame.data(), false);
                // store frequency magnitude in the outputData
                frameOutput.assign(m_sigFrame.begin(), m_sigFrame.begin() + (dftParams.fftSize / 2) + 1);
                break;
            case krotos::ShortTimeFourierTransform::realOnly:
                // apply FFT on the frame
                // TODO: Decode the complex number encoding in here to retrieve correctly
                // phase&mag info.
                m_fft.performRealOnlyForwardTransform(m_sigFrame.data(), false);
                // store the complex pairs in the outputData
                frameOutput.assign(m_sigFrame.begin(), m_sigFrame.begin() + dftParams.winSize);
                break;
            }
 
            outputData.push_back(std::move(frameOutput));
 
            // increment counter of frames
            m_numFrames++;
        }
 
        return outputData;
    }
 
    Eigen::MatrixXf ShortTimeFourierTransform::processSignal(const Eigen::VectorXf inputSignal)
    {
        m_inputSize = inputSignal.size();
 
        // Pre allocate memory for outputData
        Eigen::MatrixXf outputData(dftParams.fftSize / 2 + 1, m_inputSize / dftParams.hopSize + 1);
        outputData.setZero();
 
        for (int posin = 0; posin < m_inputSize; posin += dftParams.hopSize)
        {
            // Initialize sigFrame with zeros for each frame
            sigFrame.setZero();
 
            // Extract a signal frame and apply window function
            int frameEnd = std::min(posin + dftParams.winSize, m_inputSize);
            for (int i = 0; i < frameEnd - posin; ++i)
            {
                sigFrame[i] = inputSignal[posin + i] * dftParams.window[i];
            }
 
            // Note: No explicit zero-padding needed here since sigFrame was initialized with zeros
            // and it's already twice the size of FFT points as required.
 
            // Check which FFT we want to process the frame with
            switch (m_mode)
            {
            case krotos::ShortTimeFourierTransform::freqMagOnly: {
                // apply FFT on the frame - get back frequency magnitude
                m_fft.performFrequencyOnlyForwardTransform(sigFrame.data(), false);
                // store frequency magnitude in the outputData
                outputData.col(posin / dftParams.hopSize) =
                        Eigen::Map<Eigen::VectorXf>(sigFrame.data(), dftParams.fftSize / 2 + 1);
                break;
            }
            case krotos::ShortTimeFourierTransform::realOnly: {
                // apply FFT on the frame
                // TODO: Decode the complex number encoding in here to retrieve correctly
                // phase&mag info.
                m_fft.performRealOnlyForwardTransform(sigFrame.data(), false);
                // store the complex pairs in the outputData
                outputData.col(posin / dftParams.hopSize) =
                        Eigen::Map<Eigen::VectorXf>(sigFrame.data(), dftParams.fftSize);
                break;
            }
            }
 
            // increment counter of frames
            m_numFrames++;
        }
 
        return outputData;
    }
 
    std::vector<float> ShortTimeFourierTransform::istft(std::vector<std::vector<float>> stftMatrix)
    {
        jassert(m_mode == fftMode::realOnly);
 
        // Number of the spectral frames
        int numOfFrames = static_cast<int>(stftMatrix.size());
        jassert(numOfFrames > 0);
 
        // Size of the reconstructed signal
        int outputSize = numOfFrames * dftParams.hopSize;
 
        std::vector<float> outputData(outputSize, 0.0f);
 
        int posout = 0;
 
        // int frameSize = static_cast<int>(m_sigFrame.size());
 
        for (const auto& specFrame : stftMatrix)
        {
            // Inverse transform the spectral frame
            m_fft.performRealOnlyInverseTransform(const_cast<float*>(specFrame.data()));
 
            // Window and overlap-add the frame
            for (size_t i = 0; i < dftParams.fftSize; i++)
            {
                if (posout + i < outputSize)
                {
                    outputData[posout + i] += specFrame[i];
                }
            }
 
            posout += dftParams.hopSize;
        }
 
        return outputData;
    }
 
    void ShortTimeFourierTransform::drawSpectrogram(juce::Image& image,
                                                    const std::vector<std::vector<float>>& stftMatrix)
    {
        jassert(m_mode == fftMode::freqMagOnly);
        auto numOfFrames = stftMatrix.size();
 
        auto imageWidth = image.getWidth();
        auto imageHeight = image.getHeight();
 
        jassert(imageWidth > 0 && imageHeight > 0);
 
        float factor = (static_cast<float>(numOfFrames)) / (static_cast<float>(imageWidth));
 
        for (int x = 0; x < imageWidth; x++)
        {
            int frameIndex = static_cast<int>(factor * static_cast<float>(x));
 
            jassert(frameIndex <= numOfFrames);
 
            for (auto y = 1; y < imageHeight; ++y)
            {
                auto maxLevel =
                        juce::FloatVectorOperations::findMinAndMax(stftMatrix.at(frameIndex).data(), dftParams.fftSize);
 
                // proportion of the image that this y pixel corresponds to ?
                auto skewedProportionY = 1.0f - std::exp(std::log((float)y / (float)imageHeight) * m_scalingConstant);
                // proportion to fft Index
                auto fftDataIndex = (size_t)jlimit(
                        /*lowerLimit*/ 0,
                        /*upperLimit*/ dftParams.fftSize,
                        /*value to constrain*/ (int)(skewedProportionY * dftParams.fftSize));
 
                // maps from Magnitude of spectrum to 0...1
                auto level = jmap(/*sourceValue*/ stftMatrix.at(frameIndex).at(fftDataIndex),
                                  /*sourceRangeMin*/ 0.0f,
                                  /*sourceRangeMax*/ jmax(maxLevel.getEnd(), 1e-5f),
                                  /*targetRangeMin*/ 0.0f,
                                  /*targetRangeMax*/ 1.0f);
 
                image.setPixelAt(x, y, Colour::fromHSV(level, 1.0f, level, 1.0f));
            }
        }
    }
 
    std::vector<float> ShortTimeFourierTransform::getBinsSTL(ShortTimeFourierTransform::FrequencyRange freqRange)
    {
        float freqRes = dftParams.sampleRate / static_cast<float>(dftParams.fftSize);
        Eigen::VectorXf w1 = freqRes * Eigen::VectorXf::LinSpaced(dftParams.fftSize, 0, dftParams.fftSize - 1).array();
        float Nyq = dftParams.sampleRate / 2.0f;
        float halfRes = freqRes / 2;
        // NPT Info
        NPTSinfo(dftParams.fftSize);
 
        if (binParams.isNPTSodd)
        {
            // Adjust points on either side of Nyquist for odd NPTS.
            w1(binParams.halfNPTS - 1) = Nyq - halfRes;
            w1(binParams.halfNPTS) = Nyq + halfRes;
        }
        else
        {
            // Make sure we hit Nyquist exactly for even NPTS.
            w1(binParams.halfNPTS - 1) = Nyq;
        }
 
        // Adjust the last point.
        // Assuming NPTS is passed correctly and maps to w1.size() for the conversion
        w1(w1.size() - 1) = dftParams.sampleRate - freqRes;
 
        // Get the right grid based on range, centerdc, etc.
        Eigen::VectorXf binsVector = finalGrid(w1, Nyq, false /*TODO: Check this*/, freqRange);
        std::vector<float> binsVectorSTL(binsVector.data(), binsVector.data() + binsVector.size());
 
        return binsVectorSTL;
    }
 
    Eigen::VectorXf ShortTimeFourierTransform::getBins(ShortTimeFourierTransform::FrequencyRange freqRange)
    {
        float freqRes = dftParams.sampleRate / static_cast<float>(dftParams.fftSize);
        Eigen::VectorXf w1 = freqRes * Eigen::VectorXf::LinSpaced(dftParams.fftSize, 0, dftParams.fftSize - 1).array();
        float Nyq = dftParams.sampleRate / 2.0f;
        float halfRes = freqRes / 2;
        // NPT Info
        NPTSinfo(dftParams.fftSize);
 
        if (binParams.isNPTSodd)
        {
            // Adjust points on either side of Nyquist for odd NPTS.
            w1(binParams.halfNPTS - 1) = Nyq - halfRes;
            w1(binParams.halfNPTS) = Nyq + halfRes;
        }
        else
        {
            // Make sure we hit Nyquist exactly for even NPTS.
            w1(binParams.halfNPTS - 1) = Nyq;
        }
 
        // Adjust the last point.
        // Assuming NPTS is passed correctly and maps to w1.size() for the conversion
        w1(w1.size() - 1) = dftParams.sampleRate - freqRes;
 
        // Get the right grid based on range, centerdc, etc.
        return finalGrid(w1, Nyq, false /*TODO: Check this*/, freqRange);
    }
 
    std::vector<float> ShortTimeFourierTransform::getColumnsSTL()
    {
        // You need to have already computed stft at this point!!
        jassert(m_inputSize != 0);
        int nx = m_inputSize + /*Careful this might cause problem!*/ dftParams.winSize;
        size_t nCol = std::floor((nx - dftParams.nOverlap) / dftParams.hopSize);
 
        // Determine the number of columns of the STFT output(i.e., the S output)
        Eigen::ArrayXf colOffsets = Eigen::ArrayXf::LinSpaced(nCol, 0, (nCol - 1) * dftParams.hopSize);
        Eigen::VectorXf columns = (colOffsets + (dftParams.winSize / 2.0)).transpose() / dftParams.sampleRate;
        std::vector<float> columnsSTL(columns.data(), columns.data() + columns.size());
 
        return columnsSTL;
    }
 
    Eigen::VectorXf ShortTimeFourierTransform::getColumns()
    {
        // You need to have already computed stft at this point!!
        jassert(m_inputSize != 0);
        int nx = m_inputSize + /*Careful this might cause problem!*/ dftParams.winSize;
        size_t nCol = std::floor((nx - dftParams.nOverlap) / dftParams.hopSize);
 
        // Determine the number of columns of the STFT output(i.e., the S output)
        Eigen::ArrayXf colOffsets = Eigen::ArrayXf::LinSpaced(nCol, 0, (nCol - 1) * dftParams.hopSize);
        Eigen::VectorXf columns = (colOffsets + (dftParams.winSize / 2.0)).transpose() / dftParams.sampleRate;
 
        return columns;
    }
 
    float ShortTimeFourierTransform::getFFTBinSize() const
    {
        jassert(binParams.binSize != -1.0f);
 
        return binParams.binSize;
    }
 
    void ShortTimeFourierTransform::NPTSinfo(float NPTS)
    {
        binParams.isNPTSodd = false;
        if (static_cast<int>(NPTS) % 2 != 0)
            binParams.isNPTSodd = true;
 
        if (binParams.isNPTSodd)
            binParams.halfNPTS = (NPTS + 1) / 2;
        else
            binParams.halfNPTS = NPTS / 2 + 1;
 
        binParams.isHalfNPTSodd = false;
        if (static_cast<int>(binParams.halfNPTS) % 2 != 0)
            binParams.isHalfNPTSodd = true;
 
        if (binParams.isHalfNPTSodd)
            binParams.quarterNPTS = (binParams.halfNPTS + 1) / 2;
        else
            binParams.quarterNPTS = binParams.halfNPTS / 2 + 1;
    }
 
    Eigen::VectorXf ShortTimeFourierTransform::finalGrid(const Eigen::VectorXf& w1, float Nyq, bool centerDC,
                                                         ShortTimeFourierTransform::FrequencyRange freqRange)
    {
        jassert(binParams.halfNPTS != -1 && binParams.quarterNPTS != -1); // You need to have calculated them by now!
 
        Eigen::VectorXf w;
        switch (freqRange)
        {
        case FrequencyRange::Whole: {
            w.resize(dftParams.fftSize);
            if (centerDC)
            {
                int negEndPt = binParams.isNPTSodd ? binParams.halfNPTS : binParams.halfNPTS - 1;
                Eigen::VectorXf temp(negEndPt * 2);
                temp.head(negEndPt) = w1.segment(1, negEndPt).reverse();
                temp.tail(negEndPt) = w1.head(negEndPt);
                w = temp;
            }
            else
            {
                w = w1;
            }
            break;
        }
        case FrequencyRange::Half: {
            w = w1.head(binParams.halfNPTS);
            if (centerDC)
            {
                int negEndPt = binParams.isHalfNPTSodd ? binParams.quarterNPTS : binParams.quarterNPTS - 1;
                Eigen::VectorXf temp(negEndPt * 2);
                temp.head(negEndPt) = w1.segment(1, negEndPt).reverse();
                temp.tail(negEndPt) = w1.head(binParams.quarterNPTS);
                w = temp;
                if (dftParams.fftSize % 4 == 0)
                {
                    w(w.size() - 1) = Nyq / 2;
                }
            }
            break;
        }
        default:
            throw std::runtime_error("Unsupported frequency range.");
        }
        return w;
    }
 
} // namespace krotos