Adding basic vector and rotation logic

requirements.in

@@ -1,6 +1,7 @@
 asyncio
 fastapi
 jinja2
+numpy
 opencv-python
 uvicorn[standard]
 websockets

requirements_lock.txt

@@ -177,7 +177,9 @@ numpy==1.26.4 \
     --hash=sha256:edd8b5fe47dab091176d21bb6de568acdd906d1887a4584a15a9a96a1dca06ef \
     --hash=sha256:f870204a840a60da0b12273ef34f7051e98c3b5961b61b0c2c1be6dfd64fbcd3 \
     --hash=sha256:ffa75af20b44f8dba823498024771d5ac50620e6915abac414251bd971b4529f
-    # via opencv-python
+    # via
+    #   -r requirements.in
+    #   opencv-python
 opencv-python==4.9.0.80 \
     --hash=sha256:1a9f0e6267de3a1a1db0c54213d022c7c8b5b9ca4b580e80bdc58516c922c9e1 \
     --hash=sha256:3f16f08e02b2a2da44259c7cc712e779eff1dd8b55fdb0323e8cab09548086c0 \

src/robot/v1/BUILD.bazel

@@ -11,7 +11,22 @@ py_library(
     deps = [
         "@pypi//fastapi:pkg",
         "@pypi//jinja2:pkg",
+        "@pypi//numpy:pkg",
         "@pypi//opencv_python:pkg",
         "@pypi//websockets:pkg",
     ],
 )
+
+py_library(
+    name = "structs_src",
+    srcs = ["structs.py"],
+    deps = [
+        "@pypi//numpy:pkg",
+    ],
+)
+
+py_test(
+    name = "structs_test",
+    srcs = ["structs_test.py"],
+    deps = [":structs_src"],
+)

src/robot/v1/structs.py

@@ -1,46 +1,33 @@
 from dataclasses import dataclass
+from typing import Union
+import numpy as np
 
-@dataclass
-class Vector2D:
-    x: float
-    y: float
-
-    def __add__(self, other):
-        assert isinstance(other, Vector2D)
-
-        return Vector2D(self.x + other.x, self.y + other.y)
-
-    def __iadd__(self, other):
-        assert isinstance(other, Vector2D)
-
-        self.x += other.x
-        self.y += other.y
 
-    def __sub__(self, other):
-        assert isinstance(other, Vector2D)
+def vector2d(x: Union[int, float], y: Union[int, float]):
+    """
+    Returns a 2D vector as a numpy array.
 
-        return Vector2D(self.x - other.x, self.y - other.y)
+    Params:
+        x [Union[int, float]]: X component of the vector
+        y [Union[int, float]]: Y component of the vector
+    """
+    return np.array([x, y])
 
-    def __isub__(self, other):
-        assert isinstance(other, Vector2D)
 
-        self.x -= other.x
-        self.y -= other.y
+def rotation2d(theta: Union[int, float]):
+    """
+    Return a 2D rotation matrix as a numpy array, provided an angle in radians.
+    From https://scipython.com/book/chapter-6-numpy/examples/creating-a-rotation-matrix-in-numpy/
 
+    Params:
+        theta [Union[int, float]]: Angle in radians
+    """
+    c, s = np.cos(theta), np.sin(theta)
+    return np.array(((c, -s), (s, c)))
 
-@dataclass
-class Rotation2D:
-    theta: float
-
-
-@dataclass
-class Position2D:
-    location: Vector2D
-    rotation: Rotation2D
 
-    def __init__(self, x, y, theta):
-        self.location = Vector2D(x, y)
-        self.rotation = Rotation2D(theta)
+# TODO: Add some smallest angle function. Needs to keep sign,
+# can likely be done with modulus
 
 
 @dataclass

src/robot/v1/structs_test.py

@@ -1,78 +1,182 @@
 import unittest
+import math
+import numpy as np
+from src.robot.v1.structs import vector2d, rotation2d
+
 
 class TestVector2D(unittest.TestCase):
