Parametrization of Rotations

Matrix Parametrization

e3nn.o3.rand_matrix(*shape, requires_grad=False, dtype=None, device=None)[source]

random rotation matrix

Parameters

*shape (int) –

Returns

tensor of shape \((\mathrm{shape}, 3, 3)\)

Return type

torch.Tensor

e3nn.o3.matrix_x(angle: Tensor) Tensor[source]

matrix of rotation around X axis

Parameters

angle (torch.Tensor) – tensor of any shape \((...)\)

Returns

matrices of shape \((..., 3, 3)\)

Return type

torch.Tensor

e3nn.o3.matrix_y(angle: Tensor) Tensor[source]

matrix of rotation around Y axis

Parameters

angle (torch.Tensor) – tensor of any shape \((...)\)

Returns

matrices of shape \((..., 3, 3)\)

Return type

torch.Tensor

e3nn.o3.matrix_z(angle: Tensor) Tensor[source]

matrix of rotation around Z axis

Parameters

angle (torch.Tensor) – tensor of any shape \((...)\)

Returns

matrices of shape \((..., 3, 3)\)

Return type

torch.Tensor

Euler Angles Parametrization

e3nn.o3.identity_angles(*shape, requires_grad=False, dtype=None, device=None)[source]

angles of the identity rotation

Parameters

*shape (int) –

Returns

  • alpha (torch.Tensor) – tensor of shape \((\mathrm{shape})\)

  • beta (torch.Tensor) – tensor of shape \((\mathrm{shape})\)

  • gamma (torch.Tensor) – tensor of shape \((\mathrm{shape})\)

e3nn.o3.rand_angles(*shape, requires_grad=False, dtype=None, device=None)[source]

random rotation angles

Parameters

*shape (int) –

Returns

  • alpha (torch.Tensor) – tensor of shape \((\mathrm{shape})\)

  • beta (torch.Tensor) – tensor of shape \((\mathrm{shape})\)

  • gamma (torch.Tensor) – tensor of shape \((\mathrm{shape})\)

e3nn.o3.compose_angles(a1, b1, c1, a2, b2, c2)[source]

compose angles

Computes \((a, b, c)\) such that \(R(a, b, c) = R(a_1, b_1, c_1) \circ R(a_2, b_2, c_2)\)

Parameters
  • a1 (torch.Tensor) – tensor of shape \((...)\), (applied second)

  • b1 (torch.Tensor) – tensor of shape \((...)\), (applied second)

  • c1 (torch.Tensor) – tensor of shape \((...)\), (applied second)

  • a2 (torch.Tensor) – tensor of shape \((...)\), (applied first)

  • b2 (torch.Tensor) – tensor of shape \((...)\), (applied first)

  • c2 (torch.Tensor) – tensor of shape \((...)\), (applied first)

Returns

e3nn.o3.inverse_angles(a, b, c)[source]

angles of the inverse rotation

Parameters
Returns

Quaternion Parametrization

e3nn.o3.identity_quaternion(*shape, requires_grad=False, dtype=None, device=None)[source]

quaternion of identity rotation

Parameters

*shape (int) –

Returns

tensor of shape \((\mathrm{shape}, 4)\)

Return type

torch.Tensor

e3nn.o3.rand_quaternion(*shape, requires_grad=False, dtype=None, device=None)[source]

generate random quaternion

Parameters

*shape (int) –

Returns

tensor of shape \((\mathrm{shape}, 4)\)

Return type

torch.Tensor

e3nn.o3.compose_quaternion(q1, q2)[source]

compose two quaternions: \(q_1 \circ q_2\)

Parameters
  • q1 (torch.Tensor) – tensor of shape \((..., 4)\), (applied second)

  • q2 (torch.Tensor) – tensor of shape \((..., 4)\), (applied first)

Returns

tensor of shape \((..., 4)\)

Return type

torch.Tensor

e3nn.o3.inverse_quaternion(q)[source]

inverse of a quaternion

Works only for unit quaternions.

Parameters

q (torch.Tensor) – tensor of shape \((..., 4)\)

Returns

tensor of shape \((..., 4)\)

Return type

torch.Tensor

