ctf_estimate_from_micrograph.cpp

Go to the documentation of this file.

/***************************************************************************
 *
 * Authors:     Javier Angel Velazquez Muriel (javi@cnb.csic.es)
 *              Carlos Oscar Sanchez Sorzano
 *
 * 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'
 ***************************************************************************/

#include "ctf_estimate_from_micrograph.h"
#include "ctf_enhance_psd.h"
#include "core/xmipp_fftw.h"
#include "core/metadata_extension.h"
#include "core/transformations.h"
#include "core/xmipp_threads.h"
#include "core/xmipp_image_generic.h"
#include "data/basic_pca.h"
#include "data/normalize.h"
#include "data/numerical_tools.h"

/* Read parameters ========================================================= */
ProgCTFEstimateFromMicrograph::ProgCTFEstimateFromMicrograph()
{
    psd_mode = OnePerMicrograph; 
    PSDEstimator_mode = Periodogram;
}

void ProgCTFEstimateFromMicrograph::readParams()
{
    fn_micrograph = getParam("--micrograph");
    fn_root = getParam("--oroot");
    if (fn_root == "")
        fn_root = fn_micrograph.withoutExtension();
    pieceDim = getIntParam("--pieceDim");
    skipBorders = getIntParam("--skipBorders");
    overlap = getDoubleParam("--overlap");
    String aux = getParam("--psd_estimator");
    if (aux == "periodogram")
        PSDEstimator_mode = Periodogram;
    else
    {
        PSDEstimator_mode = ARMA;
        ARMA_prm.readParams(this);
    }
    Nsubpiece = getIntParam("--Nsubpiece");

    String mode = getParam("--mode");
    if (mode == "micrograph")
        psd_mode = OnePerMicrograph;
    else if (mode == "regions")
    {
        psd_mode = OnePerRegion;
        fn_pos = getParam("--mode", 1);
    }
    else if (mode == "particles")
    {
        psd_mode = OnePerParticle;
        fn_pos = getParam("--mode", 1);
    }
    estimate_ctf = !checkParam("--dont_estimate_ctf");
    acceleration1D = checkParam("--acceleration1D");

    if (estimate_ctf)
    {
        if (!acceleration1D)
            prmEstimateCTFFromPSD.readBasicParams(this);
        else
            prmEstimateCTFFromPSDFast.readBasicParams(this);
    }

    bootstrapN = getIntParam("--bootstrapFit");
}

void ProgCTFEstimateFromMicrograph::defineParams()
{
    addUsageLine("Estimate the CTF from a micrograph.");
    addUsageLine("The PSD of the micrograph is first estimated using periodogram averaging or ");
    addUsageLine("ARMA models ([[http://www.ncbi.nlm.nih.gov/pubmed/12623169][See article]]). ");
    addUsageLine("Then, the PSD is enhanced ([[http://www.ncbi.nlm.nih.gov/pubmed/16987671][See article]]). ");
    addUsageLine("And finally, the CTF is fitted to the PSD, being guided by the enhanced PSD ");
    addUsageLine("([[http://www.ncbi.nlm.nih.gov/pubmed/17911028][See article]]).");
    addParamsLine("   --micrograph <file>         : File with the micrograph");
    addParamsLine("  [--oroot <rootname=\"\">]    : Rootname for output");
    addParamsLine("                               : If not given, the micrograph without extensions is taken");
    addParamsLine("                               :++ rootname.psd or .psdstk contains the PSD or PSDs");
    addParamsLine("==+ PSD estimation");
    addParamsLine("  [--psd_estimator <method=periodogram>] : Method for estimating the PSD");
    addParamsLine("         where <method>");
    addParamsLine("                  periodogram");
    addParamsLine("                  ARMA");
    addParamsLine("  [--pieceDim <d=512>]       : Size of the piece");
    addParamsLine("  [--overlap <o=0.5>]        : Overlap (0=no overlap, 1=full overlap)");
    addParamsLine("  [--skipBorders <s=2>]      : Number of pieces around the border to skip");
    addParamsLine("  [--Nsubpiece <N=1>]        : Each piece is further subdivided into NxN subpieces.");
    addParamsLine("                              : This option is useful for small micrographs in which ");
    addParamsLine("                              : not many pieces of size pieceDim x pieceDim can be defined. ");
    addParamsLine("                              :++ Note that this is not the same as defining a smaller pieceDim. ");
    addParamsLine("                              :++ Defining a smaller pieceDim, would result in a small PSD, while ");
    addParamsLine("                              :++ subdividing the piece results in a large PSD, although smoother.");
    addParamsLine("  [--mode <mode=micrograph>]  : How many PSDs are to be estimated");
    addParamsLine("         where <mode>");
    addParamsLine("                  micrograph  : Single PSD for the whole micrograph");
    addParamsLine("                  regions <file=\"\"> : The micrograph is divided into a region grid ");
    addParamsLine("                              : and a PSD is computed for each one.");
    addParamsLine("                              : The file is metadata with the position of each particle within the micrograph");
    addParamsLine("                  particles <file> : One PSD per particle.");
    addParamsLine("                              : The file is metadata with the position of each particle within the micrograph");
    addParamsLine("==+ CTF fit");
    addParamsLine("  [--dont_estimate_ctf]       : Do not fit a CTF to PSDs");
    addParamsLine("  [--acceleration1D]          : Accelerate PSD estimation");
    ARMA_parameters::defineParams(this);
    ProgCTFEstimateFromPSD::defineBasicParams(this);
    ProgCTFBasicParams::defineBasicParams(this);
    addExampleLine("Estimate PSD", false);
    addExampleLine("xmipp_ctf_estimate_from_micrograph --micrograph micrograph.mrc --dont_estimate_ctf");
    addExampleLine("Estimate a single CTF for the whole micrograph", false);
    addExampleLine("xmipp_ctf_estimate_from_micrograph --micrograph micrograph.mrc --sampling_rate 1.4 --voltage 200 --spherical_aberration 2.5");
    addExampleLine("Estimate a single CTF for the whole micrograph providing a starting point for the defocus",false);
    addExampleLine("xmipp_ctf_estimate_from_micrograph --micrograph micrograph.mrc --sampling_rate 1.4 --voltage 200 --spherical_aberration 2.5 --defocusU -15000");
    addExampleLine("Estimate a CTF per region", false);
    addExampleLine("xmipp_ctf_estimate_from_micrograph --micrograph micrograph.mrc --mode regions micrograph.pos --sampling_rate 1.4 --voltage 200 --spherical_aberration 2.5 --defocusU -15000");
    addExampleLine("Estimate a CTF per particle", false);
    addExampleLine("xmipp_ctf_estimate_from_micrograph --micrograph micrograph.mrc --mode particles micrograph.pos --sampling_rate 1.4 --voltage 200 --spherical_aberration 2.5 --defocusU -15000");
}

