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

roms_tools/setup/utils.py

@@ -3,6 +3,8 @@ import numpy as np
 from typing import Union
 import pandas as pd
 import cftime
+from roms_tools.utils import partition
+from pathlib import Path
 
 
 def nan_check(field, mask) -> None:
@@ -41,6 +43,28 @@ def nan_check(field, mask) -> None:
         )
 
 
+def substitute_nans_by_fillvalue(field, fill_value=0.0) -> xr.DataArray:
+    """
+    Replace NaN values in the field with a specified fill value.
+
+    This function replaces any NaN values in the input field with the provided fill value.
+
+    Parameters
+    ----------
+    field : xr.DataArray
+        The data array in which NaN values need to be replaced. This is typically an xarray.DataArray.
+    fill_value : scalar, optional
+        The value to use for replacing NaNs. Default is 0.0.
+
+    Returns
+    -------
+    xr.DataArray
+        The data array with NaN values replaced by the specified fill value.
+    """
+
+    return field.fillna(fill_value)
+
+
 def interpolate_from_rho_to_u(field, method="additive"):
 
     """
@@ -350,3 +374,377 @@ def convert_cftime_to_datetime(data_array: np.ndarray) -> np.ndarray:
         return np.datetime64(dt, "ns")
 
     return np.vectorize(convert_datetime)(data_array)
+
+
+def get_variable_metadata():
+    """
+    Retrieves metadata for commonly used variables in the dataset.
+
+    This function returns a dictionary containing the metadata for various variables, including long names
+    and units for each variable.
+
+    Returns
+    -------
+    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"},
+        "pot_Re": {"long_name": "Tidal potential, real part", "units": "m"},
+        "pot_Im": {"long_name": "Tidal potential, complex part", "units": "m"},
+        "u_Re": {
+            "long_name": "Tidal velocity in x-direction, real part",
+            "units": "m/s",
+        },
+        "u_Im": {
+            "long_name": "Tidal velocity in x-direction, complex part",
+            "units": "m/s",
+        },
+        "v_Re": {
+            "long_name": "Tidal velocity in y-direction, real part",
+            "units": "m/s",
+        },
+        "v_Im": {
+            "long_name": "Tidal velocity in y-direction, complex part",
+            "units": "m/s",
+        },
+        "uwnd": {"long_name": "10 meter wind in x-direction", "units": "m/s"},
+        "vwnd": {"long_name": "10 meter wind in y-direction", "units": "m/s"},
+        "swrad": {
+            "long_name": "downward short-wave (solar) radiation",
+            "units": "W/m^2",
+        },
+        "lwrad": {
+            "long_name": "downward long-wave (thermal) radiation",
+            "units": "W/m^2",
+        },
+        "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"},
+        "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"},
+        "w": {"long_name": "w-flux component", "units": "m/s"},
+        "ubar": {
+            "long_name": "vertically integrated u-flux component",
+            "units": "m/s",
+        },
+        "vbar": {
+            "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"},
+        "DIC_ALT_CO2": {
+            "long_name": "dissolved inorganic carbon, alternative CO2",
+            "units": "mmol/m^3",
+        },
+        "ALK": {"long_name": "alkalinity", "units": "meq/m^3"},
+        "ALK_ALT_CO2": {
+            "long_name": "alkalinity, alternative CO2",
+            "units": "meq/m^3",
+        },
+        "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",
+        },
+        "DONr": {
+            "long_name": "refractory dissolved organic nitrogen",
+            "units": "mmol/m^3",
+        },
+        "DOPr": {
+            "long_name": "refractory dissolved organic phosphorus",
+            "units": "mmol/m^3",
+        },
+        "zooC": {"long_name": "zooplankton carbon", "units": "mmol/m^3"},
+        "spChl": {
+            "long_name": "small phytoplankton chlorophyll",
+            "units": "mg/m^3",
+        },
+        "spC": {"long_name": "small phytoplankton carbon", "units": "mmol/m^3"},
+        "spP": {
+            "long_name": "small phytoplankton phosphorous",
+            "units": "mmol/m^3",
+        },
+        "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",
+            "units": "ppmv",
+        },
+        "iron": {"long_name": "iron decomposition", "units": "nmol/cm^2/s"},
+        "dust": {"long_name": "dust decomposition", "units": "kg/m^2/s"},
+        "nox": {"long_name": "NOx decomposition", "units": "kg/m^2/s"},
+        "nhy": {"long_name": "NHy decomposition", "units": "kg/m^2/s"},
+    }
+    return d
+
+
+def get_boundary_info():
+
+    """
+    This function provides information about the boundary points for the rho, u, and v
+    variables on the grid, specifying the indices for the south, east, north, and west
+    boundaries.
+
+    Returns
+    -------
+    dict
+        A dictionary where keys are variable types ("rho", "u", "v"), and values
+        are nested dictionaries mapping directions ("south", "east", "north", "west")
+        to the corresponding boundary coordinates.
+    """
+
+    # Boundary coordinates
+    bdry_coords = {
+        "rho": {
+            "south": {"eta_rho": 0},
+            "east": {"xi_rho": -1},
+            "north": {"eta_rho": -1},
+            "west": {"xi_rho": 0},
+        },
+        "u": {
+            "south": {"eta_rho": 0},
+            "east": {"xi_u": -1},
+            "north": {"eta_rho": -1},
+            "west": {"xi_u": 0},
+        },
+        "v": {
+            "south": {"eta_v": 0},
+            "east": {"xi_rho": -1},
+            "north": {"eta_v": -1},
+            "west": {"xi_rho": 0},
+        },
+    }
+
+    return bdry_coords
+
+
+def extract_single_value(data):
+    """
+    Extracts a single value from an xarray.DataArray or numpy array.
+
+    Parameters
+    ----------
+    data : xarray.DataArray or numpy.ndarray
+        The data from which to extract the single value.
+
+    Returns
+    -------
+    scalar
+        The single value contained in the array.
+
+    Raises
+    ------
+    ValueError
+        If the data contains more than one element or is not a recognized type.
+    """
+    # Convert xarray.DataArray to numpy array if necessary
+    if isinstance(data, xr.DataArray):
+        data = data.values
+
+    # Check that the data is a numpy array and contains only one element
+    if isinstance(data, np.ndarray) and data.size == 1:
+        return data.item()
+    else:
+        raise ValueError("Data must be a single-element array or DataArray.")
+
+
+def group_dataset(ds, filepath):
+    """
+    Group the dataset into monthly or yearly subsets based on the frequency of the data.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        The dataset to be grouped.
+    filepath : str
+        The base filename for the output files.
+
+    Returns
+    -------
+    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]
+        dataset_list = [ds]
+    else:
+        if len(ds["abs_time"]) > 2:
+            # Determine the frequency of the data
+            abs_time_freq = pd.infer_freq(ds["abs_time"].to_index())
+
+            if abs_time_freq in ["D", "H", "T", "S"]:  # Daily or higher frequency
+                dataset_list, output_filenames = group_by_month(ds, filepath)
+            else:
+                dataset_list, output_filenames = group_by_year(ds, filepath)
+        else:
+            # Convert time index to datetime if not already
+            abs_time_index = ds["abs_time"].to_index()
+            # Determine if the entries are in the same month
+            first_entry = abs_time_index[0]
+            last_entry = abs_time_index[-1]
+
+            if (
+                first_entry.year == last_entry.year
+                and first_entry.month == last_entry.month
+            ):
+                # Same month
+                dataset_list, output_filenames = group_by_month(ds, filepath)
+            else:
+                # Different months, group by year
+                dataset_list, output_filenames = group_by_year(ds, filepath)
+
+    return dataset_list, output_filenames
+
+
+def group_by_month(ds, filepath):
+    """
+    Group the dataset by month and generate filenames with 'YYYYMM' format.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        The dataset to be grouped.
+    filepath : str
+        The base filename for the output files.
+
+    Returns
+    -------
+    tuple
+        A tuple containing the list of monthly datasets and corresponding output filenames.
+    """
+
+    dataset_list = []
+    output_filenames = []
+
+    # Group dataset by year
+    grouped_by_year = ds.groupby("abs_time.year")
+
+    for year, yearly_dataset in grouped_by_year:
+        # Further group each yearly group by month
+        grouped_by_month = yearly_dataset.groupby("abs_time.month")
+
+        for month, monthly_dataset in grouped_by_month:
+            dataset_list.append(monthly_dataset)
+
+            # Format: "filepath_YYYYMM.nc"
+            year_month_str = f"{year}{month:02}"
+            output_filename = f"{filepath}_{year_month_str}"
+            output_filenames.append(output_filename)
+
+    return dataset_list, output_filenames
+
+
+def group_by_year(ds, filepath):
+    """
+    Group the dataset by year and generate filenames with 'YYYY' format.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        The dataset to be grouped.
+    filepath : str
+        The base filename for the output files.
+
+    Returns
+    -------
+    tuple
+        A tuple containing the list of yearly datasets and corresponding output filenames.
+    """
+
+    dataset_list = []
+    output_filenames = []
+
+    # Group dataset by year
+    grouped_by_year = ds.groupby("abs_time.year")
+
+    for year, yearly_dataset in grouped_by_year:
+        dataset_list.append(yearly_dataset)
+
+        # Format: "filepath_YYYY.nc"
+        year_str = f"{year}"
+        output_filename = f"{filepath}_{year_str}"
+        output_filenames.append(output_filename)
+
+    return dataset_list, output_filenames
+
+
+def save_datasets(dataset_list, output_filenames, np_eta=None, np_xi=None):
+    """
+    Save the list of datasets to netCDF4 files, with optional spatial partitioning.
+
+    Parameters
+    ----------
+    dataset_list : list
+        List of datasets to be saved.
+    output_filenames : list
+        List of filenames for the output files.
+    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
+    -------
+    List[Path]
+        A list of Path objects for the filenames that were saved.
+    """
+
+    saved_filenames = []
+
+    if np_eta is None and np_xi is None:
+        # Save the dataset as a single file
+        output_filenames = [f"{filename}.nc" for filename in output_filenames]
+        xr.save_mfdataset(dataset_list, output_filenames)
+
+        saved_filenames.extend(Path(f) for f in output_filenames)
+
+    else:
+        # Partition the dataset and save each partition as a separate file
+        np_eta = np_eta or 1
+        np_xi = np_xi or 1
+
+        partitioned_datasets = []
+        partitioned_filenames = []
+        for dataset, base_filename in zip(dataset_list, output_filenames):
+            partition_indices, partitions = partition(
+                dataset, np_eta=np_eta, np_xi=np_xi
+            )
+            partition_filenames = [
+                f"{base_filename}.{index}.nc" for index in partition_indices
+            ]
+            partitioned_datasets.extend(partitions)
+            partitioned_filenames.extend(partition_filenames)
+
+        xr.save_mfdataset(partitioned_datasets, partitioned_filenames)
+
+        saved_filenames.extend(Path(f) for f in partitioned_filenames)
+
+    return saved_filenames

roms_tools/setup/vertical_coordinate.py

@@ -114,4 +114,9 @@ def compute_depth(zeta, h, hc, cs, sigma):
     s = (hc * sigma + h * cs) / (hc + h)
     z = zeta + (zeta + h) * s
 
+    if "s_rho" in z.dims:
+        z = z.transpose("s_rho", "eta_rho", "xi_rho")
+    elif "s_w" in z.dims:
+        z = z.transpose("s_w", "eta_rho", "xi_rho")
+
     return z