From a526c58399824d106e60aff7610e145a8cd98fce Mon Sep 17 00:00:00 2001 From: Javier Tausia Hoyal Date: Fri, 7 Feb 2025 13:54:04 +0100 Subject: [PATCH] [JTH] first binwaves version implemented --- bluemath_tk/waves/binwaves.py | 130 ++++++++++++++++++ bluemath_tk/wrappers/swan/swan_example.py | 26 +++- bluemath_tk/wrappers/swan/swan_wrapper.py | 143 ++++++++++++++++++++ bluemath_tk/wrappers/swash/swash_wrapper.py | 2 +- 4 files changed, 296 insertions(+), 5 deletions(-) create mode 100644 bluemath_tk/waves/binwaves.py diff --git a/bluemath_tk/waves/binwaves.py b/bluemath_tk/waves/binwaves.py new file mode 100644 index 0000000..13be53f --- /dev/null +++ b/bluemath_tk/waves/binwaves.py @@ -0,0 +1,130 @@ +import numpy as np +import xarray as xr + + +def transform_CAWCR_WS(ds): + ds = ds.rename({"frequency": "freq", "direction": "dir"}) + ds["efth"] = ds["efth"] * np.pi / 180.0 + ds["dir"] = ds["dir"] - 180.0 + ds["dir"] = np.where(ds["dir"] < 0, ds["dir"] + 360, ds["dir"]) + return ds + + +def process_kp_coefficients( + swan_ds: xr.Dataset, + spectrum_freq: np.ndarray, + spectrum_dir: np.ndarray, + latitude: float, + longitude: float, +): + """ + This function processes the propagation coefficients for all the grid points within the + SWAN simulation output (out_sim). + It takes a long time to run but it only needs to be done once per location + + p_store_kp - location to save all the kp coefficients. One file per location + out_sim - output from SWAN simulations. Dimensions: cases x lat x lon. + variables: Hs_part, Tp_part, Dir_part + spectrum - spectra dataset (freq and dir coordinates) + override - bool, True for re-calculate already solved points + rotated_mesh - if True looks for closer location in Xp and Yp + """ + + kp_matrix = np.full( + [ + len(swan_ds.case_num), + len(spectrum_freq), + len(spectrum_dir), + ], + 0.0, + ) + + swan_point = swan_ds.sel( + Xp=longitude, + Yp=latitude, + method="nearest", + ) + # TODO: Check if this is the correct way to handle NaN values + if any(np.isnan(swan_point["TPsmoo"].values)): + raise ValueError("NaN values found for variable TPsmoo_part") + + # Tp mask + swan_point_cut = swan_point[["Hsig", "TPsmoo", "Dir"]].where( + swan_point["TPsmoo"] > 0, + drop=True, + ) + + # get k,f,d + kfd = xr.Dataset( + { + "k": swan_point_cut.Hsig, + "f": 1 / swan_point_cut.TPsmoo, + "d": swan_point_cut.Dir, + } + ) + + # fill kp + for case_num in kfd.case_num.values: + kfd_c = kfd.sel(case_num=case_num) + + # get k,f,d and clean nans + k = kfd_c.k.values + f = kfd_c.f.values + d = kfd_c.d.values + + k = k[~np.isnan(f)] + d = d[~np.isnan(f)] + f = f[~np.isnan(f)] + + # set case kp at point + for c in range(len(f)): + i = np.argmin(np.abs(spectrum_freq - f[c])) + j = np.argmin(np.abs(spectrum_dir - d[c])) + kp_matrix[case_num, i, j] = k[c] + + # prepare output dataset + return xr.Dataset( + { + "kp": (["case", "freq", "dir"], kp_matrix), + "swan_freqs": (["case"], 1.0 / swan_point_cut.TPsmoo), + "swan_dirs": (["case"], swan_point_cut.Dir), + }, + coords={ + "case": swan_ds.case_num, + "freq": spectrum_freq, + "dir": spectrum_dir, + }, + ) + + +def reconstruc_spectra( + spectra_ds: xr.Dataset, + kp_coeffs: xr.Dataset, +): + EFTH = np.full( + np.shape(spectra_ds.efth.values), + 0, + ) + + for case in range(len(kp_coeffs.case)): + freq_, dir_ = ( + kp_coeffs.isel(case=case).swan_freqs.values, + kp_coeffs.isel(case=case).swan_dirs.values, + ) + efth_case = spectra_ds.sel(freq=freq_, dir=dir_, method="nearest") + kp_case = kp_coeffs.sortby("dir").isel(case=case) + + EFTH = EFTH + (efth_case.efth * kp_case.kp**2).values + + # ns_sp = off_sp.drop(("Wspeed", "Wdir", "Depth")).copy() + + return xr.Dataset( + { + "efth": (["time", "freq", "dir"], EFTH), + }, + coords={ + "time": spectra_ds.time, + "freq": spectra_ds.freq, + "dir": spectra_ds.dir, + }, + ) diff --git a/bluemath_tk/wrappers/swan/swan_example.py b/bluemath_tk/wrappers/swan/swan_example.py index 0ece90e..945b334 100644 --- a/bluemath_tk/wrappers/swan/swan_example.py +++ b/bluemath_tk/wrappers/swan/swan_example.py @@ -1,8 +1,9 @@ import xarray as xr + from bluemath_tk.topo_bathy.swan_grid import generate_grid_parameters +from bluemath_tk.waves.binwaves import process_kp_coefficients, transform_CAWCR_WS from bluemath_tk.wrappers.swan.swan_wrapper import SwanModelWrapper - # Usage example if __name__ == "__main__": # Load GEBCO bathymetry @@ -18,7 +19,7 @@ model_parameters = ( xr.open_dataset("/home/tausiaj/GitHub-GeoOcean/BlueMath/test_data/subset.nc") .to_dataframe() - .iloc[:100] + .iloc[:10] .to_dict(orient="list") ) # Create an instance of the SWAN model wrapper @@ -32,5 +33,22 @@ # List available launchers print(swan_wrapper.list_available_launchers()) # Run the model - swan_wrapper.run_cases(launcher="docker", parallel=True) - swan_wrapper.run_cases_bulk("sbatch javi_slurm.sh") + # swan_wrapper.run_cases(launcher="docker", parallel=True) + # Post-process the output files + postprocessed_ds = swan_wrapper.postprocess_cases() + print(postprocessed_ds) + # Load spectra example + spectra = xr.open_dataset( + "/home/tausiaj/GitHub-GeoOcean/BlueMath/test_data/Waves_Cantabria_356.08_43.82.nc" + ) + spectra_transformed = transform_CAWCR_WS(spectra) + # Extract binwaves kp coeffs + kp_coeffs = process_kp_coefficients( + swan_ds=postprocessed_ds, + spectrum_freq=spectra_transformed.freq.values, + spectrum_dir=spectra_transformed.dir.values, + latitude=43.3, + longitude=173.0, + ) + print(kp_coeffs) + # Reconstruct spectra with kp coeffs diff --git a/bluemath_tk/wrappers/swan/swan_wrapper.py b/bluemath_tk/wrappers/swan/swan_wrapper.py index ef0885a..8d4e484 100644 --- a/bluemath_tk/wrappers/swan/swan_wrapper.py +++ b/bluemath_tk/wrappers/swan/swan_wrapper.py @@ -1,3 +1,7 @@ +import os +from typing import List +import xarray as xr +import scipy.io as sio from .._base_wrappers import BaseModelWrapper @@ -26,6 +30,37 @@ class SwanModelWrapper(BaseModelWrapper): "docker": "docker run --rm -v .:/case_dir -w /case_dir tausiaj/swan-geoocean:41.51 swanrun -input input", } + output_variables = { + "Depth": { + "long_name": "Water depth at the point", + "units": "m", + }, + "Hsig": { + "long_name": "Significant wave height", + "units": "m", + }, + "Tm02": { + "long_name": "Mean wave period", + "units": "s", + }, + "Dir": { + "long_name": "Wave direction", + "units": "degrees", + }, + "PkDir": { + "long_name": "Peak wave direction", + "units": "degrees", + }, + "TPsmoo": { + "long_name": "Peak wave period", + "units": "s", + }, + "Dspr": { + "long_name": "Directional spread", + "units": "degrees", + }, + } + def __init__( self, templates_dir: str, @@ -48,3 +83,111 @@ def __init__( self.set_logger_name( name=self.__class__.__name__, level="DEBUG" if debug else "INFO" ) + + def list_available_output_variables(self) -> List[str]: + """ + List available output variables. + + Returns + ------- + List[str] + The available output variables. + """ + + return list(self.output_variables.keys()) + + def _convert_case_output_files_to_nc( + self, case_num: int, output_path: str, output_vars: List[str] + ) -> xr.Dataset: + """ + Convert mat file to netCDF file. + + Parameters + ---------- + case_num : int + The case number. + output_path : str + The output path. + output_vars : List[str] + The output variables to use. + + Returns + ------- + xr.Dataset + The xarray Dataset. + """ + + # Read mat file + output_dict = sio.loadmat(output_path) + + # Create Dataset + ds_output_dict = {var: (("Yp", "Xp"), output_dict[var]) for var in output_vars} + ds = xr.Dataset( + ds_output_dict, + coords={"Xp": output_dict["Xp"][0, :], "Yp": output_dict["Yp"][:, 0]}, + ) + + # assign correct coordinate case_num + ds.coords["case_num"] = case_num + + return ds + + def postprocess_case( + self, case_num: int, case_dir: str, output_vars: List[str] = None + ) -> xr.Dataset: + """ + Convert mat ouput files to netCDF file. + + Parameters + ---------- + case_num : int + The case number. + case_dir : str + The case directory. + output_vars : list, optional + The output variables to postprocess. Default is None. + + Returns + ------- + xr.Dataset + The postprocessed Dataset. + """ + + if output_vars is None: + self.logger.info("Postprocessing all available variables.") + output_vars = list(self.output_variables.keys()) + + output_nc_path = os.path.join(case_dir, "output.nc") + if not os.path.exists(output_nc_path): + # Convert tab files to netCDF file + output_path = os.path.join(case_dir, "output_main.mat") + output_nc = self._convert_case_output_files_to_nc( + case_num=case_num, + output_path=output_path, + output_vars=output_vars, + ) + output_nc.to_netcdf(os.path.join(case_dir, "output.nc")) + else: + self.logger.info("Reading existing output.nc file.") + output_nc = xr.open_dataset(output_nc_path) + + return output_nc + + def join_postprocessed_files( + self, postprocessed_files: List[xr.Dataset] + ) -> xr.Dataset: + """ + Join postprocessed files in a single Dataset. + + Parameters + ---------- + postprocessed_files : list + The postprocessed files. + + Returns + ------- + xr.Dataset + The joined Dataset. + """ + + return xr.concat(postprocessed_files, dim="case_num") diff --git a/bluemath_tk/wrappers/swash/swash_wrapper.py b/bluemath_tk/wrappers/swash/swash_wrapper.py index eaf6b71..24f8c66 100644 --- a/bluemath_tk/wrappers/swash/swash_wrapper.py +++ b/bluemath_tk/wrappers/swash/swash_wrapper.py @@ -166,7 +166,7 @@ def _convert_case_output_files_to_nc( # merge output files to one xarray.Dataset ds = xr.merge([ds_ouput, ds_run], compat="no_conflicts") - # assign correct coordinate case_id + # assign correct coordinate case_num ds.coords["case_num"] = case_num return ds