Code fragments of adapterhipims.py

• import csv
• import shutil
• from pypims import IO, flood
• import pandas as pd
• from ModelAdapter import ModelAdapter
• import os
• from util import *
• import matplotlib.pyplot as plt
• import matplotlib.colors as mcolors
• from matplotlib.colors import ListedColormap
• from mpl_toolkits.axes_grid1 import make_axes_locatable
• from matplotlib.patches import Patch
• import numpy as np
• import matplotlib.pyplot as plt
• import rasterio
• from rasterio.plot import show
• import contextily as ctx
• from matplotlib.colors import ListedColormap, Normalize
• from scipy.ndimage import binary_dilation
• from PIL import Image
• from rasterio.warp import reproject, Resampling
• class HiPIMSAdapter(ModelAdapter):
• def __init__(self, settings,envs, n_gpus):
• super().__init__(settings,envs)
• self.n_gpus = n_gpus
• self.rain_cloud_path = "Flood_Smoke_Test/rainfall_data_gan_smoke_10.tif"
• self.time_resolution = settings["time_resolution"]
• self.num_rain = 0
• self.rain_processed_path = "/Model/rain_processed.tif"
• def visualize_and_save_geotiff(self,input_path, output_path): # Read the geotiff
• with rasterio.open(input_path) as src:
• band = src.read(1)
• # Mask values outside [0.1, 4] range
• band_masked = np.ma.masked_where((band < 0.1) | (band > 4), band)
• # Create a custom colormap
• base_cmap = plt.cm.turbo
• colors = base_cmap(np.linspace(0, 1, 256))
• # Set the first color to white transparent
• colors[0] = [1, 1, 1, 0]
• custom_cmap = ListedColormap(colors)
• norm = mcolors.Normalize(vmin=0.1, vmax=4, clip=True)
• fig, ax = plt.subplots(figsize=(10, 10))
• cax = ax.imshow(band_masked, cmap=custom_cmap, norm=norm)
• # Add a colorbar
• cbar = fig.colorbar(cax, ax=ax, orientation='vertical')
• # Save the figure
• plt.savefig(output_path, dpi=300, bbox_inches='tight')
• plt.close(fig)
• #path1 = "D:/flood_out (1).tif"
• def add_grayscale_basemap(self,ax, zoom, source=ctx.providers.Esri.WorldImagery):
• xmin, xmax, ymin, ymax = ax.axis()
• basemap, extent = ctx.bounds2img(xmin, ymin, xmax, ymax, zoom=zoom, source=source)
• # Convert to grayscale
• gray_basemap = Image.fromarray(basemap).convert("L")
• rgb_gray_basemap = gray_basemap.convert("RGB")
• ax.imshow(rgb_gray_basemap, extent=extent, interpolation='bilinear')
• ax.axis((xmin, xmax, ymin, ymax))
• def plot_tif_on_osm(self,tif_paths, unified_tif_path,output_path="/flood_visualisation.png"):
• color1 = (120 / 255, 130 / 255, 252 / 255, 1)
• color3 = (130 / 255, 252 / 255, 130 / 255, 1)
• color2 = (255 / 255, 180 / 255, 120 / 255, 1)
• color4 = (255 / 255, 242 / 255, 120 / 255, 1)
• color5 = (248 / 255, 120 / 255, 120 / 255, 1)
• color6 = (225 / 255, 120 / 255, 231 / 255, 1)
• colors = [color1, color2, color3, color4, color5, color6]
• fig, ax = plt.subplots(1, figsize=(10, 10))
• divider = make_axes_locatable(ax)
• cax = divider.append_axes("right", size="5%", pad=0.1)
• # Combining Bounds
• combined_bounds = None
• tif_crs = None
• for tif_path in tif_paths:
• with rasterio.open(tif_path) as src:
• if combined_bounds is None:
• combined_bounds = src.bounds
• tif_crs = src.crs.to_string()
• else:
• bounds = src.bounds
• combined_bounds = rasterio.coords.BoundingBox( min(combined_bounds.left, bounds.left), min(combined_bounds.bottom, bounds.bottom), max(combined_bounds.right, bounds.right), max(combined_bounds.top, bounds.top) )
• ax.set_xlim(combined_bounds[0], combined_bounds[2])
• ax.set_ylim(combined_bounds[1], combined_bounds[3])
• self.add_grayscale_basemap(ax, zoom=13)
• # Displaying Unified Flood Map
• with rasterio.open(unified_tif_path) as src:
• flood_image = src.read(2)
• mask_for_white = (flood_image == -9999)
• flood_image[flood_image < 0.05] = np.nan
• flood_image[flood_image == -9999] = np.nan
• cmap = plt.get_cmap("viridis")
• norm = Normalize(vmin=0.05, vmax=2, clip=True)
• img = ax.imshow(flood_image, cmap=cmap, norm=norm, alpha=1, extent=(combined_bounds[0], combined_bounds[2], combined_bounds[1], combined_bounds[3]))
• cbar = fig.colorbar(img, cax=cax, orientation="vertical")
• cbar.set_label('Water Depth [m]', rotation=270, labelpad=20)
• white_colormap = ListedColormap(['none', 'white'])
• show(mask_for_white, transform=src.transform, ax=ax, cmap=white_colormap)
• # Displaying Shaded Municipalities
• legend_patches = []
• for tif_path, color in zip(tif_paths, colors):
• with rasterio.open(tif_path) as src:
• tif_image = src.read(1)
• tif_image[tif_image == -9999] = np.nan
• mask_colormap = ListedColormap([color])
• show(tif_image, transform=src.transform, ax=ax, cmap=mask_colormap, alpha=.15)
• legend_patches.append(Patch(color=color, label=tif_path.split("_")[-1].replace(".tif", ""), alpha=0.5))
• ax.legend(handles=legend_patches, loc='upper left', title="Municipalities", facecolor='white', framealpha=1)
• ax.axis('off')
• plt.savefig(output_path, dpi=300, bbox_inches='tight')
• plt.close(fig)
• def handle_rainfall_input(self): # Read the data from the GeoTIFF file
• with rasterio.open(self.rain_tif_path) as src:
• data = src.read()
• resample_tif(self.rain_tif_path, self.dem_path, self.rain_processed_path)
• num_data = data.shape[0]
• self.num_rain = num_data
• # Write the data to a CSV file which can be read by the model
• for n in range(num_data):
• timesteps = self.rain_time // self.time_resolution
• path = f"/input/{n}/rainfall_data.csv"
• os.makedirs(os.path.dirname(path), exist_ok=True)
• with open(path, 'w', newline='') as f:
• writer = csv.writer(f)
• for i in range(timesteps):
• flattened_data = ((data[n]) / 3600000.).flatten()
• row = [int(i * self.time_resolution), ] + flattened_data.tolist()
• writer.writerow(row)
• row = [int(i * self.time_resolution + self.time_resolution - 1), ] + flattened_data.tolist()
• writer.writerow(row)
• for i in range(timesteps, self.run_duration // self.time_resolution):
• writer.writerow( [int(i * self.time_resolution), ] + [0 for _ in range(data.shape[1] * data.shape[2])])
• # mask the rainfall data
• mask_data = np.indices((data.shape[1], data.shape[2])).reshape(2, -1).T
• mask_data = np.sum(mask_data * np.array([data.shape[2], 1]), axis=1).reshape(data.shape[1], data.shape[2])
• # Get the geographic metadata from the original GeoTIFF
• with rasterio.open(self.rain_tif_path) as src:
• crs = src.crs
• transform = src.transform
• # Write the mask data to a new GeoTIFF file
• with rasterio.open( '/rainfall_mask_all.tif', 'w', driver='GTiff', height=mask_data.shape[0], width=mask_data.shape[1], count=1, dtype=mask_data.dtype, crs=crs, transform=transform, ) as dst:
• dst.write(mask_data, 1)
• resample_tif('/rainfall_mask_all.tif', self.dem_path,"/rainfall_mask.tif")
• def run(self):
• ''' Run the hydraulic model with the given rainfall input and local DEM :return: '''
• for i in range(self.num_rain):
• DEM = IO.Raster(self.dem_path)
• rain_mask = IO.Raster("/rainfall_mask.tif")
• rain_source = pd.read_csv(f"/input/{i}/rainfall_data.csv", header=None)
• case_input = IO.InputHipims(DEM, num_of_sections=self.n_gpus, case_folder=f"/Model/output/{i}/")
• rain_source_np = rain_source.to_numpy()
• case_input.set_rainfall(rain_mask=rain_mask, rain_source=rain_source_np)
• case_input.set_runtime([1, self.run_duration, self.run_duration, self.run_duration+1])
• case_input.write_input_files()
• if self.n_gpus > 1:
• flood.run_mgpus(f"/Model/output/{i}/")
• else:
• flood.run(f"/Model/output/{i}/")
• def handle_output(self): # Read in the input GeoTIFF
• with rasterio.open(self.dem_path) as src:
• profile = src.profile
• # Get the old resolution and calculate the new resolution
• old_res = profile['transform'][0]
• new_res = self.resolution
• # Get the old width, height, and transform
• old_width = profile['width']
• old_height = profile['height']
• old_transform = profile['transform']
• # Calculate new width and height
• new_width = int(old_width * (old_res / new_res))
• new_height = int(old_height * (old_res / new_res))
• new_transform = rasterio.Affine(new_res, old_transform.b, old_transform.c, old_transform.d, -new_res, old_transform.f)
• # Update the width, height, and transform in the profile
• profile['width'] = int(new_width)
• profile['height'] = int(new_height)
• profile['transform'] = new_transform
• profile["count"] = self.num_rain
• flood_list = []
• rain_list = []
• # read each output
• for t in range(self.num_rain):
• flood = np.loadtxt(f'/Model/output/{t}/output/h_max_{self.run_duration}.asc', skiprows=6).astype(np.float16)
• # Create xnew and ynew indices with new resolution
• xnew = np.linspace(0, old_width, new_width)
• ynew = np.linspace(0, old_height, new_height)
• flood = interpolate_array(flood, xnew, ynew, new_res)
• flood_list.append(flood)
• with rasterio.open(self.rain_processed_path) as src:
• rain = src.read(t + 1)
• rain = interpolate_array(rain, xnew, ynew, new_res)
• rain_list.append(rain)
• dem = rasterio.open(self.dem_path).read(1)
• dem = interpolate_array(dem, xnew, ynew, new_res)[np.newaxis, :, :]
• # Round the data to two decimal places and set negative values to -9999
• flood = np.around(np.stack(flood_list), 2)
• rain = np.around(np.stack(rain_list), 2)
• flood[flood < 0] = -9999
• rain[rain < 0] = -9999
• dem[dem < 0] = -9999
• # Round the data to two decimal places and set negative values to -9999
• flood = np.around(np.stack(flood_list), 2)
• flood[flood < 0] = -9999
• # Modify the profile to apply compression (LZW)
• profile.update(compress='lzw')
• # remove outpu
• os.makedirs("/Stationdata", exist_ok=True)
• shutil.rmtree(os.path.join("/Model/output"))
• shutil.rmtree(os.path.join("/input"))
• # write the flood output
• with rasterio.open(f"/Stationdata/flood_out_{self.envs['STATION_NAME']}.tif", 'w', **profile) as dst:
• dst.write(flood)
• if not os.path.exists("/Model/flood_out.tif"):
• with rasterio.open("/Model/flood_out.tif", 'w', **profile) as dst:
• dst.write(flood)
• with rasterio.open("/Model/rain_out.tif", 'w', **profile) as dst:
• dst.write(rain)
• profile_dem = profile
• profile_dem["count"] = 1
• with rasterio.open("/Model/dem_out.tif", 'w', **profile_dem) as dst:
• dst.write(dem)
• else:
• with rasterio.open("/Model/flood_out_.tif", 'w', **profile) as dst:
• dst.write(flood)
• with rasterio.open("/Model/rain_out_.tif", 'w', **profile) as dst:
• dst.write(rain)
• profile_dem = profile
• profile_dem["count"] = 1
• with rasterio.open("/Model/dem_out_.tif", 'w', **profile_dem) as dst:
• dst.write(dem)
• merge_tifs(["/Model/flood_out_.tif", "/Model/flood_out.tif"], "/Model/flood_out.tif")
• merge_tifs(["/Model/rain_out_.tif", "/Model/rain_out.tif"], "/Model/rain_out.tif")
• merge_tifs(["/Model/dem_out_.tif", "/Model/dem_out.tif"], "/Model/dem_out.tif")
• self.visualize_and_save_geotiff("/Model/flood_out.tif", "/Model/flood_out.png")
• self.plot_tif_on_osm(["/Stationdata/"+f for f in os.listdir("/Stationdata")], "/Model/flood_out.tif")

Graph