Source

dem_waters_extractor / demfunctions.py

#!/usr/bin/env python
#-*- encoding:utf-8 -*-

import math
import copy

def fill_sinks(data):
    r_max = len(data)
    c_max = len(data[0])
    for r in xrange(len(data)):
        for c in xrange(len(data[r])):
            if r == 0 or r == r_max-1 or c == 0 or c == c_max-1:
                continue
            s0,s1,s2,s3,s4,s5,s6,s7 = [True for i in range(8)] # all the border is highter than this node
            if r-1>=0 and c+1<=c_max-1:
                s0 = data[r][c] <= data[r-1][c+1]
                print 'x0'
                print data[r-1][c+1]
            if c+1<=c_max-1:
                s1 = data[r][c] <= data[r][c+1]
                print 'x1'
                print data[r][c+1]
            if r+1<=r_max-1 and c+1<=c_max-1:
                s2 = data[r][c] <= data[r+1][c+1]
                print 'x2'
                print data[r+1][c+1]
            if r+1<=r_max-1:
                s3 = data[r][c] <= data[r+1][c]
                print 'x3'
                print data[r+1][c]
            if r+1<=r_max-1 and c-1>=0:
                s4 = data[r][c] <= data[r+1][c-1]
                print 'x4'
                print data[r+1][c-1]
            if c-1>=0:
                s5 = data[r][c] <= data[r][c-1]
                print 'x5'
                print data[r][c-1]
            if r-1>=0 and c-1>=0:
                s6 = data[r][c] <= data[r-1][c-1]
                print 'x6'
                print data[r-1][c-1]
            if r-1>=0:
                s7 = data[r][c] <= data[r-1][c]
                print 'x7'
                print data[r-1][c]
            if False not in [s0,s1,s2,s3,s4,s5,s6,s7]: # sink?
                # window size default to 2*2
                print '==================================='
                print 'center point:r:%s,c:%s'%(r,c)
                print '==================================='
                len_ = 0
                marks = [[r,c]]
                hidden_outset_min = None
                border = []
                while True:
                    border = []
                    print '___'
                    print 'center point:r:%s,c:%s'%(r,c)
                    print 'len_:%s'%len_
                    print '___'
                    len_ += 1
                    for i in range(len_+1):
                        if r+len_ <= r_max-1 and c+i <=c_max-1:
                            tem = [r+len_,c+i]
                            if tem not in border:
                                border.append(tem)
                        if r+len_ <= r_max-1 and c-i >=0:
                            tem= [r+len_,c-i]
                            if tem not in border:
                                border.append(tem)
                        if r-len_ >= 0 and c+i <=c_max-1:
                            tem = [r-len_,c+i]
                            if tem not in border:
                                border.append(tem)
                        if r-len_ >= 0 and c-i >=0:
                            tem = [r-len_,c-i]
                            if tem not in border:
                                border.append(tem)
                        if c+len_ <= c_max-1 and r+i <=r_max-1:
                            tem = [r+i,c+len_]
                            if tem not in border:
                                border.append(tem)
                        if c+len_ <= c_max-1 and r-i >=0:
                            tem = [r-i,c+len_]
                            if tem not in border:
                                border.append(tem)
                        if c-len_ >= 0 and r+i <=r_max-1:
                            tem = [r+i,c-len_]
                            if tem not in border:
                                border.append(tem)
                        if c-len_ >=0 and r-i >=0:
                            tem = [r-i,c-len_]
                            if tem not in border:
                                border.append(tem)
                    print 'border:%s'%border
                    if len(border) < 8*len_:
                        print 'center node: %s.\n the end of dem for this center node.\n this node is assumed to be a biggggggg sink'%[r,c]
                        break
                    hidden_outset = []
                    for node in border:
                        print 'which node:%s'%node
                        loc0 = [node[0]+1,node[1]+1]
                        loc1 = [node[0],node[1]+1]
                        loc2 = [node[0]-1,node[1]+1]
                        loc3 = [node[0]-1,node[1]]
                        loc4 = [node[0]-1,node[1]-1]
                        loc5 = [node[0],node[1]-1]
                        loc6 = [node[0]+1,node[1]-1]
                        loc7 = [node[0]+1,node[1]]
                        t0,t1,t2,t3,t4,t5,t6,t7 = [None for i in range(8)]
                        print 'marks:%s'%marks
                        if loc0 not in marks:
                            if loc0[0]>=0 and loc0[1]>=0 and loc0[0]<=r_max-1 and loc0[1]<=c_max-1:
                                print 'loc0:%s'%loc0
                                t0 = data[loc0[0]][loc0[1]]
                        if loc1 not in marks:
                            if loc1[0]>=0 and loc1[1]>=0 and loc1[0]<=r_max-1 and loc1[1]<=c_max-1:
                                print 'loc1:%s'%loc1
                                t1 = data[loc1[0]][loc1[1]]
                        if loc2 not in marks:
                            if loc2[0]>=0 and loc2[1]>=0 and loc2[0]<=r_max-1 and loc2[1]<=c_max-1:
                                print 'loc2:%s'%loc2
                                t2 = data[loc2[0]][loc2[1]]
                        if loc3 not in marks:
                            if loc3[0]>=0 and loc3[1]>=0 and loc3[0]<=r_max-1 and loc3[1]<=c_max-1:
                                print 'loc3:%s'%loc3
                                t3 = data[loc3[0]][loc3[1]]
                        if loc4 not in marks:
                            if loc4[0]>=0 and loc4[1]>=0 and loc4[0]<=r_max-1 and loc4[1]<=c_max-1:
                                print 'loc4:%s'%loc4
                                t4 = data[loc4[0]][loc4[1]]
                        if loc5 not in marks:
                            if loc5[0]>=0 and loc5[1]>=0 and loc5[0]<=r_max-1 and loc5[1]<=c_max-1:
                                print 'loc5:%s'%loc5
                                t5 = data[loc5[0]][loc5[1]]
                        if loc6 not in marks:
                            if loc6[0]>=0 and loc6[1]>=0 and loc6[0]<=r_max-1 and loc6[1]<=c_max-1:
                                print 'loc6:%s'%loc6
                                t6 = data[loc6[0]][loc6[1]]
                        if loc7 not in marks:
                            if loc7[0]>=0 and loc7[1]>=0 and loc7[0]<=r_max-1 and loc7[1]<=c_max-1:
                                print 'loc7:%s'%loc7
                                t7 = data[loc7[0]][loc7[1]]
                        if True in [data[node[0]][node[1]]>i for i in [t0,t1,t2,t3,t4,t5,t6,t7] if i!=None]:
                            # is a hidden output node
                            hidden_outset.append(node)
                        # mark belonging to sink area
                        marks.append(node)
                    if len(hidden_outset) > 0:
                        print 'current hidden_outset number:%s'%len(hidden_outset)
                    else:
                        continue
                    hidden_outset_min = hidden_outset[0]
                    for node in (hidden_outset):
                        if hidden_outset_min > data[node[0]][node[1]]:
                            hidden_outset_min = data[node[0]][node[1]]
                    print 'hidden_outset:%s'%hidden_outset
                    print 'hidden_outset_min:%s'%hidden_outset_min
                    next_border_e = 0 # escape while loop
                    for node in border:
                        if data[node[0]][node[1]] < hidden_outset_min:
                            next_border_e = 1
                            break
                    if next_border_e == 0:
                        break
                if len(border) == 8*len_:
                    if hidden_outset_min > data[r][c]:
                        print "sink"
                        for node in marks:
                            if data[node[0]][node[1]] < hidden_outset_min:
                                data[node[0]][node[1]] = hidden_outset_min
                    else:
                        print "plane"
    return data

