roms-tools 1.4.2__py3-none-any.whl → 1.6.0__py3-none-any.whl

roms_tools/setup/initial_conditions.py

@@ -12,15 +12,19 @@ from roms_tools.setup.utils import (
     substitute_nans_by_fillvalue,
     get_variable_metadata,
     save_datasets,
+    get_target_coords,
+    rotate_velocities,
+    compute_barotropic_velocity,
+    transpose_dimensions,
 )
-from roms_tools.setup.mixins import ROMSToolsMixins
+from roms_tools.setup.regrid import LateralRegrid, VerticalRegrid
 from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
 import matplotlib.pyplot as plt
 from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
-class InitialConditions(ROMSToolsMixins):
+class InitialConditions:
     """Represents initial conditions for ROMS, including physical and biogeochemical
     data.
 
@@ -85,58 +89,103 @@ class InitialConditions(ROMSToolsMixins):
     def __post_init__(self):
 
         self._input_checks()
-        lon, lat, angle, straddle = super()._get_target_lon_lat()
 
-        data = self._get_data()
-        data.choose_subdomain(
-            latitude_range=[lat.min().values, lat.max().values],
-            longitude_range=[lon.min().values, lon.max().values],
-            margin=2,
-            straddle=straddle,
-        )
-
-        vars_2d = ["zeta"]
-        vars_3d = ["temp", "salt", "u", "v"]
-        data_vars = super()._regrid_data(data, vars_2d, vars_3d, lon, lat)
-        data_vars = super()._process_velocities(data_vars, angle, "u", "v")
+        processed_fields = {}
+        processed_fields = self._process_data(processed_fields, type="physics")
 
         if self.bgc_source is not None:
-            bgc_data = self._get_bgc_data()
-            bgc_data.choose_subdomain(
-                latitude_range=[lat.min().values, lat.max().values],
-                longitude_range=[lon.min().values, lon.max().values],
-                margin=2,
-                straddle=straddle,
-            )
-
-            vars_2d = []
-            vars_3d = bgc_data.var_names.keys()
-            bgc_data_vars = super()._regrid_data(bgc_data, vars_2d, vars_3d, lon, lat)
+            processed_fields = self._process_data(processed_fields, type="bgc")
 
-            # Ensure time coordinate matches that of physical variables
-            for var in bgc_data_vars.keys():
-                bgc_data_vars[var] = bgc_data_vars[var].assign_coords(
-                    {"time": data_vars["temp"]["time"]}
-                )
-
-            # Combine data variables from physical and biogeochemical sources
-            data_vars.update(bgc_data_vars)
+        for var_name in processed_fields.keys():
+            processed_fields[var_name] = transpose_dimensions(
+                processed_fields[var_name]
+            )
 
         d_meta = get_variable_metadata()
-        ds = self._write_into_dataset(data_vars, d_meta)
+        ds = self._write_into_dataset(processed_fields, d_meta)
 
         ds = self._add_global_metadata(ds)
 
-        ds["zeta"].load()
-        # NaN values at wet points indicate that the raw data did not cover the domain, and the following will raise a ValueError
-        nan_check(ds["zeta"].squeeze(), self.grid.ds.mask_rho)
+        self._validate(ds)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
-        for var in ds.data_vars:
-            ds[var] = substitute_nans_by_fillvalue(ds[var])
+        for var_name in ds.data_vars:
+            ds[var_name] = substitute_nans_by_fillvalue(ds[var_name])
 
         object.__setattr__(self, "ds", ds)
 
+    def _process_data(self, processed_fields, type="physics"):
+
+        target_coords = get_target_coords(self.grid)
+
+        if type == "physics":
+            data = self._get_data()
+        else:
+            data = self._get_bgc_data()
+
+        data.choose_subdomain(
+            target_coords,
+            buffer_points=20,  # lateral fill needs good buffer from data margin
+        )
+
+        data.extrapolate_deepest_to_bottom()
+        data.apply_lateral_fill()
+
+        variable_info = self._set_variable_info(data, type=type)
+        var_names = variable_info.keys()
+
+        # lateral regridding
+        lateral_regrid = LateralRegrid(target_coords, data.dim_names)
+        for var_name in var_names:
+            if var_name in data.var_names.keys():
+                processed_fields[var_name] = lateral_regrid.apply(
+                    data.ds[data.var_names[var_name]]
+                )
+
+        # rotation of velocities and interpolation to u/v points
+        if "u" in variable_info and "v" in variable_info:
+            (processed_fields["u"], processed_fields["v"],) = rotate_velocities(
+                processed_fields["u"],
+                processed_fields["v"],
+                target_coords["angle"],
+                interpolate=True,
+            )
+
+        # vertical regridding
+        for location in ["rho", "u", "v"]:
+            var_names = [
+                name
+                for name, info in variable_info.items()
+                if info["location"] == location and info["is_3d"]
+            ]
+            if len(var_names) > 0:
+                vertical_regrid = VerticalRegrid(
+                    self.grid.ds[f"layer_depth_{location}"],
+                    data.ds[data.dim_names["depth"]],
+                )
+                for var_name in var_names:
+                    if var_name in processed_fields:
+                        processed_fields[var_name] = vertical_regrid.apply(
+                            processed_fields[var_name]
+                        )
+
+        # compute barotropic velocities
+        if "u" in variable_info and "v" in variable_info:
+            for var_name in ["u", "v"]:
+                processed_fields[f"{var_name}bar"] = compute_barotropic_velocity(
+                    processed_fields[var_name],
+                    self.grid.ds[f"interface_depth_{var_name}"],
+                )
+
+        if type == "bgc":
+            # Ensure time coordinate matches that of physical variables
+            for var_name in variable_info.keys():
+                processed_fields[var_name] = processed_fields[var_name].assign_coords(
+                    {"time": processed_fields["temp"]["time"]}
+                )
+
+        return processed_fields
+
     def _input_checks(self):
 
         if "name" not in self.source.keys():
@@ -201,19 +250,86 @@ class InitialConditions(ROMSToolsMixins):
 
         return data
 
-    def _write_into_dataset(self, data_vars, d_meta):
+    def _set_variable_info(self, data, type="physics"):
+        """Sets up a dictionary with metadata for variables based on the type.
+
+        The dictionary contains the following information:
+        - `location`: Where the variable resides in the grid (e.g., rho, u, or v points).
+        - `is_vector`: Whether the variable is part of a vector (True for velocity components like 'u' and 'v').
+        - `vector_pair`: For vector variables, this indicates the associated variable that forms the vector (e.g., 'u' and 'v').
+        - `is_3d`: Indicates whether the variable is 3D (True for variables like 'temp' and 'salt') or 2D (False for 'zeta').
+
+        Returns
+        -------
+        dict
+            A dictionary where the keys are variable names and the values are dictionaries of metadata
+            about each variable, including 'location', 'is_vector', 'vector_pair', and 'is_3d'.
+        """
+        default_info = {
+            "location": "rho",
+            "is_vector": False,
+            "vector_pair": None,
+            "is_3d": True,
+        }
+
+        # Define a dictionary for variable names and their associated information
+        if type == "physics":
+            variable_info = {
+                "zeta": {
+                    "location": "rho",
+                    "is_vector": False,
+                    "vector_pair": None,
+                    "is_3d": False,
+                },
+                "temp": default_info,
+                "salt": default_info,
+                "u": {
+                    "location": "u",
+                    "is_vector": True,
+                    "vector_pair": "v",
+                    "is_3d": True,
+                },
+                "v": {
+                    "location": "v",
+                    "is_vector": True,
+                    "vector_pair": "u",
+                    "is_3d": True,
+                },
+                "ubar": {
+                    "location": "u",
+                    "is_vector": True,
+                    "vector_pair": "vbar",
+                    "is_3d": False,
+                },
+                "vbar": {
+                    "location": "v",
+                    "is_vector": True,
+                    "vector_pair": "ubar",
+                    "is_3d": False,
+                },
+            }
+        elif type == "bgc":
+            variable_info = {}
+            for var_name in data.var_names.keys():
+                variable_info[var_name] = default_info
+
+        return variable_info
+
+    def _write_into_dataset(self, processed_fields, d_meta):
 
         # save in new dataset
         ds = xr.Dataset()
 
-        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"]
+        for var_name in processed_fields.keys():
+            ds[var_name] = processed_fields[var_name].astype(np.float32)
+            ds[var_name].attrs["long_name"] = d_meta[var_name]["long_name"]
+            ds[var_name].attrs["units"] = d_meta[var_name]["units"]
 
         # initialize vertical velocity to zero
         ds["w"] = xr.zeros_like(
-            self.grid.ds["interface_depth_rho"].expand_dims(time=data_vars["u"].time)
+            self.grid.ds["interface_depth_rho"].expand_dims(
+                time=processed_fields["u"].time
+            )
         ).astype(np.float32)
         ds["w"].attrs["long_name"] = d_meta["w"]["long_name"]
         ds["w"].attrs["units"] = d_meta["w"]["units"]
@@ -222,14 +338,18 @@ class InitialConditions(ROMSToolsMixins):
             "s_rho",
             "lat_rho",
             "lon_rho",
-            "layer_depth_rho",
-            "interface_depth_rho",
             "lat_u",
             "lon_u",
             "lat_v",
             "lon_v",
+            "layer_depth_rho",
+            "interface_depth_rho",
+            "layer_depth_u",
+            "interface_depth_u",
+            "layer_depth_v",
+            "interface_depth_v",
         ]
-        existing_vars = [var for var in variables_to_drop if var in ds]
+        existing_vars = [var_name for var_name in variables_to_drop if var_name in ds]
         ds = ds.drop_vars(existing_vars)
 
         ds["Cs_r"] = self.grid.ds["Cs_r"]
@@ -253,6 +373,29 @@ class InitialConditions(ROMSToolsMixins):
 
         return ds
 
+    def _validate(self, ds):
+        """Validates the dataset by checking for NaN values in SSH at wet points, which
+        would indicate missing raw data coverage over the target domain.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            The dataset to validate.
+
+        Raises
+        ------
+        ValueError
+            If NaN values are found in any of the specified variables at wet points,
+            indicating incomplete data coverage.
+
+        Notes
+        -----
+        This check is only applied to the 2D variable SSH to improve performance.
+        """
+
+        ds["zeta"].load()
+        nan_check(ds["zeta"].squeeze(), self.grid.ds.mask_rho)
+
     def _add_global_metadata(self, ds):
 
         ds.attrs["title"] = "ROMS initial conditions file created by ROMS-Tools"
@@ -276,7 +419,7 @@ class InitialConditions(ROMSToolsMixins):
 
     def plot(
         self,
-        varname,
+        var_name,
         s=None,
         eta=None,
         xi=None,
@@ -287,7 +430,7 @@ class InitialConditions(ROMSToolsMixins):
 
         Parameters
         ----------
-        varname : str
+        var_name : str
             The name of the initial conditions field to plot. Options include:
 
             - "temp": Potential temperature.
@@ -358,30 +501,30 @@ class InitialConditions(ROMSToolsMixins):
         Raises
         ------
         ValueError
-            If the specified `varname` is not one of the valid options.
-            If the field specified by `varname` is 3D and none of `s`, `eta`, or `xi` are specified.
-            If the field specified by `varname` is 2D and both `eta` and `xi` are specified.
+            If the specified `var_name` is not one of the valid options.
+            If the field specified by `var_name` is 3D and none of `s`, `eta`, or `xi` are specified.
+            If the field specified by `var_name` is 2D and both `eta` and `xi` are specified.
         """
 
