mirror of https://github.com/Nonannet/copapy.git
matrix class added
This commit is contained in:
Nicolas Kruse
parent
29f5a26278
commit
e14092a2c3
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
)
|
||||
Loading…
Reference in New Issue