/* Construct piece smoother =============================================== */
template<typename T>
void ProgCTFEstimateFromMicrograph::constructPieceSmoother(
        const MultidimArray<T> &piece,
        MultidimArray<T> &pieceSmoother)
{
    // Attenuate borders to avoid discontinuities
    pieceSmoother.resizeNoCopy(piece);
    pieceSmoother.initConstant(1);
    pieceSmoother.setXmippOrigin();
    T iHalfsize = (T)2 / YSIZE(pieceSmoother);
    const T alpha = 0.025;
    const T alpha1 = 1 - alpha;
    const T ialpha = 1.0 / alpha;

    auto computeMaskValue = [&] (T fraction) {
        return 0.5 * (1 + cos(PI * ((fraction - 1) * ialpha + 1)));
    };

    for (int i = STARTINGY(pieceSmoother);
            i <= FINISHINGY(pieceSmoother);
            i++) {
        T iFraction = fabs(i * iHalfsize);
        if (iFraction > alpha1)
        {
            T maskValue = computeMaskValue(iFraction);
            for (int j = STARTINGX(pieceSmoother);
                    j <= FINISHINGX(pieceSmoother); j++) {
                A2D_ELEM(pieceSmoother,i,j) *= maskValue;
            }
        }
    }

    for (int j = STARTINGX(pieceSmoother);
            j <= FINISHINGX(pieceSmoother);
            j++) {
        T jFraction = fabs(j * iHalfsize);
        if (jFraction > alpha1)
        {
            T maskValue = computeMaskValue(jFraction);
            for (int i = STARTINGY(pieceSmoother);
                    i <= FINISHINGY(pieceSmoother); i++) {
                A2D_ELEM(pieceSmoother,i,j) *= maskValue;
            }
        }
    }
}

/* Compute PSD by piece averaging ========================================== */
//#define DEBUG
void ProgCTFEstimateFromMicrograph::PSD_piece_by_averaging(
    MultidimArray<double> &piece, MultidimArray<double> &psd)
{
    int small_Ydim = 2 * YSIZE(piece) / Nsubpiece;
    int small_Xdim = 2 * XSIZE(piece) / Nsubpiece;
    MultidimArray<double> small_piece(small_Ydim, small_Xdim);

    int Xstep = (XSIZE(piece) - small_Xdim) / (Nsubpiece - 1);
    int Ystep = (YSIZE(piece) - small_Ydim) / (Nsubpiece - 1);
    psd.initZeros(small_piece);
#ifdef DEBUG

    Image<double> save;
    save()=piece;
    save.write("PPPpiece.xmp");
#endif

    MultidimArray<std::complex<double> > Periodogram;
    MultidimArray<double> small_psd;

    // Attenuate borders to avoid discontinuities
    MultidimArray<double> pieceSmoother;
    constructPieceSmoother(piece, pieceSmoother);

    for (int ii = 0; ii < Nsubpiece; ii++)
        for (int jj = 0; jj < Nsubpiece; jj++)
        {
            // Take the corresponding small piece from the piece
            int i0 = ii * Xstep;
            int j0 = jj * Ystep;

            int i, j, ib, jb;
            for (i = 0, ib = i0; i < small_Ydim; i++, ib++)
                for (j = 0, jb = j0; j < small_Xdim; j++, jb++)
                    DIRECT_A2D_ELEM(small_piece, i, j)=
                        DIRECT_A2D_ELEM(piece, ib, jb);
            normalize_ramp(piece);
            piece *= pieceSmoother;

#ifdef DEBUG

            save()=small_piece;
            save.write("PPPsmall_piece.xmp");
#endif

            // Compute the PSD of the small piece
            small_psd.initZeros(small_piece);
            if (PSDEstimator_mode == ARMA)
            {
                CausalARMA(small_piece, ARMA_prm);
                ARMAFilter(small_piece, small_psd, ARMA_prm);
            }
            else
            {
                FourierTransform(small_piece, Periodogram);
                FFT_magnitude(Periodogram, small_psd);
                small_psd *= small_psd;
                small_psd *= small_Ydim * small_Xdim;
            }

#ifdef DEBUG
            save()=small_psd;
            save.write("PPPsmall_psd.xmp");
#endif

            // Add to the average
            psd += small_psd;
        }

    // Compute the average of all the small pieces and enlarge
    psd *= 1.0 / (Nsubpiece * Nsubpiece);

#ifdef DEBUG

    save()=psd;
    save.write("PPPpsd1.xmp");
#endif

    CenterFFT(psd, true);
    selfScaleToSize(xmipp_transformation::BSPLINE3, psd, YSIZE(piece), XSIZE(piece));
    CenterFFT(psd, false);
    psd.threshold("below", 0, 0);

#ifdef DEBUG

    save()=psd;
    save.write("PPPpsd2.xmp");
    std::cout << "Press any key\n";
    char c;
    std::cin >> c;
#endif
}
#undef DEBUG