-        if len(self.ds[varname].squeeze().dims) == 3 and not any(
+        if len(self.ds[var_name].squeeze().dims) == 3 and not any(
             [s is not None, eta is not None, xi is not None]
         ):
             raise ValueError(
                 "For 3D fields, at least one of s, eta, or xi must be specified."
             )
 
-        if len(self.ds[varname].squeeze().dims) == 2 and all(
+        if len(self.ds[var_name].squeeze().dims) == 2 and all(
             [eta is not None, xi is not None]
         ):
             raise ValueError("For 2D fields, specify either eta or xi, not both.")
 
-        self.ds[varname].load()
-        field = self.ds[varname].squeeze()
+        self.ds[var_name].load()
+        field = self.ds[var_name].squeeze()
 
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
             interface_depth = self.grid.ds.interface_depth_rho
             layer_depth = self.grid.ds.layer_depth_rho
-            field = field.where(self.grid.ds.mask_rho)
+            mask = self.grid.ds.mask_rho
             field = field.assign_coords(
                 {"lon": self.grid.ds.lon_rho, "lat": self.grid.ds.lat_rho}
             )
@@ -389,7 +532,7 @@ class InitialConditions(ROMSToolsMixins):
         elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
             interface_depth = self.grid.ds.interface_depth_u
             layer_depth = self.grid.ds.layer_depth_u
-            field = field.where(self.grid.ds.mask_u)
+            mask = self.grid.ds.mask_u
             field = field.assign_coords(
                 {"lon": self.grid.ds.lon_u, "lat": self.grid.ds.lat_u}
             )
@@ -397,7 +540,7 @@ class InitialConditions(ROMSToolsMixins):
         elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
             interface_depth = self.grid.ds.interface_depth_v
             layer_depth = self.grid.ds.layer_depth_v
-            field = field.where(self.grid.ds.mask_v)
+            mask = self.grid.ds.mask_v
             field = field.assign_coords(
                 {"lon": self.grid.ds.lon_v, "lat": self.grid.ds.lat_v}
             )
@@ -419,45 +562,49 @@ class InitialConditions(ROMSToolsMixins):
                 title = title + f", eta_rho = {field.eta_rho[eta].item()}"
                 field = field.isel(eta_rho=eta)
                 layer_depth = layer_depth.isel(eta_rho=eta)
-                field = field.assign_coords({"layer_depth": layer_depth})
                 interface_depth = interface_depth.isel(eta_rho=eta)
+                if "s_rho" in field.dims:
+                    field = field.assign_coords({"layer_depth": layer_depth})
             elif "eta_v" in field.dims:
                 title = title + f", eta_v = {field.eta_v[eta].item()}"
                 field = field.isel(eta_v=eta)
                 layer_depth = layer_depth.isel(eta_v=eta)
-                field = field.assign_coords({"layer_depth": layer_depth})
                 interface_depth = interface_depth.isel(eta_v=eta)
+                if "s_rho" in field.dims:
+                    field = field.assign_coords({"layer_depth": layer_depth})
             else:
                 raise ValueError(
-                    f"None of the expected dimensions (eta_rho, eta_v) found in ds[{varname}]."
+                    f"None of the expected dimensions (eta_rho, eta_v) found in ds[{var_name}]."
                 )
         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)
                 layer_depth = layer_depth.isel(xi_rho=xi)
-                field = field.assign_coords({"layer_depth": layer_depth})
                 interface_depth = interface_depth.isel(xi_rho=xi)
+                if "s_rho" in field.dims:
+                    field = field.assign_coords({"layer_depth": layer_depth})
             elif "xi_u" in field.dims:
                 title = title + f", xi_u = {field.xi_u[xi].item()}"
                 field = field.isel(xi_u=xi)
                 layer_depth = layer_depth.isel(xi_u=xi)
-                field = field.assign_coords({"layer_depth": layer_depth})
                 interface_depth = interface_depth.isel(xi_u=xi)
+                if "s_rho" in field.dims:
+                    field = field.assign_coords({"layer_depth": layer_depth})
             else:
                 raise ValueError(
-                    f"None of the expected dimensions (xi_rho, xi_u) found in ds[{varname}]."
+                    f"None of the expected dimensions (xi_rho, xi_u) found in ds[{var_name}]."
                 )
 
         # chose colorbar
-        if varname in ["u", "v", "w", "ubar", "vbar", "zeta"]:
+        if var_name in ["u", "v", "w", "ubar", "vbar", "zeta"]:
             vmax = max(field.max().values, -field.min().values)
             vmin = -vmax
             cmap = plt.colormaps.get_cmap("RdBu_r")
         else:
             vmax = field.max().values
             vmin = field.min().values
-            if varname in ["temp", "salt"]:
+            if var_name in ["temp", "salt"]:
                 cmap = plt.colormaps.get_cmap("YlOrRd")
             else:
                 cmap = plt.colormaps.get_cmap("YlGn")
@@ -467,7 +614,7 @@ class InitialConditions(ROMSToolsMixins):
         if eta is None and xi is None:
             _plot(
                 self.grid.ds,
-                field=field,
+                field=field.where(mask),
                 straddle=self.grid.straddle,
                 depth_contours=depth_contours,
                 title=title,

roms_tools/setup/regrid.py

@@ -0,0 +1,143 @@
+import xarray as xr
+
+
+class LateralRegrid:
+    """Handles lateral regridding of data onto a new spatial grid."""
+
+    def __init__(self, target_coords, source_dim_names):
+        """Initialize target grid coordinates and names for grid dimensions.
+
+        Parameters
+        ----------
+        target_coords : dict
+            Dictionary containing 'lon' and 'lat' as xarray.DataArrays representing
+            the longitude and latitude values of the target grid.
+        source_dim_names : dict
+            Dictionary specifying names for the latitude and longitude dimensions,
+            typically using keys like "latitude" and "longitude" to align with the dataset conventions.
+
+        Attributes
+        ----------
+        coords : dict
+            Maps the dimension names to the corresponding latitude and longitude
+            DataArrays, providing easy access to target grid coordinates.
+        """
+
+        self.coords = {
+            source_dim_names["latitude"]: target_coords["lat"],
+            source_dim_names["longitude"]: target_coords["lon"],
+        }
+
+    def apply(self, da):
+        """Fills missing values and regrids the variable.
+
+        Parameters
+        ----------
+        da : xarray.DataArray
+            Input data to fill and regrid.
+
+        Returns
+        -------
+        xarray.DataArray
+            Regridded data with filled values.
+        """
+        regridded = da.interp(self.coords, method="linear").drop_vars(
+            list(self.coords.keys())
+        )
+        return regridded
+
+
+class VerticalRegrid:
+    """Interpolates data onto new vertical (depth) coordinates.
+
+    Parameters
+    ----------
+    target_depth_coords : xarray.DataArray
+        Depth coordinates for the target grid.
+    source_depth_coords : xarray.DataArray
+        Depth coordinates for the source grid.
+    """
+
+    def __init__(self, target_depth_coords, source_depth_coords):
+        """Initialize regridding factors for interpolation.
+
+        Parameters
+        ----------
+        target_depth_coords : xarray.DataArray
+            Depth coordinates for the target grid.
+        source_depth_coords : xarray.DataArray
+            Depth coordinates for the source grid.
+
+        Attributes
+        ----------
+        coeff : xarray.Dataset
+            Dataset containing:
+            - `is_below` : Boolean mask for depths just below target.
+            - `is_above` : Boolean mask for depths just above target.
+            - `upper_mask`, `lower_mask` : Masks for valid interpolation bounds.
+            - `factor` : Weight for blending values between levels.
+        """
+
+        self.depth_dim = source_depth_coords.dims[0]
+        source_depth = source_depth_coords
+        dims = {"dim": self.depth_dim}
+
+        dlev = source_depth - target_depth_coords
+        is_below = dlev == dlev.where(dlev >= 0).min(**dims)
+        is_above = dlev == dlev.where(dlev <= 0).max(**dims)
+        p_below = dlev.where(is_below).sum(**dims)
+        p_above = -dlev.where(is_above).sum(**dims)
+        denominator = p_below + p_above
+        denominator = denominator.where(denominator > 1e-6, 1e-6)
+        factor = p_below / denominator
+
+        upper_mask = is_above.sum(**dims) > 0
+        lower_mask = is_below.sum(**dims) > 0
+
+        self.coeff = xr.Dataset(
+            {
+                "is_below": is_below,
+                "is_above": is_above,
+                "upper_mask": upper_mask,
+                "lower_mask": lower_mask,
+                "factor": factor,
+            }
+        )
+
+    def apply(self, var, fill_nans=True):
+        """Interpolates the variable onto the new depth grid using precomputed
+        coefficients for linear interpolation between layers.
+
+        Parameters
+        ----------
+        var : xarray.DataArray
+            The input data to be regridded along the depth dimension. This should be
+            an array with the same depth coordinates as the original grid.
+        fill_nans : bool, optional
+            Whether to fill NaN values in the regridded data. If True (default),
+            forward-fill and backward-fill are applied along the 's_rho' dimension to
+            ensure there are no NaNs after interpolation.
+
+        Returns
+        -------
+        xarray.DataArray
+            The regridded data array, interpolated onto the new depth grid. NaN values
+            are replaced if `fill_nans=True`, with extrapolation allowed at the surface
+            and bottom layers to minimize gaps.
+        """
+
+        dims = {"dim": self.depth_dim}
+
+        var_below = var.where(self.coeff["is_below"]).sum(**dims)
+        var_above = var.where(self.coeff["is_above"]).sum(**dims)
+
+        result = var_below + (var_above - var_below) * self.coeff["factor"]
+        if fill_nans:
+            result = result.where(self.coeff["upper_mask"], var.isel({dims["dim"]: 0}))
+            result = result.where(self.coeff["lower_mask"], var.isel({dims["dim"]: -1}))
+        else:
+            result = result.where(self.coeff["upper_mask"]).where(
+                self.coeff["lower_mask"]
+            )
+
+        return result