148 lines
4.1 KiB
Python
148 lines
4.1 KiB
Python
import math
|
|
|
|
from tabulate import tabulate
|
|
|
|
|
|
def matmul(a: list[list[float]], b: list[list[float]]) -> list[list[float]]:
|
|
"""Multiply two non-empty rectangular matrices.
|
|
|
|
Args:
|
|
a: The left matrix.
|
|
b: The right matrix.
|
|
|
|
Returns:
|
|
The matrix product a * b.
|
|
|
|
Raises:
|
|
ValueError: If either matrix is empty or ragged, or if their
|
|
dimensions are incompatible.
|
|
"""
|
|
_validate_matrix(a, "a")
|
|
_validate_matrix(b, "b")
|
|
|
|
if len(a[0]) != len(b):
|
|
raise ValueError(
|
|
"Matrix dimensions are incompatible: the number of columns in a "
|
|
"must equal the number of rows in b."
|
|
)
|
|
|
|
return [
|
|
[
|
|
sum(a_row[index] * b[index][column] for index in range(len(b)))
|
|
for column in range(len(b[0]))
|
|
]
|
|
for a_row in a
|
|
]
|
|
|
|
|
|
def transpose(matrix: list[list[float]]) -> list[list[float]]:
|
|
"""Return the transpose of a non-empty rectangular matrix.
|
|
|
|
Args:
|
|
matrix: A matrix represented as rows of numbers.
|
|
|
|
Returns:
|
|
A new matrix whose rows are the columns of matrix.
|
|
|
|
Raises:
|
|
ValueError: If the matrix is empty or its rows have different lengths.
|
|
"""
|
|
_validate_matrix(matrix, "matrix")
|
|
return [list(column) for column in zip(*matrix)]
|
|
|
|
|
|
def rot_2D(angle: float) -> list[list[float]]:
|
|
"""Return the 2D rotation matrix for a given angle in radians.
|
|
|
|
Args:
|
|
angle: The rotation angle in radians.
|
|
|
|
Returns:
|
|
The rotation matrix for rotating points counterclockwise.
|
|
"""
|
|
return [[math.cos(angle), -math.sin(angle)], [math.sin(angle), math.cos(angle)]]
|
|
|
|
|
|
def rot_3D(angle: float, axis: str) -> list[list[float]]:
|
|
"""Return a 3D rotation matrix around the x, y, or z axis.
|
|
|
|
With ``c = cos(angle)`` and ``s = sin(angle)``, the matrices are::
|
|
|
|
x: [[1, 0, 0], [0, c, -s], [0, s, c]]
|
|
y: [[c, 0, s], [0, 1, 0], [-s, 0, c]]
|
|
z: [[c, -s, 0], [s, c, 0], [0, 0, 1]]
|
|
|
|
Args:
|
|
angle: The rotation angle in radians.
|
|
axis: The axis to rotate around: ``"x"``, ``"y"``, or ``"z"``.
|
|
|
|
Returns:
|
|
The 3D rotation matrix for the selected axis.
|
|
|
|
Raises:
|
|
ValueError: If axis is not ``"x"``, ``"y"`, or ``"z"``.
|
|
"""
|
|
cosine = math.cos(angle)
|
|
sine = math.sin(angle)
|
|
|
|
if axis == "x":
|
|
return [[1.0, 0.0, 0.0], [0.0, cosine, -sine], [0.0, sine, cosine]]
|
|
if axis == "y":
|
|
return [[cosine, 0.0, sine], [0.0, 1.0, 0.0], [-sine, 0.0, cosine]]
|
|
if axis == "z":
|
|
return [[cosine, -sine, 0.0], [sine, cosine, 0.0], [0.0, 0.0, 1.0]]
|
|
|
|
raise ValueError("axis must be 'x', 'y', or 'z'")
|
|
|
|
|
|
def project_ortho(
|
|
point: tuple[float, float, float], scale: float = 1
|
|
) -> tuple[float, float]:
|
|
"""Project a 3D point onto the xy-plane and apply a scale factor.
|
|
|
|
The z-coordinate is ignored. The returned x- and y-coordinates are each
|
|
multiplied by ``scale``.
|
|
|
|
Args:
|
|
point: The 3D point as an (x, y, z) tuple.
|
|
scale: The scale factor to apply to x and y.
|
|
|
|
Returns:
|
|
The scaled 2D point as an (x, y) tuple.
|
|
"""
|
|
return point[0] * scale, point[1] * scale
|
|
|
|
|
|
def _validate_matrix(matrix: list[list[float]], name: str) -> None:
|
|
"""Raise ValueError unless matrix is non-empty and rectangular."""
|
|
if not matrix:
|
|
raise ValueError(f"Matrix {name} must be non-empty.")
|
|
|
|
column_count = len(matrix[0])
|
|
|
|
if column_count == 0:
|
|
raise ValueError(f"Matrix {name} must have at least one column.")
|
|
|
|
if any(len(row) != column_count for row in matrix):
|
|
raise ValueError(f"Matrix {name} must have rows of equal length.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
first = [[1, 2, 3], [4, 5, 6]]
|
|
second = [[7, 8], [9, 10], [11.0, 12.0]]
|
|
|
|
assert matmul(first, second) == [[58.0, 64.0], [139.0, 154.0]]
|
|
print("Matrix product:", matmul(first, second))
|
|
|
|
try:
|
|
matmul([[1, 2]], [[1, 2]])
|
|
except ValueError as error:
|
|
print("Expected error:", error)
|
|
else:
|
|
raise AssertionError("Expected a ValueError for incompatible dimensions.")
|
|
|
|
assert transpose(first) == [[1, 4], [2, 5], [3, 6]]
|
|
print("Transpose:", transpose(first))
|
|
|
|
print(tabulate(rot_2D(1.0)))
|