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
def frm_constrained_align(vf, wf, vg, wg, b, max_freq, peak_offset=None, mask=None, constraint=None, weights=None, position=None, num_seeds=5, pytom_volume=None): """Find the best alignment (translation & rotation) of volume vg (reference) to match vf. For details, please check the paper. Parameters ---------- vf: Volume Nr. 1 pytom_volume.vol wf: Mask of vf in Fourier space. pytom.basic.structures.Wedge. If none, no missing wedge. vg: Volume Nr. 2 / Reference pytom_volume.vol wg: Mask of vg in Fourier space. pytom.basic.structures.Wedge. If none, no missing 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. 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. pytom_volume.vol / Integer. By default is half of the volume radius. mask: Mask volume for vg in real space. pytom_volume.vol constraint: Angular constraint sh_alignment.constrained_frm.AngularConstraint weights: Obsolete. position: If the translation is already known or not. If provided, no translational search will be conducted. List: [x,y,z], default None. num_seeds: Number of threads for the expectation maximization procedure. The more the better, yet slower. Integer, default is 5. 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, peak from pytom.basic.transformations import shift from pytom.basic.correlation import FLCF from pytom.basic.filter import lowpassFilter from pytom.basic.structures import Mask, SingleTiltWedge from pytom_volume import initSphere from pytom_numpy import vol2npy 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) if wg is None: wg = SingleTiltWedge(0) if peak_offset is None: peak_offset = vol(vf.sizeX(), vf.sizeY(), vf.sizeZ()) initSphere(peak_offset, vf.sizeX() / 4, 0, 0, vf.sizeX() / 2, vf.sizeY() / 2, vf.sizeZ() / 2) elif peak_offset.__class__ == 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!') # cut out the outer part which normally contains nonsense m = vol(vf.sizeX(), vf.sizeY(), vf.sizeZ()) initSphere(m, vf.sizeX() / 2, 0, 0, vf.sizeX() / 2, vf.sizeY() / 2, vf.sizeZ() / 2) vf = vf * m vg = vg * m if mask is None: mask = m else: vg = vg * mask if position is None: # if position is not given, we have to find it ourself # first roughtly determine the orientation (only according to the energy info) # get multiple candidate orientations numerator, denominator1, denominator2 = frm_correlate( vf, wf, vg, wg, b, max_freq, weights, True, None, None, False) score = numerator / (denominator1 * denominator2)**0.5 res = frm_find_topn_constrained_angles_interp( score, num_seeds, get_adaptive_bw(max_freq, b) / 16., constraint) else: # the position is given by the user vf2 = shift(vf, -position[0] + vf.sizeX() / 2, -position[1] + vf.sizeY() / 2, -position[2] + vf.sizeZ() / 2, 'fourier') score = frm_correlate(vf2, wf, vg, wg, b, max_freq, weights, ps=False) orientation, max_value = frm_find_best_constrained_angle_interp( score, constraint=constraint) return position, orientation, max_value # iteratively refine the position & orientation from pytom.tools.maths import euclidianDistance max_iter = 10 # maximal number of iterations mask2 = vol(mask.sizeX(), mask.sizeY(), mask.sizeZ()) # store the rotated mask vg2 = vol(vg.sizeX(), vg.sizeY(), vg.sizeZ()) lowpass_vf = lowpassFilter(vf, max_freq, max_freq / 10.)[0] max_position = None max_orientation = None max_value = -1.0 for i in xrange(num_seeds): old_pos = [-1, -1, -1] lm_pos = [-1, -1, -1] lm_ang = None lm_value = -1.0 orientation = res[i][0] # initial orientation for j in xrange(max_iter): rotateSpline(vg, vg2, orientation[0], orientation[1], orientation[2]) # first rotate rotateSpline(mask, mask2, orientation[0], orientation[1], orientation[2]) # rotate the mask as well vg2 = wf.apply(vg2) # then apply the wedge vg2 = lowpassFilter(vg2, max_freq, max_freq / 10.)[0] score = FLCF(lowpass_vf, vg2, mask2) # find the position pos = peak(score, peak_offset) pos, val = find_subpixel_peak_position(vol2npy(score), pos) if val > lm_value: lm_pos = pos lm_ang = orientation lm_value = val if euclidianDistance(lm_pos, old_pos) <= 1.0: # terminate this thread if lm_value > max_value: max_position = lm_pos max_orientation = lm_ang max_value = lm_value break else: old_pos = lm_pos # here we shift the target, not the reference # if you dont want the shift to change the energy landscape, use fourier shift vf2 = shift(vf, -lm_pos[0] + vf.sizeX() / 2, -lm_pos[1] + vf.sizeY() / 2, -lm_pos[2] + vf.sizeZ() / 2, 'fourier') score = frm_correlate(vf2, wf, vg, wg, b, max_freq, weights, False, denominator1, denominator2, True) orientation, val = frm_find_best_constrained_angle_interp( score, constraint=constraint) else: # no converge after the specified iteration, still get the best result as we can if lm_value > max_value: max_position = lm_pos max_orientation = lm_ang max_value = lm_value # print max_value # for show the convergence of the algorithm return max_position, max_orientation, max_value
if len(sys.argv) == 1: print(helper) sys.exit() try: dir_name, name_prefix, wedge_angle, output, bHelp = parse_script_options( sys.argv[1:], helper) except: sys.exit() if bHelp is True: print(helper) sys.exit() if name_prefix is None: name_prefix = '' wedge_angle = float(wedge_angle) w = SingleTiltWedge(wedge_angle) pl = ParticleList() all_files = os.listdir(dir_name) for fname in all_files: p = None name_suffix = fname.split('.')[-1] if len(name_prefix): if name_prefix in fname and name_suffix in ['em', 'mrc']: p = Particle(dir_name + '/' + fname) else: if name_suffix in ['em', 'mrc']: p = Particle(dir_name + '/' + fname) if p is not None: pl.append(p)
from pytom.tools.parse_script_options import parse_script_options from pytom.basic.structures import ParticleList, SingleTiltWedge helper = ScriptHelper(sys.argv[0].split('/')[-1], # script name description='Set the same wedge to all the particles in the particle list.', authors='Yuxiang Chen', options=[ScriptOption(['-p'], 'Particle list', True, False), ScriptOption(['-w'], 'Wedge Angle (degree)', True, False), ScriptOption(['-o'], 'Output particle list', True, True), ScriptOption(['-h', '--help'], 'Help.', False, True)]) if len(sys.argv) == 1: print(helper) sys.exit() try: pl_name, wedge_angle, output, bHelp = parse_script_options(sys.argv[1:], helper) except: sys.exit() if bHelp is True: print(helper) sys.exit() pl = ParticleList() pl.fromXMLFile(pl_name) w = SingleTiltWedge(int(wedge_angle)) pl.setWedgeAllParticles(w) if output is None: pl.toXMLFile(pl_name) else: pl.toXMLFile(output)