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

roms_tools/setup/surface_forcing.py

@@ -1,55 +1,61 @@
 import xarray as xr
 import pandas as pd
-import dask
 import yaml
-from datatree import DataTree
 import importlib.metadata
 from dataclasses import dataclass, field, asdict
 from roms_tools.setup.grid import Grid
 from datetime import datetime
 import numpy as np
-from typing import Dict, Union
+from typing import Dict, Union, List
 from roms_tools.setup.mixins import ROMSToolsMixins
 from roms_tools.setup.datasets import (
     ERA5Dataset,
     ERA5Correction,
     CESMBGCSurfaceForcingDataset,
 )
-from roms_tools.setup.utils import nan_check, interpolate_from_climatology
+from roms_tools.setup.utils import (
+    nan_check,
+    substitute_nans_by_fillvalue,
+    interpolate_from_climatology,
+    get_variable_metadata,
+    group_dataset,
+    save_datasets,
+)
 from roms_tools.setup.plot import _plot
-import calendar
 import matplotlib.pyplot as plt
+from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
 class SurfaceForcing(ROMSToolsMixins):
     """
-    Represents surface forcing data for ocean modeling.
+    Represents surface forcing input data for ROMS.
 
     Parameters
     ----------
     grid : Grid
         Object representing the grid information.
     start_time : datetime
-        Start time of the desired forcing data.
+        Start time of the desired surface forcing data.
     end_time : datetime
-        End time of the desired forcing data.
-    physics_source : Dict[str, Union[str, None]]
-        Dictionary specifying the source of the physical surface forcing data:
+        End time of the desired surface forcing data.
+    source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
+        Dictionary specifying the source of the surface forcing data:
         - "name" (str): Name of the data source (e.g., "ERA5").
-        - "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 False.
+        - "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.
+    type : str
+        Specifies the type of forcing data, either "physics" for physical
+        atmospheric forcing or "bgc" for biogeochemical forcing.
     correct_radiation : bool
         Whether to correct shortwave radiation. Default is False.
     use_coarse_grid: bool
         Whether to interpolate to coarsened grid. Default is 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
     ----------
@@ -59,11 +65,12 @@ class SurfaceForcing(ROMSToolsMixins):
 
     Examples
     --------
-    >>> atm_forcing = SurfaceForcing(
+    >>> surface_forcing = SurfaceForcing(
     ...     grid=grid,
     ...     start_time=datetime(2000, 1, 1),
     ...     end_time=datetime(2000, 1, 2),
-    ...     physics_source={"name": "ERA5", "path": "physics_data.nc"},
+    ...     source={"name": "ERA5", "path": "era5_data.nc"},
+    ...     type="physics",
     ...     correct_radiation=True,
     ... )
     """
@@ -71,11 +78,13 @@ class SurfaceForcing(ROMSToolsMixins):
     grid: Grid
     start_time: datetime
     end_time: datetime
-    physics_source: Dict[str, Union[str, None]]
-    bgc_source: Dict[str, Union[str, None]] = None
+    source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
+    type: str = "physics"
     correct_radiation: bool = False
     use_coarse_grid: bool = False
     model_reference_date: datetime = datetime(2000, 1, 1)
+    use_dask: bool = False
+
     ds: xr.Dataset = field(init=False, repr=False)
 
     def __post_init__(self):
@@ -92,68 +101,59 @@ class SurfaceForcing(ROMSToolsMixins):
             margin=2,
             straddle=straddle,
         )
-        vars_2d = ["uwnd", "vwnd", "swrad", "lwrad", "Tair", "qair", "rain"]
+        if self.type == "physics":
+            vars_2d = ["uwnd", "vwnd", "swrad", "lwrad", "Tair", "qair", "rain"]
+        elif self.type == "bgc":
+            vars_2d = data.var_names.keys()
         vars_3d = []
+
         data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
