def main(args): star = parse_star(args.input, keep_index=False) if args.class_2d is not None: refs = glob.glob(args.class_2d) elif args.class_3d is not None: refs = glob.glob(args.class_3d) else: refs = [] shifts = [] for r in refs: if args.class_3d is not None: refmap = EMData(r) com = Vec3f(*refmap.phase_cog()[:3]) shifts.append(com) else: stack = EMData.read_images(r) for im in stack: com = Vec2f(*im.phase_cog()[:2]) shifts.append(com) if args.class_2d is None and args.class_3d is None: for ptcl in star.rows: im = EMData.read_image(ptcl) com = im.phase_cog() ptcl["rlnOriginX"] += com[0] ptcl["rlnOriginY"] += com[1] else: for ptcl in star.rows: com = shifts[ptcl["rlnClassNumber"]] xshift, yshift = transform_com(com, ptcl) ptcl["rlnOriginX"] += xshift ptcl["rlnOriginY"] += yshift if args.zero_origin: star["rlnCoordinateX"] = star["rlnCoordinateX"] - star["rlnOriginX"] star["rlnCoordinateY"] = star["rlnCoordinateY"] - star["rlnOriginY"] star["rlnOriginX"] = 0 star["rlnOriginY"] = 0 write_star(args.output, star, reindex=True) return 0
def main(options): """ Projection subtraction program entry point. :param options: Command-line arguments parsed by ArgumentParser.parse_args() :return: Exit status """ rchop = lambda x, y: x if not x.endswith(y) or len(y) == 0 else x[:-len(y)] options.output = rchop(options.output, ".star") options.suffix = rchop(options.suffix, ".mrc") options.suffix = rchop(options.suffix, ".mrcs") star = StarFile(options.input) npart = len(star['rlnImageName']) sub_dens = EMData(options.submap) if options.wholemap is not None: dens = EMData(options.wholemap) else: print "Reference map is required." return 1 # Write star header for output.star. top_header = "\ndata_\n\nloop_\n" headings = star.keys() output_star = open("{0}.star".format(options.output), 'w') output_star.write(top_header) for i, heading in enumerate(headings): output_star.write("_{0} #{1}\n".format(heading, i + 1)) if options.recenter: # Compute difference vector between new and old mass centers. if options.wholemap is None: print "Reference map required for recentering." return 1 new_dens = dens - sub_dens # Note the sign of the shift in coordinate frame is opposite the shift in the CoM. recenter = Vec3f(*dens.phase_cog()[:3]) - Vec3f( *new_dens.phase_cog()[:3]) else: recenter = None pool = None if options.nproc > 1: # Compute subtraction in parallel. pool = Pool(processes=options.nproc) results = pool.imap( lambda x: subtract(x, dens, sub_dens, recenter=recenter, no_frc=options.no_frc, low_cutoff=options.low_cutoff, high_cutoff=options.high_cutoff), particles(star), chunksize=min(npart / options.nproc, options.maxchunk)) else: # Use serial generator. results = (subtract(x, dens, sub_dens, recenter=recenter, no_frc=options.no_frc, low_cutoff=options.low_cutoff, high_cutoff=options.high_cutoff) for x in particles(star)) # Write subtraction results to .mrcs and .star files. i = 0 nfile = 1 starpath = None mrcs = None mrcs_orig = None for r in results: if i % options.maxpart == 0: mrcsuffix = options.suffix + "_%d" % nfile nfile += 1 starpath = "{0}.mrcs".format( os.path.sep.join( os.path.relpath(mrcsuffix, options.output).split(os.path.sep)[1:])) mrcs = "{0}.mrcs".format(mrcsuffix) mrcs_orig = "{0}_original.mrcs".format(mrcsuffix) if os.path.exists(mrcs): os.remove(mrcs) if os.path.exists(mrcs_orig): os.remove(mrcs_orig) r.ptcl_norm_sub.append_image(mrcs) if options.original: r.ptcl.append_image(mrcs_orig) if logger.getEffectiveLevel( ) == logging.DEBUG: # Write additional debug output. ptcl_sub_img = r.ptcl.process("math.sub.optimal", { "ref": r.ctfproj, "actual": r.ctfproj_sub, "return_subim": True }) ptcl_lowpass = r.ptcl.process("filter.lowpass.gauss", { "apix": 1.22, "cutoff_freq": 0.05 }) ptcl_sub_lowpass = r.ptcl_norm_sub.process("filter.lowpass.gauss", { "apix": 1.22, "cutoff_freq": 0.05 }) ptcl_sub_img.write_image("poreclass_subimg.mrcs", -1) ptcl_lowpass.write_image("poreclass_lowpass.mrcs", -1) ptcl_sub_lowpass.write_image("poreclass_sublowpass.mrcs", -1) r.ctfproj.write_image("poreclass_ctfproj.mrcs", -1) r.ctfproj_sub.write_image("poreclass_ctfprojsub.mrcs", -1) assert r.meta.i == i # Assert particle order is preserved. star['rlnImageName'][i] = "{0:06d}@{1}".format( i % options.maxpart + 1, starpath) # Set new image name. r.meta.update(star) # Update StarFile with altered fields. line = ' '.join(str(star[key][i]) for key in headings) output_star.write("{0}\n".format(line)) i += 1 output_star.close() if pool is not None: pool.close() pool.join() return 0