roms-tools 2.6.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

roms_tools/setup/utils.py

@@ -1,15 +1,26 @@
-import xarray as xr
+import importlib.metadata
+from dataclasses import asdict, fields, is_dataclass
+from datetime import datetime
+from enum import StrEnum
+from pathlib import Path
+from typing import Any, Dict, Sequence, Type, Union
+
+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.
@@ -409,8 +420,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"},
@@ -423,59 +438,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",
@@ -594,8 +712,8 @@ def compute_missing_surface_bgc_variables(bgc_data):
     # Define the relationships for missing variables
     variable_relations = {
         "pco2_air_alt": ("pco2_air", 1.0),
-        "nox": (None, 1e-13),  # has incorrect unit nmol/cm2/s, we need kg/m2/s
-        "nhy": (None, 5e-12),  # has incorrect unit nmol/cm2/s, we need kg/m2/s
+        "nox": (None, 1e-12),  # kg/m2/s
+        "nhy": (None, 5e-12),  # kg/m2/s
     }
 
     # Fill in missing variables using the defined relationships
@@ -609,6 +727,181 @@ def compute_missing_surface_bgc_variables(bgc_data):
     return bgc_data
 
 
+def get_tracer_metadata_dict(include_bgc=True, with_flux_units=False):
+    """Generate a dictionary containing metadata for model tracers.
+
+    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
+        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 where keys are tracer names and values are dictionaries
+        containing 'units' and 'long_name' for each tracer.
+    """
+    if include_bgc:
+        tracer_names = [
+            "temp",
+            "salt",
+            "PO4",
+            "NO3",
+            "SiO3",
+            "NH4",
+            "Fe",
+            "Lig",
+            "O2",
+            "DIC",
+            "DIC_ALT_CO2",
+            "ALK",
+            "ALK_ALT_CO2",
+            "DOC",
+            "DON",
+            "DOP",
+            "DOPr",
+            "DONr",
+            "DOCr",
+            "zooC",
+            "spChl",
+            "spC",
+            "spP",
+            "spFe",
+            "spCaCO3",
+            "diatChl",
+            "diatC",
+            "diatP",
+            "diatFe",
+            "diatSi",
+            "diazChl",
+            "diazC",
+            "diazP",
+            "diazFe",
+        ]
+    else:
+        tracer_names = ["temp", "salt"]
+
+    metadata = get_variable_metadata()
+
+    tracer_dict = {
+        tracer: {
+            "units": metadata[tracer]["flux_units"]
+            if with_flux_units
+            else metadata[tracer]["units"],
+            "long_name": metadata[tracer]["long_name"],
+        }
+        for tracer in tracer_names
+    }
+
+    return tracer_dict
+
+
+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
+    the provided dataset.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        Dataset to which tracer metadata will be added.
+    include_bgc : bool, optional
+        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, with_flux_units)
+
+    tracer_names = list(tracer_dict.keys())
+    tracer_units = [tracer_dict[tracer]["units"] for tracer in tracer_names]
+    tracer_long_names = [tracer_dict[tracer]["long_name"] for tracer in tracer_names]
+
+    ds = ds.assign_coords(
+        tracer_name=("ntracers", tracer_names, {"long_name": "Tracer name"}),
+        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,
+            {"long_name": "Tracer long name"},
+        ),
+    )
+
+    return ds
+
+
+def get_tracer_defaults() -> dict[str, float]:
+    """Returns constant default tracer concentrations for ROMS-MARBL.
+
+    These values represent typical physical and biogeochemical tracer levels
+    (e.g., temperature, salinity, nutrients, carbon) in freshwater.
+
+    Returns
+    -------
+    dict
+        Dictionary of tracer names and their default concentrations
+    """
+    return {
+        "temp": 17.0,  # degrees C
+        "salt": 1.0,  # psu
+        "PO4": 2.7,  # mmol m-3
+        "NO3": 24.2,  # mmol m-3
+        "SiO3": 13.2,  # mmol m-3
+        "NH4": 2.2,  # mmol m-3
+        "Fe": 1.79,  # mmol m-3
+        "Lig": 3 * 1.79,  # mmol m-3, inferred from Fe
+        "O2": 187.5,  # mmol m-3
+        "DIC": 2370.0,  # mmol m-3
+        "DIC_ALT_CO2": 2370.0,  # mmol m-3
+        "ALK": 2310.0,  # meq m-3
+        "ALK_ALT_CO2": 2310.0,  # meq m-3
+        "DOC": 1e-4,  # mmol m-3
+        "DON": 1.0,  # mmol m-3
+        "DOP": 0.1,  # mmol m-3
+        "DOPr": 0.003,  # mmol m-3
+        "DONr": 0.8,  # mmol m-3
+        "DOCr": 1e-6,  # mmol m-3
+        "zooC": 2.7,  # mmol m-3
+        "spChl": 1.35,  # mg m-3
+        "spC": 6.75,  # mmol m-3
+        "spP": 1.5 * 0.03,  # mmol m-3, inferred from ?
+        "spFe": 2.7e-5,  # mmol m-3
+        "spCaCO3": 0.135,  # mmol m-3
+        "diatChl": 0.135,  # mg m-3
+        "diatC": 0.405,  # mmol m-3
+        "diatP": 1.5 * 0.02,  # mmol m-3, inferred from ?
+        "diatFe": 2.7e-6,  # mmol m-3
+        "diatSi": 0.135,  # mmol m-3
+        "diazChl": 0.015,  # mg m-3
+        "diazC": 0.075,  # mmol m-3
+        "diazP": 1.5 * 0.01,  # mmol m-3, inferred from ?
+        "diazFe": 1.5e-6,  # mmol m-3
+    }
+
+
 def extract_single_value(data):
     """Extracts a single value from an xarray.DataArray or numpy array.
 
@@ -943,6 +1236,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
@@ -955,21 +1251,34 @@ 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
@@ -982,34 +1291,165 @@ def gc_dist(lon1, lat1, lon2, lat2, input_in_degrees=True):
         )
     )
 
-    # Radius of the Earth in meters
-    r_earth = 6371315.0
+    # 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
+
+    # 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
+    dis = R_EARTH * dang
 
     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: Union[Sequence[datetime], np.ndarray],
+    model_reference_date: Union[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: Union[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,
@@ -1017,36 +1457,81 @@ 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.encoding["units"] = "days"
     time.attrs["units"] = "days"
-    ds.encoding["unlimited_dims"] = "time"
+    ds.encoding["unlimited_dims"] = time_name
 
     return ds, time
 
 
-def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
-    """Serialize a forcing object (including its grid) into a YAML file.
+class NoAliasDumper(yaml.SafeDumper):
+    def ignore_aliases(self, data):
+        return True
+
+
+def _write_to_yaml(yaml_data, filepath: Union[str, Path]) -> None:
+    """Write pre-serialized YAML data and additional metadata to a YAML file.
+
+    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
+    ----------
+    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 (as a string or Path object) where the serialized YAML file will be saved.
+
+    Returns
+    -------
+    None
+        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)
+
+    # Create YAML header with version information
+    try:
+        roms_tools_version = importlib.metadata.version("roms-tools")
+    except importlib.metadata.PackageNotFoundError:
+        roms_tools_version = "unknown"
+
+    header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
+
+    # Write to YAML file
+    with filepath.open("w") as file:
+        # Write the header first
+        file.write(header)
+        # Write the serialized YAML data
+        yaml.dump(
+            yaml_data,
+            file,
+            Dumper=NoAliasDumper,
+            default_flow_style=False,
+            sort_keys=False,
+        )
+
+
+def _to_dict(forcing_object) -> dict:
+    """Serialize a forcing object (including its grid) into a dictionary.
 
     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, and omits fields like `grid` and `ds`
