Merge pull request #36 from Nonannet/dev

Dev

src/copapy/_autograd.py

@@ -1,3 +1,5 @@
+from copapy._quaternion import quaternion
+
 from . import value, vector, tensor
 import copapy.backend as cpb
 from typing import Any, Sequence, overload
@@ -12,8 +14,10 @@ def grad(x: Any, y: vector[Any]) -> vector[float]: ...
 @overload
 def grad(x: Any, y: tensor[Any]) -> tensor[float]: ...
 @overload
+def grad(x: Any, y: quaternion) -> quaternion: ...
+@overload
 def grad(x: Any, y: Sequence[value[Any]]) -> list[unifloat]: ...
-def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[Any]) -> Any:
+def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[Any] | quaternion) -> Any:
     """Returns the partial derivative dx/dy where x needs to be a scalar
     and y might be a scalar, a list of scalars, a vector or matrix. It
     uses automatic differentiation in reverse-mode.
@@ -27,12 +31,14 @@ def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[Any
     """
     assert isinstance(x, value), f"Argument x for grad function must be a copapy value but is {type(x)}."
 
+    y_set: set[value[float] | float]
+
     if isinstance(y, value):
         y_set = {y}
     if isinstance(y, tensor):
         y_set = {v.get_scalar(0) for v in y.flatten()}
     else:
-        assert isinstance(y, Sequence) or isinstance(y, vector)
+        assert isinstance(y, Sequence) or isinstance(y, vector) or isinstance(y, quaternion)
         y_set = set(y)
 
     edges = cpb.get_all_dag_edges_between([x.net.source], (v.net.source for v in y_set if isinstance(v, value)))
@@ -125,6 +131,8 @@ def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[Any
         return grad_dict[y.net]
     if isinstance(y, vector):
         return vector(grad_dict[yi.net] if isinstance(yi, value) else 0.0 for yi in y.values)
+    if isinstance(y, quaternion):
+        return quaternion(grad_dict[yi.net] if isinstance(yi, value) else 0.0 for yi in y.values)
     if isinstance(y, tensor):
         return tensor([grad_dict[yi.net] if isinstance(yi, value) else 0.0 for yi in y.values], y.shape)
     return [grad_dict[yi.net] for yi in y]

src/copapy/_quaternion.py

@@ -1,4 +1,4 @@
-from typing import overload, Iterable, Callable, Any
+from typing import overload, Iterable, Callable, Any, Iterator
 from ._vectors import vector
 from ._tensors import tensor
 import copapy as cp
@@ -208,6 +208,9 @@ class quaternion(ArrayType[float]):
         """
         return quaternion(func(x) for x in self.values)
 
+    def __iter__(self) -> Iterator[value[float] | float]:
+        return iter(self.values)
+
     def __neg__(self) -> 'quaternion':
         return quaternion(-self.w, -self.x, -self.y, -self.z)
 
@@ -215,7 +218,7 @@ class quaternion(ArrayType[float]):
         return self.norm()
 
     def __repr__(self) -> str:
-        return f"vector({self.values})"
+        return f"quaternion(w={self.w}, x={self.x}, y={self.y}, z={self.z})"
 
     def __len__(self) -> int:
         return len(self.values)

tests/test_quaternion.py

@@ -1,7 +1,7 @@
 import math
-from copapy import quaternion, tensor
+from copapy import quaternion, tensor, vector, Target
-from copapy import vector, Target
 import copapy as cp
+import pytest
 
 
 def isclose(a, b, rel_tol=1e-9, abs_tol=0.0):
@@ -294,3 +294,40 @@ def test_satellite_attitude_correction():
     assert isclose(result_normal[0], expected_rotated[0], abs_tol=1e-6)
     assert isclose(result_normal[1], expected_rotated[1], abs_tol=1e-6)
     assert isclose(result_normal[2], expected_rotated[2], abs_tol=1e-6)
+
+
+def test_sensor_fusion():
+    # Based on Sebastian O. H. Madgwick's sensor fusion algorithm for orientation estimation.
+    # https://x-io.co.uk/open-source-imu-and-ahrs-algorithms
+
+    def update_orientation(q: quaternion, gyro: vector[float], accel: vector[float], dt: float = 0.01):
+        # Compute the cost function and its gradient
+        objective = q.rotate_vector(vector([0.0, 0.0, 1.0])) - accel.normalize()
+        cost = 0.5 * objective.dot(objective)
+        gradient = cp.grad(cost, q).normalize()
+
+        # Quaternion derivative from gyroscope measurements
+        gyro_quat = cp.quaternion(0.0, *gyro)
+        q_dot_gyro = 0.5 * (q @ gyro_quat)
+
+        # Update quaternion using gradient descent
+        q_dot = q_dot_gyro - 0.1 * gradient
+        
+        return (q + q_dot * dt).normalize()
+    
+    q: quaternion = quaternion(cp.value(0.7071), cp.value(0.7071), cp.value(0.0), cp.value(0.0))  # Initial orientation (45 degrees around X-axis)
+    gyro = vector([0.01, 0.02, 0.015])
+    accel = vector([0.0, 0.0, 1.0])
+
+    new_q = update_orientation(q, gyro, accel)
+
+    tg = Target()
+    tg.compile(new_q)
+    tg.run()
+
+    new_q_value = tg.read_value(new_q)
+
+    assert pytest.approx(new_q_value[0], abs=1e-4) == 0.7072948217391968
+    assert pytest.approx(new_q_value[1], abs=1e-4) == 0.7069186568260193
+    assert pytest.approx(new_q_value[2], abs=1e-4) == 0.7660913727013394e-05
+    assert pytest.approx(new_q_value[3], abs=1e-4) == 0.00012362639245111495