Rotations in cryo-EM

cryojax.rotations implements an engine and utility functions for rotations in cryo-EM.

Representing rotations

The engine for handling rotations in cryoJAX is the cryojax.rotations.SO3 class. This is based on the implementation in the package jaxlie.

cryojax.rotations.SO3

A rotation in 3D space, represented by the SO3 matrix lie group.

The class is almost exactly derived from jaxlie.SO3.

jaxlie was written for Yi, Brent, et al. 2021.

__init__(wxyz: Float[NDArrayLike, '... 4'] | Sequence[float])

Arguments:

• wxyz: A quaternion represented as \((q_w, q_x, q_y, q_z)\). A leading batch dimension is supported, e.g. of shape (..., 4).

apply(target: Float[Array, '... 3']) -> Float[Array, '... 3']

compose(other: typing.Self) -> typing.Self

inverse() -> typing.Self

from_x_radians(angle: cryojax.jax_util.FloatLike) -> typing.Self classmethod

Generates a x-axis rotation.

from_y_radians(angle: cryojax.jax_util.FloatLike) -> typing.Self classmethod

Generates a y-axis rotation.

from_z_radians(angle: cryojax.jax_util.FloatLike) -> typing.Self classmethod

Generates a z-axis rotation.

identity() -> typing.Self classmethod

from_matrix(matrix: Float[Array, '... 3 3']) -> typing.Self classmethod

as_matrix() -> Float[Array, '... 3 3']

exp(tangent: Float[Array, '... 3']) -> typing.Self classmethod

log() -> Float[Array, '... 3']

normalize() -> typing.Self

sample_uniform(key: PRNGKeyArray) -> typing.Self classmethod

cryojax.rotations.SO2

A rotation in 2D space, represented by the SO2 matrix lie group.

The class is almost exactly derived from jaxlie.SO2.

jaxlie was written for Yi, Brent, et al. 2021.

__init__(unit_complex: Float[NDArrayLike, '... 2'] | Sequence[float | cryojax.jax_util.NDArrayLike])

Arguments:

• unit_complex: A complex number represented as a 2-vector \((cos \theta, sin \theta)\). A leading batch dimension is supported, e.g. of shape (..., 2).

apply(target: Float[Array, '... 2']) -> Float[Array, '... 2']

compose(other: typing.Self) -> typing.Self

inverse() -> typing.Self

from_radians(theta: cryojax.jax_util.FloatLike) -> typing.Self classmethod

Construct a rotation object from a scalar angle.

as_radians() -> Float[Array, '']

Compute a scalar angle from a rotation object.

identity() -> typing.Self classmethod

from_matrix(matrix: Float[NDArrayLike, '... 2 2']) -> typing.Self classmethod

as_matrix() -> Float[Array, '... 2 2']

exp(tangent: cryojax.jax_util.FloatLike) -> typing.Self classmethod

log() -> Float[Array, '...']

normalize() -> typing.Self

sample_uniform(key: PRNGKeyArray) -> typing.Self classmethod

Utilities for converting between angular parametrizations