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

roms_tools/setup/boundary_forcing.py

@@ -1,19 +1,18 @@
 import xarray as xr
 import numpy as np
+import pandas as pd
 import yaml
+from datatree import DataTree
 import importlib.metadata
-from typing import Dict
+from typing import Dict, Union, Optional
 from dataclasses import dataclass, field, asdict
 from roms_tools.setup.grid import Grid
 from roms_tools.setup.vertical_coordinate import VerticalCoordinate
+from roms_tools.setup.mixins import ROMSToolsMixins
 from datetime import datetime
-from roms_tools.setup.datasets import Dataset
-from roms_tools.setup.fill import fill_and_interpolate
+from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
 from roms_tools.setup.utils import (
     nan_check,
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
-    extrapolate_deepest_to_bottom,
 )
 from roms_tools.setup.plot import _section_plot, _line_plot
 import calendar
@@ -22,7 +21,7 @@ import matplotlib.pyplot as plt
 
 
 @dataclass(frozen=True, kw_only=True)
-class BoundaryForcing:
+class BoundaryForcing(ROMSToolsMixins):
     """
     Represents boundary forcing for ROMS.
 
@@ -38,22 +37,39 @@ class BoundaryForcing:
         End time of the desired boundary forcing data.
     boundaries : Dict[str, bool], optional
         Dictionary specifying which boundaries are forced (south, east, north, west). Default is all True.
+    physics_source : Dict[str, Union[str, None]]
+        Dictionary specifying the source of the physical boundary forcing data:
+        - "name" (str): Name of the data source (e.g., "GLORYS").
+        - "path" (str): Path to the physical data file. Can contain wildcards.
+        - "climatology" (bool): Indicates if the physical data is climatology data. Defaults to False.
+    bgc_source : Optional[Dict[str, Union[str, None]]]
+        Dictionary specifying the source of the biogeochemical (BGC) initial condition data:
+        - "name" (str): Name of the BGC data source (e.g., "CESM_REGRIDDED").
+        - "path" (str): Path to the BGC data file. Can contain wildcards.
+        - "climatology" (bool): Indicates if the BGC data is climatology data. Defaults to True.
     model_reference_date : datetime, optional
         Reference date for the model. Default is January 1, 2000.
-    source : str, optional
-        Source of the boundary forcing data. Default is "glorys".
-    filename: str
-        Path to the source data file. Can contain wildcards.
 
     Attributes
     ----------
     ds : xr.Dataset
         Xarray Dataset containing the atmospheric forcing data.
 
-    Notes
-    -----
-    This class represents atmospheric forcing data used in ocean modeling. It provides a convenient
-    interface to work with forcing data including shortwave radiation correction and river forcing.
+    Examples
+    --------
+    >>> boundary_forcing = BoundaryForcing(
+    ...     grid=grid,
+    ...     vertical_coordinate=vertical_coordinate,
+    ...     boundaries={"south": True, "east": True, "north": False, "west": True},
+    ...     start_time=datetime(2022, 1, 1),
+    ...     end_time=datetime(2022, 1, 2),
+    ...     physics_source={"name": "GLORYS", "path": "physics_data.nc"},
+    ...     bgc_source={
+    ...         "name": "CESM_REGRIDDED",
+    ...         "path": "bgc_data.nc",
+    ...         "climatology": True,
+    ...     },
+    ... )
     """
 
     grid: Grid
@@ -68,50 +84,18 @@ class BoundaryForcing:
             "west": True,
         }
     )
+    physics_source: Dict[str, Union[str, None]]
+    bgc_source: Optional[Dict[str, Union[str, None]]] = None
     model_reference_date: datetime = datetime(2000, 1, 1)
-    source: str = "glorys"
-    filename: str
+
     ds: xr.Dataset = field(init=False, repr=False)
 
     def __post_init__(self):
 
-        lon = self.grid.ds.lon_rho
-        lat = self.grid.ds.lat_rho
-        angle = self.grid.ds.angle
-
-        if self.source == "glorys":
-            dims = {
-                "longitude": "longitude",
-                "latitude": "latitude",
-                "depth": "depth",
-                "time": "time",
-            }
-
-            varnames = {
-                "temp": "thetao",
-                "salt": "so",
-                "u": "uo",
-                "v": "vo",
-                "ssh": "zos",
-            }
-        data = Dataset(
-            filename=self.filename,
-            start_time=self.start_time,
-            end_time=self.end_time,
-            var_names=varnames.values(),
-            dim_names=dims,
-        )
-
-        # operate on longitudes between -180 and 180 unless ROMS domain lies at least 5 degrees in lontitude away from Greenwich meridian
-        lon = xr.where(lon > 180, lon - 360, lon)
-        straddle = True
-        if not self.grid.straddle and abs(lon).min() > 5:
-            lon = xr.where(lon < 0, lon + 360, lon)
-            straddle = False
+        self._input_checks()
+        lon, lat, angle, straddle = super().get_target_lon_lat()
 
