hw1cx.py

import math


# Get the value of n and initialize the matrices A and B.
with open("input.txt", "r") as matrix_info:
    n = int(matrix_info.readline())
    assert n == 2 ** math.log(n, 2), "The first line must be a power of 2."

    mat_a = []
    mat_b = []
    for line_number in xrange(2 * n):
        row_info = matrix_info.readline().split()
        row_frmt = map(float, row_info)
        if line_number < n:
            mat_a.append(row_frmt)
        else:
            mat_b.append(row_frmt)


def col_select(mat, j):
    # Select column j from matrix mat.
    return [row[j] for row in mat]


def dot_product(vec_a, vec_b):
    # Return the dot product of two vectors vec_a and vec_b
    return sum(vec_a[i] * vec_b[i] for i in xrange(len(vec_a)))


def slicer(mat, (i_i, i_f), (j_i, j_f)):
    # Return a section of a matrix mat defined by two pairs of i
    # and j indices.
    return [[mat[i][j] for j in xrange(j_i, j_f)] for i in xrange(i_i, i_f)]


def mat_add(a, b):
    # Add matrices a and b.
    c = [[[] for _ in xrange(len(a))] for _ in xrange(len(b[0]))]
    for i in xrange(len(a)):
        for j in xrange(len(b[0])):
            c[i][j] = a[i][j] + b[i][j]
    return c


def mat_sub(a, b):
    # Subtract matrix b from matrix a.
    c = [[[] for _ in xrange(len(a))] for _ in xrange(len(b[0]))]
    for i in xrange(len(a)):
        for j in xrange(len(b[0])):
            c[i][j] = a[i][j] - b[i][j]
    return c


def multiply_strassen(a, b):
    # Implements Strassen's algorithm for matrices a and b.
    c = [[[] for _ in xrange(len(a))] for _ in xrange(len(b[0]))]
    q = len(a)
    if q == 2:
        m_1 = (a[0][0] + a[1][1]) * (b[0][0] + b[1][1])
        m_2 = (a[1][0] + a[1][1]) * b[0][0]
        m_3 = a[0][0] * (b[0][1] - b[1][1])
        m_4 = a[1][1] * (b[1][0] - b[0][0])
        m_5 = (a[0][0] + a[0][1]) * b[1][1]
        m_6 = (a[1][0] - a[0][0]) * (b[0][0] + b[0][1])
        m_7 = (a[0][1] - a[1][1]) * (b[1][0] + b[1][1])
        c[0][0] = m_1 + m_4 - m_5 + m_7
        c[0][1] = m_3 + m_5
        c[1][0] = m_2 + m_4
        c[1][1] = m_1 - m_2 + m_3 + m_6
        return c
    else:
        p = q / 2
        zero_slice = (0, p)
        one_slice = (p, q)
        m_1 = multiply_strassen(mat_add(slicer(a, zero_slice, zero_slice),
                                        slicer(a, one_slice, one_slice)),
                                mat_add(slicer(b, zero_slice, zero_slice),
                                        slicer(b, one_slice, one_slice)))
        m_2 = multiply_strassen(mat_add(slicer(a, one_slice, zero_slice),
                                        slicer(a, one_slice, one_slice)),
                                slicer(b, zero_slice, zero_slice))
        m_3 = multiply_strassen(slicer(a, zero_slice, zero_slice),
                                mat_sub(slicer(b, zero_slice, one_slice),
                                        slicer(b, one_slice, one_slice)))
        m_4 = multiply_strassen(slicer(a, one_slice, one_slice),
                                mat_sub(slicer(b, one_slice, zero_slice),
                                        slicer(b, zero_slice, zero_slice)))
        m_5 = multiply_strassen(mat_add(slicer(a, zero_slice, zero_slice),
                                        slicer(a, zero_slice, one_slice)),
                                slicer(b, one_slice, one_slice))
        m_6 = multiply_strassen(mat_sub(slicer(a, one_slice, zero_slice),
                                        slicer(a, zero_slice, zero_slice)),
                                mat_add(slicer(b, zero_slice, zero_slice),
                                        slicer(b, zero_slice, one_slice)))
        m_7 = multiply_strassen(mat_sub(slicer(a, zero_slice, one_slice),
                                        slicer(a, one_slice, one_slice)),
                                mat_add(slicer(b, one_slice, zero_slice),
                                        slicer(b, one_slice, one_slice)))
        # Collect the four quadrants of C
        c_up_lf = mat_add(mat_sub(mat_add(m_1, m_4), m_5), m_7)
        c_up_rt = mat_add(m_3, m_5)
        c_dn_lf = mat_add(m_2, m_4)
        c_dn_rt = mat_add(mat_add(mat_sub(m_1, m_2), m_3), m_6)
        # Populate the C matrix
        for i in xrange(p):
            for j in xrange(p):
                c[i][j] = c_up_lf[i][j]
                c[i][p + j] = c_up_rt[i][j]
                c[p + i][j] = c_dn_lf[i][j]
                c[p + i][p + j] = c_dn_rt[i][j]
        return c


def multiply_classical(a, b):
    # Implements the classical algorithm for multiplying matrices and b.
    c = [[[] for _ in xrange(len(a))] for _ in xrange(len(b[0]))]
    for i in xrange(len(c)):
        for j in xrange(len(c[0])):
            c[i][j] = dot_product(a[i], col_select(b, j))
    return c


print "Strassen result:"
c_s = multiply_strassen(mat_a, mat_b)
for row in c_s:
    print " ".join(map(str, row))
print "Classical result:"
c_c = multiply_classical(mat_a, mat_b)
for row in c_c:
    print " ".join(map(str, row))
print "Difference matrix:"
e = mat_sub(c_s, c_c)
for row in e:
    print " ".join(map(str, row))

input_generator.py

import sys
import math
import random


"Pass the matrix size n to input_generator.py as a command line argument"
n = int(sys.argv[1])
assert n == 2 ** math.log(n, 2), "Argument passed must be a power of 2"
mat_a = [[str(random.uniform(-10, 10)) for _ in xrange(n)] for _ in xrange(n)]
mat_b = [[str(random.uniform(-10, 10)) for _ in xrange(n)] for _ in xrange(n)]

with open("input.txt", "w") as matrix_info:
    matrix_info.write(str(n) + "\n")
    for row in mat_a:
        matrix_info.write(" ".join(row) + "\n")
    for row in mat_b:
        matrix_info.write(" ".join(row) + "\n")