"""
Implements various means of generating rotation matrices.
Copyright (c) 2023-2025 European Molecular Biology Laboratory
Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
"""
import yaml
import warnings
from typing import Tuple
from os.path import join, dirname
import numpy as np
from scipy.spatial.transform import Rotation
from .types import NDArray
__all__ = [
"get_cone_rotations",
"align_vectors",
"euler_to_rotationmatrix",
"euler_from_rotationmatrix",
"get_rotation_matrices",
"align_to_axis",
]
def _sample_cone(
angle: float, sampling: float, axis: Tuple[float] = (1, 0, 0)
) -> NDArray:
"""
Sample points on a cone surface around cone_axis using golden spiral distribution.
Parameters
----------
angle : float
The half-angle of the cone in degrees.
sampling : float
Angular increment used for sampling points in degrees.
axis : tuple of floats
Vector to align the cone with.
Returns
-------
NDArray
Array of points around axis with shape n,3.
References
----------
.. [1] https://stackoverflow.com/questions/9600801/evenly-distributing-n-points-on-a-sphere
"""
theta = np.linspace(0, angle, round(angle / sampling) + 1)
number_of_points = np.ceil(
360 * np.divide(np.sin(np.radians(theta)), sampling),
)
number_of_points = int(np.sum(number_of_points + 1) + 2)
indices = np.arange(0, number_of_points, dtype=float) + 0.5
radius = np.radians(angle * np.sqrt(indices / number_of_points))
theta = np.pi * (1 + np.sqrt(5)) * indices
points = np.stack(
[
np.cos(radius),
np.cos(theta) * np.sin(radius),
np.sin(theta) * np.sin(radius),
],
axis=1,
)
rotation = Rotation.from_euler(
angles=align_vectors((1, 0, 0), axis, seq="zyz"),
seq="zyz",
degrees=True,
)
return rotation.apply(points)
[docs]
def get_cone_rotations(
cone_angle: float,
cone_sampling: float,
axis: Tuple[float] = (1, 0, 0),
axis_angle: float = 360.0,
axis_sampling: float = None,
reference: Tuple[float] = (1, 0, 0),
n_symmetry: int = 1,
seq: str = None,
) -> NDArray:
"""
Generate rotations describing the possible placements of a vector in a cone.
Parameters
----------
cone_angle : float
The half-angle of the cone in degrees.
cone_sampling : float
Angular increment used for sampling points on the cone in degrees.
axis : Tuple[float], optional
Base-vector of the cone.
axis_angle : float, optional
The total angle of rotation around the cone axis in degrees (default is 360.0).
axis_sampling : float, optional
Angular increment used for sampling points around the cone axis in degrees.
If None, it takes the value of cone_sampling.
reference : Tuple[float], optional
Returned rotations will map this point onto the cone. In practice, this is
the principal axis of the template.
n_symmetry : int, optional
Number of symmetry axis around the vector axis.
seq : str, optional
Convention for angles. By default returns rotation matrices.
Returns
-------
NDArray
An arary of rotations represented as stack of rotation matrices when convention
None (n, 3, 3) or an array of Euler angles (n, 3) for available conventions.
"""
if axis_sampling is None:
axis_sampling = cone_sampling
points = _sample_cone(angle=cone_angle, sampling=cone_sampling, axis=axis)
reference = np.asarray(reference).astype(np.float32)
reference /= np.linalg.norm(reference)
axis_angle /= n_symmetry
phi_steps = np.maximum(np.round(axis_angle / axis_sampling), 1).astype(int)
phi = np.linspace(0, axis_angle, phi_steps + 1)[:-1]
axis_rotation = Rotation.from_rotvec(axis * np.radians(phi)[:, None])
all_rotations = [
axis_rotation * Rotation.from_matrix(align_vectors(reference, x))
for x in points
]
rotations = Rotation.concatenate(all_rotations)
if seq is None:
return rotations.as_matrix()
return rotations.as_euler(seq=seq, degrees=True)
[docs]
def align_vectors(base: NDArray, target: NDArray = (0, 0, 1), seq: str = None):
"""
Compute the rotation matrix or Euler angles required to align an initial vector with a target vector.
Parameters
----------
base : NDArray
The basis vector.
target : NDArray, optional
The vector to map base to, defaults to (0,0,1).
seq : str, optional
Euler angle convention, None returns a rotation matrix instead.
Returns
-------
NDArray
Rotation matrix if seq is None otherwise Euler angles in desired convention
"""
base = np.asarray(base, dtype=np.float32)
target = np.asarray(target, dtype=np.float32)
rotation, error = Rotation.align_vectors(target, base)
if seq is None:
return rotation.as_matrix()
return rotation.as_euler(seq=seq, degrees=True)
[docs]
def euler_to_rotationmatrix(angles: Tuple[float], seq: str = "zyz") -> NDArray:
"""
Convert Euler angles to a rotation matrix.
Parameters
----------
angles : tuple
A tuple representing the Euler angles in degrees.
seq : str, optional
Euler angle convention.
Returns
-------
NDArray
The generated rotation matrix.
"""
angles = np.asarray(angles)
n_angles = len(angles)
if angles.ndim == 2:
n_angles = angles.shape[1]
rotation_matrix = Rotation.from_euler(
seq=seq[:n_angles], angles=angles, degrees=True
)
return rotation_matrix.as_matrix().astype(np.float32)
[docs]
def euler_from_rotationmatrix(rotation_matrix: NDArray, seq: str = "zyz") -> Tuple:
"""
Convert a rotation matrix to euler angles.
Parameters
----------
rotation_matrix : NDArray
A 2 x 2 or 3 x 3 rotation matrix.
seq : str, optional
Euler angle convention, zyz by default.
Returns
-------
Tuple
The generate euler angles in degrees
"""
if rotation_matrix.shape[0] == 2:
temp_matrix = np.eye(3)
temp_matrix[:2, :2] = rotation_matrix
rotation_matrix = temp_matrix
with warnings.catch_warnings():
warnings.simplefilter("ignore")
rotation = Rotation.from_matrix(rotation_matrix)
angles = rotation.as_euler(seq=seq, degrees=True).astype(np.float32)
return angles
[docs]
def get_rotation_matrices(
angular_sampling: float, dim: int = 3, use_optimized_set: bool = True
) -> NDArray:
"""
Returns rotation matrices with desired ``angular_sampling`` rate.
Parameters
----------
angular_sampling : float
The desired angular sampling in degrees.
dim : int, optional
Dimension of the rotation matrices.
use_optimized_set : bool, optional
Use optimized rotational sets, True by default and available for dim=3.
Notes
-----
For dim = 3 optimized sets are used, otherwise QR-decomposition.
Returns
-------
NDArray
Array of shape (n, d, d) containing n rotation matrices.
"""
if dim == 3 and use_optimized_set:
quaternions, *_ = _load_quaternions_by_angle(angular_sampling)
ret = Rotation.from_quat(quaternions, scalar_first=True).as_matrix()
else:
num_rotations = dim * (dim - 1) // 2
k = int((360 / angular_sampling) ** num_rotations)
As = np.random.randn(k, dim, dim)
ret, _ = np.linalg.qr(As)
dets = np.linalg.det(ret)
neg_dets = dets < 0
ret[neg_dets, :, -1] *= -1
ret[0] = np.eye(dim, dtype=ret.dtype)
return ret
def _load_quaternions_by_angle(
angular_sampling: float,
) -> Tuple[NDArray, NDArray, float]:
"""
Get orientations and weights proportional to the given angular_sampling.
Parameters
----------
angular_sampling : float
Requested angular sampling.
Returns
-------
Tuple[NDArray, NDArray, float]
Quaternions, associated weights and realized angular sampling rate.
"""
# Metadata contains (N orientations, rotational sampling, coverage as values)
with open(join(dirname(__file__), "data", "metadata.yaml"), "r") as infile:
metadata = yaml.full_load(infile)
set_diffs = {
setname: abs(angular_sampling - set_angle)
for setname, (_, set_angle, _) in metadata.items()
}
fname = min(set_diffs, key=set_diffs.get)
infile = join(dirname(__file__), "data", fname)
quat_weights = np.load(infile)
quat = quat_weights[:, :4]
weights = quat_weights[:, -1]
angle = metadata[fname][0]
return quat, weights, angle
[docs]
def align_to_axis(
coordinates: NDArray,
weights: NDArray = None,
axis: int = 2,
flip: bool = False,
eigenvector_index: int = 0,
) -> NDArray:
"""
Calculate a rotation matrix that aligns the principal axis of a point cloud
with a specified coordinate axis.
Parameters
----------
coordinates : NDArray
Array of 3D coordinates with shape (n, 3) representing the point cloud.
weights : NDArray
Coordinate weighting factors with shape (n,).
axis : int, optional
The target axis to align with, defaults to 2 (z-axis).
flip : bool, optional
Whether to align with the negative direction of the axis, default is False.
eigenvector_index : int, optional
Index of eigenvector to select, sorted by descending eigenvalues.
0 = largest eigenvalue (most variance), 1 = second largest, etc.
Default is 0 (primary principal component).
Returns
-------
NDArray
3x3 rotation matrix that aligns the principal component of the
coordinates with the specified axis.
"""
axis = int(axis)
coordinates = np.asarray(coordinates)
alignment_axis = np.array(
[0 if i != axis else 1 for i in range(coordinates.shape[1])]
)
if flip:
alignment_axis *= -1
ndim = coordinates.shape[1]
if eigenvector_index >= ndim:
raise ValueError(f"eigenvector_index has to be less than {ndim}.")
avg = np.average(coordinates, axis=0, weights=weights)
coordinates = coordinates - avg
cov_matrix = np.cov(coordinates.T, aweights=weights)
# Eigenvalues are already sorted in ascending order
eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)
eigenvector = eigenvectors[:, -(eigenvector_index + 1)]
return align_vectors(eigenvector, alignment_axis)