/* Main ==================================================================== */
//#define DEBUG
void ProgCTFEstimateFromMicrograph::run()
{
    // Open input files -----------------------------------------------------
    // Open coordinates
    MetaDataVec posFile;
    if (fn_pos != "")
        posFile.read(fn_pos);
    auto iterPosFile = posFile.ids().begin();

    // Open the micrograph --------------------------------------------------
    Image<double> M_in;
    size_t Ndim, Zdim, Ydim , Xdim; // Micrograph dimensions

    //ImageInfo imgInfo;
    //getImageInfo(fn_micrograph, imgInfo);
    //imgInfo.adim.ndim

    M_in.read(fn_micrograph,HEADER);
    M_in.getDimensions(Xdim, Ydim, Zdim, Ndim);

    // Compute the number of divisions --------------------------------------
    int div_Number = 0;
    int div_NumberX=1, div_NumberY=1;
    if (psd_mode == OnePerParticle)
        div_Number = posFile.size();
    else if (psd_mode == OnePerMicrograph)
    {
        div_NumberX = CEIL((double)Xdim / (pieceDim *(1-overlap))) - 1;
        div_NumberY = CEIL((double)Ydim / (pieceDim *(1-overlap))) - 1;
        div_Number = div_NumberX * div_NumberY;
    }
    else if (psd_mode == OnePerRegion)
    {
        div_NumberX = CEIL((double)Xdim / pieceDim);
        div_NumberY = CEIL((double)Ydim / pieceDim);
        if (div_NumberX<=2*skipBorders || div_NumberY<=2*skipBorders)
            REPORT_ERROR(ERR_ARG_INCORRECT,formatString("The micrograph is not big enough to skip %d pieces on each side",skipBorders));
        div_Number = div_NumberX * div_NumberY;
    }

    // Process each piece ---------------------------------------------------
    Image<double> psd_avg, psd_std, psd, psd2;
    MultidimArray<std::complex<double> > Periodogram;
    MultidimArray<double> piece(pieceDim, pieceDim);
    psd().resizeNoCopy(piece);
    MultidimArray<double> &mpsd = psd();
    MultidimArray<double> &mpsd2 = psd2();
    PCAMahalanobisAnalyzer pcaAnalyzer;
    MultidimArray<int> PCAmask;
    MultidimArray<float> PCAv;
    double pieceDim2 = pieceDim * pieceDim;

    //Multidimensional data variables to store the defocus obtained locally for plane fitting
    MultidimArray<double> defocusPlanefittingU(div_NumberX-2*skipBorders, div_NumberY-2*skipBorders);
    MultidimArray<double> defocusPlanefittingV(defocusPlanefittingU);
    MultidimArray<double> Xm(defocusPlanefittingU);
    MultidimArray<double> Ym(defocusPlanefittingU);

    // Attenuate borders to avoid discontinuities
    MultidimArray<double> pieceSmoother;
    constructPieceSmoother(piece, pieceSmoother);
    if (verbose)
        std::cerr << "Computing models of each piece ...\n";

    // Prepare these filenames in case they are needed
    FileName fn_psd;
    if (psd_mode == OnePerMicrograph)
        fn_psd = fn_root + ".psd";
    else
        fn_psd = fn_root + ".psdstk";
    if (fileExists(fn_psd))
        fn_psd.deleteFile();
    if (fileExists(fn_root+".ctfparam"))
        FileName(fn_root+".ctfparam").deleteFile();

    if (verbose)
        init_progress_bar(div_Number);
    int N = 1; // Index of current piece
    size_t piecei = 0, piecej = 0; // top-left corner of the current piece
    FourierTransformer transformer;
    int actualDiv_Number = 0;

    for (size_t nIm = 1; nIm <= Ndim; nIm++)
    {
        M_in.read(fn_micrograph,DATA,nIm);
        std::cout << "Micrograph number: " << nIm << std::endl;

        while (N <= div_Number)
        {
            bool skip = false;
            int blocki=0, blockj=0;
            // Compute the top-left corner of the piece ..........................
            if (psd_mode == OnePerParticle)
            {
                // Read position of the particle
                posFile.getValue(MDL_X, piecej, *iterPosFile);
                posFile.getValue(MDL_Y, piecei, *iterPosFile);

                // j,i are the selfWindow center, we need the top-left corner
                piecej -= (int) (pieceDim / 2);
                piecei -= (int) (pieceDim / 2);
                if (piecei < 0)
                    piecei = 0;
                if (piecej < 0)
                    piecej = 0;
            }
            else
            {
                int step = pieceDim;
                if (psd_mode == OnePerMicrograph)
                    step = (int) ((1 - overlap) * step);
                blocki = (N - 1) / div_NumberX;
                blockj = (N - 1) % div_NumberX;
                if (blocki < skipBorders || blockj < skipBorders
                        || blocki > (div_NumberY - skipBorders - 1)
                        || blockj > (div_NumberX - skipBorders - 1))
                    skip = true;
                piecei = blocki * step;
                piecej = blockj * step;
            }

            // test if the full piece is inside the micrograph
            if (piecei + pieceDim > Ydim)
                piecei = Ydim - pieceDim;
            if (piecej + pieceDim > Xdim)
                piecej = Xdim - pieceDim;

            if (!skip)
            {
                // Extract micrograph piece ..........................................
//              M_in().window(piece, 0, 0, piecei, piecej, 0, 0, piecei + YSIZE(piece) - 1,
//                      piecej + XSIZE(piece) - 1);
                window2D( M_in(), piece, piecei, piecej, piecei + YSIZE(piece) - 1, piecej + XSIZE(piece) - 1);
                piece.statisticsAdjust(0., 1.);
                normalize_ramp(piece);
                piece *= pieceSmoother;

                // Estimate the power spectrum .......................................
                if (Nsubpiece == 1)
                    if (PSDEstimator_mode == ARMA)
                    {
                        CausalARMA(piece, ARMA_prm);
                        ARMAFilter(piece, mpsd, ARMA_prm);
                    }
                    else
                    {
                        transformer.completeFourierTransform(piece, Periodogram);
                        FFT_magnitude(Periodogram, mpsd);
                        FOR_ALL_DIRECT_ELEMENTS_IN_MULTIDIMARRAY(mpsd)
                        DIRECT_MULTIDIM_ELEM(mpsd,n)*=DIRECT_MULTIDIM_ELEM(mpsd,n)*pieceDim2;
                    }
                else
                    PSD_piece_by_averaging(piece, mpsd);
                mpsd2.resizeNoCopy(mpsd);
                FOR_ALL_DIRECT_ELEMENTS_IN_MULTIDIMARRAY(mpsd2)
                {
                    double psdval = DIRECT_MULTIDIM_ELEM(mpsd,n);
                    DIRECT_MULTIDIM_ELEM(mpsd2,n)=psdval*psdval;
                }

                // Perform averaging if applicable ...................................
                if (psd_mode == OnePerMicrograph)
                {
                    actualDiv_Number += 1;
                    // Compute average and standard deviation
                    if (XSIZE(psd_avg()) != XSIZE(mpsd))
                    {
                        psd_avg() = mpsd;
                        psd_std() = psd2();
                    }
                    else
                    {
                        psd_avg() += mpsd;
                        psd_std() += psd2();
                    }

                    // Keep psd for the PCA
                    if (estimate_ctf)
                    {
                        if (XSIZE(PCAmask) == 0)
                        {
                            PCAmask.initZeros(mpsd);
                            Matrix1D<int> idx(2);  // Indexes for Fourier plane
                            Matrix1D<double> freq(2); // Frequencies for Fourier plane
                            size_t PCAdim = 0;
                            FOR_ALL_ELEMENTS_IN_ARRAY2D(PCAmask)
                            {
                                VECTOR_R2(idx, j, i);
                                FFT_idx2digfreq(mpsd, idx, freq);
                                double w = freq.module();
                                if (w > 0.05 && w < 0.4)
                                {
                                    A2D_ELEM(PCAmask,i,j)=1;
                                    ++PCAdim;
                                }
                            }
                            PCAv.initZeros(PCAdim);
                        }

                        size_t ii = -1;
                        FOR_ALL_DIRECT_ELEMENTS_IN_MULTIDIMARRAY(PCAmask)
                        if (DIRECT_MULTIDIM_ELEM(PCAmask,n))
                            A1D_ELEM(PCAv,++ii)=(float)DIRECT_MULTIDIM_ELEM(mpsd,n);
                        pcaAnalyzer.addVector(PCAv);
                    }
                }

                // Compute the theoretical model if not averaging ....................
                if (psd_mode != OnePerMicrograph)
                {

                    if (bootstrapN != -1)
                        REPORT_ERROR(ERR_VALUE_INCORRECT,
                                "Bootstrapping is only available for micrograph averages");

                    FileName fn_psd_piece;
                    fn_psd_piece.compose(N, fn_psd);
                    psd.write(fn_psd_piece);
                    if (psd_mode == OnePerParticle)
                        posFile.setValue(MDL_PSD, fn_psd_piece, *iterPosFile);
                    if (estimate_ctf)
                    {
                        // Estimate the CTF parameters of this piece
                        prmEstimateCTFFromPSD.fn_psd = fn_psd_piece;
                        CTFDescription ctfmodel;

                        ctfmodel.isLocalCTF = true;
                        ctfmodel.x0 = piecej;
                        ctfmodel.xF = (piecej + pieceDim-1);
                        ctfmodel.y0 = piecei;
                        ctfmodel.yF = (piecei + pieceDim-1);
                        ROUT_Adjust_CTF(prmEstimateCTFFromPSD, ctfmodel, false);

                        int idxi=blocki-skipBorders;
                        int idxj=blockj-skipBorders;
                        A2D_ELEM(defocusPlanefittingU,idxi,idxj)=ctfmodel.DeltafU;
                        A2D_ELEM(defocusPlanefittingV,idxi,idxj)=ctfmodel.DeltafV;

                        A2D_ELEM(Xm,idxi,idxj)=(piecei+pieceDim/2)*ctfmodel.Tm;
                        A2D_ELEM(Ym,idxi,idxj)=(piecej+pieceDim/2)*ctfmodel.Tm;

                        //1D//
                        prmEstimateCTFFromPSDFast.fn_psd = fn_psd_piece; //Nuevo
                        CTFDescription1D ctf1Dmodel; //Nuevo
                        ctf1Dmodel.isLocalCTF = true;
                        ctf1Dmodel.x0 = piecej;
                        ctf1Dmodel.xF = (piecej + pieceDim-1);
                        ctf1Dmodel.y0 = piecei;
                        ctf1Dmodel.yF = (piecei + pieceDim-1);
                        ROUT_Adjust_CTFFast(prmEstimateCTFFromPSDFast, ctf1Dmodel, false); //Nuevo

                        A2D_ELEM(defocusPlanefittingU,idxi,idxj)=ctf1Dmodel.Defocus;

                        A2D_ELEM(Xm,idxi,idxj)=(piecei+pieceDim/2)*ctf1Dmodel.Tm;
                        A2D_ELEM(Ym,idxi,idxj)=(piecej+pieceDim/2)*ctf1Dmodel.Tm;

                        if (psd_mode == OnePerParticle)
                            posFile.setValue(MDL_CTF_MODEL,
                                    fn_psd_piece.withoutExtension() + ".ctfparam",
                                    *iterPosFile);
                    }
                }
            }
            // Increment the division counter
            ++N;
            if (verbose)
                progress_bar(N);
            if (psd_mode == OnePerParticle)
                ++iterPosFile;
        }

        init_progress_bar(div_Number);
        N = 1;
    }
    if (verbose)
        progress_bar(div_Number);

    // If averaging, compute the CTF model ----------------------------------  //Program execution
    if (psd_mode == OnePerMicrograph)
    {
        // Compute the avg and stddev of the local PSDs
        const MultidimArray<double> &mpsd_std = psd_std();
        const MultidimArray<double> &mpsd_avg = psd_avg();
        double idiv_Number = 1.0 / actualDiv_Number;
        FOR_ALL_DIRECT_ELEMENTS_IN_MULTIDIMARRAY(mpsd_avg)
        {
            DIRECT_MULTIDIM_ELEM(mpsd_avg,n)*=idiv_Number;
            DIRECT_MULTIDIM_ELEM(mpsd_std,n)*=idiv_Number;
            DIRECT_MULTIDIM_ELEM(mpsd_std,n)-=DIRECT_MULTIDIM_ELEM(mpsd_avg,n)*
                                              DIRECT_MULTIDIM_ELEM(mpsd_avg,n);
            if (DIRECT_MULTIDIM_ELEM(mpsd_std,n)<0)
                DIRECT_MULTIDIM_ELEM(mpsd_std,n)=0;
            else
                DIRECT_MULTIDIM_ELEM(mpsd_std,n)=sqrt(DIRECT_MULTIDIM_ELEM(mpsd_std,n));
        }
        psd_avg.write(fn_psd);

        if (estimate_ctf)
        {
            // Estimate the CTF parameters
            std::cerr << "Adjusting CTF model to the PSD ...\n";
            prmEstimateCTFFromPSD.fn_psd = fn_psd;
            prmEstimateCTFFromPSDFast.fn_psd = fn_psd;
            CTFDescription ctfmodel;
            CTFDescription1D ctf1Dmodel;

            if (bootstrapN == -1)
            {
                try {
                    // No bootstrapping
                    // Compute the PCA of the local PSDs
                    pcaAnalyzer.standardarizeVariables();
                    // pcaAnalyzer.subtractAvg();
                    pcaAnalyzer.learnPCABasis(1, 10);
                } catch (XmippError &xe)
                {
                    if (xe.__errno==ERR_NUMERICAL)
                        REPORT_ERROR(ERR_NUMERICAL,"There is no variance in the PSD, check that the micrograph is not constant");
                    else
                        throw xe;
                }

#ifdef DEBUG

                Image<double> save;
                save().initZeros(psd());
                int ii=-1;
                FOR_ALL_ELEMENTS_IN_ARRAY2D(PCAmask)
                if (PCAmask(i,j))
                    save(i,j)=pcaAnalyzer.PCAbasis[0](++ii);
                save.write("PPPbasis.xmp");
#endif

                Matrix2D<double> CtY;
                pcaAnalyzer.projectOnPCABasis(CtY);
                Matrix1D<double> p;
                CtY.toVector(p);
                double pavg = p.sum(true);
                double pstd = p.sum2() / VEC_XSIZE(p) - pavg * pavg;
                pstd = (pstd < 0) ? 0 : sqrt(pstd);

                std::string psign;
                FOR_ALL_ELEMENTS_IN_MATRIX1D(p)
                if (p(i) < 0)
                    psign += "-";
                else
                    psign += "+";
                double zrandomness = checkRandomness(psign);

                if (!acceleration1D)
                {
                    ctfmodel.isLocalCTF = false;
                    ctfmodel.x0 = 0;
                    ctfmodel.xF = (Xdim-1);
                    ctfmodel.y0 = 0;
                    ctfmodel.yF = (Ydim-1);
                    ROUT_Adjust_CTF(prmEstimateCTFFromPSD,ctfmodel, false);
                }
                else
                {
                    ctf1Dmodel.isLocalCTF = false;
                    ctf1Dmodel.x0 = 0;
                    ctf1Dmodel.xF = (Xdim-1);
                    ctf1Dmodel.y0 = 0;
                    ctf1Dmodel.yF = (Ydim-1);
                    ROUT_Adjust_CTFFast(prmEstimateCTFFromPSDFast,ctf1Dmodel, false);
                }
                // Evaluate PSD variance and write into the CTF
                double stdQ = 0;
                FOR_ALL_ELEMENTS_IN_ARRAY2D(mpsd_std)
                stdQ += A2D_ELEM(mpsd_std,i,j)/A2D_ELEM(mpsd_avg,i,j);
                stdQ /= MULTIDIM_SIZE(psd_std());

                MetaDataVec MD;
                MD.read(fn_psd.withoutExtension() + ".ctfparam");
                size_t id = MD.firstRowId();
                MD.setValue(MDL_CTF_CRIT_PSDVARIANCE, stdQ, id);
                MD.setValue(MDL_CTF_CRIT_PSDPCA1VARIANCE, pstd, id);
                MD.setValue(MDL_CTF_CRIT_PSDPCARUNSTEST, zrandomness, id);
                MD.write((String)"fullMicrograph@"+fn_psd.withoutExtension() + ".ctfparam");
            }
            else
            {
                // If bootstrapping
                prmEstimateCTFFromPSD.bootstrap = true;
                prmEstimateCTFFromPSD.show_optimization = true;
                if (!acceleration1D)
                {
                MultidimArray<double> CTFs(bootstrapN, 32);
                FileName fnBase = fn_psd.withoutExtension();
                std::cerr << "Computing bootstrap ...\n";
                init_progress_bar(bootstrapN);
                for (int n = 0; n < bootstrapN; n++)
                {
                    CTFs(n, 31) = ROUT_Adjust_CTF(prmEstimateCTFFromPSD,
                                                  ctfmodel, false);

                    CTFs(n, 0) = ctfmodel.Tm;
                    CTFs(n, 1) = ctfmodel.kV;
                    CTFs(n, 2) = ctfmodel.DeltafU;
                    CTFs(n, 3) = ctfmodel.DeltafV;
                    CTFs(n, 4) = ctfmodel.azimuthal_angle;
                    CTFs(n, 5) = ctfmodel.Cs;
                    CTFs(n, 6) = ctfmodel.Ca;
                    CTFs(n, 7) = ctfmodel.espr;
                    CTFs(n, 8) = ctfmodel.ispr;
                    CTFs(n, 9) = ctfmodel.alpha;
                    CTFs(n, 10) = ctfmodel.DeltaF;
                    CTFs(n, 11) = ctfmodel.DeltaR;
                    CTFs(n, 12) = ctfmodel.Q0;
                    CTFs(n, 13) = ctfmodel.K;
                    CTFs(n, 14) = ctfmodel.gaussian_K;
                    CTFs(n, 15) = ctfmodel.sigmaU;
                    CTFs(n, 16) = ctfmodel.sigmaV;
                    CTFs(n, 17) = ctfmodel.cU;
                    CTFs(n, 18) = ctfmodel.cV;
                    CTFs(n, 19) = ctfmodel.gaussian_angle;
                    CTFs(n, 20) = ctfmodel.sqrt_K;
                    CTFs(n, 21) = ctfmodel.sqU;
                    CTFs(n, 22) = ctfmodel.sqV;
                    CTFs(n, 23) = ctfmodel.sqrt_angle;
                    CTFs(n, 24) = ctfmodel.base_line;
                    CTFs(n, 25) = ctfmodel.gaussian_K2;
                    CTFs(n, 26) = ctfmodel.sigmaU2;
                    CTFs(n, 27) = ctfmodel.sigmaV2;
                    CTFs(n, 28) = ctfmodel.cU2;
                    CTFs(n, 29) = ctfmodel.cV2;
                    CTFs(n, 30) = ctfmodel.gaussian_angle2;


                    std::string command = (std::string) "mv -i " + fnBase
                                          + ".ctfparam " + fnBase + "_bootstrap_"
                                          + integerToString(n, 4) + ".ctfparam";
                    if (system(command.c_str())==-1)
                        REPORT_ERROR(ERR_UNCLASSIFIED,"Cannot open shell");
                    command = (std::string) "mv -i " + fnBase
                              + ".ctfmodel_quadrant " + fnBase + "_bootstrap_"
                              + integerToString(n, 4) + ".ctfmodel_quadrant";
                    if (system(command.c_str())==-1)
                        REPORT_ERROR(ERR_UNCLASSIFIED,"Cannot open shell");
                    command = (std::string) "mv -i " + fnBase
                              + ".ctfmodel_halfplane " + fnBase + "_bootstrap_"
                              + integerToString(n, 4) + ".ctfmodel_halfplane";
                    if (system(command.c_str())==-1)
                        REPORT_ERROR(ERR_UNCLASSIFIED,"Cannot open shell");

                    progress_bar(n);
                }
                progress_bar(bootstrapN);

                }
                else
                {
                    MultidimArray<double> CTFs(bootstrapN, 21);
                    prmEstimateCTFFromPSDFast.bootstrap = true;
                    prmEstimateCTFFromPSDFast.show_optimization = true;
                    FileName fnBase = fn_psd.withoutExtension();
                    std::cerr << "Computing bootstrap ...\n";
                    init_progress_bar(bootstrapN);
                    for (int n = 0; n < bootstrapN; n++)
                    {
                        CTFs(n, 21) = ROUT_Adjust_CTFFast(prmEstimateCTFFromPSDFast,ctf1Dmodel, false);

                        CTFs(n, 0) = ctf1Dmodel.Tm;
                        CTFs(n, 1) = ctf1Dmodel.kV;
                        CTFs(n, 2) = ctf1Dmodel.Defocus;
                        CTFs(n, 3) = ctf1Dmodel.Cs;
                        CTFs(n, 4) = ctf1Dmodel.Ca;
                        CTFs(n, 5) = ctf1Dmodel.espr;
                        CTFs(n, 6) = ctf1Dmodel.ispr;
                        CTFs(n, 7) = ctf1Dmodel.alpha;
                        CTFs(n, 8) = ctf1Dmodel.DeltaF;
                        CTFs(n, 9) = ctf1Dmodel.DeltaR;
                        CTFs(n, 10) = ctf1Dmodel.Q0;
                        CTFs(n, 11) = ctf1Dmodel.K;
                        CTFs(n, 12) = ctf1Dmodel.gaussian_K;
                        CTFs(n, 13) = ctf1Dmodel.sigma1;
                        CTFs(n, 14) = ctf1Dmodel.Gc1;
                        CTFs(n, 15) = ctf1Dmodel.sqrt_K;
                        CTFs(n, 16) = ctf1Dmodel.sq;
                        CTFs(n, 17) = ctf1Dmodel.base_line;
                        CTFs(n, 18) = ctf1Dmodel.gaussian_K2;
                        CTFs(n, 19) = ctf1Dmodel.sigma2;
                        CTFs(n, 20) = ctf1Dmodel.Gc2;


                        std::string command = (std::string) "mv -i " + fnBase
                                              + ".ctfparam " + fnBase + "_bootstrap_"
                                              + integerToString(n, 4) + ".ctfparam";
                        if (system(command.c_str())==-1)
                            REPORT_ERROR(ERR_UNCLASSIFIED,"Cannot open shell");
                        command = (std::string) "mv -i " + fnBase
                                  + ".ctf1Dmodel_quadrant " + fnBase + "_bootstrap_"
                                  + integerToString(n, 4) + ".ctf1Dmodel_quadrant";
                        if (system(command.c_str())==-1)
                            REPORT_ERROR(ERR_UNCLASSIFIED,"Cannot open shell");
                        command = (std::string) "mv -i " + fnBase
                                  + ".ctf1Dmodel_halfplane " + fnBase + "_bootstrap_"
                                  + integerToString(n, 4) + ".ctf1Dmodel_halfplane";
                        if (system(command.c_str())==-1)
                            REPORT_ERROR(ERR_UNCLASSIFIED,"Cannot open shell");

                        progress_bar(n);
                    }
                    progress_bar(bootstrapN);
                    }
                }


            }
        }

    // Assign a CTF to each particle ----------------------------------------
    if (psd_mode == OnePerRegion && estimate_ctf)
    {
        double pU0 = 0, pU1 = 0, pU2 = 0;
        planeFit(defocusPlanefittingU, Xm, Ym, pU0, pU1, pU2);
        double pV0 = 0, pV1 = 0, pV2 = 0;
        planeFit(defocusPlanefittingV, Xm, Ym, pV0, pV1, pV2);

        MetaDataVec MDctf;
        MDctf.read(fn_root+".ctfparam");
        double Tm, downsampling;
        size_t id=MDctf.firstRowId();
        MDctf.getValue(MDL_CTF_SAMPLING_RATE,Tm,id);
        MDctf.getValue(MDL_CTF_DOWNSAMPLE_PERFORMED,downsampling,id);

        MetaDataVec MD;
        MD.setColumnFormat(false);
        id = MD.addObject();
        MD.setValue(MDL_CTF_DEFOCUS_PLANEUA, pU1, id);
        MD.setValue(MDL_CTF_DEFOCUS_PLANEUB, pU2, id);
        MD.setValue(MDL_CTF_DEFOCUS_PLANEUC, pU0, id);
        MD.setValue(MDL_CTF_DEFOCUS_PLANEVA, pV1, id);
        MD.setValue(MDL_CTF_DEFOCUS_PLANEVB, pV2, id);
        MD.setValue(MDL_CTF_DEFOCUS_PLANEVC, pV0, id);
        MD.setValue(MDL_CTF_X0, 0., id);
        MD.setValue(MDL_CTF_Y0, 0., id);
        MD.setValue(MDL_CTF_XF, (Xdim-1)*Tm*downsampling, id);
        MD.setValue(MDL_CTF_YF, (Ydim-1)*Tm*downsampling, id);
        MD.write((String)"fullMicrograph@"+fn_root+".ctfparam", MD_APPEND);

        if (fn_pos != "")
        {
            FileName fn_img, fn_psd_piece, fn_ctfparam_piece;
            int Y, X;
            for (size_t objId : posFile.ids())
            {
                posFile.getValue(MDL_IMAGE, fn_img, objId);
                posFile.getValue(MDL_X, X, objId);
                posFile.getValue(MDL_Y, Y, objId);
                auto idx_X = (int)floor((double) X / pieceDim);
                auto idx_Y = (int)floor((double) Y / pieceDim);
                int N = idx_Y * div_NumberX + idx_X + 1;

                fn_psd_piece.compose(N, fn_psd);
                fn_ctfparam_piece = fn_psd_piece.withoutExtension()
                                    + ".ctfparam";
                posFile.setValue(MDL_PSD, fn_psd_piece, objId);
                posFile.setValue(MDL_CTF_MODEL, fn_ctfparam_piece, objId);
            }
        }
    }
    posFile.write(fn_pos);
}

