movie_alignment_correlation_gpu.h

Go to the documentation of this file.

/***************************************************************************
 *
 * Authors:    David Strelak (davidstrelak@gmail.com)
 *
 * Unidad de  Bioinformatica of Centro Nacional de Biotecnologia , CSIC
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
 * 02111-1307  USA
 *
 *  All comments concerning this program package may be sent to the
 *  e-mail address 'xmipp@cnb.csic.es'
 ***************************************************************************/

#pragma once

#include "reconstruction/movie_alignment_correlation_base.h"
#include "data/fft_settings.h"
#include "reconstruction_cuda/gpu.h"
#include "reconstruction_cuda/cuda_flexalign_scale.h"
#include "reconstruction_cuda/cuda_flexalign_correlate.h"

template<typename T>
class ProgMovieAlignmentCorrelationGPU: public AProgMovieAlignmentCorrelation<T> {
protected:
    void defineParams();
    void show();
    void readParams();

private:

    auto findGoodCorrelationSize(const Dimensions &ref);

    auto findGoodPatchSize();

    auto findGoodCropSize(const Dimensions &ref);
        
    auto getCorrelationHint(const Dimensions &d);

    std::vector<FramePatchMeta<T>> getPatchesLocation(const std::pair<T, T> &borders,
            const Dimensions &patch);

    void getPatchData(const Rectangle<Point2D<T>> &patch,
            const AlignmentResult<T> &globAlignment,
            T *result);

    std::pair<T,T> getMovieBorders(const AlignmentResult<T> &globAlignment, int verbose);

    void LAOptimize();

    AlignmentResult<T> computeGlobalAlignment(const MetaData &movie,
            const Image<T> &dark,
            const Image<T> &igain) override;

    T* setFilter(float scale, const Dimensions &dims);

    struct AlignmentContext
    { 
        int verbose;
        float maxShift;
        size_t N; // number of frames 
        std::pair<T, T> scale;
        core::optional<size_t> refFrame;
        Dimensions out = Dimensions(0);
        size_t alignmentBytes() const {
            return (N * (N-1) / 2) * 2 * sizeof(float); // 2D position for each correlation
        }
    };

    LocalAlignmentResult<T> localFromGlobal(const AlignmentResult<T> &globAlignment);

    void applyShiftsComputeAverage(
            const MetaData& movieMD, const Image<T>& dark, const Image<T>& igain,
            Image<T>& initialMic, size_t& Ninitial, Image<T>& averageMicrograph,
            size_t& N, const AlignmentResult<T> &globAlignment) override;

    
    auto getOutputStreamCount();

    void storeSizes(const Dimensions &orig, const Dimensions &opt, bool applyCrop);

    std::optional<Dimensions> getStoredSizes(const Dimensions &dim, bool applyCrop);

    void GAOptimize();

    void report(float scale, 
        typename CUDAFlexAlignScale<T>::Params &SP, 
        typename CUDAFlexAlignCorrelate<T>::Params &CP, 
        int streams, int threads);

    void getCroppedFrame(const Dimensions &settings, T *output, T *src);

    class Movie final
    {
    public:
        auto set(const Dimensions &dim) {
            mDim = dim;
            mFrames.reserve(dim.n());
            for (auto n = 0; n < dim.n(); ++n) {
                // create multidim arrays, but without data 
                mFrames.emplace_back(1, 1, dim.y(), dim.x(), nullptr);
            }
        }

        auto &getDim() const {
            return mDim;
        }

        auto &getFrame(size_t index) const {
            return mFrames[index];
        }

        auto setFrameData(size_t index, T *ptr) {
            mFrames[index].data = ptr;
        }

    private:
        std::vector<MultidimArray<T>> mFrames;
        Dimensions mDim = Dimensions(0);
    };

    T* loadFrame(const MetaData& movieMD, const Image<T>& dark, const Image<T>& igain, size_t index);

    auto computeAlignment(float *positions, const AlignmentContext &context);

    void releaseAll() override;

    LocalAlignmentResult<T> computeLocalAlignment(const MetaData &movie,
        const Image<T> &dark, const Image<T> &igain,
        const AlignmentResult<T> &globAlignment) override;

    std::string getKey(const std::string &keyword, const Dimensions &dim, bool crop) const {
        std::stringstream ss;
        ss << version << " " << mGpu.value().getUUID() << keyword << dim << " " << crop;
        return ss.str();
    }        

    void applyShiftsComputeAverage(
        const MetaData& movie, const Image<T>& dark, const Image<T>& igain,
        Image<T>& initialMic, size_t& Ninitial, Image<T>& averageMicrograph,
        size_t& N, const LocalAlignmentResult<T> &alignment) override;

    // holding the loaded movie
    Movie movie;
    // number of GPU streams for global alignment
    int GAStreams = 1;
    // parameters for global alignment scaling
    typename CUDAFlexAlignScale<T>::Params GASP;
    // parameters for global alignment correlation
    typename CUDAFlexAlignCorrelate<T>::Params GACP; 
    // parameters for local alignment scaling
    typename CUDAFlexAlignScale<T>::Params LASP;
    // parameters for local alignment correlation
    typename CUDAFlexAlignCorrelate<T>::Params LACP; 

    int patchesAvg;

    std::string storage;

    core::optional<GPU> mGpu;

    std::string minMemoryStr = std::string("minMem");
    std::string optSizeXStr = std::string("optSizeX");
    std::string optSizeYStr = std::string("optSizeY");
    std::string optBatchSizeStr = std::string("optBatchSize");

    static constexpr auto version = "2.0";
};

//@}