def taper_edges(image, width):
"""
taper edges of image (or volume) with cos function
@param image: input image (or volume)
@type image: ndarray
@param width: width of edge
@type width: int
@return: image with smoothened edges, taper_mask
@rtype: array-like
@author: GvdS
"""
dims = list(image.shape) + [0]
val = xp.cos(xp.arange(1, width + 1) * xp.pi / (2. * (width)))
taperX = xp.ones((dims[0]), dtype=xp.float32)
taperY = xp.ones((dims[1]))
taperX[:width] = val[::-1]
taperX[-width:] = val
taperY[:width] = val[::-1]
taperY[-width:] = val
if dims[2] > 1:
taperZ = xp.ones((dims[2]))
taperZ[:width] = val[::-1]
taperZ[-width:] = val
Z, X, Y = xp.meshgrid(taperX, taperY, taperZ)
taper_mask = X * (X < Y) * (X < Z) + Y * (Y <= X) * (Y < Z) + Z * (
Z <= Y) * (Z <= X)
else:
X, Y = xp.meshgrid(taperY, taperX)
taper_mask = X * (X < Y) + Y * (Y <= X)
return image * taper_mask, taper_mask
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 cut_patch(projection,
ang,
pick_position,
vol_size=200,
binning=8,
dimz=None,
offset=[0, 0, 0],
projection_name=None):
from pytom.tompy.tools import taper_edges
#from pytom.gpu.initialize import xp
from pytom.voltools import transform
from pytom.tompy.transform import rotate3d, rotate_axis
from pytom.tompy.transform import cut_from_projection
from pytom.tompy.io import read
# Get the size of the original projection
dim_x = projection.shape[0]
dim_z = dim_x if dimz is None else dimz
x, y, z = pick_position
x = (x + offset[0]) * binning
y = (y + offset[1]) * binning
z = (z + offset[2]) * binning
# Get coordinates of the paricle adjusted for the tilt angle
yy = y # assume the rotation axis is around y
xx = (xp.cos(ang * xp.pi / 180) *
(x - dim_x / 2) - xp.sin(ang * xp.pi / 180) *
(z - dim_z / 2)) + dim_x / 2
# Cut the small patch out
patch = projection[max(0,
int(xp.floor(xx)) -
vol_size // 2):int(xp.floor(xx)) + vol_size // 2,
int(yy) - vol_size // 2:int(yy) + vol_size // 2, :]
shiftedPatch = xp.zeros_like(patch)
transform(patch,
output=shiftedPatch,
translation=[xx % 1, yy % 1, 0],
device=device,
interpolation='filt_bspline')
return shiftedPatch.squeeze()
def cut_patch(projection,
ang,
pick_position,
vol_size=200,
binning=8,
dimz=0,
offset=[0, 0, 0],
projection_name=None):
#from pytom.gpu.initialize import xp
from pytom.voltools import transform
from pytom.tompy.transform import rotate3d, rotate_axis
from pytom.tompy.transform import cut_from_projection
from pytom.tompy.io import read
# Get the size of the original projection
dim_x = projection.shape[0]
dim_z = dim_x if dimz is None else dimz
x, y, z = pick_position
x = (x + offset[0]) * binning
y = (y + offset[1]) * binning
z = (z + offset[2]) * binning
# Get coordinates of the paricle adjusted for the tilt angle
yy = y # assume the rotation axis is around y
xx = (xp.cos(ang * xp.pi / 180) *
(x - dim_x / 2) - xp.sin(ang * xp.pi / 180) *
(z - dim_z / 2)) + dim_x / 2
# Cut the small patch out
patch = projection[max(0,
int(xp.floor(xx)) -
vol_size // 2):int(xp.floor(xx)) + vol_size // 2,
int(yy) - vol_size // 2:int(yy) + vol_size // 2, :]
#transform(patch, output=patch, translation=[0,xx-float(xx),0], device='gpu:1')
#patch -= patch.mean()
return patch, xx, yy
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 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