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

roms_tools/setup/utils.py

@@ -1,6 +1,6 @@
 import xarray as xr
 import numpy as np
-from typing import Union, Any, Dict, Type
+from typing import Union, Any, Dict, Type, Sequence
 import pandas as pd
 import cftime
 from pathlib import Path
@@ -8,6 +8,7 @@ from datetime import datetime
 from dataclasses import fields, asdict
 import importlib.metadata
 import yaml
+from roms_tools.constants import R_EARTH
 from roms_tools.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
 
 
@@ -125,41 +126,108 @@ def assign_dates_to_climatology(ds: xr.Dataset, time_dim: str) -> xr.Dataset:
     return ds
 
 
+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]]
+    ],
+) -> xr.DataArray:
+    """Interpolates a DataArray cyclically across the start and end of the year.
+
+    This function extends the data cyclically by appending the last time step
+    (shifted back by one year) at the beginning and the first time step
+    (shifted forward by one year) at the end. It then performs linear interpolation
+    to match the specified `day_of_year` values.
+
+    Parameters
+    ----------
+    data_array : xr.DataArray
+        The input data array containing a time-like dimension.
+    time_dim_name : str
+        The name of the time dimension in the dataset.
+    day_of_year : Union[int, float, np.ndarray, xr.DataArray, Sequence[Union[int, float]]]
+        The target day(s) of the year for interpolation. This can be:
+        - A single integer or float representing the day of the year.
+        - A NumPy array or xarray DataArray containing multiple days.
+        - A list or tuple of integers or floats for multiple target days.
+
+    Returns
+    -------
+    xr.DataArray
+        The interpolated DataArray, ensuring cyclic continuity across year boundaries.
+
+    Notes
+    -----
+    - This function is useful for interpolating climatological data, where the time axis
+      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(
+        [
+            data_array[time_dim_name][-1] - 365.25,  # Shift last time backward
+            data_array[time_dim_name],
+            data_array[time_dim_name][0] + 365.25,  # Shift first time forward
+        ],
+        dim=time_dim_name,
+    )
+
+    data_array_concat = xr.concat(
+        [
+            data_array.isel(
+                **{time_dim_name: -1}
+            ),  # Append last value at the beginning
+            data_array,
+            data_array.isel(**{time_dim_name: 0}),  # Append first value at the end
+        ],
+        dim=time_dim_name,
+    )
+    data_array_concat[time_dim_name] = time_concat
+
+    # Interpolate to specified times
+    data_array_interpolated = data_array_concat.interp(
+        **{time_dim_name: day_of_year}, method="linear"
+    )
+
+    return data_array_interpolated
+
+
 def interpolate_from_climatology(
     field: Union[xr.DataArray, xr.Dataset],
     time_dim_name: str,
     time: Union[xr.DataArray, pd.DatetimeIndex],
 ) -> Union[xr.DataArray, xr.Dataset]:
-    """Temporally interpolates a field based on specified time points.
+    """Interpolates a climatological field to specified time points.
 
-    This function performs temporal interpolation on the input `field` to match the provided `time` values.
-    If the input `field` is an `xarray.Dataset`, the interpolation is applied to all its data variables individually.
+    This function interpolates the input `field` based on `day_of_year` values
+    extracted from the provided `time` points. If `field` is an `xarray.Dataset`,
+    interpolation is applied to all its data variables individually.
 
     Parameters
     ----------
     field : xarray.DataArray or xarray.Dataset
         The input field to be interpolated.
-        - If `field` is an `xarray.DataArray`, it should have a time dimension identified by `time_dim_name`.
-        - If `field` is an `xarray.Dataset`, all variables within the dataset are interpolated along the specified time dimension.
+        - If `field` is an `xarray.DataArray`, it must have a time dimension identified by `time_dim_name`.
+        - If `field` is an `xarray.Dataset`, all variables within the dataset are interpolated along `time_dim_name`.
         The time dimension is assumed to represent `day_of_year` for climatological purposes.
     time_dim_name : str
-        The name of the time dimension in the `field`. This dimension is used for interpolation.
+        The name of the time dimension in `field`. This dimension is used for interpolation.
     time : xarray.DataArray or pandas.DatetimeIndex
-        The target time points for interpolation. The time values should be compatible with the time format used in the `field`.
+        The target time points for interpolation. These are internally converted to `day_of_year`
+        before performing interpolation.
 
     Returns
     -------
     xarray.DataArray or xarray.Dataset
-        The interpolated field, with the same type as the input (`xarray.DataArray` or `xarray.Dataset`),
+        The interpolated field, maintaining the same type (`xarray.DataArray` or `xarray.Dataset`)
         but aligned to the specified `time` values.
 
     Notes
     -----
-    - The interpolation assumes the time dimension in `field` corresponds to `day_of_year`.
-      If the input time values are in a datetime format, ensure they are converted to `day_of_year` before calling this function.
-      For example, you can preprocess the time as follows:
-
-      >>> field["time"] = field["time"].dt.dayofyear
+    - This function assumes that `field` represents a climatological dataset, where time is expressed as `day_of_year` (1-365).
+    - The `time` input is automatically converted to `day_of_year`, so manual conversion is not required before calling this function.
     """
 
     def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
