roms-tools 2.7.0__py3-none-any.whl → 3.1.0__py3-none-any.whl

roms_tools/setup/utils.py

@@ -1,16 +1,27 @@
-import xarray as xr
+import importlib.metadata
+from collections.abc import Sequence
+from dataclasses import asdict, fields, is_dataclass
+from datetime import datetime
+from enum import StrEnum
+from pathlib import Path
+from typing import Any
+
+import cftime
+import numba as nb
 import numpy as np
-from typing import Union, Any, Dict, Type, Sequence
 import pandas as pd
-import cftime
-from pathlib import Path
-from datetime import datetime
-from dataclasses import fields, asdict
-import importlib.metadata
+import xarray as xr
 import yaml
+from pydantic import BaseModel
+
 from roms_tools.constants import R_EARTH
 from roms_tools.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
 
+yaml.SafeDumper.add_multi_representer(
+    StrEnum,
+    yaml.representer.SafeRepresenter.represent_str,
+)
+
 
 def nan_check(field, mask, error_message=None) -> None:
     """Checks for NaN values at wet points in the field.
@@ -37,7 +48,6 @@ def nan_check(field, mask, error_message=None) -> None:
         If the field contains NaN values at any of the wet points indicated by the mask.
         The error message will explain the potential cause and suggest ensuring the dataset's coverage.
     """
-
     # Replace values in field with 0 where mask is not 1
     da = xr.where(mask == 1, field, 0)
     if error_message is None:
@@ -68,7 +78,6 @@ def substitute_nans_by_fillvalue(field, fill_value=0.0) -> xr.DataArray:
     xr.DataArray
         The data array with NaN values replaced by the specified fill value.
     """
-
     return field.fillna(fill_value)
 
 
@@ -129,9 +138,7 @@ def assign_dates_to_climatology(ds: xr.Dataset, time_dim: str) -> xr.Dataset:
 def interpolate_cyclic_time(
     data_array: xr.DataArray,
     time_dim_name: str,
-    day_of_year: Union[
-        int, float, np.ndarray, xr.DataArray, Sequence[Union[int, float]]
-    ],
+    day_of_year: int | float | np.ndarray | xr.DataArray | Sequence[int | float],
 ) -> xr.DataArray:
     """Interpolates a DataArray cyclically across the start and end of the year.
 