-        # Restrict data to relevant subdomain to achieve better performance and to avoid discontinuous longitudes introduced by converting
-        # to a different longitude range (+- 360 degrees). Discontinues longitudes can lead to artifacts in the interpolation process that
-        # would not be detected by the nan_check function.
+        data = self._get_data()
         data.choose_subdomain(
             latitude_range=[lat.min().values, lat.max().values],
             longitude_range=[lon.min().values, lon.max().values],
@@ -119,252 +103,278 @@ class BoundaryForcing:
             straddle=straddle,
         )
 
-        # extrapolate deepest value all the way to bottom ("flooding") to prepare for 3d interpolation
-        for var in ["temp", "salt", "u", "v"]:
-            data.ds[varnames[var]] = extrapolate_deepest_to_bottom(
-                data.ds[varnames[var]], dims["depth"]
+        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)
+        object.__setattr__(data, "data_vars", data_vars)
+
+        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,
             )
 
-        # interpolate onto desired grid
-        fill_dims = [dims["latitude"], dims["longitude"]]
+            vars_2d = []
+            vars_3d = bgc_data.var_names.keys()
+            data_vars = super().regrid_data(bgc_data, vars_2d, vars_3d, lon, lat)
+            object.__setattr__(bgc_data, "data_vars", data_vars)
+        else:
+            bgc_data = None
 
-        # 2d interpolation
-        coords = {dims["latitude"]: lat, dims["longitude"]: lon}
-        mask = xr.where(data.ds[varnames["ssh"]].isel(time=0).isnull(), 0, 1)
+        d_meta = super().get_variable_metadata()
+        bdry_coords, rename = super().get_boundary_info()
 
-        ssh = fill_and_interpolate(
-            data.ds[varnames["ssh"]].astype(np.float64),
-            mask,
-            fill_dims=fill_dims,
-            coords=coords,
-            method="linear",
-        )
+        ds = self._write_into_datatree(data, bgc_data, d_meta, bdry_coords, rename)
 
-        # 3d interpolation
-        coords = {
-            dims["latitude"]: lat,
-            dims["longitude"]: lon,
-            dims["depth"]: self.vertical_coordinate.ds["layer_depth_rho"],
-        }
-        mask = xr.where(data.ds[varnames["temp"]].isel(time=0).isnull(), 0, 1)
-
-        data_vars = {}
-        # setting fillvalue_interp to None means that we allow extrapolation in the
-        # interpolation step to avoid NaNs at the surface if the lowest depth in original
-        # data is greater than zero
-
-        for var in ["temp", "salt", "u", "v"]:
-
-            data_vars[var] = fill_and_interpolate(
-                data.ds[varnames[var]].astype(np.float64),
-                mask,
-                fill_dims=fill_dims,
-                coords=coords,
-                method="linear",
-                fillvalue_interp=None,
-            )
+        for direction in ["south", "east", "north", "west"]:
+            if self.boundaries[direction]:
+                nan_check(
+                    ds["physics"][f"zeta_{direction}"].isel(bry_time=0),
+                    self.grid.ds.mask_rho.isel(**bdry_coords["rho"][direction]),
+                )
 
-        # rotate velocities to grid orientation
-        u_rot = data_vars["u"] * np.cos(angle) + data_vars["v"] * np.sin(angle)
-        v_rot = data_vars["v"] * np.cos(angle) - data_vars["u"] * np.sin(angle)
+        object.__setattr__(self, "ds", ds)
 
