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 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)
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) fr2 = vol(fr); rotateSpline(fr, fr2, phi, psi, the) fi2 = vol(fi); rotateSpline(fi, fi2, phi, psi, the) fr = np.array(vol2sf(fr, r, b)) fi = np.array(vol2sf(fi, r, b)) fr2 = np.array(vol2sf(fr2, r, b)) fi2 = np.array(vol2sf(fi2, r, b)) fr2 = fr2 * w
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)