@@ -170,33 +238,27 @@ def interpolate_from_climatology(
         else:
             if np.size(time) == 1:
                 # Convert single datetime64 object to pandas.Timestamp
-                day_of_year = pd.Timestamp(time).dayofyear
+                date = pd.Timestamp(time)
+                day_of_year = (
+                    date.dayofyear
+                    + (date.hour / 24)
+                    + (date.minute / 1440)
+                    + (date.second / 86400)
+                )
             else:
-                # Convert each datetime64 object in the array to pandas.Timestamp
-                day_of_year = np.array([pd.Timestamp(t).dayofyear for t in time])
-
-        # Concatenate across the beginning and end of the year
-        time_concat = xr.concat(
-            [
-                data_array[time_dim_name][-1] - 365.25,
-                data_array[time_dim_name],
-                365.25 + data_array[time_dim_name][0],
-            ],
-            dim=time_dim_name,
-        )
-        data_array_concat = xr.concat(
-            [
-                data_array.isel(**{time_dim_name: -1}),
-                data_array,
-                data_array.isel(**{time_dim_name: 0}),
-            ],
-            dim=time_dim_name,
-        )
-        data_array_concat[time_dim_name] = time_concat
-
-        # Interpolate to specified times
-        data_array_interpolated = data_array_concat.interp(
-            **{time_dim_name: day_of_year}, method="linear"
+                # Convert each datetime64 object in the array to pandas.Timestamp and compute fractional day of year
+                day_of_year = np.array(
+                    [
+                        pd.Timestamp(t).dayofyear
+                        + (pd.Timestamp(t).hour / 24)
+                        + (pd.Timestamp(t).minute / 1440)
+                        + (pd.Timestamp(t).second / 86400)
+                        for t in time
+                    ]
+                )
+
+        data_array_interpolated = interpolate_cyclic_time(
+            data_array, time_dim_name, day_of_year
         )
 
         if np.size(time) == 1:
@@ -428,6 +490,281 @@ def get_variable_metadata():
     return d
 
 
+def compute_missing_bgc_variables(bgc_data):
+    """Fills in missing biogeochemical (BGC) variables in the input dictionary.
+
+    This function checks for missing BGC variables in the provided dictionary and
+    computes them based on predefined relationships with existing variables. The
+    relationships specify either a multiplication factor applied to an existing
+    variable or a constant value if no related variable is available. The resulting
+    variables are added to the dictionary.
+
+    Parameters
+    ----------
+    bgc_data : dict
+        A dictionary containing biogeochemical variables as xarray DataArrays.
+        Missing variables are computed and added to this dictionary.
+
+        Assumptions:
+        - If `Fe` is part of the input dictionary, it is in units of mmol m-3.
+        - If `CHL` is part of the input dictionary, it is in units of mg m-3.
+        - If `ALK` is part of the input dictionary, it is in units of meq m-3 = mmol m-3.
+        - If `DIC` is part of the input dictionary, it is in units of mmol m-3.
+
+    Returns
+    -------
+    dict
+        The updated dictionary with missing BGC variables filled in.
+
+    Notes
+    -----
+    - If `NH4`, `DOC`, `DON`, `DOP`, `DOCr`, `DONr`, and `DOPr` are not part of the input
+      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
+        "Lig": ("Fe", 3),  # mmol m-3
+        "DIC_ALT_CO2": ("DIC", 1),  # mmol m-3
+        "ALK_ALT_CO2": ("ALK", 1),  # meq m-3 = mmol m-3
+        "DOC": (None, 10**-6),  # mmol m-3
+        "DON": (None, 1.0),  # mmol m-3
+        "DOP": (None, 0.1),  # mmol m-3
+        "DOCr": (None, 10**-6),  # mmol m-3
+        "DONr": (None, 0.8),  # mmol m-3
+        "DOPr": (None, 0.003),  # mmol m-3
+        "zooC": ("CHL", 1.35),  # mmol m-3
+        "spChl": ("CHL", 0.675),  # mg m-3
+        "spC": ("CHL", 3.375),  # mmol m-3
+        "spP": ("CHL", 0.03),  # mmol m-3
+        "spFe": ("CHL", 1.35e-5),  # mmol m-3
+        "spCaCO3": ("CHL", 0.0675),  # mmol m-3
+        "diatChl": ("CHL", 0.0675),  # mg m-3
+        "diatC": ("CHL", 0.2025),  # mmol m-3
+        "diatP": ("CHL", 0.02),  # mmol m-3
+        "diatFe": ("CHL", 1.35e-6),  # mmol m-3
+        "diatSi": ("CHL", 0.0675),  # mmol m-3
+        "diazChl": ("CHL", 0.0075),  # mg m-3
+        "diazC": ("CHL", 0.0375),  # mmol m-3
+        "diazP": ("CHL", 0.01),  # mmol m-3
+        "diazFe": ("CHL", 7.5e-7),  # mmol m-3
+    }
+
+    # Fill in missing variables using the defined relationships
+    for var_name, (base_var, factor) in variable_relations.items():
+        if var_name not in bgc_data:
+            if base_var:
+                bgc_data[var_name] = bgc_data[base_var] * factor
+            else:
+                bgc_data[var_name] = factor * xr.ones_like(bgc_data["ALK"])
+
+    bgc_data.pop("CHL", None)
+
+    return bgc_data
+
+
+def compute_missing_surface_bgc_variables(bgc_data):
+    """Fills in missing surface biogeochemical (BGC) variables in the input dictionary.
+
+    This function checks for missing surface BGC variables in the provided dictionary and
+    computes them based on predefined relationships with existing variables. The relationships
+    specify either a multiplication factor applied to an existing variable or a constant value
+    if no related variable is available. The resulting variables are added to the dictionary.
+
+    Parameters
+    ----------
+    bgc_data : dict
+        A dictionary containing surface biogeochemical variables as xarray DataArrays.
+        Missing variables are computed and added to this dictionary.
+
+        Assumptions:
+        - If `pco2_air` is part of the input dictionary, it is in units of ppmv.
+
+    Returns
+    -------
+    dict
+        The updated dictionary with missing surface BGC variables filled in.
+
+    Notes
+    -----
+    - 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),
+        "nox": (None, 1e-12),  # kg/m2/s
+        "nhy": (None, 5e-12),  # kg/m2/s
+    }
+
+    # Fill in missing variables using the defined relationships
+    for var_name, (base_var, factor) in variable_relations.items():
+        if var_name not in bgc_data:
+            if base_var:
+                bgc_data[var_name] = bgc_data[base_var] * factor
+            else:
+                bgc_data[var_name] = factor * xr.ones_like(bgc_data["pco2_air"])
+
+    return bgc_data
+
+
+def get_tracer_metadata_dict(include_bgc=True):
+    """Returns a dictionary of tracer metadata.
+
+    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.
+
+    Parameters
+    ----------
+    include_bgc : bool, optional
+        Whether to include biogeochemical tracers. Defaults to True.
+
+    Returns
+    -------
+    dict
+        A dictionary containing tracer names, units, and long names.
+    """
+    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]["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):
+    """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
+        Whether to include biogeochemical tracers. Defaults to True.
+
+    Returns
+    -------
+    xarray.Dataset
+        The dataset with added tracer metadata.
+    """
+    tracer_dict = get_tracer_metadata_dict(include_bgc)
+
+    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 unit"}),
+        tracer_long_name=(
+            "ntracers",
+            tracer_long_names,
+            {"long_name": "Tracer long name"},
+        ),
+    )
+
+    return ds
+
+
+def get_tracer_defaults():
+    """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.
 
@@ -801,11 +1138,8 @@ 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
+    dis = R_EARTH * dang
 
     return dis
 
@@ -860,16 +1194,17 @@ def convert_to_roms_time(ds, model_reference_date, climatology, time_name="time"
     return ds, time
 
 
+class NoAliasDumper(yaml.SafeDumper):
+    def ignore_aliases(self, data):
+        return True
+
+
 def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
     """Serialize a forcing object (including its grid) into 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, and 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
-    serialization.
+    the version of the `roms-tools` package.
 
     Parameters
     ----------
@@ -888,8 +1223,53 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
     # 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 object into dictionary
+    yaml_data = _to_dict(forcing_object)
+
+    # 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) -> None:
+    """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 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
+    serialization.
+
+    Parameters
+    ----------
+    forcing_object : object
+        The object that contains the forcing data, typically a dataclass with attributes
+        such as `grid`, `start_time`, `end_time`, etc.
+
+    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)}
@@ -897,7 +1277,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:
@@ -906,20 +1286,14 @@ 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)]
     filtered_field_names = [
@@ -945,11 +1319,14 @@ 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)
+    # Serialize `indices` data (conditionally)
     indices_data = getattr(forcing_object, "indices", None)
     if indices_data is not None:
         serialized_indices = {
@@ -963,23 +1340,13 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
 
         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):
@@ -1165,3 +1532,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)