Sparse Matrix class implementation without using numpy or scipy. sparse matrix represented using dictionary, tuple of matrix indices as keys and matrix elements themselves as values.

sparsematrix.py

import math
from typing import Dict, Tuple, Optional, List

class SparseMatrix:
    """Sparse Matrix class
    Args:
        data: A dictionary of (i, j) -> value
    """

    def __init__(self, data: Dict[Tuple[int, int], float] = None):
        self.data = data or {}

    def from_dense_matrix(self, matrix: List[List[float]]) -> "SparseMatrix":
        for i in range(len(matrix)):
            for j in range(len(matrix[0])):
                if matrix[i][j] != 0:
                    self[(i, j)] = matrix[i][j]
        # [self[(i, j)] for i in range(len(matrix)) for j in range(len(matrix[0])) if matrix[i][j] != 0]
        return self

    def __getitem__(self, key: Tuple[int, int]) -> float:
        return self.data.get(key, 0)

    def __setitem__(self, key: Tuple[int, int], value: float):
        self.data[key] = value

    def __delitem__(self, key: Tuple[int, int]):
        del self.data[key]

    def __iter__(self):
        return iter(self.data)

    def __len__(self):
        return len(self.data)

    def __str__(self):
        return str(self.data)

    def __repr__(self):
        return repr(self.data)

    def __add__(self, other: "SparseMatrix") -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] += self[key]
        for key in other:
            result[key] += other[key]
        return result
        # res = SparseMatrix({k: v + other.data.get(k, 0) for k, v in self.data.items()})
        # return SparseMatrix({k: other.data.get(k, 0) for k, v in other.data.items()})

    def __sub__(self, other: "SparseMatrix") -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = self[key]
        for key in other:
            result[key] -= other[key]
        return result

    # return SparseMatrix({k: v - other.data.get(k, 0) for k, v in self.data.items()})

    def __mul__(self, other: "SparseMatrix") -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            for key2 in other:
                result[(key[0], key2[1])] += self[key] * other[key2]
        return result

    def __rmul__(self, other: float) -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = self[key] * other
        return result

    def __truediv__(self, other: float) -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = self[key] / other
        return result

    def __neg__(self) -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = -self[key]
        return result

    def __pow__(self, other: float) -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = self[key] ** other
        return result

    def __rpow__(self, other: float) -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = other ** self[key]
        return result

    def __eq__(self, other: "SparseMatrix") -> bool:
        return self.data == other.data

    def __ne__(self, other: "SparseMatrix") -> bool:
        return self.data != other.data

    def __lt__(self, other: "SparseMatrix") -> bool:
        return self.data < other.data

    def __le__(self, other: "SparseMatrix") -> bool:
        return self.data <= other.data

    def __gt__(self, other: "SparseMatrix") -> bool:
        return self.data > other.data

    def __ge__(self, other: "SparseMatrix") -> bool:
        return self.data >= other.data

    def __abs__(self) -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = abs(self[key])
        return result

    def __round__(self) -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = round(self[key])
        return result

    def __floor__(self) -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = math.floor(self[key])
        return result

    def __ceil__(self) -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = math.ceil(self[key])
        return result

    def __trunc__(self) -> "SparseMatrix":
        result = SparseMatrix()
        for key in self:
            result[key] = math.trunc(self[key])
        return result


if __name__ == "__main__":
    a = SparseMatrix({(1, 2): 3.5, (2, 3): 4.0})
    b = SparseMatrix({(1, 3): 5.0, (2, 3): 6.0})
    print(a + b)
    print(a - b)
    test_dens_sparse = SparseMatrix()
    dense_matrix = [
        [0, 0, 0, 0, 0, 9],
        [0, 0, 0, 0, 9, 0],
        [0, 0, 2, 0, 0, 0],
        [1, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 3, 0],
        [0, 0, 0, 0, 3, 0],
    ]
    print(test_dens_sparse.from_dense_matrix(matrix=dense_matrix))
    # print(a * b)
    # print(a * 3)