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();
}
