Example #1
0
File: star.py Project: dzyla/pyem
def transform_star(df,
                   r,
                   t=None,
                   inplace=False,
                   rots=None,
                   invert=False,
                   rotate=True,
                   adjust_defocus=False):
    """
    Transform particle angles and origins according to a rotation
    matrix (in radians) and an optional translation vector.
    The translation may also be given as the 4th column of a 3x4 matrix,
    or as a scalar distance to be applied along the axis of rotation.
    """
    assert (r.shape[0] == 3)
    if r.shape[1] == 4 and t is None:
        t = r[:, -1]
        r = r[:, :3]
    assert (r.shape == (3, 3))
    assert t is None or np.isscalar(t) or len(t) == 3

    if inplace:
        newstar = df
    else:
        newstar = df.copy()

    if rots is None:
        rots = [
            euler2rot(*np.deg2rad(row[1]))
            for row in df[Relion.ANGLES].iterrows()
        ]

    if invert:
        r = r.T

    newrots = [ptcl.dot(r) for ptcl in rots]
    if rotate:
        angles = [np.rad2deg(rot2euler(q)) for q in newrots]
        newstar[Relion.ANGLES] = angles

    if t is not None and np.linalg.norm(t) > 0:
        if np.isscalar(t):
            if invert:
                tt = -np.vstack([np.squeeze(q[:, 2]) * t for q in rots])
            else:
                tt = np.vstack([np.squeeze(q[:, 2]) * t for q in newrots])
        else:
            if invert:
                tt = -np.vstack([q.dot(t) for q in rots])
            else:
                tt = np.vstack([q.dot(t) for q in newrots])
        newshifts = df[Relion.ORIGINS] + tt[:, :-1]
        newstar[Relion.ORIGINS] = newshifts
        if adjust_defocus:
            newstar[Relion.DEFOCUSU] += tt[:, -1] * calculate_apix(df)
            newstar[Relion.DEFOCUSV] += tt[:, -1] * calculate_apix(df)
            newstar[Relion.DEFOCUSANGLE] = np.rad2deg(
                np.arctan2(newstar[Relion.DEFOCUSV], newstar[Relion.DEFOCUSV]))

    return newstar
Example #2
0
def main(args):
    meta = parse_metadata(args.input)  # Read cryosparc metadata file.
    meta["data_input_idx"] = [
        "%.6d" % (i + 1) for i in meta["data_input_idx"]
    ]  # Reformat particle idx for Relion.

    if "data_input_relpath" not in meta.columns:
        if args.data_path is None:
            print(
                "Data path missing, use --data-path to specify particle stack path"
            )
            return 1
        meta["data_input_relpath"] = args.data_path

    meta["data_input_relpath"] = meta["data_input_idx"].str.cat(
        meta["data_input_relpath"], sep="@")  # Construct _rlnImageName field.
    # Take care of trivial mappings.
    rlnheaders = [
        general[h] for h in meta.columns
        if h in general and general[h] is not None
    ]
    star = meta[[
        h for h in meta.columns if h in general and general[h] is not None
    ]].copy()
    star.columns = rlnheaders

    if "rlnRandomSubset" in star.columns:
        star["rlnRandomSubset"] = star["rlnRandomSubset"].apply(
            lambda x: ord(x) - 64)

    if "rlnPhaseShift" in star.columns:
        star["rlnPhaseShift"] = np.rad2deg(star["rlnPhaseShift"])

    # general class assignments and other model parameters.
    phic = meta[[h for h in meta.columns if "phiC" in h
                 ]]  # Posterior probability over class assignments.
    if len(phic.columns) > 0:  # Check class assignments exist in input.
        # phic.columns = [int(h[21]) for h in meta.columns if "phiC" in h]
        phic.columns = range(len(phic.columns))
        cls = phic.idxmax(axis=1)
        for p in model:
            if model[p] is not None:
                pspec = p.split("model")[1]
                param = meta[[h for h in meta.columns if pspec in h]]
                if len(param.columns) > 0:
                    param.columns = phic.columns
                    star[model[p]] = param.lookup(param.index, cls)
        star[
            "rlnClassNumber"] = cls + 1  # Compute most probable classes and add one for Relion indexing.
    else:
        for p in model:
            if model[p] is not None and p in meta.columns:
                star[model[p]] = meta[p]
        star["rlnClassNumber"] = 1

    if args.cls is not None:
        star = select_classes(star, args.cls)

    # Convert axis-angle representation to Euler angles (degrees).
    if star.columns.intersection(angles).size == len(angles):
        star[angles] = np.rad2deg(star[angles].apply(
            lambda x: rot2euler(expmap(x)), axis=1, raw=True, broadcast=True))

    if args.minphic is not None:
        mask = np.all(phic < args.minphic, axis=1)
        if args.drop_bad:
            star.drop(star[mask].index,
                      inplace=True)  # Delete low-confidence particles.
        else:
            star.loc[
                mask,
                "rlnClassNumber"] = 0  # Set low-confidence particles to dummy class.

    if args.transform is not None:
        r = np.array(json.loads(args.transform))
        star = transform_star(star, r, inplace=True)

    # Write Relion .star file with correct headers.
    write_star(args.output, star, reindex=True)
    return 0