def lift(data,p):
    r_max = len(data)
    c_max = len(data[0])
    lifted = copy.deepcopy(data)
    modified = False
    for r in xrange(len(data)):
        for c in xrange(len(data[r])):
            if r == 0 or r == r_max-1 or c == 0 or c == c_max-1:
                continue
            s0,s1,s2,s3,s4,s5,s6,s7 = [True for i in range(8)] # all the border is highter than or equal to this node
            if r-1>=0 and c+1<=c_max-1:
                s0 = data[r][c] <= data[r-1][c+1]
            if c+1<=c_max-1:
                s1 = data[r][c] <= data[r][c+1]
            if r+1<=r_max-1 and c+1<=c_max-1:
                s2 = data[r][c] <= data[r+1][c+1]
            if r+1<=r_max-1:
                s3 = data[r][c] <= data[r+1][c]
            if r+1<=r_max-1 and c-1>=0:
                s4 = data[r][c] <= data[r+1][c-1]
            if c-1>=0:
                s5 = data[r][c] <= data[r][c-1]
            if r-1>=0 and c-1>=0:
                s6 = data[r][c] <= data[r-1][c-1]
            if r-1>=0:
                s7 = data[r][c] <= data[r-1][c]
            if False not in [s0,s1,s2,s3,s4,s5,s6,s7]: # sink or plane
                lifted[r][c] += p
                modified = True
    if modified == False:
        return lifted
    print 'life return'
    return lift(lifted,p)

