wstackinggridder.h

//
// This file was written by André Offringa and is published
// under the GPL 3 license. 
//

#ifndef WSTACKING_GRIDDER_H
#define WSTACKING_GRIDDER_H

#ifndef AVOID_CASACORE
#include "../multibanddata.h"
#endif

// Forward declare boost::mutex to avoid header inclusion
namespace boost
{
        class mutex;
}

#include <cmath>
#include <cstring>
#include <complex>
#include <vector>
#include <stack>

class ImageBufferAllocator;

class WStackingGridder
{
        public:
                enum GridModeEnum {
                        
                        NearestNeighbourGridding,
                        
                        KaiserBesselKernel,
                        
                        RectangularKernel
                };
                
                WStackingGridder(size_t width, size_t height, double pixelSizeX, double pixelSizeY, size_t fftThreadCount, ImageBufferAllocator* allocator, size_t kernelSize = 7, size_t overSamplingFactor = 63);
                
                WStackingGridder(const WStackingGridder&) = delete;
                
                WStackingGridder& operator=(const WStackingGridder&) = delete;
                
                ~WStackingGridder();
                
                void PrepareWLayers(size_t nWLayers, double maxMem, double minW, double maxW);
                
#ifndef AVOID_CASACORE

                void PrepareBand(const MultiBandData &bandData)
                {
                        _bandData = bandData;
                }
#endif // AVOID_CASACORE
                
                size_t WToLayer(double wInLambda) const
                {
                        if(_nWLayers == 1)
                                return 0;
                        else {
                                if(_isComplex)
                                        return size_t(round((wInLambda + _maxW) * (_nWLayers-1) / (_maxW + _maxW)));
                                else
                                        return size_t(round((fabs(wInLambda) - _minW) * (_nWLayers-1) / (_maxW - _minW)));
                        }
                }
                
                double LayerToW(size_t layer) const
                {
                        if(_nWLayers == 1)
                                return 0.0;
                        else {
                                if(_isComplex)
                                        return layer * (_maxW + _maxW) / (_nWLayers-1) - _maxW;
                                else
                                        return layer * (_maxW - _minW) / (_nWLayers-1) + _minW;
                        }
                }
                
                size_t NWLayers() const { return _nWLayers; }
                
                bool IsInLayerRange(double wInLambda) const
                {
                        size_t layer = WToLayer(wInLambda);
                        return layer >= layerRangeStart(_curLayerRangeIndex) && layer < layerRangeStart(_curLayerRangeIndex+1);
                }
                
                
                bool IsInLayerRange(double wStart, double wEnd) const
                {
                        size_t
                                rangeStart = layerRangeStart(_curLayerRangeIndex),
                                rangeEnd = layerRangeStart(_curLayerRangeIndex+1),
                                l1 = WToLayer(wStart);
                        if(l1 >= rangeStart && l1 < rangeEnd)
                                return true;
                        size_t l2 = WToLayer(wEnd);
                        return ((l2 >= rangeStart && l2 < rangeEnd) // l2 is within the range
                          || (l1 < rangeStart && l2 >= rangeEnd)  // l1 is before, l2 is after range
                                || (l2 < rangeStart && l1 >= rangeEnd) // l2 is before, l1 is after range
                        );
                }
                
                size_t NPasses() const
                {
                        return _nPasses;
                }
                
#ifndef AVOID_CASACORE

                void AddData(const std::complex<float>* data, size_t dataDescId, double uInM, double vInM, double wInM);
#endif
                
                void AddDataSample(std::complex<float> sample, double uInLambda, double vInLambda, double wInLambda);
                
                void StartInversionPass(size_t passIndex);
                
                void FinishInversionPass();
                
                void FinalizeImage(double multiplicationFactor, bool correctFFTFactor);
                
                void InitializePrediction(const double *image)
                {
                        initializePrediction(image, _imageData);
                }
                
                void InitializePrediction(const double *real, const double *imaginary)
                {
                        initializePrediction(real, _imageData);
                        initializePrediction(imaginary, _imageDataImaginary);
                }

                void StartPredictionPass(size_t passIndex);
                
#ifndef AVOID_CASACORE

                void SampleData(std::complex<float>* data, size_t dataDescId, double uInM, double vInM, double wInM);
#endif
                
                void SampleDataSample(std::complex<double>& value, double uInLambda, double vInLambda, double wInLambda);
                
