Exemple #1
0
        public virtual void ProcessCTF(MapHeader originalHeader, Image originalStack, bool doastigmatism, decimal scaleFactor)
        {
            if (!Directory.Exists(PowerSpectrumDir))
                Directory.CreateDirectory(PowerSpectrumDir);

            //CTF = new CTF();
            PS1D = null;
            _SimulatedBackground = null;
            _SimulatedScale = new Cubic1D(new[] { new float2(0, 1), new float2(1, 1) });

            #region Dimensions and grids

            int NFrames = originalHeader.Dimensions.Z;
            int2 DimsImage = new int2(originalHeader.Dimensions);
            int2 DimsRegion = new int2(MainWindow.Options.CTFWindow, MainWindow.Options.CTFWindow);

            float OverlapFraction = 0.5f;
            int2 DimsPositionGrid;
            int3[] PositionGrid = Helper.GetEqualGridSpacing(DimsImage, new int2(DimsRegion.X / 1, DimsRegion.Y / 1), OverlapFraction, out DimsPositionGrid);
            int NPositions = (int)DimsPositionGrid.Elements();

            int CTFGridX = Math.Min(DimsPositionGrid.X, MainWindow.Options.GridCTFX);
            int CTFGridY = Math.Min(DimsPositionGrid.Y, MainWindow.Options.GridCTFY);
            int CTFGridZ = Math.Min(NFrames, MainWindow.Options.GridCTFZ);
            GridCTF = new CubicGrid(new int3(CTFGridX, CTFGridY, CTFGridZ));
            GridCTFPhase = new CubicGrid(new int3(1, 1, CTFGridZ));

            bool CTFSpace = CTFGridX * CTFGridY > 1;
            bool CTFTime = CTFGridZ > 1;
            int3 CTFSpectraGrid = new int3(CTFSpace ? DimsPositionGrid.X : 1,
                                           CTFSpace ? DimsPositionGrid.Y : 1,
                                           CTFTime ? CTFGridZ : 1);

            int MinFreqInclusive = (int)(MainWindow.Options.CTFRangeMin * DimsRegion.X / 2);
            int MaxFreqExclusive = (int)(MainWindow.Options.CTFRangeMax * DimsRegion.X / 2);
            int NFreq = MaxFreqExclusive - MinFreqInclusive;

            float PixelSize = (float)(MainWindow.Options.CTFPixelMin + MainWindow.Options.CTFPixelMax) * 0.5f;
            float PixelDelta = (float)(MainWindow.Options.CTFPixelMax - MainWindow.Options.CTFPixelMin) * 0.5f;
            float PixelAngle = (float)MainWindow.Options.CTFPixelAngle / 180f * (float)Math.PI;

            #endregion

            #region Allocate GPU memory

            Image CTFSpectra = new Image(IntPtr.Zero, new int3(DimsRegion.X, DimsRegion.X, (int)CTFSpectraGrid.Elements()), true);
            Image CTFMean = new Image(IntPtr.Zero, new int3(DimsRegion), true);
            Image CTFCoordsCart = new Image(new int3(DimsRegion), true, true);
            Image CTFCoordsPolarTrimmed = new Image(new int3(NFreq, DimsRegion.X, 1), false, true);

            #endregion

            // Extract movie regions, create individual spectra in Cartesian coordinates and their mean.

            #region Create spectra

            GPU.CreateSpectra(originalStack.GetDevice(Intent.Read),
                              DimsImage,
                              NFrames,
                              PositionGrid,
                              NPositions,
                              DimsRegion,
                              CTFSpectraGrid,
                              CTFSpectra.GetDevice(Intent.Write),
                              CTFMean.GetDevice(Intent.Write));
            originalStack.FreeDevice(); // Won't need it in this method anymore.

            #endregion

            // Populate address arrays for later.

            #region Init addresses

            {
                float2[] CoordsData = new float2[CTFCoordsCart.ElementsSliceComplex];

                Helper.ForEachElementFT(DimsRegion, (x, y, xx, yy, r, a) => CoordsData[y * (DimsRegion.X / 2 + 1) + x] = new float2(r, a));
                CTFCoordsCart.UpdateHostWithComplex(new[] { CoordsData });

                CoordsData = new float2[NFreq * DimsRegion.X];
                Helper.ForEachElement(CTFCoordsPolarTrimmed.DimsSlice, (x, y) =>
                {
                    float Angle = ((float)y / DimsRegion.X + 0.5f) * (float)Math.PI;
                    float Ny = 1f / DimsRegion.X;
                    CoordsData[y * NFreq + x] = new float2((x + MinFreqInclusive) * Ny, Angle);
                });
                CTFCoordsPolarTrimmed.UpdateHostWithComplex(new[] { CoordsData });
            }

            #endregion

            // Retrieve average 1D spectrum from CTFMean (not corrected for astigmatism yet).

            #region Initial 1D spectrum

            {
                Image CTFAverage1D = new Image(IntPtr.Zero, new int3(DimsRegion.X / 2, 1, 1));

                GPU.CTFMakeAverage(CTFMean.GetDevice(Intent.Read),
                                   CTFCoordsCart.GetDevice(Intent.Read),
                                   (uint)CTFMean.ElementsSliceReal,
                                   (uint)DimsRegion.X,
                                   new[] { new CTF().ToStruct() },
                                   new CTF().ToStruct(),
                                   0,
                                   (uint)DimsRegion.X / 2,
                                   null,
                                   1,
                                   CTFAverage1D.GetDevice(Intent.Write));

                //CTFAverage1D.WriteMRC("CTFAverage1D.mrc");

                float[] CTFAverage1DData = CTFAverage1D.GetHost(Intent.Read)[0];
                float2[] ForPS1D = new float2[DimsRegion.X / 2];
                for (int i = 0; i < ForPS1D.Length; i++)
                    ForPS1D[i] = new float2((float)i / DimsRegion.X, (float)Math.Round(CTFAverage1DData[i], 4));
                _PS1D = ForPS1D;

                CTFAverage1D.Dispose();
            }

            #endregion

            #region Background fitting methods

            Action UpdateBackgroundFit = () =>
            {
                float2[] ForPS1D = PS1D.Skip(Math.Max(5, MinFreqInclusive / 2)).ToArray();
                Cubic1D.FitCTF(ForPS1D,
                               v => v.Select(x => CTF.Get1D(x / (float)CTF.PixelSize, true)).ToArray(),
                               CTF.GetZeros(),
                               CTF.GetPeaks(),
                               out _SimulatedBackground,
                               out _SimulatedScale);
            };

            Action<bool> UpdateRotationalAverage = keepbackground =>
            {
                float[] MeanData = CTFMean.GetHost(Intent.Read)[0];

                Image CTFMeanCorrected = new Image(new int3(DimsRegion), true);
                float[] MeanCorrectedData = CTFMeanCorrected.GetHost(Intent.Write)[0];

                // Subtract current background estimate from spectra, populate coords.
                Helper.ForEachElementFT(DimsRegion,
                                        (x, y, xx, yy, r, a) =>
                                        {
                                            int i = y * (DimsRegion.X / 2 + 1) + x;
                                            MeanCorrectedData[i] = MeanData[i] - _SimulatedBackground.Interp(r / DimsRegion.X);
                                        });

                Image CTFAverage1D = new Image(IntPtr.Zero, new int3(DimsRegion.X / 2, 1, 1));

                GPU.CTFMakeAverage(CTFMeanCorrected.GetDevice(Intent.Read),
                                   CTFCoordsCart.GetDevice(Intent.Read),
                                   (uint)CTFMeanCorrected.DimsEffective.ElementsSlice(),
                                   (uint)DimsRegion.X,
                                   new[] { CTF.ToStruct() },
                                   CTF.ToStruct(),
                                   0,
                                   (uint)DimsRegion.X / 2,
                                   null,
                                   1,
                                   CTFAverage1D.GetDevice(Intent.Write));

                //CTFAverage1D.WriteMRC("CTFAverage1D.mrc");

                float[] RotationalAverageData = CTFAverage1D.GetHost(Intent.Read)[0];
                float2[] ForPS1D = new float2[PS1D.Length];
                if (keepbackground)
                    for (int i = 0; i < ForPS1D.Length; i++)
                        ForPS1D[i] = new float2((float)i / DimsRegion.X, RotationalAverageData[i] + _SimulatedBackground.Interp((float)i / DimsRegion.X));
                else
                    for (int i = 0; i < ForPS1D.Length; i++)
                        ForPS1D[i] = new float2((float)i / DimsRegion.X, RotationalAverageData[i]);
                MathHelper.UnNaN(ForPS1D);

                _PS1D = ForPS1D;

                CTFMeanCorrected.Dispose();
                CTFAverage1D.Dispose();
            };

            #endregion

            // Fit background to currently best average (not corrected for astigmatism yet).
            {
                float2[] ForPS1D = PS1D.Skip(MinFreqInclusive).Take(Math.Max(2, NFreq / 2)).ToArray();

                int NumNodes = Math.Max(3, (int)((MainWindow.Options.CTFRangeMax - MainWindow.Options.CTFRangeMin) * 5M));
                _SimulatedBackground = Cubic1D.Fit(ForPS1D, NumNodes); // This won't fit falloff and scale, because approx function is 0

                float[] CurrentBackground = _SimulatedBackground.Interp(PS1D.Select(p => p.X).ToArray()).Skip(MinFreqInclusive).Take(NFreq / 2).ToArray();
                float[] Subtracted1D = new float[ForPS1D.Length];
                for (int i = 0; i < ForPS1D.Length; i++)
                    Subtracted1D[i] = ForPS1D[i].Y - CurrentBackground[i];
                MathHelper.NormalizeInPlace(Subtracted1D);

                float ZMin = (float)MainWindow.Options.CTFZMin;
                float ZMax = (float)MainWindow.Options.CTFZMax;
                float ZStep = (ZMax - ZMin) / 100f;

                float BestZ = 0, BestPhase = 0, BestScore = -999;
                for (float z = ZMin; z <= ZMax + 1e-5f; z += ZStep)
                {
                    for (float p = 0; p <= (MainWindow.Options.CTFDoPhase ? 1f : 0f); p += 0.01f)
                    {
                        CTF CurrentParams = new CTF
                        {
                            PixelSize = (MainWindow.Options.CTFPixelMin + MainWindow.Options.CTFPixelMax) * 0.5M,

                            Defocus = (decimal)z,
                            PhaseShift = (decimal)p,

                            Cs = MainWindow.Options.CTFCs,
                            Voltage = MainWindow.Options.CTFVoltage,
                            Amplitude = MainWindow.Options.CTFAmplitude
                        };
                        float[] SimulatedCTF = CurrentParams.Get1D(PS1D.Length, true).Skip(MinFreqInclusive).Take(Math.Max(2, NFreq / 2)).ToArray();
                        MathHelper.NormalizeInPlace(SimulatedCTF);
                        float Score = MathHelper.CrossCorrelate(Subtracted1D, SimulatedCTF);
                        if (Score > BestScore)
                        {
                            BestScore = Score;
                            BestZ = z;
                            BestPhase = p;
                        }
                    }
                }

                CTF = new CTF
                {
                    PixelSize = (MainWindow.Options.CTFPixelMin + MainWindow.Options.CTFPixelMax) * 0.5M,

                    Defocus = (decimal)BestZ,
                    PhaseShift = (decimal)BestPhase,

                    Cs = MainWindow.Options.CTFCs,
                    Voltage = MainWindow.Options.CTFVoltage,
                    Amplitude = MainWindow.Options.CTFAmplitude
                };

                UpdateRotationalAverage(true); // This doesn't have a nice background yet.
                UpdateBackgroundFit(); // Now get a reasonably nice background.
            }

            

            // Fit defocus, (phase shift), (astigmatism) to average background-subtracted spectrum, 
            // which is in polar coords at this point (for equal weighting of all frequencies).

            #region Grid search

            {
                Image CTFMeanPolarTrimmed = CTFMean.AsPolar((uint)MinFreqInclusive, (uint)(MinFreqInclusive + NFreq / 1));

                // Subtract current background.
                Image CurrentBackground = new Image(_SimulatedBackground.Interp(PS1D.Select(p => p.X).ToArray()).Skip(MinFreqInclusive).Take(NFreq / 1).ToArray());
                CTFMeanPolarTrimmed.SubtractFromLines(CurrentBackground);
                CurrentBackground.Dispose();

                /*Image WaterMask = new Image(new int3(NFreq, 1, 1));
                float[] WaterData = WaterMask.GetHost(Intent.Write)[0];
                for (int i = 0; i < NFreq; i++)
                {
                    float f = (i + MinFreqInclusive) / (float)DimsRegion.X * 2f;
                    WaterData[i] = f > 0.2f && f < 0.6f ? 0f : 1f;
                }
                //CTFMeanPolarTrimmed.MultiplyLines(WaterMask);
                WaterMask.Dispose();*/

                // Normalize for CC (not strictly needed, but it's converted for fp16 later, so let's be on the safe side of the fp16 range.
                GPU.Normalize(CTFMeanPolarTrimmed.GetDevice(Intent.Read), CTFMeanPolarTrimmed.GetDevice(Intent.Write), (uint)CTFMeanPolarTrimmed.ElementsReal, 1);
                //CTFMeanPolarTrimmed.WriteMRC("ctfmeanpolartrimmed.mrc");

                CTF StartParams = new CTF
                {
                    PixelSize = (MainWindow.Options.CTFPixelMin + MainWindow.Options.CTFPixelMax) * 0.5M,
                    PixelSizeDelta = Math.Abs(MainWindow.Options.CTFPixelMax - MainWindow.Options.CTFPixelMin),
                    PixelSizeAngle = MainWindow.Options.CTFPixelAngle,

                    Defocus = CTF.Defocus,// (MainWindow.Options.CTFZMin + MainWindow.Options.CTFZMax) * 0.5M,
                    DefocusDelta = doastigmatism ? 0 : MainWindow.Options.CTFAstigmatism,
                    DefocusAngle = doastigmatism ? 0 : MainWindow.Options.CTFAstigmatismAngle,

                    Cs = MainWindow.Options.CTFCs,
                    Voltage = MainWindow.Options.CTFVoltage,
                    Amplitude = MainWindow.Options.CTFAmplitude
                };

                CTFFitStruct FitParams = new CTFFitStruct
                {
                    //Pixelsize = new float3(-0.02e-10f, 0.02e-10f, 0.01e-10f),
                    //Pixeldelta = new float3(0.0f, 0.02e-10f, 0.01e-10f),
                    //Pixelangle = new float3(0, 2 * (float)Math.PI, 1 * (float)Math.PI / 18),

                    //Defocus = new float3((float)(MainWindow.Options.CTFZMin - StartParams.Defocus) * 1e-6f,
                    //                     (float)(MainWindow.Options.CTFZMax - StartParams.Defocus) * 1e-6f,
                    //                     0.025e-6f),
                    Defocus = new float3(-0.4e-6f,
                                         0.4e-6f,
                                         0.025e-6f),

                    Defocusdelta = doastigmatism ? new float3(0, 0.8e-6f, 0.02e-6f) : new float3(0, 0, 0),
                    Astigmatismangle = doastigmatism ? new float3(0, 2 * (float)Math.PI, 1 * (float)Math.PI / 18) : new float3(0, 0, 0),
                    Phaseshift = MainWindow.Options.CTFDoPhase ? new float3(0, (float)Math.PI, 0.025f * (float)Math.PI) : new float3(0, 0, 0)
                };

                CTFStruct ResultStruct = GPU.CTFFitMean(CTFMeanPolarTrimmed.GetDevice(Intent.Read),
                                                        CTFCoordsPolarTrimmed.GetDevice(Intent.Read),
                                                        CTFMeanPolarTrimmed.DimsSlice,
                                                        StartParams.ToStruct(),
                                                        FitParams,
                                                        doastigmatism);
                CTF.FromStruct(ResultStruct);
                CTF.Defocus = Math.Max(CTF.Defocus, MainWindow.Options.CTFZMin);

                CTFMeanPolarTrimmed.Dispose();

                UpdateRotationalAverage(true); // This doesn't have a nice background yet.
                UpdateBackgroundFit(); // Now get a reasonably nice background.

                UpdateRotationalAverage(true); // This time, with the nice background.
                UpdateBackgroundFit(); // Make the background even nicer!
            }

            #endregion

            /*for (int i = 0; i < PS1D.Length; i++)
                PS1D[i].Y -= SimulatedBackground.Interp(PS1D[i].X);
            SimulatedBackground = new Cubic1D(SimulatedBackground.Data.Select(v => new float2(v.X, 0f)).ToArray());
            OnPropertyChanged("PS1D");

            CTFSpectra.Dispose();
            CTFMean.Dispose();
            CTFCoordsCart.Dispose();
            CTFCoordsPolarTrimmed.Dispose();

            Simulated1D = GetSimulated1D();
            CTFQuality = GetCTFQuality();

            return;*/

            // Do BFGS optimization of defocus, astigmatism and phase shift,
            // using 2D simulation for comparison

            #region BFGS

            bool[] CTFSpectraConsider = new bool[CTFSpectraGrid.Elements()];
            for (int i = 0; i < CTFSpectraConsider.Length; i++)
                CTFSpectraConsider[i] = true;
            int NCTFSpectraConsider = CTFSpectraConsider.Length;

            GridCTF = new CubicGrid(GridCTF.Dimensions, (float)CTF.Defocus, (float)CTF.Defocus, Dimension.X);
            GridCTFPhase = new CubicGrid(GridCTFPhase.Dimensions, (float)CTF.PhaseShift, (float)CTF.PhaseShift, Dimension.X);

            for (int preciseFit = 2; preciseFit < 3; preciseFit++)
            {
                NFreq = (MaxFreqExclusive - MinFreqInclusive) * (preciseFit + 1) / 3;
                //if (preciseFit >= 2)
                //    NFreq = MaxFreqExclusive - MinFreqInclusive;

                Image CTFSpectraPolarTrimmed = CTFSpectra.AsPolar((uint)MinFreqInclusive, (uint)(MinFreqInclusive + NFreq));
                CTFSpectra.FreeDevice(); // This will only be needed again for the final PS1D.

                #region Create background and scale

                float[] CurrentScale = _SimulatedScale.Interp(PS1D.Select(p => p.X).ToArray());

                Image CTFSpectraScale = new Image(new int3(NFreq, DimsRegion.X, 1));
                float[] CTFSpectraScaleData = CTFSpectraScale.GetHost(Intent.Write)[0];

                // Trim polar to relevant frequencies, and populate coordinates.
                Parallel.For(0, DimsRegion.X, y =>
                {
                    float Angle = ((float)y / DimsRegion.X + 0.5f) * (float)Math.PI;
                    for (int x = 0; x < NFreq; x++)
                        CTFSpectraScaleData[y * NFreq + x] = CurrentScale[x + MinFreqInclusive];
                });
                //CTFSpectraScale.WriteMRC("ctfspectrascale.mrc");

                // Background is just 1 line since we're in polar.
                Image CurrentBackground = new Image(_SimulatedBackground.Interp(PS1D.Select(p => p.X).ToArray()).Skip(MinFreqInclusive).Take(NFreq).ToArray());

                #endregion

                CTFSpectraPolarTrimmed.SubtractFromLines(CurrentBackground);
                CurrentBackground.Dispose();

                // Normalize background-subtracted spectra.
                GPU.Normalize(CTFSpectraPolarTrimmed.GetDevice(Intent.Read),
                              CTFSpectraPolarTrimmed.GetDevice(Intent.Write),
                              (uint)CTFSpectraPolarTrimmed.ElementsSliceReal,
                              (uint)CTFSpectraGrid.Elements());
                //CTFSpectraPolarTrimmed.WriteMRC("ctfspectrapolartrimmed.mrc");

                #region Convert to fp16

                Image CTFSpectraPolarTrimmedHalf = CTFSpectraPolarTrimmed.AsHalf();
                CTFSpectraPolarTrimmed.Dispose();

                Image CTFSpectraScaleHalf = CTFSpectraScale.AsHalf();
                CTFSpectraScale.Dispose();
                Image CTFCoordsPolarTrimmedHalf = CTFCoordsPolarTrimmed.AsHalf();

                #endregion

                // Wiggle weights show how the defocus on the spectra grid is altered 
                // by changes in individual anchor points of the spline grid.
                // They are used later to compute the dScore/dDefocus values for each spectrum 
                // only once, and derive the values for each anchor point from them.
                float[][] WiggleWeights = GridCTF.GetWiggleWeights(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 1f / (CTFGridZ + 1)));
                float[][] WiggleWeightsPhase = GridCTFPhase.GetWiggleWeights(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 1f / (CTFGridZ + 1)));

                // Helper method for getting CTFStructs for the entire spectra grid.
                Func<double[], CTF, float[], float[], CTFStruct[]> EvalGetCTF = (input, ctf, defocusValues, phaseValues) =>
                {
                    decimal AlteredDelta = (decimal)input[input.Length - 2];
                    decimal AlteredAngle = (decimal)(input[input.Length - 1] * 20 / (Math.PI / 180));

                    CTF Local = ctf.GetCopy();
                    Local.DefocusDelta = AlteredDelta;
                    Local.DefocusAngle = AlteredAngle;

                    CTFStruct LocalStruct = Local.ToStruct();
                    CTFStruct[] LocalParams = new CTFStruct[defocusValues.Length];
                    for (int i = 0; i < LocalParams.Length; i++)
                    {
                        LocalParams[i] = LocalStruct;
                        LocalParams[i].Defocus = defocusValues[i] * -1e-6f;
                        LocalParams[i].PhaseShift = phaseValues[i] * (float)Math.PI;
                    }

                    return LocalParams;
                };

                // Simulate with adjusted CTF, compare to originals

                #region Eval and Gradient methods

                float BorderZ = 0.5f / CTFGridZ;

                Func<double[], double> Eval = input =>
                {
                    CubicGrid Altered = new CubicGrid(GridCTF.Dimensions, input.Take((int)GridCTF.Dimensions.Elements()).Select(v => (float)v).ToArray());
                    float[] DefocusValues = Altered.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));
                    CubicGrid AlteredPhase = new CubicGrid(GridCTFPhase.Dimensions, input.Skip((int)GridCTF.Dimensions.Elements()).Take((int)GridCTFPhase.Dimensions.Elements()).Select(v => (float)v).ToArray());
                    float[] PhaseValues = AlteredPhase.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));

                    CTFStruct[] LocalParams = EvalGetCTF(input, CTF, DefocusValues, PhaseValues);

                    float[] Result = new float[LocalParams.Length];

                    GPU.CTFCompareToSim(CTFSpectraPolarTrimmedHalf.GetDevice(Intent.Read),
                                        CTFCoordsPolarTrimmedHalf.GetDevice(Intent.Read),
                                        CTFSpectraScaleHalf.GetDevice(Intent.Read),
                                        (uint)CTFSpectraPolarTrimmedHalf.ElementsSliceReal,
                                        LocalParams,
                                        Result,
                                        (uint)LocalParams.Length);

                    float Score = 0;
                    for (int i = 0; i < Result.Length; i++)
                        if (CTFSpectraConsider[i])
                            Score += Result[i];

                    Score /= NCTFSpectraConsider;

                    if (float.IsNaN(Score) || float.IsInfinity(Score))
                        throw new Exception("Bad score.");

                    return (1.0 - Score) * 1000.0;
                };

                Func<double[], double[]> Gradient = input =>
                {
                    const float Step = 0.005f;
                    double[] Result = new double[input.Length];

                    // In 0D grid case, just get gradient for all 4 parameters.
                    // In 1+D grid case, do simple gradient for astigmatism and phase...
                    int StartComponent = input.Length - 2;
                    //int StartComponent = 0;
                    for (int i = StartComponent; i < input.Length; i++)
                    {
                        double[] UpperInput = new double[input.Length];
                        input.CopyTo(UpperInput, 0);
                        UpperInput[i] += Step;
                        double UpperValue = Eval(UpperInput);

                        double[] LowerInput = new double[input.Length];
                        input.CopyTo(LowerInput, 0);
                        LowerInput[i] -= Step;
                        double LowerValue = Eval(LowerInput);

                        Result[i] = (UpperValue - LowerValue) / (2f * Step);
                    }

                    float[] ResultPlus = new float[CTFSpectraGrid.Elements()];
                    float[] ResultMinus = new float[CTFSpectraGrid.Elements()];

                    // ..., take shortcut for defoci...
                    {
                        CubicGrid AlteredPhase = new CubicGrid(GridCTFPhase.Dimensions, input.Skip((int)GridCTF.Dimensions.Elements()).Take((int)GridCTFPhase.Dimensions.Elements()).Select(v => (float)v).ToArray());
                        float[] PhaseValues = AlteredPhase.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));

                        {
                            CubicGrid AlteredPlus = new CubicGrid(GridCTF.Dimensions, input.Take((int)GridCTF.Dimensions.Elements()).Select(v => (float)v + Step).ToArray());
                            float[] DefocusValues = AlteredPlus.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));

                            CTFStruct[] LocalParams = EvalGetCTF(input, CTF, DefocusValues, PhaseValues);

                            GPU.CTFCompareToSim(CTFSpectraPolarTrimmedHalf.GetDevice(Intent.Read),
                                                CTFCoordsPolarTrimmedHalf.GetDevice(Intent.Read),
                                                CTFSpectraScaleHalf.GetDevice(Intent.Read),
                                                (uint)CTFSpectraPolarTrimmedHalf.ElementsSliceReal,
                                                LocalParams,
                                                ResultPlus,
                                                (uint)LocalParams.Length);
                        }
                        {
                            CubicGrid AlteredMinus = new CubicGrid(GridCTF.Dimensions, input.Take((int)GridCTF.Dimensions.Elements()).Select(v => (float)v - Step).ToArray());
                            float[] DefocusValues = AlteredMinus.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));

                            CTFStruct[] LocalParams = EvalGetCTF(input, CTF, DefocusValues, PhaseValues);

                            GPU.CTFCompareToSim(CTFSpectraPolarTrimmedHalf.GetDevice(Intent.Read),
                                                CTFCoordsPolarTrimmedHalf.GetDevice(Intent.Read),
                                                CTFSpectraScaleHalf.GetDevice(Intent.Read),
                                                (uint)CTFSpectraPolarTrimmedHalf.ElementsSliceReal,
                                                LocalParams,
                                                ResultMinus,
                                                (uint)LocalParams.Length);
                        }
                        float[] LocalGradients = new float[ResultPlus.Length];
                        for (int i = 0; i < LocalGradients.Length; i++)
                            LocalGradients[i] = ResultMinus[i] - ResultPlus[i];

                        // Now compute gradients per grid anchor point using the precomputed individual gradients and wiggle factors.
                        Parallel.For(0, GridCTF.Dimensions.Elements(), i => Result[i] = MathHelper.ReduceWeighted(LocalGradients, WiggleWeights[i]) / LocalGradients.Length / (2f * Step) * 1000f);
                    }

                    // ..., and take shortcut for phases.
                    if (MainWindow.Options.CTFDoPhase)
                    {
                        CubicGrid AlteredPlus = new CubicGrid(GridCTF.Dimensions, input.Take((int)GridCTF.Dimensions.Elements()).Select(v => (float)v).ToArray());
                        float[] DefocusValues = AlteredPlus.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));

                        {
                            CubicGrid AlteredPhasePlus = new CubicGrid(GridCTFPhase.Dimensions, input.Skip((int)GridCTF.Dimensions.Elements()).Take((int)GridCTFPhase.Dimensions.Elements()).Select(v => (float)v + Step).ToArray());
                            float[] PhaseValues = AlteredPhasePlus.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));
                            CTFStruct[] LocalParams = EvalGetCTF(input, CTF, DefocusValues, PhaseValues);

                            GPU.CTFCompareToSim(CTFSpectraPolarTrimmedHalf.GetDevice(Intent.Read),
                                                CTFCoordsPolarTrimmedHalf.GetDevice(Intent.Read),
                                                CTFSpectraScaleHalf.GetDevice(Intent.Read),
                                                (uint)CTFSpectraPolarTrimmedHalf.ElementsSliceReal,
                                                LocalParams,
                                                ResultPlus,
                                                (uint)LocalParams.Length);
                        }
                        {
                            CubicGrid AlteredPhaseMinus = new CubicGrid(GridCTFPhase.Dimensions, input.Skip((int)GridCTF.Dimensions.Elements()).Take((int)GridCTFPhase.Dimensions.Elements()).Select(v => (float)v - Step).ToArray());
                            float[] PhaseValues = AlteredPhaseMinus.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));
                            CTFStruct[] LocalParams = EvalGetCTF(input, CTF, DefocusValues, PhaseValues);

                            GPU.CTFCompareToSim(CTFSpectraPolarTrimmedHalf.GetDevice(Intent.Read),
                                                CTFCoordsPolarTrimmedHalf.GetDevice(Intent.Read),
                                                CTFSpectraScaleHalf.GetDevice(Intent.Read),
                                                (uint)CTFSpectraPolarTrimmedHalf.ElementsSliceReal,
                                                LocalParams,
                                                ResultMinus,
                                                (uint)LocalParams.Length);
                        }
                        float[] LocalGradients = new float[ResultPlus.Length];
                        for (int i = 0; i < LocalGradients.Length; i++)
                            LocalGradients[i] = ResultMinus[i] - ResultPlus[i];

                        // Now compute gradients per grid anchor point using the precomputed individual gradients and wiggle factors.
                        Parallel.For(0, GridCTFPhase.Dimensions.Elements(), i => Result[i + GridCTF.Dimensions.Elements()] = MathHelper.ReduceWeighted(LocalGradients, WiggleWeightsPhase[i]) / LocalGradients.Length / (2f * Step) * 1000f);
                    }

                    foreach (var i in Result)
                        if (double.IsNaN(i) || double.IsInfinity(i))
                            throw new Exception("Bad score.");

                    return Result;
                };

                #endregion

                #region Minimize first time with potential outpiers

                double[] StartParams = new double[GridCTF.Dimensions.Elements() + GridCTFPhase.Dimensions.Elements() + 2];
                for (int i = 0; i < GridCTF.Dimensions.Elements(); i++)
                    StartParams[i] = GridCTF.FlatValues[i];
                for (int i = 0; i < GridCTFPhase.Dimensions.Elements(); i++)
                    StartParams[i + GridCTF.Dimensions.Elements()] = GridCTFPhase.FlatValues[i];
                StartParams[StartParams.Length - 2] = (double)CTF.DefocusDelta;
                StartParams[StartParams.Length - 1] = (double)CTF.DefocusAngle / 20 * (Math.PI / 180);

                // Compute correlation for individual spectra, and throw away those that are >.75 sigma worse than mean.

                #region Discard outliers

                if (CTFSpace || CTFTime)
                {
                    CubicGrid Altered = new CubicGrid(GridCTF.Dimensions, StartParams.Take((int)GridCTF.Dimensions.Elements()).Select(v => (float)v).ToArray());
                    float[] DefocusValues = Altered.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));
                    CubicGrid AlteredPhase = new CubicGrid(GridCTFPhase.Dimensions, StartParams.Skip((int)GridCTF.Dimensions.Elements()).Take((int)GridCTFPhase.Dimensions.Elements()).Select(v => (float)v).ToArray());
                    float[] PhaseValues = AlteredPhase.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));

                    CTFStruct[] LocalParams = EvalGetCTF(StartParams, CTF, DefocusValues, PhaseValues);

                    float[] Result = new float[LocalParams.Length];

                    GPU.CTFCompareToSim(CTFSpectraPolarTrimmedHalf.GetDevice(Intent.Read),
                                        CTFCoordsPolarTrimmedHalf.GetDevice(Intent.Read),
                                        CTFSpectraScaleHalf.GetDevice(Intent.Read),
                                        (uint)CTFSpectraPolarTrimmedHalf.ElementsSliceReal,
                                        LocalParams,
                                        Result,
                                        (uint)LocalParams.Length);

                    float MeanResult = MathHelper.Mean(Result);
                    float StdResult = MathHelper.StdDev(Result);
                    CTFSpectraConsider = new bool[CTFSpectraGrid.Elements()];
                    Parallel.For(0, CTFSpectraConsider.Length, i =>
                    {
                        //if (Result[i] > MeanResult - StdResult * 1.5f)
                        CTFSpectraConsider[i] = true;
                        /*else
                        {
                            CTFSpectraConsider[i] = false;
                            for (int j = 0; j < WiggleWeights.Length; j++)
                                // Make sure the spectrum's gradient doesn't affect the overall gradient.
                                WiggleWeights[j][i] = 0;
                        }*/
                    });
                    NCTFSpectraConsider = CTFSpectraConsider.Where(v => v).Count();
                }

                #endregion

                BroydenFletcherGoldfarbShanno Optimizer = new BroydenFletcherGoldfarbShanno(StartParams.Length, Eval, Gradient)
                {
                    Past = 1,
                    Delta = 1e-6,
                    MaxLineSearch = 15,
                    Corrections = 20
                };
                Optimizer.Minimize(StartParams);

                #endregion

                #region Retrieve parameters

                CTF.Defocus = (decimal)MathHelper.Mean(Optimizer.Solution.Take((int)GridCTF.Dimensions.Elements()).Select(v => (float)v));
                CTF.DefocusDelta = (decimal)Optimizer.Solution[StartParams.Length - 2];
                CTF.DefocusAngle = (decimal)(Optimizer.Solution[StartParams.Length - 1] * 20 / (Math.PI / 180));
                CTF.PhaseShift = (decimal)MathHelper.Mean(Optimizer.Solution.Skip((int)GridCTF.Dimensions.Elements()).Take((int)GridCTFPhase.Dimensions.Elements()).Select(v => (float)v));

                if (CTF.DefocusDelta < 0)
                {
                    CTF.DefocusAngle += 90;
                    CTF.DefocusDelta *= -1;
                }
                CTF.DefocusAngle = ((int)CTF.DefocusAngle + 180 * 99) % 180;

                GridCTF = new CubicGrid(GridCTF.Dimensions, Optimizer.Solution.Take((int)GridCTF.Dimensions.Elements()).Select(v => (float)v).ToArray());
                GridCTFPhase = new CubicGrid(GridCTFPhase.Dimensions, Optimizer.Solution.Skip((int)GridCTF.Dimensions.Elements()).Take((int)GridCTFPhase.Dimensions.Elements()).Select(v => (float)v).ToArray());

                #endregion

                // Dispose GPU resources manually because GC can't be bothered to do it in time.
                CTFSpectraPolarTrimmedHalf.Dispose();
                CTFCoordsPolarTrimmedHalf.Dispose();
                CTFSpectraScaleHalf.Dispose();

                #region Get nicer envelope fit

                if (preciseFit >= 2)
                {
                    if (!CTFSpace && !CTFTime)
                    {
                        UpdateRotationalAverage(true);
                    }
                    else
                    {
                        Image CTFSpectraBackground = new Image(new int3(DimsRegion), true);
                        float[] CTFSpectraBackgroundData = CTFSpectraBackground.GetHost(Intent.Write)[0];

                        // Construct background in Cartesian coordinates.
                        Helper.ForEachElementFT(DimsRegion, (x, y, xx, yy, r, a) =>
                        {
                            CTFSpectraBackgroundData[y * CTFSpectraBackground.DimsEffective.X + x] = _SimulatedBackground.Interp(r / DimsRegion.X);
                        });

                        CTFSpectra.SubtractFromSlices(CTFSpectraBackground);

                        float[] DefocusValues = GridCTF.GetInterpolatedNative(CTFSpectraGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, BorderZ));
                        CTFStruct[] LocalParams = DefocusValues.Select(v =>
                        {
                            CTF Local = CTF.GetCopy();
                            Local.Defocus = (decimal)v + 0.0M;

                            return Local.ToStruct();
                        }).ToArray();

                        Image CTFAverage1D = new Image(IntPtr.Zero, new int3(DimsRegion.X / 2, 1, 1));

                        CTF CTFAug = CTF.GetCopy();
                        CTFAug.Defocus += 0.0M;
                        GPU.CTFMakeAverage(CTFSpectra.GetDevice(Intent.Read),
                                           CTFCoordsCart.GetDevice(Intent.Read),
                                           (uint)CTFSpectra.ElementsSliceReal,
                                           (uint)DimsRegion.X,
                                           LocalParams,
                                           CTFAug.ToStruct(),
                                           0,
                                           (uint)DimsRegion.X / 2,
                                           CTFSpectraConsider.Select(v => v ? 1 : 0).ToArray(),
                                           (uint)CTFSpectraGrid.Elements(),
                                           CTFAverage1D.GetDevice(Intent.Write));

                        CTFSpectra.AddToSlices(CTFSpectraBackground);

                        float[] RotationalAverageData = CTFAverage1D.GetHost(Intent.Read)[0];
                        float2[] ForPS1D = new float2[PS1D.Length];
                        for (int i = 0; i < ForPS1D.Length; i++)
                            ForPS1D[i] = new float2((float)i / DimsRegion.X, (float)Math.Round(RotationalAverageData[i], 4) + _SimulatedBackground.Interp((float)i / DimsRegion.X));
                        MathHelper.UnNaN(ForPS1D);
                        _PS1D = ForPS1D;

                        CTFSpectraBackground.Dispose();
                        CTFAverage1D.Dispose();
                        CTFSpectra.FreeDevice();
                    }

                    CTF.Defocus = Math.Max(CTF.Defocus, MainWindow.Options.CTFZMin);
                    UpdateBackgroundFit();
                }

                #endregion
            }

            #endregion

            // Subtract background from 2D average and write it to disk. 
            // This image is used for quick visualization purposes only.

            #region PS2D update

            {
                int3 DimsAverage = new int3(DimsRegion.X, DimsRegion.X / 2, 1);
                float[] Average2DData = new float[DimsAverage.Elements()];
                float[] OriginalAverageData = CTFMean.GetHost(Intent.Read)[0];

                for (int y = 0; y < DimsAverage.Y; y++)
                {
                    int yy = y * y;
                    for (int x = 0; x < DimsAverage.Y; x++)
                    {
                        int xx = DimsRegion.X / 2 - x - 1;
                        xx *= xx;
                        float r = (float)Math.Sqrt(xx + yy) / DimsRegion.X;
                        Average2DData[y * DimsAverage.X + x] = OriginalAverageData[(y + DimsRegion.X / 2) * (DimsRegion.X / 2 + 1) + x] - SimulatedBackground.Interp(r);
                    }

                    for (int x = 0; x < DimsRegion.X / 2; x++)
                    {
                        int xx = x * x;
                        float r = (float)Math.Sqrt(xx + yy) / DimsRegion.X;
                        Average2DData[y * DimsAverage.X + x + DimsRegion.X / 2] = OriginalAverageData[(DimsRegion.X / 2 - y) * (DimsRegion.X / 2 + 1) + (DimsRegion.X / 2 - 1 - x)] - SimulatedBackground.Interp(r);
                    }
                }

                IOHelper.WriteMapFloat(PowerSpectrumPath,
                                       new HeaderMRC
                                       {
                                           Dimensions = DimsAverage,
                                           MinValue = MathHelper.Min(Average2DData),
                                           MaxValue = MathHelper.Max(Average2DData)
                                       },
                                       Average2DData);

                PS2DTemp = null;
                OnPropertyChanged("PS2D");
            }

            #endregion

            for (int i = 0; i < PS1D.Length; i++)
                PS1D[i].Y -= SimulatedBackground.Interp(PS1D[i].X);
            SimulatedBackground = new Cubic1D(SimulatedBackground.Data.Select(v => new float2(v.X, 0f)).ToArray());
            OnPropertyChanged("PS1D");

            CTFSpectra.Dispose();
            CTFMean.Dispose();
            CTFCoordsCart.Dispose();
            CTFCoordsPolarTrimmed.Dispose();

            Simulated1D = GetSimulated1D();
            CTFQuality = GetCTFQuality();

            SaveMeta();
        }