-        data_vars = super().process_velocities(data_vars, angle, interpolate=False)
-
-        if self.correct_radiation:
-            correction_data = self._get_correction_data()
-            # choose same subdomain as forcing data so that we can use same mask
-            coords_correction = {
-                correction_data.dim_names["latitude"]: data.ds[
-                    data.dim_names["latitude"]
-                ],
-                correction_data.dim_names["longitude"]: data.ds[
-                    data.dim_names["longitude"]
-                ],
-            }
-            correction_data.choose_subdomain(coords_correction, straddle=straddle)
-            # apply mask from ERA5 data
-            if "mask" in data.var_names.keys():
-                mask = xr.where(
-                    data.ds[data.var_names["mask"]].isel(time=0).isnull(), 0, 1
-                )
-                for var in correction_data.ds.data_vars:
-                    correction_data.ds[var] = xr.where(
-                        mask == 1, correction_data.ds[var], np.nan
-                    )
-            vars_2d = ["swr_corr"]
-            vars_3d = []
-            # spatial interpolation
-            data_vars_corr = super().regrid_data(
-                correction_data, vars_2d, vars_3d, lon, lat
-            )
-            # temporal interpolation
-            corr_factor = interpolate_from_climatology(
-                data_vars_corr["swr_corr"],
-                correction_data.dim_names["time"],
-                time=data_vars["swrad"].time,
+
+        if self.type == "physics":
+            data_vars = super().process_velocities(
+                data_vars, angle, "uwnd", "vwnd", interpolate=False
             )
+            if self.correct_radiation:
+                correction_data = self._get_correction_data()
+                # choose same subdomain as forcing data so that we can use same mask
+                coords_correction = {
+                    correction_data.dim_names["latitude"]: data.ds[
+                        data.dim_names["latitude"]
+                    ],
+                    correction_data.dim_names["longitude"]: data.ds[
+                        data.dim_names["longitude"]
+                    ],
+                }
+                correction_data.choose_subdomain(coords_correction, straddle=straddle)
+                # apply mask from ERA5 data
+                if "mask" in data.var_names.keys():
+                    mask = xr.where(
+                        data.ds[data.var_names["mask"]].isel(time=0).isnull(), 0, 1
+                    )
+                    for var in correction_data.ds.data_vars:
+                        correction_data.ds[var] = xr.where(
+                            mask == 1, correction_data.ds[var], np.nan
+                        )
+                vars_2d = ["swr_corr"]
+                vars_3d = []
+                # spatial interpolation
+                data_vars_corr = super().regrid_data(
+                    correction_data, vars_2d, vars_3d, lon, lat
+                )
+                # temporal interpolation
+                corr_factor = interpolate_from_climatology(
+                    data_vars_corr["swr_corr"],
+                    correction_data.dim_names["time"],
+                    time=data_vars["swrad"].time,
+                )
 
-            data_vars["swrad"] = data_vars["swrad"] * corr_factor
+                data_vars["swrad"] = data_vars["swrad"] * corr_factor
 
         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,
-            )
-
-            vars_2d = bgc_data.var_names.keys()
-            vars_3d = []
-            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
-
-        d_meta = super().get_variable_metadata()
+        d_meta = get_variable_metadata()
 
-        ds = self._write_into_datatree(data, bgc_data, d_meta)
+        ds = self._write_into_dataset(data, d_meta)
 
         if self.use_coarse_grid:
             mask = self.grid.ds["mask_coarse"].rename(
@@ -162,92 +162,77 @@ class SurfaceForcing(ROMSToolsMixins):
         else:
             mask = self.grid.ds["mask_rho"]
 
-        for var in ds["physics"].data_vars:
-            nan_check(ds["physics"][var].isel(time=0), mask)
+        # NaN values at wet points indicate that the raw data did not cover the domain, and the following will raise a ValueError
+        for var in ds.data_vars:
+            nan_check(ds[var].isel(time=0), mask)
+
+        # 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'.")
+
+        # 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'.")
 
-        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
+        # 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():
+    def _get_data(self):
+
+        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,
+        }
+
+        if self.type == "physics":
+            if self.source["name"] == "ERA5":
+                data = ERA5Dataset(**data_dict)
+                data.post_process()
+            else:
                 raise ValueError(
-                    "`bgc_source` must include a 'path' if it is provided."
+                    'Only "ERA5" is a valid option for source["name"] when type is "physics".'
                 )
