Task: Strassen's algorithm




Go: package main

import (
    "fmt"
    "log"
    "math"
)

type Matrix [][]float64

func (m Matrix) rows() int { return len(m) }
func (m Matrix) cols() int { return len(m[0]) }

func (m Matrix) add(m2 Matrix) Matrix {
    if m.rows() != m2.rows() || m.cols() != m2.cols() {
        log.Fatal("Matrices must have the same dimensions.")
    }
    c := make(Matrix, m.rows())
    for i := 0; i < m.rows(); i++ {
        c[i] = make([]float64, m.cols())
        for j := 0; j < m.cols(); j++ {
            c[i][j] = m[i][j] + m2[i][j]
        }
    }
    return c
}

func (m Matrix) sub(m2 Matrix) Matrix {
    if m.rows() != m2.rows() || m.cols() != m2.cols() {
        log.Fatal("Matrices must have the same dimensions.")
    }
    c := make(Matrix, m.rows())
    for i := 0; i < m.rows(); i++ {
        c[i] = make([]float64, m.cols())
        for j := 0; j < m.cols(); j++ {
            c[i][j] = m[i][j] - m2[i][j]
        }
    }
    return c
}

func (m Matrix) mul(m2 Matrix) Matrix {
    if m.cols() != m2.rows() {
        log.Fatal("Cannot multiply these matrices.")
    }
    c := make(Matrix, m.rows())
    for i := 0; i < m.rows(); i++ {
        c[i] = make([]float64, m2.cols())
        for j := 0; j < m2.cols(); j++ {
            for k := 0; k < m2.rows(); k++ {
                c[i][j] += m[i][k] * m2[k][j]
            }
        }
    }
    return c
}

func (m Matrix) toString(p int) string {
    s := make([]string, m.rows())
    pow := math.Pow(10, float64(p))
    for i := 0; i < m.rows(); i++ {
        t := make([]string, m.cols())
        for j := 0; j < m.cols(); j++ {
            r := math.Round(m[i][j]*pow) / pow
            t[j] = fmt.Sprintf("%g", r)
            if t[j] == "-0" {
                t[j] = "0"
            }
        }
        s[i] = fmt.Sprintf("%v", t)
    }
    return fmt.Sprintf("%v", s)
}

func params(r, c int) [4][6]int {
    return [4][6]int{
        {0, r, 0, c, 0, 0},
        {0, r, c, 2 * c, 0, c},
        {r, 2 * r, 0, c, r, 0},
        {r, 2 * r, c, 2 * c, r, c},
    }
}

func toQuarters(m Matrix) [4]Matrix {
    r := m.rows() / 2
    c := m.cols() / 2
    p := params(r, c)
    var quarters [4]Matrix
    for k := 0; k < 4; k++ {
        q := make(Matrix, r)
        for i := p[k][0]; i < p[k][1]; i++ {
            q[i-p[k][4]] = make([]float64, c)
            for j := p[k][2]; j < p[k][3]; j++ {
                q[i-p[k][4]][j-p[k][5]] = m[i][j]
            }
        }
        quarters[k] = q
    }
    return quarters
}

func fromQuarters(q [4]Matrix) Matrix {
    r := q[0].rows()
    c := q[0].cols()
    p := params(r, c)
    r *= 2
    c *= 2
    m := make(Matrix, r)
    for i := 0; i < c; i++ {
        m[i] = make([]float64, c)
    }
    for k := 0; k < 4; k++ {
        for i := p[k][0]; i < p[k][1]; i++ {
            for j := p[k][2]; j < p[k][3]; j++ {
                m[i][j] = q[k][i-p[k][4]][j-p[k][5]]
            }
        }
    }
    return m
}

func strassen(a, b Matrix) Matrix {
    if a.rows() != a.cols() || b.rows() != b.cols() || a.rows() != b.rows() {
        log.Fatal("Matrices must be square and of equal size.")
    }
    if a.rows() == 0 || (a.rows()&(a.rows()-1)) != 0 {
        log.Fatal("Size of matrices must be a power of two.")
    }
    if a.rows() == 1 {
        return a.mul(b)
    }
    qa := toQuarters(a)
    qb := toQuarters(b)
    p1 := strassen(qa[1].sub(qa[3]), qb[2].add(qb[3]))
    p2 := strassen(qa[0].add(qa[3]), qb[0].add(qb[3]))
    p3 := strassen(qa[0].sub(qa[2]), qb[0].add(qb[1]))
    p4 := strassen(qa[0].add(qa[1]), qb[3])
    p5 := strassen(qa[0], qb[1].sub(qb[3]))
    p6 := strassen(qa[3], qb[2].sub(qb[0]))
    p7 := strassen(qa[2].add(qa[3]), qb[0])
    var q [4]Matrix
    q[0] = p1.add(p2).sub(p4).add(p6)
    q[1] = p4.add(p5)
    q[2] = p6.add(p7)
    q[3] = p2.sub(p3).add(p5).sub(p7)
    return fromQuarters(q)
}

func main() {
    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.0 / 3, -1.0 / 4}, {-13.0 / 3, 19.0 / 4, -7.0 / 3, 11.0 / 24},
        {3.0 / 2, -2, 7.0 / 6, -1.0 / 4}, {-1.0 / 6, 1.0 / 4, -1.0 / 6, 1.0 / 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}}
    fmt.Println("Using 'normal' matrix multiplication:")
    fmt.Printf("  a * b = %v\n", a.mul(b))
    fmt.Printf("  c * d = %v\n", c.mul(d).toString(6))
    fmt.Printf("  e * f = %v\n", e.mul(f))
    fmt.Println("\nUsing 'Strassen' matrix multiplication:")
    fmt.Printf("  a * b = %v\n", strassen(a, b))
    fmt.Printf("  c * d = %v\n", strassen(c, d).toString(6))
    fmt.Printf("  e * f = %v\n", strassen(e, f))
}




Python: 

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

    @classmethod
    def block(cls, blocks) -> Matrix:
        
        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:
        
        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:
        
        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:
        
        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:
        
        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  

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