-        # interpolate to u- and v-points
-        u = interpolate_from_rho_to_u(u_rot)
-        v = interpolate_from_rho_to_v(v_rot)
+    def _input_checks(self):
+
+        if "name" not in self.physics_source.keys():
+            raise ValueError("`physics_source` must include a 'name'.")
+        if "path" not in self.physics_source.keys():
+            raise ValueError("`physics_source` must include a 'path'.")
+        # set self.physics_source["climatology"] to False if not provided
+        object.__setattr__(
+            self,
+            "physics_source",
+            {
+                **self.physics_source,
+                "climatology": self.physics_source.get("climatology", False),
+            },
+        )
+        if self.bgc_source is not None:
+            if "name" not in self.bgc_source.keys():
+                raise ValueError(
+                    "`bgc_source` must include a 'name' if it is provided."
+                )
+            if "path" not in self.bgc_source.keys():
+                raise ValueError(
+                    "`bgc_source` must include a 'path' if it is provided."
+                )
+            # set self.bgc_source["climatology"] to True if not provided
+            object.__setattr__(
+                self,
+                "bgc_source",
+                {
+                    **self.bgc_source,
+                    "climatology": self.bgc_source.get("climatology", True),
+                },
+            )
 
-        # 3d masks for ROMS domain
-        umask = self.grid.ds.mask_u.expand_dims({"s_rho": u.s_rho})
-        vmask = self.grid.ds.mask_v.expand_dims({"s_rho": v.s_rho})
+    def _get_data(self):
 
-        u = u * umask
-        v = v * vmask
+        if self.physics_source["name"] == "GLORYS":
+            data = GLORYSDataset(
+                filename=self.physics_source["path"],
+                start_time=self.start_time,
+                end_time=self.end_time,
+                climatology=self.physics_source["climatology"],
+            )
+        else:
+            raise ValueError(
+                'Only "GLORYS" is a valid option for physics_source["name"].'
+            )
 
-        # Compute barotropic velocity
+        return data
 
-        # thicknesses
-        dz = -self.vertical_coordinate.ds["interface_depth_rho"].diff(dim="s_w")
-        dz = dz.rename({"s_w": "s_rho"})
-        # thicknesses at u- and v-points
-        dzu = interpolate_from_rho_to_u(dz)
-        dzv = interpolate_from_rho_to_v(dz)
+    def _get_bgc_data(self):
 
-        ubar = (dzu * u).sum(dim="s_rho") / dzu.sum(dim="s_rho")
-        vbar = (dzv * v).sum(dim="s_rho") / dzv.sum(dim="s_rho")
+        if self.bgc_source["name"] == "CESM_REGRIDDED":
 
-        # Boundary coordinates for rho-points
-        bdry_coords_rho = {
-            "south": {"eta_rho": 0},
-            "east": {"xi_rho": -1},
-            "north": {"eta_rho": -1},
-            "west": {"xi_rho": 0},
-        }
-        # How to rename the dimensions at rho-points
-        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"},
-        }
+            data = CESMBGCDataset(
+                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"].'
+            )
 
-        # Boundary coordinates for u-points
-        bdry_coords_u = {
-            "south": {"eta_rho": 0},
-            "east": {"xi_u": -1},
-            "north": {"eta_rho": -1},
-            "west": {"xi_u": 0},
-        }
-        # How to rename the dimensions at u-points
-        rename_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"},
-        }
+        return data
 
-        # Boundary coordinates for v-points
-        bdry_coords_v = {
-            "south": {"eta_v": 0},
-            "east": {"xi_rho": -1},
-            "north": {"eta_v": -1},
-            "west": {"xi_rho": 0},
-        }
-        # How to rename the dimensions at v-points
-        rename_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"},
-        }
+    def _write_into_dataset(self, data, d_meta, bdry_coords, rename):
 
+        # save in new dataset
         ds = xr.Dataset()
 
         for direction in ["south", "east", "north", "west"]:
-
             if self.boundaries[direction]:
 
