diff --git a/bluemath_tk/additive/greensurge.py b/bluemath_tk/additive/greensurge.py index 2dc26af..160aa07 100644 --- a/bluemath_tk/additive/greensurge.py +++ b/bluemath_tk/additive/greensurge.py @@ -1,5 +1,7 @@ import datetime as datetime import warnings +from functools import partial +from multiprocessing import Pool, cpu_count from typing import Tuple import cartopy.crs as ccrs @@ -616,9 +618,13 @@ def GS_LinearWindDragCoef_mat( return CD.item() if was_scalar else CD -def GS_LinearWindDragCoef(Wspeed: np.ndarray,CD_Wl_abc: np.ndarray,Wl_abc: np.ndarray)-> np.ndarray: + +def GS_LinearWindDragCoef( + Wspeed: np.ndarray, CD_Wl_abc: np.ndarray, Wl_abc: np.ndarray +) -> np.ndarray: """ Calculate the linear drag coefficient based on wind speed and specified thresholds. + Parameters ---------- Wspeed : np.ndarray @@ -627,46 +633,43 @@ def GS_LinearWindDragCoef(Wspeed: np.ndarray,CD_Wl_abc: np.ndarray,Wl_abc: np.nd Coefficients for the drag coefficient calculation, should be a 1D array of length 3. Wl_abc : np.ndarray Wind speed thresholds for the drag coefficient calculation, should be a 1D array of length 3. + Returns ------- np.ndarray Calculated drag coefficient values based on the input wind speed. """ - - Wla=Wl_abc[0] - Wlb=Wl_abc[1] - Wlc=Wl_abc[2] - CDa=CD_Wl_abc[0] - CDb=CD_Wl_abc[1] - CDc=CD_Wl_abc[2] - - - #coefs lines y=ax+b - if not Wla==Wlb: - a_CDline_ab=(CDa-CDb)/(Wla-Wlb) - b_CDline_ab=CDb-a_CDline_ab*Wlb + + Wla = Wl_abc[0] + Wlb = Wl_abc[1] + Wlc = Wl_abc[2] + CDa = CD_Wl_abc[0] + CDb = CD_Wl_abc[1] + CDc = CD_Wl_abc[2] + + # coefs lines y=ax+b + if not Wla == Wlb: + a_CDline_ab = (CDa - CDb) / (Wla - Wlb) + b_CDline_ab = CDb - a_CDline_ab * Wlb else: - a_CDline_ab=0 - b_CDline_ab=CDa - if not Wlb==Wlc: - a_CDline_bc=(CDb-CDc)/(Wlb-Wlc) - b_CDline_bc=CDc-a_CDline_bc*Wlc + a_CDline_ab = 0 + b_CDline_ab = CDa + if not Wlb == Wlc: + a_CDline_bc = (CDb - CDc) / (Wlb - Wlc) + b_CDline_bc = CDc - a_CDline_bc * Wlc else: - a_CDline_bc=0 - b_CDline_bc=CDb - a_CDline_cinf=0 - b_CDline_cinf=CDc - - - - if Wspeed<=Wlb: - CD=a_CDline_ab*Wspeed+b_CDline_ab - elif (Wspeed>Wlb and Wspeed<=Wlc): - CD=a_CDline_bc*Wspeed+b_CDline_bc + a_CDline_bc = 0 + b_CDline_bc = CDb + a_CDline_cinf = 0 + b_CDline_cinf = CDc + + if Wspeed <= Wlb: + CD = a_CDline_ab * Wspeed + b_CDline_ab + elif Wspeed > Wlb and Wspeed <= Wlc: + CD = a_CDline_bc * Wspeed + b_CDline_bc else: - CD=a_CDline_cinf*Wspeed+b_CDline_cinf - - + CD = a_CDline_cinf * Wspeed + b_CDline_cinf + return CD @@ -834,13 +837,15 @@ def plot_GS_TG_validation_timeseries( vmax = np.nanmax(Z) cmap, norm = join_colormaps( - cmap1= hex_colors_water, - cmap2= hex_colors_land, - value_range1=(vmin, 0.0), value_range2=(0.0, vmax), - name="raster_cmap") - ax_map.set_facecolor("#518134") + cmap1=hex_colors_water, + cmap2=hex_colors_land, + value_range1=(vmin, 0.0), + value_range2=(0.0, vmax), + name="raster_cmap", + ) + ax_map.set_facecolor("#518134") - pm = ax_map.tripcolor( + ax_map.tripcolor( X, Y, triangles, @@ -872,8 +877,24 @@ def plot_GS_TG_validation_timeseries( ax_ts.plot(time_vals, WL_IB, c="black", label="Inverse Barometer") if WLmin is None or WLmax is None: - WLmax = max(np.nanmax(WL_SS_dyn), np.nanmax(WL_SS_GS), np.nanmax(WL_TG)) - WLmin = min(np.nanmin(WL_SS_dyn), np.nanmin(WL_SS_GS), np.nanmin(WL_TG)) + WLmax = ( + max( + np.nanmax(WL_SS_dyn), + np.nanmax(WL_SS_GS), + np.nanmax(WL_TG), + np.nanmax(WL_GS), + ) + * 1.05 + ) + WLmin = ( + min( + np.nanmin(WL_SS_dyn), + np.nanmin(WL_SS_GS), + np.nanmin(WL_TG), + np.nanmin(WL_GS), + ) + * 1.05 + ) ax_ts.set_xlim(time_vals[0], time_vals[-1]) ax_ts.set_ylim(WLmin, WLmax) @@ -899,6 +920,7 @@ def extract_pos_nearest_points_tri( Array of longitudes for which to find the nearest triangle index. lat_points : np.ndarray Array of latitudes for which to find the nearest triangle index. + Returns ------- np.ndarray @@ -906,8 +928,13 @@ def extract_pos_nearest_points_tri( """ if "node_forcing_latitude" in ds_mesh_info.variables: - lon_mesh = ds_mesh_info.node_computation_longitude.values - lat_mesh = ds_mesh_info.node_computation_latitude.values + elements = ds_mesh_info.triangle_computation_connectivity.values + lon_mesh = np.mean( + ds_mesh_info.node_computation_longitude.values[elements], axis=1 + ) + lat_mesh = np.mean( + ds_mesh_info.node_computation_latitude.values[elements], axis=1 + ) type_ds = 0 else: lon_mesh = ds_mesh_info.mesh2d_face_x.values @@ -995,3 +1022,180 @@ def pressure_to_IB(xds_presure: xr.Dataset) -> xr.Dataset: xds_presure_modified["IB"] = (("lat", "lon", "time"), IB) return xds_presure_modified + + +def compute_water_level_for_time( + time_index: int, + direction_data: np.ndarray, + wind_speed_data: np.ndarray, + direction_bins: np.ndarray, + forcing_cell_indices: np.ndarray, + greensurge_dataset: xr.Dataset, + wind_speed_reference: float, + base_drag_coeff: float, + drag_coefficients: np.ndarray, + velocity_thresholds: np.ndarray, + duration_in_steps: int, + num_output_times: int, +) -> np.ndarray: + """ + Compute the water level for a specific time index based on wind direction and speed. + + Parameters + ---------- + time_index : int + The index of the time step to compute the water level for. + direction_data : np.ndarray + 2D array of wind direction data with shape (n_cells, n_times). + wind_speed_data : np.ndarray + 2D array of wind speed data with shape (n_cells, n_times). + direction_bins : np.ndarray + 1D array of wind direction bins. + forcing_cell_indices : np.ndarray + 1D array of indices for the forcing cells. + greensurge_dataset : xr.Dataset + xarray Dataset containing the GreenSurge mesh and forcing data. + wind_speed_reference : float + Reference wind speed value for scaling. + base_drag_coeff : float + Base drag coefficient value for scaling. + drag_coefficients : np.ndarray + 1D array of drag coefficients corresponding to the velocity thresholds. + velocity_thresholds : np.ndarray + 1D array of velocity thresholds for drag coefficient calculation. + duration_in_steps : int + Total duration of the simulation in steps. + num_output_times : int + Total number of output time steps. + + Returns + ------- + np.ndarray + 2D array of computed water levels for the specified time index. + """ + + adjusted_bins = np.where(direction_bins == 0, 360, direction_bins) + n_faces = greensurge_dataset["mesh2d_s1"].isel(forcing_cell=0, direction=0).shape + water_level_accumulator = np.zeros(n_faces) + + for cell_index in forcing_cell_indices.astype(int): + current_dir = direction_data[cell_index, time_index] % 360 + closest_direction_index = np.abs(adjusted_bins - current_dir).argmin() + + water_level_case = ( + greensurge_dataset["mesh2d_s1"] + .sel(forcing_cell=cell_index, direction=closest_direction_index) + .values + ) + water_level_case = np.nan_to_num(water_level_case, nan=0) + + wind_speed_value = wind_speed_data[cell_index, time_index] + drag_coeff_value = GS_LinearWindDragCoef( + wind_speed_value, drag_coefficients, velocity_thresholds + ) + + scaling_factor = (wind_speed_value**2 / wind_speed_reference**2) * ( + drag_coeff_value / base_drag_coeff + ) + water_level_accumulator += water_level_case * scaling_factor + + step_window = min(duration_in_steps, num_output_times - time_index) + result = np.zeros((num_output_times, n_faces[1])) + if (num_output_times - time_index) > step_window: + result[time_index : time_index + step_window] += water_level_accumulator + else: + shift_counter = step_window - (num_output_times - time_index) + result[time_index : time_index + step_window - shift_counter] += ( + water_level_accumulator[: step_window - shift_counter] + ) + return result + + +def GS_windsetup_reconstruction_with_postprocess_parallel( + greensurge_dataset: xr.Dataset, + ds_gfd_metadata: xr.Dataset, + wind_direction_input: xr.Dataset, + num_workers: int = None, + velocity_thresholds: np.ndarray = np.array([0, 100, 100]), + drag_coefficients: np.ndarray = np.array([0.00063, 0.00723, 0.00723]), +) -> xr.Dataset: + """ + Reconstructs the GreenSurge wind setup using the provided wind direction input and metadata in parallel. + + Parameters + ---------- + greensurge_dataset : xr.Dataset + xarray Dataset containing the GreenSurge mesh and forcing data. + ds_gfd_metadata: xr.Dataset + xarray Dataset containing metadata for the GFD mesh. + wind_direction_input: xr.Dataset + xarray Dataset containing wind direction and speed data. + velocity_thresholds : np.ndarray + Array of velocity thresholds for drag coefficient calculation. + drag_coefficients : np.ndarray + Array of drag coefficients corresponding to the velocity thresholds. + + Returns + ------- + xr.Dataset + xarray Dataset containing the reconstructed wind setup. + """ + + if num_workers is None: + num_workers = cpu_count() + + direction_bins = ds_gfd_metadata.wind_directions.values + forcing_cell_indices = greensurge_dataset.forcing_cell.values + wind_speed_reference = ds_gfd_metadata.wind_speed.values.item() + base_drag_coeff = GS_LinearWindDragCoef( + wind_speed_reference, drag_coefficients, velocity_thresholds + ) + time_step_hours = ds_gfd_metadata.time_step_hours.values + + time_start = wind_direction_input.time.values.min() + time_end = wind_direction_input.time.values.max() + duration_in_steps = ( + int((ds_gfd_metadata.simulation_duration_hours.values) / time_step_hours) + 1 + ) + output_time_vector = np.arange( + time_start, time_end, np.timedelta64(int(60 * time_step_hours.item()), "m") + ) + num_output_times = len(output_time_vector) + + direction_data = wind_direction_input.Dir.values + wind_speed_data = wind_direction_input.W.values + + n_faces = greensurge_dataset["mesh2d_s1"].isel(forcing_cell=0, direction=0).shape[1] + + args = partial( + compute_water_level_for_time, + direction_data=direction_data, + wind_speed_data=wind_speed_data, + direction_bins=direction_bins, + forcing_cell_indices=forcing_cell_indices, + greensurge_dataset=greensurge_dataset, + wind_speed_reference=wind_speed_reference, + base_drag_coeff=base_drag_coeff, + drag_coefficients=drag_coefficients, + velocity_thresholds=velocity_thresholds, + duration_in_steps=duration_in_steps, + num_output_times=num_output_times, + ) + + with Pool(processes=num_workers) as pool: + results = list( + tqdm(pool.imap(args, range(num_output_times)), total=num_output_times) + ) + + wind_setup_output = np.sum(results, axis=0) + + ds_wind_setup = xr.Dataset( + {"WL": (["time", "nface"], wind_setup_output)}, + coords={ + "time": output_time_vector, + "nface": np.arange(n_faces), + }, + ) + ds_wind_setup.attrs["description"] = "Wind setup from GreenSurge methodology" + + return ds_wind_setup diff --git a/bluemath_tk/topo_bathy/mesh_utils.py b/bluemath_tk/topo_bathy/mesh_utils.py index 7db8049..0932be4 100644 --- a/bluemath_tk/topo_bathy/mesh_utils.py +++ b/bluemath_tk/topo_bathy/mesh_utils.py @@ -12,7 +12,6 @@ import rasterio from jigsawpy.msh_t import jigsaw_msh_t from matplotlib.axes import Axes -from matplotlib.tri import Triangulation from netCDF4 import Dataset from pyproj.enums import TransformDirection from rasterio.mask import mask @@ -20,7 +19,6 @@ from shapely.ops import transform from ..core.geo import buffer_area_for_polygon - from ..core.plotting.colors import hex_colors_land, hex_colors_water from ..core.plotting.utils import join_colormaps @@ -160,10 +158,12 @@ def plot_bathymetry(rasters_path: List[str], polygon: Polygon, ax: Axes) -> Axes vmax = np.nanmax(data[0]) cmap, norm = join_colormaps( - cmap1= hex_colors_water, - cmap2= hex_colors_land, - value_range1=(vmin, 0.0), value_range2=(0.0, vmax), - name="raster_cmap") + cmap1=hex_colors_water, + cmap2=hex_colors_land, + value_range1=(vmin, 0.0), + value_range2=(0.0, vmax), + name="raster_cmap", + ) im = ax.imshow( data[0], @@ -357,7 +357,7 @@ def plot_bathymetry_interp(mesh: jigsaw_msh_t, to_geo, ax: Axes) -> None: bnd = [crd[:, 0].min(), crd[:, 0].max(), crd[:, 1].min(), crd[:, 1].max()] triangle = mesh.msh_t.tria3["index"] - Z = np.mean(mesh.msh_t.value.flatten()[triangle],axis = 1) + Z = np.mean(mesh.msh_t.value.flatten()[triangle], axis=1) vmin = np.nanmin(Z) vmax = np.nanmax(Z) @@ -366,8 +366,8 @@ def plot_bathymetry_interp(mesh: jigsaw_msh_t, to_geo, ax: Axes) -> None: cmap2=hex_colors_land, value_range1=(vmin, 0.0), value_range2=(0.0, vmax), - name="raster_cmap" -) + name="raster_cmap", + ) im = ax.tripcolor( crd[:, 0], @@ -404,8 +404,9 @@ def simply_polygon(base_shape: Polygon, simpl_UTM: float, project) -> Polygon: ------- Polygon The simplified polygon in geographic coordinates. - Example - ------- + + Examples + -------- >>> from shapely.geometry import Polygon >>> from pyproj import Transformer >>> from shapely.ops import transform @@ -448,8 +449,9 @@ def remove_islands(base_shape: Polygon, threshold_area: float, project) -> Polyg ------- Polygon The polygon with small interior rings removed, transformed back to geographic coordinates. - Example - ------- + + Examples + -------- >>> from shapely.geometry import Polygon >>> from pyproj import Transformer >>> from shapely.ops import transform @@ -492,8 +494,9 @@ def read_adcirc_grd(grd_file: str) -> Tuple[np.ndarray, np.ndarray, List[str]]: - Elmts (np.ndarray): An array of shape (nelmts, 3) containing the element connectivity, with node indices adjusted (decremented by 1). - lines (List[str]): The remaining lines in the file after reading the nodes and elements. - Example - ------- + + Examples + -------- >>> nodes, elmts, lines = read_adcirc_grd("path/to/grid.grd") >>> print(nodes.shape, elmts.shape, len(lines)) (1000, 3) (500, 3) 10 @@ -527,8 +530,9 @@ def calculate_edges(Elmts: np.ndarray) -> np.ndarray: np.ndarray A 2D array of shape (n_edges, 2) containing the unique edges, each represented by a pair of node indices. - Example - ------- + + Examples + -------- >>> Elmts = np.array([[0, 1, 2], [1, 2, 3], [2, 0, 3]]) >>> edges = calculate_edges(Elmts) >>> print(edges) @@ -569,7 +573,10 @@ def adcirc2netcdf(Path_grd: str, netcdf_path: str) -> None: netcdf_path : str Path where the resulting NetCDF file will be saved. - TODO: Check the whole function for correctness and completeness. + Examples + -------- + >>> adcirc2netcdf("path/to/grid.grd", "path/to/output.nc") + >>> print("NetCDF file created successfully.") """ Nodes_full, Elmts_full, lines = read_adcirc_grd(Path_grd) @@ -817,18 +824,19 @@ def decode_open_boundary_data(data: List[str]) -> dict: dict A dictionary with keys corresponding to open boundary identifiers (e.g., 'open_boundary_1') and values as lists of integers representing boundary node indices. - Example - ------- + + Examples + -------- >>> data = [ - "100! 200! 300!", - "open_boundary_1", - "open_boundary_2", - "open_boundary_3", - "land boundaries", - "open_boundary_1! 10", - "open_boundary_2! 20", - "open_boundary_3! 30", - ] + ... "100! 200! 300!", + ... "open_boundary_1", + ... "open_boundary_2", + ... "open_boundary_3", + ... "land boundaries", + ... "open_boundary_1! 10", + ... "open_boundary_2! 20", + ... "open_boundary_3! 30", + ... ] >>> boundaries = decode_open_boundary_data(data) >>> print(boundaries) {'open_boundary_1': [10], 'open_boundary_2': [20], 'open_boundary_3': [30]} @@ -869,10 +877,11 @@ def compute_circumcenter(p0: np.ndarray, p1: np.ndarray, p2: np.ndarray) -> np.n Returns ------- - center : np.ndarray + np.ndarray 2D coordinates of the circumcenter. - Example - ------- + + Examples + -------- >>> p0 = np.array([0, 0]) >>> p1 = np.array([1, 0]) >>> p2 = np.array([0, 1]) @@ -912,8 +921,9 @@ def build_edge_to_cells(elements: np.ndarray) -> Dict[Tuple[int, int], List[int] ------- edge_to_cells : Dict[Tuple[int, int], List[int]] Dictionary mapping edges to the list of adjacent element indices. - Example - ------- + + Examples + -------- >>> elements = np.array([[0, 1, 2], [1, 2, 3], [2, 0, 3]]) >>> edge_to_cells = build_edge_to_cells(elements) >>> print(edge_to_cells) @@ -946,15 +956,16 @@ def detect_circumcenter_too_close( 1D arrays of x and y coordinates of the nodes. elements : np.ndarray 2D array of shape (nelmts, 3) containing the node indices for each triangle element. - aj_threshold : float - Threshold for the ratio of circumcenter distance to edge length. + aj_threshold : float, optional + Threshold for the ratio of circumcenter distance to edge length. Default is 0.1. Returns ------- bad_elements_mask : np.ndarray Boolean mask indicating which elements are problematic (True if bad). - Example - ------- + + Examples + -------- >>> X = np.array([0, 1, 0, 1]) >>> Y = np.array([0, 0, 1, 1]) >>> elements = np.array([[0, 1, 2], [1, 3, 2]]) @@ -1000,7 +1011,7 @@ def detect_circumcenter_too_close( def define_mesh_target_size( rasters: List[rasterio.io.DatasetReader], raster_resolution_meters: float, - depth_ranges: float, + depth_ranges: dict, nprocs: int = 1, ) -> ocsmesh.Hfun: """ @@ -1008,11 +1019,15 @@ def define_mesh_target_size( Parameters ---------- - rasters : list + rasters : List[rasterio.io.DatasetReader] List of raster objects. + raster_resolution_meters : float + Resolution of the raster in meters. depth_ranges : dict, optional Dictionary containing depth ranges and their corresponding mesh sizes and rates. Format: {(lower_bound, upper_bound): {'value': mesh_size, 'rate': expansion_rate}} + nprocs : int, optional + Number of processors to use for the mesh generation. Default is 1. Returns -------