-    def test_instantiation(self):
-        pass
+    def test_add(self):
+        total = vector2d(1, 2) + vector2d(3, 4)
+        self.assertEqual(total[0], 4)
+        self.assertEqual(total[1], 6)
+
+    def test_iadd(self):
+        vec_a = vector2d(1, 2)
+        vec_a += vector2d(5, 5)
+        self.assertEqual(vec_a[0], 6)
+        self.assertEqual(vec_a[1], 7)
+
+    def test_sub(self):
+        total = vector2d(5, 5) - vector2d(3, 4)
+        self.assertEqual(total[0], 2)
+        self.assertEqual(total[1], 1)
+
+    def test_isub(self):
+        vec_a = vector2d(5, 5)
+        vec_a -= vector2d(1, 2)
+        self.assertEqual(vec_a[0], 4)
+        self.assertEqual(vec_a[1], 3)
+
+    def test_45deg_rotation(self):
+        angle = math.pi / 4
+        starting_vec = vector2d(1, 2)
+
+        # Figure out the expected angles programmatically
+        r = math.sqrt(5)
+        expected_angle_cw = math.atan2(starting_vec[1], starting_vec[0]) - (math.pi / 4)
+        expected_vec_cw = vector2d(
+            r * math.cos(expected_angle_cw), r * math.sin(expected_angle_cw)
+        )
+
+        expected_angle_ccw = math.atan2(starting_vec[1], starting_vec[0]) + (
+            math.pi / 4
+        )
+        expected_vec_ccw = vector2d(
+            r * math.cos(expected_angle_ccw), r * math.sin(expected_angle_ccw)
+        )
 
-    def test_
+        for theta, expected in [(angle, expected_vec_cw), (-angle, expected_vec_ccw)]:
+            rotation = rotation2d(theta)
 
-@dataclass
-class Vector2D:
-    x: float
-    y: float
+            rotated = starting_vec.dot(rotation)
+            self.assertEqual(rotated.shape, (2,))
+            self.assertAlmostEqual(rotated[0], expected[0])
+            self.assertAlmostEqual(rotated[1], expected[1])
 
-    def __add__(self, other):
-        assert isinstance(other, Vector2D)
+    def test_90deg_rotation(self):
+        angle = math.pi / 2
+        starting_vec = vector2d(1, 2)
+        expected_vec = vector2d(2, -1)
 
-        return Vector2D(self.x + other.x, self.y + other.y)
+        for theta, expected in [(angle, expected_vec), (-angle, -expected_vec)]:
+            rotation = rotation2d(theta)
 
-    def __iadd__(self, other):
-        assert isinstance(other, Vector2D)
+            rotated = starting_vec.dot(rotation)
+            self.assertEqual(rotated.shape, (2,))
+            self.assertAlmostEqual(rotated[0], expected[0])
+            self.assertAlmostEqual(rotated[1], expected[1])
 
-        self.x += other.x
-        self.y += other.y
+    def test_180deg_rotation(self):
+        angle = math.pi
+        starting_vec = vector2d(1, 2)
+        expected_vec = vector2d(-1, -2)
 
