welch_psd.cpp

Go to the documentation of this file.

//=============================================================================================================
 
//=============================================================================================================
// INCLUDES
//=============================================================================================================
 
#include "welch_psd.h"
 
//=============================================================================================================
// EIGEN INCLUDES
//=============================================================================================================
 
#include <Eigen/Core>
//#ifndef EIGEN_FFTW_DEFAULT
//#define EIGEN_FFTW_DEFAULT
//#endif
#include <unsupported/Eigen/FFT>
 
//=============================================================================================================
// STD INCLUDES
//=============================================================================================================
 
#include <cmath>
#include <vector>
 
//=============================================================================================================
// USED NAMESPACES
//=============================================================================================================
 
using namespace UTILSLIB;
using namespace Eigen;
 
namespace
{
constexpr double WELCH_PI = 3.14159265358979323846;
}
 
//=============================================================================================================
// DEFINE MEMBER METHODS
//=============================================================================================================
 
VectorXd WelchPsd::buildWindow(int iN, WindowType window)
{
    VectorXd w(iN);
    const double pi2 = 2.0 * WELCH_PI;
 
    for (int n = 0; n < iN; ++n) {
        const double t = static_cast<double>(n) / static_cast<double>(iN);
        switch (window) {
            case Hann:
                w[n] = 0.5 * (1.0 - std::cos(pi2 * t));
                break;
            case Hamming:
                w[n] = 0.54 - 0.46 * std::cos(pi2 * t);
                break;
            case Blackman:
                w[n] = 0.42 - 0.5 * std::cos(pi2 * t) + 0.08 * std::cos(2.0 * pi2 * t);
                break;
            case FlatTop:
                w[n] = 1.0 - 1.93293488969 * std::cos(pi2 * t) + 1.28349769674 * std::cos(2.0 * pi2 * t) - 0.38763473916 * std::cos(3.0 * pi2 * t) + 0.03279543650 * std::cos(4.0 * pi2 * t);
                break;
        }
    }
    return w;
}
 
//=============================================================================================================
 
RowVectorXd WelchPsd::freqAxis(int iNfft, double dSFreq)
{
    const int iNFreqs = iNfft / 2 + 1;
    RowVectorXd freqs(iNFreqs);
    for (int k = 0; k < iNFreqs; ++k)
        freqs[k] = static_cast<double>(k) * dSFreq / static_cast<double>(iNfft);
    return freqs;
}
 
//=============================================================================================================
 
RowVectorXd WelchPsd::computeVector(const RowVectorXd& vecData,
                                    double dSFreq,
                                    int iNfft,
                                    double dOverlap,
                                    WindowType window,
                                    bool bRemoveDc)
{
    const int nIn = static_cast<int>(vecData.cols());
    if (iNfft > nIn)
        iNfft = nIn;
 
    const int iNFreqs = iNfft / 2 + 1;
    const int iStep = std::max(1, static_cast<int>(std::round(static_cast<double>(iNfft) * (1.0 - dOverlap))));
 
    const VectorXd w = buildWindow(iNfft, window);
    const double dWinPow = w.squaredNorm(); // Σ w²
 
    Eigen::FFT<double> fft;
    RowVectorXd psd = RowVectorXd::Zero(iNFreqs);
    int nSeg = 0;
 
    for (int start = 0; start + iNfft <= nIn; start += iStep) {
        // Apply window and forward FFT
        VectorXd seg = vecData.segment(start, iNfft).transpose();
        if (bRemoveDc)
            seg.array() -= seg.mean();
        seg.array() *= w.array();
        VectorXcd spec;
        fft.fwd(spec, seg);
 
        for (int k = 0; k < iNFreqs; ++k)
            psd[k] += std::norm(spec[k]); // accumulate |X[k]|²
        ++nSeg;
    }
 
    if (nSeg == 0)
        return RowVectorXd::Zero(iNFreqs);
 
    // Normalise: divide by (nSeg × window_power × sfreq)
    const double dNorm = 1.0 / (static_cast<double>(nSeg) * dWinPow * dSFreq);
    psd *= dNorm;
 
    // One-sided: double every bin with a mirror image (not DC, not an even-length Nyquist)
    const int kLast = (iNfft % 2 == 0) ? iNFreqs - 1 : iNFreqs;
    for (int k = 1; k < kLast; ++k)
        psd[k] *= 2.0;
 
    return psd;
}
 
//=============================================================================================================
 
WelchPsdResult WelchPsd::compute(const MatrixXd& matData,
                                 double dSFreq,
                                 int iNfft,
                                 double dOverlap,
                                 WindowType window,
                                 const RowVectorXi& vecPicks,
                                 bool bRemoveDc)
{
    std::vector<int> picks;
    if (vecPicks.size() > 0) {
        picks.reserve(static_cast<std::size_t>(vecPicks.size()));
        for (int i = 0; i < vecPicks.size(); ++i)
            picks.push_back(vecPicks[i]);
    } else {
        picks.reserve(static_cast<std::size_t>(matData.rows()));
        for (int i = 0; i < static_cast<int>(matData.rows()); ++i)
            picks.push_back(i);
    }
 
    const int iNFreqs = iNfft / 2 + 1;
    WelchPsdResult result;
    result.matPsd.resize(static_cast<int>(picks.size()), iNFreqs);
    result.vecFreqs = freqAxis(iNfft, dSFreq);
 
    for (int ci = 0; ci < static_cast<int>(picks.size()); ++ci)
        result.matPsd.row(ci) = computeVector(matData.row(picks[ci]),
                                              dSFreq, iNfft, dOverlap, window, bRemoveDc);
    return result;
}