def main(): arglist = [] for arg in sys.argv: arglist.append(arg) progname = os.path.basename(arglist[0]) usage = progname + """ firstvolume secondvolume maskfile outputfile --wn --step --cutoff --radius --fsc --res_overall --out_ang_res --apix --MPI Compute local resolution in real space within area outlined by the maskfile and within regions wn x wn x wn """ parser = optparse.OptionParser(usage, version=global_def.SPARXVERSION) parser.add_option( "--wn", type="int", default=7, help= "Size of window within which local real-space FSC is computed. (default 7)" ) parser.add_option( "--step", type="float", default=1.0, help="Shell step in Fourier size in pixels. (default 1.0)") parser.add_option("--cutoff", type="float", default=0.5, help="Resolution cut-off for FSC. (default 0.5)") parser.add_option( "--radius", type="int", default=-1, help= "If there is no maskfile, sphere with r=radius will be used. By default, the radius is nx/2-wn (default -1)" ) parser.add_option( "--fsc", type="string", default=None, help= "Save overall FSC curve (might be truncated). By default, the program does not save the FSC curve. (default none)" ) parser.add_option( "--res_overall", type="float", default=-1.0, help= "Overall resolution at the cutoff level estimated by the user [abs units]. (default None)" ) parser.add_option( "--out_ang_res", action="store_true", default=False, help= "Additionally creates a local resolution file in Angstroms. (default False)" ) parser.add_option( "--apix", type="float", default=1.0, help= "Pixel size in Angstrom. Effective only with --out_ang_res options. (default 1.0)" ) parser.add_option("--MPI", action="store_true", default=False, help="Use MPI version.") (options, args) = parser.parse_args(arglist[1:]) if len(args) < 3 or len(args) > 4: print("See usage " + usage) sys.exit() if global_def.CACHE_DISABLE: utilities.disable_bdb_cache() res_overall = options.res_overall if options.MPI: sys.argv = mpi.mpi_init(len(sys.argv), sys.argv) number_of_proc = mpi.mpi_comm_size(mpi.MPI_COMM_WORLD) myid = mpi.mpi_comm_rank(mpi.MPI_COMM_WORLD) main_node = 0 global_def.MPI = True cutoff = options.cutoff nk = int(options.wn) if (myid == main_node): #print sys.argv vi = utilities.get_im(sys.argv[1]) ui = utilities.get_im(sys.argv[2]) nx = vi.get_xsize() ny = vi.get_ysize() nz = vi.get_zsize() dis = [nx, ny, nz] else: dis = [0, 0, 0, 0] global_def.BATCH = True dis = utilities.bcast_list_to_all(dis, myid, source_node=main_node) if (myid != main_node): nx = int(dis[0]) ny = int(dis[1]) nz = int(dis[2]) vi = utilities.model_blank(nx, ny, nz) ui = utilities.model_blank(nx, ny, nz) if len(args) == 3: m = utilities.model_circle((min(nx, ny, nz) - nk) // 2, nx, ny, nz) outvol = args[2] elif len(args) == 4: if (myid == main_node): m = morphology.binarize(utilities.get_im(args[2]), 0.5) else: m = utilities.model_blank(nx, ny, nz) outvol = args[3] utilities.bcast_EMData_to_all(m, myid, main_node) """Multiline Comment0""" freqvol, resolut = statistics.locres(vi, ui, m, nk, cutoff, options.step, myid, main_node, number_of_proc) if (myid == 0): # Remove outliers based on the Interquartile range output_volume(freqvol, resolut, options.apix, outvol, options.fsc, options.out_ang_res, nx, ny, nz, res_overall) mpi.mpi_finalize() else: cutoff = options.cutoff vi = utilities.get_im(args[0]) ui = utilities.get_im(args[1]) nn = vi.get_xsize() nk = int(options.wn) if len(args) == 3: m = utilities.model_circle((nn - nk) // 2, nn, nn, nn) outvol = args[2] elif len(args) == 4: m = morphology.binarize(utilities.get_im(args[2]), 0.5) outvol = args[3] mc = utilities.model_blank(nn, nn, nn, 1.0) - m vf = fundamentals.fft(vi) uf = fundamentals.fft(ui) """Multiline Comment1""" lp = int(nn / 2 / options.step + 0.5) step = 0.5 / lp freqvol = utilities.model_blank(nn, nn, nn) resolut = [] for i in range(1, lp): fl = step * i fh = fl + step #print(lp,i,step,fl,fh) v = fundamentals.fft(filter.filt_tophatb(vf, fl, fh)) u = fundamentals.fft(filter.filt_tophatb(uf, fl, fh)) tmp1 = EMAN2_cppwrap.Util.muln_img(v, v) tmp2 = EMAN2_cppwrap.Util.muln_img(u, u) do = EMAN2_cppwrap.Util.infomask( morphology.square_root( morphology.threshold( EMAN2_cppwrap.Util.muln_img(tmp1, tmp2))), m, True)[0] tmp3 = EMAN2_cppwrap.Util.muln_img(u, v) dp = EMAN2_cppwrap.Util.infomask(tmp3, m, True)[0] resolut.append([i, (fl + fh) / 2.0, dp / do]) tmp1 = EMAN2_cppwrap.Util.box_convolution(tmp1, nk) tmp2 = EMAN2_cppwrap.Util.box_convolution(tmp2, nk) tmp3 = EMAN2_cppwrap.Util.box_convolution(tmp3, nk) EMAN2_cppwrap.Util.mul_img(tmp1, tmp2) tmp1 = morphology.square_root(morphology.threshold(tmp1)) EMAN2_cppwrap.Util.mul_img(tmp1, m) EMAN2_cppwrap.Util.add_img(tmp1, mc) EMAN2_cppwrap.Util.mul_img(tmp3, m) EMAN2_cppwrap.Util.add_img(tmp3, mc) EMAN2_cppwrap.Util.div_img(tmp3, tmp1) EMAN2_cppwrap.Util.mul_img(tmp3, m) freq = (fl + fh) / 2.0 bailout = True for x in range(nn): for y in range(nn): for z in range(nn): if (m.get_value_at(x, y, z) > 0.5): if (freqvol.get_value_at(x, y, z) == 0.0): if (tmp3.get_value_at(x, y, z) < cutoff): freqvol.set_value_at(x, y, z, freq) bailout = False else: bailout = False if (bailout): break #print(len(resolut)) # remove outliers output_volume(freqvol, resolut, options.apix, outvol, options.fsc, options.out_ang_res, nx, ny, nz, res_overall)
def main(): def params_3D_2D_NEW(phi, theta, psi, s2x, s2y, mirror): if mirror: m = 1 alpha, sx, sy, scalen = compose_transform2(0, s2x, s2y, 1.0, 540.0 - psi, 0, 0, 1.0) else: m = 0 alpha, sx, sy, scalen = compose_transform2(0, s2x, s2y, 1.0, 360.0 - psi, 0, 0, 1.0) return alpha, sx, sy, m progname = os.path.basename(sys.argv[0]) usage = progname + " prj_stack --ave2D= --var2D= --ave3D= --var3D= --img_per_grp= --fl=15. --aa=0.01 --sym=symmetry --CTF" parser = OptionParser(usage, version=SPARXVERSION) parser.add_option("--output_dir", type="string", default="./", help="output directory") parser.add_option("--ave2D", type="string", default=False, help="write to the disk a stack of 2D averages") parser.add_option("--var2D", type="string", default=False, help="write to the disk a stack of 2D variances") parser.add_option("--ave3D", type="string", default=False, help="write to the disk reconstructed 3D average") parser.add_option("--var3D", type="string", default=False, help="compute 3D variability (time consuming!)") parser.add_option("--img_per_grp", type="int", default=10, help="number of neighbouring projections") parser.add_option("--no_norm", action="store_true", default=False, help="do not use normalization") #parser.add_option("--radius", type="int" , default=-1 , help="radius for 3D variability" ) parser.add_option("--npad", type="int", default=2, help="number of time to pad the original images") parser.add_option("--sym", type="string", default="c1", help="symmetry") parser.add_option( "--fl", type="float", default=0.0, help= "cutoff freqency in absolute frequency (0.0-0.5). (Default - no filtration)" ) parser.add_option( "--aa", type="float", default=0.0, help= "fall off of the filter. Put 0.01 if user has no clue about falloff (Default - no filtration)" ) parser.add_option("--CTF", action="store_true", default=False, help="use CFT correction") parser.add_option("--VERBOSE", action="store_true", default=False, help="Long output for debugging") #parser.add_option("--MPI" , action="store_true", default=False, help="use MPI version") #parser.add_option("--radiuspca", type="int" , default=-1 , help="radius for PCA" ) #parser.add_option("--iter", type="int" , default=40 , help="maximum number of iterations (stop criterion of reconstruction process)" ) #parser.add_option("--abs", type="float" , default=0.0 , help="minimum average absolute change of voxels' values (stop criterion of reconstruction process)" ) #parser.add_option("--squ", type="float" , default=0.0 , help="minimum average squared change of voxels' values (stop criterion of reconstruction process)" ) parser.add_option( "--VAR", action="store_true", default=False, help="stack on input consists of 2D variances (Default False)") parser.add_option( "--decimate", type="float", default=1.0, help= "image decimate rate, a number larger (expand image) or less (shrink image) than 1. default is 1" ) parser.add_option( "--window", type="int", default=0, help= "reduce images to a small image size without changing pixel_size. Default value is zero." ) #parser.add_option("--SND", action="store_true", default=False, help="compute squared normalized differences (Default False)") parser.add_option( "--nvec", type="int", default=0, help="number of eigenvectors, default = 0 meaning no PCA calculated") parser.add_option( "--symmetrize", action="store_true", default=False, help="Prepare input stack for handling symmetry (Default False)") (options, args) = parser.parse_args() ##### from mpi import mpi_init, mpi_comm_rank, mpi_comm_size, mpi_recv, MPI_COMM_WORLD from mpi import mpi_barrier, mpi_reduce, mpi_bcast, mpi_send, MPI_FLOAT, MPI_SUM, MPI_INT, MPI_MAX from applications import MPI_start_end from reconstruction import recons3d_em, recons3d_em_MPI from reconstruction import recons3d_4nn_MPI, recons3d_4nn_ctf_MPI from utilities import print_begin_msg, print_end_msg, print_msg from utilities import read_text_row, get_image, get_im from utilities import bcast_EMData_to_all, bcast_number_to_all from utilities import get_symt # This is code for handling symmetries by the above program. To be incorporated. PAP 01/27/2015 from EMAN2db import db_open_dict # Set up global variables related to bdb cache if global_def.CACHE_DISABLE: from utilities import disable_bdb_cache disable_bdb_cache() # Set up global variables related to ERROR function global_def.BATCH = True # detect if program is running under MPI RUNNING_UNDER_MPI = "OMPI_COMM_WORLD_SIZE" in os.environ if RUNNING_UNDER_MPI: global_def.MPI = True if options.symmetrize: if RUNNING_UNDER_MPI: try: sys.argv = mpi_init(len(sys.argv), sys.argv) try: number_of_proc = mpi_comm_size(MPI_COMM_WORLD) if (number_of_proc > 1): ERROR( "Cannot use more than one CPU for symmetry prepration", "sx3dvariability", 1) except: pass except: pass if options.output_dir != "./" and not os.path.exists( options.output_dir): os.mkdir(options.output_dir) # Input #instack = "Clean_NORM_CTF_start_wparams.hdf" #instack = "bdb:data" from logger import Logger, BaseLogger_Files if os.path.exists(os.path.join(options.output_dir, "log.txt")): os.remove(os.path.join(options.output_dir, "log.txt")) log_main = Logger(BaseLogger_Files()) log_main.prefix = os.path.join(options.output_dir, "./") instack = args[0] sym = options.sym.lower() if (sym == "c1"): ERROR("There is no need to symmetrize stack for C1 symmetry", "sx3dvariability", 1) line = "" for a in sys.argv: line += " " + a log_main.add(line) if (instack[:4] != "bdb:"): if output_dir == "./": stack = "bdb:data" else: stack = "bdb:" + options.output_dir + "/data" delete_bdb(stack) junk = cmdexecute("sxcpy.py " + instack + " " + stack) else: stack = instack qt = EMUtil.get_all_attributes(stack, 'xform.projection') na = len(qt) ts = get_symt(sym) ks = len(ts) angsa = [None] * na for k in xrange(ks): #Qfile = "Q%1d"%k if options.output_dir != "./": Qfile = os.path.join(options.output_dir, "Q%1d" % k) else: Qfile = os.path.join(options.output_dir, "Q%1d" % k) #delete_bdb("bdb:Q%1d"%k) delete_bdb("bdb:" + Qfile) #junk = cmdexecute("e2bdb.py "+stack+" --makevstack=bdb:Q%1d"%k) junk = cmdexecute("e2bdb.py " + stack + " --makevstack=bdb:" + Qfile) #DB = db_open_dict("bdb:Q%1d"%k) DB = db_open_dict("bdb:" + Qfile) for i in xrange(na): ut = qt[i] * ts[k] DB.set_attr(i, "xform.projection", ut) #bt = ut.get_params("spider") #angsa[i] = [round(bt["phi"],3)%360.0, round(bt["theta"],3)%360.0, bt["psi"], -bt["tx"], -bt["ty"]] #write_text_row(angsa, 'ptsma%1d.txt'%k) #junk = cmdexecute("e2bdb.py "+stack+" --makevstack=bdb:Q%1d"%k) #junk = cmdexecute("sxheader.py bdb:Q%1d --params=xform.projection --import=ptsma%1d.txt"%(k,k)) DB.close() if options.output_dir == "./": delete_bdb("bdb:sdata") else: delete_bdb("bdb:" + options.output_dir + "/" + "sdata") #junk = cmdexecute("e2bdb.py . --makevstack=bdb:sdata --filt=Q") sdata = "bdb:" + options.output_dir + "/" + "sdata" print(sdata) junk = cmdexecute("e2bdb.py " + options.output_dir + " --makevstack=" + sdata + " --filt=Q") #junk = cmdexecute("ls EMAN2DB/sdata*") #a = get_im("bdb:sdata") a = get_im(sdata) a.set_attr("variabilitysymmetry", sym) #a.write_image("bdb:sdata") a.write_image(sdata) else: sys.argv = mpi_init(len(sys.argv), sys.argv) myid = mpi_comm_rank(MPI_COMM_WORLD) number_of_proc = mpi_comm_size(MPI_COMM_WORLD) main_node = 0 if len(args) == 1: stack = args[0] else: print(("usage: " + usage)) print(("Please run '" + progname + " -h' for detailed options")) return 1 t0 = time() # obsolete flags options.MPI = True options.nvec = 0 options.radiuspca = -1 options.iter = 40 options.abs = 0.0 options.squ = 0.0 if options.fl > 0.0 and options.aa == 0.0: ERROR("Fall off has to be given for the low-pass filter", "sx3dvariability", 1, myid) if options.VAR and options.SND: ERROR("Only one of var and SND can be set!", "sx3dvariability", myid) exit() if options.VAR and (options.ave2D or options.ave3D or options.var2D): ERROR( "When VAR is set, the program cannot output ave2D, ave3D or var2D", "sx3dvariability", 1, myid) exit() #if options.SND and (options.ave2D or options.ave3D): # ERROR("When SND is set, the program cannot output ave2D or ave3D", "sx3dvariability", 1, myid) # exit() if options.nvec > 0: ERROR("PCA option not implemented", "sx3dvariability", 1, myid) exit() if options.nvec > 0 and options.ave3D == None: ERROR("When doing PCA analysis, one must set ave3D", "sx3dvariability", myid=myid) exit() import string options.sym = options.sym.lower() # if global_def.CACHE_DISABLE: # from utilities import disable_bdb_cache # disable_bdb_cache() # global_def.BATCH = True if myid == main_node: if options.output_dir != "./" and not os.path.exists( options.output_dir): os.mkdir(options.output_dir) img_per_grp = options.img_per_grp nvec = options.nvec radiuspca = options.radiuspca from logger import Logger, BaseLogger_Files #if os.path.exists(os.path.join(options.output_dir, "log.txt")): os.remove(os.path.join(options.output_dir, "log.txt")) log_main = Logger(BaseLogger_Files()) log_main.prefix = os.path.join(options.output_dir, "./") if myid == main_node: line = "" for a in sys.argv: line += " " + a log_main.add(line) log_main.add("-------->>>Settings given by all options<<<-------") log_main.add("instack :" + stack) log_main.add("output_dir :" + options.output_dir) log_main.add("var3d :" + options.var3D) if myid == main_node: line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" #print_begin_msg("sx3dvariability") msg = "sx3dvariability" log_main.add(msg) print(line, msg) msg = ("%-70s: %s\n" % ("Input stack", stack)) log_main.add(msg) print(line, msg) symbaselen = 0 if myid == main_node: nima = EMUtil.get_image_count(stack) img = get_image(stack) nx = img.get_xsize() ny = img.get_ysize() if options.sym != "c1": imgdata = get_im(stack) try: i = imgdata.get_attr("variabilitysymmetry").lower() if (i != options.sym): ERROR( "The symmetry provided does not agree with the symmetry of the input stack", "sx3dvariability", myid=myid) except: ERROR( "Input stack is not prepared for symmetry, please follow instructions", "sx3dvariability", myid=myid) from utilities import get_symt i = len(get_symt(options.sym)) if ((nima / i) * i != nima): ERROR( "The length of the input stack is incorrect for symmetry processing", "sx3dvariability", myid=myid) symbaselen = nima / i else: symbaselen = nima 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) Tracker = {} Tracker["total_stack"] = nima if options.decimate == 1.: if options.window != 0: nx = options.window ny = options.window else: if options.window == 0: nx = int(nx * options.decimate) ny = int(ny * options.decimate) else: nx = int(options.window * options.decimate) ny = nx Tracker["nx"] = nx Tracker["ny"] = ny Tracker["nz"] = nx symbaselen = bcast_number_to_all(symbaselen) if radiuspca == -1: radiuspca = nx / 2 - 2 if myid == main_node: line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" msg = "%-70s: %d\n" % ("Number of projection", nima) log_main.add(msg) print(line, msg) img_begin, img_end = MPI_start_end(nima, number_of_proc, myid) """ if options.SND: from projection import prep_vol, prgs from statistics import im_diff from utilities import get_im, model_circle, get_params_proj, set_params_proj from utilities import get_ctf, generate_ctf from filter import filt_ctf imgdata = EMData.read_images(stack, range(img_begin, img_end)) if options.CTF: vol = recons3d_4nn_ctf_MPI(myid, imgdata, 1.0, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1) else: vol = recons3d_4nn_MPI(myid, imgdata, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1) bcast_EMData_to_all(vol, myid) volft, kb = prep_vol(vol) mask = model_circle(nx/2-2, nx, ny) varList = [] for i in xrange(img_begin, img_end): phi, theta, psi, s2x, s2y = get_params_proj(imgdata[i-img_begin]) ref_prj = prgs(volft, kb, [phi, theta, psi, -s2x, -s2y]) if options.CTF: ctf_params = get_ctf(imgdata[i-img_begin]) ref_prj = filt_ctf(ref_prj, generate_ctf(ctf_params)) diff, A, B = im_diff(ref_prj, imgdata[i-img_begin], mask) diff2 = diff*diff set_params_proj(diff2, [phi, theta, psi, s2x, s2y]) varList.append(diff2) mpi_barrier(MPI_COMM_WORLD) """ if options.VAR: #varList = EMData.read_images(stack, range(img_begin, img_end)) varList = [] this_image = EMData() for index_of_particle in xrange(img_begin, img_end): this_image.read_image(stack, index_of_particle) varList.append( image_decimate_window_xform_ctf(this_image, options.decimate, options.window, options.CTF)) else: from utilities import bcast_number_to_all, bcast_list_to_all, send_EMData, recv_EMData from utilities import set_params_proj, get_params_proj, params_3D_2D, get_params2D, set_params2D, compose_transform2 from utilities import model_blank, nearest_proj, model_circle from applications import pca from statistics import avgvar, avgvar_ctf, ccc from filter import filt_tanl from morphology import threshold, square_root from projection import project, prep_vol, prgs from sets import Set if myid == main_node: t1 = time() proj_angles = [] aveList = [] tab = EMUtil.get_all_attributes(stack, 'xform.projection') for i in xrange(nima): t = tab[i].get_params('spider') phi = t['phi'] theta = t['theta'] psi = t['psi'] x = theta if x > 90.0: x = 180.0 - x x = x * 10000 + psi proj_angles.append([x, t['phi'], t['theta'], t['psi'], i]) t2 = time() line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" msg = "%-70s: %d\n" % ("Number of neighboring projections", img_per_grp) log_main.add(msg) print(line, msg) msg = "...... Finding neighboring projections\n" log_main.add(msg) print(line, msg) if options.VERBOSE: msg = "Number of images per group: %d" % img_per_grp log_main.add(msg) print(line, msg) msg = "Now grouping projections" log_main.add(msg) print(line, msg) proj_angles.sort() proj_angles_list = [0.0] * (nima * 4) if myid == main_node: for i in xrange(nima): proj_angles_list[i * 4] = proj_angles[i][1] proj_angles_list[i * 4 + 1] = proj_angles[i][2] proj_angles_list[i * 4 + 2] = proj_angles[i][3] proj_angles_list[i * 4 + 3] = proj_angles[i][4] proj_angles_list = bcast_list_to_all(proj_angles_list, myid, main_node) proj_angles = [] for i in xrange(nima): proj_angles.append([ proj_angles_list[i * 4], proj_angles_list[i * 4 + 1], proj_angles_list[i * 4 + 2], int(proj_angles_list[i * 4 + 3]) ]) del proj_angles_list proj_list, mirror_list = nearest_proj(proj_angles, img_per_grp, range(img_begin, img_end)) all_proj = Set() for im in proj_list: for jm in im: all_proj.add(proj_angles[jm][3]) all_proj = list(all_proj) if options.VERBOSE: print("On node %2d, number of images needed to be read = %5d" % (myid, len(all_proj))) index = {} for i in xrange(len(all_proj)): index[all_proj[i]] = i mpi_barrier(MPI_COMM_WORLD) if myid == main_node: line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" msg = ("%-70s: %.2f\n" % ("Finding neighboring projections lasted [s]", time() - t2)) log_main.add(msg) print(msg) msg = ("%-70s: %d\n" % ("Number of groups processed on the main node", len(proj_list))) log_main.add(msg) print(line, msg) if options.VERBOSE: print("Grouping projections took: ", (time() - t2) / 60, "[min]") print("Number of groups on main node: ", len(proj_list)) mpi_barrier(MPI_COMM_WORLD) if myid == main_node: line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" msg = ("...... calculating the stack of 2D variances \n") log_main.add(msg) print(line, msg) if options.VERBOSE: print("Now calculating the stack of 2D variances") proj_params = [0.0] * (nima * 5) aveList = [] varList = [] if nvec > 0: eigList = [[] for i in xrange(nvec)] if options.VERBOSE: print("Begin to read images on processor %d" % (myid)) ttt = time() #imgdata = EMData.read_images(stack, all_proj) imgdata = [] for index_of_proj in xrange(len(all_proj)): #img = EMData() #img.read_image(stack, all_proj[index_of_proj]) dmg = image_decimate_window_xform_ctf( get_im(stack, all_proj[index_of_proj]), options.decimate, options.window, options.CTF) #print dmg.get_xsize(), "init" imgdata.append(dmg) if options.VERBOSE: print("Reading images on processor %d done, time = %.2f" % (myid, time() - ttt)) print("On processor %d, we got %d images" % (myid, len(imgdata))) mpi_barrier(MPI_COMM_WORLD) ''' imgdata2 = EMData.read_images(stack, range(img_begin, img_end)) if options.fl > 0.0: for k in xrange(len(imgdata2)): imgdata2[k] = filt_tanl(imgdata2[k], options.fl, options.aa) if options.CTF: vol = recons3d_4nn_ctf_MPI(myid, imgdata2, 1.0, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1) else: vol = recons3d_4nn_MPI(myid, imgdata2, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1) if myid == main_node: vol.write_image("vol_ctf.hdf") print_msg("Writing to the disk volume reconstructed from averages as : %s\n"%("vol_ctf.hdf")) del vol, imgdata2 mpi_barrier(MPI_COMM_WORLD) ''' from applications import prepare_2d_forPCA from utilities import model_blank for i in xrange(len(proj_list)): ki = proj_angles[proj_list[i][0]][3] if ki >= symbaselen: continue mi = index[ki] phiM, thetaM, psiM, s2xM, s2yM = get_params_proj(imgdata[mi]) grp_imgdata = [] for j in xrange(img_per_grp): mj = index[proj_angles[proj_list[i][j]][3]] phi, theta, psi, s2x, s2y = get_params_proj(imgdata[mj]) alpha, sx, sy, mirror = params_3D_2D_NEW( phi, theta, psi, s2x, s2y, mirror_list[i][j]) if thetaM <= 90: if mirror == 0: alpha, sx, sy, scale = compose_transform2( alpha, sx, sy, 1.0, phiM - phi, 0.0, 0.0, 1.0) else: alpha, sx, sy, scale = compose_transform2( alpha, sx, sy, 1.0, 180 - (phiM - phi), 0.0, 0.0, 1.0) else: if mirror == 0: alpha, sx, sy, scale = compose_transform2( alpha, sx, sy, 1.0, -(phiM - phi), 0.0, 0.0, 1.0) else: alpha, sx, sy, scale = compose_transform2( alpha, sx, sy, 1.0, -(180 - (phiM - phi)), 0.0, 0.0, 1.0) set_params2D(imgdata[mj], [alpha, sx, sy, mirror, 1.0]) grp_imgdata.append(imgdata[mj]) #print grp_imgdata[j].get_xsize(), imgdata[mj].get_xsize() if not options.no_norm: #print grp_imgdata[j].get_xsize() mask = model_circle(nx / 2 - 2, nx, nx) for k in xrange(img_per_grp): ave, std, minn, maxx = Util.infomask( grp_imgdata[k], mask, False) grp_imgdata[k] -= ave grp_imgdata[k] /= std del mask if options.fl > 0.0: from filter import filt_ctf, filt_table from fundamentals import fft, window2d nx2 = 2 * nx ny2 = 2 * ny if options.CTF: from utilities import pad for k in xrange(img_per_grp): grp_imgdata[k] = window2d( fft( filt_tanl( filt_ctf( fft( pad(grp_imgdata[k], nx2, ny2, 1, 0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa)), nx, ny) #grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny) #grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa) else: for k in xrange(img_per_grp): grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa) #grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny) #grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa) else: from utilities import pad, read_text_file from filter import filt_ctf, filt_table from fundamentals import fft, window2d nx2 = 2 * nx ny2 = 2 * ny if options.CTF: from utilities import pad for k in xrange(img_per_grp): grp_imgdata[k] = window2d( fft( filt_ctf(fft( pad(grp_imgdata[k], nx2, ny2, 1, 0.0)), grp_imgdata[k].get_attr("ctf"), binary=1)), nx, ny) #grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny) #grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa) ''' if i < 10 and myid == main_node: for k in xrange(10): grp_imgdata[k].write_image("grp%03d.hdf"%i, k) ''' """ if myid == main_node and i==0: for pp in xrange(len(grp_imgdata)): grp_imgdata[pp].write_image("pp.hdf", pp) """ ave, grp_imgdata = prepare_2d_forPCA(grp_imgdata) """ if myid == main_node and i==0: for pp in xrange(len(grp_imgdata)): grp_imgdata[pp].write_image("qq.hdf", pp) """ var = model_blank(nx, ny) for q in grp_imgdata: Util.add_img2(var, q) Util.mul_scalar(var, 1.0 / (len(grp_imgdata) - 1)) # Switch to std dev var = square_root(threshold(var)) #if options.CTF: ave, var = avgvar_ctf(grp_imgdata, mode="a") #else: ave, var = avgvar(grp_imgdata, mode="a") """ if myid == main_node: ave.write_image("avgv.hdf",i) var.write_image("varv.hdf",i) """ set_params_proj(ave, [phiM, thetaM, 0.0, 0.0, 0.0]) set_params_proj(var, [phiM, thetaM, 0.0, 0.0, 0.0]) aveList.append(ave) varList.append(var) if options.VERBOSE: print("%5.2f%% done on processor %d" % (i * 100.0 / len(proj_list), myid)) if nvec > 0: eig = pca(input_stacks=grp_imgdata, subavg="", mask_radius=radiuspca, nvec=nvec, incore=True, shuffle=False, genbuf=True) for k in xrange(nvec): set_params_proj(eig[k], [phiM, thetaM, 0.0, 0.0, 0.0]) eigList[k].append(eig[k]) """ if myid == 0 and i == 0: for k in xrange(nvec): eig[k].write_image("eig.hdf", k) """ del imgdata # To this point, all averages, variances, and eigenvectors are computed if options.ave2D: from fundamentals import fpol if myid == main_node: km = 0 for i in xrange(number_of_proc): if i == main_node: for im in xrange(len(aveList)): aveList[im].write_image( os.path.join(options.output_dir, options.ave2D), 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 xrange(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', map(int, members)) members = mpi_recv(nm, MPI_FLOAT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) ave.set_attr('pix_err', map(float, members)) members = mpi_recv(3, MPI_FLOAT, i, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) ave.set_attr('refprojdir', map(float, members)) """ tmpvol = fpol(ave, Tracker["nx"], Tracker["nx"], 1) tmpvol.write_image( os.path.join(options.output_dir, options.ave2D), km) km += 1 else: mpi_send(len(aveList), 1, MPI_INT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) for im in xrange(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('pix_err') 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) """ if options.ave3D: from fundamentals import fpol if options.VERBOSE: print("Reconstructing 3D average volume") ave3D = recons3d_4nn_MPI(myid, aveList, symmetry=options.sym, npad=options.npad) bcast_EMData_to_all(ave3D, myid) if myid == main_node: line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" ave3D = fpol(ave3D, Tracker["nx"], Tracker["nx"], Tracker["nx"]) ave3D.write_image( os.path.join(options.output_dir, options.ave3D)) msg = ("%-70s: %s\n" % ( "Writing to the disk volume reconstructed from averages as", options.ave3D)) log_main.add(msg) print(line, msg) del ave, var, proj_list, stack, phi, theta, psi, s2x, s2y, alpha, sx, sy, mirror, aveList if nvec > 0: for k in xrange(nvec): if options.VERBOSE: print("Reconstruction eigenvolumes", k) cont = True ITER = 0 mask2d = model_circle(radiuspca, nx, nx) while cont: #print "On node %d, iteration %d"%(myid, ITER) eig3D = recons3d_4nn_MPI(myid, eigList[k], symmetry=options.sym, npad=options.npad) bcast_EMData_to_all(eig3D, myid, main_node) if options.fl > 0.0: eig3D = filt_tanl(eig3D, options.fl, options.aa) if myid == main_node: eig3D.write_image( os.path.join(options.outpout_dir, "eig3d_%03d.hdf" % (k, ITER))) Util.mul_img(eig3D, model_circle(radiuspca, nx, nx, nx)) eig3Df, kb = prep_vol(eig3D) del eig3D cont = False icont = 0 for l in xrange(len(eigList[k])): phi, theta, psi, s2x, s2y = get_params_proj( eigList[k][l]) proj = prgs(eig3Df, kb, [phi, theta, psi, s2x, s2y]) cl = ccc(proj, eigList[k][l], mask2d) if cl < 0.0: icont += 1 cont = True eigList[k][l] *= -1.0 u = int(cont) u = mpi_reduce([u], 1, MPI_INT, MPI_MAX, main_node, MPI_COMM_WORLD) icont = mpi_reduce([icont], 1, MPI_INT, MPI_SUM, main_node, MPI_COMM_WORLD) if myid == main_node: line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" u = int(u[0]) msg = (" Eigenvector: ", k, " number changed ", int(icont[0])) log_main.add(msg) print(line, msg) else: u = 0 u = bcast_number_to_all(u, main_node) cont = bool(u) ITER += 1 del eig3Df, kb mpi_barrier(MPI_COMM_WORLD) del eigList, mask2d if options.ave3D: del ave3D if options.var2D: from fundamentals import fpol if myid == main_node: km = 0 for i in xrange(number_of_proc): if i == main_node: for im in xrange(len(varList)): tmpvol = fpol(varList[im], Tracker["nx"], Tracker["nx"], 1) tmpvol.write_image( os.path.join(options.output_dir, options.var2D), 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 xrange(nl): ave = recv_EMData(i, im + i + 70000) tmpvol = fpol(ave, Tracker["nx"], Tracker["nx"], 1) tmpvol.write_image( os.path.join(options.output_dir, options.var2D, km)) km += 1 else: mpi_send(len(varList), 1, MPI_INT, main_node, SPARX_MPI_TAG_UNIVERSAL, MPI_COMM_WORLD) for im in xrange(len(varList)): send_EMData(varList[im], main_node, im + myid + 70000) # What with the attributes?? mpi_barrier(MPI_COMM_WORLD) if options.var3D: if myid == main_node and options.VERBOSE: line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" msg = ("Reconstructing 3D variability volume") log_main.add(msg) print(line, msg) t6 = time() # radiusvar = options.radius # if( radiusvar < 0 ): radiusvar = nx//2 -3 res = recons3d_4nn_MPI(myid, varList, symmetry=options.sym, npad=options.npad) #res = recons3d_em_MPI(varList, vol_stack, options.iter, radiusvar, options.abs, True, options.sym, options.squ) if myid == main_node: from fundamentals import fpol res = fpol(res, Tracker["nx"], Tracker["nx"], Tracker["nx"]) res.write_image(os.path.join(options.output_dir, options.var3D)) if myid == main_node: line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" msg = ("%-70s: %.2f\n" % ("Reconstructing 3D variability took [s]", time() - t6)) log_main.add(msg) print(line, msg) if options.VERBOSE: print("Reconstruction took: %.2f [min]" % ((time() - t6) / 60)) if myid == main_node: line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" msg = ("%-70s: %.2f\n" % ("Total time for these computations [s]", time() - t0)) print(line, msg) log_main.add(msg) if options.VERBOSE: print("Total time for these computations: %.2f [min]" % ((time() - t0) / 60)) line = strftime("%Y-%m-%d_%H:%M:%S", localtime()) + " =>" msg = ("sx3dvariability") print(line, msg) log_main.add(msg) from mpi import mpi_finalize mpi_finalize() if RUNNING_UNDER_MPI: global_def.MPI = False global_def.BATCH = False
def helicalshiftali_MPI(stack, maskfile=None, maxit=100, CTF=False, snr=1.0, Fourvar=False, search_rng=-1): from applications import MPI_start_end from utilities import model_circle, model_blank, get_image, peak_search, get_im, pad from utilities import reduce_EMData_to_root, bcast_EMData_to_all, send_attr_dict, file_type, bcast_number_to_all, bcast_list_to_all from pap_statistics import varf2d_MPI from fundamentals import fft, ccf, rot_shift3D, rot_shift2D, fshift from utilities import get_params2D, set_params2D, chunks_distribution from utilities import print_msg, print_begin_msg, print_end_msg import os import sys from mpi import mpi_init, mpi_comm_size, mpi_comm_rank, MPI_COMM_WORLD from mpi import mpi_reduce, mpi_bcast, mpi_barrier, mpi_gatherv from mpi import MPI_SUM, MPI_FLOAT, MPI_INT from time import time from pixel_error import ordersegments from math import sqrt, atan2, tan, pi nproc = mpi_comm_size(MPI_COMM_WORLD) myid = mpi_comm_rank(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), "ehelix_MPI", 1, 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) print("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 filter import filt_ctf from 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 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', "helicalshiftali_MPI", 1, 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_reduce(sx_sum, 1, MPI_FLOAT, MPI_SUM, main_node, 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_barrier(MPI_COMM_WORLD) par_str = ["xform.align2d", "ID"] if myid == main_node: from utilities import file_type if (file_type(stack) == "bdb"): from utilities import recv_attr_dict_bdb recv_attr_dict_bdb(main_node, stack, data, par_str, 0, ldata, nproc) else: from 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 ali3d_MPI(stack, ref_vol, outdir, maskfile = None, ir = 1, ou = -1, rs = 1, xr = "4 2 2 1", yr = "-1", ts = "1 1 0.5 0.25", delta = "10 6 4 4", an = "-1", center = 0, maxit = 5, term = 95, CTF = False, fourvar = False, snr = 1.0, ref_a = "S", sym = "c1", sort=True, cutoff=999.99, pix_cutoff="0", two_tail=False, model_jump="1 1 1 1 1", restart=False, save_half=False, protos=None, oplane=None, lmask=-1, ilmask=-1, findseam=False, vertstep=None, hpars="-1", hsearch="73.0 170.0", full_output = False, compare_repro = False, compare_ref_free = "-1", ref_free_cutoff= "-1 -1 -1 -1", wcmask = None, debug = False, recon_pad = 4): from alignment import Numrinit, prepare_refrings from utilities import model_circle, get_image, drop_image, get_input_from_string from utilities import bcast_list_to_all, bcast_number_to_all, reduce_EMData_to_root, bcast_EMData_to_all from utilities import send_attr_dict from utilities import get_params_proj, file_type from fundamentals import rot_avg_image import os import types from utilities import print_begin_msg, print_end_msg, print_msg from mpi import mpi_bcast, mpi_comm_size, mpi_comm_rank, MPI_FLOAT, MPI_COMM_WORLD, mpi_barrier, mpi_reduce from mpi import mpi_reduce, MPI_INT, MPI_SUM, mpi_finalize from filter import filt_ctf from projection import prep_vol, prgs from statistics import hist_list, varf3d_MPI, fsc_mask from numpy import array, bincount, array2string, ones number_of_proc = mpi_comm_size(MPI_COMM_WORLD) myid = mpi_comm_rank(MPI_COMM_WORLD) main_node = 0 if myid == main_node: if os.path.exists(outdir): ERROR('Output directory exists, please change the name and restart the program', "ali3d_MPI", 1) os.mkdir(outdir) mpi_barrier(MPI_COMM_WORLD) if debug: from time import sleep while not os.path.exists(outdir): print "Node ",myid," waiting..." sleep(5) info_file = os.path.join(outdir, "progress%04d"%myid) finfo = open(info_file, 'w') else: finfo = None mjump = get_input_from_string(model_jump) xrng = get_input_from_string(xr) if yr == "-1": yrng = xrng else : yrng = get_input_from_string(yr) step = get_input_from_string(ts) delta = get_input_from_string(delta) ref_free_cutoff = get_input_from_string(ref_free_cutoff) pix_cutoff = get_input_from_string(pix_cutoff) lstp = min(len(xrng), len(yrng), len(step), len(delta)) if an == "-1": an = [-1] * lstp else: an = get_input_from_string(an) # make sure pix_cutoff is set for all iterations if len(pix_cutoff)<lstp: for i in xrange(len(pix_cutoff),lstp): pix_cutoff.append(pix_cutoff[-1]) # don't waste time on sub-pixel alignment for low-resolution ang incr for i in range(len(step)): if (delta[i] > 4 or delta[i] == -1) and step[i] < 1: step[i] = 1 first_ring = int(ir) rstep = int(rs) last_ring = int(ou) max_iter = int(maxit) center = int(center) nrefs = EMUtil.get_image_count( ref_vol ) nmasks = 0 if maskfile: # read number of masks within each maskfile (mc) nmasks = EMUtil.get_image_count( maskfile ) # open masks within maskfile (mc) maskF = EMData.read_images(maskfile, xrange(nmasks)) vol = EMData.read_images(ref_vol, xrange(nrefs)) nx = vol[0].get_xsize() ## make sure box sizes are the same if myid == main_node: im=EMData.read_images(stack,[0]) bx = im[0].get_xsize() if bx!=nx: print_msg("Error: Stack box size (%i) differs from initial model (%i)\n"%(bx,nx)) sys.exit() del im,bx # for helical processing: helicalrecon = False if protos is not None or hpars != "-1" or findseam is True: helicalrecon = True # if no out-of-plane param set, use 5 degrees if oplane is None: oplane=5.0 if protos is not None: proto = get_input_from_string(protos) if len(proto) != nrefs: print_msg("Error: insufficient protofilament numbers supplied") sys.exit() if hpars != "-1": hpars = get_input_from_string(hpars) if len(hpars) != 2*nrefs: print_msg("Error: insufficient helical parameters supplied") sys.exit() ## create helical parameter file for helical reconstruction if helicalrecon is True and myid == main_node: from hfunctions import createHpar # create initial helical parameter files dp=[0]*nrefs dphi=[0]*nrefs vdp=[0]*nrefs vdphi=[0]*nrefs for iref in xrange(nrefs): hpar = os.path.join(outdir,"hpar%02d.spi"%(iref)) params = False if hpars != "-1": # if helical parameters explicitly given, set twist & rise params = [float(hpars[iref*2]),float(hpars[(iref*2)+1])] dp[iref],dphi[iref],vdp[iref],vdphi[iref] = createHpar(hpar,proto[iref],params,vertstep) # get values for helical search parameters hsearch = get_input_from_string(hsearch) if len(hsearch) != 2: print_msg("Error: specify outer and inner radii for helical search") sys.exit() if last_ring < 0 or last_ring > int(nx/2)-2 : last_ring = int(nx/2) - 2 if myid == main_node: # import user_functions # user_func = user_functions.factory[user_func_name] print_begin_msg("ali3d_MPI") print_msg("Input stack : %s\n"%(stack)) print_msg("Reference volume : %s\n"%(ref_vol)) print_msg("Output directory : %s\n"%(outdir)) if nmasks > 0: print_msg("Maskfile (number of masks) : %s (%i)\n"%(maskfile,nmasks)) print_msg("Inner radius : %i\n"%(first_ring)) print_msg("Outer radius : %i\n"%(last_ring)) print_msg("Ring step : %i\n"%(rstep)) print_msg("X search range : %s\n"%(xrng)) print_msg("Y search range : %s\n"%(yrng)) print_msg("Translational step : %s\n"%(step)) print_msg("Angular step : %s\n"%(delta)) print_msg("Angular search range : %s\n"%(an)) print_msg("Maximum iteration : %i\n"%(max_iter)) print_msg("Center type : %i\n"%(center)) print_msg("CTF correction : %s\n"%(CTF)) print_msg("Signal-to-Noise Ratio : %f\n"%(snr)) print_msg("Reference projection method : %s\n"%(ref_a)) print_msg("Symmetry group : %s\n"%(sym)) print_msg("Fourier padding for 3D : %i\n"%(recon_pad)) print_msg("Number of reference models : %i\n"%(nrefs)) print_msg("Sort images between models : %s\n"%(sort)) print_msg("Allow images to jump : %s\n"%(mjump)) print_msg("CC cutoff standard dev : %f\n"%(cutoff)) print_msg("Two tail cutoff : %s\n"%(two_tail)) print_msg("Termination pix error : %f\n"%(term)) print_msg("Pixel error cutoff : %s\n"%(pix_cutoff)) print_msg("Restart : %s\n"%(restart)) print_msg("Full output : %s\n"%(full_output)) print_msg("Compare reprojections : %s\n"%(compare_repro)) print_msg("Compare ref free class avgs : %s\n"%(compare_ref_free)) print_msg("Use cutoff from ref free : %s\n"%(ref_free_cutoff)) if protos: print_msg("Protofilament numbers : %s\n"%(proto)) print_msg("Using helical search range : %s\n"%hsearch) if findseam is True: print_msg("Using seam-based reconstruction\n") if hpars != "-1": print_msg("Using hpars : %s\n"%hpars) if vertstep != None: print_msg("Using vertical step : %.2f\n"%vertstep) if save_half is True: print_msg("Saving even/odd halves\n") for i in xrange(100) : print_msg("*") print_msg("\n\n") if maskfile: if type(maskfile) is types.StringType: mask3D = get_image(maskfile) else: mask3D = maskfile else: mask3D = model_circle(last_ring, nx, nx, nx) numr = Numrinit(first_ring, last_ring, rstep, "F") mask2D = model_circle(last_ring,nx,nx) - model_circle(first_ring,nx,nx) fscmask = model_circle(last_ring,nx,nx,nx) if CTF: from filter import filt_ctf from reconstruction_rjh import rec3D_MPI_noCTF if myid == main_node: active = EMUtil.get_all_attributes(stack, 'active') list_of_particles = [] for im in xrange(len(active)): if active[im]: list_of_particles.append(im) del active nima = len(list_of_particles) else: nima = 0 total_nima = bcast_number_to_all(nima, source_node = main_node) if myid != main_node: list_of_particles = [-1]*total_nima list_of_particles = bcast_list_to_all(list_of_particles, source_node = main_node) image_start, image_end = MPI_start_end(total_nima, number_of_proc, myid) # create a list of images for each node list_of_particles = list_of_particles[image_start: image_end] nima = len(list_of_particles) if debug: finfo.write("image_start, image_end: %d %d\n" %(image_start, image_end)) finfo.flush() data = EMData.read_images(stack, list_of_particles) t_zero = Transform({"type":"spider","phi":0,"theta":0,"psi":0,"tx":0,"ty":0}) transmulti = [[t_zero for i in xrange(nrefs)] for j in xrange(nima)] for iref,im in ((iref,im) for iref in xrange(nrefs) for im in xrange(nima)): if nrefs == 1: transmulti[im][iref] = data[im].get_attr("xform.projection") else: # if multi models, keep track of eulers for all models try: transmulti[im][iref] = data[im].get_attr("eulers_txty.%i"%iref) except: data[im].set_attr("eulers_txty.%i"%iref,t_zero) scoremulti = [[0.0 for i in xrange(nrefs)] for j in xrange(nima)] pixelmulti = [[0.0 for i in xrange(nrefs)] for j in xrange(nima)] ref_res = [0.0 for x in xrange(nrefs)] apix = data[0].get_attr('apix_x') # for oplane parameter, create cylindrical mask if oplane is not None and myid == main_node: from hfunctions import createCylMask cmaskf=os.path.join(outdir, "mask3D_cyl.mrc") mask3D = createCylMask(data,ou,lmask,ilmask,cmaskf) # if finding seam of helix, create wedge masks if findseam is True: wedgemask=[] for pf in xrange(nrefs): wedgemask.append(EMData()) # wedgemask option if wcmask is not None: wcmask = get_input_from_string(wcmask) if len(wcmask) != 3: print_msg("Error: wcmask option requires 3 values: x y radius") sys.exit() # determine if particles have helix info: try: data[0].get_attr('h_angle') original_data = [] boxmask = True from hfunctions import createBoxMask except: boxmask = False # prepare particles for im in xrange(nima): data[im].set_attr('ID', list_of_particles[im]) data[im].set_attr('pix_score', int(0)) if CTF: # only phaseflip particles, not full CTF correction ctf_params = data[im].get_attr("ctf") st = Util.infomask(data[im], mask2D, False) data[im] -= st[0] data[im] = filt_ctf(data[im], ctf_params, sign = -1, binary=1) data[im].set_attr('ctf_applied', 1) # for window mask: if boxmask is True: h_angle = data[im].get_attr("h_angle") original_data.append(data[im].copy()) bmask = createBoxMask(nx,apix,ou,lmask,h_angle) data[im]*=bmask del bmask if debug: finfo.write( '%d loaded \n' % nima ) finfo.flush() if myid == main_node: # initialize data for the reference preparation function ref_data = [ mask3D, max(center,0), None, None, None, None ] # for method -1, switch off centering in user function from time import time # this is needed for gathering of pixel errors disps = [] recvcount = [] disps_score = [] recvcount_score = [] for im in xrange(number_of_proc): if( im == main_node ): disps.append(0) disps_score.append(0) else: disps.append(disps[im-1] + recvcount[im-1]) disps_score.append(disps_score[im-1] + recvcount_score[im-1]) ib, ie = MPI_start_end(total_nima, number_of_proc, im) recvcount.append( ie - ib ) recvcount_score.append((ie-ib)*nrefs) pixer = [0.0]*nima cs = [0.0]*3 total_iter = 0 volodd = EMData.read_images(ref_vol, xrange(nrefs)) voleve = EMData.read_images(ref_vol, xrange(nrefs)) if restart: # recreate initial volumes from alignments stored in header itout = "000_00" for iref in xrange(nrefs): if(nrefs == 1): modout = "" else: modout = "_model_%02d"%(iref) if(sort): group = iref for im in xrange(nima): imgroup = data[im].get_attr('group') if imgroup == iref: data[im].set_attr('xform.projection',transmulti[im][iref]) else: group = int(999) for im in xrange(nima): data[im].set_attr('xform.projection',transmulti[im][iref]) fscfile = os.path.join(outdir, "fsc_%s%s"%(itout,modout)) vol[iref], fscc, volodd[iref], voleve[iref] = rec3D_MPI_noCTF(data, sym, fscmask, fscfile, myid, main_node, index = group, npad = recon_pad) if myid == main_node: if helicalrecon: from hfunctions import processHelicalVol vstep=None if vertstep is not None: vstep=(vdp[iref],vdphi[iref]) print_msg("Old rise and twist for model %i : %8.3f, %8.3f\n"%(iref,dp[iref],dphi[iref])) hvals=processHelicalVol(vol[iref],voleve[iref],volodd[iref],iref,outdir,itout, dp[iref],dphi[iref],apix,hsearch,findseam,vstep,wcmask) (vol[iref],voleve[iref],volodd[iref],dp[iref],dphi[iref],vdp[iref],vdphi[iref])=hvals print_msg("New rise and twist for model %i : %8.3f, %8.3f\n"%(iref,dp[iref],dphi[iref])) # get new FSC from symmetrized half volumes fscc = fsc_mask( volodd[iref], voleve[iref], mask3D, rstep, fscfile) else: vol[iref].write_image(os.path.join(outdir, "vol_%s.hdf"%itout),-1) if save_half is True: volodd[iref].write_image(os.path.join(outdir, "volodd_%s.hdf"%itout),-1) voleve[iref].write_image(os.path.join(outdir, "voleve_%s.hdf"%itout),-1) if nmasks > 1: # Read mask for multiplying ref_data[0] = maskF[iref] ref_data[2] = vol[iref] ref_data[3] = fscc # call user-supplied function to prepare reference image, i.e., center and filter it vol[iref], cs,fl = ref_ali3d(ref_data) vol[iref].write_image(os.path.join(outdir, "volf_%s.hdf"%(itout)),-1) if (apix == 1): res_msg = "Models filtered at spatial frequency of:\t" res = fl else: res_msg = "Models filtered at resolution of: \t" res = apix / fl ares = array2string(array(res), precision = 2) print_msg("%s%s\n\n"%(res_msg,ares)) bcast_EMData_to_all(vol[iref], myid, main_node) # write out headers, under MPI writing has to be done sequentially mpi_barrier(MPI_COMM_WORLD) # projection matching for N_step in xrange(lstp): terminate = 0 Iter = -1 while(Iter < max_iter-1 and terminate == 0): Iter += 1 total_iter += 1 itout = "%03g_%02d" %(delta[N_step], Iter) if myid == main_node: print_msg("ITERATION #%3d, inner iteration #%3d\nDelta = %4.1f, an = %5.2f, xrange = %5.2f, yrange = %5.2f, step = %5.2f\n\n"%(N_step, Iter, delta[N_step], an[N_step], xrng[N_step],yrng[N_step],step[N_step])) for iref in xrange(nrefs): if myid == main_node: start_time = time() volft,kb = prep_vol( vol[iref] ) ## constrain projections to out of plane parameter theta1 = None theta2 = None if oplane is not None: theta1 = 90-oplane theta2 = 90+oplane refrings = prepare_refrings( volft, kb, nx, delta[N_step], ref_a, sym, numr, MPI=True, phiEqpsi = "Minus", initial_theta=theta1, delta_theta=theta2) del volft,kb if myid== main_node: print_msg( "Time to prepare projections for model %i: %s\n" % (iref, legibleTime(time()-start_time)) ) start_time = time() for im in xrange( nima ): data[im].set_attr("xform.projection", transmulti[im][iref]) if an[N_step] == -1: t1, peak, pixer[im] = proj_ali_incore(data[im],refrings,numr,xrng[N_step],yrng[N_step],step[N_step],finfo) else: t1, peak, pixer[im] = proj_ali_incore_local(data[im],refrings,numr,xrng[N_step],yrng[N_step],step[N_step],an[N_step],finfo) #data[im].set_attr("xform.projection"%iref, t1) if nrefs > 1: data[im].set_attr("eulers_txty.%i"%iref,t1) scoremulti[im][iref] = peak from pixel_error import max_3D_pixel_error # t1 is the current param, t2 is old t2 = transmulti[im][iref] pixelmulti[im][iref] = max_3D_pixel_error(t1,t2,numr[-3]) transmulti[im][iref] = t1 if myid == main_node: print_msg("Time of alignment for model %i: %s\n"%(iref, legibleTime(time()-start_time))) start_time = time() # gather scoring data from all processors from mpi import mpi_gatherv scoremultisend = sum(scoremulti,[]) pixelmultisend = sum(pixelmulti,[]) tmp = mpi_gatherv(scoremultisend,len(scoremultisend),MPI_FLOAT, recvcount_score, disps_score, MPI_FLOAT, main_node,MPI_COMM_WORLD) tmp1 = mpi_gatherv(pixelmultisend,len(pixelmultisend),MPI_FLOAT, recvcount_score, disps_score, MPI_FLOAT, main_node,MPI_COMM_WORLD) tmp = mpi_bcast(tmp,(total_nima * nrefs), MPI_FLOAT,0, MPI_COMM_WORLD) tmp1 = mpi_bcast(tmp1,(total_nima * nrefs), MPI_FLOAT,0, MPI_COMM_WORLD) tmp = map(float,tmp) tmp1 = map(float,tmp1) score = array(tmp).reshape(-1,nrefs) pixelerror = array(tmp1).reshape(-1,nrefs) score_local = array(scoremulti) mean_score = score.mean(axis=0) std_score = score.std(axis=0) cut = mean_score - (cutoff * std_score) cut2 = mean_score + (cutoff * std_score) res_max = score_local.argmax(axis=1) minus_cc = [0.0 for x in xrange(nrefs)] minus_pix = [0.0 for x in xrange(nrefs)] minus_ref = [0.0 for x in xrange(nrefs)] #output pixel errors if(myid == main_node): from statistics import hist_list lhist = 20 pixmin = pixelerror.min(axis=1) region, histo = hist_list(pixmin, lhist) if(region[0] < 0.0): region[0] = 0.0 print_msg("Histogram of pixel errors\n ERROR number of particles\n") for lhx in xrange(lhist): print_msg(" %10.3f %7d\n"%(region[lhx], histo[lhx])) # Terminate if 95% within 1 pixel error im = 0 for lhx in xrange(lhist): if(region[lhx] > 1.0): break im += histo[lhx] print_msg( "Percent of particles with pixel error < 1: %f\n\n"% (im/float(total_nima)*100)) term_cond = float(term)/100 if(im/float(total_nima) > term_cond): terminate = 1 print_msg("Terminating internal loop\n") del region, histo terminate = mpi_bcast(terminate, 1, MPI_INT, 0, MPI_COMM_WORLD) terminate = int(terminate[0]) for im in xrange(nima): if(sort==False): data[im].set_attr('group',999) elif (mjump[N_step]==1): data[im].set_attr('group',int(res_max[im])) pix_run = data[im].get_attr('pix_score') if (pix_cutoff[N_step]==1 and (terminate==1 or Iter == max_iter-1)): if (pixelmulti[im][int(res_max[im])] > 1): data[im].set_attr('pix_score',int(777)) if (score_local[im][int(res_max[im])]<cut[int(res_max[im])]) or (two_tail and score_local[im][int(res_max[im])]>cut2[int(res_max[im])]): data[im].set_attr('group',int(888)) minus_cc[int(res_max[im])] = minus_cc[int(res_max[im])] + 1 if(pix_run == 777): data[im].set_attr('group',int(777)) minus_pix[int(res_max[im])] = minus_pix[int(res_max[im])] + 1 if (compare_ref_free != "-1") and (ref_free_cutoff[N_step] != -1) and (total_iter > 1): id = data[im].get_attr('ID') if id in rejects: data[im].set_attr('group',int(666)) minus_ref[int(res_max[im])] = minus_ref[int(res_max[im])] + 1 minus_cc_tot = mpi_reduce(minus_cc,nrefs,MPI_FLOAT,MPI_SUM,0,MPI_COMM_WORLD) minus_pix_tot = mpi_reduce(minus_pix,nrefs,MPI_FLOAT,MPI_SUM,0,MPI_COMM_WORLD) minus_ref_tot = mpi_reduce(minus_ref,nrefs,MPI_FLOAT,MPI_SUM,0,MPI_COMM_WORLD) if (myid == main_node): if(sort): tot_max = score.argmax(axis=1) res = bincount(tot_max) else: res = ones(nrefs) * total_nima print_msg("Particle distribution: \t\t%s\n"%(res*1.0)) afcut1 = res - minus_cc_tot afcut2 = afcut1 - minus_pix_tot afcut3 = afcut2 - minus_ref_tot print_msg("Particle distribution after cc cutoff:\t\t%s\n"%(afcut1)) print_msg("Particle distribution after pix cutoff:\t\t%s\n"%(afcut2)) print_msg("Particle distribution after ref cutoff:\t\t%s\n\n"%(afcut3)) res = [0.0 for i in xrange(nrefs)] for iref in xrange(nrefs): if(center == -1): from utilities import estimate_3D_center_MPI, rotate_3D_shift dummy=EMData() cs[0], cs[1], cs[2], dummy, dummy = estimate_3D_center_MPI(data, total_nima, myid, number_of_proc, main_node) cs = mpi_bcast(cs, 3, MPI_FLOAT, main_node, MPI_COMM_WORLD) cs = [-float(cs[0]), -float(cs[1]), -float(cs[2])] rotate_3D_shift(data, cs) if(sort): group = iref for im in xrange(nima): imgroup = data[im].get_attr('group') if imgroup == iref: data[im].set_attr('xform.projection',transmulti[im][iref]) else: group = int(999) for im in xrange(nima): data[im].set_attr('xform.projection',transmulti[im][iref]) if(nrefs == 1): modout = "" else: modout = "_model_%02d"%(iref) fscfile = os.path.join(outdir, "fsc_%s%s"%(itout,modout)) vol[iref], fscc, volodd[iref], voleve[iref] = rec3D_MPI_noCTF(data, sym, fscmask, fscfile, myid, main_node, index=group, npad=recon_pad) if myid == main_node: print_msg("3D reconstruction time for model %i: %s\n"%(iref, legibleTime(time()-start_time))) start_time = time() # Compute Fourier variance if fourvar: outvar = os.path.join(outdir, "volVar_%s.hdf"%(itout)) ssnr_file = os.path.join(outdir, "ssnr_%s"%(itout)) varf = varf3d_MPI(data, ssnr_text_file=ssnr_file, mask2D=None, reference_structure=vol[iref], ou=last_ring, rw=1.0, npad=1, CTF=None, sign=1, sym=sym, myid=myid) if myid == main_node: print_msg("Time to calculate 3D Fourier variance for model %i: %s\n"%(iref, legibleTime(time()-start_time))) start_time = time() varf = 1.0/varf varf.write_image(outvar,-1) else: varf = None if myid == main_node: if helicalrecon: from hfunctions import processHelicalVol vstep=None if vertstep is not None: vstep=(vdp[iref],vdphi[iref]) print_msg("Old rise and twist for model %i : %8.3f, %8.3f\n"%(iref,dp[iref],dphi[iref])) hvals=processHelicalVol(vol[iref],voleve[iref],volodd[iref],iref,outdir,itout, dp[iref],dphi[iref],apix,hsearch,findseam,vstep,wcmask) (vol[iref],voleve[iref],volodd[iref],dp[iref],dphi[iref],vdp[iref],vdphi[iref])=hvals print_msg("New rise and twist for model %i : %8.3f, %8.3f\n"%(iref,dp[iref],dphi[iref])) # get new FSC from symmetrized half volumes fscc = fsc_mask( volodd[iref], voleve[iref], mask3D, rstep, fscfile) print_msg("Time to search and apply helical symmetry for model %i: %s\n\n"%(iref, legibleTime(time()-start_time))) start_time = time() else: vol[iref].write_image(os.path.join(outdir, "vol_%s.hdf"%(itout)),-1) if save_half is True: volodd[iref].write_image(os.path.join(outdir, "volodd_%s.hdf"%(itout)),-1) voleve[iref].write_image(os.path.join(outdir, "voleve_%s.hdf"%(itout)),-1) if nmasks > 1: # Read mask for multiplying ref_data[0] = maskF[iref] ref_data[2] = vol[iref] ref_data[3] = fscc ref_data[4] = varf # call user-supplied function to prepare reference image, i.e., center and filter it vol[iref], cs,fl = ref_ali3d(ref_data) vol[iref].write_image(os.path.join(outdir, "volf_%s.hdf"%(itout)),-1) if (apix == 1): res_msg = "Models filtered at spatial frequency of:\t" res[iref] = fl else: res_msg = "Models filtered at resolution of: \t" res[iref] = apix / fl del varf bcast_EMData_to_all(vol[iref], myid, main_node) if compare_ref_free != "-1": compare_repro = True if compare_repro: outfile_repro = comp_rep(refrings, data, itout, modout, vol[iref], group, nima, nx, myid, main_node, outdir) mpi_barrier(MPI_COMM_WORLD) if compare_ref_free != "-1": ref_free_output = os.path.join(outdir,"ref_free_%s%s"%(itout,modout)) rejects = compare(compare_ref_free, outfile_repro,ref_free_output,yrng[N_step], xrng[N_step], rstep,nx,apix,ref_free_cutoff[N_step], number_of_proc, myid, main_node) # retrieve alignment params from all processors par_str = ['xform.projection','ID','group'] if nrefs > 1: for iref in xrange(nrefs): par_str.append('eulers_txty.%i'%iref) if myid == main_node: from 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: ares = array2string(array(res), precision = 2) print_msg("%s%s\n\n"%(res_msg,ares)) dummy = EMData() if full_output: nimat = EMUtil.get_image_count(stack) output_file = os.path.join(outdir, "paramout_%s"%itout) foutput = open(output_file, 'w') for im in xrange(nimat): # save the parameters for each of the models outstring = "" dummy.read_image(stack,im,True) param3d = dummy.get_attr('xform.projection') g = dummy.get_attr("group") # retrieve alignments in EMAN-format pE = param3d.get_params('eman') outstring += "%f\t%f\t%f\t%f\t%f\t%i\n" %(pE["az"], pE["alt"], pE["phi"], pE["tx"], pE["ty"],g) foutput.write(outstring) foutput.close() del dummy mpi_barrier(MPI_COMM_WORLD) # mpi_finalize() if myid == main_node: print_end_msg("ali3d_MPI")
def resample( prjfile, outdir, bufprefix, nbufvol, nvol, seedbase,\ delta, d, snr, CTF, npad,\ MPI, myid, ncpu, verbose = 0 ): from utilities import even_angles from random import seed, jumpahead, shuffle import os from sys import exit nprj = EMUtil.get_image_count(prjfile) if MPI: from mpi import mpi_barrier, MPI_COMM_WORLD if myid == 0: if os.path.exists(outdir): nx = 1 else: nx = 0 else: nx = 0 ny = bcast_number_to_all(nx, source_node=0) if ny == 1: ERROR( 'Output directory exists, please change the name and restart the program', "resample", 1, myid) mpi_barrier(MPI_COMM_WORLD) if myid == 0: os.mkdir(outdir) mpi_barrier(MPI_COMM_WORLD) else: if os.path.exists(outdir): ERROR( 'Output directory exists, please change the name and restart the program', "resample", 1, 0) os.mkdir(outdir) if (verbose == 1): finfo = open(os.path.join(outdir, "progress%04d.txt" % myid), "w") else: finfo = None #print " before evenangles",myid from utilities import getvec from numpy import array, reshape refa = even_angles(delta) nrefa = len(refa) refnormal = zeros((nrefa, 3), 'float32') tetref = [0.0] * nrefa for i in range(nrefa): tr = getvec(refa[i][0], refa[i][1]) for j in range(3): refnormal[i][j] = tr[j] tetref[i] = refa[i][1] del refa vct = array([0.0] * (3 * nprj), 'float32') if myid == 0: print(" will read ", myid) tr = EMUtil.get_all_attributes(prjfile, 'xform.projection') tetprj = [0.0] * nprj for i in range(nprj): temp = tr[i].get_params("spider") tetprj[i] = temp["theta"] if (tetprj[i] > 90.0): tetprj[i] = 180.0 - tetprj[i] vct[3 * i + 0] = tr[i].at(2, 0) vct[3 * i + 1] = tr[i].at(2, 1) vct[3 * i + 2] = tr[i].at(2, 2) del tr else: tetprj = [0.0] * nprj #print " READ ",myid if MPI: #print " will bcast",myid from mpi import mpi_bcast, MPI_FLOAT, MPI_COMM_WORLD vct = mpi_bcast(vct, len(vct), MPI_FLOAT, 0, MPI_COMM_WORLD) from utilities import bcast_list_to_all tetprj = bcast_list_to_all(tetprj, myid, 0) #print " reshape ",myid vct = reshape(vct, (nprj, 3)) assignments = [[] for i in range(nrefa)] dspn = 1.25 * delta for k in range(nprj): best_s = -1.0 best_i = -1 for i in range(nrefa): if (abs(tetprj[k] - tetref[i]) <= dspn): s = abs(refnormal[i][0] * vct[k][0] + refnormal[i][1] * vct[k][1] + refnormal[i][2] * vct[k][2]) if s > best_s: best_s = s best_i = i assignments[best_i].append(k) am = len(assignments[0]) mufur = 1.0 / am for i in range(1, len(assignments)): ti = len(assignments[i]) am = min(am, ti) if (ti > 0): mufur += 1.0 / ti del tetprj, tetref dp = 1.0 - d # keep that many in each direction keep = int(am * dp + 0.5) mufur = keep * nrefa / (1.0 - mufur * keep / float(nrefa)) if myid == 0: print(" Number of projections ", nprj, ". Number of reference directions ", nrefa, ", multiplicative factor for the variance ", mufur) print(" Minimum number of assignments ", am, " Number of projections used per stratum ", keep, " Number of projections in resampled structure ", int(am * dp + 0.5) * nrefa) if am < 2 or am == keep: print("incorrect settings") exit() # FIX if (seedbase < 1): seed() jumpahead(17 * myid + 123) else: seed(seedbase) jumpahead(17 * myid + 123) volfile = os.path.join(outdir, "bsvol%04d.hdf" % myid) from random import randint niter = nvol / ncpu / nbufvol for kiter in range(niter): if (verbose == 1): finfo.write("Iteration %d: \n" % kiter) finfo.flush() iter_start = time() # the following has to be converted to resample mults=1 means take given projection., mults=0 means omit mults = [[0] * nprj for i in range(nbufvol)] for i in range(nbufvol): for l in range(nrefa): mass = assignments[l][:] shuffle(mass) mass = mass[:keep] mass.sort() #print l, " * ",mass for k in range(keep): mults[i][mass[k]] = 1 ''' lout = [] for l in xrange(len(mults[i])): if mults[i][l] == 1: lout.append(l) write_text_file(lout, os.path.join(outdir, "list%04d_%03d.txt" %(i, myid))) del lout ''' del mass rectors, fftvols, wgtvols = resample_prepare(prjfile, nbufvol, snr, CTF, npad) resample_insert(bufprefix, fftvols, wgtvols, mults, CTF, npad, finfo) del mults resample_finish(rectors, fftvols, wgtvols, volfile, kiter, nprj, finfo) rectors = None fftvols = None wgtvols = None if (verbose == 1): finfo.write("time for iteration: %10.3f\n" % (time() - iter_start)) finfo.flush()
def ali3d_MPI(stack, ref_vol, outdir, maskfile=None, ir=1, ou=-1, rs=1, xr="4 2 2 1", yr="-1", ts="1 1 0.5 0.25", delta="10 6 4 4", an="-1", center=0, maxit=5, term=95, CTF=False, fourvar=False, snr=1.0, ref_a="S", sym="c1", sort=True, cutoff=999.99, pix_cutoff="0", two_tail=False, model_jump="1 1 1 1 1", restart=False, save_half=False, protos=None, oplane=None, lmask=-1, ilmask=-1, findseam=False, vertstep=None, hpars="-1", hsearch="0.0 50.0", full_output=False, compare_repro=False, compare_ref_free="-1", ref_free_cutoff="-1 -1 -1 -1", wcmask=None, debug=False, recon_pad=4, olmask=75): from alignment import Numrinit, prepare_refrings from utilities import model_circle, get_image, drop_image, get_input_from_string from utilities import bcast_list_to_all, bcast_number_to_all, reduce_EMData_to_root, bcast_EMData_to_all from utilities import send_attr_dict from utilities import get_params_proj, file_type from fundamentals import rot_avg_image import os import types from utilities import print_begin_msg, print_end_msg, print_msg from mpi import mpi_bcast, mpi_comm_size, mpi_comm_rank, MPI_FLOAT, MPI_COMM_WORLD, mpi_barrier, mpi_reduce from mpi import mpi_reduce, MPI_INT, MPI_SUM, mpi_finalize from filter import filt_ctf from projection import prep_vol, prgs from statistics import hist_list, varf3d_MPI, fsc_mask from numpy import array, bincount, array2string, ones number_of_proc = mpi_comm_size(MPI_COMM_WORLD) myid = mpi_comm_rank(MPI_COMM_WORLD) main_node = 0 if myid == main_node: if os.path.exists(outdir): ERROR( 'Output directory exists, please change the name and restart the program', "ali3d_MPI", 1) os.mkdir(outdir) mpi_barrier(MPI_COMM_WORLD) if debug: from time import sleep while not os.path.exists(outdir): print "Node ", myid, " waiting..." sleep(5) info_file = os.path.join(outdir, "progress%04d" % myid) finfo = open(info_file, 'w') else: finfo = None mjump = get_input_from_string(model_jump) xrng = get_input_from_string(xr) if yr == "-1": yrng = xrng else: yrng = get_input_from_string(yr) step = get_input_from_string(ts) delta = get_input_from_string(delta) ref_free_cutoff = get_input_from_string(ref_free_cutoff) pix_cutoff = get_input_from_string(pix_cutoff) lstp = min(len(xrng), len(yrng), len(step), len(delta)) if an == "-1": an = [-1] * lstp else: an = get_input_from_string(an) # make sure pix_cutoff is set for all iterations if len(pix_cutoff) < lstp: for i in xrange(len(pix_cutoff), lstp): pix_cutoff.append(pix_cutoff[-1]) # don't waste time on sub-pixel alignment for low-resolution ang incr for i in range(len(step)): if (delta[i] > 4 or delta[i] == -1) and step[i] < 1: step[i] = 1 first_ring = int(ir) rstep = int(rs) last_ring = int(ou) max_iter = int(maxit) center = int(center) nrefs = EMUtil.get_image_count(ref_vol) nmasks = 0 if maskfile: # read number of masks within each maskfile (mc) nmasks = EMUtil.get_image_count(maskfile) # open masks within maskfile (mc) maskF = EMData.read_images(maskfile, xrange(nmasks)) vol = EMData.read_images(ref_vol, xrange(nrefs)) nx = vol[0].get_xsize() ## make sure box sizes are the same if myid == main_node: im = EMData.read_images(stack, [0]) bx = im[0].get_xsize() if bx != nx: print_msg( "Error: Stack box size (%i) differs from initial model (%i)\n" % (bx, nx)) sys.exit() del im, bx # for helical processing: helicalrecon = False if protos is not None or hpars != "-1" or findseam is True: helicalrecon = True # if no out-of-plane param set, use 5 degrees if oplane is None: oplane = 5.0 if protos is not None: proto = get_input_from_string(protos) if len(proto) != nrefs: print_msg("Error: insufficient protofilament numbers supplied") sys.exit() if hpars != "-1": hpars = get_input_from_string(hpars) if len(hpars) != 2 * nrefs: print_msg("Error: insufficient helical parameters supplied") sys.exit() ## create helical parameter file for helical reconstruction if helicalrecon is True and myid == main_node: from hfunctions import createHpar # create initial helical parameter files dp = [0] * nrefs dphi = [0] * nrefs vdp = [0] * nrefs vdphi = [0] * nrefs for iref in xrange(nrefs): hpar = os.path.join(outdir, "hpar%02d.spi" % (iref)) params = False if hpars != "-1": # if helical parameters explicitly given, set twist & rise params = [float(hpars[iref * 2]), float(hpars[(iref * 2) + 1])] dp[iref], dphi[iref], vdp[iref], vdphi[iref] = createHpar( hpar, proto[iref], params, vertstep) # get values for helical search parameters hsearch = get_input_from_string(hsearch) if len(hsearch) != 2: print_msg("Error: specify outer and inner radii for helical search") sys.exit() if last_ring < 0 or last_ring > int(nx / 2) - 2: last_ring = int(nx / 2) - 2 if myid == main_node: # import user_functions # user_func = user_functions.factory[user_func_name] print_begin_msg("ali3d_MPI") print_msg("Input stack : %s\n" % (stack)) print_msg("Reference volume : %s\n" % (ref_vol)) print_msg("Output directory : %s\n" % (outdir)) if nmasks > 0: print_msg("Maskfile (number of masks) : %s (%i)\n" % (maskfile, nmasks)) print_msg("Inner radius : %i\n" % (first_ring)) print_msg("Outer radius : %i\n" % (last_ring)) print_msg("Ring step : %i\n" % (rstep)) print_msg("X search range : %s\n" % (xrng)) print_msg("Y search range : %s\n" % (yrng)) print_msg("Translational step : %s\n" % (step)) print_msg("Angular step : %s\n" % (delta)) print_msg("Angular search range : %s\n" % (an)) print_msg("Maximum iteration : %i\n" % (max_iter)) print_msg("Center type : %i\n" % (center)) print_msg("CTF correction : %s\n" % (CTF)) print_msg("Signal-to-Noise Ratio : %f\n" % (snr)) print_msg("Reference projection method : %s\n" % (ref_a)) print_msg("Symmetry group : %s\n" % (sym)) print_msg("Fourier padding for 3D : %i\n" % (recon_pad)) print_msg("Number of reference models : %i\n" % (nrefs)) print_msg("Sort images between models : %s\n" % (sort)) print_msg("Allow images to jump : %s\n" % (mjump)) print_msg("CC cutoff standard dev : %f\n" % (cutoff)) print_msg("Two tail cutoff : %s\n" % (two_tail)) print_msg("Termination pix error : %f\n" % (term)) print_msg("Pixel error cutoff : %s\n" % (pix_cutoff)) print_msg("Restart : %s\n" % (restart)) print_msg("Full output : %s\n" % (full_output)) print_msg("Compare reprojections : %s\n" % (compare_repro)) print_msg("Compare ref free class avgs : %s\n" % (compare_ref_free)) print_msg("Use cutoff from ref free : %s\n" % (ref_free_cutoff)) if protos: print_msg("Protofilament numbers : %s\n" % (proto)) print_msg("Using helical search range : %s\n" % hsearch) if findseam is True: print_msg("Using seam-based reconstruction\n") if hpars != "-1": print_msg("Using hpars : %s\n" % hpars) if vertstep != None: print_msg("Using vertical step : %.2f\n" % vertstep) if save_half is True: print_msg("Saving even/odd halves\n") for i in xrange(100): print_msg("*") print_msg("\n\n") if maskfile: if type(maskfile) is types.StringType: mask3D = get_image(maskfile) else: mask3D = maskfile else: mask3D = model_circle(last_ring, nx, nx, nx) numr = Numrinit(first_ring, last_ring, rstep, "F") mask2D = model_circle(last_ring, nx, nx) - model_circle(first_ring, nx, nx) fscmask = model_circle(last_ring, nx, nx, nx) if CTF: from filter import filt_ctf from reconstruction_rjh import rec3D_MPI_noCTF if myid == main_node: active = EMUtil.get_all_attributes(stack, 'active') list_of_particles = [] for im in xrange(len(active)): if active[im]: list_of_particles.append(im) del active nima = len(list_of_particles) else: nima = 0 total_nima = bcast_number_to_all(nima, source_node=main_node) if myid != main_node: list_of_particles = [-1] * total_nima list_of_particles = bcast_list_to_all(list_of_particles, source_node=main_node) image_start, image_end = MPI_start_end(total_nima, number_of_proc, myid) # create a list of images for each node list_of_particles = list_of_particles[image_start:image_end] nima = len(list_of_particles) if debug: finfo.write("image_start, image_end: %d %d\n" % (image_start, image_end)) finfo.flush() data = EMData.read_images(stack, list_of_particles) t_zero = Transform({ "type": "spider", "phi": 0, "theta": 0, "psi": 0, "tx": 0, "ty": 0 }) transmulti = [[t_zero for i in xrange(nrefs)] for j in xrange(nima)] for iref, im in ((iref, im) for iref in xrange(nrefs) for im in xrange(nima)): if nrefs == 1: transmulti[im][iref] = data[im].get_attr("xform.projection") else: # if multi models, keep track of eulers for all models try: transmulti[im][iref] = data[im].get_attr("eulers_txty.%i" % iref) except: data[im].set_attr("eulers_txty.%i" % iref, t_zero) scoremulti = [[0.0 for i in xrange(nrefs)] for j in xrange(nima)] pixelmulti = [[0.0 for i in xrange(nrefs)] for j in xrange(nima)] ref_res = [0.0 for x in xrange(nrefs)] apix = data[0].get_attr('apix_x') # for oplane parameter, create cylindrical mask if oplane is not None and myid == main_node: from hfunctions import createCylMask cmaskf = os.path.join(outdir, "mask3D_cyl.mrc") mask3D = createCylMask(data, olmask, lmask, ilmask, cmaskf) # if finding seam of helix, create wedge masks if findseam is True: wedgemask = [] for pf in xrange(nrefs): wedgemask.append(EMData()) # wedgemask option if wcmask is not None: wcmask = get_input_from_string(wcmask) if len(wcmask) != 3: print_msg( "Error: wcmask option requires 3 values: x y radius") sys.exit() # determine if particles have helix info: try: data[0].get_attr('h_angle') original_data = [] boxmask = True from hfunctions import createBoxMask except: boxmask = False # prepare particles for im in xrange(nima): data[im].set_attr('ID', list_of_particles[im]) data[im].set_attr('pix_score', int(0)) if CTF: # only phaseflip particles, not full CTF correction ctf_params = data[im].get_attr("ctf") st = Util.infomask(data[im], mask2D, False) data[im] -= st[0] data[im] = filt_ctf(data[im], ctf_params, sign=-1, binary=1) data[im].set_attr('ctf_applied', 1) # for window mask: if boxmask is True: h_angle = data[im].get_attr("h_angle") original_data.append(data[im].copy()) bmask = createBoxMask(nx, apix, ou, lmask, h_angle) data[im] *= bmask del bmask if debug: finfo.write('%d loaded \n' % nima) finfo.flush() if myid == main_node: # initialize data for the reference preparation function ref_data = [mask3D, max(center, 0), None, None, None, None] # for method -1, switch off centering in user function from time import time # this is needed for gathering of pixel errors disps = [] recvcount = [] disps_score = [] recvcount_score = [] for im in xrange(number_of_proc): if (im == main_node): disps.append(0) disps_score.append(0) else: disps.append(disps[im - 1] + recvcount[im - 1]) disps_score.append(disps_score[im - 1] + recvcount_score[im - 1]) ib, ie = MPI_start_end(total_nima, number_of_proc, im) recvcount.append(ie - ib) recvcount_score.append((ie - ib) * nrefs) pixer = [0.0] * nima cs = [0.0] * 3 total_iter = 0 volodd = EMData.read_images(ref_vol, xrange(nrefs)) voleve = EMData.read_images(ref_vol, xrange(nrefs)) if restart: # recreate initial volumes from alignments stored in header itout = "000_00" for iref in xrange(nrefs): if (nrefs == 1): modout = "" else: modout = "_model_%02d" % (iref) if (sort): group = iref for im in xrange(nima): imgroup = data[im].get_attr('group') if imgroup == iref: data[im].set_attr('xform.projection', transmulti[im][iref]) else: group = int(999) for im in xrange(nima): data[im].set_attr('xform.projection', transmulti[im][iref]) fscfile = os.path.join(outdir, "fsc_%s%s" % (itout, modout)) vol[iref], fscc, volodd[iref], voleve[iref] = rec3D_MPI_noCTF( data, sym, fscmask, fscfile, myid, main_node, index=group, npad=recon_pad) if myid == main_node: if helicalrecon: from hfunctions import processHelicalVol vstep = None if vertstep is not None: vstep = (vdp[iref], vdphi[iref]) print_msg( "Old rise and twist for model %i : %8.3f, %8.3f\n" % (iref, dp[iref], dphi[iref])) hvals = processHelicalVol(vol[iref], voleve[iref], volodd[iref], iref, outdir, itout, dp[iref], dphi[iref], apix, hsearch, findseam, vstep, wcmask) (vol[iref], voleve[iref], volodd[iref], dp[iref], dphi[iref], vdp[iref], vdphi[iref]) = hvals print_msg( "New rise and twist for model %i : %8.3f, %8.3f\n" % (iref, dp[iref], dphi[iref])) # get new FSC from symmetrized half volumes fscc = fsc_mask(volodd[iref], voleve[iref], mask3D, rstep, fscfile) else: vol[iref].write_image( os.path.join(outdir, "vol_%s.hdf" % itout), -1) if save_half is True: volodd[iref].write_image( os.path.join(outdir, "volodd_%s.hdf" % itout), -1) voleve[iref].write_image( os.path.join(outdir, "voleve_%s.hdf" % itout), -1) if nmasks > 1: # Read mask for multiplying ref_data[0] = maskF[iref] ref_data[2] = vol[iref] ref_data[3] = fscc # call user-supplied function to prepare reference image, i.e., center and filter it vol[iref], cs, fl = ref_ali3d(ref_data) vol[iref].write_image( os.path.join(outdir, "volf_%s.hdf" % (itout)), -1) if (apix == 1): res_msg = "Models filtered at spatial frequency of:\t" res = fl else: res_msg = "Models filtered at resolution of: \t" res = apix / fl ares = array2string(array(res), precision=2) print_msg("%s%s\n\n" % (res_msg, ares)) bcast_EMData_to_all(vol[iref], myid, main_node) # write out headers, under MPI writing has to be done sequentially mpi_barrier(MPI_COMM_WORLD) # projection matching for N_step in xrange(lstp): terminate = 0 Iter = -1 while (Iter < max_iter - 1 and terminate == 0): Iter += 1 total_iter += 1 itout = "%03g_%02d" % (delta[N_step], Iter) if myid == main_node: print_msg( "ITERATION #%3d, inner iteration #%3d\nDelta = %4.1f, an = %5.2f, xrange = %5.2f, yrange = %5.2f, step = %5.2f\n\n" % (N_step, Iter, delta[N_step], an[N_step], xrng[N_step], yrng[N_step], step[N_step])) for iref in xrange(nrefs): if myid == main_node: start_time = time() volft, kb = prep_vol(vol[iref]) ## constrain projections to out of plane parameter theta1 = None theta2 = None if oplane is not None: theta1 = 90 - oplane theta2 = 90 + oplane refrings = prepare_refrings(volft, kb, nx, delta[N_step], ref_a, sym, numr, MPI=True, phiEqpsi="Minus", initial_theta=theta1, delta_theta=theta2) del volft, kb if myid == main_node: print_msg( "Time to prepare projections for model %i: %s\n" % (iref, legibleTime(time() - start_time))) start_time = time() for im in xrange(nima): data[im].set_attr("xform.projection", transmulti[im][iref]) if an[N_step] == -1: t1, peak, pixer[im] = proj_ali_incore( data[im], refrings, numr, xrng[N_step], yrng[N_step], step[N_step], finfo) else: t1, peak, pixer[im] = proj_ali_incore_local( data[im], refrings, numr, xrng[N_step], yrng[N_step], step[N_step], an[N_step], finfo) #data[im].set_attr("xform.projection"%iref, t1) if nrefs > 1: data[im].set_attr("eulers_txty.%i" % iref, t1) scoremulti[im][iref] = peak from pixel_error import max_3D_pixel_error # t1 is the current param, t2 is old t2 = transmulti[im][iref] pixelmulti[im][iref] = max_3D_pixel_error(t1, t2, numr[-3]) transmulti[im][iref] = t1 if myid == main_node: print_msg("Time of alignment for model %i: %s\n" % (iref, legibleTime(time() - start_time))) start_time = time() # gather scoring data from all processors from mpi import mpi_gatherv scoremultisend = sum(scoremulti, []) pixelmultisend = sum(pixelmulti, []) tmp = mpi_gatherv(scoremultisend, len(scoremultisend), MPI_FLOAT, recvcount_score, disps_score, MPI_FLOAT, main_node, MPI_COMM_WORLD) tmp1 = mpi_gatherv(pixelmultisend, len(pixelmultisend), MPI_FLOAT, recvcount_score, disps_score, MPI_FLOAT, main_node, MPI_COMM_WORLD) tmp = mpi_bcast(tmp, (total_nima * nrefs), MPI_FLOAT, 0, MPI_COMM_WORLD) tmp1 = mpi_bcast(tmp1, (total_nima * nrefs), MPI_FLOAT, 0, MPI_COMM_WORLD) tmp = map(float, tmp) tmp1 = map(float, tmp1) score = array(tmp).reshape(-1, nrefs) pixelerror = array(tmp1).reshape(-1, nrefs) score_local = array(scoremulti) mean_score = score.mean(axis=0) std_score = score.std(axis=0) cut = mean_score - (cutoff * std_score) cut2 = mean_score + (cutoff * std_score) res_max = score_local.argmax(axis=1) minus_cc = [0.0 for x in xrange(nrefs)] minus_pix = [0.0 for x in xrange(nrefs)] minus_ref = [0.0 for x in xrange(nrefs)] #output pixel errors if (myid == main_node): from statistics import hist_list lhist = 20 pixmin = pixelerror.min(axis=1) region, histo = hist_list(pixmin, lhist) if (region[0] < 0.0): region[0] = 0.0 print_msg( "Histogram of pixel errors\n ERROR number of particles\n" ) for lhx in xrange(lhist): print_msg(" %10.3f %7d\n" % (region[lhx], histo[lhx])) # Terminate if 95% within 1 pixel error im = 0 for lhx in xrange(lhist): if (region[lhx] > 1.0): break im += histo[lhx] print_msg("Percent of particles with pixel error < 1: %f\n\n" % (im / float(total_nima) * 100)) term_cond = float(term) / 100 if (im / float(total_nima) > term_cond): terminate = 1 print_msg("Terminating internal loop\n") del region, histo terminate = mpi_bcast(terminate, 1, MPI_INT, 0, MPI_COMM_WORLD) terminate = int(terminate[0]) for im in xrange(nima): if (sort == False): data[im].set_attr('group', 999) elif (mjump[N_step] == 1): data[im].set_attr('group', int(res_max[im])) pix_run = data[im].get_attr('pix_score') if (pix_cutoff[N_step] == 1 and (terminate == 1 or Iter == max_iter - 1)): if (pixelmulti[im][int(res_max[im])] > 1): data[im].set_attr('pix_score', int(777)) if (score_local[im][int(res_max[im])] < cut[int( res_max[im])]) or (two_tail and score_local[im][int( res_max[im])] > cut2[int(res_max[im])]): data[im].set_attr('group', int(888)) minus_cc[int(res_max[im])] = minus_cc[int(res_max[im])] + 1 if (pix_run == 777): data[im].set_attr('group', int(777)) minus_pix[int( res_max[im])] = minus_pix[int(res_max[im])] + 1 if (compare_ref_free != "-1") and (ref_free_cutoff[N_step] != -1) and (total_iter > 1): id = data[im].get_attr('ID') if id in rejects: data[im].set_attr('group', int(666)) minus_ref[int( res_max[im])] = minus_ref[int(res_max[im])] + 1 minus_cc_tot = mpi_reduce(minus_cc, nrefs, MPI_FLOAT, MPI_SUM, 0, MPI_COMM_WORLD) minus_pix_tot = mpi_reduce(minus_pix, nrefs, MPI_FLOAT, MPI_SUM, 0, MPI_COMM_WORLD) minus_ref_tot = mpi_reduce(minus_ref, nrefs, MPI_FLOAT, MPI_SUM, 0, MPI_COMM_WORLD) if (myid == main_node): if (sort): tot_max = score.argmax(axis=1) res = bincount(tot_max) else: res = ones(nrefs) * total_nima print_msg("Particle distribution: \t\t%s\n" % (res * 1.0)) afcut1 = res - minus_cc_tot afcut2 = afcut1 - minus_pix_tot afcut3 = afcut2 - minus_ref_tot print_msg("Particle distribution after cc cutoff:\t\t%s\n" % (afcut1)) print_msg("Particle distribution after pix cutoff:\t\t%s\n" % (afcut2)) print_msg("Particle distribution after ref cutoff:\t\t%s\n\n" % (afcut3)) res = [0.0 for i in xrange(nrefs)] for iref in xrange(nrefs): if (center == -1): from utilities import estimate_3D_center_MPI, rotate_3D_shift dummy = EMData() cs[0], cs[1], cs[2], dummy, dummy = estimate_3D_center_MPI( data, total_nima, myid, number_of_proc, main_node) cs = mpi_bcast(cs, 3, MPI_FLOAT, main_node, MPI_COMM_WORLD) cs = [-float(cs[0]), -float(cs[1]), -float(cs[2])] rotate_3D_shift(data, cs) if (sort): group = iref for im in xrange(nima): imgroup = data[im].get_attr('group') if imgroup == iref: data[im].set_attr('xform.projection', transmulti[im][iref]) else: group = int(999) for im in xrange(nima): data[im].set_attr('xform.projection', transmulti[im][iref]) if (nrefs == 1): modout = "" else: modout = "_model_%02d" % (iref) fscfile = os.path.join(outdir, "fsc_%s%s" % (itout, modout)) vol[iref], fscc, volodd[iref], voleve[iref] = rec3D_MPI_noCTF( data, sym, fscmask, fscfile, myid, main_node, index=group, npad=recon_pad) if myid == main_node: print_msg("3D reconstruction time for model %i: %s\n" % (iref, legibleTime(time() - start_time))) start_time = time() # Compute Fourier variance if fourvar: outvar = os.path.join(outdir, "volVar_%s.hdf" % (itout)) ssnr_file = os.path.join(outdir, "ssnr_%s" % (itout)) varf = varf3d_MPI(data, ssnr_text_file=ssnr_file, mask2D=None, reference_structure=vol[iref], ou=last_ring, rw=1.0, npad=1, CTF=None, sign=1, sym=sym, myid=myid) if myid == main_node: print_msg( "Time to calculate 3D Fourier variance for model %i: %s\n" % (iref, legibleTime(time() - start_time))) start_time = time() varf = 1.0 / varf varf.write_image(outvar, -1) else: varf = None if myid == main_node: if helicalrecon: from hfunctions import processHelicalVol vstep = None if vertstep is not None: vstep = (vdp[iref], vdphi[iref]) print_msg( "Old rise and twist for model %i : %8.3f, %8.3f\n" % (iref, dp[iref], dphi[iref])) hvals = processHelicalVol(vol[iref], voleve[iref], volodd[iref], iref, outdir, itout, dp[iref], dphi[iref], apix, hsearch, findseam, vstep, wcmask) (vol[iref], voleve[iref], volodd[iref], dp[iref], dphi[iref], vdp[iref], vdphi[iref]) = hvals print_msg( "New rise and twist for model %i : %8.3f, %8.3f\n" % (iref, dp[iref], dphi[iref])) # get new FSC from symmetrized half volumes fscc = fsc_mask(volodd[iref], voleve[iref], mask3D, rstep, fscfile) print_msg( "Time to search and apply helical symmetry for model %i: %s\n\n" % (iref, legibleTime(time() - start_time))) start_time = time() else: vol[iref].write_image( os.path.join(outdir, "vol_%s.hdf" % (itout)), -1) if save_half is True: volodd[iref].write_image( os.path.join(outdir, "volodd_%s.hdf" % (itout)), -1) voleve[iref].write_image( os.path.join(outdir, "voleve_%s.hdf" % (itout)), -1) if nmasks > 1: # Read mask for multiplying ref_data[0] = maskF[iref] ref_data[2] = vol[iref] ref_data[3] = fscc ref_data[4] = varf # call user-supplied function to prepare reference image, i.e., center and filter it vol[iref], cs, fl = ref_ali3d(ref_data) vol[iref].write_image( os.path.join(outdir, "volf_%s.hdf" % (itout)), -1) if (apix == 1): res_msg = "Models filtered at spatial frequency of:\t" res[iref] = fl else: res_msg = "Models filtered at resolution of: \t" res[iref] = apix / fl del varf bcast_EMData_to_all(vol[iref], myid, main_node) if compare_ref_free != "-1": compare_repro = True if compare_repro: outfile_repro = comp_rep(refrings, data, itout, modout, vol[iref], group, nima, nx, myid, main_node, outdir) mpi_barrier(MPI_COMM_WORLD) if compare_ref_free != "-1": ref_free_output = os.path.join( outdir, "ref_free_%s%s" % (itout, modout)) rejects = compare(compare_ref_free, outfile_repro, ref_free_output, yrng[N_step], xrng[N_step], rstep, nx, apix, ref_free_cutoff[N_step], number_of_proc, myid, main_node) # retrieve alignment params from all processors par_str = ['xform.projection', 'ID', 'group'] if nrefs > 1: for iref in xrange(nrefs): par_str.append('eulers_txty.%i' % iref) if myid == main_node: from 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: ares = array2string(array(res), precision=2) print_msg("%s%s\n\n" % (res_msg, ares)) dummy = EMData() if full_output: nimat = EMUtil.get_image_count(stack) output_file = os.path.join(outdir, "paramout_%s" % itout) foutput = open(output_file, 'w') for im in xrange(nimat): # save the parameters for each of the models outstring = "" dummy.read_image(stack, im, True) param3d = dummy.get_attr('xform.projection') g = dummy.get_attr("group") # retrieve alignments in EMAN-format pE = param3d.get_params('eman') outstring += "%f\t%f\t%f\t%f\t%f\t%i\n" % ( pE["az"], pE["alt"], pE["phi"], pE["tx"], pE["ty"], g) foutput.write(outstring) foutput.close() del dummy mpi_barrier(MPI_COMM_WORLD) # mpi_finalize() if myid == main_node: print_end_msg("ali3d_MPI")
def helicalshiftali_MPI(stack, maskfile=None, maxit=100, CTF=False, snr=1.0, Fourvar=False, search_rng=-1): from applications import MPI_start_end from utilities import model_circle, model_blank, get_image, peak_search, get_im, pad from utilities import reduce_EMData_to_root, bcast_EMData_to_all, send_attr_dict, file_type, bcast_number_to_all, bcast_list_to_all from statistics import varf2d_MPI from fundamentals import fft, ccf, rot_shift3D, rot_shift2D, fshift from utilities import get_params2D, set_params2D, chunks_distribution from utilities import print_msg, print_begin_msg, print_end_msg import os import sys from mpi import mpi_init, mpi_comm_size, mpi_comm_rank, MPI_COMM_WORLD from mpi import mpi_reduce, mpi_bcast, mpi_barrier, mpi_gatherv from mpi import MPI_SUM, MPI_FLOAT, MPI_INT from time import time from pixel_error import ordersegments from math import sqrt, atan2, tan, pi nproc = mpi_comm_size(MPI_COMM_WORLD) myid = mpi_comm_rank(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 xrange(total_nfils): inidl[i] = len(filaments[i]) linidl = sum(inidl) nima = linidl tfilaments = [] for i in xrange(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 xrange(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), "ehelix_MPI", 1,myid) # balanced load temp = chunks_distribution([[len(filaments[i]), i] for i in xrange(len(filaments))], nproc)[myid:myid+1][0] filaments = [filaments[temp[i][1]] for i in xrange(len(temp))] nfils = len(filaments) #filaments = [[0,1]] #print "filaments",filaments list_of_particles = [] indcs = [] k = 0 for i in xrange(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) print "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 xrange(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 filter import filt_ctf from 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 utilities import info for im in xrange(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', "helicalshiftali_MPI", 1,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 xrange(0,nx+2,2): for j in xrange(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 xrange(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 xrange(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 xrange(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 xrange(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 xrange(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_reduce(sx_sum, 1, MPI_FLOAT, MPI_SUM, main_node, 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 xrange(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 xrange(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_barrier(MPI_COMM_WORLD) par_str = ["xform.align2d", "ID"] if myid == main_node: from utilities import file_type if(file_type(stack) == "bdb"): from utilities import recv_attr_dict_bdb recv_attr_dict_bdb(main_node, stack, data, par_str, 0, ldata, nproc) else: from 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 filterlocal(ui, vi, m, falloff, myid, main_node, number_of_proc): from mpi import mpi_init, mpi_comm_size, mpi_comm_rank, MPI_COMM_WORLD from mpi import mpi_reduce, mpi_bcast, mpi_barrier, mpi_gatherv, mpi_send, mpi_recv from mpi import MPI_SUM, MPI_FLOAT, MPI_INT from utilities import bcast_number_to_all, bcast_list_to_all, model_blank, bcast_EMData_to_all, reduce_EMData_to_root from morphology import threshold_outside from filter import filt_tanl from fundamentals import fft, fftip if(myid == main_node): nx = vi.get_xsize() ny = vi.get_ysize() nz = vi.get_zsize() # Round all resolution numbers to two digits for x in xrange(nx): for y in xrange(ny): for z in xrange(nz): ui.set_value_at_fast( x,y,z, round(ui.get_value_at(x,y,z), 2) ) dis = [nx,ny,nz] else: falloff = 0.0 radius = 0 dis = [0,0,0] falloff = bcast_number_to_all(falloff, main_node) dis = bcast_list_to_all(dis, myid, source_node = main_node) if(myid != main_node): nx = int(dis[0]) ny = int(dis[1]) nz = int(dis[2]) vi = model_blank(nx,ny,nz) ui = model_blank(nx,ny,nz) bcast_EMData_to_all(vi, myid, main_node) bcast_EMData_to_all(ui, myid, main_node) fftip(vi) # volume to be filtered st = Util.infomask(ui, m, True) filteredvol = model_blank(nx,ny,nz) cutoff = max(st[2] - 0.01,0.0) while(cutoff < st[3] ): cutoff = round(cutoff + 0.01, 2) #if(myid == main_node): print cutoff,st pt = Util.infomask( threshold_outside(ui, cutoff - 0.00501, cutoff + 0.005), m, True) # Ideally, one would want to check only slices in question... if(pt[0] != 0.0): #print cutoff,pt[0] vovo = fft( filt_tanl(vi, cutoff, falloff) ) for z in xrange(myid, nz, number_of_proc): for x in xrange(nx): for y in xrange(ny): if(m.get_value_at(x,y,z) > 0.5): if(round(ui.get_value_at(x,y,z),2) == cutoff): filteredvol.set_value_at_fast(x,y,z,vovo.get_value_at(x,y,z)) mpi_barrier(MPI_COMM_WORLD) reduce_EMData_to_root(filteredvol, myid, main_node, MPI_COMM_WORLD) return filteredvol
def main(): def params_3D_2D_NEW(phi, theta, psi, s2x, s2y, mirror): if mirror: m = 1 alpha, sx, sy, scalen = compose_transform2(0, s2x, s2y, 1.0, 540.0-psi, 0, 0, 1.0) else: m = 0 alpha, sx, sy, scalen = compose_transform2(0, s2x, s2y, 1.0, 360.0-psi, 0, 0, 1.0) return alpha, sx, sy, m progname = os.path.basename(sys.argv[0]) usage = progname + " prj_stack --ave2D= --var2D= --ave3D= --var3D= --img_per_grp= --fl=0.2 --aa=0.1 --sym=symmetry --CTF" parser = OptionParser(usage, version=SPARXVERSION) parser.add_option("--ave2D", type="string" , default=False, help="write to the disk a stack of 2D averages") parser.add_option("--var2D", type="string" , default=False, help="write to the disk a stack of 2D variances") parser.add_option("--ave3D", type="string" , default=False, help="write to the disk reconstructed 3D average") parser.add_option("--var3D", type="string" , default=False, help="compute 3D variability (time consuming!)") parser.add_option("--img_per_grp", type="int" , default=10 , help="number of neighbouring projections") parser.add_option("--no_norm", action="store_true", default=False, help="do not use normalization") parser.add_option("--radiusvar", type="int" , default=-1 , help="radius for 3D var" ) parser.add_option("--npad", type="int" , default=2 , help="number of time to pad the original images") parser.add_option("--sym" , type="string" , default="c1" , help="symmetry") parser.add_option("--fl", type="float" , default=0.0 , help="stop-band frequency (Default - no filtration)") parser.add_option("--aa", type="float" , default=0.0 , help="fall off of the filter (Default - no filtration)") parser.add_option("--CTF", action="store_true", default=False, help="use CFT correction") parser.add_option("--VERBOSE", action="store_true", default=False, help="Long output for debugging") #parser.add_option("--MPI" , action="store_true", default=False, help="use MPI version") #parser.add_option("--radiuspca", type="int" , default=-1 , help="radius for PCA" ) #parser.add_option("--iter", type="int" , default=40 , help="maximum number of iterations (stop criterion of reconstruction process)" ) #parser.add_option("--abs", type="float" , default=0.0 , help="minimum average absolute change of voxels' values (stop criterion of reconstruction process)" ) #parser.add_option("--squ", type="float" , default=0.0 , help="minimum average squared change of voxels' values (stop criterion of reconstruction process)" ) parser.add_option("--VAR" , action="store_true", default=False, help="stack on input consists of 2D variances (Default False)") parser.add_option("--decimate", type="float", default=1.0, help="image decimate rate, a number large than 1. default is 1") parser.add_option("--window", type="int", default=0, help="reduce images to a small image size without changing pixel_size. Default value is zero.") #parser.add_option("--SND", action="store_true", default=False, help="compute squared normalized differences (Default False)") parser.add_option("--nvec", type="int" , default=0 , help="number of eigenvectors, default = 0 meaning no PCA calculated") parser.add_option("--symmetrize", action="store_true", default=False, help="Prepare input stack for handling symmetry (Default False)") (options,args) = parser.parse_args() ##### from mpi import mpi_init, mpi_comm_rank, mpi_comm_size, mpi_recv, MPI_COMM_WORLD, MPI_TAG_UB from mpi import mpi_barrier, mpi_reduce, mpi_bcast, mpi_send, MPI_FLOAT, MPI_SUM, MPI_INT, MPI_MAX from applications import MPI_start_end from reconstruction import recons3d_em, recons3d_em_MPI from reconstruction import recons3d_4nn_MPI, recons3d_4nn_ctf_MPI from utilities import print_begin_msg, print_end_msg, print_msg from utilities import read_text_row, get_image, get_im from utilities import bcast_EMData_to_all, bcast_number_to_all from utilities import get_symt # This is code for handling symmetries by the above program. To be incorporated. PAP 01/27/2015 from EMAN2db import db_open_dict if options.symmetrize : try: sys.argv = mpi_init(len(sys.argv), sys.argv) try: number_of_proc = mpi_comm_size(MPI_COMM_WORLD) if( number_of_proc > 1 ): ERROR("Cannot use more than one CPU for symmetry prepration","sx3dvariability",1) except: pass except: pass # Input #instack = "Clean_NORM_CTF_start_wparams.hdf" #instack = "bdb:data" instack = args[0] sym = options.sym if( sym == "c1" ): ERROR("Thre is no need to symmetrize stack for C1 symmetry","sx3dvariability",1) if(instack[:4] !="bdb:"): stack = "bdb:data" delete_bdb(stack) cmdexecute("sxcpy.py "+instack+" "+stack) else: stack = instack qt = EMUtil.get_all_attributes(stack,'xform.projection') na = len(qt) ts = get_symt(sym) ks = len(ts) angsa = [None]*na for k in xrange(ks): delete_bdb("bdb:Q%1d"%k) cmdexecute("e2bdb.py "+stack+" --makevstack=bdb:Q%1d"%k) DB = db_open_dict("bdb:Q%1d"%k) for i in xrange(na): ut = qt[i]*ts[k] DB.set_attr(i, "xform.projection", ut) #bt = ut.get_params("spider") #angsa[i] = [round(bt["phi"],3)%360.0, round(bt["theta"],3)%360.0, bt["psi"], -bt["tx"], -bt["ty"]] #write_text_row(angsa, 'ptsma%1d.txt'%k) #cmdexecute("e2bdb.py "+stack+" --makevstack=bdb:Q%1d"%k) #cmdexecute("sxheader.py bdb:Q%1d --params=xform.projection --import=ptsma%1d.txt"%(k,k)) DB.close() delete_bdb("bdb:sdata") cmdexecute("e2bdb.py . --makevstack=bdb:sdata --filt=Q") #cmdexecute("ls EMAN2DB/sdata*") a = get_im("bdb:sdata") a.set_attr("variabilitysymmetry",sym) a.write_image("bdb:sdata") else: sys.argv = mpi_init(len(sys.argv), sys.argv) myid = mpi_comm_rank(MPI_COMM_WORLD) number_of_proc = mpi_comm_size(MPI_COMM_WORLD) main_node = 0 if len(args) == 1: stack = args[0] else: print( "usage: " + usage) print( "Please run '" + progname + " -h' for detailed options") return 1 t0 = time() # obsolete flags options.MPI = True options.nvec = 0 options.radiuspca = -1 options.iter = 40 options.abs = 0.0 options.squ = 0.0 if options.fl > 0.0 and options.aa == 0.0: ERROR("Fall off has to be given for the low-pass filter", "sx3dvariability", 1, myid) if options.VAR and options.SND: ERROR("Only one of var and SND can be set!", "sx3dvariability", myid) exit() if options.VAR and (options.ave2D or options.ave3D or options.var2D): ERROR("When VAR is set, the program cannot output ave2D, ave3D or var2D", "sx3dvariability", 1, myid) exit() #if options.SND and (options.ave2D or options.ave3D): # ERROR("When SND is set, the program cannot output ave2D or ave3D", "sx3dvariability", 1, myid) # exit() if options.nvec > 0 : ERROR("PCA option not implemented", "sx3dvariability", 1, myid) exit() if options.nvec > 0 and options.ave3D == None: ERROR("When doing PCA analysis, one must set ave3D", "sx3dvariability", myid=myid) exit() import string options.sym = options.sym.lower() if global_def.CACHE_DISABLE: from utilities import disable_bdb_cache disable_bdb_cache() global_def.BATCH = True if myid == main_node: print_begin_msg("sx3dvariability") print_msg("%-70s: %s\n"%("Input stack", stack)) img_per_grp = options.img_per_grp nvec = options.nvec radiuspca = options.radiuspca symbaselen = 0 if myid == main_node: nima = EMUtil.get_image_count(stack) img = get_image(stack) nx = img.get_xsize() ny = img.get_ysize() if options.sym != "c1" : imgdata = get_im(stack) try: i = imgdata.get_attr("variabilitysymmetry") if(i != options.sym): ERROR("The symmetry provided does not agree with the symmetry of the input stack", "sx3dvariability", myid=myid) except: ERROR("Input stack is not prepared for symmetry, please follow instructions", "sx3dvariability", myid=myid) from utilities import get_symt i = len(get_symt(options.sym)) if((nima/i)*i != nima): ERROR("The length of the input stack is incorrect for symmetry processing", "sx3dvariability", myid=myid) symbaselen = nima/i else: symbaselen = nima 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) Tracker ={} Tracker["nx"] =nx Tracker["ny"] =ny Tracker["total_stack"]=nima if options.decimate==1.: if options.window !=0: nx = options.window ny = options.window else: if options.window ==0: nx = int(nx/options.decimate) ny = int(ny/options.decimate) else: nx = int(options.window/options.decimate) ny = nx symbaselen = bcast_number_to_all(symbaselen) if radiuspca == -1: radiuspca = nx/2-2 if myid == main_node: print_msg("%-70s: %d\n"%("Number of projection", nima)) img_begin, img_end = MPI_start_end(nima, number_of_proc, myid) """ if options.SND: from projection import prep_vol, prgs from statistics import im_diff from utilities import get_im, model_circle, get_params_proj, set_params_proj from utilities import get_ctf, generate_ctf from filter import filt_ctf imgdata = EMData.read_images(stack, range(img_begin, img_end)) if options.CTF: vol = recons3d_4nn_ctf_MPI(myid, imgdata, 1.0, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1) else: vol = recons3d_4nn_MPI(myid, imgdata, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1) bcast_EMData_to_all(vol, myid) volft, kb = prep_vol(vol) mask = model_circle(nx/2-2, nx, ny) varList = [] for i in xrange(img_begin, img_end): phi, theta, psi, s2x, s2y = get_params_proj(imgdata[i-img_begin]) ref_prj = prgs(volft, kb, [phi, theta, psi, -s2x, -s2y]) if options.CTF: ctf_params = get_ctf(imgdata[i-img_begin]) ref_prj = filt_ctf(ref_prj, generate_ctf(ctf_params)) diff, A, B = im_diff(ref_prj, imgdata[i-img_begin], mask) diff2 = diff*diff set_params_proj(diff2, [phi, theta, psi, s2x, s2y]) varList.append(diff2) mpi_barrier(MPI_COMM_WORLD) """ if options.VAR: #varList = EMData.read_images(stack, range(img_begin, img_end)) varList = [] this_image = EMData() for index_of_particle in xrange(img_begin,img_end): this_image.read_image(stack,index_of_particle) varList.append(image_decimate_window_xform_ctf(img,options.decimate,options.window,options.CTF)) else: from utilities import bcast_number_to_all, bcast_list_to_all, send_EMData, recv_EMData from utilities import set_params_proj, get_params_proj, params_3D_2D, get_params2D, set_params2D, compose_transform2 from utilities import model_blank, nearest_proj, model_circle from applications import pca from statistics import avgvar, avgvar_ctf, ccc from filter import filt_tanl from morphology import threshold, square_root from projection import project, prep_vol, prgs from sets import Set if myid == main_node: t1 = time() proj_angles = [] aveList = [] tab = EMUtil.get_all_attributes(stack, 'xform.projection') for i in xrange(nima): t = tab[i].get_params('spider') phi = t['phi'] theta = t['theta'] psi = t['psi'] x = theta if x > 90.0: x = 180.0 - x x = x*10000+psi proj_angles.append([x, t['phi'], t['theta'], t['psi'], i]) t2 = time() print_msg("%-70s: %d\n"%("Number of neighboring projections", img_per_grp)) print_msg("...... Finding neighboring projections\n") if options.VERBOSE: print "Number of images per group: ", img_per_grp print "Now grouping projections" proj_angles.sort() proj_angles_list = [0.0]*(nima*4) if myid == main_node: for i in xrange(nima): proj_angles_list[i*4] = proj_angles[i][1] proj_angles_list[i*4+1] = proj_angles[i][2] proj_angles_list[i*4+2] = proj_angles[i][3] proj_angles_list[i*4+3] = proj_angles[i][4] proj_angles_list = bcast_list_to_all(proj_angles_list, myid, main_node) proj_angles = [] for i in xrange(nima): proj_angles.append([proj_angles_list[i*4], proj_angles_list[i*4+1], proj_angles_list[i*4+2], int(proj_angles_list[i*4+3])]) del proj_angles_list proj_list, mirror_list = nearest_proj(proj_angles, img_per_grp, range(img_begin, img_end)) all_proj = Set() for im in proj_list: for jm in im: all_proj.add(proj_angles[jm][3]) all_proj = list(all_proj) if options.VERBOSE: print "On node %2d, number of images needed to be read = %5d"%(myid, len(all_proj)) index = {} for i in xrange(len(all_proj)): index[all_proj[i]] = i mpi_barrier(MPI_COMM_WORLD) if myid == main_node: print_msg("%-70s: %.2f\n"%("Finding neighboring projections lasted [s]", time()-t2)) print_msg("%-70s: %d\n"%("Number of groups processed on the main node", len(proj_list))) if options.VERBOSE: print "Grouping projections took: ", (time()-t2)/60 , "[min]" print "Number of groups on main node: ", len(proj_list) mpi_barrier(MPI_COMM_WORLD) if myid == main_node: print_msg("...... calculating the stack of 2D variances \n") if options.VERBOSE: print "Now calculating the stack of 2D variances" proj_params = [0.0]*(nima*5) aveList = [] varList = [] if nvec > 0: eigList = [[] for i in xrange(nvec)] if options.VERBOSE: print "Begin to read images on processor %d"%(myid) ttt = time() #imgdata = EMData.read_images(stack, all_proj) img = EMData() imgdata = [] for index_of_proj in xrange(len(all_proj)): img.read_image(stack, all_proj[index_of_proj]) dmg = image_decimate_window_xform_ctf(img,options.decimate,options.window,options.CTF) #print dmg.get_xsize(), "init" imgdata.append(dmg) if options.VERBOSE: print "Reading images on processor %d done, time = %.2f"%(myid, time()-ttt) print "On processor %d, we got %d images"%(myid, len(imgdata)) mpi_barrier(MPI_COMM_WORLD) ''' imgdata2 = EMData.read_images(stack, range(img_begin, img_end)) if options.fl > 0.0: for k in xrange(len(imgdata2)): imgdata2[k] = filt_tanl(imgdata2[k], options.fl, options.aa) if options.CTF: vol = recons3d_4nn_ctf_MPI(myid, imgdata2, 1.0, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1) else: vol = recons3d_4nn_MPI(myid, imgdata2, symmetry=options.sym, npad=options.npad, xysize=-1, zsize=-1) if myid == main_node: vol.write_image("vol_ctf.hdf") print_msg("Writing to the disk volume reconstructed from averages as : %s\n"%("vol_ctf.hdf")) del vol, imgdata2 mpi_barrier(MPI_COMM_WORLD) ''' from applications import prepare_2d_forPCA from utilities import model_blank for i in xrange(len(proj_list)): ki = proj_angles[proj_list[i][0]][3] if ki >= symbaselen: continue mi = index[ki] phiM, thetaM, psiM, s2xM, s2yM = get_params_proj(imgdata[mi]) grp_imgdata = [] for j in xrange(img_per_grp): mj = index[proj_angles[proj_list[i][j]][3]] phi, theta, psi, s2x, s2y = get_params_proj(imgdata[mj]) alpha, sx, sy, mirror = params_3D_2D_NEW(phi, theta, psi, s2x, s2y, mirror_list[i][j]) if thetaM <= 90: if mirror == 0: alpha, sx, sy, scale = compose_transform2(alpha, sx, sy, 1.0, phiM-phi, 0.0, 0.0, 1.0) else: alpha, sx, sy, scale = compose_transform2(alpha, sx, sy, 1.0, 180-(phiM-phi), 0.0, 0.0, 1.0) else: if mirror == 0: alpha, sx, sy, scale = compose_transform2(alpha, sx, sy, 1.0, -(phiM-phi), 0.0, 0.0, 1.0) else: alpha, sx, sy, scale = compose_transform2(alpha, sx, sy, 1.0, -(180-(phiM-phi)), 0.0, 0.0, 1.0) set_params2D(imgdata[mj], [alpha, sx, sy, mirror, 1.0]) grp_imgdata.append(imgdata[mj]) #print grp_imgdata[j].get_xsize(), imgdata[mj].get_xsize() if not options.no_norm: #print grp_imgdata[j].get_xsize() mask = model_circle(nx/2-2, nx, nx) for k in xrange(img_per_grp): ave, std, minn, maxx = Util.infomask(grp_imgdata[k], mask, False) grp_imgdata[k] -= ave grp_imgdata[k] /= std del mask if options.fl > 0.0: from filter import filt_ctf, filt_table from fundamentals import fft, window2d nx2 = 2*nx ny2 = 2*ny if options.CTF: from utilities import pad for k in xrange(img_per_grp): grp_imgdata[k] = window2d(fft( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa) ),nx,ny) #grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny) #grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa) else: for k in xrange(img_per_grp): grp_imgdata[k] = filt_tanl( grp_imgdata[k], options.fl, options.aa) #grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny) #grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa) else: from utilities import pad, read_text_file from filter import filt_ctf, filt_table from fundamentals import fft, window2d nx2 = 2*nx ny2 = 2*ny if options.CTF: from utilities import pad for k in xrange(img_per_grp): grp_imgdata[k] = window2d( fft( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1) ) , nx,ny) #grp_imgdata[k] = window2d(fft( filt_table( filt_tanl( filt_ctf(fft(pad(grp_imgdata[k], nx2, ny2, 1,0.0)), grp_imgdata[k].get_attr("ctf"), binary=1), options.fl, options.aa), fifi) ),nx,ny) #grp_imgdata[k] = filt_tanl(grp_imgdata[k], options.fl, options.aa) ''' if i < 10 and myid == main_node: for k in xrange(10): grp_imgdata[k].write_image("grp%03d.hdf"%i, k) ''' """ if myid == main_node and i==0: for pp in xrange(len(grp_imgdata)): grp_imgdata[pp].write_image("pp.hdf", pp) """ ave, grp_imgdata = prepare_2d_forPCA(grp_imgdata) """ if myid == main_node and i==0: for pp in xrange(len(grp_imgdata)): grp_imgdata[pp].write_image("qq.hdf", pp) """ var = model_blank(nx,ny) for q in grp_imgdata: Util.add_img2( var, q ) Util.mul_scalar( var, 1.0/(len(grp_imgdata)-1)) # Switch to std dev var = square_root(threshold(var)) #if options.CTF: ave, var = avgvar_ctf(grp_imgdata, mode="a") #else: ave, var = avgvar(grp_imgdata, mode="a") """ if myid == main_node: ave.write_image("avgv.hdf",i) var.write_image("varv.hdf",i) """ set_params_proj(ave, [phiM, thetaM, 0.0, 0.0, 0.0]) set_params_proj(var, [phiM, thetaM, 0.0, 0.0, 0.0]) aveList.append(ave) varList.append(var) if options.VERBOSE: print "%5.2f%% done on processor %d"%(i*100.0/len(proj_list), myid) if nvec > 0: eig = pca(input_stacks=grp_imgdata, subavg="", mask_radius=radiuspca, nvec=nvec, incore=True, shuffle=False, genbuf=True) for k in xrange(nvec): set_params_proj(eig[k], [phiM, thetaM, 0.0, 0.0, 0.0]) eigList[k].append(eig[k]) """ if myid == 0 and i == 0: for k in xrange(nvec): eig[k].write_image("eig.hdf", k) """ del imgdata # To this point, all averages, variances, and eigenvectors are computed if options.ave2D: from fundamentals import fpol if myid == main_node: km = 0 for i in xrange(number_of_proc): if i == main_node : for im in xrange(len(aveList)): aveList[im].write_image(options.ave2D, km) km += 1 else: nl = mpi_recv(1, MPI_INT, i, MPI_TAG_UB, MPI_COMM_WORLD) nl = int(nl[0]) for im in xrange(nl): ave = recv_EMData(i, im+i+70000) """ nm = mpi_recv(1, MPI_INT, i, MPI_TAG_UB, MPI_COMM_WORLD) nm = int(nm[0]) members = mpi_recv(nm, MPI_INT, i, MPI_TAG_UB, MPI_COMM_WORLD) ave.set_attr('members', map(int, members)) members = mpi_recv(nm, MPI_FLOAT, i, MPI_TAG_UB, MPI_COMM_WORLD) ave.set_attr('pix_err', map(float, members)) members = mpi_recv(3, MPI_FLOAT, i, MPI_TAG_UB, MPI_COMM_WORLD) ave.set_attr('refprojdir', map(float, members)) """ tmpvol=fpol(ave, Tracker["nx"],Tracker["nx"],Tracker["nx"]) tmpvol.write_image(options.ave2D, km) km += 1 else: mpi_send(len(aveList), 1, MPI_INT, main_node, MPI_TAG_UB, MPI_COMM_WORLD) for im in xrange(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, MPI_TAG_UB, MPI_COMM_WORLD) mpi_send(members, len(members), MPI_INT, main_node, MPI_TAG_UB, MPI_COMM_WORLD) members = aveList[im].get_attr('pix_err') mpi_send(members, len(members), MPI_FLOAT, main_node, MPI_TAG_UB, MPI_COMM_WORLD) try: members = aveList[im].get_attr('refprojdir') mpi_send(members, 3, MPI_FLOAT, main_node, MPI_TAG_UB, MPI_COMM_WORLD) except: mpi_send([-999.0,-999.0,-999.0], 3, MPI_FLOAT, main_node, MPI_TAG_UB, MPI_COMM_WORLD) """ if options.ave3D: from fundamentals import fpol if options.VERBOSE: print "Reconstructing 3D average volume" ave3D = recons3d_4nn_MPI(myid, aveList, symmetry=options.sym, npad=options.npad) bcast_EMData_to_all(ave3D, myid) if myid == main_node: ave3D=fpol(ave3D,Tracker["nx"],Tracker["nx"],Tracker["nx"]) ave3D.write_image(options.ave3D) print_msg("%-70s: %s\n"%("Writing to the disk volume reconstructed from averages as", options.ave3D)) del ave, var, proj_list, stack, phi, theta, psi, s2x, s2y, alpha, sx, sy, mirror, aveList if nvec > 0: for k in xrange(nvec): if options.VERBOSE: print "Reconstruction eigenvolumes", k cont = True ITER = 0 mask2d = model_circle(radiuspca, nx, nx) while cont: #print "On node %d, iteration %d"%(myid, ITER) eig3D = recons3d_4nn_MPI(myid, eigList[k], symmetry=options.sym, npad=options.npad) bcast_EMData_to_all(eig3D, myid, main_node) if options.fl > 0.0: eig3D = filt_tanl(eig3D, options.fl, options.aa) if myid == main_node: eig3D.write_image("eig3d_%03d.hdf"%k, ITER) Util.mul_img( eig3D, model_circle(radiuspca, nx, nx, nx) ) eig3Df, kb = prep_vol(eig3D) del eig3D cont = False icont = 0 for l in xrange(len(eigList[k])): phi, theta, psi, s2x, s2y = get_params_proj(eigList[k][l]) proj = prgs(eig3Df, kb, [phi, theta, psi, s2x, s2y]) cl = ccc(proj, eigList[k][l], mask2d) if cl < 0.0: icont += 1 cont = True eigList[k][l] *= -1.0 u = int(cont) u = mpi_reduce([u], 1, MPI_INT, MPI_MAX, main_node, MPI_COMM_WORLD) icont = mpi_reduce([icont], 1, MPI_INT, MPI_SUM, main_node, MPI_COMM_WORLD) if myid == main_node: u = int(u[0]) print " Eigenvector: ",k," number changed ",int(icont[0]) else: u = 0 u = bcast_number_to_all(u, main_node) cont = bool(u) ITER += 1 del eig3Df, kb mpi_barrier(MPI_COMM_WORLD) del eigList, mask2d if options.ave3D: del ave3D if options.var2D: from fundamentals import fpol if myid == main_node: km = 0 for i in xrange(number_of_proc): if i == main_node : for im in xrange(len(varList)): tmpvol=fpol(varList[im], Tracker["nx"], Tracker["nx"],1) tmpvol.write_image(options.var2D, km) km += 1 else: nl = mpi_recv(1, MPI_INT, i, MPI_TAG_UB, MPI_COMM_WORLD) nl = int(nl[0]) for im in xrange(nl): ave = recv_EMData(i, im+i+70000) tmpvol=fpol(ave, Tracker["nx"], Tracker["nx"],1) tmpvol.write_image(options.var2D, km) km += 1 else: mpi_send(len(varList), 1, MPI_INT, main_node, MPI_TAG_UB, MPI_COMM_WORLD) for im in xrange(len(varList)): send_EMData(varList[im], main_node, im+myid+70000)# What with the attributes?? mpi_barrier(MPI_COMM_WORLD) if options.var3D: if myid == main_node and options.VERBOSE: print "Reconstructing 3D variability volume" t6 = time() radiusvar = options.radiusvar if( radiusvar < 0 ): radiusvar = nx//2 -3 res = recons3d_4nn_MPI(myid, varList, symmetry=options.sym, npad=options.npad) #res = recons3d_em_MPI(varList, vol_stack, options.iter, radiusvar, options.abs, True, options.sym, options.squ) if myid == main_node: from fundamentals import fpol res =fpol(res, Tracker["nx"], Tracker["nx"], Tracker["nx"]) res.write_image(options.var3D) if myid == main_node: print_msg("%-70s: %.2f\n"%("Reconstructing 3D variability took [s]", time()-t6)) if options.VERBOSE: print "Reconstruction took: %.2f [min]"%((time()-t6)/60) if myid == main_node: print_msg("%-70s: %.2f\n"%("Total time for these computations [s]", time()-t0)) if options.VERBOSE: print "Total time for these computations: %.2f [min]"%((time()-t0)/60) print_end_msg("sx3dvariability") global_def.BATCH = False from mpi import mpi_finalize mpi_finalize()
def resample( prjfile, outdir, bufprefix, nbufvol, nvol, seedbase,\ delta, d, snr, CTF, npad,\ MPI, myid, ncpu, verbose = 0 ): from utilities import even_angles from random import seed, jumpahead, shuffle import os from sys import exit nprj = EMUtil.get_image_count( prjfile ) if MPI: from mpi import mpi_barrier, MPI_COMM_WORLD if myid == 0: if os.path.exists(outdir): nx = 1 else: nx = 0 else: nx = 0 ny = bcast_number_to_all(nx, source_node = 0) if ny == 1: ERROR('Output directory exists, please change the name and restart the program', "resample", 1,myid) mpi_barrier(MPI_COMM_WORLD) if myid == 0: os.mkdir(outdir) mpi_barrier(MPI_COMM_WORLD) else: if os.path.exists(outdir): ERROR('Output directory exists, please change the name and restart the program', "resample", 1,0) os.mkdir(outdir) if(verbose == 1): finfo=open( os.path.join(outdir, "progress%04d.txt" % myid), "w" ) else: finfo = None #print " before evenangles",myid from utilities import getvec from numpy import array, reshape refa = even_angles(delta) nrefa = len(refa) refnormal = zeros((nrefa,3),'float32') tetref = [0.0]*nrefa for i in xrange(nrefa): tr = getvec( refa[i][0], refa[i][1] ) for j in xrange(3): refnormal[i][j] = tr[j] tetref[i] = refa[i][1] del refa vct = array([0.0]*(3*nprj),'float32') if myid == 0: print " will read ",myid tr = EMUtil.get_all_attributes(prjfile,'xform.projection') tetprj = [0.0]*nprj for i in xrange(nprj): temp = tr[i].get_params("spider") tetprj[i] = temp["theta"] if(tetprj[i] > 90.0): tetprj[i] = 180.0 - tetprj[i] vct[3*i+0] = tr[i].at(2,0) vct[3*i+1] = tr[i].at(2,1) vct[3*i+2] = tr[i].at(2,2) del tr else: tetprj = [0.0]*nprj #print " READ ",myid if MPI: #print " will bcast",myid from mpi import mpi_bcast, MPI_FLOAT, MPI_COMM_WORLD vct = mpi_bcast(vct,len(vct),MPI_FLOAT,0,MPI_COMM_WORLD) from utilities import bcast_list_to_all tetprj = bcast_list_to_all(tetprj, myid, 0) #print " reshape ",myid vct = reshape(vct,(nprj,3)) assignments = [[] for i in xrange(nrefa)] dspn = 1.25*delta for k in xrange(nprj): best_s = -1.0 best_i = -1 for i in xrange( nrefa ): if(abs(tetprj[k] - tetref[i]) <= dspn): s = abs(refnormal[i][0]*vct[k][0] + refnormal[i][1]*vct[k][1] + refnormal[i][2]*vct[k][2]) if s > best_s: best_s = s best_i = i assignments[best_i].append(k) am = len(assignments[0]) mufur = 1.0/am for i in xrange(1,len(assignments)): ti = len(assignments[i]) am = min(am, ti) if(ti>0): mufur += 1.0/ti del tetprj,tetref dp = 1.0 - d # keep that many in each direction keep = int(am*dp +0.5) mufur = keep*nrefa/(1.0 - mufur*keep/float(nrefa)) if myid == 0: print " Number of projections ",nprj,". Number of reference directions ",nrefa,", multiplicative factor for the variance ",mufur print " Minimum number of assignments ",am," Number of projections used per stratum ", keep," Number of projections in resampled structure ",int(am*dp +0.5)*nrefa if am <2 or am == keep: print "incorrect settings" exit() # FIX if(seedbase < 1): seed() jumpahead(17*myid+123) else: seed(seedbase) jumpahead(17*myid+123) volfile = os.path.join(outdir, "bsvol%04d.hdf" % myid) from random import randint niter = nvol/ncpu/nbufvol for kiter in xrange(niter): if(verbose == 1): finfo.write( "Iteration %d: \n" % kiter ) finfo.flush() iter_start = time() # the following has to be converted to resample mults=1 means take given projection., mults=0 means omit mults = [ [0]*nprj for i in xrange(nbufvol) ] for i in xrange(nbufvol): for l in xrange(nrefa): mass = assignments[l][:] shuffle(mass) mass = mass[:keep] mass.sort() #print l, " * ",mass for k in xrange(keep): mults[i][mass[k]] = 1 ''' lout = [] for l in xrange(len(mults[i])): if mults[i][l] == 1: lout.append(l) write_text_file(lout, os.path.join(outdir, "list%04d_%03d.txt" %(i, myid))) del lout ''' del mass rectors, fftvols, wgtvols = resample_prepare( prjfile, nbufvol, snr, CTF, npad ) resample_insert( bufprefix, fftvols, wgtvols, mults, CTF, npad, finfo ) del mults resample_finish( rectors, fftvols, wgtvols, volfile, kiter, nprj, finfo ) rectors = None fftvols = None wgtvols = None if(verbose == 1): finfo.write( "time for iteration: %10.3f\n" % (time() - iter_start) ) finfo.flush()