Ejemplo n.º 1
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 = 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
Ejemplo n.º 2
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
Ejemplo n.º 3
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
Ejemplo n.º 4
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
Ejemplo n.º 5
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
Ejemplo n.º 6
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
Ejemplo n.º 7
0
Archivo: sort.py Proyecto: wwangat/pyem
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
Ejemplo n.º 8
0
Archivo: map.py Proyecto: Error-Y/pyem
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
Ejemplo n.º 9
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
Ejemplo n.º 10
0
Archivo: pose.py Proyecto: wwangat/pyem
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
Ejemplo n.º 11
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
Ejemplo n.º 12
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
Ejemplo n.º 13
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
Ejemplo n.º 14
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