示例#1
0
文件: stack.py 项目: smart111/pyem
def main(args):
    log = logging.getLogger('root')
    hdlr = logging.StreamHandler(sys.stdout)
    log.addHandler(hdlr)
    log.setLevel(logging.getLevelName(args.loglevel.upper()))
    # apix = args.apix = hdr["xlen"] / hdr["nx"]

    for fn in args.input:
        if not (fn.endswith(".star") or fn.endswith(".mrcs")
                or fn.endswith(".mrc")):
            log.error("Only .star, .mrc, and .mrcs files supported")
            return 1

    first_ptcl = 0
    dfs = []
    with mrc.ZSliceWriter(args.output) as writer:
        for fn in args.input:
            if fn.endswith(".star"):
                df = star.parse_star(fn, keep_index=False)
                star.augment_star_ucsf(df)
                df = df.sort_values([
                    star.UCSF.IMAGE_ORIGINAL_PATH,
                    star.UCSF.IMAGE_ORIGINAL_INDEX
                ])
                gb = df.groupby(star.UCSF.IMAGE_ORIGINAL_PATH)
                for name, g in gb:
                    with mrc.ZSliceReader(name) as reader:
                        for i in g[star.UCSF.IMAGE_ORIGINAL_INDEX].values:
                            writer.write(reader.read(i))
            else:
                with mrc.ZSliceReader(fn) as reader:
                    for img in reader:
                        writer.write(img)
                    df = pd.DataFrame(
                        {star.UCSF.IMAGE_ORIGINAL_INDEX: np.arange(reader.nz)})
                df[star.UCSF.IMAGE_ORIGINAL_PATH] = fn

            if args.star is not None:
                df[star.UCSF.IMAGE_INDEX] = np.arange(first_ptcl,
                                                      first_ptcl + df.shape[0])
                df[star.UCSF.IMAGE_PATH] = writer.path
                df["index"] = df[star.UCSF.IMAGE_INDEX]
                star.simplify_star_ucsf(df)
                dfs.append(df)
            first_ptcl += df.shape[0]

            if args.star is not None:
                df = pd.concat(dfs, join="inner")
                # df = pd.concat(dfs)
                # df = df.dropna(df, axis=1, how="any")
                star.write_star(args.star, df, reindex=True)

    return 0
示例#2
0
def consumer(queue, stack, apix=1.0, iothreads=None):
    log = logging.getLogger('root')
    with mrc.ZSliceWriter(stack, psz=apix) as zwriter:
        while True:
            log.debug("Get")
            i, ri = queue.get(block=True)
            log.debug("Got %d, queue for %s is size %d" %
                      (i, stack, queue.qsize()))
            if i == -1:
                break
            new_image = ri.get()
            log.debug("Result for %d was shape (%d,%d)" %
                      (i, new_image.shape[0], new_image.shape[1]))
            zwriter.write(new_image)
            queue.task_done()
            log.debug("Wrote %d to %d@%s" % (i, zwriter.i, stack))
    if iothreads is not None:
        iothreads.release()
