roms-tools 0.1.0__py3-none-any.whl → 0.20__py3-none-any.whl

roms_tools/setup/topography.py

@@ -4,6 +4,7 @@ import gcm_filters
 from scipy.interpolate import RegularGridInterpolator
 from scipy.ndimage import label
 from roms_tools.setup.datasets import fetch_topo
+from roms_tools.setup.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
 import warnings
 from itertools import count
 
@@ -19,7 +20,7 @@ def _add_topography_and_mask(
     hraw = xr.DataArray(data=hraw, dims=["eta_rho", "xi_rho"])
 
     # Mask is obtained by finding locations where ocean depth is positive
-    mask = xr.where(hraw > 0, 1, 0)
+    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)
@@ -37,6 +38,8 @@ def _add_topography_and_mask(
         "units": "land/water (0/1)",
     }
 
+    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"].attrs = {
@@ -57,7 +60,7 @@ def _make_raw_topography(lon, lat, topography_source) -> np.ndarray:
     topo_ds = fetch_topo(topography_source)
 
     # the following will depend on the topography source
-    if topography_source == "etopo5":
+    if topography_source == "ETOPO5":
         topo_lon = topo_ds["topo_lon"].copy()
         # Modify longitude values where necessary
         topo_lon = xr.where(topo_lon < 0, topo_lon + 360, topo_lon)
@@ -240,3 +243,15 @@ def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
     ds.attrs["rmax"] = rmax
 
     return ds
+
+
+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"].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)"}
+
+    return ds

roms_tools/setup/utils.py

@@ -0,0 +1,162 @@
+import xarray as xr
+
+
+def nan_check(field, mask) -> None:
+    """
+    Checks for NaN values at wet points in the field.
+
+    This function examines the interpolated input field for NaN values at positions indicated as wet points by the mask.
+    If any NaN values are found at these wet points, a ValueError is raised.
+
+    Parameters
+    ----------
+    field : array-like
+        The data array to be checked for NaN values. This is typically an xarray.DataArray or numpy array.
+
+    mask : array-like
+        A boolean mask or data array with the same shape as `field`. The wet points (usually ocean points)
+        are indicated by `1` or `True`, and land points by `0` or `False`.
+
+    Raises
+    ------
+    ValueError
+        If the field contains NaN values at any of the wet points indicated by the mask.
+        The error message will explain the potential cause and suggest ensuring the dataset's coverage.
+
+    """
+
+    # Replace values in field with 0 where mask is not 1
+    da = xr.where(mask == 1, field, 0)
+
+    # Check if any NaN values exist in the modified field
+    if da.isnull().any().values:
+        raise ValueError(
+            "NaN values found in interpolated field. This likely occurs because the ROMS grid, including "
+            "a small safety margin for interpolation, is not fully contained within the dataset's longitude/latitude range. Please ensure that the "
+            "dataset covers the entire area required by the ROMS grid."
+        )
+
+
+def interpolate_from_rho_to_u(field, method="additive"):
+
+    """
+    Interpolates the given field from rho points to u points.
+
+    This function performs an interpolation from the rho grid (cell centers) to the u grid
+    (cell edges in the xi direction). Depending on the chosen method, it either averages
+    (additive) or multiplies (multiplicative) the field values between adjacent rho points
+    along the xi dimension. It also handles the removal of unnecessary coordinate variables
+    and updates the dimensions accordingly.
+
+    Parameters
+    ----------
+    field : xr.DataArray
+        The input data array on the rho grid to be interpolated. It is assumed to have a dimension
+        named "xi_rho".
+
+    method : str, optional, default='additive'
+        The method to use for interpolation. Options are:
+        - 'additive': Average the field values between adjacent rho points.
+        - 'multiplicative': Multiply the field values between adjacent rho points. Appropriate for
+          binary masks.
+
+    Returns
+    -------
+    field_interpolated : xr.DataArray
+        The interpolated data array on the u grid with the dimension "xi_u".
+    """
+
+    if method == "additive":
+        field_interpolated = 0.5 * (field + field.shift(xi_rho=1)).isel(
+            xi_rho=slice(1, None)
+        )
+    elif method == "multiplicative":
+        field_interpolated = (field * field.shift(xi_rho=1)).isel(xi_rho=slice(1, None))
+    else:
+        raise NotImplementedError(f"Unsupported method '{method}' specified.")
+
+    if "lat_rho" in field_interpolated.coords:
+        field_interpolated.drop_vars(["lat_rho"])
+    if "lon_rho" in field_interpolated.coords:
+        field_interpolated.drop_vars(["lon_rho"])
+
+    field_interpolated = field_interpolated.swap_dims({"xi_rho": "xi_u"})
+
+    return field_interpolated
+
+
+def interpolate_from_rho_to_v(field, method="additive"):
+
+    """
+    Interpolates the given field from rho points to v points.
+
+    This function performs an interpolation from the rho grid (cell centers) to the v grid
+    (cell edges in the eta direction). Depending on the chosen method, it either averages
+    (additive) or multiplies (multiplicative) the field values between adjacent rho points
+    along the eta dimension. It also handles the removal of unnecessary coordinate variables
+    and updates the dimensions accordingly.
+
+    Parameters
+    ----------
+    field : xr.DataArray
+        The input data array on the rho grid to be interpolated. It is assumed to have a dimension
+        named "eta_rho".
+
+    method : str, optional, default='additive'
+        The method to use for interpolation. Options are:
+        - 'additive': Average the field values between adjacent rho points.
+        - 'multiplicative': Multiply the field values between adjacent rho points. Appropriate for
+          binary masks.
+
+    Returns
+    -------
+    field_interpolated : xr.DataArray
+        The interpolated data array on the v grid with the dimension "eta_v".
+    """
+
+    if method == "additive":
+        field_interpolated = 0.5 * (field + field.shift(eta_rho=1)).isel(
+            eta_rho=slice(1, None)
+        )
+    elif method == "multiplicative":
+        field_interpolated = (field * field.shift(eta_rho=1)).isel(
+            eta_rho=slice(1, None)
+        )
+    else:
+        raise NotImplementedError(f"Unsupported method '{method}' specified.")
+
+    if "lat_rho" in field_interpolated.coords:
+        field_interpolated.drop_vars(["lat_rho"])
+    if "lon_rho" in field_interpolated.coords:
+        field_interpolated.drop_vars(["lon_rho"])
+
+    field_interpolated = field_interpolated.swap_dims({"eta_rho": "eta_v"})
+
+    return field_interpolated
+
+
+def extrapolate_deepest_to_bottom(field: xr.DataArray, dim: str) -> xr.DataArray:
+    """
+    Extrapolate the deepest non-NaN values to the bottom along a specified dimension.
+
+    Parameters
+    ----------
+    field : xr.DataArray
+        The input data array containing NaN values that need to be filled. This array
+        should have at least one dimension named by `dim`.
+    dim : str
+        The name of the dimension along which to perform the interpolation and extrapolation.
+        Typically, this would be a vertical dimension such as 'depth' or 's_rho'.
+
+    Returns
+    -------
+    field_interpolated : xr.DataArray
+        A new data array with NaN values along the specified dimension filled by nearest
+        neighbor interpolation and extrapolation to the bottom. The original data array is not modified.
+
+    """
+    field_interpolated = field.interpolate_na(
+        dim=dim, method="nearest", fill_value="extrapolate"
+    )
+
+    return field_interpolated

