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
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()
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
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
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