示例#3
0
文件: project.py 项目: sunny1226/pyem
def main(args):
    log = logging.getLogger('root')
    hdlr = logging.StreamHandler(sys.stdout)
    log.addHandler(hdlr)
    log.setLevel(logging.getLevelName(args.loglevel.upper()))
    df = star.parse_star(args.input, keep_index=False)
    star.augment_star_ucsf(df)
    maxshift = np.round(np.max(np.abs(df[star.Relion.ORIGINS].values)))

    if args.map is not None:
        if args.map.endswith(".npy"):
            log.info("Reading precomputed 3D FFT of volume")
            f3d = np.load(args.map)
            log.info("Finished reading 3D FFT of volume")
            if args.size is None:
                args.size = (f3d.shape[0] - 3) // args.pfac
        else:
            vol = mrc.read(args.map, inc_header=False, compat="relion")
            if args.mask is not None:
                mask = mrc.read(args.mask, inc_header=False, compat="relion")
                vol *= mask
            if args.size is None:
                args.size = vol.shape[0]
            if args.crop is not None and args.size // 2 < maxshift + args.crop // 2:
                log.error(
                    "Some shifts are too large to crop (maximum crop is %d)" %
                    (args.size - 2 * maxshift))
                return 1
            log.info("Preparing 3D FFT of volume")
            f3d = vop.vol_ft(vol, pfac=args.pfac, threads=args.threads)
            log.info("Finished 3D FFT of volume")
    else:
        log.error("Please supply a map")
        return 1

    sz = (f3d.shape[0] - 3) // args.pfac
    apix = star.calculate_apix(df) * np.double(args.size) / sz
    sx, sy = np.meshgrid(np.fft.rfftfreq(sz), np.fft.fftfreq(sz))
    s = np.sqrt(sx**2 + sy**2)
    a = np.arctan2(sy, sx)
    log.info("Projection size is %d, unpadded volume size is %d" %
             (args.size, sz))
    log.info("Effective pixel size is %f A/px" % apix)

    if args.subtract and args.size != sz:
        log.error("Volume and projections must be same size when subtracting")
        return 1

    if args.crop is not None and args.size // 2 < maxshift + args.crop // 2:
        log.error("Some shifts are too large to crop (maximum crop is %d)" %
                  (args.size - 2 * maxshift))
        return 1

    ift = None

    with mrc.ZSliceWriter(args.output, psz=apix) as zsw:
        for i, p in df.iterrows():
            f2d = project(f3d,
                          p,
                          s,
                          sx,
                          sy,
                          a,
                          pfac=args.pfac,
                          apply_ctf=args.ctf,
                          size=args.size,
                          flip_phase=args.flip)
            if ift is None:
                ift = irfft2(f2d.copy(),
                             threads=args.threads,
                             planner_effort="FFTW_ESTIMATE",
                             auto_align_input=True,
                             auto_contiguous=True)
            proj = fftshift(
                ift(f2d.copy(),
                    np.zeros(ift.output_shape, dtype=ift.output_dtype)))
            log.debug("%f +/- %f" % (np.mean(proj), np.std(proj)))
            if args.subtract:
                with mrc.ZSliceReader(p["ucsfImagePath"]) as zsr:
                    img = zsr.read(p["ucsfImageIndex"])
                log.debug("%f +/- %f" % (np.mean(img), np.std(img)))
                proj = img - proj
            if args.crop is not None:
                orihalf = args.size // 2
                newhalf = args.crop // 2
                x = orihalf - np.int(np.round(p[star.Relion.ORIGINX]))
                y = orihalf - np.int(np.round(p[star.Relion.ORIGINY]))
                proj = proj[y - newhalf:y + newhalf, x - newhalf:x + newhalf]
            zsw.write(proj)
            log.debug(
                "%d@%s: %d/%d" %
                (p["ucsfImageIndex"], p["ucsfImagePath"], i + 1, df.shape[0]))

    if args.star is not None:
        log.info("Writing output .star file")
        if args.crop is not None:
            df = star.recenter(df, inplace=True)
        if args.subtract:
            df[star.UCSF.IMAGE_ORIGINAL_PATH] = df[star.UCSF.IMAGE_PATH]
            df[star.UCSF.IMAGE_ORIGINAL_INDEX] = df[star.UCSF.IMAGE_INDEX]
        df[star.UCSF.IMAGE_PATH] = args.output
        df[star.UCSF.IMAGE_INDEX] = np.arange(df.shape[0])
        star.simplify_star_ucsf(df)
        star.write_star(args.star, df)
    return 0
