I have created a python script that takes a 2D NumPy array of elevation values and given a seed value and an elevation it will grow a flood from that seed point. It does this by checking all adjacent cells to see if they are above or below the flood level. From there it checks all of those locations until all cells that touch a flooded cell have been checked. If a cell is below the flood level, but surrounded by higher values it doesn't get flooded.
This is the test code. I then implement it in GIS where I convert a raster to an array and process that. This works, but doesn't scale well at all if used with a larger raster, i.e. takes ages.
How can I improve this code, is this the right approach in python? I know there are GIS tools to do this out there, this is a learning task for myself.
'''
This program is designed to generate in_array flood array from in_array seed point
and an elevation to flood to. The output array is set to 0 where
the cell is calculated to be flooded and retains its elevation
everywhere else.
'''
#Support Python 3 print functionality in Python 2
from __future__ import print_function
#import generic modules
import sys
import numpy as np
import matplotlib.pyplot as plt
#global variables
in_array = np.random.rand(0)
proc_array = np.random.rand(0)
spent = []
to_process = []
'''
Main calling block
'''
def proc_main():
print('Starting...')
try:
#Run actual functionality
proc_run()
except Exception as e:
print ('Error:' + str(e))
finally:
pass
#optional wait for keypress
#input('Press Enter...')
'''
Main program functionality
'''
def proc_run():
global in_array,proc_array, spent, to_process
#in_array = np.random.rand(11,11)
in_array = np.array([[5,9,5,5,5,5,9,5,5,5,5],
[5,9,5,5,9,9,9,5,5,5,5],
[9,9,5,5,9,5,9,5,5,5,5],
[9,9,5,5,9,9,9,5,5,5,5],
[5,5,5,5,5,5,9,5,5,5,9],
[5,5,5,5,5,8,9,5,5,5,9],
[5,5,5,5,5,5,9,5,5,5,9],
[5,5,5,5,5,5,9,5,5,5,5],
[5,5,5,5,5,5,9,5,5,5,5],
[5,5,5,5,5,5,9,5,5,5,5],
[5,5,5,5,5,5,5,5,5,5,5]])
#plot array
plt.imshow(in_array, cmap=plt.cm.gray)
plt.colorbar()
plt.show()
#set sed location
loc = [5,5]
to_process.append(loc)
#water
flood_z = 8.5
#make in_array copy of array for processing
proc_array = np.copy(in_array)
#check if origin can flood
if in_array[loc[0], loc[1]] <= flood_z:
proc_array[loc[0], loc[1]] = -9999
#process seed
proc_loc(loc, flood_z)
#process whole array
for i in to_process:
if proc_array[i[0],i[1]] == -9999:
proc_loc(i, flood_z)
to_process.remove(i)
spent.append(i)
else:
print('Seed location is above flood value given')
raise
#classify array and mask input array
proc_array[proc_array > -9999] = 1
proc_array[proc_array <= -9999] = 0
out_array = in_array * proc_array
#plot mask and new array
plt.imshow(proc_array, cmap=plt.cm.gray)
plt.show()
plt.imshow(out_array, cmap=plt.cm.gray)
plt.colorbar()
plt.show()
'''
Takes an array location and an elevation and looks at the 8 surrounding
cells to see if they are below the elevation.
'''
def proc_loc(loc, flood_z):
global proc_array, spent
if loc not in spent:
rows, cols = in_array.shape
rownum = loc[0]
colnum = loc[1]
if rownum - 1 != -1:
# -1 -1 top left
if colnum - 1 != -1:
proc_array[rownum - 1, colnum - 1] = \
cal_z([rownum - 1, colnum - 1],
proc_array[rownum - 1, colnum - 1], flood_z)
#-1 0 top
proc_array[rownum - 1, colnum] = \
cal_z([rownum - 1, colnum],
proc_array[rownum - 1, colnum], flood_z)
#-1 +1 top right
if colnum + 1 < int(cols):
proc_array[rownum - 1, colnum + 1] = \
cal_z([rownum - 1, colnum + 1],
proc_array[rownum - 1, colnum + 1], flood_z)
if colnum - 1 != -1:
#0 -1 left
proc_array[rownum, colnum - 1] = \
cal_z([rownum, colnum -1],
proc_array[rownum, colnum - 1], flood_z)
if colnum + 1 < int(cols):
#right
proc_array[rownum, colnum + 1] = \
cal_z([rownum, colnum +1],
proc_array[rownum, colnum + 1], flood_z)
if rownum + 1 < int(rows):
#+1 -1 bottom left
if colnum - 1 != -1:
proc_array[rownum + 1, colnum - 1] = \
cal_z([rownum + 1, colnum - 1],
proc_array[rownum + 1, colnum - 1], flood_z)
#+1 0 bottom
proc_array[rownum + 1, colnum] = \
cal_z([rownum + 1, colnum],
proc_array[rownum + 1, colnum], flood_z)
#+1 +1 bottom right
if colnum + 1 < int(cols):
proc_array[rownum + 1, colnum + 1] = \
cal_z([rownum + 1, colnum + 1],
proc_array[rownum + 1, colnum + 1], flood_z)
'''
Takes an array location in_array flood_z value and the target elevation and checks if
the flood_z value s below the target. If it is the array location is
added to the list to be processed.
'''
def cal_z(loc, cell_z, flood_z):
global to_process
to_process.append(loc)
if cell_z <= flood_z:
return -9999
else:
return 9999
if __name__ == '__main__':
# Support Python 2 and 3 input
# If this is Python 2, use raw_input()
if sys.version_info[0] >= 3:
input = input
range = range
else:
input = raw_input
range = xrange
#run main program
proc_main()
