Skip to main content
Became Hot Network Question
edited tags
Link
Reinderien
  • 65.6k
  • 5
  • 70
  • 190
Source Link
Linny
  • 10.4k
  • 5
  • 37
  • 99

Python matrix class

To prepare for my studies in linear algebra, I've decided to write a simple matrix class. This is the initial program, with functions to find the determinant and decomposition(s) to follow. I wanted to get feedback first on my approach.

from __future__ import annotations
import operator
import copy

# Don't use [[0] * width] * height, TRAP!
# https://stackoverflow.com/a/44382900/8968906
def fromdim(width: int, height: int):
    return Matrix([[0] * width for _ in range(height)])

class Matrix:
    
    def __init__(self, arrays: list[list]) -> None:
        self.arrays = copy.deepcopy(arrays)
        self.width = len(arrays[0])
        self.height = len(arrays)

    # Public member functions

    def cols(self) -> list[list[int | float]]:
        """
        Returns a list of columns of this matrix.
        """
        return [list(x) for x in zip(*self.arrays)]
    
    def tpose(self) -> None:
        """
        Transposes the current Matrix, modifying the underlying arrays object. This
        converts a MxN array into an NxM array.

        >>> Matrix([[1, 2, 3], [4, 5, 6]])
        [
            [1, 4]
            [2, 5]
            [3, 6]
        ]
        """
        res = fromdim(self.height, self.width)
        for idx, c in enumerate(self.cols()):
            res[idx] = c
        self.arrays = res.arrays
    
    # Dunder methods

    def __add__(self, other: Matrix | float | int) -> Matrix:
        if isinstance(other, (int, float)):
            return Matrix([
                list(map(operator.add, self.arrays[rdx], [other] * self.width))
                for rdx in range(self.height)
            ])
        if isinstance(other, Matrix) and ((self.width, self.height) == (other.width, other.height)):
            return Matrix([
                list(map(operator.add, self.arrays[rdx], other.arrays[rdx]))
                for rdx in range(self.height)
            ])

    def __sub__(self, other: Matrix | float | int) -> Matrix:
        if isinstance(other, (int, float)):
            return self.__add__(-other)
        if isinstance(other, Matrix) and ((self.width, self.height) == (other.width, other.height)):
            return Matrix([
                list(map(operator.sub, self.arrays[rdx], other.arrays[rdx]))
                for rdx in range(self.height)
            ])
        
    def __mul__(self, other: Matrix | float | int) -> Matrix:
        if isinstance(other, (int, float)):
            return Matrix([
                list(map(operator.mul, self.arrays[rdx], [other] * self.width))
                for rdx in range(self.height)
            ])
        if isinstance(other, Matrix) and ((self.width, self.height) == (other.width, other.height)):
            result = Matrix(self.arrays)
            for rdx, row in enumerate(self.arrays):
                for cdx, col in enumerate(other.cols()):
                    result[rdx][cdx] = sum(r * c for r, c in zip(row, col))
            return result
    
    def __repr__(self) -> str:
        return '[\n' + '\n'.join([f"  {arr}" for arr in self.arrays]) + '\n]'

    def __neg__(self) -> Matrix:
        result = fromdim(self.width, self.height)
        for x in range(self.width):
            for y in range(self.height):
                result.arrays[x][y] = -self.arrays[x][y]
        return result
    
    def __getitem__(self, idx: int) -> list:
        return self.arrays[idx]
    
    def __setitem__(self, idx: int, item: int | float) -> None:
        self.arrays[idx] = item

Testing code:

from PyMatrix import Matrix

m = Matrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
n = Matrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
o = Matrix([[1, 2, 3], [4, 5, 6]])

print("M:", m)
print("N:", n)
print("M + N:", m + n)
print("M - N:", m - n)
print("M * N:", m * n)
o.tpose()
print(o)
print(o.cols())