import sys
import numpy as np
from matplotlib import pyplot as pl
from rw import WriteGTiff

fn = '../pozo-steep-vegetated-pcl.npy'
pts = np.load(fn)
x, y, z, c = pts[:, 0], pts[:, 1], pts[:, 2], pts[:, 5]
ix = (0.2 * (x - x.min())).astype('int')
iy = (0.2 * (y - y.min())).astype('int')

shape = (100, 100)
xb = np.arange(shape[1]+1)
yb = np.arange(shape[0]+1)
fg, ax = pl.subplots(ncols = 2, nrows = 2,
            figsize = (10.24, 10.24),
            sharex = True, sharey = True)

uc = (2, 5)
for j in range(len(uc)):
    print('Class %i' % uc[j])
    b = c == uc[j]
    cx, cy, cz = ix[b], iy[b], z[b]

    mean = np.zeros(shape)
    stdr = np.zeros(shape)
    for i in range(shape[0]):
        print('% 3d%%' % i)
        for k in range(shape[1]):
            b = (cy == i) * (cx == k)
            mean[i, k] = cz[b].mean()
            stdr[i, k] = cz[b].std()

    fname = 'pozo_5m_dem_mean_cl%i.tif' % uc[j]
    WriteGTiff(fname, mean, x.min(), y.min()+500, step = 5)
    np.save('pozo_5m_dem_mean_cl%i.npy' % uc[j], mean)
    np.save('pozo_5m_dem_stdr_cl%i.npy' % uc[j], stdr)
    ax[0, j].set_title('Class %i' % uc[j])
    im = ax[0, j].pcolormesh(xb, yb,
                  np.ma.masked_invalid(mean),
                  cmap = pl.cm.viridis_r)
    cb = fg.colorbar(im, ax = ax[0, j])
    cb.set_label('Mean elevation [m]')

    im = ax[1, j].pcolormesh(xb, yb,
                  np.ma.masked_invalid(stdr),
                  cmap = pl.cm.magma_r)
    cb = fg.colorbar(im, ax = ax[1, j])
    cb.set_label('Elevation STD')
    ax[0, j].set_aspect('equal')
    ax[1, j].set_aspect('equal')

pl.savefig('%s.png' % sys.argv[0][:-3])

[((7, 6, 7, 17), 'numpy.load', 'np.load', ({(7, 14, 7, 16): 'fn'}, {}), '(fn)', True, 'import numpy as np\n'), ((13, 5, 13, 26), 'numpy.arange', 'np.arange', ({(13, 15, 13, 25): 'shape[1] + 1'}, {}), '(shape[1] + 1)', True, 'import numpy as np\n'), ((14, 5, 14, 26), 'numpy.arange', 'np.arange', ({(14, 15, 14, 25): 'shape[0] + 1'}, {}), '(shape[0] + 1)', True, 'import numpy as np\n'), ((15, 9, 17, 41), 'matplotlib.pyplot.subplots', 'pl.subplots', (), '', True, 'from matplotlib import pyplot as pl\n'), ((53, 0, 53, 39), 'matplotlib.pyplot.savefig', 'pl.savefig', ({(53, 11, 53, 38): "('%s.png' % sys.argv[0][:-3])"}, {}), "('%s.png' % sys.argv[0][:-3])", True, 'from matplotlib import pyplot as pl\n'), ((25, 11, 25, 26), 'numpy.zeros', 'np.zeros', ({(25, 20, 25, 25): 'shape'}, {}), '(shape)', True, 'import numpy as np\n'), ((26, 11, 26, 26), 'numpy.zeros', 'np.zeros', ({(26, 20, 26, 25): 'shape'}, {}), '(shape)', True, 'import numpy as np\n'), ((36, 4, 36, 54), 'numpy.save', 'np.save', ({(36, 12, 36, 47): "('pozo_5m_dem_mean_cl%i.npy' % uc[j])", (36, 49, 36, 53): 'mean'}, {}), "('pozo_5m_dem_mean_cl%i.npy' % uc[j], mean)", True, 'import numpy as np\n'), ((37, 4, 37, 54), 'numpy.save', 'np.save', ({(37, 12, 37, 47): "('pozo_5m_dem_stdr_cl%i.npy' % uc[j])", (37, 49, 37, 53): 'stdr'}, {}), "('pozo_5m_dem_stdr_cl%i.npy' % uc[j], stdr)", True, 'import numpy as np\n'), ((40, 18, 40, 44), 'numpy.ma.masked_invalid', 'np.ma.masked_invalid', ({(40, 39, 40, 43): 'mean'}, {}), '(mean)', True, 'import numpy as np\n'), ((46, 18, 46, 44), 'numpy.ma.masked_invalid', 'np.ma.masked_invalid', ({(46, 39, 46, 43): 'stdr'}, {}), '(stdr)', True, 'import numpy as np\n')]