Axis-Angle Parametrization

e3nn.o3.rand_axis_angle(*shape, requires_grad=False, dtype=None, device=None)[source]

generate random rotation as axis-angle

Parameters

*shape (int) –

Returns

  • axis (torch.Tensor) – tensor of shape \((\mathrm{shape}, 3)\)

  • angle (torch.Tensor) – tensor of shape \((\mathrm{shape})\)

e3nn.o3.compose_axis_angle(axis1, angle1, axis2, angle2)[source]

compose \((\vec x_1, \alpha_1)\) with \((\vec x_2, \alpha_2)\)

Parameters
  • axis1 (torch.Tensor) – tensor of shape \((..., 3)\), (applied second)

  • angle1 (torch.Tensor) – tensor of shape \((...)\), (applied second)

  • axis2 (torch.Tensor) – tensor of shape \((..., 3)\), (applied first)

  • angle2 (torch.Tensor) – tensor of shape \((...)\), (applied first)

Returns

Convertions

e3nn.o3.angles_to_matrix(alpha, beta, gamma)[source]

conversion from angles to matrix

Parameters
Returns

matrices of shape \((..., 3, 3)\)

Return type

torch.Tensor

e3nn.o3.matrix_to_angles(R)[source]

conversion from matrix to angles

Parameters

R (torch.Tensor) – matrices of shape \((..., 3, 3)\)

Returns

e3nn.o3.angles_to_quaternion(alpha, beta, gamma)[source]

conversion from angles to quaternion

Parameters
Returns

matrices of shape \((..., 4)\)

Return type

torch.Tensor

e3nn.o3.matrix_to_quaternion(R)[source]

conversion from matrix \(R\) to quaternion \(q\)

Parameters

R (torch.Tensor) – tensor of shape \((..., 3, 3)\)

Returns

tensor of shape \((..., 4)\)

Return type

torch.Tensor

e3nn.o3.axis_angle_to_quaternion(xyz, angle)[source]

convertion from axis-angle to quaternion

Parameters
Returns

tensor of shape \((..., 4)\)

Return type

torch.Tensor

e3nn.o3.quaternion_to_axis_angle(q)[source]

convertion from quaternion to axis-angle

Parameters

q (torch.Tensor) – tensor of shape \((..., 4)\)

Returns

e3nn.o3.matrix_to_axis_angle(R)[source]

conversion from matrix to axis-angle

Parameters

R (torch.Tensor) – tensor of shape \((..., 3, 3)\)

Returns

e3nn.o3.angles_to_axis_angle(alpha, beta, gamma)[source]

conversion from angles to axis-angle

Parameters
Returns

e3nn.o3.axis_angle_to_matrix(axis, angle)[source]

conversion from axis-angle to matrix

Parameters
Returns

tensor of shape \((..., 3, 3)\)

Return type

torch.Tensor

e3nn.o3.quaternion_to_matrix(q)[source]

convertion from quaternion to matrix

Parameters

q (torch.Tensor) – tensor of shape \((..., 4)\)

Returns

tensor of shape \((..., 3, 3)\)

Return type

torch.Tensor

e3nn.o3.quaternion_to_angles(q)[source]

convertion from quaternion to angles

Parameters

q (torch.Tensor) – tensor of shape \((..., 4)\)

Returns

e3nn.o3.axis_angle_to_angles(axis, angle)[source]

convertion from axis-angle to angles

Parameters
Returns

Convertions to point on the sphere

e3nn.o3.angles_to_xyz(alpha, beta)[source]

convert \((\alpha, \beta)\) into a point \((x, y, z)\) on the sphere

Parameters
Returns

tensor of shape \((..., 3)\)

Return type

torch.Tensor

Examples

>>> angles_to_xyz(torch.tensor(1.7), torch.tensor(0.0)).abs()
tensor([0., 1., 0.])
e3nn.o3.xyz_to_angles(xyz)[source]

convert a point \(\vec r = (x, y, z)\) on the sphere into angles \((\alpha, \beta)\)

\[\vec r = R(\alpha, \beta, 0) \vec e_z\]
Parameters

xyz (torch.Tensor) – tensor of shape \((..., 3)\)

Returns