def read_velocity_files(self, time_tag):
"""
read h, hUx and hUy files, calculate velocity and return a grid obj
Parameters
----------
time_tag : string
time string in seconds, e.g. '7200'
Returns
-------
velocity_ras : Raster object
velocity values on the model grid.
list: list of 2D array
velocity values in x and y direction.
"""
file_tag_x = 'hUx_' + time_tag
file_tag_y = 'hUy_' + time_tag
file_tag_h = 'h_' + time_tag
grid_x = self.read_grid_file(file_tag=file_tag_x)
grid_y = self.read_grid_file(file_tag=file_tag_y)
grid_h = self.read_grid_file(file_tag=file_tag_h)
array_h = grid_h.array + 1e-6
array_u = grid_x.array / array_h
array_v = grid_y.array / array_h
array_vel = np.sqrt(array_u**2 + array_v**2)
velocity_ras = Raster(array=array_vel, header=self.header)
if hasattr(self, 'crs'):
velocity_ras.set_crs(self.crs)
return velocity_ras, [array_u, array_v]
def read_grid_file(self, file_tag='h_0', compressed=False):
"""Read asc grid files from output
Return:
Raster: a raster object of the grid
"""
if not file_tag.endswith('.asc'):
file_tag = file_tag + '.asc'
if compressed:
file_tag = file_tag + '.gz'
if self.num_of_sections == 1:
file_name = os.path.join(self.output_folder, file_tag)
print('reading file ' + file_name)
if os.path.isfile(file_name):
grid_array, _, _ = sp.arcgridread(file_name)
else:
warnings.warn(file_name + ' is not found!')
grid_array = None
else: # multi-GPU
print('reading file ' + file_tag)
grid_array = self._combine_multi_gpu_grid_data(file_tag)
if grid_array is None:
warnings.warn(file_tag + ' is not read successfully!')
grid_obj = None
else:
grid_obj = Raster(array=grid_array, header=self.header)
if hasattr(self, 'crs'):
grid_obj.set_crs(self.crs)
return grid_obj
def set_mask(self, rain_mask, dem_ras=None):
"""Set rainfall mask from a scalar or a grid (object/file)
if rain_mask is a scalar or array, dem_ras must be provided
"""
if type(rain_mask) is str:
obj_rain_mask = Raster(rain_mask)
print(rain_mask + ' read')
elif hasattr(rain_mask, 'header'):
obj_rain_mask = rain_mask
else: # numpy array or a scalar, dem_ras must be provided
rain_mask = np.array(rain_mask).astype('int32')
dem_shape = dem_ras.array.shape
if rain_mask.size > 1:
if rain_mask.shape != dem_shape:
raise ValueError('The shape of rainfall_mask array '
'is not consistent with DEM')
mask_array = np.zeros(dem_shape, dtype=np.int32) + rain_mask
obj_rain_mask = Raster(array=mask_array, header=dem_ras.header)
if hasattr(dem_ras, 'header'):
# rain mask resample to the same shape with DEM
self.mask_dict = indep_f._mask2dict(obj_rain_mask, dem_ras.header)
self.mask_header = dem_ras.header
self.subs_in = np.where(~np.isnan(dem_ras.array))
else:
self.mask_dict = indep_f._mask2dict(obj_rain_mask)
self.mask_header = obj_rain_mask.header
def plot_rainfall_map(self, method='sum', cmap='YlGnBu', **kw):
"""plot rainfall map within model domain
method: the way to calculate time series rainfall rate
"""
rain_array, unit_str = self.get_spatial_map(method)
grid_obj = Raster(array=rain_array, header=self.mask_header)
fig, ax = grid_obj.mapshow(cax_str=unit_str, cmap=cmap, **kw)
return fig, ax
def __init__(self, ras_data, dem_ras=None):
"""
"""
if type(ras_data) is str:
obj_landcover = Raster(ras_data)
elif hasattr(ras_data, 'header'):
obj_landcover = ras_data
if hasattr(dem_ras, 'header'):
# landcover resample to the same shape with DEM
self.mask_dict = indep_f._mask2dict(obj_landcover, dem_ras.header)
self.mask_header = dem_ras.header
self.subs_in = np.where(~np.isnan(dem_ras.array))
else:
self.mask_dict = indep_f._mask2dict(obj_landcover)
self.mask_header = obj_landcover.header
def _set_grid_header(self, asc_file=None):
"""set header for the grid of the output object
num_of_sections, input_folder, output_folder are required
asc_file: the file to get grid header, default is DEM.txt
"""
num_of_sections = self.num_of_sections
output_folder = self.output_folder
input_folder = self.input_folder
warning_message = ('The grid header is not set, '
'try _set_grid_header(asc_file) to set a '
'header if it is required in your next steps')
if num_of_sections == 1:
if asc_file is None:
file_name = os.path.join(input_folder, 'mesh', 'DEM.txt')
else:
file_name = os.path.join(output_folder, asc_file)
if os.path.exists(file_name):
dem_ras = Raster(file_name)
self.header = dem_ras.header
self.dem_array = dem_ras.array
else:
warnings.warn(warning_message)
else: #multi-gpu model
headers = []
header_set = True
for i in np.arange(num_of_sections):
if asc_file is None:
file_name = os.path.join(input_folder[i], 'mesh',
'DEM.txt')
else:
file_name = os.path.join(output_folder[i], asc_file)
if os.path.exists(file_name):
header = sp.arc_header_read(file_name)
headers.append(header)
header_set = True * header_set
else:
warnings.warn(warning_message)
header_set = False
if header_set:
self.header_list = headers
self.header = _header_local2global(headers)
if asc_file is None:
self.dem_array = self._combine_multi_gpu_grid_data('DEM.txt')
def get_spatial_map(rain_source,
rain_mask_obj,
figname=None,
method='sum',
cellsize=None,
shp_file=None,
dpi=200,
title=None,
**kwargs):
"""Get spatial rainfall map
Args:
rain_mask_obj: asc file name or Raster object for rain mask
cellsize: resample the rain_mask to a new grid (with larger cellsize)
method: sum|mean caculate method for each cell, sum by time or mean by time
"""
# caculate rain source
times = rain_source[:, 0]
rain_values = rain_source[:, 1:]
cax_str = 'mm/h'
if rain_values.shape[0] == 1:
rain_array = rain_values * 1000 * 3600
rain_array = np.reshape(rain_array, rain_mask_obj.shape)
rain_array[rain_array == 0] = np.nan
rain_map_obj = Raster(array=rain_array, header=rain_mask_obj.header)
else:
rain_total = np.trapz(rain_values, x=times, axis=0) * 1000 #mm
if method == 'sum':
cell_rain = rain_total #mm
cax_str = 'mm'
elif method == 'mean':
cell_rain = rain_total / (times.max() - times.min()) * 3600 #mm/h
elif method == 'max':
cell_rain = np.max(rain_values, axis=0) * 1000 * 3600 #mm/h
elif method == 'min':
cell_rain = np.min(rain_values, axis=0) * 1000 * 3600 #mm/h
elif method == 'median':
cell_rain = np.median(rain_values, axis=0) * 1000 * 3600 #mm/h
else:
raise ValueError('Cannot recognise the calculation method')
# get spatial data
if type(rain_mask_obj) is str:
rain_mask_obj = Raster(rain_mask_obj)
mask_obj = rain_mask_obj
if cellsize is not None:
if cellsize > rain_mask_obj.header['cellsize']:
mask_obj = rain_mask_obj.grid_resample_nearest(cellsize)
ind_nan = np.isnan(mask_obj.array)
rain_mask = mask_obj.array + 0
rain_mask[ind_nan] = 0
mask_ind = rain_mask.flatten(order='F').astype('int64')
rain_vect = cell_rain[mask_ind]
rain_array = np.reshape(rain_vect, mask_obj.array.shape, order='F')
rain_array[ind_nan] = np.nan
rain_map_obj = Raster(array=rain_array, header=mask_obj.header)
fig, ax = rain_map_obj.mapshow(cax_str=cax_str, **kwargs)
xbound = ax.get_xbound()
ybound = ax.get_ybound()
# draw shape file on the rainfall map
if shp_file is not None:
sf = shapefile.Reader(shp_file)
for shape in sf.shapeRecords():
x = [i[0] for i in shape.shape.points[:]]
y = [i[1] for i in shape.shape.points[:]]
ax.plot(x, y, color='r', linewidth=1)
ax.set_xbound(xbound)
ax.set_ybound(ybound)
if title is not None:
ax.set_title(title)
if figname is not None:
fig.savefig(figname, dpi=dpi, bbox_inches='tight', pad_inches=0.02)
plt.close(fig)
return rain_map_obj