-            # set self.bgc_source["climatology"] to False if not provided
-            object.__setattr__(
-                self,
-                "bgc_source",
-                {
-                    **self.bgc_source,
-                    "climatology": self.bgc_source.get("climatology", False),
-                },
-            )
 
-    def _get_data(self):
+        elif self.type == "bgc":
+            if self.source["name"] == "CESM_REGRIDDED":
 
-        if self.physics_source["name"] == "ERA5":
-            data = ERA5Dataset(
-                filename=self.physics_source["path"],
-                start_time=self.start_time,
-                end_time=self.end_time,
-                climatology=self.physics_source["climatology"],
-            )
-            data.post_process()
-        else:
-            raise ValueError(
-                'Only "ERA5" is a valid option for physics_source["name"].'
-            )
+                data = CESMBGCSurfaceForcingDataset(**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 _get_correction_data(self):
 
-        if self.physics_source["name"] == "ERA5":
-            correction_data = ERA5Correction()
+        if self.source["name"] == "ERA5":
+            correction_data = ERA5Correction(use_dask=self.use_dask)
         else:
             raise ValueError(
-                "The 'correct_radiation' feature is currently only supported for 'ERA5' as the physics source. "
-                "Please ensure your 'physics_source' is set to 'ERA5' or implement additional handling for other sources."
+                "The 'correct_radiation' feature is currently only supported for 'ERA5' as the source. "
+                "Please ensure your 'source' is set to 'ERA5' or implement additional handling for other sources."
             )
 
         return correction_data
 
-    def _get_bgc_data(self):
-
-        if self.bgc_source["name"] == "CESM_REGRIDDED":
-
-            bgc_data = CESMBGCSurfaceForcingDataset(
-                filename=self.bgc_source["path"],
-                start_time=self.start_time,
-                end_time=self.end_time,
-                climatology=self.bgc_source["climatology"],
-            )
-        else:
-            raise ValueError(
-                'Only "CESM_REGRIDDED" is a valid option for bgc_source["name"].'
-            )
-
-        return bgc_data
-
     def _write_into_dataset(self, data, d_meta):
 
         # save in new dataset
@@ -259,25 +244,35 @@ class SurfaceForcing(ROMSToolsMixins):
             ds[var].attrs["units"] = d_meta[var]["units"]
 
         if self.use_coarse_grid:
-            ds = ds.assign_coords({"lon": self.target_lon, "lat": self.target_lat})
             ds = ds.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
 
-        # Preserve absolute time coordinate for readability
-        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
             sfc_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"]})
+
             sfc_time = (
                 (ds["time"] - np.datetime64(self.model_reference_date)).astype(
                     "float64"
@@ -287,32 +282,32 @@ class SurfaceForcing(ROMSToolsMixins):
                 * 1e-9
             )
 
-        ds = ds.assign_coords({"time": sfc_time})
-        ds["time"].attrs[
-            "long_name"
-        ] = f"days since {np.datetime_as_string(np.datetime64(self.model_reference_date), unit='D')}"
-        ds["time"].encoding["units"] = "days"
-        if data.climatology:
-            ds["time"].attrs["cycle_length"] = 365.25
-
-        return ds
-
-    def _write_into_datatree(self, data, bgc_data, d_meta):
-
-        ds = self._add_global_metadata()
-
-        ds = DataTree(name="root", data=ds)
-
-        ds_physics = self._write_into_dataset(data, d_meta)
-        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)
-            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)
+        if self.type == "physics":
+            time_coords = ["time"]
+        elif self.type == "bgc":
+            time_coords = [
+                "pco2_time",
+                "iron_time",
+                "dust_time",
+                "nox_time",
+                "nhy_time",
+            ]
+        for time_coord in time_coords:
+            ds = ds.assign_coords({time_coord: sfc_time})
+            ds[time_coord].attrs[
+                "long_name"
+            ] = f"days since {str(self.model_reference_date)}"
+            ds[time_coord].encoding["units"] = "days"
+            ds[time_coord].attrs["units"] = "days"
+            if data.climatology:
+                ds[time_coord].attrs["cycle_length"] = 365.25
+
+        if self.type == "bgc":
+            ds = ds.drop_vars(["time"])
+
+        variables_to_drop = ["lat_rho", "lon_rho", "lat_coarse", "lon_coarse"]
+        existing_vars = [var for var in variables_to_drop if var in ds]
+        ds = ds.drop_vars(existing_vars)
 
         return ds
 
