# Source code for symforce.geo.unit3

```# ----------------------------------------------------------------------------
# SymForce - Copyright 2022, Skydio, Inc.
# This source code is under the Apache 2.0 license found in the LICENSE file.
# ----------------------------------------------------------------------------

from __future__ import annotations

import symforce.internal.symbolic as sf
from symforce import typing as T
from symforce.ops.interfaces import LieGroup

from .matrix import Matrix
from .matrix import Matrix24
from .matrix import Matrix42
from .matrix import Vector3
from .rot3 import Rot3

[docs]class Unit3(LieGroup):
"""
Direction in R^3, represented as a :class:`Rot3 <symforce.geo.rot3.Rot3>` that transforms
[0, 0, 1] to the desired direction.

The storage is therefore a quaternion and the tangent space is 2 dimensional.
Most operations are implemented using operations from :class:`Rot3 <symforce.geo.rot3.Rot3>`.

Note: an alternative implementation could directly store a unit vector and define its boxplus
manifold as described in Appendix B.2 of [Hertzberg 2013]. This can be done by finding the
Householder reflector of x and use it to transform the exponential map of delta, which is a
small perturbation in the tangent space (R^2). Namely::

x.retract(delta) = x [+] delta = Rx * Exp(delta), where
Exp(delta) = [sinc(||delta||) * delta, cos(||delta||)], and
Rx = (I - 2 vv^T / (v^Tv))X, v = x - e_z != 0, X is a matrix negating 2nd vector component
= I                     , x = e_z

[Hertzberg 2013] Integrating Generic Sensor Fusion Algorithms with Sound State Representations
through Encapsulation of Manifolds
"""

E_Z = Vector3.unit_z()

def __init__(self, rot3: T.Optional[Rot3] = None) -> None:
"""
Construct from a :class:`Rot3 <symforce.geo.rot3.Rot3>`, or identity if none provided.
"""
self.rot3 = rot3 if rot3 is not None else Rot3.identity()
assert isinstance(self.rot3, Rot3)

# -------------------------------------------------------------------------
# Storage concept - see symforce.ops.storage_ops
# -------------------------------------------------------------------------
def __repr__(self) -> str:
xyz = self.to_unit_vector()
return "<Unit3 xyz=[{}, {}, {}]>".format(repr(xyz[0]), repr(xyz[1]), repr(xyz[2]))

[docs]    @classmethod
def storage_dim(cls) -> int:
return Rot3.storage_dim()

[docs]    def to_storage(self) -> T.List[T.Scalar]:
return self.rot3.to_storage()

[docs]    @classmethod
def from_storage(cls, vec: T.Sequence[T.Scalar]) -> Unit3:
return cls(Rot3.from_storage(vec))

[docs]    @classmethod
def symbolic(cls, name: str, **kwargs: T.Any) -> Unit3:
return cls(Rot3.symbolic(name, **kwargs))

# -------------------------------------------------------------------------
# Group concept - see symforce.ops.group_ops
# -------------------------------------------------------------------------

[docs]    @classmethod
def identity(cls) -> Unit3:
return cls(Rot3.identity())

[docs]    def compose(self, other: Unit3) -> Unit3:
return Unit3(self.rot3.compose(other.rot3))

[docs]    def inverse(self) -> Unit3:
return Unit3(self.rot3.inverse())

# -------------------------------------------------------------------------
# Lie group implementation
# -------------------------------------------------------------------------

[docs]    @classmethod
def tangent_dim(cls) -> int:
return 2

[docs]    @classmethod
def from_tangent(cls, v: T.Sequence[T.Scalar], epsilon: T.Scalar = sf.epsilon()) -> Unit3:
return cls(Rot3.from_tangent([-v[1], v[0], sf.S.Zero], epsilon=epsilon))

[docs]    def to_tangent(self, epsilon: T.Scalar = sf.epsilon()) -> T.List[T.Scalar]:
v = self.rot3.to_tangent(epsilon=epsilon)
return [v[1], -v[0]]

[docs]    def storage_D_tangent(self) -> Matrix42:
D = self.rot3.storage_D_tangent()
return T.cast(Matrix42, Matrix.column_stack(D[:, 1], -D[:, 0]))

[docs]    def tangent_D_storage(self) -> Matrix24:
return 4 * T.cast(Matrix24, self.storage_D_tangent().T)

# -------------------------------------------------------------------------
# Helper methods
# -------------------------------------------------------------------------

[docs]    def to_rotation(self) -> Rot3:
return self.rot3

[docs]    def to_unit_vector(self) -> Vector3:
return self.rot3 * self.E_Z

[docs]    @classmethod
def from_vector(cls, a: Vector3, epsilon: T.Scalar = sf.epsilon()) -> Unit3:
"""
Return a :class:`Unit3` that points along the direction of vector ``a``

``a`` does not have to be a unit vector.
"""
u = a.normalized(epsilon=epsilon)
return cls(Rot3.from_two_unit_vectors(cls.E_Z, u, epsilon=epsilon))

[docs]    @classmethod
def random(cls, epsilon: T.Scalar = sf.epsilon()) -> Unit3:
"""
Generate a random element of :class:`Unit3`, by generating a random rotation first and then
rotating ``e_z`` to get a random direction.
"""
return cls.from_vector(Rot3.random() * cls.E_Z, epsilon=epsilon)
```