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