[JTH] first binwaves version implemented

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,
+        },
+    )

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

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")

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