def add_channel_width_to_shp(reach_shp_file,
stream_link_file,
width_data,
default_depth=1.5):
"""Add channel/reach width and default depth to ESRI shapefile"""
stream_link = RasterUtilClass.read_raster(stream_link_file)
n_rows = stream_link.nRows
n_cols = stream_link.nCols
nodata_value = stream_link.noDataValue
data_stream = stream_link.data
ch_width_dic = dict()
ch_num_dic = dict()
for i in range(n_rows):
for j in range(n_cols):
if abs(data_stream[i][j] - nodata_value) > UTIL_ZERO:
tmpid = int(data_stream[i][j])
ch_num_dic.setdefault(tmpid, 0)
ch_width_dic.setdefault(tmpid, 0)
ch_num_dic[tmpid] += 1
ch_width_dic[tmpid] += width_data[i][j]
for k in ch_num_dic:
ch_width_dic[k] /= ch_num_dic[k]
# add channel width_data field to reach shp file
ds_reach = ogr_Open(reach_shp_file, update=True)
layer_reach = ds_reach.GetLayer(0)
layer_def = layer_reach.GetLayerDefn()
i_link = layer_def.GetFieldIndex(ImportReaches2Mongo._LINKNO)
i_width = layer_def.GetFieldIndex(ImportReaches2Mongo._WIDTH)
i_depth = layer_def.GetFieldIndex(ImportReaches2Mongo._DEPTH)
if i_width < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._WIDTH, OFTReal)
layer_reach.CreateField(new_field)
if i_depth < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._DEPTH, OFTReal)
layer_reach.CreateField(new_field)
# grid_code:feature map
# ftmap = {}
layer_reach.ResetReading()
ft = layer_reach.GetNextFeature()
while ft is not None:
tmpid = ft.GetFieldAsInteger(i_link)
w = 1
if tmpid in list(ch_width_dic.keys()):
w = ch_width_dic[tmpid]
ft.SetField(ImportReaches2Mongo._WIDTH, w)
ft.SetField(ImportReaches2Mongo._DEPTH, default_depth)
layer_reach.SetFeature(ft)
ft = layer_reach.GetNextFeature()
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(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, n, group_kmetis,
group_pmetis, ns):
"""add group information to subbasin ESRI shapefile
Args:
shp_file: Subbasin Shapefile
subbasin_field_name: field name of subbasin
n: divide number
group_kmetis: kmetis
group_pmetis: pmetis
ns: a list of the nodes in the graph
Returns:
group_dic: group dict
group_dic_pmetis: pmetis dict
"""
ds_reach = ogr_Open(shp_file, update=True)
layer_reach = ds_reach.GetLayer(0)
layer_def = layer_reach.GetLayerDefn()
i_code = layer_def.GetFieldIndex(subbasin_field_name)
i_group = layer_def.GetFieldIndex(ImportReaches2Mongo._GROUP)
i_group_pmetis = layer_def.GetFieldIndex(ImportReaches2Mongo._PMETIS)
if i_group < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._GROUP, OFTInteger)
layer_reach.CreateField(new_field)
if i_group_pmetis < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._PMETIS, OFTInteger)
layer_reach.CreateField(new_field)
# grid_code:feature map
ftmap = dict()
layer_reach.ResetReading()
ft = layer_reach.GetNextFeature()
while ft is not None:
tmpid = ft.GetFieldAsInteger(i_code)
ftmap[tmpid] = ft
ft = layer_reach.GetNextFeature()
group_dic = dict()
group_dic_pmetis = dict()
i = 0
for node in ns:
group_dic[node] = group_kmetis[i]
group_dic_pmetis[node] = group_pmetis[i]
ftmap[node].SetField(ImportReaches2Mongo._GROUP, group_kmetis[i])
ftmap[node].SetField(ImportReaches2Mongo._PMETIS, group_pmetis[i])
layer_reach.SetFeature(ftmap[node])
i += 1
layer_reach.SyncToDisk()
ds_reach.Destroy()
del ds_reach
# 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)
return group_dic, group_dic_pmetis
def add_channel_width_to_shp(reach_shp_file, stream_link_file,
width_data, default_depth=1.5):
"""Add channel/reach width and default depth to ESRI shapefile"""
stream_link = RasterUtilClass.read_raster(stream_link_file)
n_rows = stream_link.nRows
n_cols = stream_link.nCols
nodata_value = stream_link.noDataValue
data_stream = stream_link.data
ch_width_dic = dict()
ch_num_dic = dict()
for i in range(n_rows):
for j in range(n_cols):
if abs(data_stream[i][j] - nodata_value) > UTIL_ZERO:
tmpid = int(data_stream[i][j])
ch_num_dic.setdefault(tmpid, 0)
ch_width_dic.setdefault(tmpid, 0)
ch_num_dic[tmpid] += 1
ch_width_dic[tmpid] += width_data[i][j]
for k in ch_num_dic:
ch_width_dic[k] /= ch_num_dic[k]
# add channel width_data field to reach shp file
ds_reach = ogr_Open(reach_shp_file, update=True)
layer_reach = ds_reach.GetLayer(0)
layer_def = layer_reach.GetLayerDefn()
i_link = layer_def.GetFieldIndex(ImportReaches2Mongo._LINKNO)
i_width = layer_def.GetFieldIndex(ImportReaches2Mongo._WIDTH)
i_depth = layer_def.GetFieldIndex(ImportReaches2Mongo._DEPTH)
if i_width < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._WIDTH, OFTReal)
layer_reach.CreateField(new_field)
if i_depth < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._DEPTH, OFTReal)
layer_reach.CreateField(new_field)
# grid_code:feature map
# ftmap = {}
layer_reach.ResetReading()
ft = layer_reach.GetNextFeature()
while ft is not None:
tmpid = ft.GetFieldAsInteger(i_link)
w = 1
if tmpid in list(ch_width_dic.keys()):
w = ch_width_dic[tmpid]
ft.SetField(ImportReaches2Mongo._WIDTH, w)
ft.SetField(ImportReaches2Mongo._DEPTH, default_depth)
layer_reach.SetFeature(ft)
ft = layer_reach.GetNextFeature()
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 raster2shp(rasterfile,
vectorshp,
layername=None,
fieldname=None,
band_num=1,
mask='default'):
"""Convert raster to ESRI shapefile"""
FileClass.remove_files(vectorshp)
FileClass.check_file_exists(rasterfile)
# this allows GDAL to throw Python Exceptions
gdal.UseExceptions()
src_ds = gdal.Open(rasterfile)
if src_ds is None:
print('Unable to open %s' % rasterfile)
sys.exit(1)
try:
srcband = src_ds.GetRasterBand(band_num)
except RuntimeError as e:
# for example, try GetRasterBand(10)
print('Band ( %i ) not found, %s' % (band_num, e))
sys.exit(1)
if mask == 'default':
maskband = srcband.GetMaskBand()
elif mask is None or mask.upper() == 'NONE':
maskband = None
else:
mask_ds = gdal.Open(mask)
maskband = mask_ds.GetRasterBand(1)
# create output datasource
if layername is None:
layername = FileClass.get_core_name_without_suffix(rasterfile)
drv = ogr_GetDriverByName(str('ESRI Shapefile'))
dst_ds = drv.CreateDataSource(vectorshp)
srs = None
if src_ds.GetProjection() != '':
srs = osr_SpatialReference()
srs.ImportFromWkt(src_ds.GetProjection())
dst_layer = dst_ds.CreateLayer(str(layername), srs=srs)
if fieldname is None:
fieldname = layername.upper()
fd = ogr_FieldDefn(str(fieldname), OFTInteger)
dst_layer.CreateField(fd)
dst_field = 0
result = gdal.Polygonize(srcband,
maskband,
dst_layer,
dst_field, ['8CONNECTED=8'],
callback=None)
return result
def raster2shp(rasterfile, vectorshp, layername=None, fieldname=None,
band_num=1, mask='default'):
"""Convert raster to ESRI shapefile"""
FileClass.remove_files(vectorshp)
FileClass.check_file_exists(rasterfile)
# this allows GDAL to throw Python Exceptions
gdal.UseExceptions()
src_ds = gdal.Open(rasterfile)
if src_ds is None:
print('Unable to open %s' % rasterfile)
sys.exit(1)
try:
srcband = src_ds.GetRasterBand(band_num)
except RuntimeError as e:
# for example, try GetRasterBand(10)
print('Band ( %i ) not found, %s' % (band_num, e))
sys.exit(1)
if mask == 'default':
maskband = srcband.GetMaskBand()
elif mask is None or mask.upper() == 'NONE':
maskband = None
else:
mask_ds = gdal.Open(mask)
maskband = mask_ds.GetRasterBand(1)
# create output datasource
if layername is None:
layername = FileClass.get_core_name_without_suffix(rasterfile)
drv = ogr_GetDriverByName(str('ESRI Shapefile'))
dst_ds = drv.CreateDataSource(vectorshp)
srs = None
if src_ds.GetProjection() != '':
srs = osr_SpatialReference()
srs.ImportFromWkt(src_ds.GetProjection())
dst_layer = dst_ds.CreateLayer(str(layername), srs=srs)
if fieldname is None:
fieldname = layername.upper()
fd = ogr_FieldDefn(str(fieldname), OFTInteger)
dst_layer.CreateField(fd)
dst_field = 0
result = gdal.Polygonize(srcband, maskband, dst_layer, dst_field,
['8CONNECTED=8'], callback=None)
return result
def add_channel_width_depth_to_shp(reach_shp_file, stream_link_file,
width_file, depth_file):
"""Calculate average channel width and depth, and add or modify the attribute table
of reach.shp
"""
stream_link = RasterUtilClass.read_raster(stream_link_file)
n_rows = stream_link.nRows
n_cols = stream_link.nCols
nodata_value = stream_link.noDataValue
data_stream = stream_link.data
width = RasterUtilClass.read_raster(width_file)
width_data = width.data
depth = RasterUtilClass.read_raster(depth_file)
depth_data = depth.data
ch_width_dic = dict()
ch_depth_dic = dict()
ch_num_dic = dict()
for i in range(n_rows):
for j in range(n_cols):
if abs(data_stream[i][j] - nodata_value) <= UTIL_ZERO:
continue
tmpid = int(data_stream[i][j])
ch_num_dic.setdefault(tmpid, 0)
ch_width_dic.setdefault(tmpid, 0)
ch_depth_dic.setdefault(tmpid, 0)
ch_num_dic[tmpid] += 1
ch_width_dic[tmpid] += width_data[i][j]
ch_depth_dic[tmpid] += depth_data[i][j]
for k in ch_num_dic:
ch_width_dic[k] /= ch_num_dic[k]
ch_depth_dic[k] /= ch_num_dic[k]
# add channel width and depth fields to reach shp file or update values if the fields exist
ds_reach = ogr_Open(reach_shp_file, update=True)
layer_reach = ds_reach.GetLayer(0)
layer_def = layer_reach.GetLayerDefn()
i_link = layer_def.GetFieldIndex(str(ImportReaches2Mongo._LINKNO))
i_width = layer_def.GetFieldIndex(str(ImportReaches2Mongo._WIDTH))
i_depth = layer_def.GetFieldIndex(str(ImportReaches2Mongo._DEPTH))
if i_width < 0:
new_field = ogr_FieldDefn(str(ImportReaches2Mongo._WIDTH), OFTReal)
layer_reach.CreateField(new_field)
if i_depth < 0:
new_field = ogr_FieldDefn(str(ImportReaches2Mongo._DEPTH), OFTReal)
layer_reach.CreateField(new_field)
layer_reach.ResetReading()
ft = layer_reach.GetNextFeature()
while ft is not None:
tmpid = ft.GetFieldAsInteger(i_link)
w = 5.
d = 1.5
if tmpid in ch_width_dic:
w = ch_width_dic[tmpid]
if tmpid in ch_depth_dic:
d = ch_depth_dic[tmpid]
ft.SetField(str(ImportReaches2Mongo._WIDTH), w)
ft.SetField(str(ImportReaches2Mongo._DEPTH), d)
layer_reach.SetFeature(ft)
ft = layer_reach.GetNextFeature()
layer_reach.SyncToDisk()
ds_reach.Destroy()
del ds_reach
def add_channel_width_depth_to_shp(reach_shp_file, stream_link_file, width_file, depth_file):
"""Calculate average channel width and depth, and add or modify the attribute table
of reach.shp
"""
stream_link = RasterUtilClass.read_raster(stream_link_file)
n_rows = stream_link.nRows
n_cols = stream_link.nCols
nodata_value = stream_link.noDataValue
data_stream = stream_link.data
width = RasterUtilClass.read_raster(width_file)
width_data = width.data
depth = RasterUtilClass.read_raster(depth_file)
depth_data = depth.data
ch_width_dic = dict()
ch_depth_dic = dict()
ch_num_dic = dict()
for i in range(n_rows):
for j in range(n_cols):
if abs(data_stream[i][j] - nodata_value) <= UTIL_ZERO:
continue
tmpid = int(data_stream[i][j])
ch_num_dic.setdefault(tmpid, 0)
ch_width_dic.setdefault(tmpid, 0)
ch_depth_dic.setdefault(tmpid, 0)
ch_num_dic[tmpid] += 1
ch_width_dic[tmpid] += width_data[i][j]
ch_depth_dic[tmpid] += depth_data[i][j]
for k in ch_num_dic:
ch_width_dic[k] /= ch_num_dic[k]
ch_depth_dic[k] /= ch_num_dic[k]
# add channel width and depth fields to reach shp file or update values if the fields exist
ds_reach = ogr_Open(reach_shp_file, update=True)
layer_reach = ds_reach.GetLayer(0)
layer_def = layer_reach.GetLayerDefn()
i_link = layer_def.GetFieldIndex(ImportReaches2Mongo._LINKNO)
i_width = layer_def.GetFieldIndex(ImportReaches2Mongo._WIDTH)
i_depth = layer_def.GetFieldIndex(ImportReaches2Mongo._DEPTH)
if i_width < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._WIDTH, OFTReal)
layer_reach.CreateField(new_field)
if i_depth < 0:
new_field = ogr_FieldDefn(ImportReaches2Mongo._DEPTH, OFTReal)
layer_reach.CreateField(new_field)
layer_reach.ResetReading()
ft = layer_reach.GetNextFeature()
while ft is not None:
tmpid = ft.GetFieldAsInteger(i_link)
w = 5.
d = 1.5
if tmpid in ch_width_dic:
w = ch_width_dic[tmpid]
if tmpid in ch_depth_dic:
d = ch_depth_dic[tmpid]
ft.SetField(ImportReaches2Mongo._WIDTH, w)
ft.SetField(ImportReaches2Mongo._DEPTH, d)
layer_reach.SetFeature(ft)
ft = layer_reach.GetNextFeature()
layer_reach.SyncToDisk()
ds_reach.Destroy()
del ds_reach
