roms-tools 1.1.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

roms_tools/setup/boundary_forcing.py

@@ -2,9 +2,8 @@ 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, Union, Optional
+from typing import Dict, Union, List
 from dataclasses import dataclass, field, asdict
 from roms_tools.setup.grid import Grid
 from roms_tools.setup.mixins import ROMSToolsMixins
@@ -12,17 +11,21 @@ from datetime import datetime
 from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
 from roms_tools.setup.utils import (
     nan_check,
+    substitute_nans_by_fillvalue,
+    get_variable_metadata,
+    get_boundary_info,
+    group_dataset,
+    save_datasets,
 )
 from roms_tools.setup.plot import _section_plot, _line_plot
-import calendar
-import dask
 import matplotlib.pyplot as plt
+from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
 class BoundaryForcing(ROMSToolsMixins):
     """
-    Represents boundary forcing for ROMS.
+    Represents boundary forcing input data for ROMS.
 
     Parameters
     ----------
@@ -34,23 +37,21 @@ class BoundaryForcing(ROMSToolsMixins):
         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:
+    source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
+        Dictionary specifying the source of the 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.
+        - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file(s). Can be a single string (with or without wildcards),
+          a single Path object, or a list of strings or Path objects containing multiple files.
+        - "climatology" (bool): Indicates if the data is climatology data. Defaults to False.
     model_reference_date : datetime, optional
         Reference date for the model. Default is January 1, 2000.
+    use_dask: bool, optional
+        Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
 
     Attributes
     ----------
     ds : xr.Dataset
-        Xarray Dataset containing the atmospheric forcing data.
+        Xarray Dataset containing the boundary forcing data.
 
     Examples
     --------
@@ -59,12 +60,8 @@ class BoundaryForcing(ROMSToolsMixins):
     ...     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,
-    ...     },
+    ...     source={"name": "GLORYS", "path": "glorys_data.nc"},
+    ...     type="physics",
     ... )
     """
 
@@ -79,9 +76,10 @@ class BoundaryForcing(ROMSToolsMixins):
             "west": True,
         }
     )
-    physics_source: Dict[str, Union[str, None]]
-    bgc_source: Optional[Dict[str, Union[str, None]]] = None
+    source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
+    type: str = "physics"
     model_reference_date: datetime = datetime(2000, 1, 1)
+    use_dask: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
 
@@ -98,111 +96,94 @@ class BoundaryForcing(ROMSToolsMixins):
             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)
-        object.__setattr__(data, "data_vars", data_vars)
+        if self.type == "physics":
+            vars_2d = ["zeta"]
+            vars_3d = ["temp", "salt", "u", "v"]
+        elif self.type == "bgc":
+            vars_2d = []
+            vars_3d = data.var_names.keys()
 
-        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,
-            )
+        data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
 
-            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
+        if self.type == "physics":
+            data_vars = super().process_velocities(data_vars, angle, "u", "v")
+        object.__setattr__(data, "data_vars", data_vars)
 
-        d_meta = super().get_variable_metadata()
-        bdry_coords, rename = super().get_boundary_info()
+        d_meta = get_variable_metadata()
+        bdry_coords = get_boundary_info()
 
-        ds = self._write_into_datatree(data, bgc_data, d_meta, bdry_coords, rename)
+        ds = self._write_into_dataset(data, d_meta, bdry_coords)
 
+        # NaN values at wet points indicate that the raw data did not cover the domain, and the following will raise a ValueError
+        # this check works only for 2D fields because for 3D I extrapolate to bottom which eliminates NaNs
         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]),
-                )
+                if type == "physics":
+                    nan_check(
+                        ds[f"zeta_{direction}"].isel(bry_time=0),
+                        self.grid.ds.mask_rho.isel(**bdry_coords["rho"][direction]),
+                    )
+                elif type == "bgc":
+                    nan_check(
+                        ds[f"ALK_{direction}"].isel(bry_time=0, s_rho=-1),
+                        self.grid.ds.mask_rho.isel(**bdry_coords["rho"][direction]),
+                    )
+
+        # 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])
 
         object.__setattr__(self, "ds", ds)
 
     def _input_checks(self):
