

In [ ]:

    
import os
import rasterio
from math import floor, ceil
import numpy as np

class VirtualRaster():
    def __init__(self, shape, transformation = None, proj4_crs = None):
        self.height = shape[0]
        self.width = shape[1]
        self.transform = transformation
        self.crs = proj4_crs
        self.bands = []
    
    def read_band(self, bidx, out=None, window=None, masked=None):
        """Read the `bidx` band into an `out` array if provided, 
        otherwise return a new array.

        Band indexes begin with 1: read_band(1) returns the first band.

        The optional `window` argument is a 2 item tuple. The first item
        is a tuple containing the indexes of the rows at which the
        window starts and stops and the second is a tuple containing the
        indexes of the columns at which the window starts and stops. For
        example, ((0, 2), (0, 2)) defines a 2x2 window at the upper left
        of the raster dataset.
        """
        band = self.bands[ bidx - 1]
        if not window:
            window = ((0,self.height),(0,self.width))
        if not out:
            window_shape = rasterio._base.window_shape(window, self.height, self.width)
            if masked:
                out = np.ma.empty(window_shape, band.dtype)
            else:
                out = np.empty(window_shape, band.dtype)
        return band.read(out, window, masked)
        
    def open(mode = 'r', base_path = None):
        #map( lambda b: map( lambda s: s.open, b.sources ),self.bands)
        for b in bands:
            for s in b.sources:
                s.open()
                
    def close():
        #map( lambda b: map( lambda s: s.open, b.sources ),self.bands)
        for b in bands:
            for s in b.sources:
                s.close()

class Band():
    def __init__(self, band_number, dtype, nodata = None):
        self.band_number = band_number
        self.dtype = dtype
        self.nodata = nodata
        self.sources = []
    
    def read(self, out, req_window, masked=None):
        # Consider using indexed dest_windows instead of brute force
        map(lambda src: src.read(out, req_window, masked), self.sources)
        return out
        
def crop_window(window, cropper_window):
    """Returns a version of window cropped against cropper_window. 
    Also returns a tuple containing two bools: (cropped_rows, cropped_cols)"""
    (changed_rows, changed_cols) = (False, False)
    ((row_start,row_end),(col_start, col_end)) = window
    if row_start < cropper_window[0][0]:
        row_start = cropper_window[0][0]
        changed_rows = True
    if col_start < cropper_window[1][0]:
        col_start = cropper_window[1][0]
        changed_cols = True
    if row_end > cropper_window[0][1]:
        row_end = cropper_window[0][1]
        changed_rows = True
    if col_end > cropper_window[1][1]:
        col_end = cropper_window[1][1]
        changed_cols = True
    return ( (row_start,row_end),(col_start,col_end) ), (changed_rows, changed_cols)

def windows_overlap(win1, win2):
        (ymin1, ymax1), (xmin1, xmax1) = win1
        (ymin2, ymax2), (xmin2, xmax2) = win2
        if ymin1 > ymax2 or ymax1 < ymin2 or xmin1 > xmax2 or xmax1 < xmin2:
            return False
        return True

