Implemented transformation-buffer

commit-id:2fbffaa2

Author: integraledelebesgue

packages/transformation-buffer/pyproject.toml

@@ -3,12 +3,27 @@ name = "transformation-buffer"
 version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
-authors = [
-    { name = "Jan Smółka", email = "[email protected]" }
-]
+authors = [{ name = "Jan Smółka", email = "[email protected]" }]
 requires-python = ">=3.12.7"
-dependencies = []
+dependencies = [
+    "jaxtyping>=0.2.37",
+    "kornia>=0.8.0",
+    "more-itertools>=10.6.0",
+    "pyserde[numpy]>=0.23.0",
+    "rustworkx>=0.16.0",
+]
+
+[dependency-groups]
+dev = [
+    "hypothesis>=6.125.3",
+    "pytest>=8.3.4",
+    "syrupy>=4.8.1",
+]
 
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
+
+[tool.pytest.ini_options]
+python_files = "tests/**/*.py"
+python_functions = "test_*"

packages/transformation-buffer/src/transformation_buffer/__init__.py

@@ -1 +1,4 @@
+from .buffer import Buffer
+from .transformation import Transformation
 
+__all__ = ['Buffer', 'Transformation']

packages/transformation-buffer/src/transformation_buffer/buffer.py

@@ -0,0 +1,132 @@
+from collections.abc import Callable, Hashable, Iterable
+from copy import deepcopy
+from functools import reduce
+from typing import Self
+
+from more_itertools import pairwise
+from rustworkx import PyDiGraph, digraph_dijkstra_shortest_paths, is_connected
+
+from transformation_buffer.rigid_model import RigidModel
+from transformation_buffer.transformation import Transformation
+
+# TODO: Implement `pyserde` serialization.
+# TODO: Implement path compression and verification.
+
+
+class Buffer[T: Hashable]:
+    __frames_of_reference: set[T]
+    __frame_names_to_node_indices: dict[T, int]
+    __connections: PyDiGraph[T, Transformation]
+
+    def __init__(self, frames_of_reference: Iterable[T] | None = None) -> None:
+        nodes = (
+            set() if frames_of_reference is None else deepcopy(set(frames_of_reference))
+        )
+        n_nodes = len(nodes)
+
+        connections = PyDiGraph[T, Transformation](
+            multigraph=True,  # Think about it later
+            node_count_hint=n_nodes,
+            edge_count_hint=n_nodes,
+        )
+
+        node_indices = connections.add_nodes_from(nodes)
+
+        self.__frames_of_reference = nodes
+        self.__frame_names_to_node_indices = dict(zip(nodes, node_indices))
+        self.__connections = connections
+
+    @property
+    def connected(self) -> bool:
+        return is_connected(self.__connections.to_undirected(multigraph=False))
+
+    @property
+    def frames_of_reference(self) -> frozenset[T]:
+        # Construction of frozenset from set of strings is *somehow* optimized.
+        # Perhaps there's no need for caching this property.
+        return frozenset(self.__frames_of_reference)
+
+    def add_object(self, model: RigidModel[T]) -> Self:
+        for from_to, transformation in model.transformations().items():
+            self[from_to] = transformation
+
+        return self
+
+    def add_frame_of_reference(self, frame: T) -> Self:
+        self.__add_frame_of_reference(frame)
+        return self
+
+    def __add_frame_of_reference(self, frame: T) -> int:
+        if (index := self.__frame_index(frame)) is not None:
+            return index
+
+        self.__frames_of_reference.add(frame)
+        index = self.__connections.add_node(frame)
+        self.__frame_names_to_node_indices[frame] = index
+
+        return index
+
+    def __frame_index(self, frame: T) -> int | None:
+        return self.__frame_names_to_node_indices.get(frame, None)
+
+    def __setitem__(self, from_to: tuple[T, T], transformation: Transformation) -> None:
+        from_id, to_id = from_to
+
+        from_index = self.__frame_index(from_id) or self.__add_frame_of_reference(from_id)
+        to_index = self.__frame_index(to_id) or self.__add_frame_of_reference(to_id)
+
+        self.__connections.add_edge(from_index, to_index, transformation)
+        self.__connections.add_edge(to_index, from_index, transformation.inverse())
+
+    def __getitem__(self, from_to: tuple[T, T]) -> Transformation | None:
+        from_id, to_id = from_to
+
+        from_index = self.__frame_index(from_id)
+        to_index = self.__frame_index(to_id)
+
+        if from_index is None or to_index is None:
+            return None
+
+        if from_index == to_index:
+            return Transformation.identity().clone()
+
+        connections = self.__connections
+
+        if self.__connections.has_edge(from_index, to_index):
+            return self.__connections.get_edge_data(from_index, to_index)
+
+        path_mapping = digraph_dijkstra_shortest_paths(
+            connections,
+            from_index,
+            to_index,
+        )
+
+        if to_index not in path_mapping:
+            return None
+
+        shortest_path = path_mapping[to_index]
+
+        return map_reduce(  # type: ignore[no-any-return]  # MyPy complains about the possible `Any`.
+            lambda nodes: connections.get_edge_data(*nodes),
+            lambda t1, t2: t1 @ t2,
+            pairwise(shortest_path),
+        )
+
+    def __contains__(self, from_to: tuple[T, T]) -> bool:
+        from_id, to_id = from_to
+
+        from_index = self.__frame_index(from_id)
+        to_index = self.__frame_index(to_id)
+
+        if from_index is None or to_index is None:
+            return False
+
+        return self.__connections.has_edge(from_index, to_index)
+
+
+def map_reduce[U, V](
+    map_function: Callable[[U], V],
+    reduce_function: Callable[[V, V], V],
+    iterable: Iterable[U],
+) -> V:
+    return reduce(reduce_function, map(map_function, iterable))

