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

roms_tools/setup/initial_conditions.py

@@ -5,7 +5,6 @@ import importlib.metadata
 from dataclasses import dataclass, field, asdict
 from typing import Optional, Dict, Union
 from roms_tools.setup.grid import Grid
-from roms_tools.setup.vertical_coordinate import VerticalCoordinate
 from datetime import datetime
 from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
 from roms_tools.setup.utils import (
@@ -25,8 +24,6 @@ class InitialConditions(ROMSToolsMixins):
     ----------
     grid : Grid
         Object representing the grid information used for the model.
-    vertical_coordinate : VerticalCoordinate
-        Object representing the vertical coordinate system.
     ini_time : datetime
         The date and time at which the initial conditions are set.
     physics_source : Dict[str, Union[str, None]]
@@ -51,7 +48,6 @@ class InitialConditions(ROMSToolsMixins):
     --------
     >>> initial_conditions = InitialConditions(
     ...     grid=grid,
-    ...     vertical_coordinate=vertical_coordinate,
     ...     ini_time=datetime(2022, 1, 1),
     ...     physics_source={"name": "GLORYS", "path": "physics_data.nc"},
     ...     bgc_source={
@@ -63,7 +59,6 @@ class InitialConditions(ROMSToolsMixins):
     """
 
     grid: Grid
-    vertical_coordinate: VerticalCoordinate
     ini_time: datetime
     physics_source: Dict[str, Union[str, None]]
     bgc_source: Optional[Dict[str, Union[str, None]]] = None
@@ -87,7 +82,7 @@ class InitialConditions(ROMSToolsMixins):
         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)
+        data_vars = super().process_velocities(data_vars, angle, "u", "v")
 
         if self.bgc_source is not None:
             bgc_data = self._get_bgc_data()
@@ -177,18 +172,18 @@ class InitialConditions(ROMSToolsMixins):
 
         if self.bgc_source["name"] == "CESM_REGRIDDED":
 
-            bgc_data = CESMBGCDataset(
+            data = CESMBGCDataset(
                 filename=self.bgc_source["path"],
                 start_time=self.ini_time,
                 climatology=self.bgc_source["climatology"],
             )
-            bgc_data.post_process()
+            data.post_process()
         else:
             raise ValueError(
                 'Only "CESM_REGRIDDED" is a valid option for bgc_source["name"].'
             )
 
-        return bgc_data
+        return data
 
     def _write_into_dataset(self, data_vars, d_meta):
 
@@ -202,26 +197,48 @@ class InitialConditions(ROMSToolsMixins):
 
         # initialize vertical velocity to zero
         ds["w"] = xr.zeros_like(
-            self.vertical_coordinate.ds["interface_depth_rho"].expand_dims(
-                time=data_vars["u"].time
-            )
+            self.grid.ds["interface_depth_rho"].expand_dims(time=data_vars["u"].time)
         ).astype(np.float32)
         ds["w"].attrs["long_name"] = d_meta["w"]["long_name"]
         ds["w"].attrs["units"] = d_meta["w"]["units"]
 
+        variables_to_drop = [
+            "s_rho",
+            "lat_rho",
+            "lon_rho",
+            "layer_depth_rho",
+            "interface_depth_rho",
+            "lat_u",
+            "lon_u",
+            "lat_v",
+            "lon_v",
+        ]
+        existing_vars = [var for var in variables_to_drop if var in ds]
+        ds = ds.drop_vars(existing_vars)
+
+        ds["sc_r"] = self.grid.ds["sc_r"]
+        ds["Cs_r"] = self.grid.ds["Cs_r"]
+
+        # Preserve absolute time coordinate for readability
+        ds = ds.assign_coords({"abs_time": ds["time"]})
+
+        # Translate the time coordinate to days since the model reference date
+        model_reference_date = np.datetime64(self.model_reference_date)
+
+        # Convert the time coordinate to the format expected by ROMS (days since model reference date)
+        ocean_time = (ds["time"] - model_reference_date).astype("float64") * 1e-9
+        ds = ds.assign_coords(ocean_time=("time", np.float32(ocean_time)))
+        ds["ocean_time"].attrs[
+            "long_name"
+        ] = f"seconds since {np.datetime_as_string(model_reference_date, unit='s')}"
+        ds["ocean_time"].attrs["units"] = "seconds"
+        ds = ds.swap_dims({"time": "ocean_time"})
+        ds = ds.drop_vars("time")
+
         return ds
 
     def _add_global_metadata(self, ds):
 
-        ds = ds.assign_coords(
-            {
-                "layer_depth_u": self.vertical_coordinate.ds["layer_depth_u"],
-                "layer_depth_v": self.vertical_coordinate.ds["layer_depth_v"],
-                "interface_depth_u": self.vertical_coordinate.ds["interface_depth_u"],
-                "interface_depth_v": self.vertical_coordinate.ds["interface_depth_v"],
-            }
-        )
-
         ds.attrs["title"] = "ROMS initial conditions file created by ROMS-Tools"
         # Include the version of roms-tools
         try:
@@ -235,25 +252,9 @@ class InitialConditions(ROMSToolsMixins):
         if self.bgc_source is not None:
             ds.attrs["bgc_source"] = self.bgc_source["name"]
 
-        # Translate the time coordinate to days since the model reference date
-        model_reference_date = np.datetime64(self.model_reference_date)
-
-        # Convert the time coordinate to the format expected by ROMS (days since model reference date)
-        ocean_time = (ds["time"] - model_reference_date).astype("float64") * 1e-9
-        ds = ds.assign_coords(ocean_time=("time", np.float32(ocean_time)))
-        ds["ocean_time"].attrs[
-            "long_name"
-        ] = f"seconds since {np.datetime_as_string(model_reference_date, unit='s')}"
-        ds["ocean_time"].attrs["units"] = "seconds"
-
-        ds.attrs["theta_s"] = self.vertical_coordinate.ds["theta_s"].item()
-        ds.attrs["theta_b"] = self.vertical_coordinate.ds["theta_b"].item()
-        ds.attrs["Tcline"] = self.vertical_coordinate.ds["Tcline"].item()
-        ds.attrs["hc"] = self.vertical_coordinate.ds["hc"].item()
-        ds["sc_r"] = self.vertical_coordinate.ds["sc_r"]
-        ds["Cs_r"] = self.vertical_coordinate.ds["Cs_r"]
-
-        ds = ds.drop_vars(["s_rho"])
+        ds.attrs["theta_s"] = self.grid.ds.attrs["theta_s"]
+        ds.attrs["theta_b"] = self.grid.ds.attrs["theta_b"]
+        ds.attrs["hc"] = self.grid.ds.attrs["hc"]
 
         return ds
 
@@ -314,13 +315,23 @@ class InitialConditions(ROMSToolsMixins):
             - "diazP": Diazotroph Phosphorus (mmol/m³).
             - "diazFe": Diazotroph Iron (mmol/m³).
         s : int, optional
-            The index of the vertical layer to plot. Default is None.
+            The index of the vertical layer (`s_rho`) to plot. If not specified, the plot
+            will represent a horizontal slice (eta- or xi- plane). Default is None.
         eta : int, optional
-            The eta-index to plot. Default is None.
+            The eta-index to plot. Used for vertical sections or horizontal slices.
+            Default is None.
         xi : int, optional
-            The xi-index to plot. Default is None.
+            The xi-index to plot. Used for vertical sections or horizontal slices.
+            Default is None.
         depth_contours : bool, optional
-            Whether to include depth contours in the plot. Default is False.
+            If True, depth contours will be overlaid on the plot, showing lines of constant
+            depth. This is typically used for plots that show a single vertical layer.
+            Default is False.
+        layer_contours : bool, optional
+            If True, contour lines representing the boundaries between vertical layers will
+            be added to the plot. This is particularly useful in vertical sections to
+            visualize the layering of the water column. For clarity, the number of layer
+            contours displayed is limited to a maximum of 10. Default is False.
 
         Returns
         -------
@@ -333,6 +344,7 @@ class InitialConditions(ROMSToolsMixins):
             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 len(self.ds[varname].squeeze().dims) == 3 and not any(
@@ -351,17 +363,38 @@ class InitialConditions(ROMSToolsMixins):
         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
+            interface_depth = self.grid.ds.interface_depth_rho
+            layer_depth = self.grid.ds.layer_depth_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"]):
-            interface_depth = self.ds.interface_depth_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)
+            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"]):
-            interface_depth = self.ds.interface_depth_v
+            interface_depth = self.grid.ds.interface_depth_v
+            layer_depth = self.grid.ds.layer_depth_v
+            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")
 
         # slice the field as desired
         title = field.long_name
         if s is not None:
             title = title + f", s_rho = {field.s_rho[s].item()}"
             field = field.isel(s_rho=s)
+            layer_depth = layer_depth.isel(s_rho=s)
+            field = field.assign_coords({"layer_depth": layer_depth})
         else:
             depth_contours = False
 
@@ -369,10 +402,14 @@ class InitialConditions(ROMSToolsMixins):
             if "eta_rho" in field.dims:
                 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)
             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)
             else:
                 raise ValueError(
@@ -382,10 +419,14 @@ class InitialConditions(ROMSToolsMixins):
             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)
             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)
             else:
                 raise ValueError(
@@ -462,11 +503,6 @@ class InitialConditions(ROMSToolsMixins):
         grid_data.pop("ds", None)  # Exclude non-serializable fields
         grid_data.pop("straddle", None)
 
-        # Serialize VerticalCoordinate data
-        vertical_coordinate_data = asdict(self.vertical_coordinate)
-        vertical_coordinate_data.pop("ds", None)  # Exclude non-serializable fields
-        vertical_coordinate_data.pop("grid", None)  # Exclude non-serializable fields
-
         # Include the version of roms-tools
         try:
             roms_tools_version = importlib.metadata.version("roms-tools")
@@ -477,7 +513,6 @@ class InitialConditions(ROMSToolsMixins):
         header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
 
         grid_yaml_data = {"Grid": grid_data}
-        vertical_coordinate_yaml_data = {"VerticalCoordinate": vertical_coordinate_data}
 
         initial_conditions_data = {
             "InitialConditions": {
@@ -492,7 +527,6 @@ class InitialConditions(ROMSToolsMixins):
 
         yaml_data = {
             **grid_yaml_data,
-            **vertical_coordinate_yaml_data,
             **initial_conditions_data,
         }
 
@@ -545,13 +579,10 @@ class InitialConditions(ROMSToolsMixins):
                 initial_conditions_data[date_string]
             )
 
-        # Create VerticalCoordinate instance from the YAML file
-        vertical_coordinate = VerticalCoordinate.from_yaml(filepath)
-        grid = vertical_coordinate.grid
+        grid = Grid.from_yaml(filepath)
 
         # Create and return an instance of InitialConditions
         return cls(
             grid=grid,
-            vertical_coordinate=vertical_coordinate,
             **initial_conditions_data,
         )

roms_tools/setup/mixins.py

@@ -1,6 +1,5 @@
 from dataclasses import dataclass
 from roms_tools.setup.grid import Grid
-from roms_tools.setup.vertical_coordinate import VerticalCoordinate
 from roms_tools.setup.fill import fill_and_interpolate
 from roms_tools.setup.utils import (
     extrapolate_deepest_to_bottom,
@@ -21,13 +20,10 @@ class ROMSToolsMixins:
     ----------
     grid : Grid
         Object representing the grid information used for the model.
-    vertical_coordinate : VerticalCoordinate
-        Object representing the vertical coordinate system. Defaults to None.
 
     """
 
     grid: Grid
-    vertical_coordinate: VerticalCoordinate = None
 
     def get_target_lon_lat(self, use_coarse_grid=False):
         """
@@ -54,11 +50,6 @@ class ROMSToolsMixins:
         """
 
         if use_coarse_grid:
-            if "lon_coarse" not in self.grid.ds:
-                raise ValueError(
-                    "Grid has not been coarsened yet. Execute grid.coarsen() first."
-                )
-
             lon = self.grid.ds.lon_coarse
             lat = self.grid.ds.lat_coarse
             angle = self.grid.ds.angle_coarse
@@ -132,9 +123,9 @@ class ROMSToolsMixins:
         if vars_3d:
             # 3d interpolation
             coords = {
+                data.dim_names["depth"]: self.grid.ds["layer_depth_rho"],
                 data.dim_names["latitude"]: lat,
                 data.dim_names["longitude"]: lon,
-                data.dim_names["depth"]: self.vertical_coordinate.ds["layer_depth_rho"],
             }
         # extrapolate deepest value all the way to bottom ("flooding")
         for var in vars_3d:
@@ -158,37 +149,50 @@ class ROMSToolsMixins:
             if data.dim_names["time"] != "time":
                 data_vars[var] = data_vars[var].rename({data.dim_names["time"]: "time"})
 
+            # transpose to correct order (time, s_rho, eta_rho, xi_rho)
+            data_vars[var] = data_vars[var].transpose(
+                "time", "s_rho", "eta_rho", "xi_rho"
+            )
+
         return data_vars
 
-    def process_velocities(self, data_vars, angle, interpolate=True):
+    def process_velocities(self, data_vars, angle, uname, vname, interpolate=True):
         """
-        Processes and rotates velocity components, and interpolates them to the appropriate grid points.
+        Process and rotate velocity components to align with the grid orientation and optionally interpolate
+        them to the appropriate grid points.
 
         This method performs the following steps:
-        1. Rotates the velocity components to align with the grid orientation using the provided angle.
-        2. Optionally interpolates the rotated velocities to the u- and v-points of the grid.
-        3. If a vertical coordinate is provided, computes the barotropic velocities by integrating
-           over the vertical dimension.
+
+        1. **Rotation**: Rotates the velocity components (e.g., `u`, `v`) to align with the grid orientation
+           using the provided angle data.
+        2. **Interpolation**: Optionally interpolates the rotated velocities from rho-points to u- and v-points
+           of the grid.
+        3. **Barotropic Velocity Calculation**: If the velocity components are 3D (with vertical coordinates),
+           computes the barotropic (depth-averaged) velocities.
 
         Parameters
         ----------
         data_vars : dict of str: xarray.DataArray
-            Dictionary containing the velocity components to be processed. Must include keys "u" and "v"
-            or "uwnd" and "vwnd".
+            Dictionary containing the velocity components to be processed. The dictionary should include keys
+            corresponding to the velocity component names (e.g., `uname`, `vname`).
         angle : xarray.DataArray
-            DataArray containing the angle used for rotating the velocity components to the grid orientation.
+            DataArray containing the grid angle values used to rotate the velocity components to the correct
+            orientation on the grid.
+        uname : str
+            The key corresponding to the zonal (east-west) velocity component in `data_vars`.
+        vname : str
+            The key corresponding to the meridional (north-south) velocity component in `data_vars`.
         interpolate : bool, optional
-            If True, interpolates the velocities to the u- and v-points. Defaults to True.
+            If True, interpolates the rotated velocity components to the u- and v-points of the grid.
+            Defaults to True.
 
         Returns
         -------
         dict of str: xarray.DataArray
-            Dictionary of processed velocity components. Includes "ubar" and "vbar" if a vertical coordinate
-            is provided.
+            A dictionary of the processed velocity components. The returned dictionary includes the rotated and,
+            if applicable, interpolated velocity components. If the input velocities are 3D (having a vertical
+            dimension), the dictionary also includes the barotropic (depth-averaged) velocities (`ubar` and `vbar`).
         """
-        # Determine the correct variable names based on the keys in data_vars
-        uname = "u" if "u" in data_vars else "uwnd"
-        vname = "v" if "v" in data_vars else "vwnd"
 
         # Rotate velocities to grid orientation
         u_rot = data_vars[uname] * np.cos(angle) + data_vars[vname] * np.sin(angle)
@@ -202,7 +206,7 @@ class ROMSToolsMixins:
             data_vars[uname] = u_rot
             data_vars[vname] = v_rot
 
-        if self.vertical_coordinate is not None:
+        if "s_rho" in data_vars[uname].dims and "s_rho" in data_vars[vname].dims:
             # 3D masks for ROMS domain
             umask = self.grid.ds.mask_u.expand_dims({"s_rho": data_vars[uname].s_rho})
             vmask = self.grid.ds.mask_v.expand_dims({"s_rho": data_vars[vname].s_rho})
@@ -211,7 +215,7 @@ class ROMSToolsMixins:
             data_vars[vname] = data_vars[vname] * vmask
 
             # Compute barotropic velocity
-            dz = -self.vertical_coordinate.ds["interface_depth_rho"].diff(dim="s_w")
+            dz = -self.grid.ds["interface_depth_rho"].diff(dim="s_w")
             dz = dz.rename({"s_w": "s_rho"})
             dzu = interpolate_from_rho_to_u(dz)
             dzv = interpolate_from_rho_to_v(dz)
@@ -240,6 +244,26 @@ class ROMSToolsMixins:
         """
 
         d = {
+            "ssh_Re": {"long_name": "Tidal elevation, real part", "units": "m"},
+            "ssh_Im": {"long_name": "Tidal elevation, complex part", "units": "m"},
+            "pot_Re": {"long_name": "Tidal potential, real part", "units": "m"},
+            "pot_Im": {"long_name": "Tidal potential, complex part", "units": "m"},
+            "u_Re": {
+                "long_name": "Tidal velocity in x-direction, real part",
+                "units": "m/s",
+            },
+            "u_Im": {
+                "long_name": "Tidal velocity in x-direction, complex part",
+                "units": "m/s",
+            },
+            "v_Re": {
+                "long_name": "Tidal velocity in y-direction, real part",
+                "units": "m/s",
+            },
+            "v_Im": {
+                "long_name": "Tidal velocity in y-direction, complex part",
+                "units": "m/s",
+            },
             "uwnd": {"long_name": "10 meter wind in x-direction", "units": "m/s"},
             "vwnd": {"long_name": "10 meter wind in y-direction", "units": "m/s"},
             "swrad": {
@@ -333,18 +357,18 @@ class ROMSToolsMixins:
         return d
 
     def get_boundary_info(self):
-        """
-        Provides boundary coordinate information and renaming conventions for grid boundaries.
 
-        This method returns two dictionaries: one specifying the boundary coordinates for different types of
-        grid variables (e.g., "rho", "u", "v"), and another specifying how to rename dimensions for these boundaries.
+        """
+        This method provides information about the boundary points for the rho, u, and v
+        variables on the grid, specifying the indices for the south, east, north, and west
+        boundaries.
 
         Returns
         -------
-        tuple of (dict, dict)
-            - A dictionary mapping variable types and directions to boundary coordinates.
-            - A dictionary mapping variable types and directions to new dimension names.
-
+        dict
+            A dictionary where keys are variable types ("rho", "u", "v"), and values
+            are nested dictionaries mapping directions ("south", "east", "north", "west")
+            to the corresponding boundary coordinates.
         """
 
         # Boundary coordinates
@@ -369,27 +393,4 @@ class ROMSToolsMixins:
             },
         }
 
-        # How to rename the dimensions
-
-        rename = {
-            "rho": {
-                "south": {"xi_rho": "xi_rho_south"},
-                "east": {"eta_rho": "eta_rho_east"},
-                "north": {"xi_rho": "xi_rho_north"},
-                "west": {"eta_rho": "eta_rho_west"},
-            },
-            "u": {
-                "south": {"xi_u": "xi_u_south"},
-                "east": {"eta_rho": "eta_u_east"},
-                "north": {"xi_u": "xi_u_north"},
-                "west": {"eta_rho": "eta_u_west"},
-            },
-            "v": {
-                "south": {"xi_rho": "xi_v_south"},
-                "east": {"eta_v": "eta_v_east"},
-                "north": {"xi_rho": "xi_v_north"},
-                "west": {"eta_v": "eta_v_west"},
-            },
-        }
-
-        return bdry_coords, rename
+        return bdry_coords

roms_tools/setup/plot.py

@@ -8,7 +8,6 @@ def _plot(
     field=None,
     depth_contours=False,
     straddle=False,
-    coarse_grid=False,
     c="red",
     title="",
     kwargs={},
@@ -21,29 +20,8 @@ def _plot(
     else:
 
         field = field.squeeze()
-
-        if coarse_grid:
-
-            field = field.rename({"eta_rho": "eta_coarse", "xi_rho": "xi_coarse"})
-            field = field.where(grid_ds.mask_coarse)
-            lon_deg = field.lon
-            lat_deg = field.lat
-
-        else:
-            if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-                field = field.where(grid_ds.mask_rho)
-                lon_deg = grid_ds["lon_rho"]
-                lat_deg = grid_ds["lat_rho"]
-            elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-                field = field.where(grid_ds.mask_u)
-                lon_deg = grid_ds["lon_u"]
-                lat_deg = grid_ds["lat_u"]
-            elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-                field = field.where(grid_ds.mask_v)
-                lon_deg = grid_ds["lon_v"]
-                lat_deg = grid_ds["lat_v"]
-            else:
-                ValueError("provided field does not have two horizontal dimension")
+        lon_deg = field.lon
+        lat_deg = field.lat
 
         # check if North or South pole are in domain
         if lat_deg.max().values > 89 or lat_deg.min().values < -89:
@@ -96,23 +74,7 @@ def _plot(
         plt.colorbar(p, label=f"{field.long_name} [{field.units}]")
 
     if depth_contours:
-        if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            if "layer_depth_rho" in field.coords:
-                depth = field.layer_depth_rho
-            else:
-                depth = field.interface_depth_rho
-        elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            if "layer_depth_u" in field.coords:
-                depth = field.layer_depth_u
-            else:
-                depth = field.interface_depth_u
-        elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            if "layer_depth_v" in field.coords:
-                depth = field.layer_depth_v
-            else:
-                depth = field.interface_depth_v
-
-        cs = ax.contour(lon_deg, lat_deg, depth, transform=proj, colors="k")
+        cs = ax.contour(lon_deg, lat_deg, field.layer_depth, transform=proj, colors="k")
         ax.clabel(cs, inline=True, fontsize=10)
 
     return fig
@@ -135,12 +97,8 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}):
         )
 
     depths_to_check = [
-        "layer_depth_rho",
-        "layer_depth_u",
-        "layer_depth_v",
-        "interface_depth_rho",
-        "interface_depth_u",
-        "interface_depth_v",
+        "layer_depth",
+        "interface_depth",
     ]
     try:
         depth_label = next(

roms_tools/setup/surface_forcing.py

@@ -95,7 +95,9 @@ class SurfaceForcing(ROMSToolsMixins):
         vars_2d = ["uwnd", "vwnd", "swrad", "lwrad", "Tair", "qair", "rain"]
         vars_3d = []
         data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
-        data_vars = super().process_velocities(data_vars, angle, interpolate=False)
+        data_vars = super().process_velocities(
+            data_vars, angle, "uwnd", "vwnd", interpolate=False
+        )
 
         if self.correct_radiation:
             correction_data = self._get_correction_data()
@@ -235,18 +237,19 @@ class SurfaceForcing(ROMSToolsMixins):
 
         if self.bgc_source["name"] == "CESM_REGRIDDED":
 
-            bgc_data = CESMBGCSurfaceForcingDataset(
+            data = CESMBGCSurfaceForcingDataset(
                 filename=self.bgc_source["path"],
                 start_time=self.start_time,
                 end_time=self.end_time,
                 climatology=self.bgc_source["climatology"],
             )
+            data.post_process()
         else:
             raise ValueError(
                 'Only "CESM_REGRIDDED" is a valid option for bgc_source["name"].'
             )
 
-        return bgc_data
+        return data
 
     def _write_into_dataset(self, data, d_meta):
 
@@ -259,7 +262,6 @@ class SurfaceForcing(ROMSToolsMixins):
             ds[var].attrs["units"] = d_meta[var]["units"]
 
         if self.use_coarse_grid:
-            ds = ds.assign_coords({"lon": self.target_lon, "lat": self.target_lat})
             ds = ds.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
 
         # Preserve absolute time coordinate for readability
@@ -295,6 +297,10 @@ class SurfaceForcing(ROMSToolsMixins):
         if data.climatology:
             ds["time"].attrs["cycle_length"] = 365.25
 
+        variables_to_drop = ["lat_rho", "lon_rho", "lat_coarse", "lon_coarse"]
+        existing_vars = [var for var in variables_to_drop if var in ds]
+        ds = ds.drop_vars(existing_vars)
+
         return ds
 
     def _write_into_datatree(self, data, bgc_data, d_meta):
@@ -392,6 +398,15 @@ class SurfaceForcing(ROMSToolsMixins):
         field = ds[varname].isel(time=time).load()
         title = field.long_name
 
+        # assign lat / lon
+        if self.use_coarse_grid:
+            field = field.rename({"eta_rho": "eta_coarse", "xi_rho": "xi_coarse"})
+            field = field.where(self.grid.ds.mask_coarse)
+        else:
+            field = field.where(self.grid.ds.mask_rho)
+
+        field = field.assign_coords({"lon": self.target_lon, "lat": self.target_lat})
+
         # choose colorbar
         if varname in ["uwnd", "vwnd"]:
             vmax = max(field.max().values, -field.min().values)
@@ -412,7 +427,6 @@ class SurfaceForcing(ROMSToolsMixins):
             self.grid.ds,
             field=field,
             straddle=self.grid.straddle,
-            coarse_grid=self.use_coarse_grid,
             title=title,
             kwargs=kwargs,
             c="g",