diff --git a/compute.py b/compute.py index 10862fd..0f42a61 100644 --- a/compute.py +++ b/compute.py @@ -76,7 +76,7 @@ def matmul( -def get_rotate(matrix: list[list[float]], angle:int) -> list[list[float]]: +def rot_2D(matrix: list[list[float]], angle:int) -> list[list[float]]: """Apply a rotation matrix to a base matrix. Parameters @@ -103,6 +103,43 @@ def get_rotate(matrix: list[list[float]], angle:int) -> list[list[float]]: print("Can't rotate this matrix! Wrong amount of rows!") return [[0, 0],[0, 0]] + +def rot_3D(angle:float, axis: str) -> list[list[float]]: + """Apply a rotation matrix to a base matrix. + + Parameters + ---------- + angle: float + The angle to be rotated. + axis: str + The axis as a string. + + Returns + ------- + list[list[float]] + Rotation matrix for *angle* by *axis*. + """ + + if axis == "x": + row_1 = [1, 0, 0] + row_2 = [0, math.cos(angle), -math.sin(angle)] + row_3 = [0, math.sin(angle), math.cos(angle)] + elif axis == "y": + row_1 = [math.cos(angle), 0, math.sin(angle)] + row_2 = [0, 1, 0] + row_3 = [-math.sin(angle), 0, math.cos(angle)] + elif axis == "z": + row_1 = [math.cos(angle), -math.sin(angle), 0] + row_2 = [math.sin(angle), math.cos(angle), 0] + row_3 = [0, 0, 1] + else: + print("Can't rotate this matrix! Wrong axis!") + row_1 = [0, 0, 0] + row_2 = [0, 0, 0] + row_3 = [0, 0, 0] + return [row_1, row_2, row_3] + + def transpose(matrix: list[list[float]]) -> list[list[float]]: """Return the transpose of a 2‑D list.