class Source():
    def __init__(self, path, source_band, source_window, destination_window, source_nodata = None):
        self.path = path
        self.source_band = source_band
        self.source_window = source_window
        self.source_shape = rasterio._base.window_shape(source_window)
        self.destination_window = destination_window
        self.destination_shape = rasterio._base.window_shape(destination_window)
        self.source_nodata = source_nodata
        self.dataset = None
        self._scale = tuple(float(src)/float(dest) for src,dest in zip(self.source_shape,self.destination_shape))

    def open(self, mode = 'r', base_path = None):
        if self.dataset is None:
            absolute_path = self.path if not base_path else os.path.join(base_path, self.path)
            self.dataset = rasterio.open(absolute_path)
    
    def close(self):
        if self.dataset:
            self.dataset.close()
    
    def _source_to_destination(self, source):
        """Transforms source pixel coordinates into destination pixel coordinates.
        Accepts either a coordinate tuple or a window"""
        if isinstance(source[0], (tuple, list)) :
            # This is a window, not a coord pair
            zipped = zip( *source )
            start = tuple( int(floor(c)) for c in self._source_to_destination(zipped[0]) )
            # vrtsources.cpp does not ceil() the end coord. Rather it floors it
            end =  tuple( int(floor(c)) for c in self._source_to_destination(zipped[1]) )
            return tuple(zip(start, end))

        dest_col = (source[1] - self.source_window[1][0]) / self._scale[1] + self.destination_window[1][0]
        dest_row = (source[0] - self.source_window[0][0]) / self._scale[0] + self.destination_window[0][0]
        return (dest_row, dest_col)
    
    def _destination_to_source(self, destination ):
        """Transforms destination pixel coordinates into source pixel coordinates.
        Accepts either a (row,col) tuple or a window like ((row_start,row_end),(col_start,col_end))"""
        if isinstance(destination[0], (tuple, list)) :
            # This is a window, not a coord pair
            zipped = zip( *destination )
            source_start = tuple( int(floor(c)) for c in self._destination_to_source(zipped[0]) )
            source_end =  tuple( int(ceil(c)) for c in self._destination_to_source(zipped[1]) )
            return tuple(zip(source_start, source_end))
                           
        source_col = (destination[1] - self.destination_window[1][0]) * self._scale[1] + self.source_window[1][0]
        source_row = (destination[0] - self.destination_window[0][0]) * self._scale[0] + self.source_window[0][0]
        return (source_row, source_col)
    
    def read(self, out, req_window, masked=None):
        """ req_window is the total requested window in destination coordinates. 
        Out is a numpy array."""
        
        # Logic is roughly copied from GDAL's vrtsources.cpp
        
        req_window_shape = rasterio._base.window_shape(req_window)
        print 'req_window_shape', req_window_shape
        
        # Does req_window overlap destination_window
        if not windows_overlap(self.destination_window, req_window):
            return
        
        # Crop req_window to not extent outside dest_window
        dest_req_window, req_window_changed = crop_window(req_window, self.destination_window)
        print 'dest_req_window, req_window_changed', dest_req_window, req_window_changed
        
        # Translate req_window into source pix coords
        src_req_window = self._destination_to_source( dest_req_window )
        print 'src_req_window', src_req_window
        
        # If the requested area does not overlap the source window
        if not windows_overlap(self.source_window, dest_req_window):
            return
        
        # Crop source req window to be within source windowed bounds
        src_req_window, src_req_window_changed = crop_window(src_req_window, self.source_window)
        print 'src_req_window, src_req_window_changed', src_req_window, src_req_window_changed
        
        # Transform the source req window back into destination pixel coordinates
        dest_req_window = self._source_to_destination(src_req_window)
        print 'dest_req_window', dest_req_window
        
        # Where to put the data in the outarray        
        # Scale between original requested window and output buffer size
        scale_req_win_to_outarray = tuple( float(a)/b for a,b in zip(out.shape, req_window_shape) )
        print 'scale_req_win_to_outarray', scale_req_win_to_outarray
        
        # Calculate resulting window into outarray
        out_start_row = int((dest_req_window[1][0]-req_window[1][0])*scale_req_win_to_outarray[1]+0.001)
        out_end_row   = int((dest_req_window[1][1]-req_window[1][0])*scale_req_win_to_outarray[1]+0.001)
        out_start_col = int((dest_req_window[0][0]-req_window[0][0])*scale_req_win_to_outarray[0]+0.001)
        out_end_col   = int((dest_req_window[0][1]-req_window[0][0])*scale_req_win_to_outarray[0]+0.001)
        
        out_window = ((out_start_row, out_end_row),(out_start_col, out_end_col))
        print 'out_window', out_window
        out_window_shape = rasterio._base.window_shape(out_window)
        
        if out_window_shape[0] < 1 or out_window_shape[1] < 1:
            return
        
        # Create tmp array with source dtype and possibly masked
        if masked:
            tmp_out = np.ma.empty(out_window_shape, src.dataset.dtypes[0]) 
        else:
            tmp_out = np.empty(out_window_shape, src.dataset.dtypes[0])
            
        # Ok. Phew. Read
        self.dataset.read_band(self.source_band, out=tmp_out, window=src_req_window, masked=masked)        
        
        # Put the data in out
        out[ [slice(*dim) for dim in out_window] ] = tmp_out
        
        return out

        
        


# In[167]:

# Create a dummy source
# dummy.asc is 4 cols x 5 rows. We take 3x3 pixels and put them at the middle of a 9x9 destination raster
path = "dummy.asc"
source_band = 1
source_window = ( (1,4),(1,4) )
destination_window = ( (3,6),(3,6) )
src = Source(path, source_band, source_window, destination_window, source_nodata = None)
src.open()
out = np.ma.zeros((9,9))
req_window = ( (0,9),(0,9) )
src.read(out, req_window, masked=True)


# In[168]:

band = Band(1,'int64')
band.sources.append(src)

dataset = VirtualRaster((9,9))
dataset.bands.append(band)


# In[170]:

dataset.read_band(1, masked=True)


# In[46]:

zipped = zip( *((2,3) , (1,2)) )
print zipped
unzipped = zip(*zipped)
print unzipped


# In[48]:

def test1(r, c):
    return (r*5,c*5)
win = ( (1,2), (2,3) )
zip( *[ test1(*coord) for coord in zip( *win ) ])


# In[23]:

(ymin, ymax), (xmin, xmax) = ( (0,5),(10,20) ) 


# In[55]:

print int(floor(8.2))
print int(ceil(8.2))


# In[58]:

tuple( float(a)/b for a,b in zip((10,20), (30,40)) )


# In[61]:

zip((0.3333333333333333, 0.5), ( (0,5),(10,20) ), ( (10,15),(110,120) ) )


# In[62]:

any([0,1])


# In[72]:

a = np.array([1,2,3,4,5,6,6])


# In[76]:

a[slice(1,2)]


# In[80]:

slice(*(1,2))


# In[124]:

int(8.0001)


# In[96]:

b = np.arange(15).reshape(3,5)
print b
window = ( (0,3),(1,4) )
print [slice(*dim) for dim in window]
print b[[slice(*dim) for dim in window]]

b[[slice(*dim) for dim in window]] = np.ones((3,3))
print b


# In[ ]: