def shiftali_MPI(stack, maskfile=None, maxit=100, CTF=False, snr=1.0, Fourvar=False, search_rng=-1, oneDx=False, search_rng_y=-1): number_of_proc = mpi.mpi_comm_size(mpi.MPI_COMM_WORLD) myid = mpi.mpi_comm_rank(mpi.MPI_COMM_WORLD) main_node = 0 ftp = file_type(stack) if myid == main_node: print_begin_msg("shiftali_MPI") max_iter = int(maxit) if myid == main_node: if ftp == "bdb": from EMAN2db import db_open_dict dummy = db_open_dict(stack, True) nima = EMUtil.get_image_count(stack) else: nima = 0 nima = bcast_number_to_all(nima, source_node=main_node) list_of_particles = list(range(nima)) image_start, image_end = MPI_start_end(nima, number_of_proc, myid) list_of_particles = list_of_particles[image_start:image_end] # read nx and ctf_app (if CTF) and broadcast to all nodes if myid == main_node: ima = EMData() ima.read_image(stack, list_of_particles[0], True) nx = ima.get_xsize() ny = ima.get_ysize() if CTF: ctf_app = ima.get_attr_default('ctf_applied', 2) del ima else: nx = 0 ny = 0 if CTF: ctf_app = 0 nx = bcast_number_to_all(nx, source_node=main_node) ny = bcast_number_to_all(ny, source_node=main_node) if CTF: ctf_app = bcast_number_to_all(ctf_app, source_node=main_node) if ctf_app > 0: ERROR("data cannot be ctf-applied", myid=myid) if maskfile == None: mrad = min(nx, ny) mask = model_circle(mrad // 2 - 2, nx, ny) else: mask = get_im(maskfile) if CTF: from sp_filter import filt_ctf from sp_morphology import ctf_img ctf_abs_sum = EMData(nx, ny, 1, False) ctf_2_sum = EMData(nx, ny, 1, False) else: ctf_2_sum = None from sp_global_def import CACHE_DISABLE if CACHE_DISABLE: data = EMData.read_images(stack, list_of_particles) else: for i in range(number_of_proc): if myid == i: data = EMData.read_images(stack, list_of_particles) if ftp == "bdb": mpi.mpi_barrier(mpi.MPI_COMM_WORLD) for im in range(len(data)): data[im].set_attr('ID', list_of_particles[im]) st = Util.infomask(data[im], mask, False) data[im] -= st[0] if CTF: ctf_params = data[im].get_attr("ctf") ctfimg = ctf_img(nx, ctf_params, ny=ny) Util.add_img2(ctf_2_sum, ctfimg) Util.add_img_abs(ctf_abs_sum, ctfimg) if CTF: reduce_EMData_to_root(ctf_2_sum, myid, main_node) reduce_EMData_to_root(ctf_abs_sum, myid, main_node) else: ctf_2_sum = None if CTF: if myid != main_node: del ctf_2_sum del ctf_abs_sum else: temp = EMData(nx, ny, 1, False) for i in range(0, nx, 2): for j in range(ny): temp.set_value_at(i, j, snr) Util.add_img(ctf_2_sum, temp) del temp total_iter = 0 # apply initial xform.align2d parameters stored in header init_params = [] for im in range(len(data)): t = data[im].get_attr('xform.align2d') init_params.append(t) p = t.get_params("2d") data[im] = rot_shift2D(data[im], p['alpha'], sx=p['tx'], sy=p['ty'], mirror=p['mirror'], scale=p['scale']) # fourier transform all images, and apply ctf if CTF for im in range(len(data)): if CTF: ctf_params = data[im].get_attr("ctf") data[im] = filt_ctf(fft(data[im]), ctf_params) else: data[im] = fft(data[im]) sx_sum = 0 sy_sum = 0 sx_sum_total = 0 sy_sum_total = 0 shift_x = [0.0] * len(data) shift_y = [0.0] * len(data) ishift_x = [0.0] * len(data) ishift_y = [0.0] * len(data) for Iter in range(max_iter): if myid == main_node: start_time = time() print_msg("Iteration #%4d\n" % (total_iter)) total_iter += 1 avg = EMData(nx, ny, 1, False) for im in data: Util.add_img(avg, im) reduce_EMData_to_root(avg, myid, main_node) if myid == main_node: if CTF: tavg = Util.divn_filter(avg, ctf_2_sum) else: tavg = Util.mult_scalar(avg, 1.0 / float(nima)) else: tavg = EMData(nx, ny, 1, False) if Fourvar: bcast_EMData_to_all(tavg, myid, main_node) vav, rvar = varf2d_MPI(myid, data, tavg, mask, "a", CTF) if myid == main_node: if Fourvar: tavg = fft(Util.divn_img(fft(tavg), vav)) vav_r = Util.pack_complex_to_real(vav) # normalize and mask tavg in real space tavg = fft(tavg) stat = Util.infomask(tavg, mask, False) tavg -= stat[0] Util.mul_img(tavg, mask) # For testing purposes: shift tavg to some random place and see if the centering is still correct #tavg = rot_shift3D(tavg,sx=3,sy=-4) tavg = fft(tavg) if Fourvar: del vav bcast_EMData_to_all(tavg, myid, main_node) sx_sum = 0 sy_sum = 0 if search_rng > 0: nwx = 2 * search_rng + 1 else: nwx = nx if search_rng_y > 0: nwy = 2 * search_rng_y + 1 else: nwy = ny not_zero = 0 for im in range(len(data)): if oneDx: ctx = Util.window(ccf(data[im], tavg), nwx, 1) p1 = peak_search(ctx) p1_x = -int(p1[0][3]) ishift_x[im] = p1_x sx_sum += p1_x else: p1 = peak_search(Util.window(ccf(data[im], tavg), nwx, nwy)) p1_x = -int(p1[0][4]) p1_y = -int(p1[0][5]) ishift_x[im] = p1_x ishift_y[im] = p1_y sx_sum += p1_x sy_sum += p1_y if not_zero == 0: if (not (ishift_x[im] == 0.0)) or (not (ishift_y[im] == 0.0)): not_zero = 1 sx_sum = mpi.mpi_reduce(sx_sum, 1, mpi.MPI_INT, mpi.MPI_SUM, main_node, mpi.MPI_COMM_WORLD) if not oneDx: sy_sum = mpi.mpi_reduce(sy_sum, 1, mpi.MPI_INT, mpi.MPI_SUM, main_node, mpi.MPI_COMM_WORLD) if myid == main_node: sx_sum_total = int(sx_sum[0]) if not oneDx: sy_sum_total = int(sy_sum[0]) else: sx_sum_total = 0 sy_sum_total = 0 sx_sum_total = bcast_number_to_all(sx_sum_total, source_node=main_node) if not oneDx: sy_sum_total = bcast_number_to_all(sy_sum_total, source_node=main_node) sx_ave = round(float(sx_sum_total) / nima) sy_ave = round(float(sy_sum_total) / nima) for im in range(len(data)): p1_x = ishift_x[im] - sx_ave p1_y = ishift_y[im] - sy_ave params2 = { "filter_type": Processor.fourier_filter_types.SHIFT, "x_shift": p1_x, "y_shift": p1_y, "z_shift": 0.0 } data[im] = Processor.EMFourierFilter(data[im], params2) shift_x[im] += p1_x shift_y[im] += p1_y # stop if all shifts are zero not_zero = mpi.mpi_reduce(not_zero, 1, mpi.MPI_INT, mpi.MPI_SUM, main_node, mpi.MPI_COMM_WORLD) if myid == main_node: not_zero_all = int(not_zero[0]) else: not_zero_all = 0 not_zero_all = bcast_number_to_all(not_zero_all, source_node=main_node) if myid == main_node: print_msg("Time of iteration = %12.2f\n" % (time() - start_time)) start_time = time() if not_zero_all == 0: break #for im in xrange(len(data)): data[im] = fft(data[im]) This should not be required as only header information is used # combine shifts found with the original parameters for im in range(len(data)): t0 = init_params[im] t1 = Transform() t1.set_params({ "type": "2D", "alpha": 0, "scale": t0.get_scale(), "mirror": 0, "tx": shift_x[im], "ty": shift_y[im] }) # combine t0 and t1 tt = t1 * t0 data[im].set_attr("xform.align2d", tt) # write out headers and STOP, under MPI writing has to be done sequentially mpi.mpi_barrier(mpi.MPI_COMM_WORLD) par_str = ["xform.align2d", "ID"] if myid == main_node: from sp_utilities import file_type if (file_type(stack) == "bdb"): from sp_utilities import recv_attr_dict_bdb recv_attr_dict_bdb(main_node, stack, data, par_str, image_start, image_end, number_of_proc) else: from sp_utilities import recv_attr_dict recv_attr_dict(main_node, stack, data, par_str, image_start, image_end, number_of_proc) else: send_attr_dict(main_node, data, par_str, image_start, image_end) if myid == main_node: print_end_msg("shiftali_MPI")
def helicalshiftali_MPI(stack, maskfile=None, maxit=100, CTF=False, snr=1.0, Fourvar=False, search_rng=-1): nproc = mpi.mpi_comm_size(mpi.MPI_COMM_WORLD) myid = mpi.mpi_comm_rank(mpi.MPI_COMM_WORLD) main_node = 0 ftp = file_type(stack) if myid == main_node: print_begin_msg("helical-shiftali_MPI") max_iter = int(maxit) if (myid == main_node): infils = EMUtil.get_all_attributes(stack, "filament") ptlcoords = EMUtil.get_all_attributes(stack, 'ptcl_source_coord') filaments = ordersegments(infils, ptlcoords) total_nfils = len(filaments) inidl = [0] * total_nfils for i in range(total_nfils): inidl[i] = len(filaments[i]) linidl = sum(inidl) nima = linidl tfilaments = [] for i in range(total_nfils): tfilaments += filaments[i] del filaments else: total_nfils = 0 linidl = 0 total_nfils = bcast_number_to_all(total_nfils, source_node=main_node) if myid != main_node: inidl = [-1] * total_nfils inidl = bcast_list_to_all(inidl, myid, source_node=main_node) linidl = bcast_number_to_all(linidl, source_node=main_node) if myid != main_node: tfilaments = [-1] * linidl tfilaments = bcast_list_to_all(tfilaments, myid, source_node=main_node) filaments = [] iendi = 0 for i in range(total_nfils): isti = iendi iendi = isti + inidl[i] filaments.append(tfilaments[isti:iendi]) del tfilaments, inidl if myid == main_node: print_msg("total number of filaments: %d" % total_nfils) if total_nfils < nproc: ERROR( 'number of CPUs (%i) is larger than the number of filaments (%i), please reduce the number of CPUs used' % (nproc, total_nfils), myid=myid) # balanced load temp = chunks_distribution([[len(filaments[i]), i] for i in range(len(filaments))], nproc)[myid:myid + 1][0] filaments = [filaments[temp[i][1]] for i in range(len(temp))] nfils = len(filaments) #filaments = [[0,1]] #print "filaments",filaments list_of_particles = [] indcs = [] k = 0 for i in range(nfils): list_of_particles += filaments[i] k1 = k + len(filaments[i]) indcs.append([k, k1]) k = k1 data = EMData.read_images(stack, list_of_particles) ldata = len(data) sxprint("ldata=", ldata) nx = data[0].get_xsize() ny = data[0].get_ysize() if maskfile == None: mrad = min(nx, ny) // 2 - 2 mask = pad(model_blank(2 * mrad + 1, ny, 1, 1.0), nx, ny, 1, 0.0) else: mask = get_im(maskfile) # apply initial xform.align2d parameters stored in header init_params = [] for im in range(ldata): t = data[im].get_attr('xform.align2d') init_params.append(t) p = t.get_params("2d") data[im] = rot_shift2D(data[im], p['alpha'], p['tx'], p['ty'], p['mirror'], p['scale']) if CTF: from sp_filter import filt_ctf from sp_morphology import ctf_img ctf_abs_sum = EMData(nx, ny, 1, False) ctf_2_sum = EMData(nx, ny, 1, False) else: ctf_2_sum = None ctf_abs_sum = None from sp_utilities import info for im in range(ldata): data[im].set_attr('ID', list_of_particles[im]) st = Util.infomask(data[im], mask, False) data[im] -= st[0] if CTF: ctf_params = data[im].get_attr("ctf") qctf = data[im].get_attr("ctf_applied") if qctf == 0: data[im] = filt_ctf(fft(data[im]), ctf_params) data[im].set_attr('ctf_applied', 1) elif qctf != 1: ERROR('Incorrectly set qctf flag', myid=myid) ctfimg = ctf_img(nx, ctf_params, ny=ny) Util.add_img2(ctf_2_sum, ctfimg) Util.add_img_abs(ctf_abs_sum, ctfimg) else: data[im] = fft(data[im]) del list_of_particles if CTF: reduce_EMData_to_root(ctf_2_sum, myid, main_node) reduce_EMData_to_root(ctf_abs_sum, myid, main_node) if CTF: if myid != main_node: del ctf_2_sum del ctf_abs_sum else: temp = EMData(nx, ny, 1, False) tsnr = 1. / snr for i in range(0, nx + 2, 2): for j in range(ny): temp.set_value_at(i, j, tsnr) temp.set_value_at(i + 1, j, 0.0) #info(ctf_2_sum) Util.add_img(ctf_2_sum, temp) #info(ctf_2_sum) del temp total_iter = 0 shift_x = [0.0] * ldata for Iter in range(max_iter): if myid == main_node: start_time = time() print_msg("Iteration #%4d\n" % (total_iter)) total_iter += 1 avg = EMData(nx, ny, 1, False) for im in range(ldata): Util.add_img(avg, fshift(data[im], shift_x[im])) reduce_EMData_to_root(avg, myid, main_node) if myid == main_node: if CTF: tavg = Util.divn_filter(avg, ctf_2_sum) else: tavg = Util.mult_scalar(avg, 1.0 / float(nima)) else: tavg = model_blank(nx, ny) if Fourvar: bcast_EMData_to_all(tavg, myid, main_node) vav, rvar = varf2d_MPI(myid, data, tavg, mask, "a", CTF) if myid == main_node: if Fourvar: tavg = fft(Util.divn_img(fft(tavg), vav)) vav_r = Util.pack_complex_to_real(vav) # normalize and mask tavg in real space tavg = fft(tavg) stat = Util.infomask(tavg, mask, False) tavg -= stat[0] Util.mul_img(tavg, mask) tavg.write_image("tavg.hdf", Iter) # For testing purposes: shift tavg to some random place and see if the centering is still correct #tavg = rot_shift3D(tavg,sx=3,sy=-4) if Fourvar: del vav bcast_EMData_to_all(tavg, myid, main_node) tavg = fft(tavg) sx_sum = 0.0 nxc = nx // 2 for ifil in range(nfils): """ # Calculate filament average avg = EMData(nx, ny, 1, False) filnima = 0 for im in xrange(indcs[ifil][0], indcs[ifil][1]): Util.add_img(avg, data[im]) filnima += 1 tavg = Util.mult_scalar(avg, 1.0/float(filnima)) """ # Calculate 1D ccf between each segment and filament average nsegms = indcs[ifil][1] - indcs[ifil][0] ctx = [None] * nsegms pcoords = [None] * nsegms for im in range(indcs[ifil][0], indcs[ifil][1]): ctx[im - indcs[ifil][0]] = Util.window(ccf(tavg, data[im]), nx, 1) pcoords[im - indcs[ifil][0]] = data[im].get_attr( 'ptcl_source_coord') #ctx[im-indcs[ifil][0]].write_image("ctx.hdf",im-indcs[ifil][0]) #print " CTX ",myid,im,Util.infomask(ctx[im-indcs[ifil][0]], None, True) # search for best x-shift cents = nsegms // 2 dst = sqrt( max((pcoords[cents][0] - pcoords[0][0])**2 + (pcoords[cents][1] - pcoords[0][1])**2, (pcoords[cents][0] - pcoords[-1][0])**2 + (pcoords[cents][1] - pcoords[-1][1])**2)) maxincline = atan2(ny // 2 - 2 - float(search_rng), dst) kang = int(dst * tan(maxincline) + 0.5) #print " settings ",nsegms,cents,dst,search_rng,maxincline,kang # ## C code for alignment. @ming results = [0.0] * 3 results = Util.helixshiftali(ctx, pcoords, nsegms, maxincline, kang, search_rng, nxc) sib = int(results[0]) bang = results[1] qm = results[2] #print qm, sib, bang # qm = -1.e23 # # for six in xrange(-search_rng, search_rng+1,1): # q0 = ctx[cents].get_value_at(six+nxc) # for incline in xrange(kang+1): # qt = q0 # qu = q0 # if(kang>0): tang = tan(maxincline/kang*incline) # else: tang = 0.0 # for kim in xrange(cents+1,nsegms): # dst = sqrt((pcoords[cents][0] - pcoords[kim][0])**2 + (pcoords[cents][1] - pcoords[kim][1])**2) # xl = dst*tang+six+nxc # ixl = int(xl) # dxl = xl - ixl # #print " A ", ifil,six,incline,kim,xl,ixl,dxl # qt += (1.0-dxl)*ctx[kim].get_value_at(ixl) + dxl*ctx[kim].get_value_at(ixl+1) # xl = -dst*tang+six+nxc # ixl = int(xl) # dxl = xl - ixl # qu += (1.0-dxl)*ctx[kim].get_value_at(ixl) + dxl*ctx[kim].get_value_at(ixl+1) # for kim in xrange(cents): # dst = sqrt((pcoords[cents][0] - pcoords[kim][0])**2 + (pcoords[cents][1] - pcoords[kim][1])**2) # xl = -dst*tang+six+nxc # ixl = int(xl) # dxl = xl - ixl # qt += (1.0-dxl)*ctx[kim].get_value_at(ixl) + dxl*ctx[kim].get_value_at(ixl+1) # xl = dst*tang+six+nxc # ixl = int(xl) # dxl = xl - ixl # qu += (1.0-dxl)*ctx[kim].get_value_at(ixl) + dxl*ctx[kim].get_value_at(ixl+1) # if( qt > qm ): # qm = qt # sib = six # bang = tang # if( qu > qm ): # qm = qu # sib = six # bang = -tang #if incline == 0: print "incline = 0 ",six,tang,qt,qu #print qm,six,sib,bang #print " got results ",indcs[ifil][0], indcs[ifil][1], ifil,myid,qm,sib,tang,bang,len(ctx),Util.infomask(ctx[0], None, True) for im in range(indcs[ifil][0], indcs[ifil][1]): kim = im - indcs[ifil][0] dst = sqrt((pcoords[cents][0] - pcoords[kim][0])**2 + (pcoords[cents][1] - pcoords[kim][1])**2) if (kim < cents): xl = -dst * bang + sib else: xl = dst * bang + sib shift_x[im] = xl # Average shift sx_sum += shift_x[indcs[ifil][0] + cents] # #print myid,sx_sum,total_nfils sx_sum = mpi.mpi_reduce(sx_sum, 1, mpi.MPI_FLOAT, mpi.MPI_SUM, main_node, mpi.MPI_COMM_WORLD) if myid == main_node: sx_sum = float(sx_sum[0]) / total_nfils print_msg("Average shift %6.2f\n" % (sx_sum)) else: sx_sum = 0.0 sx_sum = 0.0 sx_sum = bcast_number_to_all(sx_sum, source_node=main_node) for im in range(ldata): shift_x[im] -= sx_sum #print " %3d %6.3f"%(im,shift_x[im]) #exit() # combine shifts found with the original parameters for im in range(ldata): t1 = Transform() ##import random ##shix=random.randint(-10, 10) ##t1.set_params({"type":"2D","tx":shix}) t1.set_params({"type": "2D", "tx": shift_x[im]}) # combine t0 and t1 tt = t1 * init_params[im] data[im].set_attr("xform.align2d", tt) # write out headers and STOP, under MPI writing has to be done sequentially mpi.mpi_barrier(mpi.MPI_COMM_WORLD) par_str = ["xform.align2d", "ID"] if myid == main_node: from sp_utilities import file_type if (file_type(stack) == "bdb"): from sp_utilities import recv_attr_dict_bdb recv_attr_dict_bdb(main_node, stack, data, par_str, 0, ldata, nproc) else: from sp_utilities import recv_attr_dict recv_attr_dict(main_node, stack, data, par_str, 0, ldata, nproc) else: send_attr_dict(main_node, data, par_str, 0, ldata) if myid == main_node: print_end_msg("helical-shiftali_MPI")
def multalign2d_scf(image, refrings, frotim, numr, xrng=-1, yrng=-1, ou=-1): nx = image.get_xsize() ny = image.get_xsize() if ou < 0: ou = min(old_div(nx, 2) - 1, old_div(ny, 2) - 1) if yrng < 0: yrng = xrng if ou < 2: sp_global_def.ERROR("Radius of the object (ou) has to be given", "align2d_scf", 1) sci = sp_fundamentals.scf(image) first_ring = 1 # center in SPIDER convention cnx = old_div(nx, 2) + 1 cny = old_div(ny, 2) + 1 cimage = EMAN2_cppwrap.Util.Polar2Dm(sci, cnx, cny, numr, "H") EMAN2_cppwrap.Util.Frngs(cimage, numr) mimage = EMAN2_cppwrap.Util.Polar2Dm(sp_fundamentals.mirror(sci), cnx, cny, numr, "H") EMAN2_cppwrap.Util.Frngs(mimage, numr) nrx = min(2 * (xrng + 1) + 1, ((old_div((nx - 2), 2)) * 2 + 1)) nry = min(2 * (yrng + 1) + 1, ((old_div((ny - 2), 2)) * 2 + 1)) totpeak = -1.0e23 for iki in range(len(refrings)): # print "TEMPLATE ",iki # Find angle retvals = EMAN2_cppwrap.Util.Crosrng_e(refrings[iki], cimage, numr, 0, 0.0) alpha1 = ang_n(retvals["tot"], "H", numr[-1]) peak1 = retvals["qn"] retvals = EMAN2_cppwrap.Util.Crosrng_e(refrings[iki], mimage, numr, 0, 0.0) alpha2 = ang_n(retvals["tot"], "H", numr[-1]) peak2 = retvals["qn"] # print alpha1, peak1 # print alpha2, peak2 if peak1 > peak2: mirr = 0 alpha = alpha1 else: mirr = 1 alpha = -alpha2 ccf1 = EMAN2_cppwrap.Util.window( sp_fundamentals.ccf( sp_fundamentals.rot_shift2D(image, alpha, 0.0, 0.0, mirr), frotim[iki]), nrx, nry, ) p1 = sp_utilities.peak_search(ccf1) ccf2 = EMAN2_cppwrap.Util.window( sp_fundamentals.ccf( sp_fundamentals.rot_shift2D(image, alpha + 180.0, 0.0, 0.0, mirr), frotim[iki], ), nrx, nry, ) p2 = sp_utilities.peak_search(ccf2) # print p1 # print p2 peak_val1 = p1[0][0] peak_val2 = p2[0][0] if peak_val1 > peak_val2: sxs = -p1[0][4] sys = -p1[0][5] cx = int(p1[0][1]) cy = int(p1[0][2]) peak = peak_val1 else: alpha += 180.0 sxs = -p2[0][4] sys = -p2[0][5] peak = peak_val2 cx = int(p2[0][1]) cy = int(p2[0][2]) ccf1 = ccf2 # print cx,cy z = sp_utilities.model_blank(3, 3) for i in range(3): for j in range(3): z[i, j] = ccf1[i + cx - 1, j + cy - 1] # print ccf1[cx,cy],z[1,1] XSH, YSH, PEAKV = parabl(z) # print PEAKV if PEAKV > totpeak: totpeak = PEAKV iref = iki if mirr == 1: sx = -sxs + XSH else: sx = sxs - XSH sy = sys - YSH talpha = alpha tmirr = mirr # print "BETTER",sx,sy,iref,talpha,tmirr,totpeak # return alpha, sx, sys-YSH, mirr, PEAKV return sx, sy, iref, talpha, tmirr, totpeak
def main(): progname = os.path.basename(sys.argv[0]) usage = progname + " proj_stack output_averages --MPI" parser = OptionParser(usage, version=SPARXVERSION) parser.add_option("--img_per_group", type="int", default=100, help="number of images per group") parser.add_option("--radius", type="int", default=-1, help="radius for alignment") parser.add_option( "--xr", type="string", default="2 1", help="range for translation search in x direction, search is +/xr") parser.add_option( "--yr", type="string", default="-1", help= "range for translation search in y direction, search is +/yr (default = same as xr)" ) parser.add_option( "--ts", type="string", default="1 0.5", help= "step size of the translation search in both directions, search is -xr, -xr+ts, 0, xr-ts, xr, can be fractional" ) parser.add_option( "--iter", type="int", default=30, help="number of iterations within alignment (default = 30)") parser.add_option( "--num_ali", type="int", default=5, help="number of alignments performed for stability (default = 5)") parser.add_option("--thld_err", type="float", default=1.0, help="threshold of pixel error (default = 1.732)") parser.add_option( "--grouping", type="string", default="GRP", help= "do grouping of projections: PPR - per projection, GRP - different size groups, exclusive (default), GEV - grouping equal size" ) parser.add_option( "--delta", type="float", default=-1.0, help="angular step for reference projections (required for GEV method)" ) parser.add_option( "--fl", type="float", default=0.3, help="cut-off frequency of hyperbolic tangent low-pass Fourier filter") parser.add_option( "--aa", type="float", default=0.2, help="fall-off of hyperbolic tangent low-pass Fourier filter") parser.add_option("--CTF", action="store_true", default=False, help="Consider CTF correction during the alignment ") parser.add_option("--MPI", action="store_true", default=False, help="use MPI version") (options, args) = parser.parse_args() myid = mpi.mpi_comm_rank(MPI_COMM_WORLD) number_of_proc = mpi.mpi_comm_size(MPI_COMM_WORLD) main_node = 0 if len(args) == 2: stack = args[0] outdir = args[1] else: sp_global_def.ERROR("Incomplete list of arguments", "sxproj_stability.main", 1, myid=myid) return if not options.MPI: sp_global_def.ERROR("Non-MPI not supported!", "sxproj_stability.main", 1, myid=myid) return if sp_global_def.CACHE_DISABLE: from sp_utilities import disable_bdb_cache disable_bdb_cache() sp_global_def.BATCH = True img_per_grp = options.img_per_group radius = options.radius ite = options.iter num_ali = options.num_ali thld_err = options.thld_err xrng = get_input_from_string(options.xr) if options.yr == "-1": yrng = xrng else: yrng = get_input_from_string(options.yr) step = get_input_from_string(options.ts) if myid == main_node: nima = EMUtil.get_image_count(stack) img = get_image(stack) nx = img.get_xsize() ny = img.get_ysize() else: nima = 0 nx = 0 ny = 0 nima = bcast_number_to_all(nima) nx = bcast_number_to_all(nx) ny = bcast_number_to_all(ny) if radius == -1: radius = nx / 2 - 2 mask = model_circle(radius, nx, nx) st = time() if options.grouping == "GRP": if myid == main_node: sxprint(" A ", myid, " ", time() - st) proj_attr = EMUtil.get_all_attributes(stack, "xform.projection") proj_params = [] for i in range(nima): dp = proj_attr[i].get_params("spider") phi, theta, psi, s2x, s2y = dp["phi"], dp["theta"], dp[ "psi"], -dp["tx"], -dp["ty"] proj_params.append([phi, theta, psi, s2x, s2y]) # Here is where the grouping is done, I didn't put enough annotation in the group_proj_by_phitheta, # So I will briefly explain it here # proj_list : Returns a list of list of particle numbers, each list contains img_per_grp particle numbers # except for the last one. Depending on the number of particles left, they will either form a # group or append themselves to the last group # angle_list : Also returns a list of list, each list contains three numbers (phi, theta, delta), (phi, # theta) is the projection angle of the center of the group, delta is the range of this group # mirror_list: Also returns a list of list, each list contains img_per_grp True or False, which indicates # whether it should take mirror position. # In this program angle_list and mirror list are not of interest. proj_list_all, angle_list, mirror_list = group_proj_by_phitheta( proj_params, img_per_grp=img_per_grp) del proj_params sxprint(" B number of groups ", myid, " ", len(proj_list_all), time() - st) mpi_barrier(MPI_COMM_WORLD) # Number of groups, actually there could be one or two more groups, since the size of the remaining group varies # we will simply assign them to main node. n_grp = nima / img_per_grp - 1 # Divide proj_list_all equally to all nodes, and becomes proj_list proj_list = [] for i in range(n_grp): proc_to_stay = i % number_of_proc if proc_to_stay == main_node: if myid == main_node: proj_list.append(proj_list_all[i]) elif myid == main_node: mpi_send(len(proj_list_all[i]), 1, MPI_INT, proc_to_stay, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) mpi_send(proj_list_all[i], len(proj_list_all[i]), MPI_INT, proc_to_stay, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) elif myid == proc_to_stay: img_per_grp = mpi_recv(1, MPI_INT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) img_per_grp = int(img_per_grp[0]) temp = mpi_recv(img_per_grp, MPI_INT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) proj_list.append(list(map(int, temp))) del temp mpi_barrier(MPI_COMM_WORLD) sxprint(" C ", myid, " ", time() - st) if myid == main_node: # Assign the remaining groups to main_node for i in range(n_grp, len(proj_list_all)): proj_list.append(proj_list_all[i]) del proj_list_all, angle_list, mirror_list # Compute stability per projection projection direction, equal number assigned, thus overlaps elif options.grouping == "GEV": if options.delta == -1.0: ERROR( "Angular step for reference projections is required for GEV method" ) return from sp_utilities import even_angles, nearestk_to_refdir, getvec refproj = even_angles(options.delta) img_begin, img_end = MPI_start_end(len(refproj), number_of_proc, myid) # Now each processor keeps its own share of reference projections refprojdir = refproj[img_begin:img_end] del refproj ref_ang = [0.0] * (len(refprojdir) * 2) for i in range(len(refprojdir)): ref_ang[i * 2] = refprojdir[0][0] ref_ang[i * 2 + 1] = refprojdir[0][1] + i * 0.1 sxprint(" A ", myid, " ", time() - st) proj_attr = EMUtil.get_all_attributes(stack, "xform.projection") # the solution below is very slow, do not use it unless there is a problem with the i/O """ for i in xrange(number_of_proc): if myid == i: proj_attr = EMUtil.get_all_attributes(stack, "xform.projection") mpi_barrier(MPI_COMM_WORLD) """ sxprint(" B ", myid, " ", time() - st) proj_ang = [0.0] * (nima * 2) for i in range(nima): dp = proj_attr[i].get_params("spider") proj_ang[i * 2] = dp["phi"] proj_ang[i * 2 + 1] = dp["theta"] sxprint(" C ", myid, " ", time() - st) asi = Util.nearestk_to_refdir(proj_ang, ref_ang, img_per_grp) del proj_ang, ref_ang proj_list = [] for i in range(len(refprojdir)): proj_list.append(asi[i * img_per_grp:(i + 1) * img_per_grp]) del asi sxprint(" D ", myid, " ", time() - st) #from sys import exit #exit() # Compute stability per projection elif options.grouping == "PPR": sxprint(" A ", myid, " ", time() - st) proj_attr = EMUtil.get_all_attributes(stack, "xform.projection") sxprint(" B ", myid, " ", time() - st) proj_params = [] for i in range(nima): dp = proj_attr[i].get_params("spider") phi, theta, psi, s2x, s2y = dp["phi"], dp["theta"], dp[ "psi"], -dp["tx"], -dp["ty"] proj_params.append([phi, theta, psi, s2x, s2y]) img_begin, img_end = MPI_start_end(nima, number_of_proc, myid) sxprint(" C ", myid, " ", time() - st) from sp_utilities import nearest_proj proj_list, mirror_list = nearest_proj( proj_params, img_per_grp, list(range(img_begin, img_begin + 1))) #range(img_begin, img_end)) refprojdir = proj_params[img_begin:img_end] del proj_params, mirror_list sxprint(" D ", myid, " ", time() - st) else: ERROR("Incorrect projection grouping option") return ########################################################################################################### # Begin stability test from sp_utilities import get_params_proj, read_text_file #if myid == 0: # from utilities import read_text_file # proj_list[0] = map(int, read_text_file("lggrpp0.txt")) from sp_utilities import model_blank aveList = [model_blank(nx, ny)] * len(proj_list) if options.grouping == "GRP": refprojdir = [[0.0, 0.0, -1.0]] * len(proj_list) for i in range(len(proj_list)): sxprint(" E ", myid, " ", time() - st) class_data = EMData.read_images(stack, proj_list[i]) #print " R ",myid," ",time()-st if options.CTF: from sp_filter import filt_ctf for im in range(len(class_data)): # MEM LEAK!! atemp = class_data[im].copy() btemp = filt_ctf(atemp, atemp.get_attr("ctf"), binary=1) class_data[im] = btemp #class_data[im] = filt_ctf(class_data[im], class_data[im].get_attr("ctf"), binary=1) for im in class_data: try: t = im.get_attr( "xform.align2d") # if they are there, no need to set them! except: try: t = im.get_attr("xform.projection") d = t.get_params("spider") set_params2D(im, [0.0, -d["tx"], -d["ty"], 0, 1.0]) except: set_params2D(im, [0.0, 0.0, 0.0, 0, 1.0]) #print " F ",myid," ",time()-st # Here, we perform realignment num_ali times all_ali_params = [] for j in range(num_ali): if (xrng[0] == 0.0 and yrng[0] == 0.0): avet = ali2d_ras(class_data, randomize=True, ir=1, ou=radius, rs=1, step=1.0, dst=90.0, maxit=ite, check_mirror=True, FH=options.fl, FF=options.aa) else: avet = within_group_refinement(class_data, mask, True, 1, radius, 1, xrng, yrng, step, 90.0, ite, options.fl, options.aa) ali_params = [] for im in range(len(class_data)): alpha, sx, sy, mirror, scale = get_params2D(class_data[im]) ali_params.extend([alpha, sx, sy, mirror]) all_ali_params.append(ali_params) #aveList[i] = avet #print " G ",myid," ",time()-st del ali_params # We determine the stability of this group here. # stable_set contains all particles deemed stable, it is a list of list # each list has two elements, the first is the pixel error, the second is the image number # stable_set is sorted based on pixel error #from utilities import write_text_file #write_text_file(all_ali_params, "all_ali_params%03d.txt"%myid) stable_set, mir_stab_rate, average_pix_err = multi_align_stability( all_ali_params, 0.0, 10000.0, thld_err, False, 2 * radius + 1) #print " H ",myid," ",time()-st if (len(stable_set) > 5): stable_set_id = [] members = [] pix_err = [] # First put the stable members into attr 'members' and 'pix_err' for s in stable_set: # s[1] - number in this subset stable_set_id.append(s[1]) # the original image number members.append(proj_list[i][s[1]]) pix_err.append(s[0]) # Then put the unstable members into attr 'members' and 'pix_err' from sp_fundamentals import rot_shift2D avet.to_zero() if options.grouping == "GRP": aphi = 0.0 atht = 0.0 vphi = 0.0 vtht = 0.0 l = -1 for j in range(len(proj_list[i])): # Here it will only work if stable_set_id is sorted in the increasing number, see how l progresses if j in stable_set_id: l += 1 avet += rot_shift2D(class_data[j], stable_set[l][2][0], stable_set[l][2][1], stable_set[l][2][2], stable_set[l][2][3]) if options.grouping == "GRP": phi, theta, psi, sxs, sy_s = get_params_proj( class_data[j]) if (theta > 90.0): phi = (phi + 540.0) % 360.0 theta = 180.0 - theta aphi += phi atht += theta vphi += phi * phi vtht += theta * theta else: members.append(proj_list[i][j]) pix_err.append(99999.99) aveList[i] = avet.copy() if l > 1: l += 1 aveList[i] /= l if options.grouping == "GRP": aphi /= l atht /= l vphi = (vphi - l * aphi * aphi) / l vtht = (vtht - l * atht * atht) / l from math import sqrt refprojdir[i] = [ aphi, atht, (sqrt(max(vphi, 0.0)) + sqrt(max(vtht, 0.0))) / 2.0 ] # Here more information has to be stored, PARTICULARLY WHAT IS THE REFERENCE DIRECTION aveList[i].set_attr('members', members) aveList[i].set_attr('refprojdir', refprojdir[i]) aveList[i].set_attr('pixerr', pix_err) else: sxprint(" empty group ", i, refprojdir[i]) aveList[i].set_attr('members', [-1]) aveList[i].set_attr('refprojdir', refprojdir[i]) aveList[i].set_attr('pixerr', [99999.]) del class_data if myid == main_node: km = 0 for i in range(number_of_proc): if i == main_node: for im in range(len(aveList)): aveList[im].write_image(args[1], km) km += 1 else: nl = mpi_recv(1, MPI_INT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) nl = int(nl[0]) for im in range(nl): ave = recv_EMData(i, im + i + 70000) nm = mpi_recv(1, MPI_INT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) nm = int(nm[0]) members = mpi_recv(nm, MPI_INT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) ave.set_attr('members', list(map(int, members))) members = mpi_recv(nm, MPI_FLOAT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) ave.set_attr('pixerr', list(map(float, members))) members = mpi_recv(3, MPI_FLOAT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) ave.set_attr('refprojdir', list(map(float, members))) ave.write_image(args[1], km) km += 1 else: mpi_send(len(aveList), 1, MPI_INT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) for im in range(len(aveList)): send_EMData(aveList[im], main_node, im + myid + 70000) members = aveList[im].get_attr('members') mpi_send(len(members), 1, MPI_INT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) mpi_send(members, len(members), MPI_INT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) members = aveList[im].get_attr('pixerr') mpi_send(members, len(members), MPI_FLOAT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) try: members = aveList[im].get_attr('refprojdir') mpi_send(members, 3, MPI_FLOAT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) except: mpi_send([-999.0, -999.0, -999.0], 3, MPI_FLOAT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) sp_global_def.BATCH = False mpi_barrier(MPI_COMM_WORLD)
def align2d_direct3(input_images, refim, xrng=1, yrng=1, psimax=180, psistep=1, ou=-1, CTF=None): nx = input_images[0].get_xsize() if ou < 0: ou = old_div(nx, 2) - 1 mask = sp_utilities.model_circle(ou, nx, nx) nk = int(old_div(psimax, psistep)) nm = 2 * nk + 1 nc = nk + 1 refs = [None] * nm * 2 for i in range(nm): temp = sp_fundamentals.rot_shift2D(refim, (i - nc) * psistep) * mask refs[2 * i] = [ sp_fundamentals.fft(temp), sp_fundamentals.fft(sp_fundamentals.mirror(temp)), ] temp = sp_fundamentals.rot_shift2D(refim, (i - nc) * psistep + 180.0) * mask refs[2 * i + 1] = [ sp_fundamentals.fft(temp), sp_fundamentals.fft(sp_fundamentals.mirror(temp)), ] del temp results = [] mir = 0 for image in input_images: if CTF: ims = sp_filter.filt_ctf(sp_fundamentals.fft(image), image.get_attr("ctf")) else: ims = sp_fundamentals.fft(image) ama = -1.0e23 bang = 0.0 bsx = 0.0 bsy = 0.0 for i in range(nm * 2): for mirror_flag in [0, 1]: c = sp_fundamentals.ccf(ims, refs[i][mirror_flag]) w = EMAN2_cppwrap.Util.window(c, 2 * xrng + 1, 2 * yrng + 1) pp = sp_utilities.peak_search(w)[0] px = int(pp[4]) py = int(pp[5]) if pp[0] == 1.0 and px == 0 and py == 0: pass # XSH, YSH, PEAKV = 0.,0.,0. else: ww = sp_utilities.model_blank(3, 3) ux = int(pp[1]) uy = int(pp[2]) for k in range(3): for l in range(3): ww[k, l] = w[k + ux - 1, l + uy - 1] XSH, YSH, PEAKV = parabl(ww) # print i,pp[-1],XSH, YSH,px+XSH, py+YSH, PEAKV if PEAKV > ama: ama = PEAKV bsx = px + round(XSH, 2) bsy = py + round(YSH, 2) bang = i mir = mirror_flag # returned parameters have to be inverted bang = (old_div(bang, 2) - nc) * psistep + 180.0 * (bang % 2) bang, bsx, bsy, _ = sp_utilities.inverse_transform2( bang, (1 - 2 * mir) * bsx, bsy, mir) results.append([bang, bsx, bsy, mir, ama]) return results
def align2d_scf(image, refim, xrng=-1, yrng=-1, ou=-1): nx = image.get_xsize() ny = image.get_xsize() if ou < 0: ou = min(old_div(nx, 2) - 1, old_div(ny, 2) - 1) if yrng < 0: yrng = xrng if ou < 2: sp_global_def.ERROR("Radius of the object (ou) has to be given", "align2d_scf", 1) sci = sp_fundamentals.scf(image) scr = sp_fundamentals.scf(refim) first_ring = 1 # alpha1, sxs, sys, mirr, peak1 = align2d_no_mirror(scf(image), scr, last_ring=ou, mode="H") # alpha2, sxs, sys, mirr, peak2 = align2d_no_mirror(scf(mirror(image)), scr, last_ring=ou, mode="H") # alpha1, sxs, sys, mirr, peak1 = align2d_no_mirror(sci, scr, first_ring = 1, last_ring=ou, mode="H") # alpha2, sxs, sys, mirr, peak2 = align2d_no_mirror(mirror(sci), scr, first_ring = 1, last_ring=ou, mode="H") # center in SPIDER convention cnx = old_div(nx, 2) + 1 cny = old_div(ny, 2) + 1 # precalculate rings numr = Numrinit(first_ring, ou, 1, "H") wr = ringwe(numr, "H") crefim = EMAN2_cppwrap.Util.Polar2Dm(scr, cnx, cny, numr, "H") EMAN2_cppwrap.Util.Frngs(crefim, numr) EMAN2_cppwrap.Util.Applyws(crefim, numr, wr) alpha1, sxs, sys, mirr, peak1 = ornq(sci, crefim, [0.0], [0.0], 1.0, "H", numr, cnx, cny) alpha2, sxs, sys, mirr, peak2 = ornq(sp_fundamentals.mirror(sci), crefim, [0.0], [0.0], 1.0, "H", numr, cnx, cny) if peak1 > peak2: mirr = 0 alpha = alpha1 else: mirr = 1 alpha = -alpha2 nrx = min(2 * (xrng + 1) + 1, ((old_div((nx - 2), 2)) * 2 + 1)) nry = min(2 * (yrng + 1) + 1, ((old_div((ny - 2), 2)) * 2 + 1)) frotim = sp_fundamentals.fft(refim) ccf1 = EMAN2_cppwrap.Util.window( sp_fundamentals.ccf( sp_fundamentals.rot_shift2D(image, alpha, 0.0, 0.0, mirr), frotim), nrx, nry, ) p1 = sp_utilities.peak_search(ccf1) ccf2 = EMAN2_cppwrap.Util.window( sp_fundamentals.ccf( sp_fundamentals.rot_shift2D(image, alpha + 180.0, 0.0, 0.0, mirr), frotim), nrx, nry, ) p2 = sp_utilities.peak_search(ccf2) # print p1 # print p2 peak_val1 = p1[0][0] peak_val2 = p2[0][0] if peak_val1 > peak_val2: sxs = -p1[0][4] sys = -p1[0][5] cx = int(p1[0][1]) cy = int(p1[0][2]) peak = peak_val1 else: alpha += 180.0 sxs = -p2[0][4] sys = -p2[0][5] peak = peak_val2 cx = int(p2[0][1]) cy = int(p2[0][2]) ccf1 = ccf2 # print cx,cy z = sp_utilities.model_blank(3, 3) for i in range(3): for j in range(3): z[i, j] = ccf1[i + cx - 1, j + cy - 1] # print ccf1[cx,cy],z[1,1] XSH, YSH, PEAKV = parabl(z) # print sxs, sys, XSH, YSH, PEAKV, peak if mirr == 1: sx = -sxs + XSH else: sx = sxs - XSH return alpha, sx, sys - YSH, mirr, PEAKV
def filter_shrink(input_bdb_name, output_stack_path, output_bdb_name, alignYN=False, filtrad=None, shrink=None, verbose=False): """ Filters and shrinks image stack. Arguments: input_bdb_name : Input BDB stack (in the form bdb:DIRECTORY#STACK) output_stack_path : Name for output image stack (MRCS/HDF/etc.) output_bdb_name : Name for output BDB stack (in the form bdb:DIRECTORY#STACK) alignYN: (boolean) Whether to apply alignments filtrad : Filter radius, reciprocal pixels shrink : Downsampling factor combined_params_file : Output combined alignment parameters verbose : (boolean) Whether to write to screen """ input_stack = EMData.read_images(input_bdb_name) num_imgs = EMUtil.get_image_count(input_bdb_name) box_size = input_stack[0].get_attr('nx') if options.shrink: sub_rate = float(1) / options.shrink box_size = int(float(box_size) / options.shrink + 0.5) if verbose: print('sub_rate', sub_rate, type(sub_rate), box_size, options.shrink) # Initialize stack & BDB aligned_stack_obj = EMData(box_size, box_size, num_imgs) new_bdb_dict = db_open_dict(output_bdb_name) # Loop through images for img_num in range(len(input_stack)): img_orig = input_stack[img_num] try: alpha, sx, sy, mirror, scale = get_params2D(img_orig) except RuntimeError: print('\nERROR! Exiting with RuntimeError') img_prev = input_stack[img_num - 1] print('\nPrevious particle: %s %s' % (img_num - 1, img_prev.get_attr_dict())) print('\nCurrent particle: %s %s' % (img_num, img_orig.get_attr_dict())) exit() # Optionally apply alignment parameters if alignYN: img_ali = rot_shift2D(img_orig, alpha, sx, sy, mirror, scale, "quadratic") else: img_ali = img_orig if verbose and img_num == 0: print('\nimg_orig.get_attr_dict0', img_orig.get_attr_dict()) img_ali_dict = img_ali.get_attr_dict() print('\nimg_ali.get_attr_dict1\n', img_ali.get_attr_dict()) if filtrad: img_ali = filt_gaussl(img_ali, filtrad) if shrink: img_ali = resample(img_ali, sub_rate) #### (Maybe update resample_ratio) img_ali_dict = img_ali.get_attr_dict() img_ali_dict["data_path"] = os.path.join( '..', STACKFILEDIR, os.path.basename(output_stack_path)) img_ali_dict["ptcl_source_coord_id"] = img_num new_bdb_dict[img_num] = img_ali_dict aligned_stack_obj.insert_clip(img_ali, (0, 0, img_num)) # End image-loop aligned_stack_obj.write_image(output_stack_path) db_close_dict(output_bdb_name) return num_imgs
def mref_ali2d_MPI(stack, refim, outdir, maskfile=None, ir=1, ou=-1, rs=1, xrng=0, yrng=0, step=1, center=1, maxit=10, CTF=False, snr=1.0, user_func_name="ref_ali2d", rand_seed=1000): # 2D multi-reference alignment using rotational ccf in polar coordinates and quadratic interpolation from sp_utilities import model_circle, combine_params2, inverse_transform2, drop_image, get_image, get_im from sp_utilities import reduce_EMData_to_root, bcast_EMData_to_all, bcast_number_to_all from sp_utilities import send_attr_dict from sp_utilities import center_2D from sp_statistics import fsc_mask from sp_alignment import Numrinit, ringwe, search_range from sp_fundamentals import rot_shift2D, fshift from sp_utilities import get_params2D, set_params2D from random import seed, randint from sp_morphology import ctf_2 from sp_filter import filt_btwl, filt_params from numpy import reshape, shape from sp_utilities import print_msg, print_begin_msg, print_end_msg import os import sys import shutil from sp_applications import MPI_start_end from mpi import mpi_comm_size, mpi_comm_rank, MPI_COMM_WORLD from mpi import mpi_reduce, mpi_bcast, mpi_barrier, mpi_recv, mpi_send from mpi import MPI_SUM, MPI_FLOAT, MPI_INT number_of_proc = mpi_comm_size(MPI_COMM_WORLD) myid = mpi_comm_rank(MPI_COMM_WORLD) main_node = 0 # create the output directory, if it does not exist if os.path.exists(outdir): ERROR( 'Output directory exists, please change the name and restart the program', "mref_ali2d_MPI ", 1, myid) mpi_barrier(MPI_COMM_WORLD) import sp_global_def if myid == main_node: os.mkdir(outdir) sp_global_def.LOGFILE = os.path.join(outdir, sp_global_def.LOGFILE) print_begin_msg("mref_ali2d_MPI") nima = EMUtil.get_image_count(stack) image_start, image_end = MPI_start_end(nima, number_of_proc, myid) nima = EMUtil.get_image_count(stack) ima = EMData() ima.read_image(stack, image_start) first_ring = int(ir) last_ring = int(ou) rstep = int(rs) max_iter = int(maxit) if max_iter == 0: max_iter = 10 auto_stop = True else: auto_stop = False if myid == main_node: print_msg("Input stack : %s\n" % (stack)) print_msg("Reference stack : %s\n" % (refim)) print_msg("Output directory : %s\n" % (outdir)) print_msg("Maskfile : %s\n" % (maskfile)) print_msg("Inner radius : %i\n" % (first_ring)) nx = ima.get_xsize() # default value for the last ring if last_ring == -1: last_ring = nx / 2 - 2 if myid == main_node: print_msg("Outer radius : %i\n" % (last_ring)) print_msg("Ring step : %i\n" % (rstep)) print_msg("X search range : %f\n" % (xrng)) print_msg("Y search range : %f\n" % (yrng)) print_msg("Translational step : %f\n" % (step)) print_msg("Center type : %i\n" % (center)) print_msg("Maximum iteration : %i\n" % (max_iter)) print_msg("CTF correction : %s\n" % (CTF)) print_msg("Signal-to-Noise Ratio : %f\n" % (snr)) print_msg("Random seed : %i\n\n" % (rand_seed)) print_msg("User function : %s\n" % (user_func_name)) import sp_user_functions user_func = sp_user_functions.factory[user_func_name] if maskfile: import types if type(maskfile) is bytes: mask = get_image(maskfile) else: mask = maskfile else: mask = model_circle(last_ring, nx, nx) # references, do them on all processors... refi = [] numref = EMUtil.get_image_count(refim) # IMAGES ARE SQUARES! center is in SPIDER convention cnx = nx / 2 + 1 cny = cnx mode = "F" #precalculate rings numr = Numrinit(first_ring, last_ring, rstep, mode) wr = ringwe(numr, mode) # prepare reference images on all nodes ima.to_zero() for j in range(numref): # even, odd, numer of even, number of images. After frc, totav refi.append([get_im(refim, j), ima.copy(), 0]) # for each node read its share of data data = EMData.read_images(stack, list(range(image_start, image_end))) for im in range(image_start, image_end): data[im - image_start].set_attr('ID', im) if myid == main_node: seed(rand_seed) a0 = -1.0 again = True Iter = 0 ref_data = [mask, center, None, None] while Iter < max_iter and again: ringref = [] mashi = cnx - last_ring - 2 for j in range(numref): refi[j][0].process_inplace("normalize.mask", { "mask": mask, "no_sigma": 1 }) # normalize reference images to N(0,1) cimage = Util.Polar2Dm(refi[j][0], cnx, cny, numr, mode) Util.Frngs(cimage, numr) Util.Applyws(cimage, numr, wr) ringref.append(cimage) # zero refi refi[j][0].to_zero() refi[j][1].to_zero() refi[j][2] = 0 assign = [[] for i in range(numref)] # begin MPI section for im in range(image_start, image_end): alpha, sx, sy, mirror, scale = get_params2D(data[im - image_start]) # Why inverse? 07/11/2015 PAP alphai, sxi, syi, scalei = inverse_transform2(alpha, sx, sy) # normalize data[im - image_start].process_inplace("normalize.mask", { "mask": mask, "no_sigma": 0 }) # subtract average under the mask # If shifts are outside of the permissible range, reset them if (abs(sxi) > mashi or abs(syi) > mashi): sxi = 0.0 syi = 0.0 set_params2D(data[im - image_start], [0.0, 0.0, 0.0, 0, 1.0]) ny = nx txrng = search_range(nx, last_ring, sxi, xrng, "mref_ali2d_MPI") txrng = [txrng[1], txrng[0]] tyrng = search_range(ny, last_ring, syi, yrng, "mref_ali2d_MPI") tyrng = [tyrng[1], tyrng[0]] # align current image to the reference [angt, sxst, syst, mirrort, xiref, peakt] = Util.multiref_polar_ali_2d(data[im - image_start], ringref, txrng, tyrng, step, mode, numr, cnx + sxi, cny + syi) iref = int(xiref) # combine parameters and set them to the header, ignore previous angle and mirror [alphan, sxn, syn, mn] = combine_params2(0.0, -sxi, -syi, 0, angt, sxst, syst, (int)(mirrort)) set_params2D(data[im - image_start], [alphan, sxn, syn, int(mn), scale]) data[im - image_start].set_attr('assign', iref) # apply current parameters and add to the average temp = rot_shift2D(data[im - image_start], alphan, sxn, syn, mn) it = im % 2 Util.add_img(refi[iref][it], temp) assign[iref].append(im) #assign[im] = iref refi[iref][2] += 1.0 del ringref # end MPI section, bring partial things together, calculate new reference images, broadcast them back for j in range(numref): reduce_EMData_to_root(refi[j][0], myid, main_node) reduce_EMData_to_root(refi[j][1], myid, main_node) refi[j][2] = mpi_reduce(refi[j][2], 1, MPI_FLOAT, MPI_SUM, main_node, MPI_COMM_WORLD) if (myid == main_node): refi[j][2] = int(refi[j][2][0]) # gather assignements for j in range(numref): if myid == main_node: for n in range(number_of_proc): if n != main_node: import sp_global_def ln = mpi_recv(1, MPI_INT, n, sp_global_def.SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) lis = mpi_recv(ln[0], MPI_INT, n, sp_global_def.SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) for l in range(ln[0]): assign[j].append(int(lis[l])) else: import sp_global_def mpi_send(len(assign[j]), 1, MPI_INT, main_node, sp_global_def.SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) mpi_send(assign[j], len(assign[j]), MPI_INT, main_node, sp_global_def.SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) if myid == main_node: # replace the name of the stack with reference with the current one refim = os.path.join(outdir, "aqm%03d.hdf" % Iter) a1 = 0.0 ave_fsc = [] for j in range(numref): if refi[j][2] < 4: #ERROR("One of the references vanished","mref_ali2d_MPI",1) # if vanished, put a random image (only from main node!) there assign[j] = [] assign[j].append( randint(image_start, image_end - 1) - image_start) refi[j][0] = data[assign[j][0]].copy() #print 'ERROR', j else: #frsc = fsc_mask(refi[j][0], refi[j][1], mask, 1.0, os.path.join(outdir,"drm%03d%04d"%(Iter, j))) from sp_statistics import fsc frsc = fsc( refi[j][0], refi[j][1], 1.0, os.path.join(outdir, "drm%03d%04d.txt" % (Iter, j))) Util.add_img(refi[j][0], refi[j][1]) Util.mul_scalar(refi[j][0], 1.0 / float(refi[j][2])) if ave_fsc == []: for i in range(len(frsc[1])): ave_fsc.append(frsc[1][i]) c_fsc = 1 else: for i in range(len(frsc[1])): ave_fsc[i] += frsc[1][i] c_fsc += 1 #print 'OK', j, len(frsc[1]), frsc[1][0:5], ave_fsc[0:5] #print 'sum', sum(ave_fsc) if sum(ave_fsc) != 0: for i in range(len(ave_fsc)): ave_fsc[i] /= float(c_fsc) frsc[1][i] = ave_fsc[i] for j in range(numref): ref_data[2] = refi[j][0] ref_data[3] = frsc refi[j][0], cs = user_func(ref_data) # write the current average TMP = [] for i_tmp in range(len(assign[j])): TMP.append(float(assign[j][i_tmp])) TMP.sort() refi[j][0].set_attr_dict({'ave_n': refi[j][2], 'members': TMP}) del TMP refi[j][0].process_inplace("normalize.mask", { "mask": mask, "no_sigma": 1 }) refi[j][0].write_image(refim, j) Iter += 1 msg = "ITERATION #%3d %d\n\n" % (Iter, again) print_msg(msg) for j in range(numref): msg = " group #%3d number of particles = %7d\n" % ( j, refi[j][2]) print_msg(msg) Iter = bcast_number_to_all(Iter, main_node) # need to tell all if again: for j in range(numref): bcast_EMData_to_all(refi[j][0], myid, main_node) # clean up del assign # write out headers and STOP, under MPI writing has to be done sequentially (time-consumming) mpi_barrier(MPI_COMM_WORLD) if CTF and data_had_ctf == 0: for im in range(len(data)): data[im].set_attr('ctf_applied', 0) par_str = ['xform.align2d', 'assign', 'ID'] if myid == main_node: from sp_utilities import file_type if (file_type(stack) == "bdb"): from sp_utilities import recv_attr_dict_bdb recv_attr_dict_bdb(main_node, stack, data, par_str, image_start, image_end, number_of_proc) else: from sp_utilities import recv_attr_dict recv_attr_dict(main_node, stack, data, par_str, image_start, image_end, number_of_proc) else: send_attr_dict(main_node, data, par_str, image_start, image_end) if myid == main_node: print_end_msg("mref_ali2d_MPI")
def mref_ali2d(stack, refim, outdir, maskfile=None, ir=1, ou=-1, rs=1, xrng=0, yrng=0, step=1, center=1, maxit=0, CTF=False, snr=1.0, user_func_name="ref_ali2d", rand_seed=1000, MPI=False): """ Name mref_ali2d - Perform 2-D multi-reference alignment of an image series Input stack: set of 2-D images in a stack file, images have to be squares refim: set of initial reference 2-D images in a stack file maskfile: optional maskfile to be used in the alignment inner_radius: inner radius for rotational correlation > 0 outer_radius: outer radius for rotational correlation < nx/2-1 ring_step: step between rings in rotational correlation >0 x_range: range for translation search in x direction, search is +/xr y_range: range for translation search in y direction, search is +/yr translation_step: step of translation search in both directions center: center the average max_iter: maximum number of iterations the program will perform CTF: if this flag is set, the program will use CTF information provided in file headers snr: signal-to-noise ratio of the data rand_seed: the seed used for generating random numbers MPI: whether to use MPI version Output output_directory: directory name into which the output files will be written. header: the alignment parameters are stored in the headers of input files as 'xform.align2d'. """ # 2D multi-reference alignment using rotational ccf in polar coordinates and quadratic interpolation if MPI: mref_ali2d_MPI(stack, refim, outdir, maskfile, ir, ou, rs, xrng, yrng, step, center, maxit, CTF, snr, user_func_name, rand_seed) return from sp_utilities import model_circle, combine_params2, inverse_transform2, drop_image, get_image from sp_utilities import center_2D, get_im, get_params2D, set_params2D from sp_statistics import fsc from sp_alignment import Numrinit, ringwe, fine_2D_refinement, search_range from sp_fundamentals import rot_shift2D, fshift from random import seed, randint import os import sys from sp_utilities import print_begin_msg, print_end_msg, print_msg import shutil # create the output directory, if it does not exist if os.path.exists(outdir): shutil.rmtree( outdir ) #ERROR('Output directory exists, please change the name and restart the program', "mref_ali2d", 1) os.mkdir(outdir) import sp_global_def sp_global_def.LOGFILE = os.path.join(outdir, sp_global_def.LOGFILE) first_ring = int(ir) last_ring = int(ou) rstep = int(rs) max_iter = int(maxit) print_begin_msg("mref_ali2d") print_msg("Input stack : %s\n" % (stack)) print_msg("Reference stack : %s\n" % (refim)) print_msg("Output directory : %s\n" % (outdir)) print_msg("Maskfile : %s\n" % (maskfile)) print_msg("Inner radius : %i\n" % (first_ring)) ima = EMData() ima.read_image(stack, 0) nx = ima.get_xsize() # default value for the last ring if last_ring == -1: last_ring = nx / 2 - 2 print_msg("Outer radius : %i\n" % (last_ring)) print_msg("Ring step : %i\n" % (rstep)) print_msg("X search range : %i\n" % (xrng)) print_msg("Y search range : %i\n" % (yrng)) print_msg("Translational step : %i\n" % (step)) print_msg("Center type : %i\n" % (center)) print_msg("Maximum iteration : %i\n" % (max_iter)) print_msg("CTF correction : %s\n" % (CTF)) print_msg("Signal-to-Noise Ratio : %f\n" % (snr)) print_msg("Random seed : %i\n\n" % (rand_seed)) print_msg("User function : %s\n" % (user_func_name)) output = sys.stdout import sp_user_functions user_func = sp_user_functions.factory[user_func_name] if maskfile: import types if type(maskfile) is bytes: mask = get_image(maskfile) else: mask = maskfile else: mask = model_circle(last_ring, nx, nx) # references refi = [] numref = EMUtil.get_image_count(refim) # IMAGES ARE SQUARES! center is in SPIDER convention cnx = nx / 2 + 1 cny = cnx mode = "F" #precalculate rings numr = Numrinit(first_ring, last_ring, rstep, mode) wr = ringwe(numr, mode) # reference images params = [] #read all data data = EMData.read_images(stack) nima = len(data) # prepare the reference ima.to_zero() for j in range(numref): temp = EMData() temp.read_image(refim, j) # eve, odd, numer of even, number of images. After frc, totav refi.append([temp, ima.copy(), 0]) seed(rand_seed) again = True ref_data = [mask, center, None, None] Iter = 0 while Iter < max_iter and again: ringref = [] #print "numref",numref ### Reference ### mashi = cnx - last_ring - 2 for j in range(numref): refi[j][0].process_inplace("normalize.mask", { "mask": mask, "no_sigma": 1 }) cimage = Util.Polar2Dm(refi[j][0], cnx, cny, numr, mode) Util.Frngs(cimage, numr) Util.Applyws(cimage, numr, wr) ringref.append(cimage) assign = [[] for i in range(numref)] sx_sum = [0.0] * numref sy_sum = [0.0] * numref for im in range(nima): alpha, sx, sy, mirror, scale = get_params2D(data[im]) # Why inverse? 07/11/2015 PAP alphai, sxi, syi, scalei = inverse_transform2(alpha, sx, sy) # normalize data[im].process_inplace("normalize.mask", { "mask": mask, "no_sigma": 0 }) # If shifts are outside of the permissible range, reset them if (abs(sxi) > mashi or abs(syi) > mashi): sxi = 0.0 syi = 0.0 set_params2D(data[im], [0.0, 0.0, 0.0, 0, 1.0]) ny = nx txrng = search_range(nx, last_ring, sxi, xrng, "mref_ali2d") txrng = [txrng[1], txrng[0]] tyrng = search_range(ny, last_ring, syi, yrng, "mref_ali2d") tyrng = [tyrng[1], tyrng[0]] # align current image to the reference #[angt, sxst, syst, mirrort, xiref, peakt] = Util.multiref_polar_ali_2d_p(data[im], # ringref, txrng, tyrng, step, mode, numr, cnx+sxi, cny+syi) #print(angt, sxst, syst, mirrort, xiref, peakt) [angt, sxst, syst, mirrort, xiref, peakt] = Util.multiref_polar_ali_2d(data[im], ringref, txrng, tyrng, step, mode, numr, cnx + sxi, cny + syi) iref = int(xiref) # combine parameters and set them to the header, ignore previous angle and mirror [alphan, sxn, syn, mn] = combine_params2(0.0, -sxi, -syi, 0, angt, sxst, syst, int(mirrort)) set_params2D(data[im], [alphan, sxn, syn, int(mn), scale]) if mn == 0: sx_sum[iref] += sxn else: sx_sum[iref] -= sxn sy_sum[iref] += syn data[im].set_attr('assign', iref) # apply current parameters and add to the average temp = rot_shift2D(data[im], alphan, sxn, syn, mn) it = im % 2 Util.add_img(refi[iref][it], temp) assign[iref].append(im) refi[iref][2] += 1 del ringref if again: a1 = 0.0 for j in range(numref): msg = " group #%3d number of particles = %7d\n" % ( j, refi[j][2]) print_msg(msg) if refi[j][2] < 4: #ERROR("One of the references vanished","mref_ali2d",1) # if vanished, put a random image there assign[j] = [] assign[j].append(randint(0, nima - 1)) refi[j][0] = data[assign[j][0]].copy() else: max_inter = 0 # switch off fine refi. br = 1.75 # the loop has to for INter in range(max_inter + 1): # Calculate averages at least ones, meaning even if no within group refinement was requested frsc = fsc( refi[j][0], refi[j][1], 1.0, os.path.join(outdir, "drm_%03d_%04d.txt" % (Iter, j))) Util.add_img(refi[j][0], refi[j][1]) Util.mul_scalar(refi[j][0], 1.0 / float(refi[j][2])) ref_data[2] = refi[j][0] ref_data[3] = frsc refi[j][0], cs = user_func(ref_data) if center == -1: cs[0] = sx_sum[j] / len(assign[j]) cs[1] = sy_sum[j] / len(assign[j]) refi[j][0] = fshift(refi[j][0], -cs[0], -cs[1]) for i in range(len(assign[j])): im = assign[j][i] alpha, sx, sy, mirror, scale = get_params2D( data[im]) alphan, sxn, syn, mirrorn = combine_params2( alpha, sx, sy, mirror, 0.0, -cs[0], -cs[1], 0) set_params2D( data[im], [alphan, sxn, syn, int(mirrorn), scale]) # refine images within the group # Do the refinement only if max_inter>0, but skip it for the last iteration. if INter < max_inter: fine_2D_refinement(data, br, mask, refi[j][0], j) # Calculate updated average refi[j][0].to_zero() refi[j][1].to_zero() for i in range(len(assign[j])): im = assign[j][i] alpha, sx, sy, mirror, scale = get_params2D( data[im]) # apply current parameters and add to the average temp = rot_shift2D(data[im], alpha, sx, sy, mn) it = im % 2 Util.add_img(refi[j][it], temp) # write the current average TMP = [] for i_tmp in range(len(assign[j])): TMP.append(float(assign[j][i_tmp])) TMP.sort() refi[j][0].set_attr_dict({'ave_n': refi[j][2], 'members': TMP}) del TMP # replace the name of the stack with reference with the current one newrefim = os.path.join(outdir, "aqm%03d.hdf" % Iter) refi[j][0].write_image(newrefim, j) Iter += 1 msg = "ITERATION #%3d \n" % (Iter) print_msg(msg) newrefim = os.path.join(outdir, "multi_ref.hdf") for j in range(numref): refi[j][0].write_image(newrefim, j) from sp_utilities import write_headers write_headers(stack, data, list(range(nima))) print_end_msg("mref_ali2d")
def main(): progname = os.path.basename(sys.argv[0]) usage = progname + """ [options] <inputfile> <outputfile> Forms chains of 2D images based on their similarities. Functionality: Order a 2-D stack of image based on pair-wise similarity (computed as a cross-correlation coefficent). Options 1-3 require image stack to be aligned. The program will apply orientation parameters if present in headers. The ways to use the program: 1. Use option initial to specify which image will be used as an initial seed to form the chain. sp_chains.py input_stack.hdf output_stack.hdf --initial=23 --radius=25 2. If options initial is omitted, the program will determine which image best serves as initial seed to form the chain sp_chains.py input_stack.hdf output_stack.hdf --radius=25 3. Use option circular to form a circular chain. sp_chains.py input_stack.hdf output_stack.hdf --circular--radius=25 4. New circular code based on pairwise alignments sp_chains.py aclf.hdf chain.hdf circle.hdf --align --radius=25 --xr=2 --pairwiseccc=lcc.txt 5. Circular ordering based on pairwise alignments sp_chains.py vols.hdf chain.hdf mask.hdf --dd --radius=25 """ parser = OptionParser(usage, version=SPARXVERSION) parser.add_option( "--dd", action="store_true", help="Circular ordering without adjustment of orientations", default=False) parser.add_option( "--circular", action="store_true", help= "Select circular ordering (first image has to be similar to the last)", default=False) parser.add_option( "--align", action="store_true", help= "Compute all pairwise alignments and from the table of image similarities find the best chain", default=False) parser.add_option( "--initial", type="int", default=-1, help= "Specifies which image will be used as an initial seed to form the chain. (default = 0, means the first image)" ) parser.add_option( "--radius", type="int", default=-1, help="Radius of a circular mask for similarity based ordering") # import params for 2D alignment parser.add_option( "--ou", type="int", default=-1, help= "outer radius for 2D alignment < nx/2-1 (set to the radius of the particle)" ) parser.add_option( "--xr", type="int", default=0, help="range for translation search in x direction, search is +/xr (0)") parser.add_option( "--yr", type="int", default=0, help="range for translation search in y direction, search is +/yr (0)") # parser.add_option("--nomirror", action="store_true", default=False, help="Disable checking mirror orientations of images (default False)") parser.add_option("--pairwiseccc", type="string", default=" ", help="Input/output pairwise ccc file") (options, args) = parser.parse_args() sp_global_def.BATCH = True if options.dd: nargs = len(args) if nargs != 3: ERROR("Must provide name of input and two output files!") return stack = args[0] new_stack = args[1] from sp_utilities import model_circle from sp_statistics import ccc from sp_statistics import mono lend = EMUtil.get_image_count(stack) lccc = [None] * (old_div(lend * (lend - 1), 2)) for i in range(lend - 1): v1 = get_im(stack, i) if (i == 0 and nargs == 2): nx = v1.get_xsize() ny = v1.get_ysize() nz = v1.get_ysize() if options.ou < 1: radius = old_div(nx, 2) - 2 else: radius = options.ou mask = model_circle(radius, nx, ny, nz) else: mask = get_im(args[2]) for j in range(i + 1, lend): lccc[mono(i, j)] = [ccc(v1, get_im(stack, j), mask), 0] order = tsp(lccc) if (len(order) != lend): ERROR("Problem with data length") return sxprint("Total sum of cccs :", TotalDistance(order, lccc)) sxprint("ordering :", order) for i in range(lend): get_im(stack, order[i]).write_image(new_stack, i) elif options.align: nargs = len(args) if nargs != 3: ERROR("Must provide name of input and two output files!") return from sp_utilities import get_params2D, model_circle from sp_fundamentals import rot_shift2D from sp_statistics import ccc from time import time from sp_alignment import align2d, align2d_scf stack = args[0] new_stack = args[1] d = EMData.read_images(stack) if (len(d) < 6): ERROR( "Chains requires at least six images in the input stack to be executed" ) return """ # will align anyway try: ttt = d[0].get_attr('xform.params2d') for i in xrange(len(d)): alpha, sx, sy, mirror, scale = get_params2D(d[i]) d[i] = rot_shift2D(d[i], alpha, sx, sy, mirror) except: pass """ nx = d[0].get_xsize() ny = d[0].get_ysize() if options.ou < 1: radius = old_div(nx, 2) - 2 else: radius = options.ou mask = model_circle(radius, nx, ny) if (options.xr < 0): xrng = 0 else: xrng = options.xr if (options.yr < 0): yrng = xrng else: yrng = options.yr initial = max(options.initial, 0) from sp_statistics import mono lend = len(d) lccc = [None] * (old_div(lend * (lend - 1), 2)) from sp_utilities import read_text_row if options.pairwiseccc == " " or not os.path.exists( options.pairwiseccc): st = time() for i in range(lend - 1): for j in range(i + 1, lend): # j>i meaning mono entry (i,j) or (j,i) indicates T i->j (from smaller index to larger) # alpha, sx, sy, mir, peak = align2d(d[i],d[j], xrng, yrng, step=options.ts, first_ring=options.ir, last_ring=radius, mode = "F") alpha, sx, sy, mir, peak = align2d_scf(d[i], d[j], xrng, yrng, ou=radius) lccc[mono(i, j)] = [ ccc(d[j], rot_shift2D(d[i], alpha, sx, sy, mir, 1.0), mask), alpha, sx, sy, mir ] # print " %4d %10.1f"%(i,time()-st) if ((not os.path.exists(options.pairwiseccc)) and (options.pairwiseccc != " ")): write_text_row([[initial, 0, 0, 0, 0]] + lccc, options.pairwiseccc) elif (os.path.exists(options.pairwiseccc)): lccc = read_text_row(options.pairwiseccc) initial = int(lccc[0][0] + 0.1) del lccc[0] for i in range(len(lccc)): T = Transform({ "type": "2D", "alpha": lccc[i][1], "tx": lccc[i][2], "ty": lccc[i][3], "mirror": int(lccc[i][4] + 0.1) }) lccc[i] = [lccc[i][0], T] tdummy = Transform({"type": "2D"}) maxsum = -1.023 for m in range(0, lend): # initial, initial+1): indc = list(range(lend)) lsnake = [[m, tdummy, 0.0]] del indc[m] lsum = 0.0 while len(indc) > 1: maxcit = -111. for i in range(len(indc)): cuc = lccc[mono(indc[i], lsnake[-1][0])][0] if cuc > maxcit: maxcit = cuc qi = indc[i] # Here we need transformation from the current to the previous, # meaning indc[i] -> lsnake[-1][0] T = lccc[mono(indc[i], lsnake[-1][0])][1] # If direction is from larger to smaller index, the transformation has to be inverted if (indc[i] > lsnake[-1][0]): T = T.inverse() lsnake.append([qi, T, maxcit]) lsum += maxcit del indc[indc.index(qi)] T = lccc[mono(indc[-1], lsnake[-1][0])][1] if (indc[-1] > lsnake[-1][0]): T = T.inverse() lsnake.append( [indc[-1], T, lccc[mono(indc[-1], lsnake[-1][0])][0]]) sxprint(" initial image and lsum ", m, lsum) # print lsnake if (lsum > maxsum): maxsum = lsum init = m snake = [lsnake[i] for i in range(lend)] sxprint(" Initial image selected : ", init, maxsum, " ", TotalDistance([snake[m][0] for m in range(lend)], lccc)) # for q in snake: print q from copy import deepcopy trans = deepcopy([snake[i][1] for i in range(len(snake))]) sxprint([snake[i][0] for i in range(len(snake))]) """ for m in xrange(lend): prms = trans[m].get_params("2D") print " %3d %7.1f %7.1f %7.1f %2d %6.2f"%(snake[m][0], prms["alpha"], prms["tx"], prms["ty"], prms["mirror"], snake[m][2]) """ for k in range(lend - 2, 0, -1): T = snake[k][1] for i in range(k + 1, lend): trans[i] = T * trans[i] # To add - apply all transformations and do the overall centering. for m in range(lend): prms = trans[m].get_params("2D") # print(" %3d %7.1f %7.1f %7.1f %2d %6.2f"%(snake[m][0], prms["alpha"], prms["tx"], prms["ty"], prms["mirror"], snake[m][2]) ) # rot_shift2D(d[snake[m][0]], prms["alpha"], prms["tx"], prms["ty"], prms["mirror"]).write_image(new_stack, m) rot_shift2D(d[snake[m][0]], prms["alpha"], 0.0, 0.0, prms["mirror"]).write_image(new_stack, m) order = tsp(lccc) if (len(order) != lend): ERROR("Problem with data length") return sxprint(TotalDistance(order, lccc)) sxprint(order) ibeg = order.index(init) order = [order[(i + ibeg) % lend] for i in range(lend)] sxprint(TotalDistance(order, lccc)) sxprint(order) snake = [tdummy] for i in range(1, lend): # Here we need transformation from the current to the previous, # meaning order[i] -> order[i-1]] T = lccc[mono(order[i], order[i - 1])][1] # If direction is from larger to smaller index, the transformation has to be inverted if (order[i] > order[i - 1]): T = T.inverse() snake.append(T) assert (len(snake) == lend) from copy import deepcopy trans = deepcopy(snake) for k in range(lend - 2, 0, -1): T = snake[k] for i in range(k + 1, lend): trans[i] = T * trans[i] # Try to smooth the angles - complicated, I am afraid one would have to use angles forward and backwards # and find their average?? # In addition, one would have to recenter them """ trms = [] for m in xrange(lend): prms = trans[m].get_params("2D") trms.append([prms["alpha"], prms["mirror"]]) for i in xrange(3): for m in xrange(lend): mb = (m-1)%lend me = (m+1)%lend # angles order mb,m,me # calculate predicted angles mb->m """ best_params = [] for m in range(lend): prms = trans[m].get_params("2D") # rot_shift2D(d[order[m]], prms["alpha"], prms["tx"], prms["ty"], prms["mirror"]).write_image("metro.hdf", m) rot_shift2D(d[order[m]], prms["alpha"], 0.0, 0.0, prms["mirror"]).write_image(args[2], m) best_params.append( [m, order[m], prms["alpha"], 0.0, 0.0, prms["mirror"]]) # Write alignment parameters outdir = os.path.dirname(args[2]) aligndoc = os.path.join(outdir, "chains_params.txt") write_text_row(best_params, aligndoc) """ # This was an effort to get number of loops, inconclusive, to say the least from numpy import outer, zeros, float32, sqrt lend = len(d) cor = zeros(lend,float32) cor = outer(cor, cor) for i in xrange(lend): cor[i][i] = 1.0 for i in xrange(lend-1): for j in xrange(i+1, lend): cor[i,j] = lccc[mono(i,j)][0] cor[j,i] = cor[i,j] lmbd, eigvec = pca(cor) from sp_utilities import write_text_file nvec=20 print [lmbd[j] for j in xrange(nvec)] print " G" mm = [-1]*lend for i in xrange(lend): # row mi = -1.0e23 for j in xrange(nvec): qt = eigvec[j][i] if(abs(qt)>mi): mi = abs(qt) mm[i] = j for j in xrange(nvec): qt = eigvec[j][i] print round(qt,3), # eigenvector print mm[i] print for j in xrange(nvec): qt = [] for i in xrange(lend): if(mm[i] == j): qt.append(i) if(len(qt)>0): write_text_file(qt,"loop%02d.txt"%j) """ """ print [lmbd[j] for j in xrange(nvec)] print " B" mm = [-1]*lend for i in xrange(lend): # row mi = -1.0e23 for j in xrange(nvec): qt = eigvec[j][i]/sqrt(lmbd[j]) if(abs(qt)>mi): mi = abs(qt) mm[i] = j for j in xrange(nvec): qt = eigvec[j][i]/sqrt(lmbd[j]) print round(qt,3), # eigenvector print mm[i] print """ """ lend=3 cor = zeros(lend,float32) cor = outer(cor, cor) cor[0][0] =136.77 cor[0][1] = 79.15 cor[0][2] = 37.13 cor[1][0] = 79.15 cor[2][0] = 37.13 cor[1][1] = 50.04 cor[1][2] = 21.65 cor[2][1] = 21.65 cor[2][2] = 13.26 lmbd, eigvec = pca(cor) print lmbd print eigvec for i in xrange(lend): # row for j in xrange(lend): print eigvec[j][i], # eigenvector print print " B" for i in xrange(lend): # row for j in xrange(lend): print eigvec[j][i]/sqrt(lmbd[j]), # eigenvector print print " G" for i in xrange(lend): # row for j in xrange(lend): print eigvec[j][i]*sqrt(lmbd[j]), # eigenvector print """ else: nargs = len(args) if nargs != 2: ERROR("Must provide name of input and output file!") return from sp_utilities import get_params2D, model_circle from sp_fundamentals import rot_shift2D from sp_statistics import ccc from time import time from sp_alignment import align2d stack = args[0] new_stack = args[1] d = EMData.read_images(stack) try: sxprint("Using 2D alignment parameters from header.") ttt = d[0].get_attr('xform.params2d') for i in range(len(d)): alpha, sx, sy, mirror, scale = get_params2D(d[i]) d[i] = rot_shift2D(d[i], alpha, sx, sy, mirror) except: pass nx = d[0].get_xsize() ny = d[0].get_ysize() if options.radius < 1: radius = old_div(nx, 2) - 2 else: radius = options.radius mask = model_circle(radius, nx, ny) init = options.initial if init > -1: sxprint(" initial image: %d" % init) temp = d[init].copy() temp.write_image(new_stack, 0) del d[init] k = 1 lsum = 0.0 while len(d) > 1: maxcit = -111. for i in range(len(d)): cuc = ccc(d[i], temp, mask) if cuc > maxcit: maxcit = cuc qi = i # sxprint k, maxcit lsum += maxcit temp = d[qi].copy() del d[qi] temp.write_image(new_stack, k) k += 1 sxprint(lsum) d[0].write_image(new_stack, k) else: if options.circular: sxprint("Using options.circular, no alignment") # figure the "best circular" starting image maxsum = -1.023 for m in range(len(d)): indc = list(range(len(d))) lsnake = [-1] * (len(d) + 1) lsnake[0] = m lsnake[-1] = m del indc[m] temp = d[m].copy() lsum = 0.0 direction = +1 k = 1 while len(indc) > 1: maxcit = -111. for i in range(len(indc)): cuc = ccc(d[indc[i]], temp, mask) if cuc > maxcit: maxcit = cuc qi = indc[i] lsnake[k] = qi lsum += maxcit del indc[indc.index(qi)] direction = -direction for i in range(1, len(d)): if (direction > 0): if (lsnake[i] == -1): temp = d[lsnake[i - 1]].copy() # print " forw ",lsnake[i-1] k = i break else: if (lsnake[len(d) - i] == -1): temp = d[lsnake[len(d) - i + 1]].copy() # print " back ",lsnake[len(d) - i +1] k = len(d) - i break lsnake[lsnake.index(-1)] = indc[-1] # print " initial image and lsum ",m,lsum # print lsnake if (lsum > maxsum): maxsum = lsum init = m snake = [lsnake[i] for i in range(len(d))] sxprint(" Initial image selected : ", init, maxsum) sxprint(lsnake) for m in range(len(d)): d[snake[m]].write_image(new_stack, m) else: # figure the "best" starting image sxprint("Straight chain, no alignment") maxsum = -1.023 for m in range(len(d)): indc = list(range(len(d))) lsnake = [m] del indc[m] temp = d[m].copy() lsum = 0.0 while len(indc) > 1: maxcit = -111. for i in range(len(indc)): cuc = ccc(d[indc[i]], temp, mask) if cuc > maxcit: maxcit = cuc qi = indc[i] lsnake.append(qi) lsum += maxcit temp = d[qi].copy() del indc[indc.index(qi)] lsnake.append(indc[-1]) # sxprint " initial image and lsum ",m,lsum # sxprint lsnake if (lsum > maxsum): maxsum = lsum init = m snake = [lsnake[i] for i in range(len(d))] sxprint(" Initial image selected : ", init, maxsum) sxprint(lsnake) for m in range(len(d)): d[snake[m]].write_image(new_stack, m)
def main(): from sp_utilities import get_input_from_string progname = os.path.basename(sys.argv[0]) usage = progname + " stack output_average --radius=particle_radius --xr=xr --yr=yr --ts=ts --thld_err=thld_err --num_ali=num_ali --fl=fl --aa=aa --CTF --verbose --stables" parser = OptionParser(usage, version=SPARXVERSION) parser.add_option("--radius", type="int", default=-1, help=" particle radius for alignment") parser.add_option( "--xr", type="string", default="2 1", help= "range for translation search in x direction, search is +/xr (default 2,1)" ) parser.add_option( "--yr", type="string", default="-1", help= "range for translation search in y direction, search is +/yr (default = same as xr)" ) parser.add_option( "--ts", type="string", default="1 0.5", help= "step size of the translation search in both directions, search is -xr, -xr+ts, 0, xr-ts, xr, can be fractional (default: 1,0.5)" ) parser.add_option("--thld_err", type="float", default=0.7, help="threshld of pixel error (default = 0.75)") parser.add_option( "--num_ali", type="int", default=5, help="number of alignments performed for stability (default = 5)") parser.add_option("--maxit", type="int", default=30, help="number of iterations for each xr (default = 30)") parser.add_option( "--fl", type="float", default=0.45, help= "cut-off frequency of hyperbolic tangent low-pass Fourier filter (default = 0.3)" ) parser.add_option( "--aa", type="float", default=0.2, help= "fall-off of hyperbolic tangent low-pass Fourier filter (default = 0.2)" ) parser.add_option("--CTF", action="store_true", default=False, help="Use CTF correction during the alignment ") parser.add_option("--verbose", action="store_true", default=False, help="print individual pixel error (default = False)") parser.add_option( "--stables", action="store_true", default=False, help="output the stable particles number in file (default = False)") parser.add_option( "--method", type="string", default=" ", help="SHC (standard method is default when flag is ommitted)") (options, args) = parser.parse_args() if len(args) != 1 and len(args) != 2: sxprint("Usage: " + usage) sxprint("Please run \'" + progname + " -h\' for detailed options") ERROR( "Invalid number of parameters used. Please see usage information above." ) return else: if sp_global_def.CACHE_DISABLE: from sp_utilities import disable_bdb_cache disable_bdb_cache() from sp_applications import within_group_refinement, ali2d_ras from sp_pixel_error import multi_align_stability from sp_utilities import write_text_file, write_text_row sp_global_def.BATCH = True xrng = get_input_from_string(options.xr) if options.yr == "-1": yrng = xrng else: yrng = get_input_from_string(options.yr) step = get_input_from_string(options.ts) class_data = EMData.read_images(args[0]) nx = class_data[0].get_xsize() ou = options.radius num_ali = options.num_ali if ou == -1: ou = nx / 2 - 2 from sp_utilities import model_circle, get_params2D, set_params2D mask = model_circle(ou, nx, nx) if options.CTF: from sp_filter import filt_ctf for im in range(len(class_data)): # Flip phases class_data[im] = filt_ctf(class_data[im], class_data[im].get_attr("ctf"), binary=1) for im in class_data: im.set_attr("previousmax", -1.0e10) try: t = im.get_attr( "xform.align2d") # if they are there, no need to set them! except: try: t = im.get_attr("xform.projection") d = t.get_params("spider") set_params2D(im, [0.0, -d["tx"], -d["ty"], 0, 1.0]) except: set_params2D(im, [0.0, 0.0, 0.0, 0, 1.0]) all_ali_params = [] for ii in range(num_ali): ali_params = [] if options.verbose: ALPHA = [] SX = [] SY = [] MIRROR = [] if (xrng[0] == 0.0 and yrng[0] == 0.0): avet = ali2d_ras(class_data, randomize = True, ir = 1, ou = ou, rs = 1, step = 1.0, dst = 90.0, \ maxit = options.maxit, check_mirror = True, FH=options.fl, FF=options.aa) else: avet = within_group_refinement(class_data, mask, True, 1, ou, 1, xrng, yrng, step, 90.0, \ maxit = options.maxit, FH=options.fl, FF=options.aa, method = options.method) from sp_utilities import info #print " avet ",info(avet) for im in class_data: alpha, sx, sy, mirror, scale = get_params2D(im) ali_params.extend([alpha, sx, sy, mirror]) if options.verbose: ALPHA.append(alpha) SX.append(sx) SY.append(sy) MIRROR.append(mirror) all_ali_params.append(ali_params) if options.verbose: write_text_file([ALPHA, SX, SY, MIRROR], "ali_params_run_%d" % ii) """ avet = class_data[0] from sp_utilities import read_text_file all_ali_params = [] for ii in xrange(5): temp = read_text_file( "ali_params_run_%d"%ii,-1) uuu = [] for k in xrange(len(temp[0])): uuu.extend([temp[0][k],temp[1][k],temp[2][k],temp[3][k]]) all_ali_params.append(uuu) """ stable_set, mir_stab_rate, pix_err = multi_align_stability( all_ali_params, 0.0, 10000.0, options.thld_err, options.verbose, 2 * ou + 1) sxprint("%4s %20s %20s %20s %30s %6.2f" % ("", "Size of set", "Size of stable set", "Mirror stab rate", "Pixel error prior to pruning the set above threshold of", options.thld_err)) sxprint("Average stat: %10d %20d %20.2f %15.2f" % (len(class_data), len(stable_set), mir_stab_rate, pix_err)) if (len(stable_set) > 0): if options.stables: stab_mem = [[0, 0.0, 0] for j in range(len(stable_set))] for j in range(len(stable_set)): stab_mem[j] = [int(stable_set[j][1]), stable_set[j][0], j] write_text_row(stab_mem, "stable_particles.txt") stable_set_id = [] particle_pixerr = [] for s in stable_set: stable_set_id.append(s[1]) particle_pixerr.append(s[0]) from sp_fundamentals import rot_shift2D avet.to_zero() l = -1 sxprint("average parameters: angle, x-shift, y-shift, mirror") for j in stable_set_id: l += 1 sxprint(" %4d %4d %12.2f %12.2f %12.2f %1d" % (l, j, stable_set[l][2][0], stable_set[l][2][1], stable_set[l][2][2], int(stable_set[l][2][3]))) avet += rot_shift2D(class_data[j], stable_set[l][2][0], stable_set[l][2][1], stable_set[l][2][2], stable_set[l][2][3]) avet /= (l + 1) avet.set_attr('members', stable_set_id) avet.set_attr('pix_err', pix_err) avet.set_attr('pixerr', particle_pixerr) avet.write_image(args[1]) sp_global_def.BATCH = False