public void Correlate(Image tiltStack, Image reference, int size, float lowpassAngstrom, float highpassAngstrom, int3 volumeDimensions, int healpixOrder, string symmetry = "C1") { if (!Directory.Exists(CorrelationDir)) Directory.CreateDirectory(CorrelationDir); float DownscaleFactor = lowpassAngstrom / (float)(CTF.PixelSize * 2); Image DownscaledStack = tiltStack.AsScaledMassive(new int2(tiltStack.Dims) / DownscaleFactor / 2 * 2); tiltStack.FreeDevice(); float HighpassNyquist = (float)(CTF.PixelSize * 2) * DownscaleFactor / highpassAngstrom; DownscaledStack.Bandpass(HighpassNyquist, 1, false); Image ScaledReference = reference.AsScaled(reference.Dims / DownscaleFactor / 2 * 2); reference.FreeDevice(); Image PaddedReference = ScaledReference.AsPadded(new int3(size, size, size)); ScaledReference.Dispose(); PaddedReference.Bandpass(HighpassNyquist, 1, true); Projector ProjectorReference = new Projector(PaddedReference, 2); PaddedReference.Dispose(); VolumeDimensions = volumeDimensions / DownscaleFactor / 2 * 2; int SizeCropped = size / 2; int3 Grid = (VolumeDimensions + SizeCropped - 1) / SizeCropped; List<float3> GridCoords = new List<float3>(); for (int z = 0; z < Grid.Z; z++) for (int y = 0; y < Grid.Y; y++) for (int x = 0; x < Grid.X; x++) GridCoords.Add(new float3(x * SizeCropped + SizeCropped / 2, y * SizeCropped + SizeCropped / 2, z * SizeCropped + SizeCropped / 2)); float3[] HealpixAngles = Helper.GetHealpixAngles(healpixOrder, symmetry).Select(a => a * Helper.ToRad).ToArray(); Image CTFCoords = GetCTFCoords(size, (int)(size * DownscaleFactor)); float[] OutputCorr = new float[VolumeDimensions.Elements()]; int PlanForw, PlanBack, PlanForwCTF; Projector.GetPlans(new int3(size, size, size), 2, out PlanForw, out PlanBack, out PlanForwCTF); int BatchSize = 16; for (int b = 0; b < GridCoords.Count; b += BatchSize) { int CurBatch = Math.Min(BatchSize, GridCoords.Count - b); Image Subtomos = new Image(IntPtr.Zero, new int3(size, size, size * CurBatch), true, true); Image SubtomoCTFs = new Image(IntPtr.Zero, new int3(size, size, size * CurBatch), true); for (int st = 0; st < CurBatch; st++) { Image ImagesFT = GetSubtomoImages(DownscaledStack, size, GridCoords[b + st], true); Image CTFs = GetSubtomoCTFs(GridCoords[b + st], CTFCoords); //Image CTFWeights = GetSubtomoCTFs(GridCoords[b + st], CTFCoords, true, true); ImagesFT.Multiply(CTFs); // Weight and phase-flip image FTs by CTF, which still has its sign here //ImagesFT.Multiply(CTFWeights); CTFs.Abs(); // CTF has to be positive from here on since image FT phases are now flipped // CTF has to be converted to complex numbers with imag = 0, and weighted by itself float2[] CTFsComplexData = new float2[CTFs.ElementsComplex]; float[] CTFsContinuousData = CTFs.GetHostContinuousCopy(); for (int i = 0; i < CTFsComplexData.Length; i++) CTFsComplexData[i] = new float2(CTFsContinuousData[i] * CTFsContinuousData[i], 0); Image CTFsComplex = new Image(CTFsComplexData, CTFs.Dims, true); Projector ProjSubtomo = new Projector(new int3(size, size, size), 2); lock (GPU.Sync) ProjSubtomo.BackProject(ImagesFT, CTFs, GetAngleInImages(GridCoords[b + st])); Image Subtomo = ProjSubtomo.Reconstruct(false, PlanForw, PlanBack, PlanForwCTF); ProjSubtomo.Dispose(); /*Image CroppedSubtomo = Subtomo.AsPadded(new int3(SizeCropped, SizeCropped, SizeCropped)); CroppedSubtomo.WriteMRC(ParticlesDir + RootName + "_" + (b + st).ToString("D5") + ".mrc"); CroppedSubtomo.Dispose();*/ Projector ProjCTF = new Projector(new int3(size, size, size), 2); lock (GPU.Sync) ProjCTF.BackProject(CTFsComplex, CTFs, GetAngleInImages(GridCoords[b + st])); Image SubtomoCTF = ProjCTF.Reconstruct(true, PlanForw, PlanBack, PlanForwCTF); ProjCTF.Dispose(); GPU.FFT(Subtomo.GetDevice(Intent.Read), Subtomos.GetDeviceSlice(size * st, Intent.Write), Subtomo.Dims, 1); GPU.CopyDeviceToDevice(SubtomoCTF.GetDevice(Intent.Read), SubtomoCTFs.GetDeviceSlice(size * st, Intent.Write), SubtomoCTF.ElementsReal); ImagesFT.Dispose(); CTFs.Dispose(); //CTFWeights.Dispose(); CTFsComplex.Dispose(); Subtomo.Dispose(); SubtomoCTF.Dispose(); } Image BestCorrelation = new Image(IntPtr.Zero, new int3(size, size, size * CurBatch)); Image BestRot = new Image(IntPtr.Zero, new int3(size, size, size * CurBatch)); Image BestTilt = new Image(IntPtr.Zero, new int3(size, size, size * CurBatch)); Image BestPsi = new Image(IntPtr.Zero, new int3(size, size, size * CurBatch)); GPU.CorrelateSubTomos(ProjectorReference.Data.GetDevice(Intent.Read), ProjectorReference.Oversampling, ProjectorReference.Data.Dims, Subtomos.GetDevice(Intent.Read), SubtomoCTFs.GetDevice(Intent.Read), new int3(size, size, size), (uint)CurBatch, Helper.ToInterleaved(HealpixAngles), (uint)HealpixAngles.Length, MainWindow.Options.ExportParticleRadius / ((float)CTF.PixelSize * DownscaleFactor), BestCorrelation.GetDevice(Intent.Write), BestRot.GetDevice(Intent.Write), BestTilt.GetDevice(Intent.Write), BestPsi.GetDevice(Intent.Write)); for (int st = 0; st < CurBatch; st++) { Image ThisCorrelation = new Image(BestCorrelation.GetDeviceSlice(size * st, Intent.Read), new int3(size, size, size)); Image CroppedCorrelation = ThisCorrelation.AsPadded(new int3(SizeCropped, SizeCropped, SizeCropped)); //CroppedCorrelation.WriteMRC(CorrelationDir + RootName + "_" + (b + st).ToString("D5") + ".mrc"); float[] SubCorr = CroppedCorrelation.GetHostContinuousCopy(); int3 Origin = new int3(GridCoords[b + st]) - SizeCropped / 2; for (int z = 0; z < SizeCropped; z++) { int zVol = Origin.Z + z; if (zVol >= VolumeDimensions.Z) continue; for (int y = 0; y < SizeCropped; y++) { int yVol = Origin.Y + y; if (yVol >= VolumeDimensions.Y) continue; for (int x = 0; x < SizeCropped; x++) { int xVol = Origin.X + x; if (xVol >= VolumeDimensions.X) continue; OutputCorr[(zVol * VolumeDimensions.Y + yVol) * VolumeDimensions.X + xVol] = SubCorr[(z * SizeCropped + y) * SizeCropped + x]; } } } CroppedCorrelation.Dispose(); ThisCorrelation.Dispose(); } Subtomos.Dispose(); SubtomoCTFs.Dispose(); BestCorrelation.Dispose(); BestRot.Dispose(); BestTilt.Dispose(); BestPsi.Dispose(); } GPU.DestroyFFTPlan(PlanForw); GPU.DestroyFFTPlan(PlanBack); GPU.DestroyFFTPlan(PlanForwCTF); CTFCoords.Dispose(); ProjectorReference.Dispose(); DownscaledStack.Dispose(); Image OutputCorrImage = new Image(OutputCorr, VolumeDimensions); OutputCorrImage.WriteMRC(CorrelationDir + RootName + ".mrc"); OutputCorrImage.Dispose(); }
public void Reconstruct(Image tiltStack, int size, float lowpassAngstrom, int3 volumeDimensions) { if (!Directory.Exists(ReconstructionDir)) Directory.CreateDirectory(ReconstructionDir); if (File.Exists(ReconstructionDir + RootName + ".mrc")) return; VolumeDimensions = volumeDimensions; float DownscaleFactor = lowpassAngstrom / (float)(CTF.PixelSize * 2); Image DownscaledStack = tiltStack.AsScaledMassive(new int2(tiltStack.Dims) / DownscaleFactor / 2 * 2); tiltStack.FreeDevice(); GPU.Normalize(DownscaledStack.GetDevice(Intent.Read), DownscaledStack.GetDevice(Intent.Write), (uint)DownscaledStack.ElementsSliceReal, (uint)DownscaledStack.Dims.Z); int3 VolumeDimensionsCropped = volumeDimensions / DownscaleFactor / 2 * 2; int SizeCropped = size / 4; int3 Grid = (VolumeDimensionsCropped + SizeCropped - 1) / SizeCropped; List<float3> GridCoords = new List<float3>(); for (int z = 0; z < Grid.Z; z++) for (int y = 0; y < Grid.Y; y++) for (int x = 0; x < Grid.X; x++) GridCoords.Add(new float3(x * SizeCropped + SizeCropped / 2, y * SizeCropped + SizeCropped / 2, z * SizeCropped + SizeCropped / 2)); Image CTFCoords = GetCTFCoords(size, (int)(size * DownscaleFactor)); float[] OutputRec = new float[VolumeDimensionsCropped.Elements()]; int PlanForw, PlanBack, PlanForwCTF; Projector.GetPlans(new int3(size, size, size), 2, out PlanForw, out PlanBack, out PlanForwCTF); for (int p = 0; p < GridCoords.Count; p++) { Image ImagesFT = GetSubtomoImages(DownscaledStack, size, GridCoords[p] * DownscaleFactor, false, 1f / DownscaleFactor); Image CTFs = GetSubtomoCTFs(GridCoords[p] * DownscaleFactor, CTFCoords); //Image CTFWeights = GetSubtomoCTFs(GridCoords[p], CTFCoords, true, true); ImagesFT.Multiply(CTFs); // Weight and phase-flip image FTs by CTF, which still has its sign here //ImagesFT.Multiply(CTFWeights); CTFs.Abs(); // Since everything is already phase-flipped, weights are positive Projector ProjSubtomo = new Projector(new int3(size, size, size), 2); lock (GPU.Sync) ProjSubtomo.BackProject(ImagesFT, CTFs, GetAngleInImages(GridCoords[p] * DownscaleFactor)); Image Subtomo = ProjSubtomo.Reconstruct(false, PlanForw, PlanBack, PlanForwCTF); Image SubtomoCropped = Subtomo.AsPadded(new int3(SizeCropped, SizeCropped, SizeCropped)); Subtomo.Dispose(); ProjSubtomo.Dispose(); /*Image CroppedSubtomo = Subtomo.AsPadded(new int3(SizeCropped, SizeCropped, SizeCropped)); CroppedSubtomo.WriteMRC(ParticlesDir + RootName + "_" + (b + st).ToString("D5") + ".mrc"); CroppedSubtomo.Dispose();*/ ImagesFT.Dispose(); CTFs.Dispose(); //CTFWeights.Dispose(); float[] SubtomoData = SubtomoCropped.GetHostContinuousCopy(); int3 Origin = new int3(GridCoords[p]) - SizeCropped / 2; for (int z = 0; z < SizeCropped; z++) { int zVol = Origin.Z + z; if (zVol >= VolumeDimensionsCropped.Z) continue; for (int y = 0; y < SizeCropped; y++) { int yVol = Origin.Y + y; if (yVol >= VolumeDimensionsCropped.Y) continue; for (int x = 0; x < SizeCropped; x++) { int xVol = Origin.X + x; if (xVol >= VolumeDimensionsCropped.X) continue; OutputRec[(zVol * VolumeDimensionsCropped.Y + yVol) * VolumeDimensionsCropped.X + xVol] = -SubtomoData[(z * SizeCropped + y) * SizeCropped + x]; } } } SubtomoCropped.Dispose(); } GPU.DestroyFFTPlan(PlanForw); GPU.DestroyFFTPlan(PlanBack); GPU.DestroyFFTPlan(PlanForwCTF); CTFCoords.Dispose(); DownscaledStack.Dispose(); Image OutputRecImage = new Image(OutputRec, VolumeDimensionsCropped); OutputRecImage.WriteMRC(ReconstructionDir + RootName + ".mrc"); OutputRecImage.Dispose(); }