def get_vect_martrix(data,x,y):
    '''
    x: x distance of each node on column
    y: y distance of each node on row
    '''
    r_max = len(data) # exactly len(data) - 1
    c_max = len(data[0]) # exactly len(data[0]) - 1
    vect_martix = [[0 for c in range(c_max)] for r in range(r_max)]
    for r in xrange(len(data)):
        for c in xrange(len(data[r])):
            if r == 0:
                vect_martix[r][c] = 7
                continue
            if r == r_max-1:
                vect_martix[r][c] = 3
                continue
            if c == 0:
                vect_martix[r][c] = 5
                continue
            if c == c_max-1:
                vect_martix[r][c] = 1
                continue
            v0,v1,v2,v3,v4,v5,v6,v7 = [0 for i in range(8)]
            if r-1>=0 and c+1<=c_max-1:
                v0 = (data[r][c]-data[r-1][c+1])/math.pow(x**2+y**2,0.5)
            if c+1<=c_max-1:
                v1 = (data[r][c]-data[r][c+1])/x
            if r+1<=r_max-1 and c+1<=c_max-1:
                v2 = (data[r][c]-data[r+1][c+1])/math.pow(x**2+y**2,0.5)
            if r+1<=r_max-1:
                v3 = (data[r][c]-data[r+1][c])/y
            if r+1<=r_max-1 and c-1>0:
                v4 = (data[r][c]-data[r+1][c-1])/math.pow(x**2+y**2,0.5)
            if c-1>=0:
                v5 = (data[r][c]-data[r][c-1])/x
            if r-1>=0 and c-1>=0:
                v6 = (data[r][c]-data[r-1][c-1])/math.pow(x**2+y**2,0.5)
            if r-1>=0:
                v7 = (data[r][c]-data[r-1][c])/y
            which = 0
            v_max = max([v0,v1,v2,v3,v4,v5,v6,v7])
            print 'v_max:%s'%v_max
            for i in [v0,v1,v2,v3,v4,v5,v6,v7]:
                if v_max == i:
                    break
                which += 1
            if which == None:
                raise 'error'
            vect_martix[r][c] = which
    return vect_martix

def count_water(data):
    r_max = len(data) # exactly len(data) - 1
    c_max = len(data[0]) # exactly len(data[0]) - 1
    water_matrix = [[0 for c in range(c_max)] for r in range(r_max)]
    def trackpath(r,c):
        if r > r_max-1 or r < 0 or c>c_max-1 or c < 0:
            return []
        #if r>5 or r<0 or c>5 or c<0:
        #    return[]
        node = data[r][c]
        path = [[r,c],]
        if node == 0:
            temp = trackpath(r-1,c+1)
            if len(temp) >0:
                path.extend(temp)
        if node == 1:
            temp = trackpath(r,c+1)
            if len(temp) >0:
                path.extend(temp)
        if node == 2:
            temp = trackpath(r+1,c+1)
            if len(temp) >0:
                path.extend(temp)
        if node == 3:
            temp = trackpath(r+1,c)
            if len(temp) >0:
                path.extend(temp)
        if node == 4:
            temp = trackpath(r+1,c-1)
            if len(temp) >0:
                path.extend(temp)
        if node == 5:
            temp = trackpath(r,c-1)
            if len(temp) >0:
                path.extend(temp)
        if node == 6:
            temp = trackpath(r-1,c-1)
            if len(temp) >0:
                path.extend(temp)
        if node == 7:
            temp = trackpath(r-1,c)
            if len(temp) >0:
                path.extend(temp)
        return path
    for r in xrange(len(data)):
        for c in xrange(len(data)):
            path = trackpath(r,c)
            for node in path:
                water_matrix[node[0]][node[1]] += 1
    return water_matrix

def river_map(data,threshold=5):
    r_max = len(data) # exactly len(data) - 1
    c_max = len(data[0]) # exactly len(data[0]) - 1
    rivermap_matrix = [[0 for c in range(c_max)] for r in range(r_max)]
    for r in xrange(len(data)):
        for c in xrange(len(data[r])):
            if data[r][c] >=threshold:
                rivermap_matrix[r][c] = 1
    return rivermap_matrix