'''
read a raster file, calculate the topographic gradient and save this
as a new raster file

april 2013, Elco Luijendijk <elco.luijendijk at gmail.com>
'''

import numpy as np
import matplotlib.pyplot as pl

import pyGISlib

# load digital elevation file
rasterfile = 'dem_coast_netherlands.tif'
raster = pyGISlib.open_raster_file( rasterfile )

# get raster info
dimensions, origin, pixelsize, nodata, projection =\
    pyGISlib.get_raster_info( raster )

# read part of the raster
new_dimensions = (dimensions[0],dimensions[1]/2)
offset = (0,dimensions[1]/4)
raster_array = pyGISlib.read_part_of_raster(raster, offset, new_dimensions)
# calculate new origin
new_origin = [origin[0] + offset[0]*pixelsize[0], origin[1] + offset[1]*pixelsize[1]]

# uncomment the next two lines to read the entire raster file instead:
#raster_array, new_dimensions, new_origin, pixelsize, nodata, projection =\
# pyGISlib.read_raster_file(rasterfile)


# mask nodata values
raster_array = pyGISlib.mask_raster(raster_array, nodata)
print 'mean elevation = ', raster_array.mean()

# calculate the topographic gradient
grad_x, grad_y = np.gradient(raster_array)
grad_x = grad_x/pixelsize[0]
grad_y = grad_y/pixelsize[1]
grad = np.sqrt( (grad_x)**2 + (grad_y)**2 ) 
print 'mean topographic gradient ', grad.mean()

#
print 'saving raster file'
pyGISlib.write_raster_file( 'topo_gradient.tif', grad.data,
                        new_origin, pixelsize, nodata,
                        crs=projection )

# get coordinates of new raster
raster_x, raster_y = pyGISlib.get_raster_coordinates(new_dimensions, 
                                                   pixelsize, 
                                                   new_origin)

# make figures of elevation & topographic gradient
pl.figure(figsize=(8, 4))
pl.subplots_adjust(wspace=0.35, left=0.1, right=0.92)

ax1 = pl.subplot(1, 2, 1)
pl.contourf(raster_x, raster_y, raster_array)
cbar = pl.colorbar(shrink=0.5)
cbar.set_label('Elevation (m)')
pl.setp(ax1.get_xticklabels(), fontsize='xx-small')
pl.setp(ax1.get_yticklabels(), fontsize='xx-small')

ax2 = pl.subplot(1, 2, 2)
pl.contourf(raster_x, raster_y, grad)
cbar = pl.colorbar(shrink=0.5)
cbar.set_label('Topographic gradient')
pl.setp(ax2.get_xticklabels(), fontsize='xx-small')
pl.setp(ax2.get_yticklabels(), fontsize='xx-small')

print 'saving figure'
pl.savefig('pyGISlib_example_01.png')

pl.clf()