+        # Validate the 'type' parameter
+        if self.type not in ["physics", "bgc"]:
+            raise ValueError("`type` must be either 'physics' or 'bgc'.")
 
-        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
+        # Ensure 'source' dictionary contains required keys
+        if "name" not in self.source:
+            raise ValueError("`source` must include a 'name'.")
+        if "path" not in self.source:
+            raise ValueError("`source` must include a 'path'.")
+
+        # Set 'climatology' to False if not provided in 'source'
         object.__setattr__(
             self,
-            "physics_source",
-            {
-                **self.physics_source,
-                "climatology": self.physics_source.get("climatology", False),
-            },
+            "source",
+            {**self.source, "climatology": self.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),
-                },
-            )
 
     def _get_data(self):
 
-        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"].'
-            )
-
-        return data
+        data_dict = {
+            "filename": self.source["path"],
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "climatology": self.source["climatology"],
+            "use_dask": self.use_dask,
+        }
 
-    def _get_bgc_data(self):
+        if self.type == "physics":
+            if self.source["name"] == "GLORYS":
+                data = GLORYSDataset(**data_dict)
+            else:
+                raise ValueError(
+                    'Only "GLORYS" is a valid option for source["name"] when type is "physics".'
+                )
 
-        if self.bgc_source["name"] == "CESM_REGRIDDED":
+        elif self.type == "bgc":
+            if self.source["name"] == "CESM_REGRIDDED":
 
-            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"].'
-            )
+                data = CESMBGCDataset(**data_dict)
+                data.post_process()
+            else:
+                raise ValueError(
+                    'Only "CESM_REGRIDDED" is a valid option for source["name"] when type is "bgc".'
+                )
 
         return data
 
-    def _write_into_dataset(self, data, d_meta, bdry_coords, rename):
+    def _write_into_dataset(self, data, d_meta, bdry_coords):
 
         # save in new dataset
         ds = xr.Dataset()
@@ -215,21 +196,18 @@ class BoundaryForcing(ROMSToolsMixins):
                         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[
@@ -256,30 +234,49 @@ class BoundaryForcing(ROMSToolsMixins):
         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"]})
+        ds = self._add_global_metadata(ds)
 
         # Convert the time coordinate to the format expected by ROMS
         if data.climatology:
+            ds.attrs["climatology"] = str(True)
+            # Preserve absolute time coordinate for readability
+            ds = ds.assign_coords(
+                {"abs_time": np.datetime64(self.model_reference_date) + ds["time"]}
+            )
             # 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
+
+            # Convert the timedelta to nanoseconds first, then to days
             bry_time = xr.DataArray(
-                timedelta_index - offset,
+                (timedelta_index - offset).view("int64") / 3600 / 24 * 1e-9,
                 dims="time",
             )
+
         else:
+            # Preserve absolute time coordinate for readability
+            ds = ds.assign_coords({"abs_time": ds["time"]})
             # 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)
+            bry_time = (
+                (ds["time"] - np.datetime64(self.model_reference_date)).astype(
+                    "float64"
+                )
+                / 3600
+                / 24
+                * 1e-9
+            )
 
         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"
+        ] = f"days since {str(self.model_reference_date)}"
+        ds["bry_time"].encoding["units"] = "days"
+        ds["bry_time"].attrs["units"] = "days"
         ds = ds.swap_dims({"time": "bry_time"})
         ds = ds.drop_vars("time")
 
@@ -288,127 +285,41 @@ class BoundaryForcing(ROMSToolsMixins):
 
         return ds
 