示例#4
0
文件: stack.py 项目: mahabul123/pyem
def main(args):
    log = logging.getLogger('root')
    hdlr = logging.StreamHandler(sys.stdout)
    log.addHandler(hdlr)
    log.setLevel(logging.getLevelName(args.loglevel.upper()))
    # apix = args.apix = hdr["xlen"] / hdr["nx"]

    for fn in args.input:
        if not (fn.endswith(".star") or fn.endswith(".mrcs") or
                fn.endswith(".mrc") or fn.endswith(".par")):
            log.error("Only .star, .mrc, .mrcs, and .par files supported")
            return 1

    first_ptcl = 0
    dfs = []
    with mrc.ZSliceWriter(args.output) as writer:
        for fn in args.input:
            if fn.endswith(".star"):
                df = star.parse_star(fn, augment=True)
                if args.cls is not None:
                    df = star.select_classes(df, args.cls)
                star.set_original_fields(df, inplace=True)
                if args.resort:
                    df = df.sort_values([star.UCSF.IMAGE_ORIGINAL_PATH,
                                         star.UCSF.IMAGE_ORIGINAL_INDEX])
                for idx, row in df.iterrows():
                    if args.stack_path is not None:
                        input_stack_path = os.path.join(args.stack_path, row[star.UCSF.IMAGE_ORIGINAL_PATH])
                    else:
                        input_stack_path = row[star.UCSF.IMAGE_ORIGINAL_PATH]
                    with mrc.ZSliceReader(input_stack_path) as reader:
                        i = row[star.UCSF.IMAGE_ORIGINAL_INDEX]
                        writer.write(reader.read(i))
            elif fn.endswith(".par"):
                if args.stack_path is None:
                    log.error(".par file input requires --stack-path")
                    return 1
                df = metadata.par2star(metadata.parse_fx_par(fn), data_path=args.stack_path)
                # star.set_original_fields(df, inplace=True)  # Redundant.
                star.augment_star_ucsf(df)
            elif fn.endswith(".csv"):
                return 1
            elif fn.endswith(".cs"):
                return 1
            else:
                if fn.endswith(".mrcs"):
                    with mrc.ZSliceReader(fn) as reader:
                        for img in reader:
                            writer.write(img)
                        df = pd.DataFrame(
                            {star.UCSF.IMAGE_ORIGINAL_INDEX: np.arange(reader.nz)})
                    df[star.UCSF.IMAGE_ORIGINAL_PATH] = fn
                else:
                    print("Unrecognized input file type")
                    return 1
            if args.star is not None:
                df[star.UCSF.IMAGE_INDEX] = np.arange(first_ptcl,
                                                      first_ptcl + df.shape[0])
                if args.abs_path:
                    df[star.UCSF.IMAGE_PATH] = writer.path
                else:
                    df[star.UCSF.IMAGE_PATH] = os.path.relpath(writer.path, os.path.dirname(args.star))
                df["index"] = df[star.UCSF.IMAGE_INDEX]
                star.simplify_star_ucsf(df)
                dfs.append(df)
            first_ptcl += df.shape[0]

    if args.star is not None:
        df = pd.concat(dfs, join="inner")
        # df = pd.concat(dfs)
        # df = df.dropna(df, axis=1, how="any")
        if not args.relion2:  # Relion 3.1 style output.
            df = star.remove_deprecated_relion2(df, inplace=True)
            star.write_star(args.star, df, resort_records=False, optics=True)
        else:
            df = star.remove_new_relion31(df, inplace=True)
            star.write_star(args.star, df, resort_records=False, optics=False)
    return 0
示例#5
0
def main(args):
    log = logging.getLogger('root')
    hdlr = logging.StreamHandler(sys.stdout)
    log.addHandler(hdlr)
    log.setLevel(logging.getLevelName(args.loglevel.upper()))
    df = star.parse_star(args.input, keep_index=False)
    star.augment_star_ucsf(df)
    if args.map is not None:
        vol = mrc.read(args.map, inc_header=False, compat="relion")
        if args.mask is not None:
            mask = mrc.read(args.mask, inc_header=False, compat="relion")
            vol *= mask
    else:
        print("Please supply a map")
        return 1

    f3d = vop.vol_ft(vol, pfac=args.pfac, threads=args.threads)
    sz = f3d.shape[0] // 2 - 1
    sx, sy = np.meshgrid(np.fft.rfftfreq(sz), np.fft.fftfreq(sz))
    s = np.sqrt(sx**2 + sy**2)
    a = np.arctan2(sy, sx)

    ift = None

    with mrc.ZSliceWriter(args.output) as zsw:
        for i, p in df.iterrows():
            f2d = project(f3d,
                          p,
                          s,
                          sx,
                          sy,
                          a,
                          apply_ctf=args.ctf,
                          size=args.size)
            if ift is None:
                ift = irfft2(f2d.copy(),
                             threads=cpu_count(),
                             planner_effort="FFTW_ESTIMATE",
                             auto_align_input=True,
                             auto_contiguous=True)
            proj = fftshift(
                ift(f2d.copy(), np.zeros(vol.shape[:-1], dtype=vol.dtype)))
            log.debug("%f +/- %f" % (np.mean(proj), np.std(proj)))
            if args.subtract:
                with mrc.ZSliceReader(p["ucsfImagePath"]) as zsr:
                    img = zsr.read(p["ucsfImageIndex"])
                log.debug("%f +/- %f" % (np.mean(img), np.std(img)))
                proj = img - proj
            zsw.write(proj)
            log.info(
                "%d@%s: %d/%d" %
                (p["ucsfImageIndex"], p["ucsfImagePath"], i + 1, df.shape[0]))

    if args.star is not None:
        if args.subtract:
            df[star.UCSF.IMAGE_ORIGINAL_PATH] = df[star.UCSF.IMAGE_PATH]
            df[star.UCSF.IMAGE_ORIGINAL_INDEX] = df[star.UCSF.IMAGE_INDEX]
        df[star.UCSF.IMAGE_PATH] = args.output
        df[star.UCSF.IMAGE_INDEX] = np.arange(df.shape[0])
        star.simplify_star_ucsf(df)
        star.write_star(args.star, df)
    return 0