def main(args): if args.threshold is None: print("Please provide a binarization threshold") return 1 data, hdr = read(args.input, inc_header=True) mask = binarize_volume(data, args.threshold, minvol=args.minvol, fill=args.fill) if args.base_map is not None: base_map = read(args.base_map, inc_header=False) base_mask = binarize_volume(base_map, args.threshold, minvol=args.minvol, fill=args.fill) total_width = args.extend + args.edge_width excl_mask = binary_dilate(mask, total_width, strel=args.relion) base_mask = binary_dilate(base_mask, args.extend, strel=args.relion) base_mask = base_mask &~ excl_mask if args.overlap > 0: incl_mask = binary_dilate(base_mask, args.overlap, strel=args.relion) & excl_mask base_mask = base_mask | incl_mask mask = base_mask elif args.extend > 0: mask = binary_dilate(mask, args.extend, strel=args.relion) if args.close: se = binary_sphere(args.extend, False) mask = binary_closing(mask, structure=se, iterations=1) final = mask.astype(np.single) if args.edge_width != 0: dt = distance_transform_edt(~mask) # Compute *outward* distance transform of mask. idx = (dt <= args.edge_width) & (dt > 0) # Identify edge points by distance from mask. x = np.arange(1, args.edge_width + 1) # Domain of the edge profile. if "sin" in args.edge_profile: y = np.sin(np.linspace(np.pi/2, 0, args.edge_width + 1)) # Range of the edge profile. f = interp1d(x, y[1:]) final[idx] = f(dt[idx]) # Insert edge heights interpolated at distance transform values. write(args.output, final, psz=hdr["xlen"] / hdr["nx"]) return 0
def main(args): log = logging.getLogger('root') hdlr = logging.StreamHandler(sys.stdout) log.addHandler(hdlr) log.setLevel(logging.getLevelName(args.loglevel.upper())) pyfftw.interfaces.cache.enable() vol1 = mrc.read(args.volume1, inc_header=False, compat="relion") vol2 = mrc.read(args.volume2, inc_header=False, compat="relion") if args.mask is not None: mask = mrc.read(args.mask, inc_header=False, compat="relion") vol1 *= mask vol2 *= mask f3d1 = fft.rfftn(vol1, threads=args.threads) f3d2 = fft.rfftn(vol2, threads=args.threads) nside = 2**args.healpix_order x, y, z = pix2vec(nside, np.arange(12 * nside ** 2)) xhalf = x >= 0 hp = np.column_stack([x[xhalf], y[xhalf], z[xhalf]]) t0 = time.time() fcor = calc_dfsc(f3d1, f3d2, hp, np.deg2rad(args.arc)) log.info("Computed CFSC in %0.2f s" % (time.time() - t0)) fsc = calc_fsc(f3d1, f3d2) t0 = time.time() log.info("Computed GFSC in %0.2f s" % (time.time() - t0)) freqs = np.fft.rfftfreq(f3d1.shape[0]) np.save(args.output, np.row_stack([freqs, fsc, fcor])) return 0
def main(args): x = mrc.read(args.input[0]) sigma = np.zeros(x.shape) mu = x.copy() for i, f in enumerate(args.input[1:]): x = mrc.read(f) olddif = x - mu mu += (x - mu) / (i + 1) sigma += olddif * (x - mu) sigma_sq = np.power(sigma, 2) mrc.write(args.output, sigma_sq) if args.mean is not None: mrc.write(args.mean, mu) return 0
def main(args): x = mrc.read(args.input[0]) m2 = np.zeros(x.shape) mu = x.copy() for i, f in enumerate(args.input[1:]): x = mrc.read(f) olddif = x - mu mu += (x - mu) / (i + 1) m2 += olddif * (x - mu) var = m2 / len(args.input) mrc.write(args.output, var) if args.mean is not None: mrc.write(args.mean, mu) return 0
def main(args): log = logging.getLogger(__name__) hdlr = logging.StreamHandler(sys.stdout) log.addHandler(hdlr) log.setLevel(logging.getLevelName(args.loglevel.upper())) data = {} hdr = {} for i, inp in enumerate(args.input[1:]): d, h = read(inp, inc_header=True) if args.normalize: d = vop.normalize(d) data[ascii_lowercase[i]] = d hdr[ascii_lowercase[i]] = h if args.eval: final = eval(args.input[0], globals(), data) else: final = ne.evaluate(args.input[0], local_dict=data) if args.apix is None: args.apix = hdr[ascii_lowercase[0]]['xlen'] / hdr[ ascii_lowercase[0]]['nx'] write(args.output, final.astype(np.single), psz=args.apix) return 0
def main(args): if args.threshold is None: print("Please provide a binarization threshold") return 1 data, hdr = read(args.input, inc_header=True) mask = data >= args.threshold if args.minvol is not None: mask = binary_volume_opening(mask, args.minvol) if args.fill: mask = binary_fill_holes(mask) if args.extend is not None: se = binary_sphere(args.extend, False) mask = binary_dilation(mask, structure=se, iterations=1) if args.close: se = binary_sphere(args.extend, False) mask = binary_closing(mask, structure=se, iterations=1) final = mask.astype(np.single) if args.edge_width is not None: dt = distance_transform_edt( ~mask) # Compute *outward* distance transform of mask. idx = (dt <= args.edge_width) & ( dt > 0) # Identify edge points by distance from mask. x = np.arange(1, args.edge_width + 1) # Domain of the edge profile. if "sin" in args.edge_profile: y = np.sin(np.linspace(np.pi / 2, 0, args.edge_width + 1)) # Range of the edge profile. f = interp1d(x, y[1:]) final[idx] = f( dt[idx] ) # Insert edge heights interpolated at distance transform values. write(args.output, final, psz=hdr["xlen"] / hdr["nx"]) return 0
def main(args): pyfftw.interfaces.cache.enable() refmap = mrc.read(args.key, compat="relion") df = star.parse_star(args.input, keep_index=False) star.augment_star_ucsf(df) refmap_ft = vop.vol_ft(refmap, threads=args.threads) apix = star.calculate_apix(df) sz = refmap_ft.shape[0] // 2 - 1 sx, sy = np.meshgrid(rfftfreq(sz), fftfreq(sz)) s = np.sqrt(sx**2 + sy**2) r = s * sz r = np.round(r).astype(np.int64) r[r > sz // 2] = sz // 2 + 1 a = np.arctan2(sy, sx) def1 = df["rlnDefocusU"].values def2 = df["rlnDefocusV"].values angast = df["rlnDefocusAngle"].values phase = df["rlnPhaseShift"].values kv = df["rlnVoltage"].values ac = df["rlnAmplitudeContrast"].values cs = df["rlnSphericalAberration"].values xshift = df["rlnOriginX"].values yshift = df["rlnOriginY"].values score = np.zeros(df.shape[0]) # TODO parallelize for i, row in df.iterrows(): xcor = particle_xcorr(row, refmap_ft) if args.top is None: args.top = df.shape[0] top = df.iloc[np.argsort(score)][:args.top] star.simplify_star_ucsf(top) star.write_star(args.output, top) return 0
def main(args): log = logging.getLogger(__name__) hdlr = logging.StreamHandler(sys.stdout) log.addHandler(hdlr) log.setLevel(logging.getLevelName(args.loglevel.upper())) data, hdr = read(args.input, inc_header=True) if args.half2 is not None: half2, hdr_half2 = read(args.input, inc_header=True) if data.shape == half2.shape: data += half2 else: log.error("--half2 map is not the same shape as input map!") return 1 final = None box = np.array([hdr[a] for a in ["nx", "ny", "nz"]]) center = box // 2 if args.fft: if args.final_mask is not None: final_mask = read(args.final_mask) data *= final_mask data_ft = vop.vol_ft(data.T, pfac=args.pfac, threads=args.threads) np.save(args.output, data_ft) return 0 if args.transpose is not None: try: tax = [np.int64(a) for a in args.transpose.split(",")] data = np.transpose(data, axes=tax) except: log.error( "Transpose axes must be comma-separated list of three integers" ) return 1 if args.flip is not None: if args.flip.isnumeric(): args.flip = int(args.flip) else: args.flip = vop.label_to_axis(args.flip) data = np.flip(data, axis=args.flip) if args.apix is None: args.apix = hdr["xlen"] / hdr["nx"] log.info("Using computed pixel size of %f Angstroms" % args.apix) if args.normalize: if args.diameter is not None: if args.diameter > 1.0: args.diameter /= args.apix * 2 # Convert Angstrom diameter to pixel radius. if args.reference is not None: ref, refhdr = read(args.reference, inc_header=True) final, mu, sigma = vop.normalize(data, ref=ref, return_stats=True, rmask=args.diameter) else: final, mu, sigma = vop.normalize(data, return_stats=True, rmask=args.diameter) log.info("Mean: %f, Standard deviation: %f" % (mu, sigma)) if args.apix_out is not None: if args.scale is not None: log.warn("--apix-out supersedes --scale") args.scale = args.apix / args.apix_out elif args.scale is not None: args.apix_out = args.apix / args.scale elif args.boxsize is not None: args.scale = box[0] / np.double(args.boxsize) if args.apix_out is None: args.apix_out = args.apix if args.boxsize is None: if args.scale is None: args.boxsize = box[0] args.scale = 1 else: args.boxsize = np.int(box[0] * args.scale) log.info("Volume will be scaled by %f to size %d @ %f A/px" % (args.scale, args.boxsize, args.apix_out)) if args.target and args.transform: log.warn( "Target pose transformation will be applied after explicit matrix") if args.euler is not None and (args.target is not None or args.transform is not None): log.warn( "Euler transformation will be applied after target pose transformation" ) if args.translate is not None and (args.euler is not None or args.target is not None or args.transform is not None): log.warn("Translation will be applied after other transformations") if args.origin is not None: try: args.origin = np.array( [np.double(tok) for tok in args.origin.split(",")]) / args.apix assert np.all(args.origin < box) except: log.error( "Origin must be comma-separated list of x,y,z coordinates and lie within the box" ) return 1 else: args.origin = center log.info("Origin set to box center, %s" % (args.origin * args.apix)) if not (args.target is None and args.euler is None and args.transform is None and args.boxsize is None) \ and vop.ismask(data) and args.spline_order != 0: log.warn( "Input looks like a mask, --spline-order 0 (nearest neighbor) is recommended" ) if args.transform is not None: try: args.transform = np.array(json.loads(args.transform)) except: log.error("Transformation matrix must be in JSON/Numpy format") return 1 r = args.transform[:, :3] if args.transform.shape[1] == 4: t = args.transform[:, -1] / args.apix t = r.dot(args.origin) + t - args.origin t = -r.T.dot(t) else: t = 0 log.debug("Final rotation: %s" % str(r).replace("\n", "\n" + " " * 16)) log.debug("Final translation: %s (%f px)" % (str(t), np.linalg.norm(t))) data = vop.resample_volume(data, r=r, t=t, ori=None, order=args.spline_order, invert=args.invert) if args.target is not None: try: args.target = np.array( [np.double(tok) for tok in args.target.split(",")]) / args.apix except: log.error( "Standard pose target must be comma-separated list of x,y,z coordinates" ) return 1 args.target -= args.origin args.target = np.where(np.abs(args.target) < 1, 0, args.target) ori = None if args.origin is center else args.origin - center r = vec2rot(args.target) t = np.linalg.norm(args.target) log.info("Euler angles are %s deg and shift is %f px" % (np.rad2deg(rot2euler(r)), t)) log.debug("Final rotation: %s" % str(r).replace("\n", "\n" + " " * 16)) log.debug("Final translation: %s (%f px)" % (str(t), np.linalg.norm(t))) data = vop.resample_volume(data, r=r, t=args.target, ori=ori, order=args.spline_order, invert=args.invert) if args.euler is not None: try: args.euler = np.deg2rad( np.array([np.double(tok) for tok in args.euler.split(",")])) except: log.error( "Eulers must be comma-separated list of phi,theta,psi angles") return 1 r = euler2rot(*args.euler) offset = args.origin - 0.5 offset = offset - r.T.dot(offset) data = affine_transform(data, r.T, offset=offset, order=args.spline_order) if args.translate is not None: try: args.translate = np.array( [np.double(tok) for tok in args.translate.split(",")]) / args.apix except: log.error( "Translation vector must be comma-separated list of x,y,z coordinates" ) return 1 args.translate -= args.origin data = shift(data, -args.translate, order=args.spline_order) if final is None: final = data if args.final_mask is not None: final_mask = read(args.final_mask) final *= final_mask if args.scale != 1 or args.boxsize != box[0]: final = vop.resample_volume(final, scale=args.scale, output_shape=args.boxsize, order=args.spline_order) write(args.output, final, psz=args.apix_out) 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) 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())) os.environ["OMP_NUM_THREADS"] = str(args.threads) os.environ["OPENBLAS_NUM_THREADS"] = str(args.threads) os.environ["MKL_NUM_THREADS"] = str(args.threads) os.environ["NUMBA_NUM_THREADS"] = str(args.threads) outdir = os.path.dirname(args.output) outbase = os.path.basename(args.output) dfs = [star.parse_star(inp, keep_index=False) for inp in args.input] size_err = np.array( args.input)[np.where(~np.equal([df.shape[0] for df in dfs[1:]], dfs[0].shape[0]))] if len(size_err) > 0: log.error( "All files must have same number of particles. Offending files:\n%s" % ", ".join(size_err)) return 1 dfo = dfs[0] dfn = dfs[1] oq = geom.e2q_vec(np.deg2rad(dfo[star.Relion.ANGLES].values)) nq = geom.e2q_vec(np.deg2rad(dfn[star.Relion.ANGLES].values)) oqu = geom.normq(oq) nqu = geom.normq(nq) resq = geom.qtimes(geom.qconj(oqu), nqu) mu = geom.meanq(resq) resqu = geom.normq(resq, mu) si_mult = np.random.choice(resqu.shape[0] / args.multimer, args.sample / args.multimer, replace=False) si = np.array([ si_mult[i] * args.multimer + k for i in range(si_mult.shape[0]) for k in range(args.multimer) ]) not_si = np.setdiff1d(np.arange(resqu.shape[0], dtype=np.int), si) samp = resqu[si, :].copy() t = time.time() d = geom.pdistq(samp, np.zeros((samp.shape[0], samp.shape[0]), dtype=np.double)) log.info("Sample pairwise distances calculated in %0.3f s" % (time.time() - t)) g = geom.double_center(d, inplace=False) t = time.time() vals, vecs = np.linalg.eigh(g) log.info("Sample Gram matrix decomposed in %0.3f s" % (time.time() - t)) np.save(args.output + "_evals.npy", vals) np.save(args.output + "_evecs.npy", vecs) x = vecs[:, [-1, -2, -3]].dot(np.diag(np.sqrt(vals[[-1, -2, -3]]))) np.save(args.output + "_xtrain.npy", x) test = resqu[not_si].copy() t = time.time() ga = geom.cdistq(test, samp, np.zeros((test.shape[0], samp.shape[0]), dtype=np.single)) log.info("Test pairwise distances calculated in %0.3f s" % (time.time() - t)) ga = geom.double_center(ga, reference=d, inplace=True) xa = ga.dot(x) / vals[[-1, -2, -3]].reshape(1, 3) np.save(args.output + "_xtest.npy", xa) vol, hdr = mrc.read(args.volume, inc_header=True) psz = hdr["xlen"] / hdr["nx"] for pc in range(2): keyq = geom.findkeyq(test, xa, nkey=10, pc_cyl_ptile=args.outlier_radius, pc_ptile=args.outlier_length, pc=pc) keyq_exp = geom.qslerp_mult_balanced(keyq, 10) volbase = os.path.basename( args.volume).rstrip(".mrc") + "_kpc%d" % pc + "_%.4d.mrc" util.write_q_series(vol, keyq_exp, os.path.join(outdir, volbase), psz=psz, order=args.spline_order) return 0
def main(args): log = logging.getLogger(__name__) log.setLevel(logging.INFO) hdlr = logging.StreamHandler(sys.stdout) if args.quiet: hdlr.setLevel(logging.ERROR) elif args.verbose: hdlr.setLevel(logging.INFO) else: hdlr.setLevel(logging.WARN) log.addHandler(hdlr) data, hdr = read(args.input, inc_header=True) final = None box = np.array([hdr[a] for a in ["nx", "ny", "nz"]]) center = box // 2 if args.fft: data_ft = vop.vol_ft(data.T, threads=args.threads) np.save(args.output, data_ft) return 0 if args.transpose is not None: try: tax = [np.int64(a) for a in args.transpose.split(",")] data = np.transpose(data, axes=tax) except: log.error( "Transpose axes must be comma-separated list of three integers" ) return 1 if args.normalize: if args.reference is not None: ref, refhdr = read(args.reference, inc_header=True) final, mu, sigma = vop.normalize(data, ref=ref, return_stats=True) else: final, mu, sigma = vop.normalize(data, return_stats=True) final = (data - mu) / sigma if args.verbose: log.info("Mean: %f, Standard deviation: %f" % (mu, sigma)) if args.apix is None: args.apix = hdr["xlen"] / hdr["nx"] log.info("Using computed pixel size of %f Angstroms" % args.apix) if args.target and args.matrix: log.warn( "Target pose transformation will be applied after explicit matrix") if args.euler is not None and (args.target is not None or args.matrix is not None): log.warn( "Euler transformation will be applied after target pose transformation" ) if args.translate is not None and (args.euler is not None or args.target is not None or args.matrix is not None): log.warn("Translation will be applied after other transformations") if args.origin is not None: try: args.origin = np.array( [np.double(tok) for tok in args.origin.split(",")]) / args.apix assert np.all(args.origin < box) except: log.error( "Origin must be comma-separated list of x,y,z coordinates and lie within the box" ) return 1 else: args.origin = center log.info("Origin set to box center, %s" % (args.origin * args.apix)) if not (args.target is None and args.euler is None and args.matrix is None and args.boxsize is None) \ and vop.ismask(data) and args.spline_order != 0: log.warn( "Input looks like a mask, --spline-order 0 (nearest neighbor) is recommended" ) if args.matrix is not None: try: r = np.array(json.loads(args.matrix)) except: log.error("Matrix format is incorrect") return 1 data = vop.resample_volume(data, r=r, t=None, ori=None, order=args.spline_order) if args.target is not None: try: args.target = np.array( [np.double(tok) for tok in args.target.split(",")]) / args.apix except: log.error( "Standard pose target must be comma-separated list of x,y,z coordinates" ) return 1 args.target -= args.origin args.target = np.where(np.abs(args.target) < 1, 0, args.target) ori = None if args.origin is center else args.origin - args.center r = vec2rot(args.target) t = np.linalg.norm(args.target) log.info("Euler angles are %s deg and shift is %f px" % (np.rad2deg(rot2euler(r)), t)) data = vop.resample_volume(data, r=r, t=args.target, ori=ori, order=args.spline_order, invert=args.target_invert) if args.euler is not None: try: args.euler = np.deg2rad( np.array([np.double(tok) for tok in args.euler.split(",")])) except: log.error( "Eulers must be comma-separated list of phi,theta,psi angles") return 1 r = euler2rot(*args.euler) offset = args.origin - 0.5 offset = offset - r.T.dot(offset) data = affine_transform(data, r.T, offset=offset, order=args.spline_order) if args.translate is not None: try: args.translate = np.array( [np.double(tok) for tok in args.translate.split(",")]) / args.apix except: log.error( "Translation vector must be comma-separated list of x,y,z coordinates" ) return 1 args.translate -= args.origin data = shift(data, -args.translate, order=args.spline_order) if args.boxsize is not None: args.boxsize = np.double(args.boxsize) data = zoom(data, args.boxsize / box, order=args.spline_order) args.apix = args.apix * box[0] / args.boxsize if final is None: final = data if args.final_mask is not None: final_mask = read(args.final_mask) final *= final_mask write(args.output, final, psz=args.apix) return 0
def main(args): """ Projection subtraction program entry point. :param args: Command-line arguments parsed by ArgumentParser.parse_args() :return: Exit status """ log = logging.getLogger('root') hdlr = logging.StreamHandler(sys.stdout) log.addHandler(hdlr) log.setLevel(logging.getLevelName(args.loglevel.upper())) if args.dest is None and args.suffix == "": args.dest = "" args.suffix = "_subtracted" log.info("Reading particle .star file") df = star.parse_star(args.input, keep_index=False) star.augment_star_ucsf(df) if not args.original: df[star.UCSF.IMAGE_ORIGINAL_PATH] = df[star.UCSF.IMAGE_PATH] df[star.UCSF.IMAGE_ORIGINAL_INDEX] = df[star.UCSF.IMAGE_INDEX] df.sort_values(star.UCSF.IMAGE_ORIGINAL_PATH, inplace=True, kind="mergesort") gb = df.groupby(star.UCSF.IMAGE_ORIGINAL_PATH) df[star.UCSF.IMAGE_INDEX] = gb.cumcount() df[star.UCSF.IMAGE_PATH] = df[star.UCSF.IMAGE_ORIGINAL_PATH].map( lambda x: os.path.join( args.dest, args.prefix + os.path.basename(x).replace( ".mrcs", args.suffix + ".mrcs"))) if args.submap_ft is None: log.info("Reading volume") submap = mrc.read(args.submap, inc_header=False, compat="relion") if args.submask is not None: log.info("Masking volume") submask = mrc.read(args.submask, inc_header=False, compat="relion") submap *= submask log.info("Preparing 3D FFT of volume") submap_ft = vop.vol_ft(submap, pfac=args.pfac, threads=min(args.threads, cpu_count())) log.info("Finished 3D FFT of volume") else: log.info("Loading 3D FFT from %s" % args.submap_ft) submap_ft = np.load(args.submap_ft) log.info("Loaded 3D FFT from %s" % args.submap_ft) sz = (submap_ft.shape[0] - 3) // args.pfac maxshift = np.round(np.max(np.abs(df[star.Relion.ORIGINS].values))) if args.crop is not None and sz < 2 * maxshift + args.crop: log.error("Some shifts are too large to crop (maximum crop is %d)" % (sz - 2 * maxshift)) return 1 sx, sy = np.meshgrid(np.fft.rfftfreq(sz), np.fft.fftfreq(sz)) s = np.sqrt(sx**2 + sy**2) r = s * sz r = np.round(r).astype(np.int64) r[r > sz // 2] = sz // 2 + 1 nr = np.max(r) + 1 a = np.arctan2(sy, sx) if args.refmap is not None: coefs_method = 1 if args.refmap_ft is None: refmap = mrc.read(args.refmap, inc_header=False, compat="relion") refmap_ft = vop.vol_ft(refmap, pfac=args.pfac, threads=min(args.threads, cpu_count())) else: log.info("Loading 3D FFT from %s" % args.refmap_ft) refmap_ft = np.load(args.refmap_ft) log.info("Loaded 3D FFT from %s" % args.refmap_ft) else: coefs_method = 0 refmap_ft = np.empty(submap_ft.shape, dtype=submap_ft.dtype) apix = star.calculate_apix(df) log.info("Computed pixel size is %f A" % apix) log.debug("Grouping particles by output stack") gb = df.groupby(star.UCSF.IMAGE_PATH) iothreads = threading.BoundedSemaphore(args.io_thread_pairs) qsize = args.io_queue_length fftthreads = args.fft_threads def init(): global tls tls = threading.local() log.info("Instantiating thread pool with %d workers" % args.threads) pool = Pool(processes=args.threads, initializer=init) threads = [] log.info("Performing projection subtraction") try: for fname, particles in gb: log.debug("Instantiating queue") queue = Queue.Queue(maxsize=qsize) log.debug("Create producer for %s" % fname) prod = threading.Thread(target=producer, args=(pool, queue, submap_ft, refmap_ft, fname, particles, sx, sy, s, a, apix, coefs_method, r, nr, fftthreads, args.crop, args.pfac)) log.debug("Create consumer for %s" % fname) cons = threading.Thread(target=consumer, args=(queue, fname, apix, iothreads)) threads.append((prod, cons)) iothreads.acquire() log.debug("iotheads at %d" % iothreads._Semaphore__value) log.debug("Start consumer for %s" % fname) cons.start() log.debug("Start producer for %s" % fname) prod.start() except KeyboardInterrupt: log.debug("Main thread wants out!") for pair in threads: for thread in pair: try: thread.join() except RuntimeError as e: log.debug(e) pool.close() pool.join() pool.terminate() log.info("Finished projection subtraction") log.info("Writing output .star file") if args.crop is not None: df = star.recenter(df, inplace=True) star.simplify_star_ucsf(df) star.write_star(args.output, 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())) 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
def main(args): """ Projection subtraction program entry point. :param args: Command-line arguments parsed by ArgumentParser.parse_args() :return: Exit status """ log = logging.getLogger('root') hdlr = logging.StreamHandler(sys.stdout) log.addHandler(hdlr) log.setLevel(logging.getLevelName(args.loglevel.upper())) log.debug("Reading particle .star file") df = parse_star(args.input, keep_index=False) df.reset_index(inplace=True) df["rlnImageOriginalName"] = df["rlnImageName"] df["ucsfOriginalParticleIndex"], df["ucsfOriginalImagePath"] = \ df["rlnImageOriginalName"].str.split("@").str df["ucsfOriginalParticleIndex"] = pd.to_numeric( df["ucsfOriginalParticleIndex"]) df.sort_values("rlnImageOriginalName", inplace=True, kind="mergesort") gb = df.groupby("ucsfOriginalImagePath") df["ucsfParticleIndex"] = gb.cumcount() + 1 df["ucsfImagePath"] = df["ucsfOriginalImagePath"].map( lambda x: os.path.join( args.dest, args.prefix + os.path.basename(x).replace( ".mrcs", args.suffix + ".mrcs"))) df["rlnImageName"] = df["ucsfParticleIndex"].map( lambda x: "%.6d" % x).str.cat(df["ucsfImagePath"], sep="@") log.debug("Read particle .star file") if args.submap_ft is None: submap = mrc.read(args.submap, inc_header=False, compat="relion") submap_ft = vol_ft(submap, threads=min(args.threads, cpu_count())) else: log.debug("Loading %s" % args.submap_ft) submap_ft = np.load(args.submap_ft) log.debug("Loaded %s" % args.submap_ft) sz = submap_ft.shape[0] // 2 - 1 sx, sy = np.meshgrid(np.fft.rfftfreq(sz), np.fft.fftfreq(sz)) s = np.sqrt(sx**2 + sy**2) r = s * sz r = np.round(r).astype(np.int64) r[r > sz // 2] = sz // 2 + 1 nr = np.max(r) + 1 a = np.arctan2(sy, sx) if args.refmap is not None: coefs_method = 1 if args.refmap_ft is None: refmap = mrc.read(args.refmap, inc_header=False, compat="relion") refmap_ft = vol_ft(refmap, threads=min(args.threads, cpu_count())) else: log.debug("Loading %s" % args.refmap_ft) refmap_ft = np.load(args.refmap_ft) log.debug("Loaded %s" % args.refmap_ft) else: coefs_method = 0 refmap_ft = np.empty(submap_ft.shape, dtype=submap_ft.dtype) apix = calculate_apix(df) log.debug("Constructing particle metadata references") # npart = df.shape[0] idx = df["ucsfOriginalParticleIndex"].values stack = df["ucsfOriginalImagePath"].values.astype(np.str, copy=False) def1 = df["rlnDefocusU"].values def2 = df["rlnDefocusV"].values angast = df["rlnDefocusAngle"].values phase = df["rlnPhaseShift"].values kv = df["rlnVoltage"].values ac = df["rlnAmplitudeContrast"].values cs = df["rlnSphericalAberration"].values az = df["rlnAngleRot"].values el = df["rlnAngleTilt"].values sk = df["rlnAnglePsi"].values xshift = df["rlnOriginX"].values yshift = df["rlnOriginY"].values new_idx = df["ucsfParticleIndex"].values new_stack = df["ucsfImagePath"].values.astype(np.str, copy=False) log.debug("Grouping particles by output stack") gb = df.groupby("ucsfImagePath") iothreads = threading.BoundedSemaphore(args.io_thread_pairs) qsize = args.io_queue_length fftthreads = args.fft_threads # pyfftw.interfaces.cache.enable() log.debug("Instantiating worker pool") pool = Pool(processes=args.threads) threads = [] try: for fname, particles in gb.indices.iteritems(): log.debug("Instantiating queue") queue = Queue.Queue(maxsize=qsize) log.debug("Create producer for %s" % fname) prod = threading.Thread( target=producer, args=(pool, queue, submap_ft, refmap_ft, fname, particles, idx, stack, sx, sy, s, a, apix, def1, def2, angast, phase, kv, ac, cs, az, el, sk, xshift, yshift, new_idx, new_stack, coefs_method, r, nr, fftthreads)) log.debug("Create consumer for %s" % fname) cons = threading.Thread(target=consumer, args=(queue, fname, apix, fftthreads, iothreads)) threads.append((prod, cons)) iothreads.acquire() log.debug("iotheads at %d" % iothreads._Semaphore__value) log.debug("Start consumer for %s" % fname) cons.start() log.debug("Start producer for %s" % fname) prod.start() except KeyboardInterrupt: log.debug("Main thread wants out!") for pair in threads: for thread in pair: try: thread.join() except RuntimeError as e: log.debug(e) pool.close() pool.join() pool.terminate() df.drop([c for c in df.columns if "ucsf" in c or "eman" in c], axis=1, inplace=True) df.set_index("index", inplace=True) df.sort_index(inplace=True, kind="mergesort") write_star(args.output, df, reindex=True) return 0