-    def _write_into_datatree(self, data, bgc_data, d_meta, bdry_coords, rename):
-
-        ds = self._add_global_metadata()
-        ds["sc_r"] = self.grid.ds["sc_r"]
-        ds["Cs_r"] = self.grid.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)
-
-        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"]:
+    def _get_coordinates(self, direction, point):
+        """
+        Retrieve layer and interface depth coordinates for a specified grid boundary.
 
-            if self.boundaries[direction]:
+        This method extracts the layer depth and interface depth coordinates along
+        a specified boundary (north, south, east, or west) and for a specified point
+        type (rho, u, or v) from the grid dataset.
 
-                lat_rho = self.grid.ds.lat_rho.isel(
-                    **bdry_coords["rho"][direction]
-                ).rename(**rename["rho"][direction])
-                lon_rho = self.grid.ds.lon_rho.isel(
-                    **bdry_coords["rho"][direction]
-                ).rename(**rename["rho"][direction])
-                layer_depth_rho = (
-                    self.grid.ds["layer_depth_rho"]
-                    .isel(**bdry_coords["rho"][direction])
-                    .rename(**rename["rho"][direction])
-                )
-                interface_depth_rho = (
-                    self.grid.ds["interface_depth_rho"]
-                    .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]
-                )
-                lon_u = self.grid.ds.lon_u.isel(**bdry_coords["u"][direction]).rename(
-                    **rename["u"][direction]
-                )
-                layer_depth_u = (
-                    self.grid.ds["layer_depth_u"]
-                    .isel(**bdry_coords["u"][direction])
-                    .rename(**rename["u"][direction])
-                )
-                interface_depth_u = (
-                    self.grid.ds["interface_depth_u"]
-                    .isel(**bdry_coords["u"][direction])
-                    .rename(**rename["u"][direction])
-                )
+        Parameters
+        ----------
+        direction : str
+            The direction of the boundary to retrieve coordinates for. Valid options
+            are "north", "south", "east", and "west".
+        point : str
+            The type of grid point to retrieve coordinates for. Valid options are
+            "rho" for the grid's central points, "u" for the u-flux points, and "v"
+            for the v-flux points.
 
-                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]
-                )
-                layer_depth_v = (
-                    self.grid.ds["layer_depth_v"]
-                    .isel(**bdry_coords["v"][direction])
-                    .rename(**rename["v"][direction])
-                )
-                interface_depth_v = (
-                    self.grid.ds["interface_depth_v"]
-                    .isel(**bdry_coords["v"][direction])
-                    .rename(**rename["v"][direction])
-                )
+        Returns
+        -------
+        xarray.DataArray, xarray.DataArray
+            The layer depth and interface depth coordinates for the specified grid
+            boundary and point type.
+        """
 
-                ds = ds.assign_coords(
-                    {
-                        f"layer_depth_rho_{direction}": layer_depth_rho,
-                        f"layer_depth_u_{direction}": layer_depth_u,
-                        f"layer_depth_v_{direction}": layer_depth_v,
-                        f"interface_depth_rho_{direction}": interface_depth_rho,
-                        f"interface_depth_u_{direction}": interface_depth_u,
-                        f"interface_depth_v_{direction}": interface_depth_v,
-                        f"lat_rho_{direction}": lat_rho,
-                        f"lat_u_{direction}": lat_u,
-                        f"lat_v_{direction}": lat_v,
-                        f"lon_rho_{direction}": lon_rho,
-                        f"lon_u_{direction}": lon_u,
-                        f"lon_v_{direction}": lon_v,
-                    }
-                )
+        bdry_coords = get_boundary_info()
 
-        # 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)
+        layer_depth = self.grid.ds[f"layer_depth_{point}"].isel(
+            **bdry_coords[point][direction]
+        )
+        interface_depth = self.grid.ds[f"interface_depth_{point}"].isel(
+            **bdry_coords[point][direction]
+        )
 
-        return ds
+        return layer_depth, interface_depth
 
     def _add_global_metadata(self, ds=None):
 
@@ -423,6 +334,7 @@ class BoundaryForcing(ROMSToolsMixins):
         ds.attrs["roms_tools_version"] = roms_tools_version
         ds.attrs["start_time"] = str(self.start_time)
         ds.attrs["end_time"] = str(self.end_time)
+        ds.attrs["source"] = self.source["name"]
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
 
         ds.attrs["theta_s"] = self.grid.ds.attrs["theta_s"]
