rot_3d added

matrixmania/__init__.py

@@ -1 +1 @@
-from .compute import matmul, transpose, rot_2D
+from .compute import matmul, transpose, rot_2D, rot_3D

matrixmania/compute.py

@@ -66,6 +66,59 @@ def rot_2D(theta: float) -> List[List[float]]:
         [sin_t,  cos_t]
     ]
 
+def rot_3D(angle: float, axis: str) -> List[List[float]]:
+    """
+    Returns a 3D rotation matrix for a given angle and axis.
+
+    Rotates points in 3D space around the specified axis by the given angle.
+    The rotation follows the right-hand rule (counterclockwise when looking
+    along the axis of rotation).
+
+    Rotation matrices:
+        x-axis:
+            [[1,      0,       0],
+             [0,  cosθ,  -sinθ],
+             [0,  sinθ,   cosθ]]
+
+        y-axis:
+            [[ cosθ, 0, sinθ],
+             [    0, 1,    0],
+             [-sinθ, 0, cosθ]]
+
+        z-axis:
+            [[cosθ, -sinθ, 0],
+             [sinθ,  cosθ, 0],
+             [   0,     0, 1]]
+
+    :param angle: Rotation angle in radians.
+    :param axis: Axis of rotation ('x', 'y', or 'z').
+    :return: A 3×3 rotation matrix as a list of lists.
+    :raises ValueError: If the axis is not 'x', 'y', or 'z'.
+    """
+    cos_t = math.cos(angle)
+    sin_t = math.sin(angle)
+
+    if axis == 'x':
+        return [
+            [1, 0, 0],
+            [0, cos_t, -sin_t],
+            [0, sin_t, cos_t]
+        ]
+    elif axis == 'y':
+        return [
+            [cos_t, 0, sin_t],
+            [0, 1, 0],
+            [-sin_t, 0, cos_t]
+        ]
+    elif axis == 'z':
+        return [
+            [cos_t, -sin_t, 0],
+            [sin_t, cos_t, 0],
+            [0, 0, 1]
+        ]
+    else:
+        raise ValueError("Axis must be 'x', 'y', or 'z'.")
+
 
 def main():
     # Test matrix multiplication