                void SampleDataSample(std::complex<float>& value, double uInLambda, double vInLambda, double wInLambda)
                {
                        std::complex<double> doubleValue;
                        SampleDataSample(doubleValue, uInLambda, vInLambda, wInLambda);
                        value = doubleValue;
                }
                
                double *RealImage() { return _imageData[0]; }
                
                double *ImaginaryImage() { return _imageDataImaginary[0]; }
                
                size_t NFFTThreads() const { return _nFFTThreads; }
                
                void SetNFFTThreads(size_t nfftThreads) { _nFFTThreads = nfftThreads; }
                
                enum GridModeEnum GridMode() const { return _gridMode; }
                
                void SetGridMode(enum GridModeEnum mode) {
                        if(mode != _gridMode)
                        {
                                _gridMode = mode;
                                if(_gridMode != NearestNeighbourGridding)
                                        makeKernels();
                        }
                }
                
                bool IsComplex() const { return _isComplex; }
                
                void SetIsComplex(bool isComplex) { _isComplex = isComplex; }
                
                //void SetImageConjugatePart(bool imageConjugatePart) { _imageConjugatePart = imageConjugatePart; }
                
                void SetDenormalPhaseCentre(double dl, double dm) { _phaseCentreDL = dl; _phaseCentreDM = dm; }
                
                void GetGriddingCorrectionImage(double* image) const;
                
                void ReplaceRealImageBuffer(double* newBuffer);
                
                void ReplaceImaginaryImageBuffer(double* newBuffer);
                
                const std::complex<double>* GetGriddedUVLayer(size_t layerIndex) const
                {
                        return _layeredUVData[layerIndex];
                }
                
                void GetKaiserBesselKernel(double* kernel, size_t n, bool multiplyWithSinc);
                
                size_t Width() const { return _width; }
                
                size_t Height() const { return _height; }
                
                double PixelSizeX() const { return _pixelSizeX; }
                
                double PixelSizeY() const { return _pixelSizeY; }
                
                ImageBufferAllocator* Allocator() const { return _imageBufferAllocator; }
        private:
                size_t layerRangeStart(size_t layerRangeIndex) const
                {
                        return (_nWLayers * layerRangeIndex) / _nPasses;
                }
                template<bool IsComplexImpl>
                void projectOnImageAndCorrect(const std::complex<double> *source, double w, size_t threadIndex);
                template<bool IsComplexImpl>
                void copyImageToLayerAndInverseCorrect(std::complex<double> *dest, double w);
                void initializeSqrtLMLookupTable();
                void initializeSqrtLMLookupTableForSampling();
                void initializeLayeredUVData(size_t n);
                void freeLayeredUVData() { initializeLayeredUVData(0); }
                void fftToImageThreadFunction(boost::mutex *mutex, std::stack<size_t> *tasks, size_t threadIndex);
                void fftToUVThreadFunction(boost::mutex *mutex, std::stack<size_t> *tasks);
                void finalizeImage(double multiplicationFactor, std::vector<double*>& dataArray);
                void initializePrediction(const double *image, std::vector<double*>& dataArray);
                
                void makeKernels();
                void makeKaiserBesselKernel(std::vector<double> &kernel, double alpha, size_t overSamplingFactor, bool withSinc=true);
                
                void makeRectangularKernel(std::vector<double> &kernel, size_t overSamplingFactor);
                
                double bessel0(double x, double precision);
                template<bool Inverse>
                void correctImageForKernel(double *image) const;
                
                const size_t _width, _height;
                const double _pixelSizeX, _pixelSizeY;
                size_t _nWLayers, _nPasses, _curLayerRangeIndex;
                double _minW, _maxW, _phaseCentreDL, _phaseCentreDM;
                bool _isComplex, _imageConjugatePart;
#ifndef AVOID_CASACORE
                MultiBandData _bandData;
#endif
                
                enum GridModeEnum _gridMode;
                size_t _overSamplingFactor, _kernelSize;
                std::vector<double> _1dKernel;
                std::vector<std::vector<double>> _griddingKernels;
                
                std::vector<std::complex<double>*> _layeredUVData;
                std::vector<double*> _imageData, _imageDataImaginary;
                std::vector<double> _sqrtLMLookupTable;
                size_t _nFFTThreads;
                ImageBufferAllocator* _imageBufferAllocator;
};

#endif