-                ds[f"zeta_{direction}"] = (
-                    ssh.isel(**bdry_coords_rho[direction])
-                    .rename(**rename_rho[direction])
-                    .astype(np.float32)
-                )
-                ds[f"zeta_{direction}"].attrs[
-                    "long_name"
-                ] = f"{direction}ern boundary sea surface height"
-                ds[f"zeta_{direction}"].attrs["units"] = "m"
-
-                ds[f"temp_{direction}"] = (
-                    data_vars["temp"]
-                    .isel(**bdry_coords_rho[direction])
-                    .rename(**rename_rho[direction])
-                    .astype(np.float32)
-                )
-                ds[f"temp_{direction}"].attrs[
-                    "long_name"
-                ] = f"{direction}ern boundary potential temperature"
-                ds[f"temp_{direction}"].attrs["units"] = "Celsius"
-
-                ds[f"salt_{direction}"] = (
-                    data_vars["salt"]
-                    .isel(**bdry_coords_rho[direction])
-                    .rename(**rename_rho[direction])
-                    .astype(np.float32)
-                )
-                ds[f"salt_{direction}"].attrs[
-                    "long_name"
-                ] = f"{direction}ern boundary salinity"
-                ds[f"salt_{direction}"].attrs["units"] = "PSU"
-
-                ds[f"u_{direction}"] = (
-                    u.isel(**bdry_coords_u[direction])
-                    .rename(**rename_u[direction])
-                    .astype(np.float32)
-                )
-                ds[f"u_{direction}"].attrs[
-                    "long_name"
-                ] = f"{direction}ern boundary u-flux component"
-                ds[f"u_{direction}"].attrs["units"] = "m/s"
-
-                ds[f"v_{direction}"] = (
-                    v.isel(**bdry_coords_v[direction])
-                    .rename(**rename_v[direction])
-                    .astype(np.float32)
-                )
-                ds[f"v_{direction}"].attrs[
-                    "long_name"
-                ] = f"{direction}ern boundary v-flux component"
-                ds[f"v_{direction}"].attrs["units"] = "m/s"
-
-                ds[f"ubar_{direction}"] = (
-                    ubar.isel(**bdry_coords_u[direction])
-                    .rename(**rename_u[direction])
-                    .astype(np.float32)
-                )
-                ds[f"ubar_{direction}"].attrs[
-                    "long_name"
-                ] = f"{direction}ern boundary vertically integrated u-flux component"
-                ds[f"ubar_{direction}"].attrs["units"] = "m/s"
-
-                ds[f"vbar_{direction}"] = (
-                    vbar.isel(**bdry_coords_v[direction])
-                    .rename(**rename_v[direction])
-                    .astype(np.float32)
-                )
-                ds[f"vbar_{direction}"].attrs[
-                    "long_name"
-                ] = f"{direction}ern boundary vertically integrated v-flux component"
-                ds[f"vbar_{direction}"].attrs["units"] = "m/s"
+                for var in data.data_vars.keys():
+                    if var in ["u", "ubar"]:
+                        ds[f"{var}_{direction}"] = (
+                            data.data_vars[var]
+                            .isel(**bdry_coords["u"][direction])
+                            .rename(**rename["u"][direction])
+                            .astype(np.float32)
+                        )
+                    elif var in ["v", "vbar"]:
+                        ds[f"{var}_{direction}"] = (
+                            data.data_vars[var]
+                            .isel(**bdry_coords["v"][direction])
+                            .rename(**rename["v"][direction])
+                            .astype(np.float32)
+                        )
+                    else:
+                        ds[f"{var}_{direction}"] = (
+                            data.data_vars[var]
+                            .isel(**bdry_coords["rho"][direction])
+                            .rename(**rename["rho"][direction])
+                            .astype(np.float32)
+                        )
+                    ds[f"{var}_{direction}"].attrs[
+                        "long_name"
+                    ] = f"{direction}ern boundary {d_meta[var]['long_name']}"
+                    ds[f"{var}_{direction}"].attrs["units"] = d_meta[var]["units"]
+
+        # Gracefully handle dropping variables that might not be present
+        variables_to_drop = [
+            "s_rho",
+            "layer_depth_rho",
+            "layer_depth_u",
+            "layer_depth_v",
+            "interface_depth_rho",
+            "interface_depth_u",
+            "interface_depth_v",
+            "lat_rho",
+            "lon_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)
+
+        # Preserve absolute time coordinate for readability
+        ds = ds.assign_coords({"abs_time": ds["time"]})
+
+        # Convert the time coordinate to the format expected by ROMS
+        if data.climatology:
+            # Convert to pandas TimedeltaIndex
+            timedelta_index = pd.to_timedelta(ds["time"].values)
+            # Determine the start of the year for the base_datetime
+            start_of_year = datetime(self.model_reference_date.year, 1, 1)
+            # Calculate the offset from midnight of the new year
+            offset = self.model_reference_date - start_of_year
+            bry_time = xr.DataArray(
+                timedelta_index - offset,
+                dims="time",
+            )
+        else:
+            # TODO: Check if we need to convert from 12:00:00 to 00:00:00 as in matlab scripts
+            bry_time = ds["time"] - np.datetime64(self.model_reference_date)
+
+        ds = ds.assign_coords({"bry_time": bry_time})
+        ds["bry_time"].attrs[
+            "long_name"
+        ] = f"nanoseconds since {np.datetime_as_string(np.datetime64(self.model_reference_date), unit='ns')}"
+        ds["bry_time"].encoding["units"] = "nanoseconds"
+        ds = ds.swap_dims({"time": "bry_time"})
+        ds = ds.drop_vars("time")
+
+        if data.climatology:
+            ds["bry_time"].attrs["cycle_length"] = 365.25
+
+        return ds
+
+    def _write_into_datatree(self, data, bgc_data, d_meta, bdry_coords, rename):
+
+        ds = self._add_global_metadata()
+        ds["sc_r"] = self.vertical_coordinate.ds["sc_r"]
+        ds["Cs_r"] = self.vertical_coordinate.ds["Cs_r"]
+
+        ds = DataTree(name="root", data=ds)
+
+        ds_physics = self._write_into_dataset(data, d_meta, bdry_coords, rename)
+        ds_physics = self._add_coordinates(bdry_coords, rename, ds_physics)
+        ds_physics = self._add_global_metadata(ds_physics)
+        ds_physics.attrs["physics_source"] = self.physics_source["name"]
+
+        ds_physics = DataTree(name="physics", parent=ds, data=ds_physics)
+
+        if bgc_data:
+            ds_bgc = self._write_into_dataset(bgc_data, d_meta, bdry_coords, rename)
+            ds_bgc = self._add_coordinates(bdry_coords, rename, ds_bgc)
+            ds_bgc = self._add_global_metadata(ds_bgc)
+            ds_bgc.attrs["bgc_source"] = self.bgc_source["name"]
+            ds_bgc = DataTree(name="bgc", parent=ds, data=ds_bgc)
 
-                # assign the correct depth coordinates
+        return ds
+
+    def _add_coordinates(self, bdry_coords, rename, ds=None):
+
+        if ds is None:
+            ds = xr.Dataset()
+
+        for direction in ["south", "east", "north", "west"]:
+
+            if self.boundaries[direction]:
 
                 lat_rho = self.grid.ds.lat_rho.isel(
-                    **bdry_coords_rho[direction]
-                ).rename(**rename_rho[direction])
+                    **bdry_coords["rho"][direction]
+                ).rename(**rename["rho"][direction])
                 lon_rho = self.grid.ds.lon_rho.isel(
-                    **bdry_coords_rho[direction]
-                ).rename(**rename_rho[direction])
+                    **bdry_coords["rho"][direction]
+                ).rename(**rename["rho"][direction])
                 layer_depth_rho = (
                     self.vertical_coordinate.ds["layer_depth_rho"]
-                    .isel(**bdry_coords_rho[direction])
-                    .rename(**rename_rho[direction])
+                    .isel(**bdry_coords["rho"][direction])
+                    .rename(**rename["rho"][direction])
                 )
                 interface_depth_rho = (
                     self.vertical_coordinate.ds["interface_depth_rho"]
-                    .isel(**bdry_coords_rho[direction])
-                    .rename(**rename_rho[direction])
+                    .isel(**bdry_coords["rho"][direction])
+                    .rename(**rename["rho"][direction])
                 )
 
