roms-tools 1.0.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

roms_tools/setup/tides.py

@@ -15,11 +15,12 @@ from roms_tools.setup.utils import (
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
 )
+from roms_tools.setup.mixins import ROMSToolsMixins
 import matplotlib.pyplot as plt
 
 
 @dataclass(frozen=True, kw_only=True)
-class TidalForcing:
+class TidalForcing(ROMSToolsMixins):
     """
     Represents tidal forcing data used in ocean modeling.
 
@@ -59,30 +60,13 @@ class TidalForcing:
     ds: xr.Dataset = field(init=False, repr=False)
 
     def __post_init__(self):
-        if "name" not in self.source.keys():
-            raise ValueError("`source` must include a 'name'.")
-        if "path" not in self.source.keys():
-            raise ValueError("`source` must include a 'path'.")
-        if self.source["name"] == "TPXO":
-            data = TPXODataset(filename=self.source["path"])
-        else:
-            raise ValueError('Only "TPXO" is a valid option for source["name"].')
+
+        self._input_checks()
+        lon, lat, angle, straddle = super().get_target_lon_lat()
+
+        data = self._get_data()
 
         data.check_number_constituents(self.ntides)
-        # operate on longitudes between -180 and 180 unless ROMS domain lies at least 5 degrees in longitude away from Greenwich meridian
-        lon = self.grid.ds.lon_rho
-        lat = self.grid.ds.lat_rho
-        angle = self.grid.ds.angle
-
-        lon = xr.where(lon > 180, lon - 360, lon)
-        straddle = True
-        if not self.grid.straddle and abs(lon).min() > 5:
-            lon = xr.where(lon < 0, lon + 360, lon)
-            straddle = False
-
-        # Restrict data to relevant subdomain to achieve better performance and to avoid discontinuous longitudes introduced by converting
-        # to a different longitude range (+- 360 degrees). Discontinues longitudes can lead to artifacts in the interpolation process that
-        # would not be detected by the nan_check function.
         data.choose_subdomain(
             latitude_range=[lat.min().values, lat.max().values],
             longitude_range=[lon.min().values, lon.max().values],
@@ -121,56 +105,64 @@ class TidalForcing:
                 method="linear",
             )
 
-        # Rotate to grid orientation
-        u_Re = data_vars["u_Re"] * np.cos(angle) + data_vars["v_Re"] * np.sin(angle)
-        v_Re = data_vars["v_Re"] * np.cos(angle) - data_vars["u_Re"] * np.sin(angle)
-        u_Im = data_vars["u_Im"] * np.cos(angle) + data_vars["v_Im"] * np.sin(angle)
-        v_Im = data_vars["v_Im"] * np.cos(angle) - data_vars["u_Im"] * np.sin(angle)
+        data_vars = super().process_velocities(
+            data_vars, angle, "u_Re", "v_Re", interpolate=False
+        )
+        data_vars = super().process_velocities(
+            data_vars, angle, "u_Im", "v_Im", interpolate=False
+        )
 
         # Convert to barotropic velocity
-        u_Re = u_Re / self.grid.ds.h
-        v_Re = v_Re / self.grid.ds.h
-        u_Im = u_Im / self.grid.ds.h
-        v_Im = v_Im / self.grid.ds.h
+        for varname in ["u_Re", "v_Re", "u_Im", "v_Im"]:
+            data_vars[varname] = data_vars[varname] / self.grid.ds.h
 
         # Interpolate from rho- to velocity points
-        u_Re = interpolate_from_rho_to_u(u_Re)
-        v_Re = interpolate_from_rho_to_v(v_Re)
-        u_Im = interpolate_from_rho_to_u(u_Im)
-        v_Im = interpolate_from_rho_to_v(v_Im)
+        for uname in ["u_Re", "u_Im"]:
+            data_vars[uname] = interpolate_from_rho_to_u(data_vars[uname])
+        for vname in ["v_Re", "v_Im"]:
+            data_vars[vname] = interpolate_from_rho_to_v(data_vars[vname])
+
+        d_meta = super().get_variable_metadata()
+        ds = self._write_into_dataset(data_vars, d_meta)
+        ds["omega"] = tides["omega"]
+
+        ds = self._add_global_metadata(ds)
+
+        for var in ["ssh_Re", "u_Re", "v_Im"]:
+            nan_check(ds[var].isel(ntides=0), self.grid.ds.mask_rho)
+
+        object.__setattr__(self, "ds", ds)
+
+    def _input_checks(self):
+
+        if "name" not in self.source.keys():
+            raise ValueError("`source` must include a 'name'.")
+        if "path" not in self.source.keys():
+            raise ValueError("`source` must include a 'path'.")
+
+    def _get_data(self):
+
+        if self.source["name"] == "TPXO":
+            data = TPXODataset(filename=self.source["path"])
+        else:
+            raise ValueError('Only "TPXO" is a valid option for source["name"].')
+        return data
+
+    def _write_into_dataset(self, data_vars, d_meta):
 
         # save in new dataset
         ds = xr.Dataset()
 
-        # ds["omega"] = tides["omega"]
-
-        ds["ssh_Re"] = data_vars["ssh_Re"].astype(np.float32)
-        ds["ssh_Im"] = data_vars["ssh_Im"].astype(np.float32)
-        ds["ssh_Re"].attrs["long_name"] = "Tidal elevation, real part"
-        ds["ssh_Im"].attrs["long_name"] = "Tidal elevation, complex part"
-        ds["ssh_Re"].attrs["units"] = "m"
-        ds["ssh_Im"].attrs["units"] = "m"
-
-        ds["pot_Re"] = data_vars["pot_Re"].astype(np.float32)
-        ds["pot_Im"] = data_vars["pot_Im"].astype(np.float32)
-        ds["pot_Re"].attrs["long_name"] = "Tidal potential, real part"
-        ds["pot_Im"].attrs["long_name"] = "Tidal potential, complex part"
-        ds["pot_Re"].attrs["units"] = "m"
-        ds["pot_Im"].attrs["units"] = "m"
-
-        ds["u_Re"] = u_Re.astype(np.float32)
-        ds["u_Im"] = u_Im.astype(np.float32)
-        ds["u_Re"].attrs["long_name"] = "Tidal velocity in x-direction, real part"
-        ds["u_Im"].attrs["long_name"] = "Tidal velocity in x-direction, complex part"
-        ds["u_Re"].attrs["units"] = "m/s"
-        ds["u_Im"].attrs["units"] = "m/s"
-
-        ds["v_Re"] = v_Re.astype(np.float32)
-        ds["v_Im"] = v_Im.astype(np.float32)
-        ds["v_Re"].attrs["long_name"] = "Tidal velocity in y-direction, real part"
-        ds["v_Im"].attrs["long_name"] = "Tidal velocity in y-direction, complex part"
-        ds["v_Re"].attrs["units"] = "m/s"
-        ds["v_Im"].attrs["units"] = "m/s"
+        for var in data_vars.keys():
+            ds[var] = data_vars[var].astype(np.float32)
+            ds[var].attrs["long_name"] = d_meta[var]["long_name"]
+            ds[var].attrs["units"] = d_meta[var]["units"]
+
+        ds = ds.drop_vars(["lat_rho", "lon_rho"])
+
+        return ds
+
+    def _add_global_metadata(self, ds):
 
         ds.attrs["title"] = "ROMS tidal forcing created by ROMS-Tools"
         # Include the version of roms-tools
@@ -185,10 +177,7 @@ class TidalForcing:
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
         ds.attrs["allan_factor"] = self.allan_factor
 
-        object.__setattr__(self, "ds", ds)
-
-        for var in ["ssh_Re", "u_Re", "v_Im"]:
-            nan_check(self.ds[var].isel(ntides=0), self.grid.ds.mask_rho)
+        return ds
 
     def plot(self, varname, ntides=0) -> None:
         """
