Data commit

Task/Strassens-algorithm/Python/strassens-algorithm.py

@@ -0,0 +1,166 @@
+"""Matrix multiplication using Strassen's algorithm. Requires Python >= 3.7."""
+
+from __future__ import annotations
+from itertools import chain
+from typing import List
+from typing import NamedTuple
+from typing import Optional
+
+
+class Shape(NamedTuple):
+    rows: int
+    cols: int
+
+
+class Matrix(List):
+    """A matrix implemented as a two-dimensional list."""
+
+    @classmethod
+    def block(cls, blocks) -> Matrix:
+        """Return a new Matrix assembled from nested blocks."""
+        m = Matrix()
+        for hblock in blocks:
+            for row in zip(*hblock):
+                m.append(list(chain.from_iterable(row)))
+
+        return m
+
+    def dot(self, b: Matrix) -> Matrix:
+        """Return a new Matrix that is the product of this matrix and matrix `b`.
+        Uses 'simple' or 'naive' matrix multiplication."""
+        assert self.shape.cols == b.shape.rows
+        m = Matrix()
+        for row in self:
+            new_row = []
+            for c in range(len(b[0])):
+                col = [b[r][c] for r in range(len(b))]
+                new_row.append(sum(x * y for x, y in zip(row, col)))
+            m.append(new_row)
+        return m
+
+    def __matmul__(self, b: Matrix) -> Matrix:
+        return self.dot(b)
+
+    def __add__(self, b: Matrix) -> Matrix:
+        """Matrix addition."""
+        assert self.shape == b.shape
+        rows, cols = self.shape
+        return Matrix(
+            [[self[i][j] + b[i][j] for j in range(cols)] for i in range(rows)]
+        )
+
+    def __sub__(self, b: Matrix) -> Matrix:
+        """Matrix subtraction."""
+        assert self.shape == b.shape
+        rows, cols = self.shape
+        return Matrix(
+            [[self[i][j] - b[i][j] for j in range(cols)] for i in range(rows)]
+        )
+
+    def strassen(self, b: Matrix) -> Matrix:
+        """Return a new Matrix that is the product of this matrix and matrix `b`.
+        Uses strassen algorithm."""
+        rows, cols = self.shape
+
+        assert rows == cols, "matrices must be square"
+        assert self.shape == b.shape, "matrices must be the same shape"
+        assert rows and (rows & rows - 1) == 0, "shape must be a power of 2"
+
+        if rows == 1:
+            return self.dot(b)
+
+        p = rows // 2  # partition
+
+        a11 = Matrix([n[:p] for n in self[:p]])
+        a12 = Matrix([n[p:] for n in self[:p]])
+        a21 = Matrix([n[:p] for n in self[p:]])
+        a22 = Matrix([n[p:] for n in self[p:]])
+
+        b11 = Matrix([n[:p] for n in b[:p]])
+        b12 = Matrix([n[p:] for n in b[:p]])
+        b21 = Matrix([n[:p] for n in b[p:]])
+        b22 = Matrix([n[p:] for n in b[p:]])
+
+        m1 = (a11 + a22).strassen(b11 + b22)
+        m2 = (a21 + a22).strassen(b11)
+        m3 = a11.strassen(b12 - b22)
+        m4 = a22.strassen(b21 - b11)
+        m5 = (a11 + a12).strassen(b22)
+        m6 = (a21 - a11).strassen(b11 + b12)
+        m7 = (a12 - a22).strassen(b21 + b22)
+
+        c11 = m1 + m4 - m5 + m7
+        c12 = m3 + m5
+        c21 = m2 + m4
+        c22 = m1 - m2 + m3 + m6
+
+        return Matrix.block([[c11, c12], [c21, c22]])
+
+    def round(self, ndigits: Optional[int] = None) -> Matrix:
+        return Matrix([[round(i, ndigits) for i in row] for row in self])
+
+    @property
+    def shape(self) -> Shape:
+        cols = len(self[0]) if self else 0
+        return Shape(len(self), cols)
+
+
+def examples():
+    a = Matrix(
+        [
+            [1, 2],
+            [3, 4],
+        ]
+    )
+    b = Matrix(
+        [
+            [5, 6],
+            [7, 8],
+        ]
+    )
+    c = Matrix(
+        [
+            [1, 1, 1, 1],
+            [2, 4, 8, 16],
+            [3, 9, 27, 81],
+            [4, 16, 64, 256],
+        ]
+    )
+    d = Matrix(
+        [
+            [4, -3, 4 / 3, -1 / 4],
+            [-13 / 3, 19 / 4, -7 / 3, 11 / 24],
+            [3 / 2, -2, 7 / 6, -1 / 4],
+            [-1 / 6, 1 / 4, -1 / 6, 1 / 24],
+        ]
+    )
+    e = Matrix(
+        [
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+            [9, 10, 11, 12],
+            [13, 14, 15, 16],
+        ]
+    )
+    f = Matrix(
+        [
+            [1, 0, 0, 0],
+            [0, 1, 0, 0],
+            [0, 0, 1, 0],
+            [0, 0, 0, 1],
+        ]
+    )
+
+    print("Naive matrix multiplication:")
+    print(f"  a * b = {a @ b}")
+    print(f"  c * d = {(c @ d).round()}")
+    print(f"  e * f = {e @ f}")
+
+    print("Strassen's matrix multiplication:")
+    print(f"  a * b = {a.strassen(b)}")
+    print(f"  c * d = {c.strassen(d).round()}")
+    print(f"  e * f = {e.strassen(f)}")
+
+
+if __name__ == "__main__":
+    examples()