packages/transformation-buffer/src/transformation_buffer/rigid_model.py

@@ -0,0 +1,52 @@
+from math import pi
+from typing import Hashable, Protocol
+
+import torch
+from attrs import frozen
+
+from transformation_buffer.transformation import Transformation
+
+
+class RigidModel[T: Hashable](Protocol):
+    def transformations(self) -> dict[tuple[T, T], Transformation]: ...
+
+
+@frozen
+class Cube[T: Hashable]:
+    width: float
+    walls: tuple[T, T, T, T, T, T]
+
+    # Cube Walls:
+    #    4
+    #  5 1 6
+    #    2
+    #    3
+
+    def transformations(self) -> dict[tuple[T, T], Transformation]:
+        first, second, third, forth, fifth, sixth = self.walls
+
+        half_pi = pi / 2.0
+        half_width = self.width / 2.0
+
+        down = Transformation.active(
+            intrinsic_euler_angles=torch.tensor((-half_pi, 0.0, 0.0)),
+            translation=torch.tensor((0.0, half_width, -half_width)),
+        )
+
+        left = Transformation.active(
+            intrinsic_euler_angles=torch.tensor((0.0, half_pi, 0.0)),
+            translation=torch.tensor((half_width, 0.0, -half_width)),
+        )
+
+        right = Transformation.active(
+            intrinsic_euler_angles=torch.tensor((0.0, -half_pi, 0.0)),
+            translation=torch.tensor((-half_width, 0.0, -half_width)),
+        )
+
+        return {
+            (first, second): down.clone(),
+            (second, third): down.clone(),
+            (third, forth): down,
+            (first, fifth): left,
+            (first, sixth): right,
+        }

packages/transformation-buffer/src/transformation_buffer/transformation.py

