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();
}
public Tuple<Image, Image> MakeReconstructionOneTomogram(Image tiltStack, int subset, int size, float3[] particleOrigins, float3[] particleOrigins2, float3[] particleAngles, float3[] particleAngles2, int[] particleSubset)
{
Projector MapProjector = new Projector(new int3(size, size, size), 2);
Projector WeightProjector = new Projector(new int3(size, size, size), 2);
List<int> ParticleIDs = new List<int>();
for (int i = 0; i < particleSubset.Length; i++)
if (particleSubset[i] == subset)
ParticleIDs.Add(i);
int NParticles = ParticleIDs.Count;
particleOrigins = ParticleIDs.Select(i => particleOrigins[i]).ToArray();
particleOrigins2 = ParticleIDs.Select(i => particleOrigins2[i]).ToArray();
particleAngles = ParticleIDs.Select(i => particleAngles[i]).ToArray();
particleAngles2 = ParticleIDs.Select(i => particleAngles2[i]).ToArray();
Image CTFCoords = GetCTFCoords(size, size);
for (int angleID = 0; angleID < NTilts; angleID++)
{
float DoseID = IndicesSortedDose[angleID] / (float)(NTilts - 1);
for (int b = 0; b < NParticles; b += 1024)
{
int NBatch = Math.Min(1024, NParticles - b);
float3[] ParticleOriginsInterp = new float3[NBatch];
float3[] ParticleAnglesInterp = new float3[NBatch];
for (int n = 0; n < NBatch; n++)
{
float3 OriginDiff = particleOrigins2[b + n] - particleOrigins[b + n];
float3 AngleDiff = particleAngles2[b + n] - particleAngles[b + n];
ParticleOriginsInterp[n] = particleOrigins[b + n] + OriginDiff * DoseID;
ParticleAnglesInterp[n] = particleAngles[b + n] + AngleDiff * DoseID;
}
Image ParticleImages = GetImagesOneAngle(tiltStack, size, ParticleOriginsInterp, angleID, true);
Image ParticleCTFs = GetCTFsOneAngle(tiltStack, CTFCoords, ParticleOriginsInterp, angleID, true);
ParticleImages.Multiply(ParticleCTFs);
//ParticleCTFs.Multiply(ParticleCTFs);
ParticleCTFs.Abs();
MapProjector.BackProject(ParticleImages,
ParticleCTFs,
GetAnglesInOneAngle(ParticleOriginsInterp,
ParticleAnglesInterp,
angleID));
ParticleImages.Dispose();
ParticleCTFs.Dispose();
// Now reconstruct the weights which will be needed during optimization later
ParticleCTFs = GetCTFsOneAngle(tiltStack, CTFCoords, ParticleOriginsInterp, angleID, true);
// CTF has to be converted to complex numbers with imag = 0
float2[] CTFsComplexData = new float2[ParticleCTFs.ElementsComplex];
float[] CTFWeightsData = new float[ParticleCTFs.ElementsComplex];
float[] CTFsContinuousData = ParticleCTFs.GetHostContinuousCopy();
for (int i = 0; i < CTFsComplexData.Length; i++)
{
CTFsComplexData[i] = new float2(Math.Abs(CTFsContinuousData[i] * CTFsContinuousData[i]), 0);
CTFWeightsData[i] = Math.Abs(CTFsContinuousData[i]);
}
Image CTFsComplex = new Image(CTFsComplexData, ParticleCTFs.Dims, true);
Image CTFWeights = new Image(CTFWeightsData, ParticleCTFs.Dims, true);
WeightProjector.BackProject(CTFsComplex,
CTFWeights,
GetAnglesInOneAngle(ParticleOriginsInterp,
ParticleAnglesInterp,
angleID));
ParticleCTFs.Dispose();
CTFsComplex.Dispose();
CTFWeights.Dispose();
}
}
CTFCoords.Dispose();
//MapProjector.Weights.WriteMRC($"d_weights{angleID:D3}.mrc");
Image Reconstruction = MapProjector.Reconstruct(false);
MapProjector.Dispose();
foreach (var slice in WeightProjector.Weights.GetHost(Intent.ReadWrite))
for (int i = 0; i < slice.Length; i++)
slice[i] = Math.Min(slice[i], 1);
Image ReconstructionWeights = WeightProjector.Reconstruct(true);
WeightProjector.Dispose();
return new Tuple<Image, Image>(Reconstruction, ReconstructionWeights);
}
public void AddToReconstructionOneAngle(Image tiltStack,
int subset,
int size,
float3[] particleOrigins,
float3[] particleOrigins2,
float3[] particleAngles,
float3[] particleAngles2,
int[] particleSubset,
int angleID,
Projector mapProjector,
Projector weightProjector)
{
List<int> ParticleIDs = new List<int>();
for (int i = 0; i < particleSubset.Length; i++)
if (particleSubset[i] == subset)
ParticleIDs.Add(i);
int NParticles = ParticleIDs.Count;
particleOrigins = ParticleIDs.Select(i => particleOrigins[i]).ToArray();
particleOrigins2 = ParticleIDs.Select(i => particleOrigins2[i]).ToArray();
particleAngles = ParticleIDs.Select(i => particleAngles[i]).ToArray();
particleAngles2 = ParticleIDs.Select(i => particleAngles2[i]).ToArray();
Image CTFCoords = GetCTFCoords(size, size);
float DoseID = IndicesSortedDose[angleID] / (float)(NTilts - 1);
for (int b = 0; b < NParticles; b += 1024)
{
int NBatch = Math.Min(1024, NParticles - b);
float3[] ParticleOriginsInterp = new float3[NBatch];
float3[] ParticleAnglesInterp = new float3[NBatch];
for (int n = 0; n < NBatch; n++)
{
float3 OriginDiff = particleOrigins2[b + n] - particleOrigins[b + n];
float3 AngleDiff = particleAngles2[b + n] - particleAngles[b + n];
ParticleOriginsInterp[n] = particleOrigins[b + n] + OriginDiff * DoseID;
ParticleAnglesInterp[n] = particleAngles[b + n] + AngleDiff * DoseID;
}
Image ParticleImages = GetImagesOneAngle(tiltStack, size, ParticleOriginsInterp, angleID, true);
Image ParticleCTFs = GetCTFsOneAngle(tiltStack, CTFCoords, ParticleOriginsInterp, angleID, false);
ParticleImages.Multiply(ParticleCTFs);
//ParticleCTFs.Multiply(ParticleCTFs);
ParticleCTFs.Abs();
mapProjector.BackProject(ParticleImages,
ParticleCTFs,
GetAnglesInOneAngle(ParticleOriginsInterp,
ParticleAnglesInterp,
angleID));
ParticleImages.Dispose();
ParticleCTFs.Dispose();
// Now reconstruct the weights which will be needed during optimization later
ParticleCTFs = GetCTFsOneAngle(tiltStack, CTFCoords, ParticleOriginsInterp, angleID, false);
// CTF has to be converted to complex numbers with imag = 0
float2[] CTFsComplexData = new float2[ParticleCTFs.ElementsComplex];
float[] CTFWeightsData = new float[ParticleCTFs.ElementsComplex];
float[] CTFsContinuousData = ParticleCTFs.GetHostContinuousCopy();
for (int i = 0; i < CTFsComplexData.Length; i++)
{
CTFsComplexData[i] = new float2(Math.Abs(CTFsContinuousData[i] * CTFsContinuousData[i]), 0);
CTFWeightsData[i] = Math.Abs(CTFsContinuousData[i]);
}
Image CTFsComplex = new Image(CTFsComplexData, ParticleCTFs.Dims, true);
Image CTFWeights = new Image(CTFWeightsData, ParticleCTFs.Dims, true);
weightProjector.BackProject(CTFsComplex,
CTFWeights,
GetAnglesInOneAngle(ParticleOriginsInterp,
ParticleAnglesInterp,
angleID));
ParticleCTFs.Dispose();
CTFsComplex.Dispose();
CTFWeights.Dispose();
}
CTFCoords.Dispose();
}
public void GetSubtomo(Image tiltStack, float3 coords, float3 angles, Image ctfCoords, out Image subtomo, out Image subtomoCTF, int planForw = -1, int planBack = -1, int planForwCTF = -1)
{
int Size = ctfCoords.Dims.X;
float3[] ImageAngles = GetAngleInImages(coords);
Image ImagesFT = GetSubtomoImages(tiltStack, Size, coords, true);
Image CTFs = GetSubtomoCTFs(coords, ctfCoords, true, false, false);
//Image CTFWeights = GetSubtomoCTFs(coords, 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, ImageAngles);
subtomo = ProjSubtomo.Reconstruct(false, planForw, planBack, planForwCTF);
ProjSubtomo.Dispose();
GPU.NormParticles(subtomo.GetDevice(Intent.Read),
subtomo.GetDevice(Intent.Write),
subtomo.Dims,
(uint)(MainWindow.Options.ExportParticleRadius / CTF.PixelSize),
false,
1);
//subtomo = new Image(new int3(1, 1, 1));
Projector ProjCTF = new Projector(new int3(Size, Size, Size), 2);
lock (GPU.Sync)
ProjCTF.BackProject(CTFsComplex, CTFs, ImageAngles);
subtomoCTF = ProjCTF.Reconstruct(true, planForw, planBack, planForwCTF);
ProjCTF.Dispose();
//subtomoCTF = new Image(new int3(1, 1, 1));
ImagesFT.Dispose();
CTFs.Dispose();
//CTFWeights.Dispose();
CTFsComplex.Dispose();
}
