cdl_diversity.py

from __future__ import division

from skbio.diversity.alpha import simpson
import rasterio as rio
from sys import argv
import numpy as np
import affine

scale = 20

with rio.open(argv[1]) as src:
  dstw = int(src.width/scale)
  dsth = int(src.height/scale)
  
  print dstw, src.width, scale
  print dsth, src.height, scale
  
  small = np.zeros((dsth, dstw), dtype=np.uint16)
  
  for x in range(dstw):
    for y in range(dsth):
      
      w = ((y*scale, (y+1)*scale), (x*scale, (x+1)*scale))
      chunk = src.read(1, window=w)

      if chunk.shape != (scale, scale):
        continue
      
      chunk = chunk.flatten()
      chunk = chunk[np.where(chunk != 0)]
      _, counts = np.unique(chunk, return_counts=True)

      small[y, x] = simpson(counts) * (2**16 - 1)
    print x


with rio.open(
  argv[2], 
  'w', 
  driver='GTiff', 
  width=dstw, 
  height=dsth, 
  count=1, 
  dtype=np.uint16) as dst:

  dst.write(small, 1)