roms_tools/setup/vertical_coordinate.py

@@ -0,0 +1,494 @@
+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):
+    """
+    Compute the S-coordinate stretching curves according to Shchepetkin and McWilliams (2009).
+
+    Parameters
+    ----------
+    sigma : np.ndarray or float
+        The sigma-coordinate values.
+    theta_s : float
+        The surface control parameter.
+    theta_b : float
+        The bottom control parameter.
+
+    Returns
+    -------
+    C : np.ndarray or float
+        The stretching curve values.
+
+    Raises
+    ------
+    ValueError
+        If theta_s or theta_b are not within the valid range.
+    """
+    if not (0 < theta_s <= 10):
+        raise ValueError("theta_s must be between 0 and 10.")
+    if not (0 < theta_b <= 4):
+        raise ValueError("theta_b must be between 0 and 4.")
+
+    C = (1 - np.cosh(theta_s * sigma)) / (np.cosh(theta_s) - 1)
+    C = (np.exp(theta_b * C) - 1) / (1 - np.exp(-theta_b))
+
+    return C
+
+
+def sigma_stretch(theta_s, theta_b, N, type):
+    """
+    Compute sigma and stretching curves based on the type and parameters.
+
+    Parameters
+    ----------
+    theta_s : float
+        The surface control parameter.
+    theta_b : float
+        The bottom control parameter.
+    N : int
+        The number of vertical levels.
+    type : str
+        The type of sigma ('w' for vertical velocity points, 'r' for rho-points).
+
+    Returns
+    -------
+    cs : xr.DataArray
+        The stretching curve values.
+    sigma : xr.DataArray
+        The sigma-coordinate values.
+
+    Raises
+    ------
+    ValueError
+        If the type is not 'w' or 'r'.
+    """
+    if type == "w":
+        k = xr.DataArray(np.arange(N + 1), dims="s_w")
+        sigma = (k - N) / N
+    elif type == "r":
+        k = xr.DataArray(np.arange(1, N + 1), dims="s_rho")
+        sigma = (k - N - 0.5) / N
+    else:
+        raise ValueError(
+            "Type must be either 'w' for vertical velocity points or 'r' for rho-points."
+        )
+
+    cs = compute_cs(sigma, theta_s, theta_b)
+
+    return cs, sigma
+
+
+def compute_depth(zeta, h, hc, cs, sigma):
+    """
+    Compute the depth at different sigma levels.
+
+    Parameters
+    ----------
+    zeta : xr.DataArray
+        The sea surface height.
+    h : xr.DataArray
+        The depth of the sea bottom.
+    hc : float
+        The critical depth.
+    cs : xr.DataArray
+        The stretching curve values.
+    sigma : xr.DataArray
+        The sigma-coordinate values.
+
+    Returns
+    -------
+    z : xr.DataArray
+        The depth at different sigma levels.
+
+    Raises
+    ------
+    ValueError
+        If theta_s or theta_b are less than or equal to zero.
+    """
+
+    # Expand dimensions
+    sigma = sigma.expand_dims(dim={"eta_rho": h.eta_rho, "xi_rho": h.xi_rho})
+    cs = cs.expand_dims(dim={"eta_rho": h.eta_rho, "xi_rho": h.xi_rho})
+
+    s = (hc * sigma + h * cs) / (hc + h)
+    z = zeta + (zeta + h) * s
+
+    return z