def subPixelPeak(inp, k=1, ignore_border=1): from pytom.voltools import transform is2d = len(inp.shape.squeeze()) == 2 if is2d: inp2 = inp.squeeze()[ignore_border:-ignore_border, ignore_border:-ignore_border] scale = (k, k, 1) out = xp.array(inp.squeeze(), dtype=xp.float32) x, y = xp.array(xp.unravel_index(inp2.argmax(), inp2.shape)) + ignore_border out = xp.expand_dims(out, 2) translation = [inp.shape[0] // 2 - x, inp.shape[1] // 2 - y, 0] else: inp2 = inp[ignore_border:-ignore_border, ignore_border:-ignore_border, ignore_border:-ignore_border] scale = (k, k, k) out = xp.array(inp, dtype=xp.float32) x, y, z = xp.array(xp.unravel_index(inp2.argmax(), inp2.shape)) + ignore_border translation = [ inp.shape[0] // 2 - x, inp.shape[1] // 2 - y, inp.shape[2] // 2 - z ] zoomed = xp.zeros_like(out, dtype=xp.float32) transform(out, output=zoomed, scale=scale, translation=translation, device='gpu:0', interpolation='filt_bspline') shifts = [x2, y2, z2] = xp.unravel_index(zoomed.argmax(), zoomed.shape) return [zoomed[x2, y2, z2], shifts[3 - is2d]]
def normalised_cross_correlation(first, second, filter_mask=None, device=0): """ Do a cross correlation based on numpy @param first: The first dataset (numpy 2D) @type first: numpy array 2D @param second: The second dataset (numpy 2D) @type second: numpy array 2D @param filter_mask: a filter which is used to filter image 'first' @type filter_mask: numpy array 2D @return: The cross correlation result @returntype: numpy array 2D @requires: the shape of first to be equal to the shape of second, and equal t the shape of the filter (if used of course) """ assert first.shape == second.shape assert len(first.shape) == 2 if not (filter_mask is None): assert first.shape == filter_mask.shape # if filter_mask is None: ffirst = xp.fft.fft2(xp.array(first)) # else: # ffirst = xp.fft.fftshift(xp.fft.fftshift(xp.fft.fftn(first)) * filter_mask) ccmap = xp.real( xp.fft.fftshift( xp.fft.ifftn(xp.multiply(xp.fft.fftn(second), xp.conj(ffirst))))) / first.size return ccmap
def cut_from_projection(proj, center, size, device=2): """Cut out a subregion out from a 2D projection. @param proj: 2D projection. @param center: cutting center. @param size: cutting size. @return: subprojection. """ import pytom.voltools as vt if len(proj.shape) > 2 and proj.shape[2] > 1: raise Exception( 'We assume that projections come from a 3D object, thus your projection cannot be a 3D object itself' ) import numpy as np from scipy.ndimage import map_coordinates try: proj = proj.squeeze().get() except: proj = proj.squeeze() #v = # adjusted to python3 grid = np.mgrid[center[0] - size[0] // 2:center[0] + size[0] - size[0] // 2 - 1:size[0] * 1j, center[1] - size[1] // 2:center[1] + size[1] - size[1] // 2 - 1:size[1] * 1j] v = map_coordinates(proj, grid, order=2) return xp.array(v)
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 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 exact_filter(tilt_angles, tiltAngle, sX, sY, sliceWidth=1, arr=[]): """ exactFilter: Generates the exact weighting function required for weighted backprojection - y-axis is tilt axis Reference : Optik, Exact filters for general geometry three dimensional reconstuction, vol.73,146,1986. @param tilt_angles: list of all the tilt angles in one tilt series @param titlAngle: tilt angle for which the exact weighting function is calculated @param sizeX: size of weighted image in X @param sizeY: size of weighted image in Y @return: filter volume """ import numpy as xp # Calculate the relative angles in radians. diffAngles = (xp.array(tilt_angles) - tiltAngle) * xp.pi / 180. # Closest angle to tiltAngle (but not tiltAngle) sets the maximal frequency of overlap (Crowther's frequency). # Weights only need to be calculated up to this frequency. sampling = xp.min(xp.abs(diffAngles)[xp.abs(diffAngles) > 0.001]) crowtherFreq = min(sX // 2, xp.int32(xp.ceil(sliceWidth / xp.sin(sampling)))) arrCrowther = xp.matrix( xp.abs(xp.arange(-crowtherFreq, min(sX // 2, crowtherFreq + 1)))) # Calculate weights wfuncCrowther = 1. / (xp.clip( 1 - xp.array(xp.matrix(xp.abs(xp.sin(diffAngles))).T * arrCrowther)**2, 0, 2)).sum(axis=0) # Create full with weightFunc wfunc = xp.ones((sX, sY), dtype=xp.float32) # row_stack is not implemented in cupy weightingFunc = xp.column_stack(([ (wfuncCrowther), ] * (sY))).T wfunc[:, sX // 2 - crowtherFreq:sX // 2 + min(sX // 2, crowtherFreq + 1)] = weightingFunc return wfunc
def design_fsc_filter(fsc, fildim=None, fsc_criterion=0.143): """ design spectral filter to weight by SNR of frequency @param fsc: input fsc @type fsc: 1-d list @param fildim: filter dimension @type fildim: int @return: filter @rtype: list @author: FF """ from math import sqrt from pytom.basic.resolution import getResolutionBandFromFSC if not fildim: fildim = len(fsc) nband = len(fsc) if fsc_criterion != 0.0: resolutionBand = getResolutionBandFromFSC(fsc, criterion=fsc_criterion) smooth = max(resolutionBand / 5, 2) else: resolutionBand = len(fsc) smooth = 1 #print "filter: fsc_criterion %2.3f, resolutionBand %d" % (fsc_criterion, resolutionBand) # filter by sqrt(FSC) fsc_fil = len(fsc) * [0.] for ii in range(0, len(fsc)): fscval = fsc[ii] if fscval > 0.: #fsc_fil[ii] = sqrt(fsc[ii]) if ii <= resolutionBand: fsc_fil[ii] = sqrt(fsc[ii]) elif (ii > resolutionBand) and (ii <= resolutionBand + smooth): fsc_fil[ii] = sqrt(fsc[ii]) * (( (resolutionBand + smooth) - ii) / smooth)**2 # squared filter else: fsc_fil[ii] = 0. else: fsc_fil[ii] = 0. #design filter fil = xp.array(fildim * [0.]) if nband != len(fil): shrinkfac = 1. / (len(fil) / nband) for ii in range(len(fil) - 1): # linear resample fsc if nband ~= size(fil) if nband != len(fil): ilow = int(xp.floor(shrinkfac * ii)) dlow = shrinkfac * ii - ilow ihi = ilow + 1 dhi = ihi - shrinkfac * ii fil[ii] = fsc_fil[ilow] * dhi + fsc_fil[ihi] * dlow else: fil[ii] = fsc_fil[ii] return fil
def read_em(filename, order='F', keepnumpy=False, deviceID=None): """Read EM file. Now only support read the type float32 on little-endian machines. @param filename: file name to read. @return The data from the EM file in ndarray """ f = open(filename, 'r') try: dt_header = np.dtype('int32') header = np.fromfile(f, dt_header, 128) x = header[1] y = header[2] z = header[3] # read the data type dt = int(hex(header[0])[2]) default_type = False if dt == 1: # byte raise Exception("Data type not supported yet!") elif dt == 2: # short dt_data = np.dtype('<i2') elif dt == 4: # long dt_data = np.dtype('<i4') elif dt == 5: # float32 dt_data = np.dtype('<f4') # little-endian, float32 default_type = True elif dt == 8: # float complex raise Exception("Data type not supported yet!") elif dt == 9: # double dt_data = np.dtype('<f8') elif dt == 10: # double complex raise Exception("Data type not supported yet!") else: raise Exception("Data type not supported yet!") v = np.fromfile(f, dt_data, x * y * z) finally: f.close() if keepnumpy: volume = np.array(v.reshape((x, y, z), order=order), dtype='float32').copy() # fortran-order array else: # if the input data is not the default type, convert if not deviceID == None: id = int(deviceID.split(":")[1]) xp.cuda.Device(id).use() volume = xp.array(v.reshape((x, y, z), order=order), dtype='float32').copy() # fortran-order array return volume
def add_transformed_particle_to_sum(particle, sum, rotation=None, rotation_order='rzxz', translation=None, scale=None): from pytom.voltools import transform rot = transform(particle, rotation=rotation, translation=translation, scale=scale, rotation_order=rotation_order, device=device, interpolation='filt_bspline') sum += xp.array(rot) return sum
def read_mrc(filename, order='F', keepnumpy=False, deviceID=None): import numpy as np f = open(filename, 'r') try: dt_header = np.dtype('int32') header = np.fromfile(f, dt_header, 256) x = header[0] y = header[1] z = header[2] if header[23]: extended_header = np.fromfile(f, np.float32, header[23] // 4) # read the data type dt = header[3] default_type = False if dt == 0: # byte dt_data = np.dtype('int8') elif dt == 1: # short dt_data = np.dtype('int16') elif dt == 2: # long dt_data = np.dtype('float32') # little-endian, float32 default_type = True elif dt == 4: # float32 dt_data = np.dtype('complex32') elif dt == 6: # float complex dt_data = np.dtype('uint16') else: raise Exception("Data type not supported yet!") v = np.fromfile(f, dt_data, x * y * z) finally: f.close() if keepnumpy: volume = np.array(v.reshape((x, y, z), order=order), dtype='float32').copy() # fortran-order array else: if not deviceID == None: id = int(deviceID.split(":")[1]) xp.cuda.Device(id).use() volume = xp.array(v.reshape((x, y, z), order=order), dtype=xp.float32).copy() # fortran-order array return volume
def find_sub_pixel_max_value_voltools(inp, k=.1, ignore_border=50): """ To find the highest point in a 2D array, with subpixel accuracy based on 1D spline interpolation. The algorithm is based on a matlab script "tom_peak.m" @param inp: A 2D numpy array containing the data points. @type inp: numpy array 2D @param k: The smoothing factor used in the spline interpolation, must be 1 <= k <= 5. @type k: int @return: A list of all points of maximal value in the structure of tuples with the x position, the y position and the value. @returntype: list """ # assert len(ixp.shape) == 2 # assert isinstance(k, int) and 1 <= k <= 5 import cupy import numpy as xp from scipy.interpolate import InterpolatedUnivariateSpline from pytom.voltools import transform inp2 = inp[ignore_border:-ignore_border, ignore_border:-ignore_border] out = cupy.array(inp, dtype=cupy.float32) x, y = xp.array(xp.unravel_index(inp2.argmax(), inp2.shape)) + ignore_border out = cupy.expand_dims(out, 2) zoomed = cupy.zeros_like(out, dtype=cupy.float32) transform(out, output=zoomed, scale=(k, k, 1), translation=[inp.shape[0] // 2 - x, inp.shape[1] // 2 - y, 0], device='gpu:0', interpolation='filt_bspline') z = zoomed.get().squeeze() x2, y2 = xp.unravel_index(z.argmax(), z.shape) return [ x + (x2 - z.shape[0] // 2) * k - z.shape[0] // 2, y + (y2 - z.shape[1] // 2) * k - z.shape[0] // 2, z, x, y, x2, y2 ]
def FSC(volume1, volume2, numberBands=None, mask=None, verbose=False, filename=None, num_procs=1): """ FSC - Calculates the Fourier Shell Correlation for two volumes @param volume1: volume one @type volume1: L{pytom_volume.vol} @param volume2: volume two @type volume2: L{pytom_volume.vol} @param numberBands: number of shells for FSC @type numberBands: int @param filename: write FSC to ascii file if specified @type filename: string @return: Returns a list of cc values @author: Thomas Hrabe @rtype: list[floats] """ from pytom.tompy.correlation import bandCC from pytom.basic.structures import Mask from pytom.tompy.io import read, write from pytom.tompy.tools import volumesSameSize import time t = time.time() if not volumesSameSize(volume1, volume2): raise RuntimeError('Volumes must have the same size!') numberBands = volume1.shape[0] // 2 if numberBands is None else numberBands if not mask is None: if mask.__class__ == xp.array([0]).__class__: volume1 = volume1 * mask volume2 = volume2 * mask elif mask.__class__ == Mask: mask = mask.getVolume() volume1 = volume1 * mask volume2 = volume2 * mask elif mask.__class__ == str: mask = read(mask) volume1 = volume1 * mask volume2 = volume2 * mask else: raise RuntimeError( 'FSC: Mask must be a volume OR a Mask object OR a string path to a mask!' ) fscResult = [] band = [-1, -1] increment = int(volume1.shape[0] / 2 * 1 / numberBands) import time fvolume1 = xp.fft.fftn(volume1) fvolume2 = xp.fft.fftn(volume2) for n, i in enumerate(range(0, volume1.shape[0] // 2, increment)): band[0] = i band[1] = i + increment if verbose: print('Band : ', band) res = bandCC(fvolume1, fvolume2, band, verbose) if i == 0 and increment == 1: #force a 1 for correlation of the zero frequency res = 1 if verbose: print('Correlation ', res) fscResult.append(res) if filename: f = open(filename, 'w') for item in fscResult: f.write("%s\n" % item) f.close() return fscResult
def determineResolution(fsc, resolutionCriterion, verbose=False, randomizedFSC=None, gpu=False): """ determineResolution: Determines frequency and band where correlation drops below the resolutionCriterion. Uses linear interpolation between two positions @param fsc: The fsc list determined by L{pytom.basic.correlation.FSC} @param resolutionCriterion: A value between 0 and 1 @return: [resolution,interpolatedBand,numberBands] @author: Thomas Hrabe @todo: Add test! """ fsc = xp.array(fsc) numberBands = len(fsc) band = numberBands if randomizedFSC is None: randomizedFSC = xp.ones_like(fsc) * (fsc.min() - 0.1) for i in range(numberBands): if fsc[i] < resolutionCriterion and fsc[i] > randomizedFSC[i]: band = i - 1 #select the band that is still larger than criterion break if verbose: print('Band detected at ', band) if band == -1: raise RuntimeError( "Please check your resolution criterion or you FSC!") elif band < numberBands: fsc1 = fsc[band] fsc2 = fsc[band + 1] rfsc1 = randomizedFSC[band] rfsc2 = randomizedFSC[band + 1] try: if fsc2 < rfsc2: interpolatedBand = (fsc1 - rfsc1) / (rfsc2 - rfsc1 + fsc1 - fsc2) pass else: interpolatedBand = (resolutionCriterion - fsc1) / (fsc2 - fsc1) + band except ZeroDivisionError: interpolatedBand = band else: interpolatedBand = band if verbose: print('Band interpolated to ', interpolatedBand) resolution = (interpolatedBand + 1) / float(numberBands) if resolution < 0: resolution = 1 interpolatedBand = numberBands print( 'Warning: PyTom determined a resolution < 0 for your data. Please check "mass" in data is positive or negative for all cubes.' ) print('Warning: Setting resolution to 1 and ', interpolatedBand) print('') return [resolution, interpolatedBand, numberBands]
def subPixelMax3D(volume, k=.01, ignore_border=50, interpolation='filt_bspline', plan=None, profile=True, num_threads=1024, zoomed=None, fast_sum=None, max_id=None): """ Function to find the highest point in a 3D array, with subpixel accuracy using cubic spline interpolation. @param inp: A 3D numpy/cupy array containing the data points. @type inp: numpy/cupy array 3D @param k: The interpolation factor used in the spline interpolation, k < 1 is zoomed in, k>1 zoom out. @type k: float @return: A list of maximal value in the interpolated volume and a list of x position, the y position and the value. @returntype: list """ from pytom.voltools import transform from pytom.tompy.io import write ox, oy, oz = volume.shape ib = ignore_border cropped_volume = volume[ib:ox - ib, ib:oy - ib, ib:oz - ib].astype(xp.float32) if profile: stream = xp.cuda.Stream.null t_start = stream.record() # x,y,z = xp.array(maxIndex(cropped_volume)) + ignore_border x, y, z = xp.array( xp.unravel_index(cropped_volume.argmax(), cropped_volume.shape)) + ignore_border dx, dy, dz = volume.shape translation = [dx // 2 - x, dy // 2 - y, dz // 2 - z] if profile: t_end = stream.record() t_end.synchronize() time_took = xp.cuda.get_elapsed_time(t_start, t_end) print(f'initial find max time: \t{time_took:.3f}ms') t_start = stream.record() b = border = max(0, int(volume.shape[0] // 2 - 4 / k)) zx, zy, zz = volume.shape out = volume[b:zx - b, b:zy - b, b:zz - b] transform(out, output=zoomed, scale=(k, k, k), device=device, translation=translation, interpolation=interpolation) if profile: t_end = stream.record() t_end.synchronize() time_took = xp.cuda.get_elapsed_time(t_start, t_end) print(f'transform finished in \t{time_took:.3f}ms') t_start = stream.record() nblocks = int(xp.ceil(zoomed.size / num_threads / 2)) argmax(( nblocks, 1, ), (num_threads, 1, 1), (zoomed, fast_sum, max_id, zoomed.size), shared_mem=8 * num_threads) x2, y2, z2 = xp.unravel_index(max_id[fast_sum.argmax()], zoomed.shape) peakValue = zoomed[x2][y2][z2] peakShift = [ x + (x2 - zoomed.shape[0] // 2) * k - volume.shape[0] // 2, y + (y2 - zoomed.shape[1] // 2) * k - volume.shape[1] // 2, z + (z2 - zoomed.shape[2] // 2) * k - volume.shape[2] // 2 ] if profile: t_end = stream.record() t_end.synchronize() time_took = xp.cuda.get_elapsed_time(t_start, t_end) print(f'argmax finished in \t{time_took:.3f}ms') t_start = stream.record() print() return [peakValue, peakShift]
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 euclidian_distance(pos1, pos2): return xp.linalg.norm(xp.array(pos1) - xp.array(pos2))
def volumesSameSize(v0, v1): if len(v0.shape) != len(v1.shape): return False return xp.array([ abs(v0.shape[pos] - v1.shape[pos]) == 0 for pos in range(len(v0.shape)) ]).sum() == len(v0.shape)
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 profile2FourierVol(profile, dim=None, reduced=False): """ create Volume from 1d radial profile, e.g., to modulate signal with \ specific function such as CTF or FSC. Simple linear interpolation is used\ for sampling. @param profile: profile @type profile: 1-d L{pytom_volume.vol} or 1-d python array @param dim: dimension of (cubic) output @type dim: L{int} @param reduced: If true reduced Fourier representation (N/2+1, N, N) is generated. @type reduced: L{bool} @return: 3-dim complex volume with spherically symmetrical profile @rtype: L{pytom_volume.vol} @author: FF """ if dim is None: try: dim = [ 2 * profile.shape[0], ] * 3 except: dim = [ 2 * len(profile), ] * 3 is3D = (len(dim) == 3) nx, ny = dim[:2] if reduced: if is3D: nz = int(dim[2] // 2) + 1 else: ny = int(ny // 2) + 1 else: if is3D: nz = dim[2] try: r_max = profile.shape[0] - 1 except: r_max = len(profile) - 1 if len(dim) == 3: if reduced: X, Y, Z = xp.meshgrid(xp.arange(-nx // 2, nx // 2 + nx % 2), xp.arange(-ny // 2, ny // 2 + ny % 2), xp.arange(0, nz)) else: X, Y, Z = xp.meshgrid(xp.arange(-nx // 2, nx // 2 + nx % 2), xp.arange(-ny // 2, ny // 2 + ny % 2), xp.arange(-nz // 2, nz // 2 + nz % 2)) R = xp.sqrt(X**2 + Y**2 + Z**2) else: if reduced: X, Y = xp.meshgrid(xp.arange(-nx // 2, ny // 2 + ny % 2), xp.arange(0, ny)) else: X, Y = xp.meshgrid(xp.arange(-nx // 2, nx // 2 + nx % 2), xp.arange(-ny // 2, ny // 2 + ny % 2)) R = xp.sqrt(X**2 + Y**2) IR = xp.floor(R).astype(xp.int64) valIR_l1 = IR.copy() valIR_l2 = valIR_l1 + 1 val_l1, val_l2 = xp.zeros_like(X, dtype=xp.float64), xp.zeros_like( X, dtype=xp.float64) l1 = R - IR.astype(xp.float32) l2 = 1 - l1 try: profile = xp.array(profile) except: import numpy profile = xp.array(numpy.array(profile)) for n in xp.arange(r_max): val_l1[valIR_l1 == n] = profile[n] val_l2[valIR_l2 == n + 1] = profile[n + 1] val_l1[IR == r_max] = profile[n + 1] val_l2[IR == r_max] = profile[n + 1] val_l1[R > r_max] = 0 val_l2[R > r_max] = 0 fkernel = l2 * val_l1 + l1 * val_l2 if reduced: fkernel = xp.fft.fftshift(fkernel, axes=(0, 1)) else: fkernel = xp.fft.fftshift(fkernel) return fkernel
def polish_particles(particle_list_filename, projection_directory, averaged_subtomogram, binning, offset, projections, tilt_angles, fsc_path='', peak_border=75, outputDirectory='./', create_graphics=False, number_of_particles=0, verbose=False, gpuID=-1): """ To polish a particle list based on (an) initial subtomogram(s). :param particle_list_filename: the filename of the particlelist :type particle_list_filename: str :param projection_directory: the directory of the projections :type projection_directory: str :param averaged_subtomogram: to give a path to an averaged subtomogram to be used instead of subtomograms of all particles separately :type averaged_subtomogram: str :param binning: the binning factor used :type binning: int :param offset: the offset used (x, y, z) :type offset: list(int, int, int) :param projections: a list with filenames of projections :type projections: list(str) :param tilt_angles: the list of tiltangles used :type tilt_angles: list(int) :param create_graphics: to create plots of major parts of the algorithm, mainly used for debugging and initial creation :type create_graphics: bool :param number_of_particles: to use a subset of the particles for the particle polishing :type number_of_particles: int :param skip_alignment: skips the alignment phase, does not do particle polishing :type skip_alignment: bool :return: nothing, it writes everything to disk :returntype: void """ assert number_of_particles == -1 or number_of_particles > 0 assert binning > 0 assert vol_size > 0 assert vol_size % 2 == 0 assert isinstance(projections, list) assert isinstance(vol_size, int) assert isinstance(binning, int) assert isinstance(offset, list) and len(offset) == 3 assert isinstance(offset[0], int) and isinstance( offset[1], int) and isinstance(offset[2], int) assert isinstance(tilt_angles, list) assert isinstance(particle_list_filename, str) assert isinstance(projection_directory, str) assert isinstance(create_graphics, bool) assert isinstance(averaged_subtomogram, str) assert isinstance(number_of_particles, int) assert isinstance(skip_alignment, bool) import os, time from pytom.tompy.io import read_size, read from pytom.gui.guiFunctions import fmtLAR, headerLocalAlignmentResults, LOCAL_ALIGNMENT_RESULTS import pytom.voltools as vt # load particle list from pytom.basic.structures import ParticleList particlelist = ParticleList() particlelist.fromXMLFile(particle_list_filename) particle_list_name = os.path.splitext( os.path.basename(str(particle_list_filename)))[0] if number_of_particles > 0: particlelist = particlelist[:number_of_particles] if verbose: print(len(particlelist)) print("{:s}> Creating the input array".format(gettime())) dimz = read_size(particlelist[0].getPickPosition().getOriginFilename(), 'z') * binning vol_size = 200 input_to_processes = [] data = {} for projectioname in projections: data[projectioname] = xp.array(read(projectioname)) # template1 = read(averaged_subtomogram, order='F') #template1 = mrcfile.open(averaged_subtomogram,permissive=True).data.copy().T.copy() template = vt.StaticVolume(template1, interpolation='filt_bspline', device=device) # output = [] results_file = os.path.join(outputDirectory, f"resultsPolish_{particle_list_name}.txt") results = [] for particle_number, particle in enumerate(particlelist): rot = (particle.getRotation().getZ1(), particle.getRotation().getX(), particle.getRotation().getZ2()) # loop over tiltrange, take patch and cross correlate with reprojected subtomogram for img, ang in zip(projections, tilt_angles): pick_position = particle.getPickPosition().toVector() result = run_single_tilt_angle(template, ang, offset, vol_size, pick_position, rot, particle.getFilename(), particle_number, binning, data, create_graphics, fsc_path, dimz, peak_border, img, averaged_subtomogram) results.append(tuple(result)) try: np.savetxt(results_file, np.array(results, dtype=LOCAL_ALIGNMENT_RESULTS), fmt=fmtLAR, header=headerLocalAlignmentResults) except Exception as e: print(e) for res in results: print('{:7d} {:15.3f} {:15.3f} {:15.3f} {:15.3f} {:15.10f} {:s}'. format(*res)) break if verbose: print("{:s}> Ran the processes".format(gettime()))