@@ -443,7 +355,7 @@ class BoundaryForcing(ROMSToolsMixins):
         Parameters
         ----------
         varname : str
-            The name of the initial conditions field to plot. Options include:
+            The name of the boundary forcing field to plot. Options include:
             - "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"].
@@ -486,8 +398,9 @@ class BoundaryForcing(ROMSToolsMixins):
         time : int, optional
             The time index to plot. Default is 0.
         layer_contours : bool, optional
-            Whether to include layer contours in the plot. This can help visualize the depth levels
-            of the field. Default is False.
+            If True, contour lines representing the boundaries between vertical layers will
+            be added to the plot. For clarity, the number of layer
+            contours displayed is limited to a maximum of 10. Default is False.
 
         Returns
         -------
@@ -500,19 +413,25 @@ class BoundaryForcing(ROMSToolsMixins):
             If the specified varname is not one of the valid options.
         """
 
-        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."
-                )
+        if varname not in self.ds:
+            raise ValueError(f"Variable '{varname}' is not found in dataset.")
 
-        field = ds[varname].isel(bry_time=time).load()
+        field = self.ds[varname].isel(bry_time=time).load()
         title = field.long_name
 
+        if "s_rho" in field.dims:
+            if varname.startswith(("u_", "ubar_")):
+                point = "u"
+            elif varname.startswith(("v_", "vbar_")):
+                point = "v"
+            else:
+                point = "rho"
+            direction = varname.split("_")[-1]
+
+            layer_depth, interface_depth = self._get_coordinates(direction, point)
+
+            field = field.assign_coords({"layer_depth": layer_depth})
+
         # chose colorbar
         if varname.startswith(("u", "v", "ubar", "vbar", "zeta")):
             vmax = max(field.max().values, -field.min().values)
@@ -530,27 +449,6 @@ class BoundaryForcing(ROMSToolsMixins):
 
         if len(field.dims) == 2:
             if layer_contours:
-                depths_to_check = [
-                    "interface_depth_rho",
-                    "interface_depth_u",
-                    "interface_depth_v",
-                ]
-                try:
-                    interface_depth = next(
-                        ds[depth_label]
-                        for depth_label in ds.coords
-                        if any(
-                            depth_label.startswith(prefix) for prefix in depths_to_check
-                        )
-                        and (
-                            set(ds[depth_label].dims) - {"s_w"}
-                            == set(field.dims) - {"s_rho"}
-                        )
-                    )
-                except StopIteration:
-                    raise ValueError(
-                        f"None of the expected depths ({', '.join(depths_to_check)}) have dimensions matching field.dims"
-                    )
                 # restrict number of layer_contours to 10 for the sake of plot clearity
                 nr_layers = len(interface_depth["s_w"])
                 selected_layers = np.linspace(
@@ -567,92 +465,67 @@ class BoundaryForcing(ROMSToolsMixins):
         else:
             _line_plot(field, title=title)
 
-    def save(self, filepath: str, time_chunk_size: int = 1) -> None:
+    def save(
+        self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
+    ) -> None:
         """
-        Save the interpolated boundary forcing fields to netCDF4 files.
+        Save the boundary forcing fields to netCDF4 files.
+
+        This method saves the dataset by grouping it into subsets based on the data frequency. The subsets are then written
+        to one or more netCDF4 files. The filenames of the output files reflect the temporal coverage of the data.
+
+        There are two modes of saving the dataset:
+
+        1. **Single File Mode (default)**:
+           - If both `np_eta` and `np_xi` are `None`, the entire dataset, divided by temporal subsets, is saved as a single netCDF4 file
+             with the base filename specified by `filepath.nc`.
 
-        This method groups the dataset by year and month, chunks the data by the specified
-        time chunk size, and saves each chunked subset to a separate netCDF4 file named
-        according to the year, month, and day range if not a complete month of data is included.
+        2. **Partitioned Mode**:
+           - If either `np_eta` or `np_xi` is specified, the dataset is divided into spatial tiles along the eta-axis and xi-axis.
+           - Each spatial tile is saved as a separate netCDF4 file.
 
         Parameters
         ----------