-                lat_u = self.grid.ds.lat_u.isel(**bdry_coords_u[direction]).rename(
-                    **rename_u[direction]
+                lat_u = self.grid.ds.lat_u.isel(**bdry_coords["u"][direction]).rename(
+                    **rename["u"][direction]
                 )
-                lon_u = self.grid.ds.lon_u.isel(**bdry_coords_u[direction]).rename(
-                    **rename_u[direction]
+                lon_u = self.grid.ds.lon_u.isel(**bdry_coords["u"][direction]).rename(
+                    **rename["u"][direction]
                 )
                 layer_depth_u = (
                     self.vertical_coordinate.ds["layer_depth_u"]
-                    .isel(**bdry_coords_u[direction])
-                    .rename(**rename_u[direction])
+                    .isel(**bdry_coords["u"][direction])
+                    .rename(**rename["u"][direction])
                 )
                 interface_depth_u = (
                     self.vertical_coordinate.ds["interface_depth_u"]
-                    .isel(**bdry_coords_u[direction])
-                    .rename(**rename_u[direction])
+                    .isel(**bdry_coords["u"][direction])
+                    .rename(**rename["u"][direction])
                 )
 
-                lat_v = self.grid.ds.lat_v.isel(**bdry_coords_v[direction]).rename(
-                    **rename_v[direction]
+                lat_v = self.grid.ds.lat_v.isel(**bdry_coords["v"][direction]).rename(
+                    **rename["v"][direction]
                 )
-                lon_v = self.grid.ds.lon_v.isel(**bdry_coords_v[direction]).rename(
-                    **rename_v[direction]
+                lon_v = self.grid.ds.lon_v.isel(**bdry_coords["v"][direction]).rename(
+                    **rename["v"][direction]
                 )
                 layer_depth_v = (
                     self.vertical_coordinate.ds["layer_depth_v"]
-                    .isel(**bdry_coords_v[direction])
-                    .rename(**rename_v[direction])
+                    .isel(**bdry_coords["v"][direction])
+                    .rename(**rename["v"][direction])
                 )
                 interface_depth_v = (
                     self.vertical_coordinate.ds["interface_depth_v"]
-                    .isel(**bdry_coords_v[direction])
-                    .rename(**rename_v[direction])
+                    .isel(**bdry_coords["v"][direction])
+                    .rename(**rename["v"][direction])
                 )
 
                 ds = ds.assign_coords(
@@ -384,46 +394,31 @@ class BoundaryForcing:
                     }
                 )
 
-                ds = ds.drop_vars(
-                    [
-                        "layer_depth_rho",
-                        "layer_depth_u",
-                        "layer_depth_v",
-                        "interface_depth_rho",
-                        "interface_depth_u",
-                        "interface_depth_v",
-                        "lat_rho",
-                        "lon_rho",
-                        "lat_u",
-                        "lon_u",
-                        "lat_v",
-                        "lon_v",
-                        "s_rho",
-                    ]
-                )
-
-        # deal with time dimension
-        if dims["time"] != "time":
-            ds = ds.rename({dims["time"]: "time"})
-
-        # Translate the time coordinate to days since the model reference date
-        # TODO: Check if we need to convert from 12:00:00 to 00:00:00 as in matlab scripts
-        model_reference_date = np.datetime64(self.model_reference_date)
-
-        # Convert the time coordinate to the format expected by ROMS (days since model reference date)
-        bry_time = ds["time"] - model_reference_date
-        ds = ds.assign_coords(bry_time=("time", bry_time.data))
-        ds["bry_time"].attrs[
-            "long_name"
-        ] = f"time since {np.datetime_as_string(model_reference_date, unit='D')}"
-
-        ds["theta_s"] = self.vertical_coordinate.ds["theta_s"]
-        ds["theta_b"] = self.vertical_coordinate.ds["theta_b"]
-        ds["Tcline"] = self.vertical_coordinate.ds["Tcline"]
-        ds["hc"] = self.vertical_coordinate.ds["hc"]
-        ds["sc_r"] = self.vertical_coordinate.ds["sc_r"]
-        ds["Cs_r"] = self.vertical_coordinate.ds["Cs_r"]
-
+        # Gracefully handle dropping variables that might not be present
+        variables_to_drop = [
+            "s_rho",
+            "layer_depth_rho",
+            "layer_depth_u",
+            "layer_depth_v",
+            "interface_depth_rho",
+            "interface_depth_u",
+            "interface_depth_v",
+            "lat_rho",
+            "lon_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)
+
+        return ds
+
+    def _add_global_metadata(self, ds=None):
+
+        if ds is None:
+            ds = xr.Dataset()
         ds.attrs["title"] = "ROMS boundary forcing file created by ROMS-Tools"
         # Include the version of roms-tools
         try:
@@ -434,15 +429,13 @@ class BoundaryForcing:
         ds.attrs["start_time"] = str(self.start_time)
         ds.attrs["end_time"] = str(self.end_time)
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
-        ds.attrs["source"] = self.source
 
-        object.__setattr__(self, "ds", ds)
+        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()
 
-        for direction in ["south", "east", "north", "west"]:
-            nan_check(
-                ds[f"zeta_{direction}"].isel(time=0),
-                self.grid.ds.mask_rho.isel(**bdry_coords_rho[direction]),
-            )
+        return ds
 
     def plot(
         self,
@@ -457,14 +450,45 @@ class BoundaryForcing:
         ----------
         varname : str
             The name of the initial conditions field to plot. Options include:
-            - "temp_{direction}": Potential temperature.
-            - "salt_{direction}": Salinity.
-            - "zeta_{direction}": Sea surface height.
-            - "u_{direction}": u-flux component.
-            - "v_{direction}": v-flux component.
-            - "ubar_{direction}": Vertically integrated u-flux component.
-            - "vbar_{direction}": Vertically integrated v-flux component.
-            where {direction} can be one of ["south", "east", "north", "west"].
+            - "temp_{direction}": Potential temperature, where {direction} can be one of ["south", "east", "north", "west"].
+            - "salt_{direction}": Salinity, where {direction} can be one of ["south", "east", "north", "west"].
+            - "zeta_{direction}": Sea surface height, where {direction} can be one of ["south", "east", "north", "west"].
+            - "u_{direction}": u-flux component, where {direction} can be one of ["south", "east", "north", "west"].
+            - "v_{direction}": v-flux component, where {direction} can be one of ["south", "east", "north", "west"].
+            - "ubar_{direction}": Vertically integrated u-flux component, where {direction} can be one of ["south", "east", "north", "west"].
+            - "vbar_{direction}": Vertically integrated v-flux component, where {direction} can be one of ["south", "east", "north", "west"].
+            - "PO4_{direction}": Dissolved Inorganic Phosphate (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "NO3_{direction}": Dissolved Inorganic Nitrate (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "SiO3_{direction}": Dissolved Inorganic Silicate (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "NH4_{direction}": Dissolved Ammonia (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "Fe_{direction}": Dissolved Inorganic Iron (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "Lig_{direction}": Iron Binding Ligand (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "O2_{direction}": Dissolved Oxygen (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "DIC_{direction}": Dissolved Inorganic Carbon (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "DIC_ALT_CO2_{direction}": Dissolved Inorganic Carbon, Alternative CO2 (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "ALK_{direction}": Alkalinity (meq/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "ALK_ALT_CO2_{direction}": Alkalinity, Alternative CO2 (meq/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "DOC_{direction}": Dissolved Organic Carbon (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "DON_{direction}": Dissolved Organic Nitrogen (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "DOP_{direction}": Dissolved Organic Phosphorus (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "DOPr_{direction}": Refractory Dissolved Organic Phosphorus (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "DONr_{direction}": Refractory Dissolved Organic Nitrogen (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "DOCr_{direction}": Refractory Dissolved Organic Carbon (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "zooC_{direction}": Zooplankton Carbon (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "spChl_{direction}": Small Phytoplankton Chlorophyll (mg/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "spC_{direction}": Small Phytoplankton Carbon (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "spP_{direction}": Small Phytoplankton Phosphorous (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "spFe_{direction}": Small Phytoplankton Iron (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "spCaCO3_{direction}": Small Phytoplankton CaCO3 (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "diatChl_{direction}": Diatom Chlorophyll (mg/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "diatC_{direction}": Diatom Carbon (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "diatP_{direction}": Diatom Phosphorus (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "diatFe_{direction}": Diatom Iron (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "diatSi_{direction}": Diatom Silicate (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "diazChl_{direction}": Diazotroph Chlorophyll (mg/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "diazC_{direction}": Diazotroph Carbon (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "diazP_{direction}": Diazotroph Phosphorus (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
+            - "diazFe_{direction}": Diazotroph Iron (mmol/m³), where {direction} can be one of ["south", "east", "north", "west"].
         time : int, optional
             The time index to plot. Default is 0.
         layer_contours : bool, optional
@@ -482,8 +506,17 @@ class BoundaryForcing:
             If the specified varname is not one of the valid options.
         """
 
-        field = self.ds[varname].isel(time=time).load()
+        if varname in self.ds["physics"]:
+            ds = self.ds["physics"]
+        else:
+            if "bgc" in self.ds and varname in self.ds["bgc"]:
+                ds = self.ds["bgc"]
+            else:
+                raise ValueError(
+                    f"Variable '{varname}' is not found in 'physics' or 'bgc' datasets."
+                )
 
+        field = ds[varname].isel(bry_time=time).load()
         title = field.long_name
 
         # chose colorbar
@@ -494,7 +527,10 @@ class BoundaryForcing:
         else:
             vmax = field.max().values
             vmin = field.min().values
-            cmap = plt.colormaps.get_cmap("YlOrRd")
+            if varname.startswith(("temp", "salt")):
+                cmap = plt.colormaps.get_cmap("YlOrRd")
+            else:
+                cmap = plt.colormaps.get_cmap("YlGn")
         cmap.set_bad(color="gray")
         kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
@@ -507,13 +543,13 @@ class BoundaryForcing:
                 ]
                 try:
                     interface_depth = next(
-                        self.ds[depth_label]
-                        for depth_label in self.ds.coords
+                        ds[depth_label]
+                        for depth_label in ds.coords
                         if any(
                             depth_label.startswith(prefix) for prefix in depths_to_check
                         )
                         and (
-                            set(self.ds[depth_label].dims) - {"s_w"}
+                            set(ds[depth_label].dims) - {"s_w"}
                             == set(field.dims) - {"s_rho"}
                         )
                     )
@@ -563,40 +599,50 @@ class BoundaryForcing:
         filenames = []
         writes = []
 
-        # Group dataset by year
-        gb = self.ds.groupby("time.year")
-
-        for year, group_ds in gb:
-            # Further group each yearly group by month
-            sub_gb = group_ds.groupby("time.month")
-
-            for month, ds in sub_gb:
-                # Chunk the dataset by the specified time chunk size
-                ds = ds.chunk({"time": time_chunk_size})
-                datasets.append(ds)
-
-                # Determine the number of days in the month
-                num_days_in_month = calendar.monthrange(year, month)[1]
-                first_day = ds.time.dt.day.values[0]
-                last_day = ds.time.dt.day.values[-1]
-
-                # Create filename based on whether the dataset contains a full month
-                if first_day == 1 and last_day == num_days_in_month:
-                    # Full month format: "filepath.YYYYMM.nc"
-                    year_month_str = f"{year}{month:02}"
-                    filename = f"{filepath}.{year_month_str}.nc"
+        for node in ["physics", "bgc"]:
+            if node in self.ds:
+                ds = self.ds[node].to_dataset()
+                # copy vertical coordinate variables from parent to children because I believe this is info that ROMS needs
+                for var in self.ds.data_vars:
+                    ds[var] = self.ds[var]
+                if hasattr(ds["bry_time"], "cycle_length"):
+                    filename = f"{filepath}_{node}_clim.nc"
+                    filenames.append(filename)
+                    datasets.append(ds)
                 else:
-                    # Partial month format: "filepath.YYYYMMDD-DD.nc"
-                    year_month_day_str = f"{year}{month:02}{first_day:02}-{last_day:02}"
-                    filename = f"{filepath}.{year_month_day_str}.nc"
-                filenames.append(filename)
+                    # Group dataset by year
+                    gb = ds.groupby("abs_time.year")
+
+                    for year, group_ds in gb:
+                        # Further group each yearly group by month
+                        sub_gb = group_ds.groupby("abs_time.month")
+
+                        for month, ds in sub_gb:
+                            # Chunk the dataset by the specified time chunk size
+                            ds = ds.chunk({"bry_time": time_chunk_size})
+                            datasets.append(ds)
+
+                            # Determine the number of days in the month
+                            num_days_in_month = calendar.monthrange(year, month)[1]
+                            first_day = ds.abs_time.dt.day.values[0]
+                            last_day = ds.abs_time.dt.day.values[-1]
+
+                            # Create filename based on whether the dataset contains a full month
+                            if first_day == 1 and last_day == num_days_in_month:
+                                # Full month format: "filepath_physics_YYYYMM.nc"
+                                year_month_str = f"{year}{month:02}"
+                                filename = f"{filepath}_{node}_{year_month_str}.nc"
+                            else:
+                                # Partial month format: "filepath_physics_YYYYMMDD-DD.nc"
+                                year_month_day_str = (
+                                    f"{year}{month:02}{first_day:02}-{last_day:02}"
+                                )
+                                filename = f"{filepath}_{node}_{year_month_day_str}.nc"
+                            filenames.append(filename)
 
         print("Saving the following files:")
-        for filename in filenames:
-            print(filename)
-
         for ds, filename in zip(datasets, filenames):
-
+            print(filename)
             # Prepare the dataset for writing to a netCDF file without immediately computing
             write = ds.to_netcdf(filename, compute=False)
             writes.append(write)
@@ -637,12 +683,12 @@ class BoundaryForcing:
 
         boundary_forcing_data = {
             "BoundaryForcing": {
-                "filename": self.filename,
                 "start_time": self.start_time.isoformat(),
                 "end_time": self.end_time.isoformat(),
-                "model_reference_date": self.model_reference_date.isoformat(),
-                "source": self.source,
                 "boundaries": self.boundaries,
+                "physics_source": self.physics_source,
+                "bgc_source": self.bgc_source,
+                "model_reference_date": self.model_reference_date.isoformat(),
             }
         }