# -*- coding: utf-8 -*-
"""
Created on Mon Jan 18 13:53:04 2016

@author: md316445
"""
import numpy as np
a = np.arange(9).reshape(3,3)

def gen_mask(k, w, s=5):
    diag = np.identity(s)
    adiag = np.roll(np.fliplr(np.identity(s)), 1, axis=1)
    roll = np.roll(diag+adiag, w, axis=0)
    return np.roll(roll, k, axis=1)

def diags(A):
    out = np.zeros(A.shape)
    for i in range(A.shape[0]):
        for j in range(A.shape[1]):
            out[i, j]=np.sum(A*gen_mask(j, i, A.shape[0]))
    return out
            
print(a)
print(diags(a))