def mean(volume, mask, fVolume=0, fMask=0, numberOfMaskVoxels=-1): """ suggestion frido: rename to mean_moving_mask mean: Determines the mean of volume under mask. Each assigned value is determined through the value of the voxels surrounding the current voxel that are covered by the mask. @param volume: The volume of interest @param mask: A mask under which the mean is determined @param fVolume: Optional - the fouriertransformed volume (fft must not be repeated) @param fMask: Optional - the fouriertransformed mask (fft must not be repeated) @param numberOfMaskVoxels: Optional - the number of voxels != 0 """ import pytom_volume if not fVolume.__class__ == pytom_volume.vol_comp: from pytom.basic.fourier import fft fVolume = fft(volume) if not fMask.__class__ == pytom_volume.vol_comp: from pytom.basic.fourier import fft fMask = fft(mask) if numberOfMaskVoxels < 0: numberOfMaskVoxels = pytom_volume.numberSetVoxels(mask) fMeanVol = fVolume * fMask from pytom.basic.fourier import ifft meanVolume = ifft(fMeanVol) from pytom.basic.fourier import iftshift iftshift(meanVolume) meanVolume.shiftscale(0, 1.0 / (numberOfMaskVoxels * meanVolume.numelem())) return meanVolume
def frm_fourier_constrained_vol(vf, mf, vg, mg): radius = vf.sizeX() / 2 - 3 # intepolation in the outer part is nonsense b = 32 from pytom.basic.fourier import fft, ifft, ftshift, iftshift vf = ftshift(reducedToFull(fft(iftshift(vf, inplace=False))), inplace=False) vg = ftshift(reducedToFull(fft(iftshift(vg, inplace=False))), inplace=False) vfr = real(vf) vfi = imag(vf) vgr = real(vg) vgi = imag(vg) res = np.zeros((2 * b, 2 * b, 2 * b)) for r in xrange(1, radius + 1): corr = frm_fourier_constrained_corr(vol2sf(vfr, r, b), vol2sf(vfi, r, b), mf, vol2sf(vgr, r, b), vol2sf(vgi, r, b), mg, True) res += corr * (r**2) return res
def bandCF(volume, reference, band=[0, 100]): """ bandCF: @param volume: The volume @param reference: The reference @param band: [a,b] - specify the lower and upper end of band. [0,1] if not set. @return: First parameter - The correlation of the two volumes in the specified ring. Second parameter - The bandpass filter used. @rtype: List - [L{pytom_volume.vol},L{pytom_freqweight.weight}] @author: Thomas Hrabe @todo: does not work yet -> test is disabled """ if gpu: import cupy as xp else: import numpy as xp import pytom_volume from math import sqrt from pytom.basic import fourier from pytom.basic.filter import bandpassFilter from pytom.basic.correlation import nXcf vf = bandpassFilter(volume, band[0], band[1], fourierOnly=True) rf = bandpassFilter(reference, band[0], band[1], vf[1], fourierOnly=True) v = pytom_volume.reducedToFull(vf[0]) r = pytom_volume.reducedToFull(rf[0]) absV = pytom_volume.abs(v) absR = pytom_volume.abs(r) pytom_volume.power(absV, 2) pytom_volume.power(absR, 2) sumV = abs(pytom_volume.sum(absV)) sumR = abs(pytom_volume.sum(absR)) if sumV == 0: sumV = 1 if sumR == 0: sumR = 1 pytom_volume.conjugate(rf[0]) fresult = vf[0] * rf[0] #transform back to real space result = fourier.ifft(fresult) fourier.iftshift(result) result.shiftscale(0, 1 / float(sqrt(sumV * sumR))) return [result, vf[1]]
def xcf(volume, template, mask=None, stdV=None): """ XCF: returns the non-normalised cross correlation function. The xcf result is scaled only by the square of the number of elements. @param volume : The search volume @type volume: L{pytom_volume.vol} @param template : The template searched (this one will be used for conjugate complex multiplication) @type template: L{pytom_volume.vol} @param mask: changed: will be used if specified @type mask: L{pytom_volume.vol} @param stdV: Will be unused, only for compatibility reasons with FLCF @return: XCF volume @rtype: L{pytom_volume.vol} @author: Thomas Hrabe """ import pytom_volume from pytom.basic import fourier if mask != None: volume = volume * mask template = template * mask #determine fourier transforms of volumes if volume.__class__ == pytom_volume.vol: from pytom.basic.fourier import fft fvolume = fft(volume) else: fvolume = volume if template.__class__ == pytom_volume.vol: from pytom.basic.fourier import fft ftemplate = fft(template) else: ftemplate = template #perform element wise - conjugate multiplication pytom_volume.conjugate(ftemplate) fresult = fvolume * ftemplate #transform back to real space result = fourier.ifft(fresult) fourier.iftshift(result) n = result.numelem() if mask: n1 = pytom_volume.sum(mask) # real -> FFT requires n1, FFT -> real n result.shiftscale(0,1/float(n1*n)) else: result.shiftscale(0,1/float(n*n)) return result
def fourier_rotate(v, phi, psi, the): vf = ftshift(reducedToFull(fft(iftshift(v, inplace=False))), inplace=False) vfr = real(vf) vfi = imag(vf) rr = vol(vfr) ii = vol(vfi) rotateSpline(vfr, rr, phi, psi, the) rotateSpline(vfi, ii, phi, psi, the) # print vfr(12,12,12),rr(12,12,12) vv = mergeRealImag(rr, ii) return ftshift(ifft(fullToReduced(iftshift(vv, inplace=False))), inplace=False) / v.numelem()
def mean0std1(volume, copyFlag=False): """ mean0std1: normalises input volume to mean 0 and std 1. Procedure is performed inplace if copyFlag is unspecified!!! @param volume: Data containing either an image or a volume @param copyFlag: If True a copy of volume will be returned. False unless specified otherwise. @return: If copyFlag == True, then return a normalised copy. @author: Thomas Hrabe """ import pytom_volume from math import sqrt from pytom.tools.maths import epsilon if volume.__class__ == pytom_volume.vol_comp: from pytom.basic.fourier import ifft, iftshift volume = iftshift(ifft(volume)) if not copyFlag: volumeMean = pytom_volume.mean(volume) volume.shiftscale(-volumeMean, 1) volumeStd = sqrt(pytom_volume.variance(volume, False)) if volumeStd < epsilon: raise ValueError( 'pytom_normalise.mean0std1 : The standard deviation is too low for division! ' + str(volumeStd)) volume.shiftscale(0, 1. / float(volumeStd)) else: volumeCopy = pytom_volume.vol(volume.sizeX(), volume.sizeY(), volume.sizeZ()) volumeCopy.copyVolume(volume) volumeMean = pytom_volume.mean(volumeCopy) volumeCopy.shiftscale(-1. * volumeMean, 1) volumeStd = sqrt(pytom_volume.variance(volumeCopy, False)) if volumeStd < epsilon: raise ValueError( 'pytom_normalise.mean0std1 : The standard deviation is too low for division! ' + str(volumeStd)) #volumeCopy.shiftscale(-1.*volumeMean,1) volumeCopy.shiftscale(0, 1 / float(volumeStd)) return volumeCopy
def std(volume, mask, meanVolume=0, fMask=0, numberOfMaskVoxels=-1): """ suggestion frido: rename to std_moving_mask std: Determines the std of volume under moving mask. Each assigned value is determined through the value of the voxels surrounding the current voxel that are covered by the mask. @param volume: The volume of interest @param mask: A mask under which the std is determined @param meanVolume: Optional - the meanVolume determined by mean. (mean must not be recalculated) @param fMask: Optional - the fouriertransformed mask (fft must not be repeated) @param numberOfMaskVoxels: Optional - the number of voxels != 0 """ from pytom.basic.fourier import fft, ifft, iftshift from pytom_volume import power import pytom_volume if not fMask.__class__ == pytom_volume.vol_comp: fMask = fft(mask) if not meanVolume.__class__ == pytom_volume.vol: meanVolume = mean(volume, mask) if numberOfMaskVoxels < 0: numberOfMaskVoxels = pytom_volume.numberSetVoxels(mask) volumeSqrd = volume * volume fVolumeSqrd = fft(volumeSqrd) fVolumeSqrd = fVolumeSqrd * fMask varianceVolume = iftshift(ifft(fVolumeSqrd)) varianceVolume.shiftscale( 0, 1.0 / (numberOfMaskVoxels * varianceVolume.numelem())) power(meanVolume, 2) varianceVolume = varianceVolume - meanVolume power(varianceVolume, 0.5) return varianceVolume
def transformFourierSpline(volume,z1,z2,x,shiftX,shiftY,shiftZ,twice=False): """ transformFourierSpline: Rotate and shift a volume in fourierspace @param volume: @param z1: @param z2: @param x: @param shiftX: Shift after rotation @param shiftY: Shift after rotation @param shiftZ: Shift after rotation @param twice: Zero pad volume into a twice sized volume and perform calculation there. @return: The transformed volume. @author: Yuxiang Chen and Thomas Hrabe """ from pytom.basic.fourier import fft, ifft, ftshift, iftshift from pytom_volume import vol, pasteCenter, subvolume, transformFourierSpline if z1 == 0 and z2 == 0 and x == 0: return vol(volume) if twice: # paste into a twice sized volume v = vol(volume.sizeX()*2, volume.sizeY()*2, volume.sizeZ()*2) pasteCenter(volume, v) else: v = volume fvol = fft(iftshift(v, inplace=False)) # these steps have to be done in python level because of the fft resF = transformFourierSpline(fvol,z1,z2,x,shiftX,shiftY,shiftZ) res = ftshift(ifft(resF),inplace=False) / v.numelem() # reverse it back if twice: # cut the center part back res = subvolume(res, (volume.sizeX()+1)/2, (volume.sizeY()+1)/2, (volume.sizeZ()+1)/2, volume.sizeX(), volume.sizeY(), volume.sizeZ()) return res
def meanUnderMask(volume, mask, p): """ meanUnderMask: calculate the mean volume under the given mask (Both should have the same size) @param volume: input volume @type volume: L{pytom_volume.vol} @param mask: mask @type mask: L{pytom_volume.vol} @param p: non zero value numbers in the mask @type p: L{int} or L{float} @return: the calculated mean volume under mask @rtype: L{pytom_volume.vol} @author: Yuxiang Chen """ size = volume.numelem() from pytom.basic.fourier import fft, ifft, iftshift from pytom_volume import conjugate # for some reason, this has to be conjugated. (Otherwise the asym mask won't work) fMask = fft(mask) conjugate(fMask) result = iftshift(ifft(fMask*fft(volume)))/(size*p) return result
from pytom_volume import * from pytom.basic.fourier import fft, ifft, ftshift, iftshift v = read('/fs/home/ychen/matlab/template/binning/temp80SRibosome_bin2.em') # v = read('/fs/home/ychen/matlab/template/temp80SRibosome.em') fv = fft(v) r1 = real(fv) i1 = imag(fv) fv2 = fullToReduced(reducedToFull(fv)) r2 = real(fv2) i2 = imag(fv2) (r1 - r2).info('') (i1 - i2).info('') fv3 = fullToReduced( iftshift(ftshift(reducedToFull(fv), inplace=False), inplace=False)) r3 = real(fv3) i3 = imag(fv3) (r1 - r3).info('') (i1 - i3).info('')
def frm_correlate(vf, wf, vg, wg, b, max_freq, weights=None, ps=False, denominator1=None, denominator2=None, return_score=True): """Calculate the correlation of two volumes as a function of Euler angle. Parameters ---------- vf: Volume Nr. 1 pytom_volume.vol wf: Mask of vf in Fourier space. pytom.basic.structures.Wedge vg: Volume Nr. 2 / Reference pytom_volume.vol wg: Mask of vg in Fourier space. pytom.basic.structures.Wedge b: Bandwidth range of spherical harmonics. None -> [4, 64] List -> [b_min, b_max] Integer -> [b, b] max_freq: Maximal frequency involved in calculation. Integer. weights: Obsolete. ps: Calculation based on only the power spectrum of two volumes or not. Boolean. Default is False. denominator1: If the denominator1 is provided or not. If yes, do not have to re-calculate it again. This field is used out of computation effeciency consideration. Default is None, not provided. denominator2: If the denominator2 is provided or not. If yes, do not have to re-calculate it again. This field is used out of computation effeciency consideration. Default is None, not provided. return_score: Return the correlation score or return the intermediate result (numerator, denominator1, denominator2). Boolean, default is True. Returns ------- If return_score is set to True, return the correlation function; otherwise return the intermediate result. """ if not weights: # weights, not used yet weights = [1 for i in xrange(max_freq)] from pytom.basic.fourier import fft, ifft, ftshift, iftshift from pytom_volume import vol, reducedToFull, abs, real, imag, rescale from vol2sf import vol2sf from math import log, ceil, pow # IMPORTANT!!! Should firstly do the IFFTSHIFT on the volume data (NOT FFTSHIFT since for odd-sized data it matters!), # and then followed by the FFT. vf = ftshift(reducedToFull(fft(iftshift(vf, inplace=False))), inplace=False) vg = ftshift(reducedToFull(fft(iftshift(vg, inplace=False))), inplace=False) if ps: # power spectrum only ff = abs(vf) ff = real(ff) gg = abs(vg) gg = real(gg) else: # use spline intepolation on the real/imaginary parts. Can be done better, but now it suffices. vfr = real(vf) vfi = imag(vf) vgr = real(vg) vgi = imag(vg) numerator = None if denominator1 is not None and denominator2 is not None: to_calculate = 1 elif denominator1 is None and denominator2 is not None: to_calculate = 2 else: to_calculate = 0 _last_bw = 0 # might be a better idea to start from 2 due to the bad interpolation around 0 frequency! # this can be better solved by NFFT! for r in xrange(1, max_freq + 1): # calculate the appropriate bw bw = get_adaptive_bw(r, b) # construct the wedge masks accordingly # now this part has been shifted to Pytom # if _last_bw != bw: # # mf = create_wedge_sf(wf[0], wf[1], bw) # # mg = create_wedge_sf(wg[0], wg[1], bw) # mf = wf.toSphericalFunc(bw) # mg = wg.toSphericalFunc(bw) mf = wf.toSphericalFunc(bw, r) mg = wg.toSphericalFunc(bw, r) if ps: corr1, corr2, corr3 = sph_correlate_ps(vol2sf(ff, r, bw), mf, vol2sf(gg, r, bw), mg, to_calculate) else: corr1, corr2, corr3 = sph_correlate_fourier( vol2sf(vfr, r, bw), vol2sf(vfi, r, bw), mf, vol2sf(vgr, r, bw), vol2sf(vgi, r, bw), mg, to_calculate) if _last_bw != bw: # size is different, have to do enlarge if numerator is None: numerator = np.zeros((2 * bw, 2 * bw, 2 * bw), dtype='double') if to_calculate == 1: pass elif to_calculate == 2: denominator1 = np.zeros((2 * bw, 2 * bw, 2 * bw), dtype='double') else: denominator1 = np.zeros((2 * bw, 2 * bw, 2 * bw), dtype='double') denominator2 = np.zeros((2 * bw, 2 * bw, 2 * bw), dtype='double') else: numerator = enlarge2(numerator) if to_calculate == 1: pass elif to_calculate == 2: denominator1 = enlarge2(denominator1) else: denominator1 = enlarge2(denominator1) denominator2 = enlarge2(denominator2) numerator += corr1 * (r**2) * weights[r - 1] if to_calculate == 1: pass elif to_calculate == 2: denominator1 += corr2 * (r**2) * weights[r - 1] else: denominator1 += corr2 * (r**2) * weights[r - 1] denominator2 += corr3 * (r**2) * weights[r - 1] _last_bw = bw if return_score: res = numerator / (denominator1 * denominator2)**0.5 return res else: return (numerator, denominator1, denominator2)
b = 16 r = 4 m = np.ones(4*b**2) # w = m w = create_wedge_sf(-60, 60, b) dist = [] dist2 = [] dist3 = [] for i in xrange(100): phi = np.random.randint(360) psi = np.random.randint(360) the = np.random.randint(180) fv1 = ftshift(reducedToFull(fft(iftshift(v, inplace=False))), inplace=False) # 1. rotate in real space and use the frm_fourier_corr to find the angle # This is the least accurate way, since the interpolation happens in real space. # rotateSpline(v, v2, phi, psi, the) # fv2 = ftshift(reducedToFull(fft(iftshift(v2)))) # res = frm_fourier_corr(vol2sf(real(fv2), r, b), vol2sf(imag(fv2), r, b), vol2sf(real(fv1), r, b), vol2sf(imag(fv1), r, b)) # 2. rotate real and imag parts seperately and feed into the frm_fourier_corr fr = real(fv1) fi = imag(fv1) # rotateSpline(v, v2, phi, psi, the) # fv2 = ftshift(reducedToFull(fft(iftshift(v2)))) # fr2 = real(fv2) # fi2 = imag(fv2)
def FLCF(volume, template, mask=None, stdV=None): ''' Created on Apr 13, 2010 FLCF: compute the fast local correlation function This functions works only when the mask is in the middle. @param volume: target volume @type volume: L{pytom_volume.vol} @param template: template to be searched. It can have smaller size then target volume. @type template: L{pytom_volume.vol} @param mask: template mask. If not given, a default sphere mask will be generated which has the same size with the given template. @type mask: L{pytom_volume.vol} @param stdV: standard deviation of the target volume under mask, which do not need to be calculated again when the mask is identical. @type stdV: L{pytom_volume.vol} @return: the local correlation function @rtype: L{pytom_volume.vol} @author: Yuxiang Chen ''' from pytom_volume import vol, pasteCenter from pytom.basic.fourier import fft, ifft, iftshift from pytom_volume import conjugate from pytom.basic.structures import Mask from pytom_volume import sum from pytom.basic.files import write_em if volume.__class__ != vol and template.__class__ != vol: raise RuntimeError('Wrong input type!') if volume.sizeX()<template.sizeX() or volume.sizeY()<template.sizeY() or volume.sizeZ()<template.sizeZ(): raise RuntimeError('Template size is bigger than the target volume') # generate the mask if mask.__class__ != vol: from pytom_volume import initSphere mask = vol(template.sizeX(), template.sizeY(), template.sizeZ()) mask.setAll(0) initSphere(mask, template.sizeX()/2,0,0,template.sizeX()/2, template.sizeY()/2, template.sizeZ()/2) else: if template.sizeX()!=mask.sizeX() and template.sizeY()!=mask.sizeY() and template.sizeZ()!=mask.sizeZ(): raise RuntimeError('Template and mask size are not consistent!') # calculate the non-zeros p = sum(mask) # normalize the template under mask meanT = meanValueUnderMask(template, mask, p) stdT = stdValueUnderMask(template, mask, meanT, p) temp = (template - meanT)/stdT temp = temp * mask # construct both the template and the mask which has the same size as target volume tempV = temp if volume.sizeX() != temp.sizeX() or volume.sizeY() != temp.sizeY() or volume.sizeZ() != temp.sizeZ(): tempV = vol(volume.sizeX(), volume.sizeY(), volume.sizeZ()) tempV.setAll(0) pasteCenter(temp, tempV) maskV = mask if volume.sizeX() != mask.sizeX() or volume.sizeY() != mask.sizeY() or volume.sizeZ() != mask.sizeZ(): maskV = vol(volume.sizeX(), volume.sizeY(), volume.sizeZ()) maskV.setAll(0) pasteCenter(mask, maskV) # calculate the mean and std of volume if stdV.__class__ != vol: meanV = meanUnderMask(volume, maskV, p) stdV = stdUnderMask(volume, maskV, p, meanV) size = volume.numelem() fT = fft(tempV) conjugate(fT) result = iftshift(ifft(fT*fft(volume)))/stdV result.shiftscale(0,1/(size*p)) return result
def frm_fourier_adaptive_wedge_vol_rscore(vf, wf, vg, wg, b, radius=None, weights=None): """Obsolete. """ if not radius: # set the radius radius = vf.sizeX()/2 if not weights: # set the weights weights = [1 for i in range(radius)] if not b: # set the bandwidth adaptively b_min = 4 b_max = 128 elif b.__class__ == tuple or b.__class__ == list: b_min = b[0] b_max = b[1] elif isinstance(b, int): # fixed bandwidth b_min = b b_max = b else: raise RuntimeError("Argument b is not valid!") from pytom.basic.fourier import fft, ifft, ftshift, iftshift from pytom_volume import vol, reducedToFull, real, imag, rescale from .vol2sf import vol2sf from pytom_numpy import vol2npy from math import log, ceil, pow # IMPORTANT!!! Should firstly do the IFFTSHIFT on the volume data (NOT FFTSHIFT since for odd-sized data it matters!), # and then followed by the FFT. vf = ftshift(reducedToFull(fft(iftshift(vf, inplace=False))), inplace=False) vg = ftshift(reducedToFull(fft(iftshift(vg, inplace=False))), inplace=False) vfr = real(vf) vfi = imag(vf) vgr = real(vg) vgi = imag(vg) get_bw = lambda x: int(pow(2, int(ceil(log(2*x, 2))))) res = None _last_bw = 0 for r in range(1, radius+1): # calculate the appropriate bw bw = get_bw(r) if bw < b_min: bw = b_min if bw > b_max: bw = b_max # construct the wedge masks accordingly if _last_bw != bw: mf = create_wedge_sf(wf[0], wf[1], bw) mg = create_wedge_sf(wg[0], wg[1], bw) corr = frm_fourier_constrained_corr(vol2sf(vfr, r, bw), vol2sf(vfi, r, bw), mf, vol2sf(vgr, r, bw), vol2sf(vgi, r, bw), mg, True, False, True) if _last_bw != bw: if res is None: res = np.zeros((2*bw, 2*bw, 2*bw), dtype='double') else: res = enlarge2(res) res += corr*(r**2)*weights[r-1] _last_bw = bw return res
def frm_determine_orientation(vf, wf, vg, wg, b, radius=None, weights=None, r_score=False, norm=False): """Auxiliary function for xu_align_vol. Find the angle to rotate vg to match vf, using only their power spectrums. Parameters ---------- vf: The volume you want to match. pytom_volume.vol wf: The single tilt wedge information of volume vf. [missing_wedge_angle1, missing_wedge_angle2]. Note this is defined different with frm_align im frm.py! vg: The reference volume. pytom_volume.vol wg: The single tilt wedge information of volume vg. [missing_wedge_angle1, missing_wedge_angle2]. Note this is defined different with frm_align im frm.py! b: The adaptive bandwidth of spherical harmonics. List [min_bandwidth, max_bandwidth], min_bandwidth, max_bandwidth in the range [4, 64]. Or integer, which would then mean to use fixed bandwidth: min_bandwidth = max_bandwidth = integer. radius: The maximal radius in the Fourier space, which is equal to say the maximal frequency involved in calculation. Integer. By default is half of the volume size. weights: Obsolete. r_score: Obsolete. norm: Obsolete. Returns ------- The angle (Euler angle, ZXZ convention [Phi, Psi, Theta]) to rotate vg to match vf. """ if not radius: # set the radius radius = vf.sizeX()/2 if not weights: # set the weights weights = [1 for i in range(radius)] if not b: # set the bandwidth adaptively b_min = 4 b_max = 128 elif b.__class__ == tuple or b.__class__ == list: b_min = b[0] b_max = b[1] elif isinstance(b, int): # fixed bandwidth b_min = b b_max = b else: raise RuntimeError("Argument b is not valid!") from pytom.basic.fourier import fft, ifft, ftshift, iftshift from pytom_volume import vol, reducedToFull, rescale, abs, real from .vol2sf import vol2sf from pytom_numpy import vol2npy from math import log, ceil, pow # IMPORTANT!!! Should firstly do the IFFTSHIFT on the volume data (NOT FFTSHIFT since for odd-sized data it matters!), # and then followed by the FFT. vf = ftshift(reducedToFull(fft(iftshift(vf, inplace=False))), inplace=False) vg = ftshift(reducedToFull(fft(iftshift(vg, inplace=False))), inplace=False) ff = abs(vf) ff = real(ff) gg = abs(vg) gg = real(gg) get_bw = lambda x: int(pow(2, int(ceil(log(2*x, 2))))) numerator = None denominator1 = None denominator2 = None _last_bw = 0 for r in range(1, radius+1): # calculate the appropriate bw bw = get_bw(r) if bw < b_min: bw = b_min if bw > b_max: bw = b_max # construct the wedge masks accordingly if _last_bw != bw: mf = create_wedge_sf(wf[0], wf[1], bw) mg = create_wedge_sf(wg[0], wg[1], bw) corr1, corr2, corr3 = frm_constrained_corr(vol2sf(ff, r, bw), mf, vol2sf(gg, r, bw), mg, norm, return_score=False) if _last_bw != bw: if numerator is None: numerator = np.zeros((2*bw, 2*bw, 2*bw), dtype='double') denominator1 = np.zeros((2*bw, 2*bw, 2*bw), dtype='double') denominator2 = np.zeros((2*bw, 2*bw, 2*bw), dtype='double') else: numerator = enlarge2(numerator) denominator1 = enlarge2(denominator1) denominator2 = enlarge2(denominator2) numerator += corr1*(r**2)*weights[r-1] denominator1 += corr2*(r**2)*weights[r-1] denominator2 += corr3*(r**2)*weights[r-1] _last_bw = bw res = numerator/(denominator1 * denominator2)**0.5 return frm_find_topn_angles_interp2(res)
def bart_align_vol(vf, wf, vg, wg, b, radius=None, peak_offset=None): """Implementation of Bartesaghi's approach for alignment. For detail, please check the paper. Parameters ---------- vf: The volume you want to match. pytom_volume.vol wf: The single tilt wedge information of volume vf. [missing_wedge_angle1, missing_wedge_angle2]. Note this is defined different with frm_align im frm.py! vg: The reference volume. pytom_volume.vol wg: The single tilt wedge information of volume vg. [missing_wedge_angle1, missing_wedge_angle2]. Note this is defined different with frm_align im frm.py! b: The bandwidth of spherical harmonics. Integer in the range [4, 64] radius: The maximal radius in the Fourier space, which is equal to say the maximal frequency involved in calculation. Integer. By default is half of the volume size. peak_offset: The maximal offset which allows the peak of the score to be. Or simply speaking, the maximal distance allowed to shift vg to match vf. This parameter is needed to prevent shifting the reference volume out of the frame. Integer. By default is half of the volume size. Returns ------- The best translation and rotation (Euler angle, ZXZ convention [Phi, Psi, Theta]) to transform vg to match vf. (best_translation, best_rotation, correlation_score) """ from pytom_volume import vol, rotateSpline, max, peak from pytom.basic.correlation import nXcf from pytom.basic.filter import lowpassFilter from pytom.basic.structures import WedgeInfo from pytom_volume import initSphere if not radius: # set the radius radius = vf.sizeX()/2 if peak_offset is None: peak_offset = vol(vf.sizeX(), vf.sizeY(), vf.sizeZ()) initSphere(peak_offset, vf.sizeX()/2, 0,0, vf.sizeX()/2,vf.sizeY()/2,vf.sizeZ()/2) elif isinstance(peak_offset, int): peak_radius = peak_offset peak_offset = vol(vf.sizeX(), vf.sizeY(), vf.sizeZ()) initSphere(peak_offset, peak_radius, 0,0, vf.sizeX()/2,vf.sizeY()/2,vf.sizeZ()/2) elif peak_offset.__class__ == vol: pass else: raise RuntimeError('Peak offset is given wrong!') from pytom.basic.fourier import fft, ifft, ftshift, iftshift from pytom_volume import vol, reducedToFull, rescale, abs, real from .vol2sf import vol2sf from pytom_numpy import vol2npy from math import log, ceil, pow # IMPORTANT!!! Should firstly do the IFFTSHIFT on the volume data (NOT FFTSHIFT since for odd-sized data it matters!), # and then followed by the FFT. ff = abs(ftshift(reducedToFull(fft(iftshift(vf, inplace=False))), inplace=False)) ff = real(ff) gg = abs(ftshift(reducedToFull(fft(iftshift(vg, inplace=False))), inplace=False)) gg = real(gg) sf = None sg = None mf = create_wedge_sf(wf[0], wf[1], b) mg = create_wedge_sf(wg[0], wg[1], b) for r in range(3, radius+1): # Should start from 3 since the interpolation in the first 2 bands is not accurate. if sf is None: sf = vol2sf(ff, r, b) sg = vol2sf(gg, r, b) else: sf += vol2sf(ff, r, b) sg += vol2sf(gg, r, b) corr = frm_constrained_corr(sf, mf, sg, mg) ang, val = frm_find_best_angle_interp(corr) tmp = vol(vg.sizeX(),vg.sizeY(),vg.sizeZ()) rotateSpline(vg, tmp, ang[0], ang[1], ang[2]) wedge_f = WedgeInfo(90+wf[0], 90-wf[1]) wedge_g = WedgeInfo(90+wg[0], 90-wg[1]) cc = nXcf(lowpassFilter(wedge_g.apply(vf), radius, 0)[0], lowpassFilter(wedge_f.apply(tmp), radius, 0)[0]) pos = peak(cc, peak_offset) pos, score = find_subpixel_peak_position(vol2npy(cc), pos) return (pos, ang, score)