Add density-dependent interpolation of BGC tracers to ROMS-Tools BC and IC

docs/boundary_forcing.ipynb

@@ -3140,6 +3140,7 @@
     "    end_time=end_time,\n",
     "    source={\"name\": \"CESM_REGRIDDED\", \"path\": cesm_bgc_path, \"climatology\": True},\n",
     "    type=\"bgc\",\n",
+    "    physics_forcing=boundary_forcing,  # use density-space interpolation\n",
     "    use_dask=True,\n",
     ")"
    ]
@@ -11209,6 +11210,7 @@
     "    end_time=end_time,\n",
     "    source={\"name\": \"UNIFIED\", \"path\": unified_bgc_path, \"climatology\": True},\n",
     "    type=\"bgc\",\n",
+    "    physics_forcing=boundary_forcing,  # use density-space interpolation\n",
     "    use_dask=True,\n",
     ")"
    ]

docs/end_to_end.ipynb

@@ -11215,6 +11215,7 @@
     "    end_time=end_time,\n",
     "    source={\"name\": \"UNIFIED\", \"path\": unified_bgc_path, \"climatology\": True},\n",
     "    type=\"bgc\",\n",
+    "    physics_forcing=boundary_forcing,  # for density-space BGC interpolation\n",
     "    use_dask=True\n",
     ")"
    ]

docs/releases.md

@@ -8,6 +8,7 @@
 
 ### New Features
 
