def distance(p, ref, freq, mask, binning):
from pytom.basic.correlation import nxcc
from pytom_volume import vol, initSphere, read, pasteCenter
from pytom.basic.filter import lowpassFilter
from pytom.basic.transformations import resize
v = p.getTransformedVolume(binning)
w = p.getWedge()
r = ref.getVolume()
a = lowpassFilter(w.apply(v, p.getRotation().invert()), freq)[0]
b = lowpassFilter(w.apply(r, p.getRotation().invert()), freq)[0]
if not mask:
mask = vol(r)
initSphere(mask,
r.sizeX() // 2 - 3, 3, 0,
r.sizeX() // 2,
r.sizeY() // 2,
r.sizeZ() // 2)
else:
#THE MASK is binning (sampled every n-points). This does lead to a reduction of the smoothing of the edges.
maskBin = read(mask, 0, 0, 0, 0, 0, 0, 0, 0, 0, binning, binning,
binning)
if a.sizeX() != maskBin.sizeX() or a.sizeY() != maskBin.sizeY(
) or a.sizeZ() != maskBin.sizeZ():
mask = vol(a.sizeX(), a.sizeY(), a.sizeZ())
mask.setAll(0)
pasteCenter(maskBin, mask)
else:
mask = maskBin
s = nxcc(a, b, mask)
d2 = 2 * (1 - s)
return d2
def focus_score(p, ref, freq, diff_mask, binning):
from pytom.basic.correlation import nxcc
from pytom.basic.filter import lowpassFilter
v = p.getTransformedVolume(binning)
w = p.getWedge()
r = ref.getVolume()
a = lowpassFilter(w.apply(v, p.getRotation().invert()), freq)[0]
b = lowpassFilter(w.apply(r, p.getRotation().invert()), freq)[0]
s = nxcc(a, b, diff_mask.getVolume())
return s
def score_noalign_proxy(p, ref, freq, offset, binning, mask):
from pytom.basic.structures import Shift, Rotation
from pytom.basic.correlation import nxcc
from pytom.basic.filter import lowpassFilter
v = p.getTransformedVolume(binning)
w = p.getWedge()
r = ref.getVolume()
a = lowpassFilter(w.apply(v, p.getRotation().invert()), freq)[0]
b = lowpassFilter(w.apply(r, p.getRotation().invert()), freq)[0]
score = nxcc(a, b)
return (p.getShift(), p.getRotation(), score, p.getFilename())
def sag_fine_grained_alignment(vf,
wf,
vg,
wg,
max_freq,
ang=[0, 0, 0],
loc=[0, 0, 0],
mask=None,
B,
alpha,
maxIter,
lambda1):
"""SAG-based fine-grained alignment between experimental data and reference data.
Parameters
vf: Experimental data
pytom_volume.vol
wf: Mask of vf in Fourier space.
pytom.basic.structures.Wedge. If none, no missing wedge.
vg: Reference data
pytom_volume.vol
wg: Mask of vg in Fourier space.
pytom.basic.structures.Wedge. If none, no missing wedge.
max_freq: Maximal frequency involved in calculation.
Integer.
ang: Initial rotation angle
loc: Initial translation value
mask: Mask volume in real space.
pytom_volume.vol
B: Batch number
alpha: Step size
maxIter: The max iteration number
lambda1: Regularization parameter
Returns
-------
(Optimal rotation angle and translation value.
(best_translation, best_rotation, correlation_score)
"""
from pytom_volume import vol, rotateSpline, peak, sum, power
from pytom.basic.transformations import shift
from pytom.basic.filter import lowpassFilter
from pytom.basic.structures import Mask, SingleTiltWedge, Rotation
from pytom_volume import initSphere
from pytom_numpy import vol2npy
import math
import random
if vf.sizeX() != vg.sizeX() or vf.sizeY() != vg.sizeY() or vf.sizeZ(
) != vg.sizeZ():
raise RuntimeError('Two volumes must have the same size!')
if wf is None:
wf = SingleTiltWedge(0)
else:
vf = wf.apply(vf)
if wg is None:
wg = SingleTiltWedge(0)
else:
vg = wg.apply(vg)
if mask is None:
m = vol(vf.sizeX(), vf.sizeY(), vf.sizeZ())
initSphere(m,
vf.sizeX() / 2, 0, 0,
vf.sizeX() / 2,
vf.sizeY() / 2,
vf.sizeZ() / 2)
mask = m
old_value = -1
max_pos = [-1, -1, -1]
max_ang = None
max_value = -1.0
ang_epsilon = np.ones(3) * (math.pi * (1.0 / 180))
loc_epsilon = np.ones(3) * 1.0
n = vf.sizeX()
vf0_n = vol2npy(vf)
if maxIter is None:
maxIter = n / 2
iVals = np.int32(np.ceil((n - B) * np.random.random(maxIter)))
if lambda1 is None:
lambda1 = 1 / n
eps = np.finfo(np.float32).eps
Lmax = 0.25 * np.max(np.sum(vf0_n**2)) + lambda1
if alpha is None:
alpha = 1 / Lmax
d = np.zeros(6)
g = np.zeros([n, 6])
covered = np.int32(np.zeros(n))
nCovered = 0
grad = np.zeros(6)
deri = np.zeros(6)
vg2 = vol(vf.sizeX(), vf.sizeY(), vf.sizeZ())
mask2 = vol(mask.sizeX(), mask.sizeY(), mask.sizeZ())
for i in range(n):
if (covered[i] != 0):
nCovered += 1
for k in range(maxIter):
i = iVals[k] - 1
if k == 0:
rotateSpline(vg, vg2, ang[0], ang[1], ang[2])
rotateSpline(mask, mask2, ang[0], ang[1], ang[2])
vg2 = wf.apply(vg2)
vg2 = lowpassFilter(vg2, max_freq, max_freq / 10.)[0]
vg2_s = transform_single_vol(vg2, mask2)
vf2 = shift(vf,
-loc[0] + vf.sizeX() / 2,
-loc[1] + vf.sizeY() / 2,
-loc[2] + vf.sizeZ() / 2,
imethod='spline')
vf2 = lowpassFilter(vf2, max_freq, max_freq / 10.)[0]
vf2 = wg.apply(vf2, Rotation(ang))
vf2_s = transform_single_vol(vf2, mask2)
i = 0
ri = np.sum(
vol2npy(vf2_s)[i:i + B, :, :] - vol2npy(vg2_s)[i:i + B, :, :])
vg2_p = vol(n, n, n)
vg2_m = vol(n, n, n)
mask2_p = vol(n, n, n)
mask2_m = vol(n, n, n)
for dim_i in range(3):
if abs(ang_epsilon[dim_i]) > eps:
ang_epsilon_t = np.zeros(3)
ang_epsilon_t[dim_i] = ang_epsilon[dim_i]
angle = ang + ang_epsilon_t
rotateSpline(vg, vg2_p, angle[0], angle[1], angle[2])
rotateSpline(mask, mask2_p, angle[0], angle[1], angle[2])
vg2_p = wf.apply(vg2_p)
vg2_p = lowpassFilter(vg2_p, max_freq, max_freq / 10.)[0]
vg2_pf = transform_single_vol(vg2_p, mask2_p)
angle = ang - ang_epsilon_t
rotateSpline(vg, vg2_m, angle[0], angle[1], angle[2])
rotateSpline(mask, mask2_m, angle[0], angle[1], angle[2])
vg2_m = wf.apply(vg2_m)
vg2_m = lowpassFilter(vg2_m, max_freq, max_freq / 10.)[0]
vg2_mf = transform_single_vol(vg2_m, mask2_m)
vg2_ang_deri = (vg2_pf - vg2_mf) / (2 * ang_epsilon[dim_i])
vg2_ang_deri_n = vol2npy(vg2_ang_deri)
deri[dim_i] = np.sum(vg2_ang_deri_n[i:i + B, :, :])
del vg2_pf, vg2_mf, vg2_ang_deri, vg2_ang_deri_n, angle
del vg2_p, vg2_m, mask2_p, mask2_m
vf1_ps = vol(n, n, n)
vf1_ms = vol(n, n, n)
ang_f = [ang[0], ang[1], ang[2]]
for dim_i in range(3):
if abs(loc_epsilon[dim_i]) > eps:
loc_epsilon_t = np.zeros(3)
loc_epsilon_t[dim_i] = ang_epsilon[dim_i]
vf1_ps.copyVolume(vf)
vf1_ms.copyVolume(vf)
loc_r = loc + loc_epsilon_t
vf1_tp = shift(vf1_ps, -loc_r[0] + vf1_ps.sizeX() / 2,
-loc_r[1] + vf1_ps.sizeY() / 2,
-loc_r[2] + vf1_ps.sizeZ() / 2, 'spline')
vf1_tp = lowpassFilter(vf1_tp, max_freq, max_freq / 10.)[0]
vf1_tp = wg.apply(vf1_tp, Rotation(ang_f))
loc_r = loc - loc_epsilon_t
vf1_tm = shift(vf1_ms, -loc_r[0] + vf1_ms.sizeX() / 2,
-loc_r[1] + vf1_ms.sizeY() / 2,
-loc_r[2] + vf1_ms.sizeZ() / 2, 'spline')
vf1_tm = lowpassFilter(vf1_tm, max_freq, max_freq / 10.)[0]
vf1_tm = wg.apply(vf1_tm, Rotation(ang_f))
vf1_tpf = transform_single_vol(vf1_tp, mask2)
vf1_tmf = transform_single_vol(vf1_tm, mask2)
vf1_loc_deri = (vf1_tpf - vf1_tmf) / (2 * ang_epsilon[dim_i])
vf1_loc_deri_n = vol2npy(vf1_loc_deri)
deri[dim_i + 3] = np.sum(vf1_loc_deri_n[i:i + B, :, :])
del vf1_tp, vf1_tm, vf1_tpf, vf1_tmf, vf1_loc_deri, vf1_loc_deri_n
del vf1_ps, vf1_ms, ang_f
for dim_i in range(6):
grad[dim_i] = ri * deri[dim_i] / B
for dim_i in range(6):
d[dim_i] += grad[dim_i] - np.sum(g[i:i + B, dim_i])
for dim_i in range(6):
g[i:i + B, dim_i] = grad[dim_i]
for j0 in range(i, i + B + 1):
if (covered[j0] == 0):
covered[j0] = 1
nCovered += 1
for dim_i in range(6):
opt_beta[dim_i] -= alpha * d[dim_i] / nCovered
ang = opt_beta[:3]
loc = opt_beta[3:]
rotateSpline(vg, vg2, ang[0], ang[1], ang[2])
rotateSpline(mask, mask2, ang[0], ang[1], ang[2])
vg2 = wf.apply(vg2)
vg2 = lowpassFilter(vg2, max_freq, max_freq / 10.)[0]
vg2_s = transform_single_vol(vg2, mask2)
vf2 = shift(vf,
-loc[0] + vf.sizeX() / 2,
-loc[1] + vf.sizeY() / 2,
-loc[2] + vf.sizeZ() / 2,
imethod='spline')
vf2 = lowpassFilter(vf2, max_freq, max_freq / 10.)[0]
vf2 = wg.apply(vf2, Rotation(ang))
vf2_s = transform_single_vol(vf2, mask2)
ri = np.sum(
vol2npy(vf2_s)[i:i + B, :, :] - vol2npy(vg2_s)[i:i + B, :, :])
from pytom.basic.correlation import nxcc
val = nxcc(vf2_s, vg2_s, mask)
if val > max_value:
max_pos = loc
max_ang = ang
max_value = val
if abs(max_value - old_value) <= eps:
break
else:
old_value = max_value
del vg2_s, vf2, vf2_s
del d, g, grad, deri
return max_pos, max_ang, max_value
# apply wedge
v2 = wedge.apply(v2)
t = timing()
# old method
angles.reset()
t.start()
tmp = vol(v)
peak = 0.0
angle = angles.nextRotation()
while angle != [None, None, None]:
rotateSpline(v, tmp, angle[0], angle[1], angle[2])
tmp = wedge.apply(tmp)
res = nxcc(v2, tmp, mask)
if res >= peak:
peak = res
ang = angle
angle = angles.nextRotation()
t1 = t.end()
dist_old = rotation_distance([phi, psi, the], ang)
# if dist_old > 30:
# print [phi, psi, the], ang
diff_old.append(dist_old)
total_time1 += t1
# new method
t.start()
res = frm_fourier_constrained_vol(v2, w, v, m)
def run(self, verbose=False):
from pytom_volume import read, sum
from pytom.basic.filter import lowpassFilter
from pytom.basic.correlation import nxcc
from pytom.basic.structures import Rotation
from pytom.tools.ProgressBar import FixedProgBar
while True:
# get the job
job = self.get_job()
try:
pairs = job["Pairs"]
pl_filename = job["ParticleList"]
except:
if verbose:
print(self.node_name + ': end')
break # get some non-job message, break it
from pytom.basic.structures import ParticleList
pl = ParticleList('.')
pl.fromXMLFile(pl_filename)
if verbose:
prog = FixedProgBar(0, len(pairs)-1, self.node_name+':')
i = 0
# run the job
result = {}
last_filename = None
binning = int(job["Binning"])
mask = read(job["Mask"], 0, 0, 0, 0, 0, 0, 0, 0, 0, binning, binning, binning)
for pair in pairs:
if verbose:
prog.update(i)
i += 1
g = pl[pair[0]]
f = pl[pair[1]]
vf = f.getTransformedVolume(binning)
wf = f.getWedge().getWedgeObject()
wf_rotation = f.getRotation().invert()
# wf.setRotation(Rotation(-rotation[1],-rotation[0],-rotation[2]))
# wf_vol = wf.returnWedgeVolume(vf.sizeX(), vf.sizeY(), vf.sizeZ(), True, -rotation[1],-rotation[0],-rotation[2])
vf = lowpassFilter(vf, job["Frequency"], 0)[0]
if g.getFilename() != last_filename:
vg = g.getTransformedVolume(binning)
wg = g.getWedge().getWedgeObject()
wg_rotation = g.getRotation().invert()
# wg.setRotation(Rotation(-rotation[1],-rotation[0],-rotation[2]))
# wg_vol = wg.returnWedgeVolume(vg.sizeX(), vg.sizeY(), vg.sizeZ(), True, -rotation[1],-rotation[0],-rotation[2])
vg = lowpassFilter(vg, job["Frequency"], 0)[0]
last_filename = g.getFilename()
score = nxcc( wg.apply(vf, wg_rotation), wf.apply(vg, wf_rotation), mask)
# overlapped_wedge_vol = wf_vol * wg_vol
# scaling = float(overlapped_wedge_vol.numelem())/sum(overlapped_wedge_vol)
# score *= scaling
result[pair] = score
# send back the result
self.send_result(result)
pytom_mpi.finalise()
from pytom.basic.correlation import nxcc
from pytom.basic.files import read
print(nxcc(read('GPU/vol_1_3.mrc'), read('GPU/vol_3_1.mrc'), read('CPU/mask.em')))
print(nxcc(read('CPU/vol_1_3.em'), read('CPU/vol_3_1.em'), read('CPU/mask.em')))
from pytom.tompy.correlation import nxcc
from pytom.tompy.io import read as readN
print(nxcc(readN('GPU/vol_1_3.mrc'), readN('GPU/vol_3_1.mrc'), readN('CPU/mask.em')))
print(nxcc(readN('CPU/vol_1_3.em'), readN('CPU/vol_3_1.em'), readN('CPU/mask.em')))
# old method
angles.reset()
t.start()
# res = extractPeaks(v2, v, angles, None, m, True, None, verboseMode=False)
# pos = peak(res[0], m)
# ang_idx = res[1].getV(pos[0], pos[1], pos[2])
# ang = angles.getRotations()[int(ang_idx)]
tmp = vol(v)
peak = 0.0
angle = angles.nextRotation()
while angle != [None, None, None]:
rotateSpline(v, tmp, angle[0], angle[1], angle[2])
res = nxcc(v2, tmp, m)
if res >= peak:
peak = res
ang = angle
angle = angles.nextRotation()
t1 = t.end()
dist_old = rotation_distance([phi, psi, the], ang)
diff_old.append(dist_old)
total_time1 += t1
# new method
t.start()
ang2 = frm_get_best_angle(v2, v, 32, None, False)
t2 = t.end()