Пример #1
0
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
Пример #2
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