Exemple #2
0
        public void ProcessParticleShift(MapHeader originalHeader, Image originalStack, Star stardata, Image refft, Image maskft, int dimbox, decimal scaleFactor)
        {
            // Deal with dimensions and grids.

            int NFrames = originalHeader.Dimensions.Z;
            int2 DimsImage = new int2(originalHeader.Dimensions);
            int2 DimsRegion = new int2(dimbox, dimbox);

            decimal SubdivisionRatio = 4M;
            List<int3> PyramidSizes = new List<int3>();
            PyramidSizes.Add(new int3(MainWindow.Options.GridMoveX, MainWindow.Options.GridMoveX, Math.Min(NFrames, MainWindow.Options.GridMoveZ)));
            while (true)
            {
                int3 Previous = PyramidSizes.Last();
                int NewZ = Math.Min((int)Math.Round(Previous.Z / SubdivisionRatio), Previous.Z - 1);
                if (NewZ < 2)
                    break;

                PyramidSizes.Add(new int3(Previous.X * 2, Previous.Y * 2, NewZ));
            }

            PyramidShiftX.Clear();
            PyramidShiftY.Clear();

            float3[] PositionsGrid, PositionsGridPerFrame;
            float2[] PositionsExtraction, PositionsShift;
            float3[] ParticleAngles;
            List<int> RowIndices = new List<int>();
            {
                string[] ColumnNames = stardata.GetColumn("rlnMicrographName");
                for (int i = 0; i < ColumnNames.Length; i++)
                    if (ColumnNames[i].Contains(RootName))
                        RowIndices.Add(i);

                string[] ColumnOriginX = stardata.GetColumn("rlnCoordinateX");
                string[] ColumnOriginY = stardata.GetColumn("rlnCoordinateY");
                string[] ColumnShiftX = stardata.GetColumn("rlnOriginX");
                string[] ColumnShiftY = stardata.GetColumn("rlnOriginY");
                string[] ColumnAngleRot = stardata.GetColumn("rlnAngleRot");
                string[] ColumnAngleTilt = stardata.GetColumn("rlnAngleTilt");
                string[] ColumnAnglePsi = stardata.GetColumn("rlnAnglePsi");

                PositionsGrid = new float3[RowIndices.Count];
                PositionsGridPerFrame = new float3[RowIndices.Count * NFrames];
                PositionsExtraction = new float2[RowIndices.Count];
                PositionsShift = new float2[RowIndices.Count * NFrames];
                ParticleAngles = new float3[RowIndices.Count];

                {
                    int i = 0;
                    foreach (var nameIndex in RowIndices)
                    {
                        float OriginX = float.Parse(ColumnOriginX[nameIndex]);
                        float OriginY = float.Parse(ColumnOriginY[nameIndex]);
                        float ShiftX = float.Parse(ColumnShiftX[nameIndex]);
                        float ShiftY = float.Parse(ColumnShiftY[nameIndex]);

                        PositionsExtraction[i] = new float2(OriginX - ShiftX, OriginY - ShiftY);
                        PositionsGrid[i] = new float3((OriginX - ShiftX) / DimsImage.X, (OriginY - ShiftY) / DimsImage.Y, 0.5f);
                        for (int z = 0; z < NFrames; z++)
                        {
                            PositionsGridPerFrame[z * RowIndices.Count + i] = new float3(PositionsGrid[i].X,
                                                                                          PositionsGrid[i].Y,
                                                                                          (float)z / (NFrames - 1));

                            PositionsShift[z * RowIndices.Count + i] = GetShiftFromPyramid(PositionsGridPerFrame[z * RowIndices.Count + i]);
                        }
                        ParticleAngles[i] = new float3(float.Parse(ColumnAngleRot[nameIndex]) * Helper.ToRad,
                                                       float.Parse(ColumnAngleTilt[nameIndex]) * Helper.ToRad,
                                                       float.Parse(ColumnAnglePsi[nameIndex]) * Helper.ToRad);

                        i++;
                    }
                }
            }
            int NPositions = PositionsGrid.Length;
            if (NPositions == 0)
                return;

            int MinFreqInclusive = (int)(MainWindow.Options.MovementRangeMin * DimsRegion.X / 2);
            int MaxFreqExclusive = (int)(MainWindow.Options.MovementRangeMax * DimsRegion.X / 2);
            int NFreq = MaxFreqExclusive - MinFreqInclusive;

            int CentralFrame = NFrames / 2;

            int MaskExpansions = 4; // Math.Max(1, PyramidSizes[0].Z / 3);
            int[] MaskSizes = new int[MaskExpansions];

            // Allocate memory and create all prerequisites:
            int MaskLength;
            Image ShiftFactors;
            Image Phases;
            Image Projections;
            Image Shifts;
            Image InvSigma;
            {
                List<long> Positions = new List<long>();
                List<float2> Factors = new List<float2>();
                List<float2> Freq = new List<float2>();
                int Min2 = MinFreqInclusive * MinFreqInclusive;
                int Max2 = MaxFreqExclusive * MaxFreqExclusive;

                for (int y = 0; y < DimsRegion.Y; y++)
                {
                    int yy = y - DimsRegion.X / 2;
                    for (int x = 0; x < DimsRegion.X / 2 + 1; x++)
                    {
                        int xx = x - DimsRegion.X / 2;
                        int r2 = xx * xx + yy * yy;
                        if (r2 >= Min2 && r2 < Max2)
                        {
                            Positions.Add(y * (DimsRegion.X / 2 + 1) + x);
                            Factors.Add(new float2((float)xx / DimsRegion.X * 2f * (float)Math.PI,
                                                   (float)yy / DimsRegion.X * 2f * (float)Math.PI));

                            float Angle = (float)Math.Atan2(yy, xx);
                            float r = (float)Math.Sqrt(r2);
                            Freq.Add(new float2(r, Angle));
                        }
                    }
                }

                // Addresses for CTF simulation
                Image CTFCoordsCart = new Image(new int3(DimsRegion), true, true);
                {
                    float2[] CoordsData = new float2[CTFCoordsCart.ElementsSliceComplex];

                    Helper.ForEachElementFT(DimsRegion, (x, y, xx, yy, r, a) => CoordsData[y * (DimsRegion.X / 2 + 1) + x] = new float2(r / DimsRegion.X, a));
                    CTFCoordsCart.UpdateHostWithComplex(new[] { CoordsData });
                    CTFCoordsCart.RemapToFT();
                }
                float[] ValuesDefocus = GridCTF.GetInterpolatedNative(PositionsGrid);
                CTFStruct[] PositionsCTF = ValuesDefocus.Select(v =>
                {
                    CTF Altered = CTF.GetCopy();
                    Altered.PixelSizeDelta = 0;
                    Altered.Defocus = (decimal)v;
                    //Altered.Bfactor = -MainWindow.Options.MovementBfactor;
                    return Altered.ToStruct();
                }).ToArray();

                // Sort everyone with ascending distance from center.
                List<KeyValuePair<float, int>> FreqIndices = Freq.Select((v, i) => new KeyValuePair<float, int>(v.X, i)).ToList();
                FreqIndices.Sort((a, b) => a.Key.CompareTo(b.Key));
                int[] SortedIndices = FreqIndices.Select(v => v.Value).ToArray();

                Helper.Reorder(Positions, SortedIndices);
                Helper.Reorder(Factors, SortedIndices);
                Helper.Reorder(Freq, SortedIndices);

                long[] RelevantMask = Positions.ToArray();
                ShiftFactors = new Image(Helper.ToInterleaved(Factors.ToArray()));
                MaskLength = RelevantMask.Length;

                // Get mask sizes for different expansion steps.
                for (int i = 0; i < MaskExpansions; i++)
                {
                    float CurrentMaxFreq = MinFreqInclusive + (MaxFreqExclusive - MinFreqInclusive) / (float)MaskExpansions * (i + 1);
                    MaskSizes[i] = Freq.Count(v => v.X * v.X < CurrentMaxFreq * CurrentMaxFreq);
                }

                Phases = new Image(IntPtr.Zero, new int3(MaskLength, NPositions, NFrames), false, true, false);
                Projections = new Image(IntPtr.Zero, new int3(MaskLength, NPositions, NFrames), false, true, false);
                InvSigma = new Image(IntPtr.Zero, new int3(MaskLength, 1, 1));

                Image ParticleMasksFT = maskft.AsProjections(ParticleAngles, DimsRegion, MainWindow.Options.ProjectionOversample);
                Image ParticleMasks = ParticleMasksFT.AsIFFT();
                ParticleMasksFT.Dispose();
                ParticleMasks.RemapFromFT();

                Parallel.ForEach(ParticleMasks.GetHost(Intent.ReadWrite), slice =>
                {
                    for (int i = 0; i < slice.Length; i++)
                        slice[i] = (Math.Max(2f, Math.Min(25f, slice[i])) - 2) / 23f;
                });

                Image ProjectionsSparse = refft.AsProjections(ParticleAngles, DimsRegion, MainWindow.Options.ProjectionOversample);

                Image InvSigmaSparse = new Image(new int3(DimsRegion), true);
                {
                    int GroupNumber = int.Parse(stardata.GetRowValue(RowIndices[0], "rlnGroupNumber"));
                    //Star SigmaTable = new Star("D:\\rado27\\RefineWarppolish\\run1_model.star", "data_model_group_" + GroupNumber);
                    Star SigmaTable = new Star(MainWindow.Options.ModelStarPath, "data_model_group_" + GroupNumber);
                    float[] SigmaValues = SigmaTable.GetColumn("rlnSigma2Noise").Select(v => float.Parse(v)).ToArray();

                    float[] Sigma2NoiseData = InvSigmaSparse.GetHost(Intent.Write)[0];
                    Helper.ForEachElementFT(new int2(DimsRegion.X, DimsRegion.Y), (x, y, xx, yy, r, angle) =>
                    {
                        int ir = (int)r;
                        float val = 0;
                        if (ir < SigmaValues.Length)
                        {
                            if (SigmaValues[ir] != 0f)
                                val = 1f / SigmaValues[ir];
                        }
                        Sigma2NoiseData[y * (DimsRegion.X / 2 + 1) + x] = val;
                    });
                    float MaxSigma = MathHelper.Max(Sigma2NoiseData);
                    for (int i = 0; i < Sigma2NoiseData.Length; i++)
                        Sigma2NoiseData[i] /= MaxSigma;

                    InvSigmaSparse.RemapToFT();
                }
                //InvSigmaSparse.WriteMRC("d_sigma2noise.mrc");

                float PixelSize = (float)CTF.PixelSize;
                float PixelDelta = (float)CTF.PixelSizeDelta;
                float PixelAngle = (float)CTF.PixelSizeAngle * Helper.ToRad;

                GPU.CreateParticleShift(originalStack.GetDevice(Intent.Read),
                                        DimsImage,
                                        NFrames,
                                        Helper.ToInterleaved(PositionsExtraction),
                                        Helper.ToInterleaved(PositionsShift),
                                        NPositions,
                                        DimsRegion,
                                        RelevantMask,
                                        (uint)RelevantMask.Length,
                                        ParticleMasks.GetDevice(Intent.Read),
                                        ProjectionsSparse.GetDevice(Intent.Read),
                                        PositionsCTF,
                                        CTFCoordsCart.GetDevice(Intent.Read),
                                        InvSigmaSparse.GetDevice(Intent.Read),
                                        PixelSize + PixelDelta / 2,
                                        PixelSize - PixelDelta / 2,
                                        PixelAngle,
                                        Phases.GetDevice(Intent.Write),
                                        Projections.GetDevice(Intent.Write),
                                        InvSigma.GetDevice(Intent.Write));

                InvSigmaSparse.Dispose();
                ParticleMasks.Dispose();
                ProjectionsSparse.Dispose();
                CTFCoordsCart.Dispose();
                originalStack.FreeDevice();
                Shifts = new Image(new float[NPositions * NFrames * 2]);
            }

            #region Fit movement

            {

                int NPyramidPoints = 0;
                float[][][] WiggleWeights = new float[PyramidSizes.Count][][];
                for (int p = 0; p < PyramidSizes.Count; p++)
                {
                    CubicGrid WiggleGrid = new CubicGrid(PyramidSizes[p]);
                    NPyramidPoints += (int)PyramidSizes[p].Elements();

                    WiggleWeights[p] = WiggleGrid.GetWiggleWeights(PositionsGridPerFrame);
                }

                double[] StartParams = new double[NPyramidPoints * 2];

                for (int m = 3; m < MaskExpansions; m++)
                {
                    for (int currentGrid = 0; currentGrid < PyramidSizes.Count; currentGrid++)
                    {
                        Action<double[]> SetPositions = input =>
                        {
                            // Construct CubicGrids and get interpolated shift values.
                            float[] AlteredX = new float[PositionsGridPerFrame.Length];
                            float[] AlteredY = new float[PositionsGridPerFrame.Length];

                            int Offset = 0;
                            foreach (var size in PyramidSizes)
                            {
                                int Elements = (int)size.Elements();
                                CubicGrid GridX = new CubicGrid(size, input.Skip(Offset).Take(Elements).Select(v => (float)v).ToArray());
                                AlteredX = MathHelper.Plus(AlteredX, GridX.GetInterpolatedNative(PositionsGridPerFrame));

                                CubicGrid GridY = new CubicGrid(size, input.Skip(NPyramidPoints + Offset).Take(Elements).Select(v => (float)v).ToArray());
                                AlteredY = MathHelper.Plus(AlteredY, GridY.GetInterpolatedNative(PositionsGridPerFrame));

                                Offset += Elements;
                            }

                            // Finally, set the shift values in the device array.
                            float[] ShiftData = Shifts.GetHost(Intent.Write)[0];
                            for (int i = 0; i < PositionsGridPerFrame.Length; i++)
                            {
                                ShiftData[i * 2] = AlteredX[i];
                                ShiftData[i * 2 + 1] = AlteredY[i];
                            }
                        };

                        Func<double[], double> Eval = input =>
                        {
                            SetPositions(input);

                            float[] Diff = new float[NPositions * NFrames];
                            GPU.ParticleShiftGetDiff(Phases.GetDevice(Intent.Read),
                                                     Projections.GetDevice(Intent.Read),
                                                     ShiftFactors.GetDevice(Intent.Read),
                                                     InvSigma.GetDevice(Intent.Read),
                                                     (uint)MaskLength,
                                                     (uint)MaskSizes[m],
                                                     Shifts.GetDevice(Intent.Read),
                                                     Diff,
                                                     (uint)NPositions,
                                                     (uint)NFrames);

                            //for (int i = 0; i < Diff.Length; i++)
                            //Diff[i] = Diff[i] * 100f;

                            double Score = Diff.Sum();
                            //Debug.WriteLine(Score);
                            return Score;
                        };

                        Func<double[], double[]> Grad = input =>
                        {
                            SetPositions(input);

                            float[] Diff = new float[NPositions * NFrames * 2];
                            GPU.ParticleShiftGetGrad(Phases.GetDevice(Intent.Read),
                                                     Projections.GetDevice(Intent.Read),
                                                     ShiftFactors.GetDevice(Intent.Read),
                                                     InvSigma.GetDevice(Intent.Read),
                                                     (uint)MaskLength,
                                                     (uint)MaskSizes[m],
                                                     Shifts.GetDevice(Intent.Read),
                                                     Diff,
                                                     (uint)NPositions,
                                                     (uint)NFrames);

                            //for (int i = 0; i < Diff.Length; i++)
                                //Diff[i] = Diff[i] * 100f;

                            float[] DiffX = new float[NPositions * NFrames], DiffY = new float[NPositions * NFrames];
                            for (int i = 0; i < DiffX.Length; i++)
                            {
                                DiffX[i] = Diff[i * 2];
                                DiffY[i] = Diff[i * 2 + 1];
                            }

                            double[] Result = new double[input.Length];
                            int Offset = 0;
                            for (int p = 0; p < PyramidSizes.Count; p++)
                            {
                                //if (p == currentGrid)
                                    Parallel.For(0, (int)PyramidSizes[p].Elements(), i =>
                                    {
                                        Result[Offset + i] = MathHelper.ReduceWeighted(DiffX, WiggleWeights[p][i]);
                                        Result[NPyramidPoints + Offset + i] = MathHelper.ReduceWeighted(DiffY, WiggleWeights[p][i]);
                                    });

                                Offset += (int)PyramidSizes[p].Elements();
                            }
                            return Result;
                        };

                        BroydenFletcherGoldfarbShanno Optimizer = new BroydenFletcherGoldfarbShanno(StartParams.Length, Eval, Grad);
                        //Optimizer.Corrections = 20;
                        Optimizer.Minimize(StartParams);
                    }

                    {
                        PyramidShiftX.Clear();
                        PyramidShiftY.Clear();
                        int Offset = 0;
                        foreach (var size in PyramidSizes)
                        {
                            int Elements = (int)size.Elements();
                            CubicGrid GridX = new CubicGrid(size, StartParams.Skip(Offset).Take(Elements).Select(v => (float)v).ToArray());
                            PyramidShiftX.Add(GridX);

                            CubicGrid GridY = new CubicGrid(size, StartParams.Skip(NPyramidPoints + Offset).Take(Elements).Select(v => (float)v).ToArray());
                            PyramidShiftY.Add(GridY);

                            Offset += Elements;
                        }
                    }
                }
            }

            #endregion

            ShiftFactors.Dispose();
            Phases.Dispose();
            Projections.Dispose();
            Shifts.Dispose();
            InvSigma.Dispose();

            SaveMeta();
        }