WfAnalyzer.h

#pragma once

// A waveform analyzer class for the Fadc250 raw waveform data
// It finds the local maxima (peaks) over the spectrum, and detemines the peak properties
// It is intended to use simple algorithms since we do not expect heavily convoluted spectra
//
// History:
// v1.0.0: a working version with basic functionalities, Chao Peng, 2020/08/04

#include <vector>
#include <algorithm>

// calculate the mean and standard deviation for an array
namespace wfa {

// data structures
struct Peak {
    size_t pos, left, right;
    double height, integral;

    bool Inside(size_t i) {
        return (i >= pos - left) && (i <= pos + right);
    }
};

struct Pedestal {
    double mean, err;
};

struct WFData {
    Pedestal ped;
    std::vector<Peak> peaks;
};

// calculate mean and standard deviation of an array
template<typename T>
inline void _calc_mean_err(double &mean, double &err, const T *source, size_t npts)
{
    if (npts == 0) { return; }

    mean = 0.;
    err = 0.;

    for (size_t i = 0; i < npts; ++i) {
       mean += source[i];
    }
    mean /= static_cast<double>(npts);

    for (size_t i = 0; i < npts; ++i) {
       err += (source[i] - mean ) * (source[i] - mean);
    }
    err = std::sqrt(err/static_cast<double>(npts));
}

template<typename T>
inline void _linear_regr(double &p0, double &p1, double &err, const T *x, const T *y, size_t npts)
{
    // no data points
    if (npts == 0) {
        return;
    // too few data points
    } else if (npts < 3) {
        p1 = 0.;
        _calc_mean_err(p0, err, y, npts);
        return;
    }

    double x_mean, x_err;
    _calc_mean_err(x_mean, x_err, x, npts);
    double y_mean, y_err;
    _calc_mean_err(y_mean, y_err, y, npts);

    double xysum = 0.;
    for (size_t i = 0; i < npts; ++i) {
        xysum += (x[i] - x_mean) * (y[i] - y_mean);
    }

    double r = xysum/(x_err*y_err*npts);

    p1 = r*y_err/x_err;
    p0 = y_mean - p1*x_mean;

    err = 0.;
    for (size_t i = 0; i < npts; ++i) {
        double diff = (p0 + p1*x[i] - y[i]);
        err += diff*diff;
    }
    err = std::sqrt(err/static_cast<double>(npts));
}

// analyzer class
class Analyzer
{
public:
    Analyzer(size_t resolution = 2, double threshold = 10.0)
        : thres(threshold), res(resolution) {}
    ~Analyzer() {}

