matrix class added

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",

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)
+    )