Example #3
0
File: map.py Project: 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
Example #4
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
Example #5
0
 def test_rot2euler(self):
     e1test = util.rot2euler(r1)
     self.assertTrue(np.array_equal(e1, e1test))
     e2test = util.rot2euler(r3)
     self.assertTrue(np.allclose(e3, e2test))
Example #6
0
def main(args):
    meta = parse_metadata(args.input)  # Read cryosparc metadata file.
    meta["data_input_idx"] = [
        "%.6d" % (i + 1) for i in meta["data_input_idx"]
    ]  # Reformat particle idx for Relion.

    if "data_input_relpath" not in meta.columns:
        if args.data_path is None:
            print(
                "Data path missing, use --data-path to specify particle stack path"
            )
            return 1
        meta["data_input_relpath"] = args.data_path

    meta["data_input_relpath"] = meta["data_input_idx"].str.cat(
        meta["data_input_relpath"], sep="@")  # Construct rlnImageName field.
    # Take care of trivial mappings.
    rlnheaders = [
        general[h] for h in meta.columns
        if h in general and general[h] is not None
    ]
    df = meta[[
        h for h in meta.columns if h in general and general[h] is not None
    ]].copy()
    df.columns = rlnheaders

    if "rlnRandomSubset" in df.columns:
        df["rlnRandomSubset"] = df["rlnRandomSubset"].apply(
            lambda x: ord(x) - 64)

    if "rlnPhaseShift" in df.columns:
        df["rlnPhaseShift"] = np.rad2deg(df["rlnPhaseShift"])

    # Class assignments and other model parameters.
    phic = meta[[h for h in meta.columns if "phiC" in h
                 ]]  # Posterior probability over class assignments.
    if len(phic.columns) > 0:  # Check class assignments exist in input.
        # phic.columns = [int(h[21]) for h in meta.columns if "phiC" in h]
        phic.columns = range(len(phic.columns))
        cls = phic.idxmax(axis=1)
        for p in model:
            if model[p] is not None:
                pspec = p.split("model")[1]
                param = meta[[h for h in meta.columns if pspec in h]]
                if len(param.columns) > 0:
                    param.columns = phic.columns
                    df[model[p]] = param.lookup(param.index, cls)
        df["rlnClassNumber"] = cls + 1  # Add one for Relion indexing.
    else:
        for p in model:
            if model[p] is not None and p in meta.columns:
                df[model[p]] = meta[p]
        df["rlnClassNumber"] = 1

    if args.cls is not None:
        df = star.select_classes(df, args.cls)

    # Convert axis-angle representation to Euler angles (degrees).
    if df.columns.intersection(star.Relion.ANGLES).size == len(
            star.Relion.ANGLES):
        df[star.Relion.ANGLES] = np.rad2deg(df[star.Relion.ANGLES].apply(
            lambda x: rot2euler(expmap(x)), axis=1, raw=True, broadcast=True))

    if args.minphic is not None:
        mask = np.all(phic < args.minphic, axis=1)
        if args.keep_bad:
            df.loc[mask, "rlnClassNumber"] = 0
        else:
            df.drop(df[mask].index, inplace=True)

    if args.copy_micrograph_coordinates is not None:
        coord_star = pd.concat(
            (star.parse_star(inp, keep_index=False)
             for inp in glob(args.copy_micrograph_coordinates)),
            join="inner")
        df = star.smart_merge(df,
                              coord_star,
                              fields=star.Relion.MICROGRAPH_COORDS)

    if args.transform is not None:
        r = np.array(json.loads(args.transform))
        df = star.transform_star(df, r, inplace=True)

    # Write Relion .star file with correct headers.
    star.write_star(args.output, df, reindex=True)
    return 0