ls_mlkit.util.manifold.so3_utils module

SO3 Utils

ls_mlkit.util.manifold.so3_utils.calculate_igso3(*, num_sigma: int, num_omega: int, min_sigma: float, max_sigma: float, discrete_omega=None, discrete_sigma=None) → dict[str, Tensor]

calculate_igso3 pre-computes numerical approximations to the IGSO3 cdfs and score norms and expected squared score norms.

Parameters:

• num_sigma – number of different sigmas for which to compute igso3 quantities.

• num_omega – number of point in the discretization in the angle of rotation.

• min_sigma – the upper and lower ranges for the angle of rotation on which to consider the IGSO3 distribution. This cannot be too low or it will create numerical instability.

• max_sigma – the upper and lower ranges for the angle of rotation on which to consider the IGSO3 distribution. This cannot be too low or it will create numerical instability.

ls_mlkit.util.manifold.so3_utils.d_logf_d_omega(omega: Tensor, c: Tensor, L: int = 2000) → Tensor

Score function of IGSO(3) distribution.

$$ frac{d}{domega} log f(omega, c, L) $$ :param omega: (…,), i.e. the angle of rotation associated with rotation matrix :param c: (1,), variance parameter of IGSO(3), maps onto time in Brownian motion :param L: (1,)Truncation level

Returns:

(…, )

ls_mlkit.util.manifold.so3_utils.exponential_map(A: Tensor) → Tensor

Exponential map from vector space of $$mathfrak{so}(3)$$ to SO(3), this is the matrix $$mathfrak{so}(3) mapsto SO(3)$$ $$ begin{align*} theta &= sqrt{A_{0,1}^2 + A_{0,2}^2 + A_{1,2}^2}\ B &= A / theta\ exp(A) &= I + sin(theta) B + (1 - cos(theta)) B^2 end{align*} $$ :param A.shape: (…, 3, 3)

Returns:

(…, 3, 3)

Return type:

shape

ls_mlkit.util.manifold.so3_utils.f_igso3(omega: Tensor, c: Tensor, L: int = 2000) → Tensor

Truncated sum of IGSO(3) distribution.

This function approximates the power series in equation 5 of “DENOISING DIFFUSION PROBABILISTIC MODELS ON SO(3) FOR ROTATIONAL ALIGNMENT” Leach et al. 2022

This expression diverges from the expression in Leach in that here, $$sigma = sqrt{2} * epsilon$$, if eps_leach were the scale parameter of the IGSO(3).

With this reparameterization, IGSO(3) agrees with the Brownian motion on SO(3) with t=sigma^2 when defined for the canonical inner product on SO3, $$<u, v>_{SO3} = Tr(u v^T)/2$$

Parameters:

• omega – (…,), i.e. the angle of rotation associated with rotation matrix

• c – (1,), variance parameter of IGSO(3), maps onto time in Brownian motion

• L – (1,)Truncation level

Returns:

$$ sum_{i=0}^{L-1} (2l+1) e^{-l(l+1)c/2} sin (omega (l +1/2)) / sin(omega/2) $$

ls_mlkit.util.manifold.so3_utils.flatten_batch_dimension(x: Tensor, ndim_micro_shape: int) → Tensor

ls_mlkit.util.manifold.so3_utils.get_macro_shape(x: Tensor, ndim_micro_shape: int) → Tuple[int, ...]

ls_mlkit.util.manifold.so3_utils.igso3_density(Rt: Tensor, c: Tensor, L: int = 2000) → Tensor

IGSO3 density with respect to the volume form on SO(3) :param Rt: (…, 3, 3), rotation matrix :param c: (1,), variance parameter of IGSO(3), maps onto time in Brownian motion :param L: (1,)Truncation level

Returns:

(…, )

ls_mlkit.util.manifold.so3_utils.igso3_density_angle(omega: Tensor, c: Tensor, L: int = 2000) → Tensor

