def add_group_field(shp_file, subbasin_field_name, group_metis_dict):
"""add group information to subbasin ESRI shapefile
Args:
shp_file: Subbasin Shapefile
subbasin_field_name: field name of subbasin
group_metis_dict: returned by func`metis_partition`
"""
if not group_metis_dict:
return
ds_reach = ogr_Open(shp_file, update=True)
layer_reach = ds_reach.GetLayer(0)
layer_def = layer_reach.GetLayerDefn()
icode = layer_def.GetFieldIndex(str(subbasin_field_name))
igrp = layer_def.GetFieldIndex(str(ImportReaches2Mongo._GROUP))
ikgrp = layer_def.GetFieldIndex(str(ImportReaches2Mongo._KMETIS))
ipgrp = layer_def.GetFieldIndex(str(ImportReaches2Mongo._PMETIS))
if igrp < 0:
new_field = ogr_FieldDefn(str(ImportReaches2Mongo._GROUP),
OFTInteger)
layer_reach.CreateField(new_field)
if ikgrp < 0:
new_field = ogr_FieldDefn(str(ImportReaches2Mongo._KMETIS),
OFTInteger)
layer_reach.CreateField(new_field)
if ipgrp < 0:
new_field = ogr_FieldDefn(str(ImportReaches2Mongo._PMETIS),
OFTInteger)
layer_reach.CreateField(new_field)
ftmap = dict()
layer_reach.ResetReading()
ft = layer_reach.GetNextFeature()
while ft is not None:
tmpid = ft.GetFieldAsInteger(icode)
ftmap[tmpid] = ft
ft = layer_reach.GetNextFeature()
groups = group_metis_dict[1]['group']
for i, n in enumerate(groups):
for node, d in group_metis_dict.items():
ftmap[node].SetField(str(ImportReaches2Mongo._GROUP), n)
ftmap[node].SetField(str(ImportReaches2Mongo._KMETIS),
d['kmetis'][i])
ftmap[node].SetField(str(ImportReaches2Mongo._PMETIS),
d['pmetis'][i])
layer_reach.SetFeature(ftmap[node])
# copy the reach file to new file
prefix = os.path.splitext(shp_file)[0]
dstfile = prefix + "_" + str(n) + ".shp"
FileClass.copy_files(shp_file, dstfile)
layer_reach.SyncToDisk()
ds_reach.Destroy()
del ds_reach
def add_group_field(shp_file, subbasin_field_name, group_metis_dict):
"""add group information to subbasin ESRI shapefile
Args:
shp_file: Subbasin Shapefile
subbasin_field_name: field name of subbasin
group_metis_dict: returned by func`metis_partition`
"""
if not group_metis_dict:
return
ds_reach = ogr_Open(shp_file, update=True)
layer_reach = ds_reach.GetLayer(0)
layer_def = layer_reach.GetLayerDefn()
icode = layer_def.GetFieldIndex(subbasin_field_name)
igrp = layer_def.GetFieldIndex(ImportReaches2Mongo._GROUP)
ikgrp = layer_def.GetFieldIndex(ImportReaches2Mongo._KMETIS)
ipgrp = layer_def.GetFieldIndex(ImportReaches2Mongo._PMETIS)
if igrp < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._GROUP, OFTInteger)
layer_reach.CreateField(new_field)
if ikgrp < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._KMETIS, OFTInteger)
layer_reach.CreateField(new_field)
if ipgrp < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._PMETIS, OFTInteger)
layer_reach.CreateField(new_field)
ftmap = dict()
layer_reach.ResetReading()
ft = layer_reach.GetNextFeature()
while ft is not None:
tmpid = ft.GetFieldAsInteger(icode)
ftmap[tmpid] = ft
ft = layer_reach.GetNextFeature()
groups = group_metis_dict[1]['group']
for i, n in enumerate(groups):
for node, d in group_metis_dict.items():
ftmap[node].SetField(ImportReaches2Mongo._GROUP, n)
ftmap[node].SetField(ImportReaches2Mongo._KMETIS, d['kmetis'][i])
ftmap[node].SetField(ImportReaches2Mongo._PMETIS, d['pmetis'][i])
layer_reach.SetFeature(ftmap[node])
# copy the reach file to new file
prefix = os.path.splitext(shp_file)[0]
dstfile = prefix + "_" + str(n) + ".shp"
FileClass.copy_files(shp_file, dstfile)
layer_reach.SyncToDisk()
ds_reach.Destroy()
del ds_reach
def post_process_of_delineated_data(cfg):
"""Do some necessary transfer for subbasin, stream, and flow direction raster."""
# inputs
stream_net_file = cfg.taudems.streamnet_shp
subbasin_file = cfg.taudems.subbsn_m
flow_dir_file_tau = cfg.taudems.d8flow_m
stream_raster_file = cfg.taudems.stream_m
# outputs
# -- shapefile
shp_dir = cfg.dirs.geoshp
UtilClass.mkdir(shp_dir)
# ---- outlet, copy from DirNameUtils.TauDEM
FileClass.copy_files(cfg.taudems.outlet_m, cfg.vecs.outlet)
# ---- reaches
output_reach_file = cfg.vecs.reach
# ---- subbasins
subbasin_vector_file = cfg.vecs.subbsn
# -- raster file
output_subbasin_file = cfg.spatials.subbsn
output_flow_dir_file = cfg.spatials.d8flow
output_stream_link_file = cfg.spatials.stream_link
output_hillslope_file = cfg.spatials.hillslope
id_map = StreamnetUtil.serialize_streamnet(stream_net_file,
output_reach_file)
RasterUtilClass.raster_reclassify(subbasin_file, id_map,
output_subbasin_file, GDT_Int32)
StreamnetUtil.assign_stream_id_raster(stream_raster_file,
output_subbasin_file,
output_stream_link_file)
# Convert D8 encoding rule to ArcGIS
D8Util.convert_code(flow_dir_file_tau, output_flow_dir_file)
# convert raster to shapefile (for subbasin and basin)
print('Generating subbasin vector...')
VectorUtilClass.raster2shp(output_subbasin_file, subbasin_vector_file,
'subbasin', FieldNames.subbasin_id)
mask_file = cfg.spatials.mask
basin_vector = cfg.vecs.bsn
print('Generating basin vector...')
VectorUtilClass.raster2shp(mask_file, basin_vector, 'basin',
FieldNames.basin)
# delineate hillslope
DelineateHillslope.downstream_method_whitebox(output_stream_link_file,
flow_dir_file_tau,
output_hillslope_file)
def post_process_of_delineated_data(cfg):
"""Do some necessary transfer for subbasin, stream, and flow direction raster."""
# inputs
stream_net_file = cfg.taudems.streamnet_shp
subbasin_file = cfg.taudems.subbsn_m
flow_dir_file_tau = cfg.taudems.d8flow_m
stream_raster_file = cfg.taudems.stream_m
# outputs
# -- shapefile
shp_dir = cfg.dirs.geoshp
UtilClass.mkdir(shp_dir)
# ---- outlet, copy from DirNameUtils.TauDEM
FileClass.copy_files(cfg.taudems.outlet_m, cfg.vecs.outlet)
# ---- reaches
output_reach_file = cfg.vecs.reach
# ---- subbasins
subbasin_vector_file = cfg.vecs.subbsn
# -- raster file
output_subbasin_file = cfg.spatials.subbsn
output_flow_dir_file = cfg.spatials.d8flow
output_stream_link_file = cfg.spatials.stream_link
output_hillslope_file = cfg.spatials.hillslope
id_map = StreamnetUtil.serialize_streamnet(stream_net_file, output_reach_file)
RasterUtilClass.raster_reclassify(subbasin_file, id_map, output_subbasin_file, GDT_Int32)
StreamnetUtil.assign_stream_id_raster(stream_raster_file, output_subbasin_file,
output_stream_link_file)
# Convert D8 encoding rule to ArcGIS
D8Util.convert_code(flow_dir_file_tau, output_flow_dir_file)
# convert raster to shapefile (for subbasin and basin)
print('Generating subbasin vector...')
VectorUtilClass.raster2shp(output_subbasin_file, subbasin_vector_file, 'subbasin',
FieldNames.subbasin_id)
mask_file = cfg.spatials.mask
basin_vector = cfg.vecs.bsn
print('Generating basin vector...')
VectorUtilClass.raster2shp(mask_file, basin_vector, 'basin', FieldNames.basin)
# delineate hillslope
DelineateHillslope.downstream_method_whitebox(output_stream_link_file, flow_dir_file_tau,
output_hillslope_file)
def serialize_streamnet(streamnet_file, output_reach_file):
"""Eliminate reach with zero length and return the reach ID map.
Args:
streamnet_file: original stream net ESRI shapefile
output_reach_file: serialized stream net, ESRI shapefile
Returns:
id pairs {origin: newly assigned}
"""
FileClass.copy_files(streamnet_file, output_reach_file)
ds_reach = ogr_Open(output_reach_file, update=True)
layer_reach = ds_reach.GetLayer(0)
layer_def = layer_reach.GetLayerDefn()
i_link = layer_def.GetFieldIndex(FLD_LINKNO)
i_link_downslope = layer_def.GetFieldIndex(FLD_DSLINKNO)
i_len = layer_def.GetFieldIndex(REACH_LENGTH)
old_id_list = []
# there are some reaches with zero length.
# this program will remove these zero-length reaches
# output_dic is used to store the downstream reaches of these zero-length
# reaches
output_dic = {}
ft = layer_reach.GetNextFeature()
while ft is not None:
link_id = ft.GetFieldAsInteger(i_link)
reach_len = ft.GetFieldAsDouble(i_len)
if link_id not in old_id_list:
if reach_len < DELTA:
downstream_id = ft.GetFieldAsInteger(i_link_downslope)
output_dic[link_id] = downstream_id
else:
old_id_list.append(link_id)
ft = layer_reach.GetNextFeature()
old_id_list.sort()
id_map = {}
for i, old_id in enumerate(old_id_list):
id_map[old_id] = i + 1
# print(id_map)
# change old ID to new ID
layer_reach.ResetReading()
ft = layer_reach.GetNextFeature()
while ft is not None:
link_id = ft.GetFieldAsInteger(i_link)
if link_id not in id_map:
layer_reach.DeleteFeature(ft.GetFID())
ft = layer_reach.GetNextFeature()
continue
ds_id = ft.GetFieldAsInteger(i_link_downslope)
ds_id = output_dic.get(ds_id, ds_id)
ds_id = output_dic.get(ds_id, ds_id)
ft.SetField(FLD_LINKNO, id_map[link_id])
if ds_id in id_map:
ft.SetField(FLD_DSLINKNO, id_map[ds_id])
else:
# print(ds_id)
ft.SetField(FLD_DSLINKNO, -1)
layer_reach.SetFeature(ft)
ft = layer_reach.GetNextFeature()
ds_reach.ExecuteSQL("REPACK reach")
layer_reach.SyncToDisk()
ds_reach.Destroy()
del ds_reach
return id_map
def serialize_streamnet(streamnet_file, output_reach_file):
"""Eliminate reach with zero length and return the reach ID map.
Args:
streamnet_file: original stream net ESRI shapefile
output_reach_file: serialized stream net, ESRI shapefile
Returns:
id pairs {origin: newly assigned}
"""
FileClass.copy_files(streamnet_file, output_reach_file)
ds_reach = ogr_Open(output_reach_file, update=True)
layer_reach = ds_reach.GetLayer(0)
layer_def = layer_reach.GetLayerDefn()
i_link = layer_def.GetFieldIndex(FLD_LINKNO)
i_link_downslope = layer_def.GetFieldIndex(FLD_DSLINKNO)
i_len = layer_def.GetFieldIndex(REACH_LENGTH)
old_id_list = []
# there are some reaches with zero length.
# this program will remove these zero-length reaches
# output_dic is used to store the downstream reaches of these zero-length
# reaches
output_dic = {}
ft = layer_reach.GetNextFeature()
while ft is not None:
link_id = ft.GetFieldAsInteger(i_link)
reach_len = ft.GetFieldAsDouble(i_len)
if link_id not in old_id_list:
if reach_len < DELTA:
downstream_id = ft.GetFieldAsInteger(i_link_downslope)
output_dic[link_id] = downstream_id
else:
old_id_list.append(link_id)
ft = layer_reach.GetNextFeature()
old_id_list.sort()
id_map = {}
for i, old_id in enumerate(old_id_list):
id_map[old_id] = i + 1
# print(id_map)
# change old ID to new ID
layer_reach.ResetReading()
ft = layer_reach.GetNextFeature()
while ft is not None:
link_id = ft.GetFieldAsInteger(i_link)
if link_id not in id_map:
layer_reach.DeleteFeature(ft.GetFID())
ft = layer_reach.GetNextFeature()
continue
ds_id = ft.GetFieldAsInteger(i_link_downslope)
ds_id = output_dic.get(ds_id, ds_id)
ds_id = output_dic.get(ds_id, ds_id)
ft.SetField(FLD_LINKNO, id_map[link_id])
if ds_id in id_map:
ft.SetField(FLD_DSLINKNO, id_map[ds_id])
else:
# print(ds_id)
ft.SetField(FLD_DSLINKNO, -1)
layer_reach.SetFeature(ft)
ft = layer_reach.GetNextFeature()
ds_reach.ExecuteSQL(str('REPACK reach'))
layer_reach.SyncToDisk()
ds_reach.Destroy()
del ds_reach
return id_map