-    def __sub__(self, other):
-        assert isinstance(other, Vector2D)
+        for theta, expected in [(angle, expected_vec), (-angle, expected_vec)]:
+            rotation = rotation2d(theta)
+
+            rotated = starting_vec.dot(rotation)
+            self.assertEqual(rotated.shape, (2,))
+            self.assertAlmostEqual(rotated[0], expected[0])
+            self.assertAlmostEqual(rotated[1], expected[1])
+
+    def test_270deg_rotation(self):
+        angle = (3 * math.pi) / 2
+        starting_vec = vector2d(1, 2)
+        expected_vec = vector2d(-2, 1)
+
+        for theta, expected in [(angle, expected_vec), (-angle, -expected_vec)]:
+            rotation = rotation2d(theta)
+
+            rotated = starting_vec.dot(rotation)
+            self.assertEqual(rotated.shape, (2,))
+            self.assertAlmostEqual(rotated[0], expected[0])
+            self.assertAlmostEqual(rotated[1], expected[1])
+
+    def test_360deg_rotation(self):
+        angle = 2 * math.pi
+        starting_vec = vector2d(1, 2)
+        expected_vec = vector2d(1, 2)
+
+        for theta, expected in [(angle, expected_vec), (-angle, expected_vec)]:
+            rotation = rotation2d(theta)
+
+            rotated = starting_vec.dot(rotation)
+            self.assertEqual(rotated.shape, (2,))
+            self.assertAlmostEqual(rotated[0], expected[0])
+            self.assertAlmostEqual(rotated[1], expected[1])
+
+
+if __name__ == "__main__":
+    unittest.main()
+
+
+# @dataclass
+# class Vector2D:
+#     x: float
+#     y: float
+
+#     def __add__(self, other):
+#         assert isinstance(other, Vector2D)
+
+#         return Vector2D(self.x + other.x, self.y + other.y)
+
+#     def __iadd__(self, other):
+#         assert isinstance(other, Vector2D)
+
+#         self.x += other.x
+#         self.y += other.y
 
-        return Vector2D(self.x - other.x, self.y - other.y)
+#     def __sub__(self, other):
+#         assert isinstance(other, Vector2D)
+
+#         return Vector2D(self.x - other.x, self.y - other.y)
 
-    def __isub__(self, other):
-        assert isinstance(other, Vector2D)
+#     def __isub__(self, other):
+#         assert isinstance(other, Vector2D)
 
-        self.x -= other.x
-        self.y -= other.y
+#         self.x -= other.x
+#         self.y -= other.y
 
 
-@dataclass
-class Rotation2D:
-    theta: float
+# @dataclass
+# class Rotation2D:
+#     theta: float
 
 
-@dataclass
-class Position2D:
-    location: Vector2D
-    rotation: Rotation2D
+# @dataclass
+# class Position2D:
+#     location: Vector2D
+#     rotation: Rotation2D
 
-    def __init__(self, x, y, theta):
-        self.location = Vector2D(x, y)
-        self.rotation = Rotation2D(theta)
+#     def __init__(self, x, y, theta):
+#         self.location = Vector2D(x, y)
+#         self.rotation = Rotation2D(theta)
 
 
-@dataclass
-class MemorySegment:
-    byte_offset: int
-    byte_size: int
+# @dataclass
+# class MemorySegment:
+#     byte_offset: int
+#     byte_size: int
 
 
-@dataclass
-class DynamixelMotorAddressConfig:
-    torque_enable: MemorySegment
-    led_enable: MemorySegment
-    d_gain: MemorySegment
-    i_gain: MemorySegment
-    p_gain: MemorySegment
-    goal_position: MemorySegment
-    moving_speed: MemorySegment
-    torque_limit: MemorySegment
-    present_position: MemorySegment
-    present_speed: MemorySegment
-    present_load: MemorySegment
-    present_input_voltage: MemorySegment
-    present_temperature: MemorySegment
-    registered: MemorySegment
-    moving: MemorySegment
-    lock: MemorySegment
-    punch: MemorySegment
-    realtime_tick: MemorySegment
-    goal_acceleration: MemorySegment
+# @dataclass
+# class DynamixelMotorAddressConfig:
+#     torque_enable: MemorySegment
+#     led_enable: MemorySegment
+#     d_gain: MemorySegment
+#     i_gain: MemorySegment
+#     p_gain: MemorySegment
+#     goal_position: MemorySegment
+#     moving_speed: MemorySegment
+#     torque_limit: MemorySegment
+#     present_position: MemorySegment
+#     present_speed: MemorySegment
+#     present_load: MemorySegment
+#     present_input_voltage: MemorySegment
+#     present_temperature: MemorySegment
+#     registered: MemorySegment
+#     moving: MemorySegment
+#     lock: MemorySegment
+#     punch: MemorySegment
+#     realtime_tick: MemorySegment
+#     goal_acceleration: MemorySegment