-        filepath : str
-            The base path and filename for the output files. The files will be named with
-            the format "filepath.YYYYMM.nc" if a full month of data is included, or
-            "filepath.YYYYMMDD-DD.nc" otherwise.
-        time_chunk_size : int, optional
-            Number of time slices to include in each chunk along the time dimension. Default is 1,
-            meaning each chunk contains one time slice.
+        filepath : Union[str, Path]
+            The base path and filename for the output files. The format of the filenames depends on whether partitioning is used
+            and the temporal range of the data. For partitioned datasets, files will be named with an additional index, e.g.,
+            `"filepath_YYYYMM.0.nc"`, `"filepath_YYYYMM.1.nc"`, etc.
+        np_eta : int, optional
+            The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed.
+        np_xi : int, optional
+            The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed.
 
         Returns
         -------
-        None
+        List[Path]
+            A list of Path objects for the filenames that were saved.
         """
-        datasets = []
-        filenames = []
-        writes = []
-
-        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:
-                    # 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 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)
-
-        # Perform the actual write operations in parallel
-        dask.compute(*writes)
-
-    def to_yaml(self, filepath: str) -> None:
+
+        # Ensure filepath is a Path object
+        filepath = Path(filepath)
+
+        # Remove ".nc" suffix if present
+        if filepath.suffix == ".nc":
+            filepath = filepath.with_suffix("")
+
+        dataset_list, output_filenames = group_dataset(self.ds.load(), str(filepath))
+        saved_filenames = save_datasets(
+            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
+        )
+
+        return saved_filenames
+
+    def to_yaml(self, filepath: Union[str, Path]) -> None:
         """
         Export the parameters of the class to a YAML file, including the version of roms-tools.
 
         Parameters
         ----------
-        filepath : str
+        filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
+        filepath = Path(filepath)
+
         # Serialize Grid data
         grid_data = asdict(self.grid)
         grid_data.pop("ds", None)  # Exclude non-serializable fields
@@ -674,8 +547,8 @@ class BoundaryForcing(ROMSToolsMixins):
                 "start_time": self.start_time.isoformat(),
                 "end_time": self.end_time.isoformat(),
                 "boundaries": self.boundaries,
-                "physics_source": self.physics_source,
-                "bgc_source": self.bgc_source,
+                "source": self.source,
+                "type": self.type,
                 "model_reference_date": self.model_reference_date.isoformat(),
             }
         }
@@ -685,29 +558,34 @@ class BoundaryForcing(ROMSToolsMixins):
             **boundary_forcing_data,
         }
 
-        with open(filepath, "w") as file:
+        with filepath.open("w") as file:
             # Write header
             file.write(header)
             # Write YAML data
             yaml.dump(yaml_data, file, default_flow_style=False)
 
     @classmethod
-    def from_yaml(cls, filepath: str) -> "BoundaryForcing":
+    def from_yaml(
+        cls, filepath: Union[str, Path], use_dask: bool = False
+    ) -> "BoundaryForcing":
         """
         Create an instance of the BoundaryForcing class from a YAML file.
 
         Parameters
         ----------
-        filepath : str
+        filepath : Union[str, Path]
             The path to the YAML file from which the parameters will be read.
+        use_dask: bool, optional
+            Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
 
         Returns
         -------
         BoundaryForcing
             An instance of the BoundaryForcing class.
         """
+        filepath = Path(filepath)
         # Read the entire file content
-        with open(filepath, "r") as file:
+        with filepath.open("r") as file:
             file_content = file.read()
 
         # Split the content into YAML documents
@@ -735,7 +613,4 @@ class BoundaryForcing(ROMSToolsMixins):
         grid = Grid.from_yaml(filepath)
 
         # Create and return an instance of InitialConditions
-        return cls(
-            grid=grid,
-            **boundary_forcing_data,
-        )
+        return cls(grid=grid, **boundary_forcing_data, use_dask=use_dask)