@@ -0,0 +1,175 @@
+from dataclasses import dataclass
+from typing import Literal, Self
+
+import torch
+from jaxtyping import Float
+from serde import field, serde
+
+
+def _serialize_tensor(tensor: torch.Tensor) -> str:
+    return f'torch.tensor({tensor.tolist()}, dtype={tensor.dtype})'
+
+
+def _deserialize_tensor(input: str) -> torch.Tensor:
+    import torch as __torch
+
+    result = eval(input, {}, {'torch': __torch})
+    del __torch
+
+    return result  # type: ignore[no-any-return]
+
+
+type Axis = Literal['x', 'y', 'z']
+
+
+def rotation_around(
+    axis: Axis,
+    angle: float | torch.Tensor,
+) -> Float[torch.Tensor, '3 3']:
+    a = torch.tensor(angle)
+    sin = torch.sin(a)
+    cos = torch.cos(a)
+
+    match axis:
+        case 'x':
+            return torch.tensor(
+                [
+                    [1.0, 0.0, 0.0],
+                    [0.0, cos, -sin],
+                    [0.0, sin, cos],
+                ]
+            )
+
+        case 'y':
+            return torch.tensor(
+                [
+                    [cos, 0.0, sin],
+                    [0.0, 1.0, 0.0],
+                    [-sin, 0.0, cos],
+                ]
+            )
+
+        case 'z':
+            return torch.tensor(
+                [
+                    [cos, -sin, 0.0],
+                    [sin, cos, 0.0],
+                    [0.0, 0.0, 1.0],
+                ]
+            )
+
+
+def translation_along(axis: Axis, distance: float) -> Float[torch.Tensor, '3']:
+    t = torch.zeros(3)
+
+    match axis:
+        case 'x':
+            t[0] = distance
+        case 'y':
+            t[1] = distance
+        case 'z':
+            t[2] = distance
+
+    return t
+
+
+@serde
+@dataclass(eq=False)
+class Transformation:
+    rotation_and_translation: Float[torch.Tensor, '4 4'] = field(
+        serializer=_serialize_tensor,
+        deserializer=_deserialize_tensor,
+    )
+
+    def __init__(
+        self,
+        rotation_and_translation: Float[torch.Tensor, '4 4'],
+    ) -> None:
+        self.rotation_and_translation = rotation_and_translation.clone()
+
+    def __eq__(self, other: object, /) -> bool:
+        return isinstance(other, Transformation) and bool(
+            self.rotation_and_translation.eq(other.rotation_and_translation).all()
+        )
+
+    def __matmul__(self, other: 'Transformation') -> 'Transformation':
+        return Transformation(
+            other.rotation_and_translation @ self.rotation_and_translation
+        )
+
+    def inverse(self) -> 'Transformation':
+        return Transformation(self.rotation_and_translation.inverse())
+
+    @staticmethod
+    def approx_eq(
+        t1: 'Transformation',
+        t2: 'Transformation',
+        absolute_tolerance: float = 1e-8,
+    ) -> bool:
+        return bool(
+            torch.isclose(
+                t1.rotation_and_translation,
+                t2.rotation_and_translation,
+                atol=absolute_tolerance,
+            ).all()
+        )
+
+    # Cannot write `IDENTITY: ClassVar[Self] = Transformation.from_parts`.
+    # Class properties are also deprecated.
+    # Python... :v
+    @classmethod
+    def identity(cls) -> Self:
+        return cls.from_parts(
+            torch.eye(3, 3, dtype=torch.float32),
+            torch.zeros(3, dtype=torch.float32),
+        )
+
+    @classmethod
+    def from_parts(
+        cls,
+        rotation: Float[torch.Tensor, '3 3'],
+        translation: Float[torch.Tensor, '3'],
+    ) -> Self:
+        rotation_and_translation = torch.zeros(
+            (4, 4),
+            dtype=rotation.dtype,
+            device=rotation.device,
+        )
+        rotation_and_translation[:3, :3] = rotation
+        rotation_and_translation[:3, 3] = translation
+        rotation_and_translation[3, 3] = 1.0
+
+        return cls(rotation_and_translation)
+
+    @classmethod
+    def active(
+        cls,
+        intrinsic_euler_angles: Float[torch.Tensor, '3'] | None = None,
+        translation: Float[torch.Tensor, '3'] | None = None,
+    ) -> Self:
+        rotation = (
+            (
+                rotation_around('x', intrinsic_euler_angles[0])
+                @ rotation_around('y', intrinsic_euler_angles[1])
+                @ rotation_around('z', intrinsic_euler_angles[2])
+            )
+            if intrinsic_euler_angles is not None
+            else torch.eye(3)
+        )
+
+        translation = translation if translation is not None else torch.zeros(3)
+
+        return cls.from_parts(rotation, translation)
+
+    def clone(self) -> 'Transformation':
+        return Transformation(self.rotation_and_translation.clone())
+
+    def cpu(self) -> 'Transformation':
+        return Transformation(self.rotation_and_translation.cpu())
+
+    def to(
+        self,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> 'Transformation':
+        return Transformation(self.rotation_and_translation.to(device, dtype, copy=True))