@@ -329,13 +324,14 @@ class SurfaceForcing(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["physics_source"] = self.physics_source["name"]
-        if self.bgc_source is not None:
-            ds.attrs["bgc_source"] = self.bgc_source["name"]
+        ds.attrs["source"] = self.source["name"]
         ds.attrs["correct_radiation"] = str(self.correct_radiation)
         ds.attrs["use_coarse_grid"] = str(self.use_coarse_grid)
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
 
+        ds.attrs["type"] = self.type
+        ds.attrs["source"] = self.source["name"]
+
         return ds
 
     def plot(self, varname, time=0) -> None:
@@ -371,7 +367,7 @@ class SurfaceForcing(ROMSToolsMixins):
         Raises
         ------
         ValueError
-            If the specified varname is not one of the valid options.
+            If the specified varname is not found in dataset.
 
 
         Examples
@@ -379,19 +375,21 @@ class SurfaceForcing(ROMSToolsMixins):
         >>> atm_forcing.plot("uwnd", time=0)
         """
 
-        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(time=time).load()
+        field = self.ds[varname].isel(time=time).load()
         title = field.long_name
 
+        # assign lat / lon
+        if self.use_coarse_grid:
+            field = field.rename({"eta_rho": "eta_coarse", "xi_rho": "xi_coarse"})
+            field = field.where(self.grid.ds.mask_coarse)
+        else:
+            field = field.where(self.grid.ds.mask_rho)
+
+        field = field.assign_coords({"lon": self.target_lon, "lat": self.target_lat})
+
         # choose colorbar
         if varname in ["uwnd", "vwnd"]:
             vmax = max(field.max().values, -field.min().values)
@@ -412,96 +410,72 @@ class SurfaceForcing(ROMSToolsMixins):
             self.grid.ds,
             field=field,
             straddle=self.grid.straddle,
-            coarse_grid=self.use_coarse_grid,
             title=title,
             kwargs=kwargs,
             c="g",
         )
 
-    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 surface forcing fields to netCDF4 files.
+        Save the surface 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()
-                if hasattr(ds["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({"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
@@ -524,15 +498,13 @@ class SurfaceForcing(ROMSToolsMixins):
             "SurfaceForcing": {
                 "start_time": self.start_time.isoformat(),
                 "end_time": self.end_time.isoformat(),
-                "physics_source": self.physics_source,
+                "source": self.source,
+                "type": self.type,
                 "correct_radiation": self.correct_radiation,
                 "use_coarse_grid": self.use_coarse_grid,
                 "model_reference_date": self.model_reference_date.isoformat(),
             }
         }
-        # Include bgc_source if it's not None
-        if self.bgc_source is not None:
-            surface_forcing_data["SurfaceForcing"]["bgc_source"] = self.bgc_source
 
         # Merge YAML data while excluding empty sections
         yaml_data = {
@@ -540,29 +512,34 @@ class SurfaceForcing(ROMSToolsMixins):
             **surface_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) -> "SurfaceForcing":
+    def from_yaml(
+        cls, filepath: Union[str, Path], use_dask: bool = False
+    ) -> "SurfaceForcing":
         """
         Create an instance of the SurfaceForcing 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
         -------
         SurfaceForcing
             An instance of the SurfaceForcing 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
@@ -590,7 +567,4 @@ class SurfaceForcing(ROMSToolsMixins):
         grid = Grid.from_yaml(filepath)
 
         # Create and return an instance of SurfaceForcing
-        return cls(
-            grid=grid,
-            **surface_forcing_data,
-        )
+        return cls(grid=grid, **surface_forcing_data, use_dask=use_dask)