    WFData Analyze(const uint32_t *samples, size_t nsamples)
    {
        WFData data;
        if (!nsamples) { return data; }

        auto buffer = SmoothSpectrum(samples, nsamples, res);

        // search local maxima
        std::vector<Peak> candidates;
        candidates.reserve(nsamples/2);
        // get trend
        auto trend = [] (double v1, double v2, double thr = 0.1) {
            return std::abs(v1 - v2) < thr ? 0 : (v1 > v2 ? 1 : -1);
        };
        for (size_t i = 1; i < nsamples - 1; ++i) {
            int tr1 = trend(buffer[i], buffer[i - 1]);
            int tr2 = trend(buffer[i], buffer[i + 1]);
            // peak at the rising (declining) edge
            if ((tr1 * tr2 >= 0) && (std::abs(tr1) > 0)) {
                size_t left = 1, right = 1;
                // search the peak range
                while ((i > left + 1) && (trend(buffer[i - left], buffer[i - left - 1]) == tr1)) {
                    left ++;
                }
                while ((i + right < nsamples - 1) && (trend(buffer[i + right], buffer[i + right + 1])*tr1 >= 0)) {
                    right ++;
                }

                double base = buffer[i - left] * right / static_cast<double>(left + right)
                            + buffer[i + right] * left / static_cast<double>(left + right);
                double height = std::abs(buffer[i] - base);
                if (height > thres) {
                    candidates.push_back(Peak{i, i - left, i + right, buffer[i] - base, 0});
                }
            }
        }

        auto in_peaks = [] (const std::vector<Peak> &peaks, size_t pos) {
            for (auto &peak : peaks) {
                // determine window size
                size_t width = std::min(peak.pos - peak.left, peak.right - peak.pos);
                if (std::abs(peak.pos - static_cast<double>(pos)) <= width) { return true; }
            }
            return false;
        };

        size_t nped = 0;
        auto ybuf = buffer;
        for (size_t i = res + 1; i < nsamples - res - 1; ++i) {
            if (!in_peaks(candidates, i)) {
                ybuf[nped] = buffer[i];
                nped ++;
            }
        }

        // calculate pedestal
        if (nped <= nsamples/4) {
            data.ped = CalcPedestal(&ybuf[0], nsamples, 1.5, 3);
        } else {
            data.ped = CalcPedestal(&ybuf[0], nped, 1.2, 3);
        }

        for (auto &peak : candidates) {
            // real sample height
            double sample_height = samples[peak.pos] - data.ped.mean;
            if ((sample_height * peak.height < 0.) ||
                (std::abs(sample_height) < thres) ||
                (std::abs(sample_height) < 3.0*data.ped.err)) {
                continue;
            }

            // integrate it over the peak range
            peak.integral = buffer[peak.pos] - data.ped.mean;
            // i may go below 0
            for (int i = peak.pos - 1; i >= static_cast<int>(peak.left); --i) {
                double val = buffer[i] - data.ped.mean;
                // stop when it touches or acrosses the baseline
                if (std::abs(val) < data.ped.err || val * sample_height < 0.) {
                    peak.left = i; break;
                }
                peak.integral += val;
            }
            for (size_t i = peak.pos + 1; i <= peak.right; ++i) {
                double val = buffer[i] - data.ped.mean;
                if (std::abs(val) < data.ped.err || val * sample_height < 0.) {
                    peak.right = i; break;
                }
                peak.integral += val;
            }

            // determine the real sample peak
            peak.height = sample_height;
            auto update_peak = [] (Peak &peak, double val, size_t pos) {
                if (std::abs(val) > std::abs(peak.height)) {
                    peak.pos = pos;
                    peak.height = val;
                }
            };
            size_t sample_pos = peak.pos;
            for (size_t i = 1; i < res; ++i) {
                if (sample_pos > i) {
                    update_peak(peak, samples[sample_pos - i] - data.ped.mean, sample_pos - i);
                }
                if (sample_pos + i < nsamples) {
                    update_peak(peak, samples[sample_pos + i] - data.ped.mean, sample_pos + i);
                }
            }
            // fill to results
            data.peaks.emplace_back(peak);
        }

        /*
        std::sort(peaks.begin(), peaks.end(),
                  [] (const Peak &p1, const Peak &p2) { return p1.height > p2.height; });
        */
        return data;
    }

private:
    double thres;
    size_t res;


public:
    // static methods
    template<typename T>
    static std::vector<double> SmoothSpectrum(const T *samples, size_t nsamples, size_t res)
    {
        std::vector<double> buffer(nsamples);
        for (size_t i = 0; i < nsamples; ++i) {
            double val = samples[i];
            double weights = 1.0;
            for (size_t j = 1; j < res; ++j) {
                if (j >= i || j + i >= nsamples) { continue; }
                double weight = 1.0 - j/static_cast<double>(res + 1);
                val += weight*(samples[i - j] + samples[i + j]);
                weights += 2.*weight;
            }
            buffer[i] = val/weights;
        }
        return buffer;
    }

    template<typename T>
    static Pedestal CalcPedestal(T *ybuf, size_t npts, double thres = 1.0, int max_iters = 3)
    {
        Pedestal res;
        _calc_mean_err(res.mean, res.err, ybuf, npts);
        // interatively fit
        while (max_iters-- > 0) {
            size_t count = 0;
            for (size_t i = 0; i < npts; ++i) {
                if (std::abs(ybuf[i] - res.mean) < res.err*thres) {
                    ybuf[count] = ybuf[i];
                    count++;
                }
            }
            if (count == npts) { break; }
            npts = count;
            _calc_mean_err(res.mean, res.err, ybuf, npts);
        }

        return res;
    }

};

};  // namespace wfa