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
• 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)
• 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 output
• shutil.rmtree(os.path.join("/Model/output"))
• shutil.rmtree(os.path.join("/input"))
• # write the flood output
• 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")

Graph