+    `grid` attribute into a dictionary. It omits fields like `grid` and `ds`
     that are not serializable or meant to be excluded.
 
     The function also converts datetime fields to ISO format strings for proper
@@ -1057,20 +1542,13 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
     forcing_object : object
         The object that contains the forcing data, typically a dataclass with attributes
         such as `grid`, `start_time`, `end_time`, etc.
-    filepath : Union[str, Path]
-        The path 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.
+    dict
     """
 
-    # Convert the filepath to a Path object
-    filepath = Path(filepath)
-
-    # Step 1: Serialize Grid data
-    # Check if the forcing_object has a grid attribute
+    # 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)}
@@ -1078,7 +1556,7 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
         grid_data = asdict(forcing_object.parent_grid)
         grid_yaml_data = {"ParentGrid": _pop_grid_data(grid_data)}
 
-    # Step 2: Ensure Paths are Strings
+    # Ensure Paths are Strings
     def ensure_paths_are_strings(obj, key):
         attr = getattr(obj, key, None)
         if attr is not None and "path" in attr:
@@ -1087,22 +1565,22 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
                 attr["path"] = [str(p) if isinstance(p, Path) else p for p in paths]
             elif isinstance(paths, Path):
                 attr["path"] = str(paths)
+            elif isinstance(paths, dict):
+                for key, path in paths.items():
+                    attr["path"][key] = str(path)
 
     ensure_paths_are_strings(forcing_object, "source")
     ensure_paths_are_strings(forcing_object, "bgc_source")
 
-    # Step 3: Get ROMS Tools version
-    try:
-        roms_tools_version = importlib.metadata.version("roms-tools")
-    except importlib.metadata.PackageNotFoundError:
-        roms_tools_version = "unknown"
-
-    # Create YAML header with version information
-    header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
-
-    # Step 4: Prepare Forcing Data
+    # Prepare Forcing Data
     forcing_data = {}
-    field_names = [field.name for field in fields(forcing_object)]
+    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")
+
     filtered_field_names = [
         param
         for param in field_names
@@ -1115,7 +1593,6 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
             "climatology",
             "verbose",
             "straddle",
-            "indices",
         )
     ]
 
@@ -1126,41 +1603,20 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
         # If the field is a datetime object, convert it to ISO format
         if isinstance(value, datetime):
             value = value.isoformat()
+        # Convert list of datetimes to list of ISO strings
+        elif isinstance(value, list) and all(isinstance(v, datetime) for v in value):
+            value = [v.isoformat() for v in value]
 
         # Add the field and its value to the forcing_data dictionary
         forcing_data[field_name] = value
 
-    # Step 5: 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
-
-    # Step 6: Combine Grid and Forcing Data into a single dictionary for the final YAML content
+    # 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
         forcing_object.__class__.__name__: forcing_data,  # Include the serialized forcing object data
     }
 
-    # Step 7: Write to YAML file
-    with filepath.open("w") as file:
-        # Write the header first
-        file.write(header)
-        # Write the serialized YAML data
-        yaml.dump(
-            yaml_data,
-            file,
-            default_flow_style=False,
-            sort_keys=False,
-        )
+    return yaml_data
 
 
 def _pop_grid_data(grid_data):
@@ -1346,3 +1802,21 @@ def wrap_longitudes(grid_ds, straddle):
             )
 
     return grid_ds
+
+
+def to_float(val):
+    """Convert a value or list of values to float.
+
+    Parameters
+    ----------
+    val : float, int, or list of float/int
+        A numeric value or a list of numeric values.
+
+    Returns
+    -------
+    float or list of float
+        The input value(s) converted to float type.
+    """
+    if isinstance(val, list):
+        return [float(v) for v in val]
+    return float(val)