/* Fast estimate of PSD --------------------------------------------------- */
class ThreadFastEstimateEnhancedPSDParams
{
public:
    ImageGeneric *I;
    MultidimArray<double> *PSD, *pieceSmoother;
    MultidimArray<int> *pieceMask;
    Mutex *mutex;
    int Nprocessed;
};

void threadFastEstimateEnhancedPSD(ThreadArgument &thArg)
{
    auto *args = (ThreadFastEstimateEnhancedPSDParams*) thArg.workClass;
    int Nthreads = thArg.getNumberOfThreads();
    int id = thArg.thread_id;
    ImageGeneric &I = *(args->I);
    const MultidimArrayGeneric& mI = I();
    size_t IXdim, IYdim, IZdim;
    I.getDimensions(IXdim, IYdim, IZdim);
    MultidimArray<double> &pieceSmoother = *(args->pieceSmoother);
    MultidimArray<int> &pieceMask = *(args->pieceMask);
    MultidimArray<double> localPSD, piece;
    MultidimArray<std::complex<double> > Periodogram;
    piece.initZeros(pieceMask);
    localPSD.initZeros(*(args->PSD));

    FourierTransformer transformer;
    transformer.setReal(piece);

    int pieceNumber = 0;
    int Nprocessed = 0;
    double pieceDim2 = XSIZE(piece) * XSIZE(piece);
    for (size_t i = 0; i < (IYdim - YSIZE(piece)); i+=YSIZE(piece))
        for (size_t j = 0; j < (IXdim - XSIZE(piece)); j+=XSIZE(piece), pieceNumber++)
        {
            if ((pieceNumber + 1) % Nthreads != id)
                continue;
            Nprocessed++;

            // Extract micrograph piece ..........................................
            for (size_t k = 0; k < YSIZE(piece); k++)
                for (size_t l = 0; l < XSIZE(piece); l++)
                    DIRECT_A2D_ELEM(piece, k, l)= mI(i+k, j+l);
            piece.statisticsAdjust(0., 1.);
            normalize_ramp(piece, &pieceMask);
            piece *= pieceSmoother;

            // Estimate the power spectrum .......................................
            transformer.FourierTransform();
            transformer.getCompleteFourier(Periodogram);
            FOR_ALL_DIRECT_ELEMENTS_IN_MULTIDIMARRAY(localPSD)
            {
                const auto *ptr = (double*) &DIRECT_MULTIDIM_ELEM(Periodogram, n);
                double re=*ptr;
                double im=*(ptr+1);
                double magnitude2=re*re+im*im;
                DIRECT_MULTIDIM_ELEM(localPSD,n)+=magnitude2*pieceDim2;
            }
        }

    // Gather results
    args->mutex->lock();
    args->Nprocessed += Nprocessed;
    *(args->PSD) += localPSD;
    args->mutex->unlock();
}

