def get_surface_route(surface_route_start):
global acc_ds, dir_ds, route_ds, slope_ds, water_value, river_th, surface_route_value
# 提取各水体的坡面流路
for route_start_point in surface_route_start:
# 初始化流路追踪数组
judge_route = [route_start_point]
while len(judge_route) > 0:
current_point = judge_route.pop()
xoff = current_point[0]
yoff = current_point[1]
# 获取此点其汇流累积量
acc_off = cu.off_transform(xoff, yoff, slope_ds, acc_ds)
acc_value = cu.get_raster_float_value(acc_ds, acc_off[0], acc_off[1])
# 获取此点在结果数据的值
ol_value = cu.get_raster_int_value(slope_ds, xoff, yoff)
# 若流向的点不为水体且不为河道则继续
if ol_value != water_value and acc_value < river_th:
# 记录此点为坡面流路
cu.set_raster_int_value(route_ds, xoff, yoff, surface_route_value)
# 获取此点流向
dir_value = cu.get_raster_int_value(dir_ds, xoff, yoff)
# 获取其流向的点
to_point = cu.get_to_point(xoff, yoff, dir_value)
if len(to_point) > 0:
to_p_a_off = cu.off_transform(to_point[0], to_point[1], slope_ds, acc_ds)
if cu.in_data(to_p_a_off[0], to_p_a_off[1], acc_ds.RasterXSize, acc_ds.RasterYSize):
# 加入判断数组
judge_route.append(to_point)
def spill_point_dir(inner_ras_indexes, dem_ds):
# 内边界的高程
inner_dem = []
for ras_index in inner_ras_indexes:
dem_value = cu.get_raster_float_value(dem_ds, ras_index[0], ras_index[1])
inner_dem.append(dem_value)
# 获得最低位置的索引
spill_pt_index = list(map(inner_dem.index, heapq.nsmallest(1, inner_dem)))[0]
spill_pt = inner_ras_indexes[spill_pt_index]
# 获得溢出方向
spill_dir = get_spill_dir(inner_ras_indexes, spill_pt_index)
# 得到接收点索引
reception_pt = cu.get_to_point(spill_pt[0], spill_pt[1], spill_dir)
# # 输出到tif
# temp_path = r'G:\Graduation\Program\Data\41\endorheic_area0\outlet_dir.tif'
# file_format = "GTiff"
# driver = gdal.GetDriverByName(file_format)
# temp_ds = driver.Create(temp_path, dem_ds.RasterXSize, dem_ds.RasterYSize, 1, gdal.GDT_Int16, options=['COMPRESS=DEFLATE'])
# temp_ds.SetGeoTransform(dem_ds.GetGeoTransform())
# temp_ds.SetProjection(dem_ds.GetProjection())
# temp_ds.GetRasterBand(1).SetNoDataValue(-1)
# cu.set_raster_int_value(temp_ds, spill_pt[0], spill_pt[1], spill_dir)
# temp_ds = None
return spill_pt, spill_dir, reception_pt
def outline_upstream_search(xoff, yoff, water_buf, upstream_inflow):
global dir_ds, slope_ds, water_value
# 创建搜索数组
search_inflows = [[xoff, yoff]]
# 开始搜索
while len(search_inflows) > 0:
# 取出初始像元
cell_off = search_inflows.pop()
xoff = cell_off[0]
yoff = cell_off[1]
# 搜索周边像元
ol_dir_array = cu.get_8_dir(xoff, yoff)
for ol_n_cell in ol_dir_array:
n_xoff = ol_n_cell[0]
n_yoff = ol_n_cell[1]
# 判断像元是否在数据集内且是否为有效流向
judge_in_data = cu.in_data(n_xoff, n_yoff, dir_ds.RasterXSize, dir_ds.RasterYSize)
if judge_in_data:
# 判断是否在水体内
ol_data = cu.get_raster_int_value(slope_ds, n_xoff, n_yoff)
# 若不在水体则继续
if ol_data != water_value:
dir_data_value = cu.get_raster_int_value(dir_ds, n_xoff, n_yoff)
to_point = cu.get_to_point(n_xoff, n_yoff, dir_data_value)
if len(to_point) != 0:
# 若为上游点
if to_point[0] == xoff and to_point[1] == yoff:
# 判断是否在外边界上，若是则记录
if ol_n_cell in water_buf:
# 则记录为外边界上的上游坡面入流点
upstream_inflow.append([n_xoff, n_yoff])
# 继续遍历相邻像元
search_inflows.append(ol_n_cell)
def get_to_point_ol_data(xoff, yoff):
global dir_ds, slope_ds, no_data_value
dir_data_value = cu.get_raster_int_value(dir_ds, xoff, yoff)
to_point = cu.get_to_point(xoff, yoff, dir_data_value)
if len(to_point) < 1:
return no_data_value
else:
return cu.get_raster_int_value(slope_ds, to_point[0], to_point[1])
def get_to_point_ol_data(xoff, yoff):
global dataset_dir, dataset_ol, no_data_value
dir_data_value = cu.get_raster_int_value(dataset_dir, xoff, yoff)
to_point = cu.get_to_point(xoff, yoff, dir_data_value)
if len(to_point) < 1:
return no_data_value
else:
return cu.get_raster_int_value(dataset_ol, to_point[0], to_point[1])
def surface_merge(upstream_inflow, new_slope_surface_id):
global slope_ds, dir_ds, water_buffer_value, water_value
# 定义合并后的坡面分组
slope_surface_unit = []
# 遍历上游流入点对坡面与水体相交处的入流口分组
for upstream_point in upstream_inflow:
# 判断坡面入流点已参与合并
not_in_unit = 1
for index in range(0, len(slope_surface_unit), 1):
if upstream_point in slope_surface_unit[index]:
not_in_unit = 0
break
# 若坡面入流点未参与合并
if not_in_unit:
# 创建新坡面集合并添加新坡面的首个水体入流点
slope_surface = [upstream_point]
# 获取新坡面集合的入流点集
slope_surface = get_new_slope_surface(upstream_point[0], upstream_point[1], upstream_inflow, slope_surface)
# 将新坡面放入合并后的坡面集合中
slope_surface_unit.append(slope_surface)
# 获取合并后新坡面的id
new_slope_surface_id += 1
# 以各合并坡面的入口点集开始合并坡面
for index in range(0, len(slope_surface_unit), 1):
# print(slope_surface_unit[index])
# 获取合并后新坡面与水体的交界点
s_s_inflows = slope_surface_unit[index]
for inflow_point in s_s_inflows:
# 初始化需要更新的像元集
to_update_points = [inflow_point]
# 循环直到待更新数组为空
while len(to_update_points) > 0:
# 取出第一个点
to_update_point = to_update_points.pop()
# 更新此点id
cu.set_raster_int_value(slope_ds, to_update_point[0], to_update_point[1], new_slope_surface_id)
# 获取此点周边需要更新的像元集
# 获取邻近像元
neighbor_points = cu.get_8_dir(to_update_point[0], to_update_point[1])
# 判断各像元是否需要更新
for point in neighbor_points:
point_x = point[0]
point_y = point[1]
judge_in_data = cu.in_data(point_x, point_y, slope_ds.RasterXSize, slope_ds.RasterYSize)
if judge_in_data:
# 若不为水体则继续
ol_data = cu.get_raster_int_value(slope_ds, point_x, point_y)
if ol_data != water_value:
# 若是当前点的上游点
dir_point = cu.off_transform(point_x, point_y, slope_ds, dir_ds)
dir_value = cu.get_raster_int_value(dir_ds, dir_point[0], dir_point[1])
to_point = cu.get_to_point(point_x, point_y, dir_value)
if to_point == to_update_point:
# 加入待更新数组
to_update_points.append(point)
return new_slope_surface_id
def down_point_coord(b_point, dir_index):
current_dir_path = find_data_by_point(b_point, dir_index)
cu_dir_ds = gdal.Open(current_dir_path)
dir_off = cu.coord_to_off(b_point, cu_dir_ds)
dir_value = cu.get_raster_int_value(cu_dir_ds, dir_off[0], dir_off[1])
to_off = cu.get_to_point(dir_off[0], dir_off[1], dir_value)
to_coord = cu.off_to_coord(to_off, cu_dir_ds)
cu_dir_ds = None
return to_coord
def judge_from_water(xoff, yoff):
global dir_ds, slope_ds, water_value
neighbor_points = cu.get_8_dir(xoff, yoff)
for point in neighbor_points:
dir_value = cu.get_raster_int_value(dir_ds, point[0], point[1])
to_point = cu.get_to_point(point[0], point[1], dir_value)
ol_value = cu.get_raster_int_value(slope_ds, point[0], point[1])
if ol_value == water_value and to_point[0] == xoff and to_point[1] == yoff:
return 1
return 0
def get_no_inflow_cells(dir_tif, no_inflow_tif):
dir_ds = gdal.Open(dir_tif)
# 创建无流入点数据
print("Create No inflow file...")
file_format = "GTiff"
driver = gdal.GetDriverByName(file_format)
full_geotransform = dir_ds.GetGeoTransform()
ni_ds = driver.Create(no_inflow_tif,
dir_ds.RasterXSize,
dir_ds.RasterYSize,
1,
gdal.GDT_Int32,
options=['COMPRESS=DEFLATE'])
ni_ds.SetGeoTransform(full_geotransform)
ni_ds.SetProjection(dir_ds.GetProjection())
ni_ds.GetRasterBand(1).SetNoDataValue(0)
# 遍历基础数据记录无流入点
print("Get No Inflow Data...")
for i in range(dir_ds.RasterYSize):
for j in range(dir_ds.RasterXSize):
c_point = [j, i]
# 邻近像元索引
n_points = cu.get_8_dir(c_point[0], c_point[1])
# 初始化流入标记
inflow_flag = 0
for n_point in n_points:
# 判断是否在数据内
in_data = cu.in_data(n_point[0], n_point[1],
dir_ds.RasterXSize, dir_ds.RasterYSize)
if in_data:
# 判断是否流向原像元点
n_dir = cu.get_raster_int_value(dir_ds, n_point[0],
n_point[1])
to_point = cu.get_to_point(n_point[0], n_point[1], n_dir)
if to_point == c_point:
inflow_flag = 1
break
# 若无流入则记录
if not inflow_flag:
cu.set_raster_int_value(ni_ds, c_point[0], c_point[1], 1)
def spill_point_dir(inner_ras_indexes, dem_ds):
global temp_ds
# 内边界的高程
inner_dem = []
for ras_index in inner_ras_indexes:
dem_value = cu.get_raster_float_value(dem_ds, ras_index[0],
ras_index[1])
inner_dem.append(dem_value)
# 获得最低位置的索引
spill_pt_index = list(map(inner_dem.index, heapq.nsmallest(1,
inner_dem)))[0]
spill_pt = inner_ras_indexes[spill_pt_index]
# 获得溢出方向
spill_dir = get_spill_dir(inner_ras_indexes, spill_pt_index)
# 得到接收点索引
reception_pt = cu.get_to_point(spill_pt[0], spill_pt[1], spill_dir)
cu.set_raster_int_value(temp_ds, spill_pt[0], spill_pt[1], spill_dir)
return spill_pt, spill_dir, reception_pt
def recode_from_river(watershed_id, river_x, river_y, watershed_tif_path, dir_tif_path, water_tif_path, rivers_index):
global water_value
# 创建数据集
watershed_ds = gdal.OpenEx(watershed_tif_path, 1)
dir_ds = gdal.Open(dir_tif_path)
water_ds = gdal.Open(water_tif_path)
# 初始化需要赋值的像元集合
update_cells = [[river_x, river_y]]
# 更新区域内像元值
# print(">>> update cell:", end='')
while len(update_cells) > 0:
# 取出要更新的像元索引
update_cell = update_cells.pop()
# print(update_cell, end='')
# 更新像元值
cu.set_raster_int_value(watershed_ds, update_cell[0], update_cell[1], watershed_id)
# print('update: ', update_cell, '->', watershed_id)
# 得到邻近像元集合
neighbor_cells = cu.get_8_dir(update_cell[0], update_cell[1])
# 搜索上游像元
for neighbor_cell in neighbor_cells:
n_x = neighbor_cell[0]
n_y = neighbor_cell[1]
# 判断邻近点是否在数据内
if cu.in_data(n_x, n_y, watershed_ds.RasterXSize, watershed_ds.RasterYSize):
# 若不为河段并不为湖泊/水库（即若为子流域）
water_off = cu.off_transform(n_x, n_y, watershed_ds, water_ds)
in_water = cu.is_water_cell(water_ds, water_off[0], water_off[1], water_value)
if neighbor_cell not in rivers_index and not in_water:
dir_value = cu.get_raster_int_value(dir_ds, n_x, n_y)
to_point = cu.get_to_point(n_x, n_y, dir_value)
# 若为上游点
if to_point == update_cell:
# 加入更新数组
update_cells.append(neighbor_cell)
watershed_ds = None
dir_ds = None
def able_update_acc(b_point, x_size, y_size, acc_index, dir_index):
global nc
# 获取acc整体数据边界
total_lt = list(acc_index.values())[0][0][:]
total_rb = list(acc_index.values())[0][0][:]
for value in acc_index.values():
if total_lt[0] >= value[0][0]:
total_lt[0] = value[0][0]
if total_lt[1] <= value[0][1]:
total_lt[1] = value[0][1]
if total_rb[0] <= value[1][0]:
total_rb[0] = value[1][0]
if total_rb[1] >= value[1][1]:
total_rb[1] = value[1][1]
# 初始化上游像元汇流累积量(初始化为0)
up_cells_acc = 0
# 初始化标记存在上游像元(初始化为不存在)
up_cell_flag = 0
# 初始化标记存在上游像元acc暂为nodata(初始化为存在nodata上游点)
up_nodata_flag = 0
# 获取周边像元坐标(左上角)从邻近像元搜索
n_cells = cu.get_8_dir_coord(b_point[0], b_point[1], x_size, y_size)
for n_cell in n_cells:
# 得到所处数据块
acc_path = find_data_by_point(n_cell, acc_index)
dir_path = find_data_by_point(n_cell, dir_index)
if acc_path != "" and dir_path != "":
# 判断是否为上游点
dir_ds = gdal.Open(dir_path)
# 根据坐标转为索引
n_off = cu.coord_to_off(n_cell, dir_ds)
# 获得此处流向值
dir_value = cu.get_raster_int_value(dir_ds, n_off[0], n_off[1])
# 获得下游像元索引
d_off = cu.get_to_point(n_off[0], n_off[1], dir_value)
# 若为有下游像元
if d_off:
# 获取该下游像元的坐标
d_coord = cu.off_to_coord(d_off, dir_ds)
# 若该下游像元坐标与当前坐标相同则此邻近像元为上游像元
if d_coord == b_point:
up_cell_flag = 1
# 获取此上游像元处的汇流累积量值
acc_ds = gdal.Open(acc_path)
n_acc_off = cu.coord_to_off(n_cell, acc_ds)
n_acc = cu.get_raster_float_value(acc_ds, n_acc_off[0],
n_acc_off[1])
# 若此acc值为nodata则标记存在上游像元acc暂为nodata
if int(n_acc) <= 0:
up_nodata_flag = 1
break
# 否则记录当前值到总上游累积量
else:
up_cells_acc += n_acc
acc_ds = None
dir_ds = None
# 若存在上游像元
if up_cell_flag:
# 若上游像元均有有效acc值
if up_nodata_flag == 0:
# 更新上游acc总值和自身累积量值到当前像元
current_acc = up_cells_acc + 1.0
return current_acc
# 存在上游像元暂为nodata
else:
return 0
# 若不存在上游像元
else:
# 若考虑边界污染
if nc:
# 若在边界上
if in_edge(b_point, total_lt, total_rb):
return 0
else:
# 否则作为汇流起始点
return 1.0
# 若不考虑边界污染
else:
return 1.0
def watershed_recode(river_record_path, watershed_tif_path, dir_tif_path, water_tif_path):
# 创建数据集
watershed_ds = gdal.Open(watershed_tif_path)
dir_ds = gdal.Open(dir_tif_path)
# 读取河系信息到内存
print("Reading rivers info...")
# 初始化河系像元索引
rivers_index = []
rivers_info = {}
with open(river_record_path, 'r') as river_file:
for line in river_file.readlines():
river_info_str = line.strip('\n')
river_info = river_info_str.split(',')
# 将river像元的x,y索引作为字典的key
river_point_str = river_info[0] + ',' + river_info[1]
# 将river像元的流向和汇流累积量作为字典的value
river_info_detail = float(river_info[2])
rivers_info[river_point_str] = river_info_detail
rivers_index.append([int(river_info[0]), int(river_info[1])])
# 创建河段集合
print("Building rivers group...")
river_reaches = {}
# 继续创建标识
flag = 1
# 开始创建
while flag:
flag = 0
# 取汇流累积量最小的像元索引
min_acc_river = min(rivers_info, key=rivers_info.get)
# 获取子流域标识值
current_cell = min_acc_river.split(',')
watershed_value = cu.get_raster_int_value(watershed_ds, int(current_cell[0]), int(current_cell[1]))
# 若不在子流域则继续寻找
if watershed_value < 0:
del rivers_info[min_acc_river]
if rivers_info:
flag = 1
# 若在子流域则作为起点
else:
# 初始化要遍历的河流起点数组
river_starts = [min_acc_river]
# 若河流起点数组不为空则遍历河段
while len(river_starts) > 0:
# 开始追踪河段
start = river_starts.pop()
# 创建河段索引记录数组
river_reach = []
# 获得此像元索引
start_xy = start.split(',')
# 获取此河段对应子流域Id标识
watershed_id = cu.get_raster_int_value(watershed_ds, int(start_xy[0]), int(start_xy[1]))
# 初始化追踪数组
tracking = [start]
# 开始追踪河段的像元
while len(tracking) > 0:
# 取出像元并记录
cell_str = tracking.pop()
del rivers_info[cell_str]
river_reach.append(cell_str)
# 获得像元xy索引
cell_xy = cell_str.split(',')
n_x = cell_xy[0]
n_y = cell_xy[1]
# 获得流向的像元
dir = cu.get_raster_int_value(dir_ds, int(n_x), int(n_y))
to_point = cu.get_to_point(int(n_x), int(n_y), dir)
to_point_str = ','.join(str(k) for k in to_point)
# 若下游像元在原河流数组
if to_point_str in rivers_info.keys():
# 获取此点在子流域数据的值
watershed_value = cu.get_raster_int_value(watershed_ds, to_point[0], to_point[1])
# 若在子流域内
if watershed_value >= 0:
# 若在同子流域内则记录为河段
if watershed_value == watershed_id:
tracking.append(to_point_str)
# 若不在同一子流域则记录为下一河段起点
else:
river_starts.append(to_point_str)
# 若不在子流域内则搜索下游河段的起点
else:
# 初始化搜索数组
find_next_start = [to_point]
# 开始搜索
while len(find_next_start) > 0:
# 取出像元
river_cell = find_next_start.pop()
# 在原河流数组中清除记录
river_cell_str = ','.join(str(k) for k in river_cell)
del rivers_info[river_cell_str]
# 获得下游像元
dir = cu.get_raster_int_value(dir_ds, river_cell[0], river_cell[1])
next_point = cu.get_to_point(dir, river_cell[0], river_cell[1])
# 若下游像元在原河流数组
next_point_str = ','.join(str(k) for k in next_point)
if next_point_str in rivers_info.keys():
# 获取此点在子流域数据的值
watershed_value = cu.get_raster_int_value(watershed_ds, next_point[0],
next_point[1])
# 若在子流域内
if watershed_value >= 0:
# 则记录为下一河段起点
river_starts.append(next_point_str)
# 若不在子流域内则继续搜索下游河段的起点
else:
find_next_start.append(next_point)
# 将河段存入集合
if watershed_id in river_reaches.keys():
update_array = river_reaches[watershed_id]
update_array.append(river_reach)
river_reaches[watershed_id] = update_array
else:
river_reaches[watershed_id] = [river_reach]
# 若河流像元未全部遍历则继续
if len(rivers_info) > 0:
flag = 1
print("Searching rivers in watersheds need recode...")
# 记录当前子流域标识最大值
max_ws_id = 0
# 初始化在需要重新编码子流域中的河段集合
recode_ws_rivers = []
# 筛选出涉及处理的子流域
for watershed, watershed_rivers in river_reaches.items():
watershed_id = int(watershed)
if watershed_id > max_ws_id:
max_ws_id = watershed_id
# 若存在多个河段则需要进行重新唯一编码
if len(watershed_rivers) > 1:
# print("%s, %d" % (watershed, len(watershed_rivers)))
# print(watershed_rivers)
# 找出每组中除最长河段的其他河段
max_len = 0
max_river_start = ''
for river in watershed_rivers:
if len(river) > max_len:
max_len = len(river)
max_river_start = river[0]
for river in watershed_rivers:
if len(river) <= max_len and max_river_start not in river:
recode_ws_rivers.append(river)
print("Recoding watersheds by rivers...")
# 在需要重新编码的子流域中的河段
watershed_count = 0
for river in recode_ws_rivers:
max_ws_id += 1
watershed_count += 1
total_count = len(recode_ws_rivers)
print("\r> updating watersheds.....................%d/%d" % (watershed_count, total_count), end='')
river_count = 0
print('')
for river_cell in river:
river_count += 1
total_r_count = len(river)
print("\r>> update by river cell: %d/%d" % (river_count, total_r_count), end='')
cell_xy = river_cell.split(',')
recode_from_river(max_ws_id, int(cell_xy[0]), int(cell_xy[1]), watershed_tif_path, dir_tif_path,
water_tif_path, rivers_index)
print(".")
watershed_ds = None
dir_ds = None
def get_trace_points(dir_tif, flag_tif, trace_tif, seaside_txt=None, final_txt=None):
coastline_value = 0
final_value = -1
dir_ds = gdal.Open(dir_tif)
f_ds = gdal.Open(flag_tif)
# 创建trace start数据
print("Create Trace file...")
file_format = "GTiff"
driver = gdal.GetDriverByName(file_format)
full_geotransform = dir_ds.GetGeoTransform()
trace_ds = driver.Create(trace_tif, dir_ds.RasterXSize, dir_ds.RasterYSize, 1, gdal.GDT_Byte, options=['COMPRESS=DEFLATE'])
trace_ds.SetGeoTransform(full_geotransform)
trace_ds.SetProjection(dir_ds.GetProjection())
trace_ds.GetRasterBand(1).SetNoDataValue(0)
# 创建记录文件
seaside_flag = 0
seaside_f = None
if seaside_txt is not None:
print("Record Final Points.")
seaside_flag = 1
if os.path.exists(seaside_txt):
os.remove(seaside_txt)
if seaside_flag:
seaside_f = open(seaside_txt, "a")
final_flag = 0
final_f = None
if final_txt is not None:
print("Record Final Points.")
final_flag = 1
if os.path.exists(final_txt):
os.remove(final_txt)
if final_flag:
final_f = open(final_txt, "a")
# 遍历基础数据记录seaside
print("Get Trace Data...")
for i in range(f_ds.RasterYSize):
for j in range(f_ds.RasterXSize):
# 获取流向的值
flag_value = cu.get_raster_int_value(f_ds, j, i)
# 如果是coastline
if flag_value == coastline_value:
# 获取周边像元索引
neibor_cells = cu.get_8_dir(j, i)
# 遍历
for n_cell in neibor_cells:
# 判断是否在数据内
in_data = cu.in_data(n_cell[0], n_cell[1], dir_ds.RasterXSize, dir_ds.RasterYSize)
# 若在数据内
if in_data:
# 获取流向值
dir_value = cu.get_raster_int_value(dir_ds, n_cell[0], n_cell[1])
# 获取流向的像元索引
n_to_point = cu.get_to_point(n_cell[0], n_cell[1], dir_value)
# 若为当前上游像元
if n_to_point == [j, i]:
# 记录为seaside
cu.set_raster_int_value(trace_ds, n_cell[0], n_cell[1], 1)
# 记录到文件
if seaside_flag:
# 获得栅格像元左上角坐标
seaside_record_item = cu.off_to_coord([j, i], f_ds)
seaside_record_str = ','.join(str(k) for k in seaside_record_item)
seaside_f.write(seaside_record_str + '\n')
# 如果是final point
elif flag_value == final_value:
# 记录为final points
cu.set_raster_int_value(trace_ds, j, i, 1)
# 记录到文件
if final_flag:
final_record_item = cu.off_to_coord([j, i], f_ds)
final_record_str = ','.join(str(k) for k in final_record_item)
final_f.write(final_record_str + '\n')
dir_ds = None
f_ds = None
trace_ds = None
seaside_f = None
final_f = None