$$((1-cos(omega)) / pi ) f$$ :param omega: (…,), i.e. the angle of rotation associated with rotation matrix :param c: (1,), variance parameter of IGSO(3), maps onto time in Brownian motion :param L: (1,)Truncation level

Returns:

(…, )

ls_mlkit.util.manifold.so3_utils.igso3_score(R: Tensor, c: Tensor, L: int = 2000) → Tensor

grad_R log IGSO3(R; I_3, c) $$ nabla_R log IG_{SO3}(R; I_3, c) = R frac{log(R)}{omega(R)} frac{d}{domega} log f(omega, c, L) $$

Parameters:

• R – (…, 3, 3), rotation matrix

• c – (1,), variance parameter of IGSO(3), maps onto time in Brownian motion

• L – (1,)Truncation level

Returns:

(…, 3, 3)

ls_mlkit.util.manifold.so3_utils.inverse_transform_sampling(shape: Tuple[int, ...], cdf: Tensor, discrete_omega: Tensor) → Tensor

Sample uses the inverse cdf to sample an angle of rotation from IGSO(3)

Parameters:

• shape – shape of the sampled angles of rotation.

• cdf – (num_omega,), cdf of the IGSO(3) distribution

• discrete_omega – (num_omega, ), discrete angles of rotation

Returns:

sampled angles of rotation. (*shape,)

ls_mlkit.util.manifold.so3_utils.logarithmic_map(R: Tensor) → Tensor

Logarithmic map from SO(3) to so(3), this is the matrix logarithm $$SO(3) mapsto mathfrak{so}(3)$$ $$ begin{align*} theta &= arccos(frac{Tr(R)-1}{2})\ log(R) &= frac{theta}{2sin(theta)} (R - R^T) end{align*} $$ :param R.shape: (…, 3, 3)

Returns:

(…, 3, 3)

Return type:

shape

ls_mlkit.util.manifold.so3_utils.rotation_matrix_to_angle(R: Tensor) → Tensor

$$ theta = arccos(frac{Tr(R)-1}{2}) $$ :param R.shape: (…, 3, 3)

Returns:

(…, )

Return type:

shape

ls_mlkit.util.manifold.so3_utils.skew_symmetric_to_angle(A: Tensor) → Tensor

$$ begin{pmatrix} 0 & -z & y \ z & 0 & -x \ -y & x & 0 end{pmatrix} to sqrt{x^2 + y^2 + z^2} $$ :param A.shape: (…, 3, 3)

Returns:

(…, )

Return type:

shape

ls_mlkit.util.manifold.so3_utils.skew_symmetric_to_vector(hat_v: Tensor) → Tensor

Map from skew-symmetric matrix to vector $$mathfrak{so}(3) mapsto mathbb{R}^3$$ $$ begin{pmatrix} 0 & -z & y \ z & 0 & -x \ -y & x & 0 end{pmatrix} to (x,y,z) $$ :param hat_v.shape: (…, 3, 3)

Returns:

(…, 3)

Return type:

shape

ls_mlkit.util.manifold.so3_utils.trace(A: Tensor) → Tensor

Parameters:

A.shape – (…, 3, 3)

Returns:

(…, )

Return type:

shape

ls_mlkit.util.manifold.so3_utils.unit_skew_symmetric(A: Tensor) → Tensor

get the unit skew-symmetric matrix :param A.shape: (…, 3, 3)

Returns:

(…, 3, 3)

Return type:

shape

ls_mlkit.util.manifold.so3_utils.vector_to_skew_symmetric(v: Tensor) → Tensor

Hat map from vector space $$mathbb{R}^3$$ to Lie algebra $$mathfrak{so}(3)$$ $$ (x,y,z) to begin{pmatrix} 0 & -z & y \ z & 0 & -x \ -y & x & 0 end{pmatrix} $$ :param v.shape: (…, 3)

Returns:

(…, 3, 3)

Return type:

shape