@@ -163,7 +170,6 @@ def interpolate_cyclic_time(
       represents a repeating annual cycle.
     - The `day_of_year` values should be within the range [1, 365] or [1, 366] for leap years.
     """
-
     # Concatenate across the beginning and end of the year
     time_concat = xr.concat(
         [
@@ -195,10 +201,10 @@ def interpolate_cyclic_time(
 
 
 def interpolate_from_climatology(
-    field: Union[xr.DataArray, xr.Dataset],
+    field: xr.DataArray | xr.Dataset,
     time_dim_name: str,
-    time: Union[xr.DataArray, pd.DatetimeIndex],
-) -> Union[xr.DataArray, xr.Dataset]:
+    time: xr.DataArray | pd.DatetimeIndex,
+) -> xr.DataArray | xr.Dataset:
     """Interpolates a climatological field to specified time points.
 
     This function interpolates the input `field` based on `day_of_year` values
@@ -231,7 +237,6 @@ def interpolate_from_climatology(
     """
 
     def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
-
         if isinstance(time, xr.DataArray):
             # Extract day of year from xarray.DataArray
             day_of_year = time.dt.dayofyear
@@ -308,7 +313,7 @@ def get_time_type(data_array: xr.DataArray) -> str:
     )
 
     # Check if any of the coordinate values are of cftime, datetime, or integer type
-    if isinstance(data_array.values, (np.ndarray, list)):
+    if isinstance(data_array.values, np.ndarray | list):
         # Check if the data type is numpy datetime64, indicating standard datetime objects
         if data_array.values.dtype == "datetime64[ns]":
             return "datetime"
@@ -375,7 +380,6 @@ def get_variable_metadata():
     dict of str: dict
         Dictionary where keys are variable names and values are dictionaries with "long_name" and "units" keys.
     """
-
     d = {
         "ssh_Re": {"long_name": "Tidal elevation, real part", "units": "m"},
         "ssh_Im": {"long_name": "Tidal elevation, complex part", "units": "m"},
@@ -410,8 +414,12 @@ def get_variable_metadata():
         "Tair": {"long_name": "air temperature at 2m", "units": "degrees Celsius"},
         "qair": {"long_name": "absolute humidity at 2m", "units": "kg/kg"},
         "rain": {"long_name": "total precipitation", "units": "cm/day"},
-        "temp": {"long_name": "potential temperature", "units": "degrees Celsius"},
-        "salt": {"long_name": "salinity", "units": "PSU"},
+        "temp": {
+            "long_name": "potential temperature",
+            "units": "degrees Celsius",
+            "flux_units": "degrees Celsius/s",
+        },
+        "salt": {"long_name": "salinity", "units": "PSU", "flux_units": "PSU/s"},
         "zeta": {"long_name": "sea surface height", "units": "m"},
         "u": {"long_name": "u-flux component", "units": "m/s"},
         "v": {"long_name": "v-flux component", "units": "m/s"},
@@ -424,59 +432,162 @@ def get_variable_metadata():
             "long_name": "vertically integrated v-flux component",
             "units": "m/s",
         },
-        "PO4": {"long_name": "dissolved inorganic phosphate", "units": "mmol/m^3"},
-        "NO3": {"long_name": "dissolved inorganic nitrate", "units": "mmol/m^3"},
-        "SiO3": {"long_name": "dissolved inorganic silicate", "units": "mmol/m^3"},
-        "NH4": {"long_name": "dissolved ammonia", "units": "mmol/m^3"},
-        "Fe": {"long_name": "dissolved inorganic iron", "units": "mmol/m^3"},
-        "Lig": {"long_name": "iron binding ligand", "units": "mmol/m^3"},
-        "O2": {"long_name": "dissolved oxygen", "units": "mmol/m^3"},
-        "DIC": {"long_name": "dissolved inorganic carbon", "units": "mmol/m^3"},
+        "PO4": {
+            "long_name": "dissolved inorganic phosphate",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "NO3": {
+            "long_name": "dissolved inorganic nitrate",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "SiO3": {
+            "long_name": "dissolved inorganic silicate",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "NH4": {
+            "long_name": "dissolved ammonia",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "Fe": {
+            "long_name": "dissolved inorganic iron",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "Lig": {
+            "long_name": "iron binding ligand",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "O2": {
+            "long_name": "dissolved oxygen",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "DIC": {
+            "long_name": "dissolved inorganic carbon",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
         "DIC_ALT_CO2": {
             "long_name": "dissolved inorganic carbon, alternative CO2",
             "units": "mmol/m^3",
+            "flux_units": "meq/s",
         },
-        "ALK": {"long_name": "alkalinity", "units": "meq/m^3"},
+        "ALK": {"long_name": "alkalinity", "units": "meq/m^3", "flux_units": "meq/s"},
         "ALK_ALT_CO2": {
             "long_name": "alkalinity, alternative CO2",
             "units": "meq/m^3",
+            "flux_units": "meq/s",
+        },
+        "DOC": {
+            "long_name": "dissolved organic carbon",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "DON": {
+            "long_name": "dissolved organic nitrogen",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "DOP": {
+            "long_name": "dissolved organic phosphorus",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
         },
-        "DOC": {"long_name": "dissolved organic carbon", "units": "mmol/m^3"},
-        "DON": {"long_name": "dissolved organic nitrogen", "units": "mmol/m^3"},
-        "DOP": {"long_name": "dissolved organic phosphorus", "units": "mmol/m^3"},
         "DOCr": {
             "long_name": "refractory dissolved organic carbon",
             "units": "mmol/m^3",
+            "flux_units": "mmol/s",
         },
         "DONr": {
             "long_name": "refractory dissolved organic nitrogen",
             "units": "mmol/m^3",
+            "flux_units": "mmol/s",
         },
         "DOPr": {
             "long_name": "refractory dissolved organic phosphorus",
             "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "zooC": {
+            "long_name": "zooplankton carbon",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
         },
-        "zooC": {"long_name": "zooplankton carbon", "units": "mmol/m^3"},
         "spChl": {
             "long_name": "small phytoplankton chlorophyll",
             "units": "mg/m^3",
+            "flux_units": "mg/s",
+        },
+        "spC": {
+            "long_name": "small phytoplankton carbon",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
         },
-        "spC": {"long_name": "small phytoplankton carbon", "units": "mmol/m^3"},
         "spP": {
             "long_name": "small phytoplankton phosphorous",
             "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "spFe": {
+            "long_name": "small phytoplankton iron",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "spCaCO3": {
+            "long_name": "small phytoplankton CaCO3",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "diatChl": {
+            "long_name": "diatom chloropyll",
+            "units": "mg/m^3",
+            "flux_units": "mg/s",
+        },
+        "diatC": {
+            "long_name": "diatom carbon",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "diatP": {
+            "long_name": "diatom phosphorus",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "diatFe": {
+            "long_name": "diatom iron",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "diatSi": {
+            "long_name": "diatom silicate",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "diazChl": {
+            "long_name": "diazotroph chloropyll",
+            "units": "mg/m^3",
+            "flux_units": "mg/s",
+        },
+        "diazC": {
+            "long_name": "diazotroph carbon",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "diazP": {
+            "long_name": "diazotroph phosphorus",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
+        },
+        "diazFe": {
+            "long_name": "diazotroph iron",
+            "units": "mmol/m^3",
+            "flux_units": "mmol/s",
         },
-        "spFe": {"long_name": "small phytoplankton iron", "units": "mmol/m^3"},
-        "spCaCO3": {"long_name": "small phytoplankton CaCO3", "units": "mmol/m^3"},
-        "diatChl": {"long_name": "diatom chloropyll", "units": "mg/m^3"},
-        "diatC": {"long_name": "diatom carbon", "units": "mmol/m^3"},
-        "diatP": {"long_name": "diatom phosphorus", "units": "mmol/m^3"},
-        "diatFe": {"long_name": "diatom iron", "units": "mmol/m^3"},
-        "diatSi": {"long_name": "diatom silicate", "units": "mmol/m^3"},
-        "diazChl": {"long_name": "diazotroph chloropyll", "units": "mg/m^3"},
-        "diazC": {"long_name": "diazotroph carbon", "units": "mmol/m^3"},
-        "diazP": {"long_name": "diazotroph phosphorus", "units": "mmol/m^3"},
-        "diazFe": {"long_name": "diazotroph iron", "units": "mmol/m^3"},
         "pco2_air": {"long_name": "atmospheric pCO2", "units": "ppmv"},
         "pco2_air_alt": {
             "long_name": "atmospheric pCO2, alternative CO2",
@@ -522,7 +633,6 @@ def compute_missing_bgc_variables(bgc_data):
       dictionary, they are filled with constant values.
     - `CHL` is removed from the dictionary after the necessary calculations.
     """
-
     # Define the relationships for missing variables
     variable_relations = {
         "NH4": (None, 10**-6),  # mmol m-3
@@ -591,7 +701,6 @@ def compute_missing_surface_bgc_variables(bgc_data):
     -----
     - If `nox` and `nhy` are not part of the input dictionary, the are assigned constant values.
     """
-
     # Define the relationships for missing variables
     variable_relations = {
         "pco2_air_alt": ("pco2_air", 1.0),
@@ -610,21 +719,27 @@ def compute_missing_surface_bgc_variables(bgc_data):
     return bgc_data
 
 
-def get_tracer_metadata_dict(include_bgc=True):
-    """Returns a dictionary of tracer metadata.
+def get_tracer_metadata_dict(include_bgc=True, with_flux_units=False):
+    """Generate a dictionary containing metadata for model tracers.
 
-    This function constructs a dictionary containing the names, units, and long names
-    of the tracers, based on whether biogeochemical tracers should be included or not.
+    The returned dictionary maps tracer names to their associated units and long names.
+    Optionally includes biogeochemical tracers and can toggle between concentration and flux units.
 
     Parameters
     ----------
     include_bgc : bool, optional
-        Whether to include biogeochemical tracers. Defaults to True.
+        If True (default), includes biogeochemical tracers in the output.
+        If False, returns only physical tracers (e.g., temperature, salinity).
+
+    with_flux_units : bool, optional
+        If True, uses units appropriate for tracer fluxes (e.g., mmol/s).
+        If False (default), uses units appropriate for tracer concentrations (e.g., mmol/m³).
 
     Returns
     -------
     dict
-        A dictionary containing tracer names, units, and long names.
+        A dictionary where keys are tracer names and values are dictionaries
+        containing 'units' and 'long_name' for each tracer.
     """
     if include_bgc:
         tracer_names = [
@@ -670,7 +785,9 @@ def get_tracer_metadata_dict(include_bgc=True):
 
     tracer_dict = {
         tracer: {
-            "units": metadata[tracer]["units"],
+            "units": metadata[tracer]["flux_units"]
+            if with_flux_units
+            else metadata[tracer]["units"],
             "long_name": metadata[tracer]["long_name"],
         }
         for tracer in tracer_names
@@ -679,7 +796,7 @@ def get_tracer_metadata_dict(include_bgc=True):
     return tracer_dict
 
 
-def add_tracer_metadata_to_ds(ds, include_bgc=True):
+def add_tracer_metadata_to_ds(ds, include_bgc=True, with_flux_units=False):
     """Adds tracer metadata to a dataset.
 
     This function adds tracer metadata (name, unit, long name) as coordinates to
@@ -690,14 +807,18 @@ def add_tracer_metadata_to_ds(ds, include_bgc=True):
     ds : xarray.Dataset
         Dataset to which tracer metadata will be added.
     include_bgc : bool, optional
-        Whether to include biogeochemical tracers. Defaults to True.
+        If True (default), includes biogeochemical tracers in the output.
+        If False, returns only physical tracers (e.g., temperature, salinity).
+    with_flux_units : bool, optional
+        If True, uses units appropriate for tracer fluxes (e.g., mmol/s).
+        If False (default), uses units appropriate for tracer concentrations (e.g., mmol/m³).
 
     Returns
     -------
     xarray.Dataset
         The dataset with added tracer metadata.
     """
-    tracer_dict = get_tracer_metadata_dict(include_bgc)
+    tracer_dict = get_tracer_metadata_dict(include_bgc, with_flux_units)
 
     tracer_names = list(tracer_dict.keys())
     tracer_units = [tracer_dict[tracer]["units"] for tracer in tracer_names]
@@ -705,7 +826,15 @@ def add_tracer_metadata_to_ds(ds, include_bgc=True):
 
     ds = ds.assign_coords(
         tracer_name=("ntracers", tracer_names, {"long_name": "Tracer name"}),
-        tracer_unit=("ntracers", tracer_units, {"long_name": "Tracer unit"}),
+        tracer_unit=(
+            "ntracers",
+            tracer_units,
+            {
+                "long_name": "Tracer flux unit"
+                if with_flux_units
+                else "Tracer concentration unit"
+            },
+        ),
         tracer_long_name=(
             "ntracers",
             tracer_long_names,
@@ -716,7 +845,7 @@ def add_tracer_metadata_to_ds(ds, include_bgc=True):
     return ds
 
 
-def get_tracer_defaults():
+def get_tracer_defaults() -> dict[str, float]:
     """Returns constant default tracer concentrations for ROMS-MARBL.
 
     These values represent typical physical and biogeochemical tracer levels
@@ -810,7 +939,6 @@ def group_dataset(ds, filepath):
     tuple
         A tuple containing the list of grouped datasets and corresponding output filenames.
     """
-
     if hasattr(ds, "climatology"):
         output_filename = f"{filepath}_clim"
         output_filenames = [output_filename]
@@ -863,7 +991,6 @@ def group_by_month(ds, filepath):
     tuple
         A tuple containing the list of monthly datasets and corresponding output filenames.
     """
-
     dataset_list = []
     output_filenames = []
 
@@ -900,7 +1027,6 @@ def group_by_year(ds, filepath):
     tuple
         A tuple containing the list of yearly datasets and corresponding output filenames.
     """
-
     dataset_list = []
     output_filenames = []
 
@@ -1014,7 +1140,6 @@ def rotate_velocities(
       - u_rot = u * cos(angle) + v * sin(angle)
       - v_rot = v * cos(angle) - u * sin(angle)
     """
-
     # Rotate velocities to grid orientation
     u_rot = u * np.cos(angle) + v * np.sin(angle)
     v_rot = v * np.cos(angle) - u * np.sin(angle)
@@ -1051,7 +1176,6 @@ def compute_barotropic_velocity(
     Computed as:
       - `vel_bar` = sum(dz * vel) / sum(dz)
     """
-
     # Layer thickness
     dz = -interface_depth.diff(dim="s_w")
     dz = dz.rename({"s_w": "s_rho"})
@@ -1099,6 +1223,9 @@ def gc_dist(lon1, lat1, lon2, lat2, input_in_degrees=True):
     Latitude and longitude are assumed to be in degrees by default. If `input_in_degrees` is set to `False`,
     the input is assumed to already be in radians.
 
+    This function is a wrapper for two numba-vectorized versions of the function, one each for degrees and radians.
+    The wrapper is additionally needed to be able to use kwargs.
+
     Parameters
     ----------
     lon1, lat1 : float
@@ -1111,22 +1238,62 @@ def gc_dist(lon1, lat1, lon2, lat2, input_in_degrees=True):
 
     Returns
     -------
-    dis : float
+    dist : float
         The great circle distance between the two points in meters.
 
     Notes
     -----
     The radius of the Earth is taken to be 6371315 meters.
     """
+    if input_in_degrees:
+        return _gc_dist_degrees(lon1, lat1, lon2, lat2)
+    return _gc_dist_radians(lon1, lat1, lon2, lat2)
 
+
+@nb.vectorize(
+    [nb.float64(nb.float64, nb.float64, nb.float64, nb.float64)], nopython=True
+)
+def _gc_dist_degrees(lon1, lat1, lon2, lat2):
+    """Calculate the great circle distance, given lat and lon in degrees.
+
+    Returns
+    -------
+    Great circle distance in meters
+    """
     # Convert degrees to radians
-    if input_in_degrees:
-        d2r = np.pi / 180
-        lon1 = lon1 * d2r
-        lat1 = lat1 * d2r
-        lon2 = lon2 * d2r
-        lat2 = lat2 * d2r
+    d2r = np.pi / 180
+    lon1 = lon1 * d2r
+    lat1 = lat1 * d2r
+    lon2 = lon2 * d2r
+    lat2 = lat2 * d2r
+
+    # Difference in latitudes and longitudes
+    dlat = lat2 - lat1
+    dlon = lon2 - lon1
+
+    # Haversine formula
+    dang = 2 * np.arcsin(
+        np.sqrt(
+            np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
+        )
+    )
+
+    # Distance in meters
+    dis = R_EARTH * dang
 
+    return dis
+
+
+@nb.vectorize(
+    [nb.float64(nb.float64, nb.float64, nb.float64, nb.float64)], nopython=True
+)
+def _gc_dist_radians(lon1, lat1, lon2, lat2):
+    """Calculate the great circle distance, given lat and lon in radians.
+
+    Returns
+    -------
+    Great circle distance in meters
+    """
     # Difference in latitudes and longitudes
     dlat = lat2 - lat1
     dlon = lon2 - lon1
@@ -1144,25 +1311,130 @@ def gc_dist(lon1, lat1, lon2, lat2, input_in_degrees=True):
     return dis
 
 
-def convert_to_roms_time(ds, model_reference_date, climatology, time_name="time"):
+@nb.njit(
+    [
+        nb.float64[:, :](
+            nb.float64[:, :],
+            nb.float64[:, :],
+            nb.int32[:, :],
+        )
+    ],
+    parallel=True,
+)
+def min_dist_to_land(
+    lon: np.ndarray,
+    lat: np.ndarray,
+    mask: np.ndarray,
+):
+    """Calculate the distance between one set of points (lon1, lat1) to the closest of
+    another set of points (lon2, lat2).
+
+    Parameters
+    ----------
+    lon : np.ndarray
+        2-D Array of longitudes (in degrees) for all points on the grid
+    lat : np.ndarray
+        2-D Arrays of latitudes (in degrees) for all points on the grid
+    mask: np.ndarray
+        2-D integer array where ocean points have value 1 and land points are 0
+
+    Returns
+    -------
+    2-D Array of the same shape as lon and lat, which will be filled with the resulting distance values
+    to the nearest non-nan lon2, lat2 point
+    """
+    # get flattened ocean/land indices
+    ocean = (mask == 1).ravel()
+    land = (mask == 0).ravel()
+
+    # get flattened and separate lon/lat arrays for ocean and land
+    ocean_lon = lon.ravel()[ocean]
+    ocean_lat = lat.ravel()[ocean]
+    land_lon = lon.ravel()[land]
+    land_lat = lat.ravel()[land]
+
+    # keep track of the alignment between the full 2-D grid and the 1-D ocean indices
+    # (nonzero() returns a tuple of the i, j indices where mask is 1)
+    ocean_indices = mask.nonzero()
+
+    # create a results array that will hold the distances from each ocean point to the nearest land point
+    # initially fill arrays with zeros, as we will not do this calculation for land points, and land points
+    # have zero distance to land by definition.
+    result = np.zeros_like(lon)
+
+    # iterate in parallel and do the distance calculation, taking the min for each ocean point without needing to
+    # allocate a huge array for the entire calculation space
+    for i in nb.prange(ocean_lon.shape[0]):
+        result[ocean_indices[0][i], ocean_indices[1][i]] = np.min(
+            _gc_dist_degrees(ocean_lon[i], ocean_lat[i], land_lon, land_lat)
+        )
+
+    return result
+
+
+def convert_to_relative_days(
+    times: Sequence[datetime] | np.ndarray,
+    model_reference_date: datetime | np.datetime64,
+) -> np.ndarray:
+    """Convert absolute datetimes to model-relative time in days.
 
+    Parameters
+    ----------
+    times : sequence of datetime or np.ndarray
+        Absolute times to convert.
+    model_reference_date : datetime or np.datetime64
+        Reference date from which to compute relative days.
+
+    Returns
+    -------
+    np.ndarray
+        Times relative to the reference date, in days.
+    """
+    times = np.array(times, dtype="datetime64[ns]")
+    ref = np.datetime64(model_reference_date, "ns")
+    rel_times = (times - ref) / np.timedelta64(1, "D")
+
+    return rel_times
+
+
+def add_time_info_to_ds(
+    ds: xr.Dataset,
+    model_reference_date: datetime | np.datetime64,
+    climatology: bool,
+    time_name: str = "time",
+) -> tuple[xr.Dataset, xr.DataArray]:
+    """Add relative and absolute time coordinates to a dataset.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        The dataset to which time information will be added.
+    model_reference_date : datetime or np.datetime64
+        The reference date for computing relative time.
+    climatology : bool
+        Whether the time data is climatological (cyclical over the year).
+    time_name : str
+        Name of the time coordinate in the dataset.
+
+    Returns
+    -------
+    tuple[xr.Dataset, xr.DataArray]
+        Updated dataset with time information and the relative time array.
+    """
     if climatology:
         ds.attrs["climatology"] = str(True)
         month = xr.DataArray(range(1, 13), dims=time_name)
         month.attrs["long_name"] = "Month index (1-12)"
         ds = ds.assign_coords({"month": month})
-        # Preserve absolute time coordinate for readability
+
+        # Absolute time (for readability only)
         abs_time = np.datetime64(model_reference_date) + ds[time_name]
-        # Convert to pandas TimedeltaIndex
-        timedelta_index = pd.to_timedelta(ds[time_name].values)
 
-        # Determine the start of the year for the base_datetime
+        # Custom relative time logic for climatology
+        timedelta_index = pd.to_timedelta(ds[time_name].values)
         start_of_year = datetime(model_reference_date.year, 1, 1)
-
-        # Calculate the offset from midnight of the new year
         offset = model_reference_date - start_of_year
 
-        # Convert the timedelta to nanoseconds first, then to days
         time = xr.DataArray(
             (timedelta_index - offset).view("int64") / 3600 / 24 * 1e-9,
             dims=time_name,
@@ -1170,26 +1442,22 @@ def convert_to_roms_time(ds, model_reference_date, climatology, time_name="time"
         time.attrs["cycle_length"] = 365.25
 
     else:
-        # Preserve absolute time coordinate for readability
         abs_time = ds[time_name]
 
-        time = (
-            (ds[time_name] - np.datetime64(model_reference_date)).astype("float64")
-            / 3600
-            / 24
-            * 1e-9
+        time = xr.DataArray(
+            convert_to_relative_days(ds[time_name].values, model_reference_date),
+            dims=time_name,
         )
 
-    attrs = [key for key in abs_time.attrs]
-    for attr in attrs:
-        del abs_time.attrs[attr]
+    # Clean up and assign attributes
+    abs_time.attrs.clear()
     abs_time.attrs["long_name"] = "absolute time"
     ds = ds.assign_coords({"abs_time": abs_time})
 
-    time.attrs["long_name"] = f"relative time: days since {str(model_reference_date)}"
+    time.attrs["long_name"] = f"relative time: days since {model_reference_date!s}"
     time.encoding["units"] = "days"
     time.attrs["units"] = "days"
-    ds.encoding["unlimited_dims"] = "time"
+    ds.encoding["unlimited_dims"] = time_name
 
     return ds, time
 
@@ -1199,33 +1467,28 @@ class NoAliasDumper(yaml.SafeDumper):
         return True
 
 
-def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
-    """Serialize a forcing object (including its grid) into a YAML file.
+def write_to_yaml(yaml_data, filepath: str | Path) -> None:
+    """Write pre-serialized YAML data and additional metadata to a YAML file.
 
-    This function serializes a dataclass object (forcing_object) and its associated
-    `grid` attribute into a YAML file. It includes additional metadata, such as
-    the version of the `roms-tools` package.
+    This function writes the provided pre-serialized YAML data along with metadata, such as the version
+    of the `roms-tools` package, to the specified file. The metadata header is written first, followed by
+    the provided YAML data.
 
     Parameters
     ----------
-    forcing_object : object
-        The object that contains the forcing data, typically a dataclass with attributes
-        such as `grid`, `start_time`, `end_time`, etc.
+    yaml_data : dict or str
+        The pre-serialized YAML data to be written to the file. This data may include the forcing object and grid.
     filepath : Union[str, Path]
-        The path where the serialized YAML file will be saved.
+        The path (as a string or Path object) where the serialized YAML file will be saved.
 
     Returns
     -------
     None
-        The function writes the serialized data directly to a YAML file at the specified path.
+        This function does not return anything. It writes the provided YAML data directly to the specified file.
     """
-
     # Convert the filepath to a Path object
     filepath = Path(filepath)
 
-    # Serialize object into dictionary
-    yaml_data = _to_dict(forcing_object)
-
     # Create YAML header with version information
     try:
         roms_tools_version = importlib.metadata.version("roms-tools")
@@ -1248,7 +1511,7 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
         )
 
 
-def _to_dict(forcing_object) -> None:
+def to_dict(forcing_object, exclude: list[str] | None = None) -> dict:
     """Serialize a forcing object (including its grid) into a dictionary.
 
     This function serializes a dataclass object (forcing_object) and its associated
@@ -1263,19 +1526,20 @@ def _to_dict(forcing_object) -> None:
     forcing_object : object
         The object that contains the forcing data, typically a dataclass with attributes
         such as `grid`, `start_time`, `end_time`, etc.
+    exclude : list[str], optional
+        List of keys to exclude from the serialized output. Defaults to empty list. The field "ds" is always excluded by default.
 
     Returns
     -------
     dict
     """
-
     # Serialize Grid data
     if hasattr(forcing_object, "grid") and forcing_object.grid is not None:
         grid_data = asdict(forcing_object.grid)
-        grid_yaml_data = {"Grid": _pop_grid_data(grid_data)}
+        grid_yaml_data = {"Grid": pop_grid_data(grid_data)}
     elif hasattr(forcing_object, "parent_grid"):
         grid_data = asdict(forcing_object.parent_grid)
-        grid_yaml_data = {"ParentGrid": _pop_grid_data(grid_data)}
+        grid_yaml_data = {"ParentGrid": pop_grid_data(grid_data)}
 
     # Ensure Paths are Strings
     def ensure_paths_are_strings(obj, key):
@@ -1295,22 +1559,18 @@ def _to_dict(forcing_object) -> None:
 
     # Prepare Forcing Data
     forcing_data = {}
-    field_names = [field.name for field in fields(forcing_object)]
-    filtered_field_names = [
-        param
-        for param in field_names
-        if param
-        not in (
-            "grid",
-            "parent_grid",
-            "ds",
-            "use_dask",
-            "climatology",
-            "verbose",
-            "straddle",
-            "indices",
-        )
-    ]
+    if isinstance(forcing_object, BaseModel):
+        field_names = forcing_object.model_fields
+    elif is_dataclass(forcing_object):
+        field_names = [field.name for field in fields(forcing_object)]
+    else:
+        raise TypeError("Forcing object is not a dataclass or pydantic model")
+
+    if exclude is None:
+        exclude = []
+    exclude = ["grid", "parent_grid", "ds", *exclude]
+
+    filtered_field_names = [param for param in field_names if param not in exclude]
 
     for field_name in filtered_field_names:
         # Retrieve the value of each field using getattr
@@ -1326,20 +1586,6 @@ def _to_dict(forcing_object) -> None:
         # Add the field and its value to the forcing_data dictionary
         forcing_data[field_name] = value
 
-    # Serialize `indices` data (conditionally)
-    indices_data = getattr(forcing_object, "indices", None)
-    if indices_data is not None:
-        serialized_indices = {
-            "_convention": "eta_rho, xi_rho"
-        }  # Add convention metadata
-
-        for key, value in indices_data.items():
-            serialized_indices[key] = [
-                f"{tup[0]}, {tup[1]}" for tup in value
-            ]  # Comma-separated string
-
-        forcing_data["indices"] = serialized_indices
-
     # Combine Grid and Forcing Data into a single dictionary for the final YAML content
     yaml_data = {
         **grid_yaml_data,  # Add the grid data to the final YAML structure
@@ -1349,7 +1595,7 @@ def _to_dict(forcing_object) -> None:
     return yaml_data
 
 
-def _pop_grid_data(grid_data):
+def pop_grid_data(grid_data):
     grid_data.pop("ds", None)  # Remove 'ds' attribute (non-serializable)
     grid_data.pop("straddle", None)
     grid_data.pop("verbose", None)
@@ -1357,7 +1603,7 @@ def _pop_grid_data(grid_data):
     return grid_data
 
 
-def _from_yaml(forcing_object: Type, filepath: Union[str, Path]) -> Dict[str, Any]:
+def from_yaml(forcing_object: type, filepath: str | Path) -> dict[str, Any]:
     """Extract the configuration data for a given forcing object from a YAML file.
 
     This function reads a YAML file, searches for the configuration data associated
@@ -1444,7 +1690,6 @@ def handle_boundaries(field):
     field : numpy.ndarray or xarray.DataArray
         The input field with adjusted boundary values.
     """
-
     field[0, :] = field[1, :]
     field[-1, :] = field[-2, :]
     field[:, 0] = field[:, 1]
@@ -1471,7 +1716,6 @@ def get_boundary_coords():
           - "v" variables (e.g., `eta_v`)
           - "vector" variables with lists of indices for multiple grid points (e.g., `eta_rho`, `xi_rho`).
     """
-
     bdry_coords = {
         "rho": {
             "south": {"eta_rho": 0},