flatten and unflatten matrix in Sage via numpy

flatten_unflatten_matrix .sage

import numpy as np
from sage.stats.distributions.discrete_gaussian_lattice import DiscreteGaussianDistributionLatticeSampler

n = 256
D = DiscreteGaussianDistributionLatticeSampler(ZZ^n, 3.0)

arr = [list(D()) for i in range(n)]

flat = flatten(list(arr))

print arr
print flat

def unflatten(flat):
    return [flat[i:i+sqrt(len(flat))] for i in np.arange(0, len(flat), sqrt(len(flat)))]

unflat = unflatten(flat)

print np.array_equal(arr, np.asarray(unflat))