diff --git a/bluemath_tk/additive/greensurge.py b/bluemath_tk/additive/greensurge.py index 6b912f0..2dc26af 100644 --- a/bluemath_tk/additive/greensurge.py +++ b/bluemath_tk/additive/greensurge.py @@ -496,7 +496,7 @@ def GS_windsetup_reconstruction_with_postprocess( 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_mat( + base_drag_coeff = GS_LinearWindDragCoef( wind_speed_reference, drag_coefficients, velocity_thresholds ) time_step_hours = ds_gfd_metadata.time_step_hours.values @@ -533,7 +533,7 @@ def GS_windsetup_reconstruction_with_postprocess( 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_mat( + drag_coeff_value = GS_LinearWindDragCoef( wind_speed_value, drag_coefficients, velocity_thresholds ) @@ -616,6 +616,59 @@ 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: + """ + Calculate the linear drag coefficient based on wind speed and specified thresholds. + Parameters + ---------- + Wspeed : np.ndarray + Wind speed values (1D array). + CD_Wl_abc : np.ndarray + 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 + 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 + 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 + else: + CD=a_CDline_cinf*Wspeed+b_CDline_cinf + + + return CD + def plot_GS_vs_dynamic_windsetup( ds_WL_GS_WindSetUp: xr.Dataset, @@ -718,7 +771,7 @@ def plot_GS_TG_validation_timeseries( ds_WL_dynamic_WindSetUp: xr.Dataset, tide_gauge: xr.Dataset, ds_GFD_info: xr.Dataset, - figsize: tuple = (15, 7), + figsize: tuple = (20, 7), WLmin: float = None, WLmax: float = None, ) -> None: @@ -777,16 +830,27 @@ def plot_GS_TG_validation_timeseries( Y = ds_WL_dynamic_WindSetUp.mesh2d_node_y.values triangles = ds_WL_dynamic_WindSetUp.mesh2d_face_nodes.values.astype(int) - 1 Z = np.mean(ds_WL_dynamic_WindSetUp.mesh2d_node_z.values[triangles], axis=1) + vmin = np.nanmin(Z) + vmax = np.nanmax(Z) - ax_map.tripcolor( + 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") + + pm = ax_map.tripcolor( X, Y, triangles, facecolors=Z, - cmap="RdYlBu_r", + cmap=cmap, + norm=norm, shading="flat", transform=ccrs.PlateCarree(), ) + ax_map.scatter( lon_obs, lat_obs, diff --git a/bluemath_tk/core/plotting/utils.py b/bluemath_tk/core/plotting/utils.py index 65fe8e1..da8d3d7 100644 --- a/bluemath_tk/core/plotting/utils.py +++ b/bluemath_tk/core/plotting/utils.py @@ -3,8 +3,7 @@ import matplotlib.colors as colors import matplotlib.pyplot as plt import numpy as np -from matplotlib.colors import Colormap - +from matplotlib.colors import Colormap, ListedColormap, BoundaryNorm def get_list_of_colors_for_colormap( cmap: Union[str, Colormap], num_colors: int @@ -54,69 +53,65 @@ def create_cmap_from_colors( return colors.LinearSegmentedColormap.from_list(name, rgb_colors, N=256) - def join_colormaps( cmap1: Union[str, List[str], Colormap], cmap2: Union[str, List[str], Colormap], name: str = "joined_cmap", range1: Tuple[float, float] = (0.0, 1.0), range2: Tuple[float, float] = (0.0, 1.0), -) -> colors.ListedColormap: + value_range1: Tuple[float, float] = None, + value_range2: Tuple[float, float] = None, +) -> Tuple[ListedColormap, BoundaryNorm]: """ - Join two colormaps into one. Each input can be either a colormap name, a list of colors, - or an existing colormap. Optionally crop each colormap to a specific range. + Join two colormaps into one, with value ranges specified for each. Parameters ---------- - cmap1 : Union[str, List[str], Colormap] - First colormap (name, color list, or colormap object). - cmap2 : Union[str, List[str], Colormap] - Second colormap (name, color list, or colormap object). - name : str, optional - Name for the resulting colormap. Default is "joined_cmap". - range1 : Tuple[float, float], optional - Range of colors to use from first colormap (start, end) between 0 and 1. - Default is (0.0, 1.0) using the full range. - range2 : Tuple[float, float], optional - Range of colors to use from second colormap (start, end) between 0 and 1. - Default is (0.0, 1.0) using the full range. + cmap1, cmap2 : Union[str, List[str], Colormap] + Input colormaps (name, list of hex codes, or Colormap object). + name : str + Name of the output colormap. + range1, range2 : Tuple[float, float] + Portion of each colormap to use (from 0 to 1). + value_range1, value_range2 : Tuple[float, float] + Value ranges in the data domain corresponding to each colormap. Returns ------- - colors.ListedColormap - A new colormap that combines both input colormaps. - - Examples - -------- - >>> # Join two colormaps using only middle 80% of each - >>> cmap = join_colormaps("viridis", "plasma", range1=(0.1, 0.9), range2=(0.1, 0.9)) - >>> # Join colormaps with different ranges - >>> cmap = join_colormaps("viridis", "plasma", range1=(0.0, 0.5), range2=(0.5, 1.0)) + ListedColormap + Merged colormap object. + BoundaryNorm + Normalization for mapping data to colors. """ - # Convert inputs to colormaps if they aren't already + # Convert cmap1 to a Colormap if needed if isinstance(cmap1, str): - if cmap1.startswith("#"): - cmap1 = create_cmap_from_colors([cmap1]) - else: - cmap1 = plt.get_cmap(cmap1) + cmap1 = plt.get_cmap(cmap1) elif isinstance(cmap1, list): - cmap1 = create_cmap_from_colors(cmap1) + cmap1 = colors.LinearSegmentedColormap.from_list("cmap1", cmap1) if isinstance(cmap2, str): - if cmap2.startswith("#"): - cmap2 = create_cmap_from_colors([cmap2]) - else: - cmap2 = plt.get_cmap(cmap2) + cmap2 = plt.get_cmap(cmap2) elif isinstance(cmap2, list): - cmap2 = create_cmap_from_colors(cmap2) + cmap2 = colors.LinearSegmentedColormap.from_list("cmap2", cmap2) - # Create the joined colormap with specified ranges + # Get colors from each colormap colors1 = cmap1(np.linspace(range1[0], range1[1], 128)) colors2 = cmap2(np.linspace(range2[0], range2[1], 128)) newcolors = np.vstack((colors1, colors2)) - return colors.ListedColormap(newcolors, name=name) + # Create corresponding boundaries in data space + if value_range1 is not None and value_range2 is not None: + + bounds1 = np.linspace(value_range1[0], value_range1[1], 129) + bounds2 = np.linspace(value_range2[0], value_range2[1], 129) + all_bounds = np.concatenate([bounds1[:-1], bounds2]) + + norm = BoundaryNorm(boundaries=all_bounds, ncolors=len(newcolors)) + + return colors.ListedColormap(newcolors, name=name), norm + else: + return colors.ListedColormap(newcolors, name=name) if __name__ == "__main__": diff --git a/bluemath_tk/topo_bathy/mesh_utils.py b/bluemath_tk/topo_bathy/mesh_utils.py index 8728a2f..7db8049 100644 --- a/bluemath_tk/topo_bathy/mesh_utils.py +++ b/bluemath_tk/topo_bathy/mesh_utils.py @@ -21,6 +21,9 @@ 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 + def plot_mesh_edge( msh_t: jigsaw_msh_t, ax: Axes = None, to_geo: callable = None, **kwargs @@ -153,9 +156,19 @@ def plot_bathymetry(rasters_path: List[str], polygon: Polygon, ax: Axes) -> Axes cols, rows = np.meshgrid(np.arange(width), np.arange(height)) xs, ys = rasterio.transform.xy(transform, rows, cols) + vmin = np.nanmin(data[0]) + 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") + im = ax.imshow( data[0], - cmap="gist_earth", + cmap=cmap, + norm=norm, extent=(np.min(xs), np.max(xs), np.min(ys), np.max(ys)), ) cbar = plt.colorbar(im, ax=ax) @@ -343,13 +356,26 @@ def plot_bathymetry_interp(mesh: jigsaw_msh_t, to_geo, ax: Axes) -> None: crd[:, 0], crd[:, 1] = to_geo(crd[:, 0], crd[:, 1]) bnd = [crd[:, 0].min(), crd[:, 0].max(), crd[:, 1].min(), crd[:, 1].max()] - im = ax.tricontourf( - Triangulation( - crd[:, 0], - crd[:, 1], - triangles=mesh.msh_t.tria3["index"], - ), - mesh.msh_t.value.flatten(), + triangle = mesh.msh_t.tria3["index"] + Z = np.mean(mesh.msh_t.value.flatten()[triangle],axis = 1) + vmin = np.nanmin(Z) + 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" +) + + im = ax.tripcolor( + crd[:, 0], + crd[:, 1], + triangle, + facecolors=Z, + cmap=cmap, + norm=norm, ) ax.set_title("Interpolated Bathymetry") gl = ax.gridlines(draw_labels=True) @@ -974,8 +1000,8 @@ def detect_circumcenter_too_close( def define_mesh_target_size( rasters: List[rasterio.io.DatasetReader], raster_resolution_meters: float, - nprocs: int, depth_ranges: float, + nprocs: int = 1, ) -> ocsmesh.Hfun: """ Define the mesh target size based on depth ranges and their corresponding values.