def subPixelShiftsImage(volume, scale=[1, 1, 1], cutout=20, shifts=True):
from pytom.voltools import transform
import numpy as np
if shifts:
volC = xp.array(
[vol.shape[0] // 2, volume.shape[1] // 2, vol.shape[2] // 2])
else:
volC = xp.array([0, 0, 0], dtype=xp.int32)
cx, cy, cz = center = find_coords_max_ccmap(volume)
cutvol = volume[cx - cutout // 2:cx + cutout // 2,
cy - cutout // 2:cy + cutout // 2, :]
nx, ny, nz = [int(np.ceil(cutout / s)) for s in scale]
nC = xp.array([nx // 2, ny // 2, 0])
centered = xp.zeros((nx, ny, 1), dtype=xp.float32)
centered[nx // 2 - cutout // 2:nx // 2 + cutout // 2,
ny // 2 - cutout // 2:ny // 2 + cutout // 2, :] = cutvol
interpolated_peak = transform(centered,
scale=scale,
device=device,
interpolation='filt_bspline')
interpol_center = find_coords_max_ccmap(interpolated_peak)
shifts = xp.array(center) - volC + (xp.array(interpol_center) -
xp.array(nC)) * xp.array(scale)
return [interpolated_peak.max(), shifts]
def subPixelShifts(volume, scale=[1, 1, 1], cutout=20, shifts=True):
from pytom.voltools import transform
import numpy as np
volC = xp.array(
[volume.shape[0] // 2, volume.shape[1] // 2, volume.shape[2] // 2])
cx, cy, cz = center = find_coords_max_ccmap(volume)
x, y, z = volume.shape
if cx - cutout // 2 < 0 or cy - cutout // 2 < 0 or cz - cutout // 2 < 0 or cx + cutout // 2 >= x or cy + cutout // 2 >= y or cz + cutout // 2 >= z:
return [
volume.max(),
xp.array([cx - x // 2, cy - y // 2, cz - z // 2])
]
cutvol = volume[cx - cutout // 2:cx + cutout // 2,
cy - cutout // 2:cy + cutout // 2,
cz - cutout // 2:cz + cutout // 2]
nx, ny, nz = [int(np.ceil(cutout / s)) for s in scale]
nC = xp.array([nx // 2, ny // 2, nz // 2])
centered = xp.zeros((nx, ny, nz), dtype=xp.float32)
centered[nx // 2 - cutout // 2:nx // 2 + cutout // 2,
ny // 2 - cutout // 2:ny // 2 + cutout // 2,
nz // 2 - cutout // 2:nz // 2 + cutout // 2] = cutvol
interpolated_peak = xp.zeros_like(centered)
transform(centered,
scale=scale,
device=device,
interpolation='filt_bspline',
output=interpolated_peak)
interpol_center = find_coords_max_ccmap(interpolated_peak)
shifts = xp.array(center) - volC + (xp.array(interpol_center) -
xp.array(nC)) * xp.array(scale)
return [interpolated_peak.max(), shifts]
def create_sphere(size,
radius=-1,
sigma=0,
num_sigma=2,
center=None,
gpu=False):
"""Create a 3D sphere volume.
@param size: size of the resulting volume.
@param radius: radius of the sphere inside the volume.
@param sigma: sigma of the Gaussian.
@param center: center of the sphere.
@return: sphere inside a volume.
"""
if size.__class__ == float or len(size) == 1:
size = (size, size, size)
assert len(size) == 3
if center is None:
center = [size[0] // 2, size[1] // 2, size[2] // 2]
if radius == -1:
radius = xp.min(size) // 2
sphere = xp.zeros(size, dtype=xp.float32)
[x, y, z] = xp.mgrid[0:size[0], 0:size[1], 0:size[2]]
r = xp.sqrt((x - center[0])**2 + (y - center[1])**2 + (z - center[2])**2)
sphere[r <= radius] = 1
if sigma > 0:
ind = xp.logical_and(r > radius, r < radius + num_sigma * sigma)
sphere[ind] = xp.exp(-((r[ind] - radius) / sigma)**2 / 2)
return sphere
def maxIndex(volume, num_threads=1024):
nblocks = int(xp.ceil(volume.size / num_threads / 2))
fast_sum = -1000000 * xp.ones((nblocks), dtype=xp.float32)
max_id = xp.zeros((nblocks), dtype=xp.int32)
argmax((
nblocks,
1,
), (num_threads, 1, 1), (volume, fast_sum, max_id, volume.size),
shared_mem=16 * num_threads)
mm = min(max_id[fast_sum.argmax()], volume.size - 1)
indices = xp.unravel_index(mm, volume.shape)
return indices
def prepare_mask(v, threshold, smooth):
"""Prepare a mask according to the given volume.
Everything above the given threshold will be set to 1.
@param v: input volume.
@param threshold: threshold.
@param smooth: sigma of Gaussian.
@return: mask.
"""
from pytom.tompy.filter import gaussian3d
ind = xp.where(v > threshold)
mask = xp.zeros(v.shape)
mask[ind] = 1
return gaussian3d(mask, smooth)
def circle_filter(sizeX, sizeY, radiusCutoff):
"""
circleFilter: NEEDS Documentation
@param sizeX: NEEDS Documentation
@param sizeY: NEEDS Documentation
@param radiusCutoff: NEEDS Documentation
"""
X, Y = xp.meshgrid(
xp.arange(-sizeX // 2 + sizeX % 2, sizeX // 2 + sizeX % 2),
xp.arange(-sizeY // 2 + sizeY % 2, sizeY // 2 + sizeY % 2))
R = xp.sqrt(X**2 + Y**2)
filter = xp.zeros((sizeX, sizeY), dtype=xp.float32)
filter[R <= radiusCutoff] = 1
return filter
def rotation_matrix_z(angle):
"""Return the 3x3 rotation matrix around z axis.
@param angle: rotation angle around z axis (in degree).
@return: rotation matrix.
"""
angle = xp.deg2rad(angle)
mtx = xp.matrix(xp.zeros((3, 3)))
mtx[0, 0] = xp.cos(angle)
mtx[1, 0] = xp.sin(angle)
mtx[1, 1] = xp.cos(angle)
mtx[0, 1] = -xp.sin(angle)
mtx[2, 2] = 1
return mtx
def ellipse_filter(sizeX, sizeY, radiusCutoffX, radiusCutoffY):
"""
circleFilter: NEEDS Documentation
@param sizeX: NEEDS Documentation
@param sizeY: NEEDS Documentation
@param radiusCutoff: NEEDS Documentation
"""
X, Y = xp.meshgrid(
xp.arange(-sizeY // 2 + sizeY % 2, sizeY // 2 + sizeY % 2),
xp.arange(-sizeX // 2 + sizeX % 2, sizeX // 2 + sizeX % 2))
R = xp.sqrt((X / radiusCutoffX)**2 + (Y / radiusCutoffY)**2)
filter = xp.zeros((sizeX, sizeY), dtype=xp.float32)
#print(filter.shape, R.shape)
filter[R <= 1] = 1
return filter
def find_sub_pixel_voltools(inp, k=0.1, border=75, max_shift=15):
import pytom.voltools as vt
inp = xp.ascontiguousarray(inp)
# find the position of the initial maximum
dimx, dimy = inp.squeeze().shape
initial_max = xp.unravel_index(
inp[border:-border, border:-border].argmax(),
(dimx - border * 2, dimy - border * 2))
ix, iy = [v + border for v in initial_max]
# Create an array with specific size so the zoom fits in (size = max_shift * 2 /k)
# Center of array is initial max
out = int(xp.around(max_shift * 2 / k))
model = xp.zeros((out, out), dtype=xp.float32)
model[out // 2 - max_shift:out // 2 + max_shift, out // 2 -
max_shift:out // 2 + max_shift] = inp[ix - max_shift:ix + max_shift,
iy - max_shift:iy + max_shift]
zoomedVol = xp.expand_dims(model, 2)
# Scale image
vt.transform(zoomedVol,
scale=(k, k, 1),
interpolation='filt_bspline',
device=device,
output=zoomedVol)
# fig, ax = subplots(1, 2, figsize=(10, 5))
# ax[0].imshow(inp.get().squeeze())
# ax[1].imshow(zoomedVol.squeeze().get())
# show()
# Find new max and update the initial max according to the shift
transformed_max = xp.unravel_index(zoomedVol.argmax(), zoomedVol.shape)
interpolX, interpolY = ix + (transformed_max[0] - out // 2) * k, iy + (
transformed_max[1] - out // 2) * k
return interpolX, interpolY, zoomedVol.squeeze()
def fourier_reduced2full(data, isodd=False):
"""Return an Hermitian symmetried data.
"""
# get the original shape
sx = data.shape[0]
sy = data.shape[1]
if isodd:
sz = (data.shape[2] - 1) * 2 + 1
else:
sz = (data.shape[2] - 1) * 2
res = xp.zeros((sx, sy, sz), dtype=data.dtype)
res[:, :, 0:data.shape[2]] = data
# calculate the coodinates accordingly
szz = sz - data.shape[2]
x, y, z = xp.indices((sx, sy, szz))
ind = [xp.mod(sx - x, sx), xp.mod(sy - y, sy), szz - z]
# do the complex conjugate of the second part
res[:, :, data.shape[2]:] = xp.conj(data[tuple(ind)])
return res
def resizeFourier(fvol, factor, isodd=False):
"""
resize Fourier transformed by factor
@param fvol: Fourier transformed of a volume - reduced complex
@type fvol: L{pytom_volume.vol_comp}
@param factor: a factor > 0. Factors < 1 will de-magnify the volume, factors > 1 will magnify.
@type factor: float
@return: resized Fourier volume (deruced complex)
@rtype: L{pytom_volume.vol_comp}
@author: FF
"""
fvol = fvol.squeeze()
if factor == 1: return fvol
oldFNx = fvol.shape[0]
oldFNy = fvol.shape[1]
# new dims in real and Fourier space
newNx = newFNx = int(xp.floor(oldFNx * factor + 0.5))
newNy = newFNy = int(xp.floor(oldFNy * factor + 0.5))
# check 3D images
if len(fvol.shape) == 3:
oldNz = int((fvol.shape[2] - 1) * 2 + 1 * isodd)
newNz = int(xp.floor(oldNz * factor + 0.5))
newFNz = newNz // 2 + 1
oldFNz = fvol.shape[2]
scf = 1 #oldNz **3 / (newNx * newNy * newNz)
newxIsEven = newFNx % 2
newyIsEven = newFNy % 2
fvol_center_scaled = xp.fft.fftshift(fvol, axes=(0, 1)) * scf
newfvol = xp.zeros((newFNx, newFNy, newFNz), dtype=fvol.dtype)
if factor >= 1:
# Effectively zero-padding
newfvol[newFNx // 2 - oldFNx // 2 + newxIsEven:newFNx // 2 +
oldFNx // 2 + isodd + newxIsEven, newFNy // 2 -
oldFNy // 2 + newyIsEven:newFNy // 2 + oldFNy // 2 +
isodd + newyIsEven, :oldFNz] = fvol_center_scaled
# newFNz // 2 - oldFNz // 2 :newFNz // 2 + oldFNz // 2 + 1] = fvol_center_scaled
else:
# Effectively cropping
newfvol = fvol_center_scaled[oldFNx // 2 - newFNx // 2 -
newxIsEven:oldFNx // 2 + newFNx // 2 +
isodd, oldFNy // 2 - newFNy // 2 -
newyIsEven:oldFNy // 2 + newFNy // 2 +
isodd, :newFNz]
newfvol = xp.fft.fftshift(newfvol, axes=(0, 1))
else:
newNz = 1
scf = 1. / (newNx * newNy * newNz)
newFNz = 1
newFNy = newNy #// 2 + 1
fvol_center_scaled = xp.fft.fftshift(fvol, axes=(0))
newfvol = xp.zeros((newFNx, newFNy), dtype=fvol.dtype) * scf
if factor >= 1:
newfvol[newFNx // 2 - oldFNx // 2:newFNx // 2 + oldFNx // 2 +
isodd, :oldFNx] = fvol_center_scaled
else:
newfvol = fvol_center_scaled[oldFNx // 2 -
newFNx // 2:oldFNx // 2 +
newFNx // 2 + isodd, :newFNy]
newfvol = xp.fft.fftshift(newfvol, axes=(0))
#newfvol = xp.expand_dims(newfvol,2)
return newfvol
def alignVolumesAndFilterByFSC(vol1,
vol2,
mask=None,
nband=None,
iniRot=None,
iniTrans=None,
interpolation='linear',
fsc_criterion=0.143,
verbose=0):
"""
align two volumes, compute their FSC, and filter by FSC
@param vol1: volume 1
@param vol2: volume 2
@mask: mask volume
@type mask: L{pytom_volume.vol}
@param nband: Number of bands
@type nband: L{int}
@param iniRot: initial guess for rotation
@param iniTrans: initial guess for translation
@param interpolation: interpolation type - 'linear' (default) or 'spline'
@param fsc_criterion: filter -> 0 according to resolution criterion
@type fsc_criterion: float
@param verbose: verbose level (0=mute, 1 some output, 2=talkative)
@type verbose: int
@type interpolation: str
@return: (filvol1, filvol2, fsc, fsc_fil, optiRot, optiTrans) i.e., filtered volumes, their FSC, the corresponding\
filter that was applied to the volumes, and the optimal rotation and translation of vol2 with respect to vol1\
note: filvol2 is NOT rotated and translated!
@author: FF
"""
from pytom_volume import transformSpline, vol
from pytom.tompy.correlation import FSC
from pytom.tompy.filter import filter_volume_by_profile
from pytom.tompy.structures import Alignment
from pytom.tompy.correlation import nxcc
from pytom.voltools import transform
assert isinstance(object=vol1, class_or_type_or_tuple=vol
), "alignVolumesAndFilterByFSC: vol1 must be of type vol"
assert isinstance(object=vol2, class_or_type_or_tuple=vol
), "alignVolumesAndFilterByFSC: vol2 must be of type vol"
# filter volumes prior to alignment according to SNR
fsc = FSC(volume1=vol1, volume2=vol2, numberBands=nband)
fil = design_fsc_filter(fsc=fsc, fildim=int(vol2.shape[2] // 2))
#filter only one volume so that resulting CCC is weighted by SNR only once
filvol2 = filter_volume_by_profile(volume=vol2, profile=fil)
# align vol2 to vol1
if verbose == 2:
alignment = Alignment(vol1=vol1,
vol2=filvol2,
score=nxcc,
mask=mask,
iniRot=iniRot,
iniTrans=iniTrans,
opti='fmin_powell',
interpolation=interpolation,
verbose=verbose)
else:
alignment = Alignment(vol1=vol1,
vol2=filvol2,
score=nxcc,
mask=mask,
iniRot=iniRot,
iniTrans=iniTrans,
opti='fmin_powell',
interpolation=interpolation,
verbose=False)
optiScore, optiRot, optiTrans = alignment.localOpti(iniRot=iniRot,
iniTrans=iniTrans)
if verbose:
from pytom.angles.angleFnc import differenceAngleOfTwoRotations
from pytom.basic.structures import Rotation
diffAng = differenceAngleOfTwoRotations(rotation1=Rotation(0, 0, 0),
rotation2=optiRot)
print(
"Alignment densities: Rotations: %2.3f, %2.3f, %2.3f; Translations: %2.3f, %2.3f, %2.3f "
% (optiRot[0], optiRot[1], optiRot[2], optiTrans[0], optiTrans[1],
optiTrans[2]))
print("Orientation difference: %2.3f deg" % diffAng)
vol2_alig = xp.zeros(vol2.shape, dtype=xp.float32)
transform(vol2,
output=vol2_alig,
rotation=[optiRot[0], optiRot[2], optiRot[1]],
rotation_order='rzxz',
center=[
int(vol2.shape[0] // 2),
int(vol2.shape[1] // 2),
int(vol2.shape[2] // 2)
],
interpolation='filt_bspline',
translation=[optiTrans[0], optiTrans[1], optiTrans[2]],
device=device)
# finally compute FSC and filter of both volumes
if not nband:
nband = int(vol2.sizeX() / 2)
fsc = FSC(volume1=vol1, volume2=vol2_alig, numberBands=nband)
fil = design_fsc_filter(fsc=fsc,
fildim=int(vol2.shape[0] // 2),
fsc_criterion=fsc_criterion)
filvol1 = filter_volume_by_profile(volume=vol1, profile=fil)
#filvol2 = filter_volume_by_profile( volume=vol2_alig, profile=fil)
filvol2 = filter_volume_by_profile(volume=vol2, profile=fil)
return (filvol1, filvol2, fsc, fil, optiRot, optiTrans)
def create_asymmetric_wedge(angle1,
angle2,
cutoffRadius,
sizeX,
sizeY,
sizeZ,
smooth,
rotation=None):
'''This function returns an asymmetric wedge object.
@param angle1: angle of wedge1 in degrees
@type angle1: int
@param angle2: angle of wedge2 in degrees
@type angle2: int
@param cutOffRadius: radius from center beyond which the wedge is set to zero.
@type cutOffRadius: int
@param sizeX: the size of the box in x-direction.
@type sizeX: int
@param sizeY: the size of the box in y-direction.
@type sizeY: int
@param sizeZ: the size of the box in z-direction.
@type sizeZ: int
@param smooth: smoothing parameter that defines the amount of smoothing at the edge of the wedge.
@type smooth: float
@return: 3D array determining the wedge object.
@rtype: ndarray of xp.float64'''
range_angle1Smooth = smooth / xp.sin(angle1 * xp.pi / 180.)
range_angle2Smooth = smooth / xp.sin(angle2 * xp.pi / 180.)
wedge = xp.zeros((sizeX, sizeY, sizeZ // 2 + 1))
if rotation is None:
z, y, x = xp.meshgrid(
xp.arange(-sizeX // 2 + sizeX % 2, sizeX // 2 + sizeX % 2),
xp.arange(-sizeY // 2 + sizeY % 2, sizeY // 2 + sizeY % 2),
xp.arange(0, sizeZ // 2 + 1))
else:
cx, cy, cz = [s // 2 for s in (sizeX, sizeY, sizeZ)]
grid = xp.mgrid[-cx:sizeX - cx, -cy:sizeY - cy, :sizeZ // 2 + 1]
phi, the, psi = rotation
phi = -float(phi) * xp.pi / 180.0
the = -float(the) * xp.pi / 180.0
psi = -float(psi) * xp.pi / 180.0
sin_alpha = xp.sin(phi)
cos_alpha = xp.cos(phi)
sin_beta = xp.sin(the)
cos_beta = xp.cos(the)
sin_gamma = xp.sin(psi)
cos_gamma = xp.cos(psi)
# Calculate inverse rotation matrix
Inv_R = xp.zeros((3, 3), dtype='float32')
Inv_R[0, 0] = cos_alpha * cos_gamma - cos_beta * sin_alpha \
* sin_gamma
Inv_R[0, 1] = -cos_alpha * sin_gamma - cos_beta * sin_alpha \
* cos_gamma
Inv_R[0, 2] = sin_beta * sin_alpha
Inv_R[1, 0] = sin_alpha * cos_gamma + cos_beta * cos_alpha \
* sin_gamma
Inv_R[1, 1] = -sin_alpha * sin_gamma + cos_beta * cos_alpha \
* cos_gamma
Inv_R[1, 2] = -sin_beta * cos_alpha
Inv_R[2, 0] = sin_beta * sin_gamma
Inv_R[2, 1] = sin_beta * cos_gamma
Inv_R[2, 2] = cos_beta
temp = grid.reshape((3, grid.size // 3))
temp = xp.dot(Inv_R, temp)
grid = xp.reshape(temp, grid.shape)
y = grid[0, :, :, :]
z = grid[1, :, :, :]
x = grid[2, :, :, :]
r = xp.sqrt(x**2 + y**2 + z**2)
wedge[xp.tan(angle1 * xp.pi / 180) < y / x] = 1
wedge[xp.tan(-angle2 * xp.pi / 180) > y / x] = 1
wedge[sizeX // 2, :, 0] = 1
if smooth:
area = xp.abs(x -
(y / xp.tan(angle1 * xp.pi / 180))) <= range_angle1Smooth
strip = 1 - (xp.abs(x - (y / xp.tan(angle1 * xp.pi / 180.))) *
xp.sin(angle1 * xp.pi / 180.) / smooth)
wedge += (strip * area * (1 - wedge) * (y > 0))
area2 = xp.abs(x +
(y /
xp.tan(angle2 * xp.pi / 180))) <= range_angle2Smooth
strip2 = 1 - (xp.abs(x + (y / xp.tan(angle2 * xp.pi / 180.))) *
xp.sin(angle2 * xp.pi / 180.) / smooth)
wedge += (strip2 * area2 * (1 - wedge) * (y <= 0))
wedge[r > cutoffRadius] = 0
return xp.fft.fftshift(wedge, axes=(0, 1))
def scale(volume, factor, interpolation='Spline'):
"""
scale: Scale a volume by a factor in REAL SPACE - see also resize function for more accurate operation in Fourier \
space
@param volume: input volume
@type volume: L{pytom_volume.vol}
@param factor: a factor > 0. Factors < 1 will de-magnify the volume, factors > 1 will magnify.
@type factor: L{float}
@param interpolation: Can be Spline (default), Cubic or Linear
@type interpolation: L{str}
@return: The scaled volume
@author: Thomas Hrabe
"""
from pytom.voltools import transform
from pytom.tompy.tools import paste_in_center
if factor <= 0:
raise RuntimeError('Scaling factor must be > 0!')
interpolation_dict = {
'Spline': 'filt_bspline',
'Linear': 'linear',
'Cubic': 'filt_bspline',
'filt_bspline': 'filt_bspline',
'linear': 'linear'
}
interpolation = interpolation_dict[interpolation]
volume = volume.squeeze()
sizeX = volume.shape[0]
sizeY = volume.shape[1]
sizeZ = 1
newSizeX = int(xp.floor(sizeX * float(factor) + 0.5))
newSizeY = int(xp.floor(sizeY * float(factor) + 0.5))
newSizeZ = 1
if len(volume.shape) == 3:
sizeZ = volume.shape[2]
newSizeZ = int(xp.floor(sizeZ * factor + 0.5))
scaleF = [1 / factor, 1 / factor, 1 / factor]
else:
scaleF = [1 / factor, 1 / factor, 1]
volume = xp.expand_dims(volume, 2)
if factor == 1:
rescaledVolume = volume
elif factor > 1:
newVolume = xp.zeros((newSizeX, newSizeY, newSizeZ),
dtype=volume.dtype)
newVolume = paste_in_center(volume, newVolume)
rescaledVolume = xp.zeros_like(newVolume)
transform(newVolume,
scale=scaleF,
output=rescaledVolume,
device=device,
interpolation=interpolation)
else:
rescaledVolumeFull = xp.zeros_like(volume)
transform(volume,
scale=scaleF,
output=rescaledVolumeFull,
device=device,
interpolation=interpolation)
rescaledVolume = xp.zeros((newSizeX, newSizeY, newSizeZ),
dtype=volume.dtype)
rescaledVolume = paste_in_center(rescaledVolumeFull, rescaledVolume)
return rescaledVolume
ScriptOption(['-v', '--verbose'], 'print intermediate output',
False, True)
])
try:
proj_dir, outdir, coordinateBinning, gpu, metafile, alignmentfile, size, verbose = ll = parse_script_options(
sys.argv[1:], helper)
except Exception as e:
print(e)
print(helper)
sys.exit()
coordinateBinning = int(coordinateBinning) if coordinateBinning else 1
metadata = loadstar(metafile, dtype=DATATYPE_METAFILE)
tilt_angles = metadata['TiltAngle']
size = [464, 464, 464] if size is None else list(map(int, size.split(',')))
patches = xp.zeros((size[0], size[1], len(tilt_angles)), dtype=xp.float32)
images = []
tt = time()
cur = 0
missed = 0
for i in range(patches.shape[2]):
temp_image = alignImageUsingAlignmentResultFile(
alignmentfile,
i,
weighting=-1,
circleFilter=True,
binning=coordinateBinning)
patches[:, :, i] = temp_image[:, :]
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
def alignImagesUsingAlignmentResultFile(alignmentResultsFile,
weighting=None,
lowpassFilter=0.9,
binning=1,
circleFilter=False):
import os
from pytom.basic.files import read as readCVol
from pytom_numpy import vol2npy, npy2vol
from pytom.gui.guiFunctions import fmtAR, headerAlignmentResults, datatype, datatypeAR, loadstar
from pytom.reconstruction.reconstructionStructures import Projection, ProjectionList
from pytom.tompy.io import read, write, read_size
from pytom.tompy.tools import taper_edges, create_circle
from pytom.tompy.filter import circle_filter, ramp_filter, exact_filter
import pytom.voltools as vt
from pytom.gpu.initialize import xp, device
print("Create aligned images from alignResults.txt")
alignmentResults = loadstar(alignmentResultsFile, dtype=datatypeAR)
imageList = alignmentResults['FileName']
tilt_angles = alignmentResults['TiltAngle']
imdim = read_size(imageList[0], 'x')
if binning > 1:
imdim = int(float(imdim) / float(binning) + .5)
else:
imdim = imdim
sliceWidth = imdim
if (weighting != None) and (float(weighting) < -0.001):
weightSlice = xp.fft.fftshift(ramp_filter(imdim, imdim))
if circleFilter:
circleFilterRadius = imdim // 2
circleSlice = xp.fft.fftshift(
circle_filter(imdim, imdim, circleFilterRadius))
else:
circleSlice = xp.ones((imdim, imdim))
# design lowpass filter
if lowpassFilter:
if lowpassFilter > 1.:
lowpassFilter = 1.
print("Warning: lowpassFilter > 1 - set to 1 (=Nyquist)")
# weighting filter: arguments: (()dimx, dimy), cutoff radius, sigma
lpf = xp.fft.fftshift(
create_circle((imdim, imdim),
lowpassFilter * (imdim // 2),
sigma=0.4 * lowpassFilter * (imdim // 2)))
projectionList = ProjectionList()
for n, image in enumerate(imageList):
atx = alignmentResults['AlignmentTransX'][n] / binning
aty = alignmentResults['AlignmentTransY'][n] / binning
rot = alignmentResults['InPlaneRotation'][n]
mag = 1 / alignmentResults['Magnification'][n]
projection = Projection(imageList[n],
tiltAngle=tilt_angles[n],
alignmentTransX=atx,
alignmentTransY=aty,
alignmentRotation=rot,
alignmentMagnification=mag)
projectionList.append(projection)
stack = xp.zeros((imdim, imdim, len(imageList)), dtype=xp.float32)
phiStack = xp.zeros((1, 1, len(imageList)), dtype=xp.float32)
thetaStack = xp.zeros((1, 1, len(imageList)), dtype=xp.float32)
offsetStack = xp.zeros((1, 2, len(imageList)), dtype=xp.float32)
for (ii, projection) in enumerate(projectionList):
print(f'Align {projection._filename}')
image = read(str(projection._filename))[::binning, ::binning].squeeze()
if lowpassFilter:
image = xp.abs((xp.fft.ifftn(xp.fft.fftn(image) * lpf)))
tiltAngle = projection._tiltAngle
# normalize to contrast - subtract mean and norm to mean
immean = image.mean()
image = (image - immean) / immean
# smoothen borders to prevent high contrast oscillations
image = taper_edges(image, imdim // 30)[0]
# transform projection according to tilt alignment
transX = projection._alignmentTransX / binning
transY = projection._alignmentTransY / binning
rot = float(projection._alignmentRotation)
mag = float(projection._alignmentMagnification)
inputImage = xp.expand_dims(image, 2).copy()
outputImage = xp.zeros_like(inputImage, dtype=xp.float32)
vt.transform(inputImage.astype(xp.float32),
rotation=[0, 0, rot],
rotation_order='rxyz',
output=outputImage,
device=device,
translation=[transX, transY, 0],
scale=[mag, mag, 1],
interpolation='filt_bspline')
image = outputImage.squeeze()
# smoothen once more to avoid edges
image = taper_edges(image, imdim // 30)[0]
# analytical weighting
if (weighting != None) and (weighting < 0):
# image = (ifft(complexRealMult(fft(image), w_func)) / (image.sizeX() * image.sizeY() * image.sizeZ()))
image = xp.fft.ifftn(
xp.fft.fftn(image) * weightSlice * circleSlice)
elif (weighting != None) and (weighting > 0):
weightSlice = xp.fft.fftshift(
exact_filter(tilt_angles, tiltAngle, imdim, imdim, sliceWidth))
image = xp.fft.ifftn(
xp.fft.fftn(image) * weightSlice * circleSlice)
thetaStack[0, 0, ii] = int(round(projection.getTiltAngle()))
offsetStack[0, :, ii] = xp.array([
int(round(projection.getOffsetX())),
int(round(projection.getOffsetY()))
])
stack[:, :, ii] = image
arrays = []
for fname, arr in (('stack.mrc', stack), ('offsetStack.mrc', offsetStack),
('thetaStack.mrc', thetaStack), ('phiStack.mrc',
phiStack)):
if 'gpu' in device:
arr = arr.get()
import numpy as np
res = npy2vol(np.array(arr, dtype='float32', order='F'), 3)
arrays.append(res)
#
# write('stack.mrc', stack)
# stack = readCVol('stack.mrc')
# write('offsetstack.mrc', offsetStack)
# offsetStack = readCVol('offsetstack.mrc')
# write('thetastack.mrc', thetaStack)
# thetaStack = readCVol('thetastack.mrc')
# write('phistack.mrc', phiStack)
# phiStack = readCVol('phistack.mrc')
#
# os.remove('stack.mrc')
# os.remove('offsetstack.mrc')
# os.remove('thetastack.mrc')
# os.remove('psistack.mrc')
return arrays
def alignImageUsingAlignmentResultFile(alignmentResultsFile,
indexImage,
weighting=None,
lowpassFilter=0.9,
binning=1,
circleFilter=False):
import pytom_freqweight
from pytom_numpy import vol2npy
from pytom.gui.guiFunctions import fmtAR, headerAlignmentResults, datatype, datatypeAR, loadstar
from pytom.reconstruction.reconstructionStructures import Projection, ProjectionList
from pytom.tompy.io import read, write, read_size
from pytom.tompy.tools import taper_edges, create_circle
from pytom.tompy.filter import circle_filter, ramp_filter, exact_filter, ellipse_filter
import pytom.voltools as vt
from pytom.gpu.initialize import xp, device
# print("Create aligned images from alignResults.txt")
alignmentResults = loadstar(alignmentResultsFile, dtype=datatypeAR)
imageList = alignmentResults['FileName']
tilt_angles = alignmentResults['TiltAngle']
imdimX = read_size(imageList[0], 'x')
imdimY = read_size(imageList[0], 'y')
if binning > 1:
imdimX = int(float(imdimX) / float(binning) + .5)
imdimY = int(float(imdimY) / float(binning) + .5)
sliceWidth = imdimX
if (weighting != None) and (float(weighting) < -0.001):
weightSlice = xp.fft.fftshift(ramp_filter(imdimY, imdimX))
if circleFilter:
circleFilterRadiusX = imdimX // 2
circleFilterRadiusY = imdimY // 2
circleSlice = xp.fft.fftshift(
ellipse_filter(imdimX, imdimY, circleFilterRadiusX,
circleFilterRadiusY))
else:
circleSlice = xp.ones((imdimX, imdimY))
# design lowpass filter
if lowpassFilter:
if lowpassFilter > 1.:
lowpassFilter = 1.
print("Warning: lowpassFilter > 1 - set to 1 (=Nyquist)")
# weighting filter: arguments: (()dimx, dimy), cutoff radius, sigma
# lpf = xp.fft.fftshift(create_circle((imdimX,imdimY),lowpassFilter*(imdim//2), sigma=0.4*lowpassFilter*(imdim//2)))
projectionList = ProjectionList()
for n, image in enumerate(imageList):
atx = alignmentResults['AlignmentTransX'][n]
aty = alignmentResults['AlignmentTransY'][n]
rot = alignmentResults['InPlaneRotation'][n]
mag = 1 / (alignmentResults['Magnification'][n])
projection = Projection(imageList[n],
tiltAngle=tilt_angles[n],
alignmentTransX=atx,
alignmentTransY=aty,
alignmentRotation=rot,
alignmentMagnification=mag)
projectionList.append(projection)
imdim = min(imdimY, imdimX)
for (ii, projection) in enumerate(projectionList):
if not ii == indexImage:
continue
from pytom.tompy.transform import resize
# print(f'read {projection._filename}')
image = read(str(projection._filename)).squeeze()
if binning > 1:
image = resize(image, 1 / binning)
#write(f'test/image_{ii}.mrc', image, tilt_angle=tilt_angles[ii])
tiltAngle = projection._tiltAngle
# 1 -- normalize to contrast - subtract mean and norm to mean
immean = image.mean()
image = (image - immean) / immean
# 2 -- smoothen borders to prevent high contrast oscillations
image = taper_edges(image, imdim // 30)[0]
# 3 -- square if needed
if 0 and imdimY != imdimX:
newImage = xp.zeros((imdim, imdim, 1), dtype=xp.float32)
pasteCenter(image, newImage)
image = newImage
# 4 -- transform projection according to tilt alignment
transX = projection._alignmentTransX / binning
transY = projection._alignmentTransY / binning
rot = float(projection._alignmentRotation)
mag = float(projection._alignmentMagnification)
inputImage = xp.expand_dims(image, 2).copy()
outputImage = xp.zeros_like(inputImage, dtype=xp.float32)
vt.transform(
inputImage.astype(xp.float32),
rotation=[0, 0, rot],
rotation_order='rxyz',
output=outputImage,
center=[inputImage.shape[0] // 2, inputImage.shape[1] // 2, 0],
device=device,
translation=[transX, transY, 0],
scale=[mag, mag, 1],
interpolation='filt_bspline')
del image
image = outputImage.squeeze()
# 5 -- Optional Low Pass Filter
if lowpassFilter:
from pytom.tompy.filter import bandpass_circle
image = bandpass_circle(
image,
high=lowpassFilter * (min(imdimX, imdimY) // 2),
sigma=0.4 * lowpassFilter * (min(imdimX, imdimY) // 2))
# image = xp.abs((xp.fft.ifftn(xp.fft.fftn(image) * lpf)))
# 6 -- smoothen once more to avoid edges
image = taper_edges(image, imdim // 30)[0]
# 7 -- analytical weighting
if (weighting != None) and (weighting < 0):
# image = (ifft(complexRealMult(fft(image), w_func)) / (image.sizeX() * image.sizeY() * image.sizeZ()))
image = xp.fft.ifftn(
xp.fft.fftn(image) * weightSlice.T * circleSlice).real
elif (weighting != None) and (weighting > 0):
weightSlice = xp.fft.fftshift(
exact_filter(tilt_angles, tiltAngle, imdim, imdim, sliceWidth))
image = xp.fft.ifftn(
xp.fft.fftn(image) * weightSlice * circleSlice).real
del inputImage, outputImage, circleSlice
write(f'inputImage_{ii}.mrc', image)
return image.astype(xp.float32)
def rotate3d(data, phi=0, psi=0, the=0, center=None, order=2):
"""Rotate a 3D data using ZXZ convention (phi: z1, the: x, psi: z2).
@param data: data to be rotated.
@param phi: 1st rotate around Z axis, in degree.
@param psi: 3rd rotate around Z axis, in degree.
@param the: 2nd rotate around X axis, in degree.
@param center: rotation center.
@return: the data after rotation.
"""
# Figure out the rotation center
import numpy as xp
if center is None:
cx = data.shape[0] // 2
cy = data.shape[1] // 2
cz = data.shape[2] // 2
else:
assert len(center) == 3
(cx, cy, cz) = center
if phi == 0 and psi == 0 and the == 0:
return data
# Transfer the angle to Euclidean
phi = -float(phi) * xp.pi / 180.0
the = -float(the) * xp.pi / 180.0
psi = -float(psi) * xp.pi / 180.0
sin_alpha = xp.sin(phi)
cos_alpha = xp.cos(phi)
sin_beta = xp.sin(the)
cos_beta = xp.cos(the)
sin_gamma = xp.sin(psi)
cos_gamma = xp.cos(psi)
# Calculate inverse rotation matrix
Inv_R = xp.zeros((3, 3), dtype='float32')
Inv_R[0, 0] = cos_alpha * cos_gamma - cos_beta * sin_alpha \
* sin_gamma
Inv_R[0, 1] = -cos_alpha * sin_gamma - cos_beta * sin_alpha \
* cos_gamma
Inv_R[0, 2] = sin_beta * sin_alpha
Inv_R[1, 0] = sin_alpha * cos_gamma + cos_beta * cos_alpha \
* sin_gamma
Inv_R[1, 1] = -sin_alpha * sin_gamma + cos_beta * cos_alpha \
* cos_gamma
Inv_R[1, 2] = -sin_beta * cos_alpha
Inv_R[2, 0] = sin_beta * sin_gamma
Inv_R[2, 1] = sin_beta * cos_gamma
Inv_R[2, 2] = cos_beta
grid = xp.mgrid[-cx:data.shape[0] - cx, -cy:data.shape[1] - cy,
-cz:data.shape[2] - cz]
temp = grid.reshape((3, grid.size // 3))
temp = xp.dot(Inv_R, temp)
grid = xp.reshape(temp, grid.shape)
grid[0] += cx
grid[1] += cy
grid[2] += cz
# Interpolation
from scipy.ndimage import map_coordinates
d = map_coordinates(data, grid, order=order)
return d