@@ -228,6 +217,25 @@ class TidalForcing:
         """
 
         field = self.ds[varname].isel(ntides=ntides).compute()
+        if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
+            field = field.where(self.grid.ds.mask_rho)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_rho, "lat": self.grid.ds.lat_rho}
+            )
+
+        elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
+            field = field.where(self.grid.ds.mask_u)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_u, "lat": self.grid.ds.lat_u}
+            )
+
+        elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
+            field = field.where(self.grid.ds.mask_v)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_v, "lat": self.grid.ds.lat_v}
+            )
+        else:
+            ValueError("provided field does not have two horizontal dimension")
 
         title = "%s, ntides = %i" % (field.long_name, self.ds[varname].ntides[ntides])
 

roms_tools/setup/topography.py

@@ -10,8 +10,42 @@ from itertools import count
 
 
 def _add_topography_and_mask(
-    ds, topography_source, smooth_factor, hmin, rmax
+    ds, topography_source, hmin, smooth_factor=8.0, rmax=0.2
 ) -> xr.Dataset:
+    """
+    Adds topography and a land/water mask to the dataset based on the provided topography source.
+
+    This function performs the following operations:
+    1. Interpolates topography data onto the desired grid.
+    2. Applies a mask based on ocean depth.
+    3. Smooths the topography globally to reduce grid-scale instabilities.
+    4. Fills enclosed basins with land.
+    5. Smooths the topography locally to ensure the steepness ratio satisfies the rmax criterion.
+    6. Adds topography metadata.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to which topography and the land/water mask will be added.
+    topography_source : str
+        The source of the topography data.
+    hmin : float
+        The minimum allowable depth for the topography.
+    smooth_factor : float, optional
+        The smoothing factor used in the domain-wide Gaussian smoothing of the
+        topography. Smaller values result in less smoothing, while larger
+        values produce more smoothing. The default is 8.0.
+    rmax : float, optional
+        The maximum allowable steepness ratio for the topography smoothing.
+        This parameter controls the local smoothing of the topography. Smaller values result in
+        smoother topography, while larger values preserve more detail. The default is 0.2.
+
+    Returns
+    -------
+    xr.Dataset
+        The dataset with added topography, mask, and metadata.
+    """
+
     lon = ds.lon_rho.values
     lat = ds.lat_rho.values
 
@@ -23,16 +57,11 @@ def _add_topography_and_mask(
     mask = xr.where(hraw > 0, 1.0, 0.0)
 
     # smooth topography domain-wide with Gaussian kernel to avoid grid scale instabilities
-    ds["hraw"] = _smooth_topography_globally(hraw, mask, smooth_factor)
-    ds["hraw"].attrs = {
-        "long_name": "Working bathymetry at rho-points",
-        "source": f"Raw bathymetry from {topography_source} (smoothing diameter {smooth_factor})",
-        "units": "meter",
-    }
+    hraw = _smooth_topography_globally(hraw, mask, smooth_factor)
 
     # fill enclosed basins with land
     mask = _fill_enclosed_basins(mask.values)
-    ds["mask_rho"] = xr.DataArray(mask, dims=("eta_rho", "xi_rho"))
+    ds["mask_rho"] = xr.DataArray(mask.astype(np.int32), dims=("eta_rho", "xi_rho"))
     ds["mask_rho"].attrs = {
         "long_name": "Mask at rho-points",
         "units": "land/water (0/1)",
@@ -41,7 +70,7 @@ def _add_topography_and_mask(
     ds = _add_velocity_masks(ds)
 
     # smooth topography locally to satisfy r < rmax
-    ds["h"] = _smooth_topography_locally(ds["hraw"] * ds["mask_rho"], hmin, rmax)
+    ds["h"] = _smooth_topography_locally(hraw * ds["mask_rho"], hmin, rmax)
     ds["h"].attrs = {
         "long_name": "Final bathymetry at rho-points",
         "units": "meter",
@@ -238,9 +267,7 @@ def _compute_rfactor(h):
 
 def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
     ds.attrs["topography_source"] = topography_source
-    ds.attrs["smooth_factor"] = smooth_factor
     ds.attrs["hmin"] = hmin
-    ds.attrs["rmax"] = rmax
 
     return ds
 
@@ -248,8 +275,12 @@ def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
 def _add_velocity_masks(ds):
 
     # add u- and v-masks
-    ds["mask_u"] = interpolate_from_rho_to_u(ds["mask_rho"], method="multiplicative")
-    ds["mask_v"] = interpolate_from_rho_to_v(ds["mask_rho"], method="multiplicative")
+    ds["mask_u"] = interpolate_from_rho_to_u(
+        ds["mask_rho"], method="multiplicative"
+    ).astype(np.int32)
+    ds["mask_v"] = interpolate_from_rho_to_v(
+        ds["mask_rho"], method="multiplicative"
+    ).astype(np.int32)
 
     ds["mask_u"].attrs = {"long_name": "Mask at u-points", "units": "land/water (0/1)"}
     ds["mask_v"].attrs = {"long_name": "Mask at v-points", "units": "land/water (0/1)"}

roms_tools/setup/vertical_coordinate.py

@@ -1,382 +1,5 @@
 import numpy as np
 import xarray as xr
-import yaml
-import importlib.metadata
-from dataclasses import dataclass, field, asdict
-from roms_tools.setup.grid import Grid
-from roms_tools.setup.utils import (
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
-)
-from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
-import matplotlib.pyplot as plt
-
-
-@dataclass(frozen=True, kw_only=True)
-class VerticalCoordinate:
-    """
-    Represents vertical coordinate for ROMS.
-
-    Parameters
-    ----------
-    grid : Grid
-        Object representing the grid information.
-    N : int
-        The number of vertical levels.
-    theta_s : float
-        The surface control parameter. Must satisfy 0 < theta_s <= 10.
-    theta_b : float
-        The bottom control parameter. Must satisfy 0 < theta_b <= 4.
-    hc : float
-        The critical depth.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        Xarray Dataset containing the atmospheric forcing data.
-    """
-
-    grid: Grid
-    N: int
-    theta_s: float
-    theta_b: float
-    hc: float
-
-    ds: xr.Dataset = field(init=False, repr=False)
-
-    def __post_init__(self):
-
-        h = self.grid.ds.h
-
-        cs_r, sigma_r = sigma_stretch(self.theta_s, self.theta_b, self.N, "r")
-        zr = compute_depth(h * 0, h, self.hc, cs_r, sigma_r)
-        cs_w, sigma_w = sigma_stretch(self.theta_s, self.theta_b, self.N, "w")
-        zw = compute_depth(h * 0, h, self.hc, cs_w, sigma_w)
-
-        ds = xr.Dataset()
-
-        ds["theta_s"] = np.float32(self.theta_s)
-        ds["theta_s"].attrs["long_name"] = "S-coordinate surface control parameter"
-        ds["theta_s"].attrs["units"] = "nondimensional"
-
-        ds["theta_b"] = np.float32(self.theta_b)
-        ds["theta_b"].attrs["long_name"] = "S-coordinate bottom control parameter"
-        ds["theta_b"].attrs["units"] = "nondimensional"
-
-        ds["Tcline"] = np.float32(self.hc)
-        ds["Tcline"].attrs["long_name"] = "S-coordinate surface/bottom layer width"
-        ds["Tcline"].attrs["units"] = "m"
-
-        ds["hc"] = np.float32(self.hc)
-        ds["hc"].attrs["long_name"] = "S-coordinate parameter critical depth"
-        ds["hc"].attrs["units"] = "m"
-
-        ds["sc_r"] = sigma_r.astype(np.float32)
-        ds["sc_r"].attrs["long_name"] = "S-coordinate at rho-points"
-        ds["sc_r"].attrs["units"] = "nondimensional"
-
-        ds["Cs_r"] = cs_r.astype(np.float32)
-        ds["Cs_r"].attrs["long_name"] = "S-coordinate stretching curves at rho-points"
-        ds["Cs_r"].attrs["units"] = "nondimensional"
-
-        depth = -zr
-        depth.attrs["long_name"] = "Layer depth at rho-points"
-        depth.attrs["units"] = "m"
-        ds = ds.assign_coords({"layer_depth_rho": depth.astype(np.float32)})
-
-        depth_u = interpolate_from_rho_to_u(depth).astype(np.float32)
-        depth_u.attrs["long_name"] = "Layer depth at u-points"
-        depth_u.attrs["units"] = "m"
-        ds = ds.assign_coords({"layer_depth_u": depth_u})
-
-        depth_v = interpolate_from_rho_to_v(depth).astype(np.float32)
-        depth_v.attrs["long_name"] = "Layer depth at v-points"
-        depth_v.attrs["units"] = "m"
-        ds = ds.assign_coords({"layer_depth_v": depth_v})
-
-        depth = -zw
-        depth.attrs["long_name"] = "Interface depth at rho-points"
-        depth.attrs["units"] = "m"
-        ds = ds.assign_coords({"interface_depth_rho": depth.astype(np.float32)})
-
-        depth_u = interpolate_from_rho_to_u(depth).astype(np.float32)
-        depth_u.attrs["long_name"] = "Interface depth at u-points"
-        depth_u.attrs["units"] = "m"
-        ds = ds.assign_coords({"interface_depth_u": depth_u})
-
-        depth_v = interpolate_from_rho_to_v(depth).astype(np.float32)
-        depth_v.attrs["long_name"] = "Interface depth at v-points"
-        depth_v.attrs["units"] = "m"
-        ds = ds.assign_coords({"interface_depth_v": depth_v})
-
-        ds = ds.drop_vars(["eta_rho", "xi_rho"])
-
-        ds.attrs["title"] = "ROMS vertical coordinate file created by ROMS-Tools"
-        # Include the version of roms-tools
-        try:
-            roms_tools_version = importlib.metadata.version("roms-tools")
-        except importlib.metadata.PackageNotFoundError:
-            roms_tools_version = "unknown"
-        ds.attrs["roms_tools_version"] = roms_tools_version
-
-        object.__setattr__(self, "ds", ds)
-
-    def plot(
-        self,
-        varname="layer_depth_rho",
-        s=None,
-        eta=None,
-        xi=None,
-    ) -> None:
-        """
-        Plot the vertical coordinate system for a given eta-, xi-, or s-slice.
-
-        Parameters
-        ----------
-        varname : str, optional
-            The field to plot. Options are "depth_rho", "depth_u", "depth_v".
-        s: int, optional
-            The s-index to plot. Default is None.
-        eta : int, optional
-            The eta-index to plot. Default is None.
-        xi : int, optional
-            The xi-index to plot. Default is None.
-
-        Returns
-        -------
-        None
-            This method does not return any value. It generates and displays a plot.
-
-        Raises
-        ------
-        ValueError
-            If the specified varname is not one of the valid options.
-            If none of s, eta, xi are specified.
-        """
-
-        if not any([s is not None, eta is not None, xi is not None]):
-            raise ValueError("At least one of s, eta, or xi must be specified.")
-
-        self.ds[varname].load()
-        field = self.ds[varname].squeeze()
-
-        if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            interface_depth = self.ds.interface_depth_rho
-        elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            interface_depth = self.ds.interface_depth_u
-        elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            interface_depth = self.ds.interface_depth_v
-
-        # slice the field as desired
-        title = field.long_name
-        if s is not None:
-            if "s_rho" in field.dims:
-                title = title + f", s_rho = {field.s_rho[s].item()}"
-                field = field.isel(s_rho=s)
-            elif "s_w" in field.dims:
-                title = title + f", s_w = {field.s_w[s].item()}"
-                field = field.isel(s_w=s)
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions (s_rho, s_w) found in ds[{varname}]."
-                )
-
-        if eta is not None:
-            if "eta_rho" in field.dims:
-                title = title + f", eta_rho = {field.eta_rho[eta].item()}"
-                field = field.isel(eta_rho=eta)
-                interface_depth = interface_depth.isel(eta_rho=eta)
-            elif "eta_v" in field.dims:
-                title = title + f", eta_v = {field.eta_v[eta].item()}"
-                field = field.isel(eta_v=eta)
-                interface_depth = interface_depth.isel(eta_v=eta)
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions (eta_rho, eta_v) found in ds[{varname}]."
-                )
-        if xi is not None:
-            if "xi_rho" in field.dims:
-                title = title + f", xi_rho = {field.xi_rho[xi].item()}"
-                field = field.isel(xi_rho=xi)
-                interface_depth = interface_depth.isel(xi_rho=xi)
-            elif "xi_u" in field.dims:
-                title = title + f", xi_u = {field.xi_u[xi].item()}"
-                field = field.isel(xi_u=xi)
-                interface_depth = interface_depth.isel(xi_u=xi)
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions (xi_rho, xi_u) found in ds[{varname}]."
-                )
-
-        if eta is None and xi is None:
-            vmax = field.max().values
-            vmin = field.min().values
-            cmap = plt.colormaps.get_cmap("YlGnBu")
-            cmap.set_bad(color="gray")
-            kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
-            _plot(
-                self.grid.ds,
-                field=field,
-                straddle=self.grid.straddle,
-                depth_contours=True,
-                title=title,
-                kwargs=kwargs,
-                c="g",
-            )
-        else:
-            if len(field.dims) == 2:
-                cmap = plt.colormaps.get_cmap("YlGnBu")
-                cmap.set_bad(color="gray")
-                kwargs = {"vmax": 0.0, "vmin": 0.0, "cmap": cmap, "add_colorbar": False}
-
-                _section_plot(
-                    xr.zeros_like(field),
-                    interface_depth=interface_depth,
-                    title=title,
-                    kwargs=kwargs,
-                )
-            else:
-                if "s_rho" in field.dims or "s_w" in field.dims:
-                    _profile_plot(field, title=title)
-                else:
-                    _line_plot(field, title=title)
-
-    def save(self, filepath: str) -> None:
-        """
-        Save the vertical coordinate information to a netCDF4 file.
-
-        Parameters
-        ----------
-        filepath
-        """
-        self.ds.to_netcdf(filepath)
-
-    def to_yaml(self, filepath: str) -> None:
-        """
-        Export the parameters of the class to a YAML file, including the version of roms-tools.
-
-        Parameters
-        ----------
-        filepath : str
-            The path to the YAML file where the parameters will be saved.
-        """
-        # Serialize Grid data
-        grid_data = asdict(self.grid)
-        grid_data.pop("ds", None)  # Exclude non-serializable fields
-        grid_data.pop("straddle", None)
-
-        # Include the version of roms-tools
-        try:
-            roms_tools_version = importlib.metadata.version("roms-tools")
-        except importlib.metadata.PackageNotFoundError:
-            roms_tools_version = "unknown"
-
-        # Create header
-        header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
-
-        grid_yaml_data = {"Grid": grid_data}
-
-        # Combine all sections
-        vertical_coordinate_data = {
-            "VerticalCoordinate": {
-                "N": self.N,
-                "theta_s": self.theta_s,
-                "theta_b": self.theta_b,
-                "hc": self.hc,
-            }
-        }
-
-        # Merge YAML data while excluding empty sections
-        yaml_data = {
-            **grid_yaml_data,
-            **vertical_coordinate_data,
-        }
-
-        with open(filepath, "w") as file:
-            # Write header
-            file.write(header)
-            # Write YAML data
-            yaml.dump(yaml_data, file, default_flow_style=False)
-
-    @classmethod
-    def from_file(cls, filepath: str) -> "VerticalCoordinate":
-        """
-        Create a VerticalCoordinate instance from an existing file.
-
-        Parameters
-        ----------
-        filepath : str
-            Path to the file containing the vertical coordinate information.
-
-        Returns
-        -------
-        VerticalCoordinate
-            A new instance of VerticalCoordinate populated with data from the file.
-        """
-        # Load the dataset from the file
-        ds = xr.open_dataset(filepath)
-
-        # Create a new VerticalCoordinate instance without calling __init__ and __post_init__
-        vertical_coordinate = cls.__new__(cls)
-
-        # Set the dataset for the vertical_corodinate instance
-        object.__setattr__(vertical_coordinate, "ds", ds)
-
-        # Manually set the remaining attributes by extracting parameters from dataset
-        object.__setattr__(vertical_coordinate, "N", ds.sizes["s_rho"])
-        object.__setattr__(vertical_coordinate, "theta_s", ds["theta_s"].values.item())
-        object.__setattr__(vertical_coordinate, "theta_b", ds["theta_b"].values.item())
-        object.__setattr__(vertical_coordinate, "hc", ds["hc"].values.item())
-        object.__setattr__(vertical_coordinate, "grid", None)
-
-        return vertical_coordinate
-
-    @classmethod
-    def from_yaml(cls, filepath: str) -> "VerticalCoordinate":
-        """
-        Create an instance of the VerticalCoordinate class from a YAML file.
-
-        Parameters
-        ----------
-        filepath : str
-            The path to the YAML file from which the parameters will be read.
-
-        Returns
-        -------
-        VerticalCoordinate
-            An instance of the VerticalCoordinate class.
-        """
-        # Read the entire file content
-        with open(filepath, "r") as file:
-            file_content = file.read()
-
-        # Split the content into YAML documents
-        documents = list(yaml.safe_load_all(file_content))
-
-        vertical_coordinate_data = None
-
-        # Process the YAML documents
-        for doc in documents:
-            if doc is None:
-                continue
-            if "VerticalCoordinate" in doc:
-                vertical_coordinate_data = doc["VerticalCoordinate"]
-                break
-
-        if vertical_coordinate_data is None:
-            raise ValueError(
-                "No VerticalCoordinate configuration found in the YAML file."
-            )
-
-        # Create Grid instance from the YAML file
-        grid = Grid.from_yaml(filepath)
-
-        # Create and return an instance of TidalForcing
-        return cls(
-            grid=grid,
-            **vertical_coordinate_data,
-        )
 
 
 def compute_cs(sigma, theta_s, theta_b):
@@ -491,4 +114,9 @@ def compute_depth(zeta, h, hc, cs, sigma):
     s = (hc * sigma + h * cs) / (hc + h)
     z = zeta + (zeta + h) * s
 
+    if "s_rho" in z.dims:
+        z = z.transpose("s_rho", "eta_rho", "xi_rho")
+    elif "s_w" in z.dims:
+        z = z.transpose("s_w", "eta_rho", "xi_rho")
+
     return z