public void ProcessShift(MapHeader originalHeader, Image originalStack, decimal scaleFactor) { // Deal with dimensions and grids. int NFrames = originalHeader.Dimensions.Z; int2 DimsImage = new int2(originalHeader.Dimensions); int2 DimsRegion = new int2(768, 768); float OverlapFraction = 0.0f; int2 DimsPositionGrid; int3[] PositionGrid = Helper.GetEqualGridSpacing(DimsImage, DimsRegion, OverlapFraction, out DimsPositionGrid); //PositionGrid = new[] { new int3(0, 0, 0) }; //DimsPositionGrid = new int2(1, 1); int NPositions = PositionGrid.Length; int ShiftGridX = 1; int ShiftGridY = 1; int ShiftGridZ = Math.Min(NFrames, MainWindow.Options.GridMoveZ); GridMovementX = new CubicGrid(new int3(ShiftGridX, ShiftGridY, ShiftGridZ)); GridMovementY = new CubicGrid(new int3(ShiftGridX, ShiftGridY, ShiftGridZ)); int LocalGridX = Math.Min(DimsPositionGrid.X, MainWindow.Options.GridMoveX); int LocalGridY = Math.Min(DimsPositionGrid.Y, MainWindow.Options.GridMoveY); int LocalGridZ = Math.Min(2, NFrames); GridLocalX = new CubicGrid(new int3(LocalGridX, LocalGridY, LocalGridZ)); GridLocalY = new CubicGrid(new int3(LocalGridX, LocalGridY, LocalGridZ)); int3 ShiftGrid = new int3(DimsPositionGrid.X, DimsPositionGrid.Y, NFrames); 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 = Math.Max(1, ShiftGridZ / 3); int[] MaskSizes = new int[MaskExpansions]; // Allocate memory and create all prerequisites: int MaskLength; Image ShiftFactors; Image Phases; Image PhasesAverage; Image Shifts; { 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; 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; 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)); } } } // 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); float Bfac = (float)MainWindow.Options.MovementBfactor * 0.25f / PixelSize / DimsRegion.X; float2[] BfacWeightsData = Freq.Select(v => (float)Math.Exp(v.X * Bfac)).Select(v => new float2(v, v)).ToArray(); Image BfacWeights = new Image(Helper.ToInterleaved(BfacWeightsData), false, false, false); 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 * 2, DimsPositionGrid.X * DimsPositionGrid.Y, NFrames), false, false, false); GPU.CreateShift(originalStack.GetDevice(Intent.Read), new int2(originalHeader.Dimensions), originalHeader.Dimensions.Z, PositionGrid, PositionGrid.Length, DimsRegion, RelevantMask, (uint)MaskLength, Phases.GetDevice(Intent.Write)); Phases.MultiplyLines(BfacWeights); BfacWeights.Dispose(); originalStack.FreeDevice(); PhasesAverage = new Image(IntPtr.Zero, new int3(MaskLength, NPositions, 1), false, true, false); Shifts = new Image(new float[NPositions * NFrames * 2]); } #region Fit global movement { int MinXSteps = 1, MinYSteps = 1; int MinZSteps = Math.Min(NFrames, 3); int3 ExpansionGridSize = new int3(MinXSteps, MinYSteps, MinZSteps); float[][] WiggleWeights = new CubicGrid(ExpansionGridSize).GetWiggleWeights(ShiftGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 0f)); double[] StartParams = new double[ExpansionGridSize.Elements() * 2]; for (int m = 0; m < MaskExpansions; m++) { double[] LastAverage = null; Action<double[]> SetPositions = input => { // Construct CubicGrids and get interpolated shift values. CubicGrid AlteredGridX = new CubicGrid(ExpansionGridSize, input.Where((v, i) => i % 2 == 0).Select(v => (float)v).ToArray()); float[] AlteredX = AlteredGridX.GetInterpolatedNative(new int3(DimsPositionGrid.X, DimsPositionGrid.Y, NFrames), new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 0f)); CubicGrid AlteredGridY = new CubicGrid(ExpansionGridSize, input.Where((v, i) => i % 2 == 1).Select(v => (float)v).ToArray()); float[] AlteredY = AlteredGridY.GetInterpolatedNative(new int3(DimsPositionGrid.X, DimsPositionGrid.Y, NFrames), new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 0f)); // Let movement start at 0 in the central frame. /*float2[] CenterFrameOffsets = new float2[NPositions]; for (int i = 0; i < NPositions; i++) CenterFrameOffsets[i] = new float2(AlteredX[CentralFrame * NPositions + i], AlteredY[CentralFrame * NPositions + i]);*/ // Finally, set the shift values in the device array. float[] ShiftData = Shifts.GetHost(Intent.Write)[0]; Parallel.For(0, AlteredX.Length, i => { ShiftData[i * 2] = AlteredX[i];// - CenterFrameOffsets[i % NPositions].X; ShiftData[i * 2 + 1] = AlteredY[i];// - CenterFrameOffsets[i % NPositions].Y; }); }; Action<double[]> DoAverage = input => { if (LastAverage == null || input.Where((t, i) => t != LastAverage[i]).Any()) { SetPositions(input); GPU.ShiftGetAverage(Phases.GetDevice(Intent.Read), PhasesAverage.GetDevice(Intent.Write), ShiftFactors.GetDevice(Intent.Read), (uint)MaskLength, (uint)MaskSizes[m], Shifts.GetDevice(Intent.Read), (uint)NPositions, (uint)NFrames); if (LastAverage == null) LastAverage = new double[input.Length]; Array.Copy(input, LastAverage, input.Length); } }; Func<double[], double> Eval = input => { DoAverage(input); float[] Diff = new float[NPositions * NFrames]; GPU.ShiftGetDiff(Phases.GetDevice(Intent.Read), PhasesAverage.GetDevice(Intent.Read), ShiftFactors.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; return Diff.Sum(); }; Func<double[], double[]> Grad = input => { DoAverage(input); float[] GradX = new float[NPositions * NFrames], GradY = new float[NPositions * NFrames]; float[] Diff = new float[NPositions * NFrames * 2]; GPU.ShiftGetGrad(Phases.GetDevice(Intent.Read), PhasesAverage.GetDevice(Intent.Read), ShiftFactors.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; for (int i = 0; i < GradX.Length; i++) { GradX[i] = Diff[i * 2]; GradY[i] = Diff[i * 2 + 1]; } double[] Result = new double[input.Length]; Parallel.For(0, input.Length / 2, i => { Result[i * 2] = MathHelper.ReduceWeighted(GradX, WiggleWeights[i]); Result[i * 2 + 1] = MathHelper.ReduceWeighted(GradY, WiggleWeights[i]); }); return Result; }; /*Func<double[], double[]> Grad = input => { DoAverage(input); float[] GradX = new float[NPositions * NFrames], GradY = new float[NPositions * NFrames]; float Step = 0.002f; { double[] InputXP = new double[input.Length]; for (int i = 0; i < input.Length; i++) if (i % 2 == 0) InputXP[i] = input[i] + Step; else InputXP[i] = input[i]; SetPositions(InputXP); float[] DiffXP = new float[NPositions * NFrames]; GPU.ShiftGetDiff(Phases.GetDevice(Intent.Read), PhasesAverage.GetDevice(Intent.Read), ShiftFactors.GetDevice(Intent.Read), (uint)MaskLength, (uint)MaskSizes[m], Shifts.GetDevice(Intent.Read), DiffXP, (uint)NPositions, (uint)NFrames); double[] InputXM = new double[input.Length]; for (int i = 0; i < input.Length; i++) if (i % 2 == 0) InputXM[i] = input[i] - Step; else InputXM[i] = input[i]; SetPositions(InputXM); float[] DiffXM = new float[NPositions * NFrames]; GPU.ShiftGetDiff(Phases.GetDevice(Intent.Read), PhasesAverage.GetDevice(Intent.Read), ShiftFactors.GetDevice(Intent.Read), (uint)MaskLength, (uint)MaskSizes[m], Shifts.GetDevice(Intent.Read), DiffXM, (uint)NPositions, (uint)NFrames); for (int i = 0; i < GradX.Length; i++) GradX[i] = (DiffXP[i] - DiffXM[i]) / (Step * 2); } { double[] InputYP = new double[input.Length]; for (int i = 0; i < input.Length; i++) if (i % 2 == 1) InputYP[i] = input[i] + Step; else InputYP[i] = input[i]; SetPositions(InputYP); float[] DiffYP = new float[NPositions * NFrames]; GPU.ShiftGetDiff(Phases.GetDevice(Intent.Read), PhasesAverage.GetDevice(Intent.Read), ShiftFactors.GetDevice(Intent.Read), (uint)MaskLength, (uint)MaskSizes[m], Shifts.GetDevice(Intent.Read), DiffYP, (uint)NPositions, (uint)NFrames); double[] InputYM = new double[input.Length]; for (int i = 0; i < input.Length; i++) if (i % 2 == 1) InputYM[i] = input[i] - Step; else InputYM[i] = input[i]; SetPositions(InputYM); float[] DiffYM = new float[NPositions * NFrames]; GPU.ShiftGetDiff(Phases.GetDevice(Intent.Read), PhasesAverage.GetDevice(Intent.Read), ShiftFactors.GetDevice(Intent.Read), (uint)MaskLength, (uint)MaskSizes[m], Shifts.GetDevice(Intent.Read), DiffYM, (uint)NPositions, (uint)NFrames); for (int i = 0; i < GradY.Length; i++) GradY[i] = (DiffYP[i] - DiffYM[i]) / (Step * 2); } double[] Result = new double[input.Length]; Parallel.For(0, input.Length / 2, i => { Result[i * 2] = MathHelper.ReduceWeighted(GradX, WiggleWeights[i]); Result[i * 2 + 1] = MathHelper.ReduceWeighted(GradY, WiggleWeights[i]); }); return Result; };*/ BroydenFletcherGoldfarbShanno Optimizer = new BroydenFletcherGoldfarbShanno(StartParams.Length, Eval, Grad); Optimizer.Corrections = 20; Optimizer.Minimize(StartParams); float MeanX = MathHelper.Mean(Optimizer.Solution.Where((v, i) => i % 2 == 0).Select(v => (float)v)); float MeanY = MathHelper.Mean(Optimizer.Solution.Where((v, i) => i % 2 == 1).Select(v => (float)v)); for (int i = 0; i < ExpansionGridSize.Elements(); i++) { Optimizer.Solution[i * 2] -= MeanX; Optimizer.Solution[i * 2 + 1] -= MeanY; } // Store coarse values in grids. GridMovementX = new CubicGrid(ExpansionGridSize, Optimizer.Solution.Where((v, i) => i % 2 == 0).Select(v => (float)v).ToArray()); GridMovementY = new CubicGrid(ExpansionGridSize, Optimizer.Solution.Where((v, i) => i % 2 == 1).Select(v => (float)v).ToArray()); if (m < MaskExpansions - 1) { // Refine sampling. ExpansionGridSize = new int3((int)Math.Round((float)(ShiftGridX - MinXSteps) / (MaskExpansions - 1) * (m + 1) + MinXSteps), (int)Math.Round((float)(ShiftGridY - MinYSteps) / (MaskExpansions - 1) * (m + 1) + MinYSteps), (int)Math.Round((float)(ShiftGridZ - MinZSteps) / (MaskExpansions - 1) * (m + 1) + MinZSteps)); WiggleWeights = new CubicGrid(ExpansionGridSize).GetWiggleWeights(ShiftGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 0f)); // Resize the grids to account for finer sampling. GridMovementX = GridMovementX.Resize(ExpansionGridSize); GridMovementY = GridMovementY.Resize(ExpansionGridSize); // Construct start parameters for next optimization iteration. StartParams = new double[ExpansionGridSize.Elements() * 2]; for (int i = 0; i < ExpansionGridSize.Elements(); i++) { StartParams[i * 2] = GridMovementX.FlatValues[i]; StartParams[i * 2 + 1] = GridMovementY.FlatValues[i]; } } } } #endregion // Center the global shifts /*{ float2[] AverageShifts = new float2[ShiftGridZ]; for (int i = 0; i < AverageShifts.Length; i++) AverageShifts[i] = new float2(MathHelper.Mean(GridMovementX.GetSliceXY(i)), MathHelper.Mean(GridMovementY.GetSliceXY(i))); float2 CenterShift = MathHelper.Mean(AverageShifts); GridMovementX = new CubicGrid(GridMovementX.Dimensions, GridMovementX.FlatValues.Select(v => v - CenterShift.X).ToArray()); GridMovementY = new CubicGrid(GridMovementY.Dimensions, GridMovementY.FlatValues.Select(v => v - CenterShift.Y).ToArray()); }*/ #region Fit local movement /*{ int MinXSteps = LocalGridX, MinYSteps = LocalGridY; int MinZSteps = LocalGridZ; int3 ExpansionGridSize = new int3(MinXSteps, MinYSteps, MinZSteps); float[][] WiggleWeights = new CubicGrid(ExpansionGridSize).GetWiggleWeights(ShiftGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 0f)); double[] StartParams = new double[ExpansionGridSize.Elements() * 2]; for (int m = MaskExpansions - 1; m < MaskExpansions; m++) { double[] LastAverage = null; Action<double[]> SetPositions = input => { // Construct CubicGrids and get interpolated shift values. float[] GlobalX = GridMovementX.GetInterpolatedNative(new int3(DimsPositionGrid.X, DimsPositionGrid.Y, NFrames), new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 0f)); CubicGrid AlteredGridX = new CubicGrid(ExpansionGridSize, input.Where((v, i) => i % 2 == 0).Select(v => (float)v).ToArray()); float[] AlteredX = AlteredGridX.GetInterpolatedNative(new int3(DimsPositionGrid.X, DimsPositionGrid.Y, NFrames), new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 0f)); AlteredX = MathHelper.Plus(GlobalX, AlteredX); float[] GlobalY = GridMovementY.GetInterpolatedNative(new int3(DimsPositionGrid.X, DimsPositionGrid.Y, NFrames), new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 0f)); CubicGrid AlteredGridY = new CubicGrid(ExpansionGridSize, input.Where((v, i) => i % 2 == 1).Select(v => (float)v).ToArray()); float[] AlteredY = AlteredGridY.GetInterpolatedNative(new int3(DimsPositionGrid.X, DimsPositionGrid.Y, NFrames), new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 0f)); AlteredY = MathHelper.Plus(GlobalY, AlteredY); // Let movement start at 0 in the central frame. float2[] CenterFrameOffsets = new float2[NPositions]; for (int i = 0; i < NPositions; i++) CenterFrameOffsets[i] = new float2(AlteredX[CentralFrame * NPositions + i], AlteredY[CentralFrame * NPositions + i]); // Finally, set the shift values in the device array. float[] ShiftData = Shifts.GetHost(Intent.Write)[0]; Parallel.For(0, AlteredX.Length, i => { ShiftData[i * 2] = AlteredX[i] - CenterFrameOffsets[i % NPositions].X; ShiftData[i * 2 + 1] = AlteredY[i] - CenterFrameOffsets[i % NPositions].Y; }); }; Action<double[]> DoAverage = input => { if (LastAverage == null || input.Where((t, i) => t != LastAverage[i]).Any()) { SetPositions(input); GPU.ShiftGetAverage(Phases.GetDevice(Intent.Read), PhasesAverage.GetDevice(Intent.Write), ShiftFactors.GetDevice(Intent.Read), (uint)MaskLength, (uint)MaskSizes[m], Shifts.GetDevice(Intent.Read), (uint)NPositions, (uint)NFrames); if (LastAverage == null) LastAverage = new double[input.Length]; Array.Copy(input, LastAverage, input.Length); } }; Func<double[], double> Eval = input => { DoAverage(input); float[] Diff = new float[NPositions * NFrames]; GPU.ShiftGetDiff(Phases.GetDevice(Intent.Read), PhasesAverage.GetDevice(Intent.Read), ShiftFactors.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; return MathHelper.Mean(Diff); }; Func<double[], double[]> Grad = input => { DoAverage(input); float[] Diff = new float[NPositions * NFrames * 2]; GPU.ShiftGetGrad(Phases.GetDevice(Intent.Read), PhasesAverage.GetDevice(Intent.Read), ShiftFactors.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]; Parallel.For(0, input.Length / 2, i => { Result[i * 2] = MathHelper.ReduceWeighted(DiffX, WiggleWeights[i]); Result[i * 2 + 1] = MathHelper.ReduceWeighted(DiffY, WiggleWeights[i]); }); return Result; }; BroydenFletcherGoldfarbShanno Optimizer = new BroydenFletcherGoldfarbShanno(StartParams.Length, Eval, Grad); Optimizer.Corrections = 20; Optimizer.Minimize(StartParams); float MeanX = MathHelper.Mean(Optimizer.Solution.Where((v, i) => i % 2 == 0).Select(v => (float)v)); float MeanY = MathHelper.Mean(Optimizer.Solution.Where((v, i) => i % 2 == 1).Select(v => (float)v)); for (int i = 0; i < ExpansionGridSize.Elements(); i++) { Optimizer.Solution[i * 2] -= MeanX; Optimizer.Solution[i * 2 + 1] -= MeanY; } // Store coarse values in grids. GridLocalX = new CubicGrid(ExpansionGridSize, Optimizer.Solution.Where((v, i) => i % 2 == 0).Select(v => (float)v).ToArray()); GridLocalY = new CubicGrid(ExpansionGridSize, Optimizer.Solution.Where((v, i) => i % 2 == 1).Select(v => (float)v).ToArray()); if (m < MaskExpansions - 1) { // Refine sampling. ExpansionGridSize = new int3((int)Math.Round((float)(LocalGridX - MinXSteps) / (MaskExpansions - 1) * (m + 1) + MinXSteps), (int)Math.Round((float)(LocalGridY - MinYSteps) / (MaskExpansions - 1) * (m + 1) + MinYSteps), (int)Math.Round((float)(LocalGridZ - MinZSteps) / (MaskExpansions - 1) * (m + 1) + MinZSteps)); WiggleWeights = new CubicGrid(ExpansionGridSize).GetWiggleWeights(ShiftGrid, new float3(DimsRegion.X / 2f / DimsImage.X, DimsRegion.Y / 2f / DimsImage.Y, 0f)); // Resize the grids to account for finer sampling. GridLocalX = GridLocalX.Resize(ExpansionGridSize); GridLocalY = GridLocalY.Resize(ExpansionGridSize); // Construct start parameters for next optimization iteration. StartParams = new double[ExpansionGridSize.Elements() * 2]; for (int i = 0; i < ExpansionGridSize.Elements(); i++) { StartParams[i * 2] = GridLocalX.FlatValues[i]; StartParams[i * 2 + 1] = GridLocalY.FlatValues[i]; } } } }*/ #endregion ShiftFactors.Dispose(); Phases.Dispose(); PhasesAverage.Dispose(); Shifts.Dispose(); // Center the local shifts /*{ float2[] AverageShifts = new float2[LocalGridZ]; for (int i = 0; i < AverageShifts.Length; i++) AverageShifts[i] = new float2(MathHelper.Mean(GridLocalX.GetSliceXY(i)), MathHelper.Mean(GridLocalY.GetSliceXY(i))); float2 CenterShift = MathHelper.Mean(AverageShifts); GridLocalX = new CubicGrid(GridLocalX.Dimensions, GridLocalX.FlatValues.Select(v => v - CenterShift.X).ToArray()); GridLocalY = new CubicGrid(GridLocalY.Dimensions, GridLocalY.FlatValues.Select(v => v - CenterShift.Y).ToArray()); }*/ SaveMeta(); }