+* Density-space vertical interpolation for BGC tracers in `InitialConditions` and `BoundaryForcing` via new `use_density_interpolation` parameter (default `True`). For `BoundaryForcing`, pass the physics `BoundaryForcing` as `physics_forcing=` to activate ([#606](https://github.com/CWorthy-ocean/roms-tools/pull/606)).
 * `make_edata` changed to `make_nesting_info`
 * `to_yaml` and `from_yaml` were adjusted to handle child grids after they've been modified ([#573](https://github.com/CWorthy-ocean/roms-tools/pull/573))
 * Nesting now supports optional baroclinic pressure fluxes via metadata ([#568](https://github.com/CWorthy-ocean/roms-tools/pull/568))
@@ -41,6 +42,7 @@
 * Document `close_narrow_channels` option in `Grid` and `update_mask()`; update notebook examples
 * Both the surface forcing and datasets notebooks are updated to reflect `restoring` function and WOA data ([#589](https://github.com/CWorthy-ocean/roms-tools/pull/589))
 * The cdr notebook is updated to reflect interpolation option. Default is same as ROMS, no interpolation ([#601](https://github.com/CWorthy-ocean/roms-tools/pull/601))
+* Notebooks updated to document density-space BGC interpolation ([#606](https://github.com/CWorthy-ocean/roms-tools/pull/606))
 
 ### Bugfixes
 

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
     "bottleneck",
     "regionmask>=0.11.0",
     "xgcm>=0.9.0",
+    "gsw",
     "numba>=0.61.2",
     "pydantic>2,<3",
 ]

roms_tools/setup/boundary_forcing.py

@@ -22,10 +22,12 @@
 from roms_tools.regrid import LateralRegridToROMS, VerticalRegrid
 from roms_tools.setup.utils import (
     RawDataSource,
+    _compute_bgc_source_density,
     add_time_info_to_ds,
     check_and_set_boundaries,
     compute_barotropic_velocity,
     compute_missing_bgc_variables,
+    compute_potential_density,
     from_yaml,
     get_boundary_coords,
     get_target_coords,
@@ -37,6 +39,7 @@
     write_to_yaml,
 )
 from roms_tools.utils import (
+    interpolate_cyclic_time,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
     rotate_velocities,
@@ -46,6 +49,50 @@
 from roms_tools.vertical_coordinate import compute_depth
 
 
+def _interpolate_phys_to_bgc_time(
+    phys_da: xr.DataArray,
+    time_dim: str,
+    bgc_time_coord: xr.DataArray,
+    bgc_climatology: bool,
+) -> xr.DataArray:
+    """Linearly interpolate a physics DataArray onto the BGC time coordinate.
+
+    Parameters
+    ----------
+    phys_da : xr.DataArray
+        Physics data with a ``datetime64`` time dimension named ``time_dim``.
+    time_dim : str
+        Name of the time dimension in ``phys_da``.
+    bgc_time_coord : xr.DataArray
+        Target time coordinate from the BGC dataset (1-D).
+    bgc_climatology : bool
+        Whether the BGC dataset is a climatology. If True, ``bgc_time_coord``
+        is expected to be ``timedelta64`` from the start of the year (as set by
+        ``assign_dates_to_climatology``), and physics times are mapped to
+        fractional day-of-year before a cyclic linear interpolation. If False,
+        a straight ``xr.DataArray.interp(method="linear")`` is performed in
+        ``datetime64`` space.
+
+    Returns
+    -------
+    xr.DataArray
+        ``phys_da`` interpolated to ``bgc_time_coord``, with time dimension
+        still named ``time_dim`` and coordinate set to ``bgc_time_coord``.
+    """
+    if bgc_climatology:
+        # Map both axes to fractional day-of-year, then cyclic-interpolate.
+        bgc_doy = (bgc_time_coord / np.timedelta64(1, "D")).values + 1.0
+        phys_doy = phys_da[time_dim].dt.dayofyear.values.astype(float)
+        phys_for_interp = phys_da.assign_coords(
+            {time_dim: xr.DataArray(phys_doy, dims=[time_dim])}
+        )
+        result = interpolate_cyclic_time(phys_for_interp, time_dim, time_dim, bgc_doy)
+        return result.assign_coords({time_dim: bgc_time_coord.values})
+
+    # Non-climatology: standard datetime64 linear interpolation
+    return phys_da.interp({time_dim: bgc_time_coord}, method="linear")
+
+
 @dataclass(kw_only=True)
 class BoundaryForcing:
     """Represents boundary forcing input data for ROMS.
@@ -104,6 +151,18 @@ class BoundaryForcing:
         Indicates whether to skip validation checks in the processed data. When set to True,
         the validation process that ensures no NaN values exist at wet points
         in the processed dataset is bypassed. Defaults to False.
+    use_density_interpolation : bool, optional
+        If True (default), BGC tracers are vertically interpolated in density space
+        rather than depth space when ``physics_forcing`` is provided. Preserves
+        water-mass properties. Only used when ``type='bgc'``.
+
+        Interpolation uses ``xgcm.Grid.transform`` with the linear method inside
+        the source density range and edge-value (nearest-neighbor) extrapolation
+        outside (``mask_edges=False``).
+    physics_forcing : BoundaryForcing, optional
+        A physics ``BoundaryForcing`` object (``type='physics'``) whose T/S fields
+        supply the density coordinate for BGC tracer interpolation. When None and
+        ``use_density_interpolation=True``, falls back to depth-based interpolation.
 
 
     Examples
@@ -143,6 +202,10 @@ class BoundaryForcing:
     """Optional initial bounding slice when loading source data (Dask); see dataset classes."""
     bypass_validation: bool = False
     """Whether to skip validation checks in the processed data."""
+    use_density_interpolation: bool = True
+    """Whether to interpolate BGC tracers in density space rather than depth space."""
+    physics_forcing: "BoundaryForcing | None" = None
+    """Physics BoundaryForcing object supplying T/S for density-based BGC interpolation."""
 
     ds: xr.Dataset = field(init=False, repr=False)
     """An xarray Dataset containing post-processed variables ready for input into
@@ -156,9 +219,20 @@ def __post_init__(self):
         # Initialize depth coordinates
         self.adjust_depth_for_sea_surface_height = False
         self.ds_depth_coords = xr.Dataset()
+        self._phys_bdry_depth_data: dict = {}
 
         self._input_checks()
 
+        if (
+            self.type == "bgc"
+            and self.use_density_interpolation
+            and self.physics_forcing is None
+        ):
+            logging.info(
+                "use_density_interpolation=True but no physics_forcing provided. "
+                "BGC tracers will be interpolated in depth space instead."
+            )
+
         target_coords = get_target_coords(self.grid)
 
         data = self._get_data()
@@ -329,6 +403,16 @@ def __post_init__(self):
                     zeta_u = 0
                     zeta_v = 0
 
+                # Save physics T/S at source depth levels for density-based BGC interpolation
+                if self.type == "physics" and self.use_density_interpolation:
+                    if "temp" in processed_fields and "salt" in processed_fields:
+                        self._phys_bdry_depth_data[direction] = {
+                            "temp": processed_fields["temp"],
+                            "salt": processed_fields["salt"],
+                            "depth_dim": bdry_data.dim_names["depth"],
+                            "depth_coord": bdry_data.ds[bdry_data.dim_names["depth"]],
+                        }
+
                 for location in ["rho", "u", "v"]:
                     # Filter var_names by location and check for 3D variables
                     filtered_vars = [
@@ -353,17 +437,43 @@ def __post_init__(self):
                         vertical_regrid = VerticalRegrid(
                             bdry_data.ds, source_dim=bdry_data.dim_names["depth"]
                         )
+
+                        use_density = (
+                            self.type == "bgc"
+                            and self.use_density_interpolation
+                            and location == "rho"
+                            and self.physics_forcing is not None
+                            and direction in self.physics_forcing._phys_bdry_depth_data
+                        )
+
+                        if use_density:
+                            source_density, target_density = (
+                                self._compute_bgc_density_coords(
+                                    direction,
+                                    bdry_data,
+                                    processed_fields,
+                                    filtered_vars,
+                                )
+                            )
+
                         for var_name in filtered_vars:
                             if var_name in processed_fields:
-                                processed_fields[var_name] = vertical_regrid.apply(
-                                    processed_fields[var_name],
-                                    source_depth_coords=bdry_data.ds[
-                                        bdry_data.dim_names["depth"]
-                                    ],
-                                    target_depth_coords=self.ds_depth_coords[
-                                        f"layer_depth_{location}_{direction}"
-                                    ],
-                                )
+                                if use_density:
+                                    processed_fields[var_name] = vertical_regrid.apply(
+                                        processed_fields[var_name],
+                                        source_depth_coords=source_density,
+                                        target_depth_coords=target_density,
+                                    )
+                                else:
+                                    processed_fields[var_name] = vertical_regrid.apply(
+                                        processed_fields[var_name],
+                                        source_depth_coords=bdry_data.ds[
+                                            bdry_data.dim_names["depth"]
+                                        ],
+                                        target_depth_coords=self.ds_depth_coords[
+                                            f"layer_depth_{location}_{direction}"
+                                        ],
+                                    )
 
                 # compute barotropic velocities
                 if "u" in var_names and "v" in var_names:
@@ -403,6 +513,91 @@ def __post_init__(self):
 
         self.ds = ds
 
+    def _compute_bgc_density_coords(
+        self,
+        direction: str,
+        bdry_data,
+        processed_fields: dict,
+        filtered_vars: list,
+    ) -> tuple[xr.DataArray, xr.DataArray]:
+        """Build source/target density coordinates for BGC density-space interpolation.
+
+        Source density is computed from the physics BC's saved depth-level T/S and
+        regridded onto the BGC depth axis. Target density is computed from the
+        physics BC's sigma-level T/S, interpolated to the BGC time coordinate.
+
+        Both arrays include the small monotonicity perturbation that
+        ``vertical_regrid.apply`` requires for a strictly monotonic coordinate.
+        """
+        assert self.physics_forcing is not None
+        phys_data = self.physics_forcing._phys_bdry_depth_data[direction]
+        phys_temp, phys_salt = phys_data["temp"], phys_data["salt"]
+        bgc_climatology = bool(self.source["climatology"])
+
+        # The BGC time axis (if any) is shared across all 3D BGC vars; take it from
+        # the first one. dim_names["time"] gives the source-data convention name.
+        bgc_time_dim = bdry_data.dim_names.get("time")
+        first_bgc = processed_fields.get(filtered_vars[0]) if filtered_vars else None
+        bgc_time_coord = (
+            first_bgc[bgc_time_dim]
+            if first_bgc is not None
+            and bgc_time_dim in (first_bgc.dims if first_bgc is not None else ())
+            else None
+        )
+
+        def _align_time(da: xr.DataArray, time_dim: str) -> xr.DataArray:
+            """Align ``da``'s ``time_dim`` to the BGC time axis, or collapse it."""
+            if time_dim not in da.dims:
+                return da
+            if bgc_time_coord is not None:
+                return _interpolate_phys_to_bgc_time(
+                    da, time_dim, bgc_time_coord, bgc_climatology
+                )
+            return da.mean(time_dim)
+
+        # --- Source density: physics depth-level T/S → BGC depth axis ---
+        phys_temp = _align_time(phys_temp, "time")
+        phys_salt = _align_time(phys_salt, "time")
+        bgc_depth_dim = bdry_data.dim_names["depth"]
+        source_density = _compute_bgc_source_density(
+            phys_temp,
+            phys_salt,
+            phys_data["depth_dim"],
+            phys_data["depth_coord"],
+            bdry_data.ds[bgc_depth_dim],
+            bgc_depth_dim,
+        )
+
+        # --- Target density: physics sigma-level T/S, aligned to BGC time ---
+        # Physics BC dataset uses "bry_time" as the time dim with an "abs_time"
+        # datetime64 companion coord. Swap to the datetime view before time-aligning.
+        temp_sigma = self.physics_forcing.ds[f"temp_{direction}"]
+        salt_sigma = self.physics_forcing.ds[f"salt_{direction}"]
+        if "abs_time" in temp_sigma.coords:
+            temp_sigma = temp_sigma.swap_dims({"bry_time": "abs_time"}).rename(
+                {"abs_time": "time"}
+            )
+            salt_sigma = salt_sigma.swap_dims({"bry_time": "abs_time"}).rename(
+                {"abs_time": "time"}
+            )
+            temp_sigma = _align_time(temp_sigma, "time")
+            salt_sigma = _align_time(salt_sigma, "time")
+        else:
+            temp_sigma = _align_time(temp_sigma, "bry_time")
+            salt_sigma = _align_time(salt_sigma, "bry_time")
+
+        target_density = compute_potential_density(temp_sigma, salt_sigma)
+        # Add a small perturbation along the vertical dim to ensure monotonicity.
+        vertical_dim = next(d for d in target_density.dims if d.startswith("s_"))
+        n_s = target_density.sizes[vertical_dim]
+        target_density = target_density + xr.DataArray(
+            np.arange(n_s) * 1e-7, dims=[vertical_dim]
+        )
+        # xgcm.transform requires a single chunk along the target vertical dim.
+        target_density = target_density.chunk({vertical_dim: -1})
+
+        return source_density, target_density
+
     def _input_checks(self) -> None:
         """Validate and normalize user-provided input parameters."""
         # -------------------------------------------------------
@@ -1091,6 +1286,7 @@ def to_yaml(self, filepath: str | Path) -> None:
                 "ds_depth_coords",
                 "adjust_depth_for_sea_surface_height",
                 "use_dask",
+                "physics_forcing",
             ],
         )
         write_to_yaml(forcing_dict, filepath)