void fastEstimateEnhancedPSD(const FileName &fnMicrograph, double downsampling,
                             MultidimArray<double> &enhancedPSD, int numberOfThreads)
{
    size_t Xdim, Ydim, Zdim, Ndim;
    getImageSizeFromFilename(fnMicrograph, Xdim, Ydim, Zdim, Ndim);
    int minSize = 2 * (std::max(Xdim, Ydim) / 10);
    minSize = (int)std::min((double) std::min(Xdim, Ydim), NEXT_POWER_OF_2(minSize));
    minSize = std::min(1024, minSize);

    /*
     ProgCTFEstimateFromMicrograph prog1;
     prog1.fn_micrograph=fnMicrograph;
     prog1.fn_root=fnMicrograph.withoutExtension()+"_tmp";
     prog1.pieceDim=(int)(minSize*downsampling);
     prog1.PSDEstimator_mode=ProgCTFEstimateFromMicrograph::Periodogram;
     prog1.Nsubpiece=1;
     prog1.psd_mode=ProgCTFEstimateFromMicrograph::OnePerMicrograph;
     prog1.estimate_ctf=false;
     prog1.bootstrapN=-1;
     prog1.verbose=1;
     prog1.overlap=0;
     prog1.run();
     */
    // Prepare auxiliary variables
    ImageGeneric I;
    I.read(fnMicrograph);

    MultidimArray<double> PSD;
    PSD.initZeros(minSize, minSize);

    MultidimArray<int> pieceMask;
    pieceMask.resizeNoCopy(PSD);
    pieceMask.initConstant(1);

    MultidimArray<double> pieceSmoother;
    ProgCTFEstimateFromMicrograph::constructPieceSmoother(PSD, pieceSmoother);

    // Prepare thread arguments
    Mutex mutex;
    ThreadFastEstimateEnhancedPSDParams args;
    args.I = &I;
    args.PSD = &PSD;
    args.pieceMask = &pieceMask;
    args.pieceSmoother = &pieceSmoother;
    args.Nprocessed = 0;
    args.mutex = &mutex;
    auto thMgr = std::unique_ptr<ThreadManager>(std::make_unique<ThreadManager>(numberOfThreads, &args));
    thMgr->run(threadFastEstimateEnhancedPSD);
    if (args.Nprocessed != 0)
        *(args.PSD) /= args.Nprocessed;

    ProgCTFEnhancePSD prog2;
    prog2.filter_w1 = 0.02;
    prog2.filter_w2 = 0.2;
    prog2.decay_width = 0.02;
    prog2.mask_w1 = 0.005;
    prog2.mask_w2 = 0.5;

    prog2.applyFilter(*(args.PSD));
    enhancedPSD = *(args.PSD);

    auto downXdim = (int) ((double)XSIZE(enhancedPSD) / downsampling);
    int firstIndex = FIRST_XMIPP_INDEX(downXdim);
    int lastIndex = LAST_XMIPP_INDEX(downXdim);
    enhancedPSD.setXmippOrigin();
    enhancedPSD.selfWindow(firstIndex, firstIndex, lastIndex, lastIndex);
}

// explicit instantiation
template void ProgCTFEstimateFromMicrograph::constructPieceSmoother<float>(MultidimArray<float> const&, MultidimArray<float>&);
template void ProgCTFEstimateFromMicrograph::constructPieceSmoother<double>(MultidimArray<double> const&, MultidimArray<double>&);