roms-tools 2.2.1__py3-none-any.whl → 2.4.0__py3-none-any.whl

roms_tools/setup/tides.py

@@ -3,9 +3,13 @@ import xarray as xr
 import numpy as np
 from typing import Dict, Union, List
 import importlib.metadata
+import matplotlib.pyplot as plt
+from pathlib import Path
 from dataclasses import dataclass, field
-from roms_tools.setup.grid import Grid
-from roms_tools.setup.plot import _plot
+from roms_tools import Grid
+from roms_tools.plot import _plot
+from roms_tools.regrid import LateralRegrid
+from roms_tools.utils import save_datasets
 from roms_tools.setup.datasets import TPXODataset
 from roms_tools.setup.utils import (
     nan_check,
@@ -13,16 +17,12 @@ from roms_tools.setup.utils import (
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
     get_variable_metadata,
-    save_datasets,
     get_target_coords,
     rotate_velocities,
     get_vector_pairs,
     _to_yaml,
     _from_yaml,
 )
-from roms_tools.setup.regrid import LateralRegrid
-import matplotlib.pyplot as plt
-from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -319,6 +319,8 @@ class TidalForcing:
         >>> tidal_forcing.plot("ssh_Re", nc=0)
         """
 
+        if var_name not in self.ds:
+            raise ValueError(f"Variable '{var_name}' is not found in dataset.")
         field = self.ds[var_name].isel(ntides=ntides)
 
         if self.use_dask:
@@ -328,25 +330,28 @@ class TidalForcing:
                 field = field.load()
 
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            field = field.where(self.grid.ds.mask_rho)
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_rho, "lat": self.grid.ds.lat_rho}
-            )
+            lon_deg = self.grid.ds["lon_rho"]
+            lat_deg = self.grid.ds["lat_rho"]
+            mask = self.grid.ds["mask_rho"]
 
         elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            field = field.where(self.grid.ds.mask_u)
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_u, "lat": self.grid.ds.lat_u}
-            )
+            lon_deg = self.grid.ds["lon_u"]
+            lat_deg = self.grid.ds["lat_u"]
+            mask = self.grid.ds["mask_u"]
 
         elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            field = field.where(self.grid.ds.mask_v)
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_v, "lat": self.grid.ds.lat_v}
-            )
+            lon_deg = self.grid.ds["lon_v"]
+            lat_deg = self.grid.ds["lat_v"]
+            mask = self.grid.ds["mask_v"]
+
         else:
             ValueError("provided field does not have two horizontal dimension")
 
+        field = field.where(mask)
+        if self.grid.straddle:
+            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
+
         title = "%s, ntides = %i" % (field.long_name, self.ds[var_name].ntides[ntides])
 
         vmax = max(field.max(), -field.min())
@@ -357,44 +362,25 @@ class TidalForcing:
         kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
         _plot(
-            self.grid.ds,
             field=field,
-            straddle=self.grid.straddle,
+            title=title,
             c="g",
             kwargs=kwargs,
-            title=title,
         )
 
-    def save(
-        self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
-    ) -> None:
+    def save(self, filepath: Union[str, Path]) -> None:
         """Save the tidal forcing information to a netCDF4 file.
 
-        This method supports saving the dataset in two modes:
-
-          1. **Single File Mode (default)**:
-
-            If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single netCDF4 file
-            with the base filename specified by `filepath.nc`.
-
-          2. **Partitioned Mode**:
-
-            - If either `np_eta` or `np_xi` is specified, the dataset is divided into spatial tiles along the eta-axis and xi-axis.
-            - Each spatial tile is saved as a separate netCDF4 file.
-
         Parameters
         ----------
         filepath : Union[str, Path]
-            The base path or filename where the dataset should be saved.
-        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.
+            The path or filename where the dataset will be saved. If a directory is specified,
+            the file will be saved with a default name within that directory.
 
         Returns
         -------
-        List[Path]
-            A list of Path objects for the filenames that were saved.
+        Path
+            A `Path` object representing the location of the saved file.
         """
 
         # Ensure filepath is a Path object
@@ -404,17 +390,11 @@ class TidalForcing:
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
 
-        if self.use_dask:
-            from dask.diagnostics import ProgressBar
-
-            with ProgressBar():
-                self.ds.load()
-
         dataset_list = [self.ds]
         output_filenames = [str(filepath)]
 
         saved_filenames = save_datasets(
-            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
+            dataset_list, output_filenames, use_dask=self.use_dask
         )
 
         return saved_filenames
@@ -464,7 +444,7 @@ class TidalForcing:
             grid=grid,
             **tidal_forcing_params,
             use_dask=use_dask,
-            bypass_validation=bypass_validation
+            bypass_validation=bypass_validation,
         )
 
     def _correct_tides(self, data):

roms_tools/setup/topography.py

@@ -6,8 +6,8 @@ import gcm_filters
 from roms_tools.setup.utils import handle_boundaries
 import warnings
 from itertools import count
+from roms_tools.regrid import LateralRegrid
 from roms_tools.setup.datasets import ETOPO5Dataset, SRTM15Dataset
-from roms_tools.setup.regrid import LateralRegrid
 
 
 def _add_topography(

roms_tools/setup/utils.py

@@ -3,12 +3,12 @@ import numpy as np
 from typing import Union, Any, Dict, Type
 import pandas as pd
 import cftime
-from roms_tools.utils import partition
 from pathlib import Path
 from datetime import datetime
 from dataclasses import fields, asdict
 import importlib.metadata
 import yaml
+from roms_tools.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
 
 
 def nan_check(field, mask, error_message=None) -> None:
@@ -71,100 +71,6 @@ def substitute_nans_by_fillvalue(field, fill_value=0.0) -> xr.DataArray:
     return field.fillna(fill_value)
 
 
-def interpolate_from_rho_to_u(field, method="additive"):
-    """Interpolates the given field from rho points to u points.
-
-    This function performs an interpolation from the rho grid (cell centers) to the u grid
-    (cell edges in the xi direction). Depending on the chosen method, it either averages
-    (additive) or multiplies (multiplicative) the field values between adjacent rho points
-    along the xi dimension. It also handles the removal of unnecessary coordinate variables
-    and updates the dimensions accordingly.
-
-    Parameters
-    ----------
-    field : xr.DataArray
-        The input data array on the rho grid to be interpolated. It is assumed to have a dimension
-        named "xi_rho".
-
-    method : str, optional, default='additive'
-        The method to use for interpolation. Options are:
-        - 'additive': Average the field values between adjacent rho points.
-        - 'multiplicative': Multiply the field values between adjacent rho points. Appropriate for
-          binary masks.
-
-    Returns
-    -------
-    field_interpolated : xr.DataArray
-        The interpolated data array on the u grid with the dimension "xi_u".
-    """
-
-    if method == "additive":
-        field_interpolated = 0.5 * (field + field.shift(xi_rho=1)).isel(
-            xi_rho=slice(1, None)
-        )
-    elif method == "multiplicative":
-        field_interpolated = (field * field.shift(xi_rho=1)).isel(xi_rho=slice(1, None))
-    else:
-        raise NotImplementedError(f"Unsupported method '{method}' specified.")
-
-    vars_to_drop = ["lat_rho", "lon_rho", "eta_rho", "xi_rho"]
-    for var in vars_to_drop:
-        if var in field_interpolated.coords:
-            field_interpolated = field_interpolated.drop_vars(var)
-
-    field_interpolated = field_interpolated.swap_dims({"xi_rho": "xi_u"})
-
-    return field_interpolated
-
-
-def interpolate_from_rho_to_v(field, method="additive"):
-    """Interpolates the given field from rho points to v points.
-
-    This function performs an interpolation from the rho grid (cell centers) to the v grid
-    (cell edges in the eta direction). Depending on the chosen method, it either averages
-    (additive) or multiplies (multiplicative) the field values between adjacent rho points
-    along the eta dimension. It also handles the removal of unnecessary coordinate variables
-    and updates the dimensions accordingly.
-
-    Parameters
-    ----------
-    field : xr.DataArray
-        The input data array on the rho grid to be interpolated. It is assumed to have a dimension
-        named "eta_rho".
-
-    method : str, optional, default='additive'
-        The method to use for interpolation. Options are:
-        - 'additive': Average the field values between adjacent rho points.
-        - 'multiplicative': Multiply the field values between adjacent rho points. Appropriate for
-          binary masks.
-
-    Returns
-    -------
-    field_interpolated : xr.DataArray
-        The interpolated data array on the v grid with the dimension "eta_v".
-    """
-
-    if method == "additive":
-        field_interpolated = 0.5 * (field + field.shift(eta_rho=1)).isel(
-            eta_rho=slice(1, None)
-        )
-    elif method == "multiplicative":
-        field_interpolated = (field * field.shift(eta_rho=1)).isel(
-            eta_rho=slice(1, None)
-        )
-    else:
-        raise NotImplementedError(f"Unsupported method '{method}' specified.")
-
-    vars_to_drop = ["lat_rho", "lon_rho", "eta_rho", "xi_rho"]
-    for var in vars_to_drop:
-        if var in field_interpolated.coords:
-            field_interpolated = field_interpolated.drop_vars(var)
-
-    field_interpolated = field_interpolated.swap_dims({"eta_rho": "eta_v"})
-
-    return field_interpolated
-
-
 def one_dim_fill(da: xr.DataArray, dim: str, direction="forward") -> xr.DataArray:
     """Fill NaN values in a DataArray along a specified dimension.
 
@@ -224,23 +130,36 @@ def interpolate_from_climatology(
     time_dim_name: str,
     time: Union[xr.DataArray, pd.DatetimeIndex],
 ) -> Union[xr.DataArray, xr.Dataset]:
-    """Interpolates the given field temporally based on the specified time points.
+    """Temporally interpolates a field based on specified time points.
 
-    If `field` is an xarray.Dataset, this function applies the interpolation to all data variables in the dataset.
+    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.
 
     Parameters
     ----------
     field : xarray.DataArray or xarray.Dataset
-        The field data to be interpolated. Can be a single DataArray or a 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.
+        The time dimension is assumed to represent `day_of_year` for climatological purposes.
     time_dim_name : str
-        The name of the dimension in `field` that represents time.
+        The name of the time dimension in the `field`. This dimension is used for interpolation.
     time : xarray.DataArray or pandas.DatetimeIndex
-        The target time points for interpolation.
+        The target time points for interpolation. The time values should be compatible with the time format used in the `field`.
 
     Returns
     -------
     xarray.DataArray or xarray.Dataset
-        The field values interpolated to the specified time points. The type matches the input type.
+        The interpolated field, with the same type as the input (`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
     """
 
     def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
@@ -250,11 +169,11 @@ def interpolate_from_climatology(
             day_of_year = time.dt.dayofyear
         else:
             if np.size(time) == 1:
-                day_of_year = time.timetuple().tm_yday
+                # Convert single datetime64 object to pandas.Timestamp
+                day_of_year = pd.Timestamp(time).dayofyear
             else:
-                day_of_year = np.array([t.timetuple().tm_yday for t in time])
-
-        data_array[time_dim_name] = data_array[time_dim_name].dt.days
+                # 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(
@@ -294,6 +213,7 @@ def interpolate_from_climatology(
             for var, data_array in field.data_vars.items()
         }
         return xr.Dataset(interpolated_data_vars, attrs=field.attrs)
+
     else:
         raise TypeError("Input 'field' must be an xarray.DataArray or xarray.Dataset.")
 
@@ -661,59 +581,6 @@ def group_by_year(ds, filepath):
     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
-
-
 def get_target_coords(grid, use_coarse_grid=False):
     """Retrieves longitude and latitude coordinates from the grid, adjusting them based
     on longitude range.
@@ -745,8 +612,6 @@ def get_target_coords(grid, use_coarse_grid=False):
         mask = grid.ds.get("mask_coarse")
         if mask is not None:
             mask = mask.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
-            mask_u = interpolate_from_rho_to_u(mask, method="multiplicative")
-            mask_v = interpolate_from_rho_to_v(mask, method="multiplicative")
 
         lat_psi = grid.ds.get("lat_psi_coarse")
         lon_psi = grid.ds.get("lon_psi_coarse")
@@ -756,8 +621,6 @@ def get_target_coords(grid, use_coarse_grid=False):
         lon = grid.ds.lon_rho
         angle = grid.ds.angle
         mask = grid.ds.get("mask_rho")
-        mask_u = grid.ds.get("mask_u")
-        mask_v = grid.ds.get("mask_v")
         lat_psi = grid.ds.get("lat_psi")
         lon_psi = grid.ds.get("lon_psi")
 
@@ -780,8 +643,6 @@ def get_target_coords(grid, use_coarse_grid=False):
         "lon_psi": lon_psi,
         "angle": angle,
         "mask": mask,
-        "mask_u": mask_u,
-        "mask_v": mask_v,
         "straddle": straddle,
     }
 
@@ -863,45 +724,6 @@ def compute_barotropic_velocity(
     return vel_bar
 
 
-def transpose_dimensions(da: xr.DataArray) -> xr.DataArray:
-    """Transpose the dimensions of an xarray.DataArray to ensure that 'time', any
-    dimension starting with 's_', 'eta_', and 'xi_' are ordered first, followed by the
-    remaining dimensions in their original order.
-
-    Parameters
-    ----------
-    da : xarray.DataArray
-        The input DataArray whose dimensions are to be reordered.
-
-    Returns
-    -------
-    xarray.DataArray
-        The DataArray with dimensions reordered so that 'time', 's_*', 'eta_*',
-        and 'xi_*' are first, in that order, if they exist.
-    """
-
-    # List of preferred dimension patterns
-    preferred_order = ["time", "s_", "eta_", "xi_"]
-
-    # Get the existing dimensions in the DataArray
-    dims = list(da.dims)
-
-    # Collect dimensions that match any of the preferred patterns
-    matched_dims = []
-    for pattern in preferred_order:
-        # Find dimensions that start with the pattern
-        matched_dims += [dim for dim in dims if dim.startswith(pattern)]
-
-    # Create a new order: first the matched dimensions, then the rest
-    remaining_dims = [dim for dim in dims if dim not in matched_dims]
-    new_order = matched_dims + remaining_dims
-
-    # Transpose the DataArray to the new order
-    transposed_da = da.transpose(*new_order)
-
-    return transposed_da
-
-
 def get_vector_pairs(variable_info):
     """Extracts all unique vector pairs from the variable_info dictionary.
 
@@ -1079,22 +901,20 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
 
         grid_yaml_data = {**parent_grid_yaml_data, **child_grid_yaml_data}
 
-    # Ensure forcing_object.source.path is a string (convert if it's a pathlib object)
-    if (
-        hasattr(forcing_object, "source")
-        and forcing_object.source is not None
-        and "path" in forcing_object.source
-    ):
-        forcing_object.source["path"] = str(forcing_object.source["path"])
-    if (
-        hasattr(forcing_object, "bgc_source")
-        and forcing_object.bgc_source is not None
-        and "path" in forcing_object.bgc_source
-    ):
-        forcing_object.bgc_source["path"] = str(forcing_object.bgc_source["path"])
-
-    # Step 2: Get ROMS Tools version
-    # Fetch the version of the 'roms-tools' package for inclusion in the YAML header
+    # Step 2: 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:
+            paths = attr["path"]
+            if isinstance(paths, list):
+                attr["path"] = [str(p) if isinstance(p, Path) else p for p in paths]
+            elif isinstance(paths, Path):
+                attr["path"] = str(paths)
+
+    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:
@@ -1103,8 +923,7 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
     # Create YAML header with version information
     header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
 
-    # Step 3: Prepare Forcing Data
-    # Prepare the forcing object fields, excluding 'grid' and 'ds'
+    # Step 4: Prepare Forcing Data
     forcing_data = {}
     field_names = [field.name for field in fields(forcing_object)]
     filtered_field_names = [
@@ -1133,14 +952,13 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
         # Add the field and its value to the forcing_data dictionary
         forcing_data[field_name] = value
 
-    # Step 4: Combine Grid and Forcing Data
-    # Combine grid and forcing data into a single dictionary for the final YAML content
+    # Step 5: 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 5: Write to YAML file
+    # Step 6: Write to YAML file
     with filepath.open("w") as file:
         # Write the header first
         file.write(header)
@@ -1186,6 +1004,8 @@ def _from_yaml(forcing_object: Type, filepath: Union[str, Path]) -> Dict[str, An
     ValueError
         If no configuration for the specified class name is found in the YAML file.
     """
+    # Ensure filepath is a Path object
+    filepath = Path(filepath)
 
     # Read the entire file content
     with filepath.open("r") as file: