Source code for tidy3d.components.transformation
"""Defines geometric transformation classes"""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Union
import pydantic.v1 as pd
import numpy as np
from .base import Tidy3dBaseModel, cached_property
from .types import Coordinate, TensorReal, ArrayFloat2D, Axis
from ..constants import RADIAN
from ..exceptions import ValidationError
class AbstractRotation(ABC, Tidy3dBaseModel):
"""Abstract rotation of vectors and tensors."""
@cached_property
@abstractmethod
def matrix(self) -> TensorReal:
"""Rotation matrix."""
@cached_property
@abstractmethod
def isidentity(self) -> bool:
"""Check whether rotation is identity."""
def rotate_vector(self, vector: ArrayFloat2D) -> ArrayFloat2D:
"""Rotate a vector/point or a list of vectors/points.
Parameters
----------
points : ArrayLike[float]
Array of shape ``(3, ...)``.
Returns
-------
Coordinate
Rotated vector.
"""
if self.isidentity:
return vector
if len(vector.shape) == 1:
return self.matrix @ vector
return np.tensordot(self.matrix, vector, axes=1)
def rotate_tensor(self, tensor: TensorReal) -> TensorReal:
"""Rotate a tensor.
Parameters
----------
tensor : ArrayLike[float]
Array of shape ``(3, 3)``.
Returns
-------
TensorReal
Rotated tensor.
"""
if self.isidentity:
return tensor
return np.matmul(self.matrix, np.matmul(tensor, self.matrix.T))
[docs]
class RotationAroundAxis(AbstractRotation):
"""Rotation of vectors and tensors around a given vector."""
axis: Union[Axis, Coordinate] = pd.Field(
0,
title="Axis of Rotation",
description="A vector that specifies the axis of rotation, or a single int: 0, 1, or 2, "
"indicating x, y, or z.",
)
angle: float = pd.Field(
0.0,
title="Angle of Rotation",
description="Angle of rotation in radians.",
units=RADIAN,
)
@pd.validator("axis", always=True)
def _convert_axis_index_to_vector(cls, val):
if not isinstance(val, tuple):
axis = [0.0, 0.0, 0.0]
axis[val] = 1.0
val = tuple(axis)
return val
@pd.validator("axis")
def _guarantee_nonzero_axis(cls, val):
norm = np.linalg.norm(val)
if np.isclose(norm, 0):
raise ValidationError(
"The norm of vector 'axis' cannot be zero. Please provide a proper rotation axis."
)
return val
@cached_property
def isidentity(self) -> bool:
"""Check whether rotation is identity."""
return np.isclose(self.angle % (2 * np.pi), 0)
@cached_property
def matrix(self) -> TensorReal:
"""Rotation matrix."""
if self.isidentity:
return np.eye(3)
norm = np.linalg.norm(self.axis)
n = self.axis / norm
c = np.cos(self.angle)
s = np.sin(self.angle)
R = np.zeros((3, 3))
tan_dim = [[1, 2], [2, 0], [0, 1]]
for dim in range(3):
R[dim, dim] = c + n[dim] ** 2 * (1 - c)
R[dim, tan_dim[dim][0]] = n[dim] * n[tan_dim[dim][0]] * (1 - c) - n[tan_dim[dim][1]] * s
R[dim, tan_dim[dim][1]] = n[dim] * n[tan_dim[dim][1]] * (1 - c) + n[tan_dim[dim][0]] * s
return R
RotationType = Union[RotationAroundAxis]