diff --git a/src/copapy/__init__.py b/src/copapy/__init__.py index 2670500..e3fd7d7 100644 --- a/src/copapy/__init__.py +++ b/src/copapy/__init__.py @@ -1,6 +1,7 @@ from ._target import Target from ._basic_types import NumLike, variable, generic_sdb, iif from ._vectors import vector, distance, scalar_projection, angle_between, rotate_vector, vector_projection +from ._matrices import matrix, identity, zeros, ones, diagonal from ._math import sqrt, abs, sin, cos, tan, asin, acos, atan, atan2, log, exp, pow, get_42, clamp, min, max __all__ = [ @@ -10,6 +11,11 @@ __all__ = [ "generic_sdb", "iif", "vector", + "matrix", + "identity", + "zeros", + "ones", + "diagonal", "sqrt", "abs", "sin", diff --git a/src/copapy/_matrices.py b/src/copapy/_matrices.py new file mode 100644 index 0000000..11ecc7d --- /dev/null +++ b/src/copapy/_matrices.py @@ -0,0 +1,291 @@ +from . import variable +from ._vectors import vector +from ._mixed import mixed_sum +from typing import Generic, TypeVar, Iterable, Any, overload, TypeAlias, Callable, Iterator + +MatNumLike: TypeAlias = 'matrix[int] | matrix[float] | variable[int] | variable[float] | int | float' +MatIntLike: TypeAlias = 'matrix[int] | variable[int] | int' +MatFloatLike: TypeAlias = 'matrix[float] | variable[float] | float' +TT = TypeVar("TT", int, float) +U = TypeVar("U", int, float) + + +class matrix(Generic[TT]): + """Mathematical matrix class supporting basic operations and interactions with variables. + """ + def __init__(self, values: Iterable[Iterable[TT | variable[TT]]]): + """Create a matrix with given values and variables. + + Args: + values: iterable of iterable of constant values and variables + """ + rows = tuple(tuple(row) for row in values) + if rows: + row_len = len(rows[0]) + assert all(len(row) == row_len for row in rows), "All rows must have the same length" + self.values: tuple[tuple[variable[TT] | TT, ...], ...] = tuple(rows) + self.rows = len(self.values) + self.cols = len(self.values[0]) if self.values else 0 + + def __repr__(self) -> str: + return f"matrix({self.values})" + + def __len__(self) -> int: + return self.rows + + def __getitem__(self, index: int) -> tuple[variable[TT] | TT, ...]: + return self.values[index] + + def __iter__(self) -> Iterator[tuple[variable[TT] | TT, ...]]: + return iter(self.values) + + def __neg__(self) -> 'matrix[TT]': + return matrix((-a for a in row) for row in self.values) + + @overload + def __add__(self: 'matrix[int]', other: MatFloatLike) -> 'matrix[float]': ... + @overload + def __add__(self: 'matrix[int]', other: MatIntLike) -> 'matrix[int]': ... + @overload + def __add__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ... + @overload + def __add__(self, other: MatNumLike) -> 'matrix[int] | matrix[float]': ... + def __add__(self, other: MatNumLike) -> Any: + if isinstance(other, matrix): + assert self.rows == other.rows and self.cols == other.cols, \ + "Matrices must have the same dimensions" + return matrix( + tuple(a + b for a, b in zip(row1, row2)) + for row1, row2 in zip(self.values, other.values) + ) + return matrix( + tuple(a + other for a in row) + for row in self.values + ) + + @overload + def __radd__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ... + @overload + def __radd__(self: 'matrix[int]', other: variable[int] | int) -> 'matrix[int]': ... + def __radd__(self, other: Any) -> Any: + return self + other + + @overload + def __sub__(self: 'matrix[int]', other: MatFloatLike) -> 'matrix[float]': ... + @overload + def __sub__(self: 'matrix[int]', other: MatIntLike) -> 'matrix[int]': ... + @overload + def __sub__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ... + @overload + def __sub__(self, other: MatNumLike) -> 'matrix[int] | matrix[float]': ... + def __sub__(self, other: MatNumLike) -> Any: + if isinstance(other, matrix): + assert self.rows == other.rows and self.cols == other.cols, \ + "Matrices must have the same dimensions" + return matrix( + tuple(a - b for a, b in zip(row1, row2)) + for row1, row2 in zip(self.values, other.values) + ) + return matrix( + tuple(a - other for a in row) + for row in self.values + ) + + @overload + def __rsub__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ... + @overload + def __rsub__(self: 'matrix[int]', other: variable[int] | int) -> 'matrix[int]': ... + def __rsub__(self, other: MatNumLike) -> Any: + if isinstance(other, matrix): + assert self.rows == other.rows and self.cols == other.cols, \ + "Matrices must have the same dimensions" + return matrix( + tuple(b - a for a, b in zip(row1, row2)) + for row1, row2 in zip(self.values, other.values) + ) + return matrix( + tuple(other - a for a in row) + for row in self.values + ) + + @overload + def __mul__(self: 'matrix[int]', other: MatFloatLike) -> 'matrix[float]': ... + @overload + def __mul__(self: 'matrix[int]', other: MatIntLike) -> 'matrix[int]': ... + @overload + def __mul__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ... + @overload + def __mul__(self, other: MatNumLike) -> 'matrix[int] | matrix[float]': ... + def __mul__(self, other: MatNumLike) -> Any: + """Element-wise multiplication""" + if isinstance(other, matrix): + assert self.rows == other.rows and self.cols == other.cols, \ + "Matrices must have the same dimensions" + return matrix( + tuple(a * b for a, b in zip(row1, row2)) + for row1, row2 in zip(self.values, other.values) + ) + return matrix( + tuple(a * other for a in row) + for row in self.values + ) + + @overload + def __rmul__(self: 'matrix[float]', other: MatNumLike) -> 'matrix[float]': ... + @overload + def __rmul__(self: 'matrix[int]', other: variable[int] | int) -> 'matrix[int]': ... + def __rmul__(self, other: MatNumLike) -> Any: + return self * other + + def __truediv__(self, other: MatNumLike) -> 'matrix[float]': + """Element-wise division""" + if isinstance(other, matrix): + assert self.rows == other.rows and self.cols == other.cols, \ + "Matrices must have the same dimensions" + return matrix( + tuple(a / b for a, b in zip(row1, row2)) + for row1, row2 in zip(self.values, other.values) + ) + return matrix( + tuple(a / other for a in row) + for row in self.values + ) + + def __rtruediv__(self, other: MatNumLike) -> 'matrix[float]': + if isinstance(other, matrix): + assert self.rows == other.rows and self.cols == other.cols, \ + "Matrices must have the same dimensions" + return matrix( + tuple(b / a for a, b in zip(row1, row2)) + for row1, row2 in zip(self.values, other.values) + ) + return matrix( + tuple(other / a for a in row) + for row in self.values + ) + + @overload + def __matmul__(self: 'matrix[TT]', other: 'vector[TT]') -> 'vector[TT]': ... + @overload + def __matmul__(self: 'matrix[TT]', other: 'matrix[TT]') -> 'matrix[TT]': ... + def __matmul__(self: 'matrix[TT]', other: 'matrix[TT] | vector[TT]') -> 'matrix[TT] | vector[TT]': + """Matrix multiplication using @ operator""" + if isinstance(other, vector): + assert self.cols == len(other.values), \ + f"Matrix columns ({self.cols}) must match vector length ({len(other.values)})" + vec_result = (mixed_sum(a * b for a, b in zip(row, other.values)) for row in self.values) + return vector(vec_result) + else: + assert isinstance(other, matrix), "Cannot multiply matrix with {type(other)}" + assert self.cols == other.rows, \ + f"Matrix columns ({self.cols}) must match other matrix rows ({other.rows})" + result: list[list[TT | variable[TT]]] = [] + for row in self.values: + new_row: list[TT | variable[TT]] = [] + for col_idx in range(other.cols): + col = tuple(other.values[i][col_idx] for i in range(other.rows)) + element = sum(a * b for a, b in zip(row, col)) + new_row.append(element) + result.append(new_row) + return matrix(result) + + def transpose(self) -> 'matrix[TT]': + """Return the transpose of the matrix.""" + if not self.values: + return matrix([]) + return matrix( + tuple(self.values[i][j] for i in range(self.rows)) + for j in range(self.cols) + ) + + @property + def T(self): + return self.transpose() + + def row(self, index: int) -> vector[TT]: + """Get a row as a vector.""" + assert 0 <= index < self.rows, f"Row index {index} out of bounds" + return vector(self.values[index]) + + def col(self, index: int) -> vector[TT]: + """Get a column as a vector.""" + assert 0 <= index < self.cols, f"Column index {index} out of bounds" + return vector(self.values[i][index] for i in range(self.rows)) + + @overload + def trace(self: 'matrix[TT]') -> TT | variable[TT]: ... + @overload + def trace(self: 'matrix[int]') -> int | variable[int]: ... + @overload + def trace(self: 'matrix[float]') -> float | variable[float]: ... + def trace(self) -> Any: + """Calculate the trace (sum of diagonal elements).""" + assert self.rows == self.cols, "Trace is only defined for square matrices" + return mixed_sum(self.values[i][i] for i in range(self.rows)) + + @overload + def sum(self: 'matrix[TT]') -> TT | variable[TT]: ... + @overload + def sum(self: 'matrix[int]') -> int | variable[int]: ... + @overload + def sum(self: 'matrix[float]') -> float | variable[float]: ... + def sum(self) -> Any: + """Calculate the sum of all elements.""" + return mixed_sum(a for row in self.values for a in row) + + def map(self, func: Callable[[Any], variable[U] | U]) -> 'matrix[U]': + """Applies a function to each element of the matrix and returns a new matrix.""" + return matrix( + tuple(func(a) for a in row) + for row in self.values + ) + + def homogenize(self) -> 'matrix[TT]': + """Convert all elements to variables if any element is a variable.""" + if any(isinstance(val, variable) for row in self.values for val in row): + return matrix( + tuple(variable(val) if not isinstance(val, variable) else val for val in row) + for row in self.values + ) + else: + return self + + +# Utility functions for matrices + +def identity(size: int) -> matrix[int]: + """Create an identity matrix of given size.""" + return matrix( + tuple(1 if i == j else 0 for j in range(size)) + for i in range(size) + ) + + +def zeros(rows: int, cols: int) -> matrix[int]: + """Create a zero matrix of given dimensions.""" + return matrix( + tuple(0 for _ in range(cols)) + for _ in range(rows) + ) + + +def ones(rows: int, cols: int) -> matrix[int]: + """Create a matrix of ones with given dimensions.""" + return matrix( + tuple(1 for _ in range(cols)) + for _ in range(rows) + ) + + +@overload +def diagonal(vec: 'vector[int]') -> matrix[int]: ... +@overload +def diagonal(vec: 'vector[float]') -> matrix[float]: ... +def diagonal(vec: vector[Any]) -> matrix[Any]: + """Create a diagonal matrix from a vector.""" + size = len(vec) + + return matrix( + tuple(vec[i] if i == j else 0 for j in range(size)) + for i in range(size) + )