neuer commit

matrixmania/compute.py

@@ -49,9 +49,37 @@ def transpose(matrix:List[List[int]]) -> List[List[int]]:
 
 
 def rot_2D(angle: float) -> List[List[float]]:
-    rot_matrix = [[math.cos(angle), -math.sin(angle)],
-                  [math.sin(angle), math.cos(angle)]]
-    return rot_matrix
+    '''
+    calculates the rotation matrix in 2D
+    :param angle: the angle
+    :return: rotation matrix
+    '''
+    return [[math.cos(angle), -math.sin(angle)],
+            [math.sin(angle), math.cos(angle)]]
+
+def rot_3D(angle: float, axis: str) -> List[List[float]]:
+    '''
+    calculates the rotation matrix for one 3D axis
+    :param angle: the angle
+    :param axis: the axis
+    :return: rotation matrix for the wanted axis
+    '''
+    axis = axis.lower()
+    if axis == "x":
+        return [[1, 0, 0],
+                [0, math.cos(angle), -math.sin(angle)],
+                [0, math.sin(angle), math.cos(angle)]]
+    elif axis == "y":
+        return [[math.cos(angle), 0, math.sin(angle)],
+                [0, 1, 0],
+                [-math.sin(angle), 0, math.cos(angle)]]
+    elif axis == "z":
+        return [[math.cos(angle), -math.sin(angle), 0],
+                [math.sin(angle), math.cos(angle), 0],
+                [0, 0, 1]]
+    else:
+        raise ValueError("Axis not valid")
+
 
 
 if __name__ == "__main__":