def evalScore(self, rot_trans):
"""
evaluate score for given rotation and translation - NEGATIVE cc because\
optimizer MINIMIZES
@param rot_trans: rotation and translation vector (6-dim: z1,z2,x \
angles, x,y,z, translations)
@type rot_trans: L{list}
@author: FF
"""
from pytom_volume import transformSpline, transform
if self.interpolation.lower() == 'spline':
transformSpline(self.vol2, self.rotvol2, rot_trans[0],
rot_trans[1], rot_trans[2], self.centX, self.centY,
self.centZ, 0, 0, 0, rot_trans[3] / self.binning,
rot_trans[4] / self.binning,
rot_trans[5] / self.binning)
else:
transform(self.vol2, self.rotvol2, rot_trans[0], rot_trans[1],
rot_trans[2], self.centX, self.centY, self.centZ, 0, 0,
0, rot_trans[3] / self.binning,
rot_trans[4] / self.binning, rot_trans[5] / self.binning)
self.val = -1. * (self.score(volume=self.vol1,
template=self.rotvol2,
mask=self.mask,
volumeIsNormalized=True))
return self.val
def writeCroppedParticles(particleListName, output, center, cubesize):
"""
@param particleListName: name of particle list
@type particleListName: str
@param output: Name of output particles
@type output: str
@param center: center of output particles in template orientation
@type center: list
@param cubesize: Size of output particles in pixel
@type cubesize: int
"""
from pytom.basic.structures import ParticleList, Particle, Shift
from pytom_volume import transformSpline as transform
from pytom_volume import subvolume, vol
pl = ParticleList()
pl.fromXMLFile(filename=particleListName)
#copy particle list for list of cropped particles
pl_new = pl.copy()
pvol = pl[0].getVolume()
sizeX = pvol.sizeX()
sizeY = pvol.sizeY()
sizeZ = pvol.sizeZ()
pvol_ali = vol(sizeX, sizeY, sizeZ)
subV = vol(cubesize, cubesize, cubesize)
sub_startX = center[0]-cubesize/2
sub_startY = center[1]-cubesize/2
sub_startZ = center[2]-cubesize/2
if (sub_startX < 0) or (sub_startY < 0) or (sub_startZ < 0):
raise ValueError('cubesize too large :(')
for (ipart, part) in enumerate(pl):
pvol_ali.setAll(0)
subV.setAll(0)
pvol = part.getVolume()
rot = part.getRotation()
rotinvert = rot.invert()
shiftV = part.getShift()
transform(pvol, pvol_ali, rotinvert[0], rotinvert[1], rotinvert[2],
sizeX/2, sizeY/2, sizeZ/2, -shiftV[0], -shiftV[1], -shiftV[2], 0, 0, 0)
# box out subvolume
subV = subvolume(pvol_ali, sub_startX, sub_startY, sub_startZ, cubesize, cubesize, cubesize)
transform(subV, subV, rot[0], rot[1], rot[2], cubesize/2, cubesize/2, cubesize/2, 0, 0, 0, 0, 0, 0)
fname = part.getFilename()
idx = fname.split('_')[-1].split('.')[0]
nfname = output+'_'+idx+'.em'
print("write file " + nfname)
subV.write(nfname)
pl_new[ipart].setFilename(newFilename=nfname)
pl_new[ipart].setShift(shift=Shift(0,0,0))
return pl_new
def sum_sub_pl(self, pl, name_prefix):
"""This is a sub-routine for average_sub_pl.
"""
from pytom_volume import vol
from pytom_volume import transformSpline as transform
from pytom.basic.normalise import mean0std1
result = None
wedgeSum = None
for p in pl:
particle = p.getVolume()
mean0std1(particle)
wedgeInfo = p.getWedge()
if result is None:
sizeX = particle.sizeX()
sizeY = particle.sizeY()
sizeZ = particle.sizeZ()
newParticle = vol(sizeX, sizeY, sizeZ)
# make consistent for python3
centerX = sizeX // 2
centerY = sizeY // 2
centerZ = sizeZ // 2
result = vol(sizeX, sizeY, sizeZ)
result.setAll(0)
wedgeSum = wedgeInfo.returnWedgeVolume(sizeX, sizeY, sizeZ)
wedgeSum.setAll(0)
# create wedge weighting
rotation = p.getRotation()
wedge = wedgeInfo.returnWedgeVolume(sizeX, sizeY, sizeZ, False,
rotation.invert())
wedgeSum = wedgeSum + wedge
# shift and rotate particle
shift = p.getShift()
newParticle.setAll(0)
transform(particle, newParticle, -rotation[1], -rotation[0],
-rotation[2], centerX, centerY, centerZ, -shift[0],
-shift[1], -shift[2], 0, 0, 0)
result = result + newParticle
# write them back to disk
result.write(name_prefix + '-PreWedge.em')
wedgeSum.write(name_prefix + '-WedgeSumUnscaled.em')
def rotate(volume,rotation,x=None,z2=None,imethod='spline',twice=False):
"""
rotate: Rotates a volume by a specified rotation
@param volume: The volume
@param rotation: The rotation (L{pytom.basic.structures.Rotation}) OR the z1 rotation.
@type rotation: Should be of type L{pytom.basic.structures.Rotation}. If not and a scalar is provided, you must provide x, z2 as the remaining two ZXZ rotation angles.
@param imethod: Interpolation method. Can be: linear, cubic, spline, fourierSpline
@param twice: Zero pad volume into a twice sized volume and perform calculation there.
@return: The rotated volume.
@author: Yuxiang Chen and Thomas Hrabe
"""
from pytom.basic.structures import Rotation
if rotation.__class__ == Rotation and x==None and z2==None:
z1 = rotation.getZ1()
x = rotation.getX()
z2 = rotation.getZ2()
elif (rotation.__class__ != float and rotation.__class__ != int) or (x.__class__ != float and x.__class__ != int) or (z2.__class__ != float and z2.__class__ != int):
raise TypeError('Rotation parameter must be a Rotation object or provide z1,x,z2 as floats!')
else:
z1 = rotation
if imethod == 'fourier':
return transformFourierSpline(volume,z1,z2,x,0,0,0,twice)
else:
from pytom_volume import vol
if imethod == 'linear':
from pytom_volume import transform
elif imethod == 'cubic':
from pytom_volume import transformCubic as transform
elif imethod == 'spline':
from pytom_volume import transformSpline as transform
centerX = volume.sizeX()/2
centerY = volume.sizeY()/2
centerZ = volume.sizeZ()/2
res = vol(volume.sizeX(),volume.sizeY(),volume.sizeZ())
transform(volume,res,z1,z2,x,centerX,centerY,centerZ,0,0,0,0,0,0)
return res
def shift(volume,shiftX,shiftY,shiftZ,imethod='fourier',twice=False):
"""
shift: Performs a shift on a volume
@param volume: the volume
@param shiftX: shift in x direction
@param shiftY: shift in y direction
@param shiftZ: shift in z direction
@param imethod: Select interpolation method. Real space : linear, cubic, spline . Fourier space: fourier
@param twice: Zero pad volume into a twice sized volume and perform calculation there.
@return: The shifted volume.
@author: Yuxiang Chen and Thomas Hrabe
"""
if imethod == 'fourier':
from pytom_volume import vol_comp,reducedToFull,fullToReduced,shiftFourier
from pytom.basic.fourier import fft,ifft,ftshift, iftshift
fvolume = fft(volume)
fullFVolume = reducedToFull(fvolume)
destFourier = vol_comp(fullFVolume.sizeX(),fullFVolume.sizeY(),fullFVolume.sizeZ())
shiftFourier(fullFVolume,destFourier,shiftX,shiftY,shiftZ)
resFourier = fullToReduced(destFourier)
return ifft(resFourier)/volume.numelem()
else:
from pytom_volume import vol
if imethod == 'linear':
from pytom_volume import transform
elif imethod == 'cubic':
from pytom_volume import transformCubic as transform
elif imethod == 'spline':
from pytom_volume import transformSpline as transform
# now results should be consistent with python2
centerX = int(volume.sizeX()/2)
centerY = int(volume.sizeY()/2)
centerZ = int(volume.sizeZ()/2)
res = vol(volume.sizeX(),volume.sizeY(),volume.sizeZ())
transform(volume,res,0,0,0,centerX,centerY,centerZ,shiftX,shiftY,shiftZ,0,0,0)
return res
def _rotateWedgeReference(reference,rotation,wedgeInfo,mask,rotationCenter):
"""
_rotateShiftWedgeParticle: Wrapper for Rotation, Shift and WedgeWeighting of Reference
@param reference: The reference
@param rotation: The rotation
@type rotation: L{pytom.basic.structures.Rotation}
@param wedgeInfo: Wedge info object
@type wedgeInfo: L{pytom.basic.structures.Wedge}
@param mask: The mask object (a volume) or None
@type mask: L{pytom_volume.vol}
@return:
@change: support mask == None, FF
"""
from pytom_volume import vol, transform as transform #developers: your can also import transformSpline for more accurate rotation!
rotatedVolume = vol(reference.sizeX(),reference.sizeY(),reference.sizeZ())
transform(reference,rotatedVolume,rotation[0],rotation[1],rotation[2],rotationCenter[0],rotationCenter[1],rotationCenter[2],0,0,0,0,0,0)
if mask:
return wedgeInfo.apply(rotatedVolume) * mask
else:
return wedgeInfo.apply(rotatedVolume)
def paverage(particleList, norm, binning, verbose, outdir='./'):
from pytom_volume import read, vol
from pytom_volume import transformSpline as transform
from pytom.basic.structures import Particle
from pytom.basic.normalise import mean0std1
from pytom.tools.ProgressBar import FixedProgBar
from pytom.basic.transformations import resize
if len(particleList) == 0:
raise RuntimeError('The particlelist provided is empty. Aborting!')
if verbose:
progressBar = FixedProgBar(0, len(particleList), 'Particles averaged ')
progressBar.update(0)
numberAlignedParticles = 0
result = None
wedgeSum = None
newParticle = None
for particleObject in particleList:
particle = read(particleObject.getFilename(), 0, 0, 0, 0, 0, 0, 0, 0,
0, 1, 1, 1)
if binning != 1:
particle, particlef = resize(volume=particle,
factor=1. / binning,
interpolation='Fourier')
if norm:
mean0std1(particle)
wedgeInfo = particleObject.getWedge()
if result is None: # initialization
sizeX = particle.sizeX()
sizeY = particle.sizeY()
sizeZ = particle.sizeZ()
newParticle = vol(sizeX, sizeY, sizeZ)
centerX = sizeX // 2
centerY = sizeY // 2
centerZ = sizeZ // 2
result = vol(sizeX, sizeY, sizeZ)
result.setAll(0.0)
wedgeSum = wedgeInfo.returnWedgeVolume(sizeX, sizeY, sizeZ)
wedgeSum.setAll(0)
# create spectral wedge weighting
rotation = particleObject.getRotation()
wedge = wedgeInfo.returnWedgeVolume(sizeX, sizeY, sizeZ, False,
rotation.invert())
wedgeSum += wedge
# shift and rotate particle
shiftV = particleObject.getShift()
newParticle.setAll(0)
transform(particle, newParticle, -rotation[1], -rotation[0],
-rotation[2], centerX, centerY, centerZ,
-shiftV[0] // binning, -shiftV[1] // binning,
-shiftV[2] // binning, 0, 0, 0)
result += newParticle
if verbose:
numberAlignedParticles = numberAlignedParticles + 1
progressBar.update(numberAlignedParticles)
# write to the disk
fname_result = os.path.join(outdir, 'avg_{}.em'.format(mpi.rank))
fname_wedge = os.path.join(outdir, 'wedge_{}.em'.format(mpi.rank))
result.write(fname_result)
result = Particle(fname_result)
wedgeSum.write(fname_wedge)
wedgeSum = Particle(fname_wedge)
return (result, wedgeSum)
def average2(particleList, weighting=False, norm=False, determine_resolution=False,
mask=None, binning=1, verbose=False):
"""
2nd version of average function. Will not write the averages to the disk. Also support internal \
resolution determination.
"""
from pytom_volume import read, vol, complexDiv, complexRealMult
from pytom_volume import transformSpline as transform
from pytom.basic.fourier import fft, ifft, convolute
from pytom.basic.normalise import mean0std1
from pytom.tools.ProgressBar import FixedProgBar
from pytom.basic.filter import lowpassFilter, rotateWeighting
from math import exp
if len(particleList) == 0:
raise RuntimeError('The particlelist provided is empty. Aborting!')
if verbose:
progressBar = FixedProgBar(0,len(particleList),'Particles averaged ')
progressBar.update(0)
numberAlignedParticles = 0
even = None
odd = None
wedgeSum_even = None
wedgeSum_odd = None
newParticle = None
is_odd = True
for particleObject in particleList:
particle = read(particleObject.getFilename(), 0,0,0,0,0,0,0,0,0, binning,binning,binning)
if norm:
mean0std1(particle)
wedgeInfo = particleObject.getWedge()
# apply its wedge to itself
particle = wedgeInfo.apply(particle)
if odd is None: # initialization
sizeX = particle.sizeX()
sizeY = particle.sizeY()
sizeZ = particle.sizeZ()
newParticle = vol(sizeX,sizeY,sizeZ)
centerX = sizeX/2
centerY = sizeY/2
centerZ = sizeZ/2
odd = vol(sizeX,sizeY,sizeZ)
odd.setAll(0.0)
even = vol(sizeX,sizeY,sizeZ)
even.setAll(0.0)
wedgeSum_odd = wedgeInfo.returnWedgeVolume(sizeX,sizeY,sizeZ)
wedgeSum_odd.setAll(0)
wedgeSum_even = wedgeInfo.returnWedgeVolume(sizeX,sizeY,sizeZ)
wedgeSum_even.setAll(0)
# create spectral wedge weighting
rotation = particleObject.getRotation()
rotinvert = rotation.invert()
if analytWedge:
# > original buggy version
wedge = wedgeInfo.returnWedgeVolume(sizeX,sizeY,sizeZ,False, rotinvert)
# < original buggy version
else:
# > FF: interpol bugfix
wedge = rotateWeighting( weighting=wedgeInfo.returnWedgeVolume(sizeX,sizeY,sizeZ,False),
z1=rotinvert[0], z2=rotinvert[1], x=rotinvert[2], mask=None,
isReducedComplex=True, returnReducedComplex=True)
# < FF
# > TH bugfix
#wedgeVolume = wedgeInfo.returnWedgeVolume(wedgeSizeX=sizeX, wedgeSizeY=sizeY, wedgeSizeZ=sizeZ,
# humanUnderstandable=True, rotation=rotinvert)
#wedge = rotate(volume=wedgeVolume, rotation=rotinvert, imethod='linear')
# < TH
if is_odd:
wedgeSum_odd = wedgeSum_odd + wedge
else:
wedgeSum_even = wedgeSum_even + wedge
# shift and rotate particle
shiftV = particleObject.getShift()
newParticle.setAll(0)
transform(particle,newParticle,-rotation[1],-rotation[0],-rotation[2],
centerX,centerY,centerZ,-shiftV[0]/binning,
-shiftV[1]/binning,-shiftV[2]/binning,0,0,0)
if is_odd:
if weighting:
weight = 1. - particleObject.getScore().getValue()
#weight = weight**2
weight = exp(-1.*weight)
odd = odd + newParticle * weight
else:
odd = odd + newParticle
else:
if weighting:
weight = 1. - particleObject.getScore().getValue()
#weight = weight**2
weight = exp(-1.*weight)
even = even + newParticle * weight
else:
even = even + newParticle
is_odd = not is_odd
if verbose:
numberAlignedParticles = numberAlignedParticles + 1
progressBar.update(numberAlignedParticles)
# determine resolution if needed
fsc = None
if determine_resolution:
# apply spectral weighting to sum
f_even = fft(even)
w_even = complexDiv(f_even, wedgeSum_even)
w_even = ifft(w_even)
w_even.shiftscale(0.0,1/float(sizeX*sizeY*sizeZ))
f_odd = fft(odd)
w_odd = complexDiv(f_odd, wedgeSum_odd)
w_odd = ifft(w_odd)
w_odd.shiftscale(0.0,1/float(sizeX*sizeY*sizeZ))
from pytom.basic.correlation import FSC
fsc = FSC(w_even, w_odd, sizeX/2, mask, verbose=False)
# add together
result = even+odd
wedgeSum = wedgeSum_even+wedgeSum_odd
invert_WedgeSum( invol=wedgeSum, r_max=sizeX/2-2., lowlimit=.05*len(particleList), lowval=.05*len(particleList))
#wedgeSum.write(averageName[:len(averageName)-3] + '-WedgeSumInverted.em')
result = convolute(v=result, k=wedgeSum, kernel_in_fourier=True)
# do a low pass filter
#result = lowpassFilter(result, sizeX/2-2, (sizeX/2-1)/10.)[0]
return (result, fsc)
def average( particleList, averageName, showProgressBar=False, verbose=False,
createInfoVolumes=False, weighting=False, norm=False):
"""
average : Creates new average from a particleList
@param particleList: The particles
@param averageName: Filename of new average
@param verbose: Prints particle information. Disabled by default.
@param createInfoVolumes: Create info data (wedge sum, inverted density) too? False by default.
@param weighting: apply weighting to each average according to its correlation score
@param norm: apply normalization for each particle
@return: A new Reference object
@rtype: L{pytom.basic.structures.Reference}
@author: Thomas Hrabe
@change: limit for wedgeSum set to 1% or particles to avoid division by small numbers - FF
"""
from pytom_volume import read,vol,reducedToFull,limit, complexRealMult
from pytom.basic.filter import lowpassFilter, rotateWeighting
from pytom_volume import transformSpline as transform
from pytom.basic.fourier import convolute
from pytom.basic.structures import Reference
from pytom.basic.normalise import mean0std1
from pytom.tools.ProgressBar import FixedProgBar
from math import exp
import os
if len(particleList) == 0:
raise RuntimeError('The particle list is empty. Aborting!')
if showProgressBar:
progressBar = FixedProgBar(0,len(particleList),'Particles averaged ')
progressBar.update(0)
numberAlignedParticles = 0
result = []
wedgeSum = []
newParticle = None
# pre-check that scores != 0
if weighting:
wsum = 0.
for particleObject in particleList:
wsum += particleObject.getScore().getValue()
if wsum < 0.00001:
weighting = False
print("Warning: all scores have been zero - weighting not applied")
for particleObject in particleList:
if verbose:
print(particleObject)
if not os.path.exists(particleObject.getFilename()): continue
particle = read(particleObject.getFilename())
if norm: # normalize the particle
mean0std1(particle) # happen inplace
wedgeInfo = particleObject.getWedge()
# apply its wedge to itself
particle = wedgeInfo.apply(particle)
if result == []:
sizeX = particle.sizeX()
sizeY = particle.sizeY()
sizeZ = particle.sizeZ()
newParticle = vol(sizeX,sizeY,sizeZ)
centerX = sizeX/2
centerY = sizeY/2
centerZ = sizeZ/2
result = vol(sizeX,sizeY,sizeZ)
result.setAll(0.0)
if analytWedge:
wedgeSum = wedgeInfo.returnWedgeVolume(wedgeSizeX=sizeX, wedgeSizeY=sizeY, wedgeSizeZ=sizeZ)
else:
# > FF bugfix
wedgeSum = wedgeInfo.returnWedgeVolume(sizeX,sizeY,sizeZ)
# < FF
# > TH bugfix
#wedgeSum = vol(sizeX,sizeY,sizeZ)
# < TH
#wedgeSum.setAll(0)
assert wedgeSum.sizeX() == sizeX and wedgeSum.sizeY() == sizeY and wedgeSum.sizeZ() == sizeZ/2+1, \
"wedge initialization result in wrong dims :("
wedgeSum.setAll(0)
### create spectral wedge weighting
rotation = particleObject.getRotation()
rotinvert = rotation.invert()
if analytWedge:
# > analytical buggy version
wedge = wedgeInfo.returnWedgeVolume(sizeX,sizeY,sizeZ,False, rotinvert)
else:
# > FF: interpol bugfix
wedge = rotateWeighting( weighting=wedgeInfo.returnWedgeVolume(sizeX,sizeY,sizeZ,False),
z1=rotinvert[0], z2=rotinvert[1], x=rotinvert[2], mask=None,
isReducedComplex=True, returnReducedComplex=True)
# < FF
# > TH bugfix
#wedgeVolume = wedgeInfo.returnWedgeVolume(wedgeSizeX=sizeX, wedgeSizeY=sizeY, wedgeSizeZ=sizeZ,
# humanUnderstandable=True, rotation=rotinvert)
#wedge = rotate(volume=wedgeVolume, rotation=rotinvert, imethod='linear')
# < TH
### shift and rotate particle
shiftV = particleObject.getShift()
newParticle.setAll(0)
transform(particle,newParticle,-rotation[1],-rotation[0],-rotation[2],
centerX,centerY,centerZ,-shiftV[0],-shiftV[1],-shiftV[2],0,0,0)
if weighting:
weight = 1.-particleObject.getScore().getValue()
#weight = weight**2
weight = exp(-1.*weight)
result = result + newParticle * weight
wedgeSum = wedgeSum + wedge * weight
else:
result = result + newParticle
wedgeSum = wedgeSum + wedge
if showProgressBar:
numberAlignedParticles = numberAlignedParticles + 1
progressBar.update(numberAlignedParticles)
###apply spectral weighting to sum
result = lowpassFilter(result, sizeX/2-1, 0.)[0]
#if createInfoVolumes:
result.write(averageName[:len(averageName)-3]+'-PreWedge.em')
wedgeSum.write(averageName[:len(averageName)-3] + '-WedgeSumUnscaled.em')
invert_WedgeSum( invol=wedgeSum, r_max=sizeX/2-2., lowlimit=.05*len(particleList), lowval=.05*len(particleList))
if createInfoVolumes:
wedgeSum.write(averageName[:len(averageName)-3] + '-WedgeSumInverted.em')
result = convolute(v=result, k=wedgeSum, kernel_in_fourier=True)
# do a low pass filter
#result = lowpassFilter(result, sizeX/2-2, (sizeX/2-1)/10.)[0]
result.write(averageName)
if createInfoVolumes:
resultINV = result * -1
#write sign inverted result to disk (good for chimera viewing ... )
resultINV.write(averageName[:len(averageName)-3]+'-INV.em')
newReference = Reference(averageName,particleList)
return newReference
def averageGPU(particleList,
averageName,
showProgressBar=False,
verbose=False,
createInfoVolumes=False,
weighting=False,
norm=False,
gpuId=None,
profile=True):
"""
average : Creates new average from a particleList
@param particleList: The particles
@param averageName: Filename of new average
@param verbose: Prints particle information. Disabled by default.
@param createInfoVolumes: Create info data (wedge sum, inverted density) too? False by default.
@param weighting: apply weighting to each average according to its correlation score
@param norm: apply normalization for each particle
@return: A new Reference object
@rtype: L{pytom.basic.structures.Reference}
@author: Thomas Hrabe
@change: limit for wedgeSum set to 1% or particles to avoid division by small numbers - FF
"""
import time
from pytom.tompy.io import read, write, read_size
from pytom.tompy.filter import bandpass as lowpassFilter, rotateWeighting, applyFourierFilter, applyFourierFilterFull, create_wedge
from pytom.voltools import transform, StaticVolume
from pytom.basic.structures import Reference
from pytom.tompy.normalise import mean0std1
from pytom.tompy.tools import volumesSameSize, invert_WedgeSum, create_sphere
from pytom.tompy.transform import fourier_full2reduced, fourier_reduced2full
from cupyx.scipy.fftpack.fft import fftn as fftnP
from cupyx.scipy.fftpack.fft import ifftn as ifftnP
from cupyx.scipy.fftpack.fft import get_fft_plan
from pytom.tools.ProgressBar import FixedProgBar
from multiprocessing import RawArray
import numpy as np
import cupy as xp
if not gpuId is None:
device = f'gpu:{gpuId}'
xp.cuda.Device(gpuId).use()
else:
print(gpuId)
raise Exception('Running gpu code on non-gpu device')
print(device)
cstream = xp.cuda.Stream()
if profile:
stream = xp.cuda.Stream.null
t_start = stream.record()
# from pytom.tools.ProgressBar import FixedProgBar
from math import exp
import os
if len(particleList) == 0:
raise RuntimeError('The particle list is empty. Aborting!')
if showProgressBar:
progressBar = FixedProgBar(0, len(particleList), 'Particles averaged ')
progressBar.update(0)
numberAlignedParticles = 0
# pre-check that scores != 0
if weighting:
wsum = 0.
for particleObject in particleList:
wsum += particleObject.getScore().getValue()
if wsum < 0.00001:
weighting = False
print("Warning: all scores have been zero - weighting not applied")
import time
sx, sy, sz = read_size(particleList[0].getFilename())
wedgeInfo = particleList[0].getWedge().convert2numpy()
print('angle: ', wedgeInfo.getWedgeAngle())
wedgeZero = xp.fft.fftshift(
xp.array(wedgeInfo.returnWedgeVolume(sx, sy, sz, True).get(),
dtype=xp.float32))
# wedgeZeroReduced = fourier_full2reduced(wedgeZero)
wedge = xp.zeros_like(wedgeZero, dtype=xp.float32)
wedgeSum = xp.zeros_like(wedge, dtype=xp.float32)
print('init texture')
wedgeText = StaticVolume(xp.fft.fftshift(wedgeZero),
device=device,
interpolation='filt_bspline')
newParticle = xp.zeros((sx, sy, sz), dtype=xp.float32)
centerX = sx // 2
centerY = sy // 2
centerZ = sz // 2
result = xp.zeros((sx, sy, sz), dtype=xp.float32)
fftplan = get_fft_plan(wedge.astype(xp.complex64))
n = 0
total = len(particleList)
# total = int(np.floor((11*1024**3 - mempool.total_bytes())/(sx*sy*sz*4)))
# total = 128
#
#
# particlesNP = np.zeros((total, sx, sy, sz),dtype=np.float32)
# particles = []
# mask = create_sphere([sx,sy,sz], sx//2-6, 2)
# raw = RawArray('f', int(particlesNP.size))
# shared_array = np.ctypeslib.as_array(raw)
# shared_array[:] = particlesNP.flatten()
# procs = allocateProcess(particleList, shared_array, n, total, wedgeZero.size)
# del particlesNP
if profile:
t_end = stream.record()
t_end.synchronize()
time_took = xp.cuda.get_elapsed_time(t_start, t_end)
print(f'startup time {n:5d}: \t{time_took:.3f}ms')
t_start = stream.record()
for particleObject in particleList:
rotation = particleObject.getRotation()
rotinvert = rotation.invert()
shiftV = particleObject.getShift()
# if n % total == 0:
# while len(procs):
# procs =[proc for proc in procs if proc.is_alive()]
# time.sleep(0.1)
# print(0.1)
# # del particles
# # xp._default_memory_pool.free_all_blocks()
# # pinned_mempool.free_all_blocks()
# particles = xp.array(shared_array.reshape(total, sx, sy, sz), dtype=xp.float32)
# procs = allocateProcess(particleList, shared_array, n, total, size=wedgeZero.size)
# #pinned_mempool.free_all_blocks()
# #print(mempool.total_bytes()/1024**3)
particle = read(particleObject.getFilename(), deviceID=device)
#particle = particles[n%total]
if norm: # normalize the particle
mean0std1(particle) # happen inplace
# apply its wedge to
#particle = applyFourierFilter(particle, wedgeZeroReduced)
#particle = (xp.fft.ifftn( xp.fft.fftn(particle) * wedgeZero)).real
particle = (ifftnP(fftnP(particle, plan=fftplan) * wedgeZero,
plan=fftplan)).real
### create spectral wedge weighting
wedge *= 0
wedgeText.transform(
rotation=[rotinvert[0], rotinvert[2], rotinvert[1]],
rotation_order='rzxz',
output=wedge)
#wedge = xp.fft.fftshift(fourier_reduced2full(create_wedge(30, 30, 21, 42, 42, 42, rotation=[rotinvert[0],rotinvert[2], rotinvert[1]])))
# if analytWedge:
# # > analytical buggy version
# wedge = wedgeInfo.returnWedgeVolume(sx, sy, sz, True, rotinvert)
# else:
# # > FF: interpol bugfix
# wedge = rotateWeighting(weighting=wedgeInfo.returnWedgeVolume(sx, sy, sz, True), rotation=[rotinvert[0], rotinvert[2], rotinvert[1]])
# # < FF
# # > TH bugfix
# # wedgeVolume = wedgeInfo.returnWedgeVolume(wedgeSizeX=sizeX, wedgeSizeY=sizeY, wedgeSizeZ=sizeZ,
# # humanUnderstandable=True, rotation=rotinvert)
# # wedge = rotate(volume=wedgeVolume, rotation=rotinvert, imethod='linear')
# # < TH
### shift and rotate particle
newParticle *= 0
transform(particle,
output=newParticle,
rotation=[-rotation[1], -rotation[2], -rotation[0]],
center=[centerX, centerY, centerZ],
translation=[-shiftV[0], -shiftV[1], -shiftV[2]],
device=device,
interpolation='filt_bspline',
rotation_order='rzxz')
#write(f'trash/GPU_{n}.em', newParticle)
# print(rotation.toVector())
# break
result += newParticle
wedgeSum += xp.fft.fftshift(wedge)
# if showProgressBar:
# numberAlignedParticles = numberAlignedParticles + 1
# progressBar.update(numberAlignedParticles)
if n % total == 0:
if profile:
t_end = stream.record()
t_end.synchronize()
time_took = xp.cuda.get_elapsed_time(t_start, t_end)
print(f'total time {n:5d}: \t{time_took:.3f}ms')
t_start = stream.record()
cstream.synchronize()
n += 1
print('averaged particles')
###apply spectral weighting to sum
result = lowpassFilter(result, high=sx / 2 - 1, sigma=0)
# if createInfoVolumes:
write(averageName[:len(averageName) - 3] + '-PreWedge.em', result)
write(averageName[:len(averageName) - 3] + '-WedgeSumUnscaled.em',
fourier_full2reduced(wedgeSum))
wedgeSumINV = invert_WedgeSum(wedgeSum,
r_max=sx // 2 - 2.,
lowlimit=.05 * len(particleList),
lowval=.05 * len(particleList))
wedgeSumINV = wedgeSumINV
#print(wedgeSum.mean(), wedgeSum.std())
if createInfoVolumes:
write(averageName[:len(averageName) - 3] + '-WedgeSumInverted.em',
xp.fft.fftshift(wedgeSumINV))
result = applyFourierFilterFull(result, xp.fft.fftshift(wedgeSumINV))
# do a low pass filter
result = lowpassFilter(result, sx / 2 - 2, (sx / 2 - 1) / 10.)[0]
write(averageName, result)
if createInfoVolumes:
resultINV = result * -1
# write sign inverted result to disk (good for chimera viewing ... )
write(averageName[:len(averageName) - 3] + '-INV.em', resultINV)
newReference = Reference(averageName, particleList)
return newReference