roms-tools 1.5.0__py3-none-any.whl → 1.6.1__py3-none-any.whl

roms_tools/_version.py

@@ -1,2 +1,2 @@
 # Do not change! Do not track in version control!
-__version__ = "1.5.0"
+__version__ = "1.6.1"

roms_tools/setup/boundary_forcing.py

@@ -3,24 +3,24 @@ import numpy as np
 import pandas as pd
 import yaml
 import importlib.metadata
+import warnings
 from typing import Dict, Union, List
 from dataclasses import dataclass, field, asdict
 from roms_tools.setup.grid import Grid
-from roms_tools.setup.fill import _lateral_fill
-from roms_tools.setup.regrid import _lateral_regrid, _vertical_regrid
+from roms_tools.setup.regrid import LateralRegrid, VerticalRegrid
 from datetime import datetime
 from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
 from roms_tools.setup.utils import (
-    nan_check,
-    substitute_nans_by_fillvalue,
     get_variable_metadata,
     group_dataset,
     save_datasets,
     get_target_coords,
     rotate_velocities,
     compute_barotropic_velocity,
-    _extrapolate_deepest_to_bottom,
     transpose_dimensions,
+    one_dim_fill,
+    nan_check,
+    substitute_nans_by_fillvalue,
 )
 from roms_tools.setup.plot import _section_plot, _line_plot
 import matplotlib.pyplot as plt
@@ -60,6 +60,10 @@ class BoundaryForcing:
 
     model_reference_date : datetime, optional
         Reference date for the model. Default is January 1, 2000.
+    apply_2d_horizontal_fill: bool, optional
+        Indicates whether to perform a two-dimensional horizontal fill on the source data prior to regridding to boundaries.
+        If `False`, a one-dimensional horizontal fill is performed separately on each of the four regridded boundaries.
+        Defaults to `False`.
     use_dask: bool, optional
         Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
 
@@ -89,6 +93,7 @@ class BoundaryForcing:
     source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
     type: str = "physics"
     model_reference_date: datetime = datetime(2000, 1, 1)
+    apply_2d_horizontal_fill: bool = False
     use_dask: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
@@ -99,31 +104,42 @@ class BoundaryForcing:
         target_coords = get_target_coords(self.grid)
 
         data = self._get_data()
-        data.choose_subdomain(
-            latitude_range=[
-                target_coords["lat"].min().values,
-                target_coords["lat"].max().values,
-            ],
-            longitude_range=[
-                target_coords["lon"].min().values,
-                target_coords["lon"].max().values,
-            ],
-            margin=2,
-            straddle=target_coords["straddle"],
-        )
 
-        variable_info = self._set_variable_info(data)
-
-        data_vars = {}
-        data_vars = _extrapolate_deepest_to_bottom(data_vars, data)
-        data_vars = _lateral_fill(data_vars, data)
+        if self.apply_2d_horizontal_fill:
+            data.choose_subdomain(
+                target_coords,
+                buffer_points=20,  # lateral fill needs good buffer from data margin
+            )
+            data.extrapolate_deepest_to_bottom()
+            data.apply_lateral_fill()
 
+        variable_info = self._set_variable_info(data)
         bdry_coords = get_boundary_info()
         ds = xr.Dataset()
+
         for direction in ["south", "east", "north", "west"]:
             if self.boundaries[direction]:
 
-                bdry_data_vars = data_vars.copy()
+                bdry_target_coords = {
+                    "lat": target_coords["lat"].isel(
+                        **bdry_coords["vector"][direction]
+                    ),
+                    "lon": target_coords["lon"].isel(
+                        **bdry_coords["vector"][direction]
+                    ),
+                    "straddle": target_coords["straddle"],
+                }
+
+                bdry_data = data.choose_subdomain(
+                    bdry_target_coords,
+                    buffer_points=3,
+                    return_copy=True,
+                )
+
+                if not self.apply_2d_horizontal_fill:
+                    bdry_data.extrapolate_deepest_to_bottom()
+
+                processed_fields = {}
 
                 # lateral regridding of vector fields
                 vector_var_names = [
@@ -132,9 +148,15 @@ class BoundaryForcing:
                 if len(vector_var_names) > 0:
                     lon = target_coords["lon"].isel(**bdry_coords["vector"][direction])
                     lat = target_coords["lat"].isel(**bdry_coords["vector"][direction])
-                    bdry_data_vars = _lateral_regrid(
-                        data, lon, lat, bdry_data_vars, vector_var_names
+                    lateral_regrid = LateralRegrid(
+                        {"lat": lat, "lon": lon}, bdry_data.dim_names
                     )
+                    for var_name in vector_var_names:
+                        if var_name in bdry_data.var_names.keys():
+                            processed_fields[var_name] = lateral_regrid.apply(
+                                bdry_data.ds[bdry_data.var_names[var_name]]
+                            )
+
                 # lateral regridding of tracer fields
                 tracer_var_names = [
                     name
@@ -144,30 +166,38 @@ class BoundaryForcing:
                 if len(tracer_var_names) > 0:
                     lon = target_coords["lon"].isel(**bdry_coords["rho"][direction])
                     lat = target_coords["lat"].isel(**bdry_coords["rho"][direction])
-                    bdry_data_vars = _lateral_regrid(
-                        data, lon, lat, bdry_data_vars, tracer_var_names
+                    lateral_regrid = LateralRegrid(
+                        {"lat": lat, "lon": lon}, bdry_data.dim_names
                     )
+                    for var_name in tracer_var_names:
+                        if var_name in bdry_data.var_names.keys():
+                            processed_fields[var_name] = lateral_regrid.apply(
+                                bdry_data.ds[bdry_data.var_names[var_name]]
+                            )
 
                 # rotation of velocities and interpolation to u/v points
                 if "u" in variable_info and "v" in variable_info:
                     angle = target_coords["angle"].isel(
                         **bdry_coords["vector"][direction]
                     )
-                    (bdry_data_vars["u"], bdry_data_vars["v"],) = rotate_velocities(
-                        bdry_data_vars["u"],
-                        bdry_data_vars["v"],
+                    (processed_fields["u"], processed_fields["v"],) = rotate_velocities(
+                        processed_fields["u"],
+                        processed_fields["v"],
                         angle,
                         interpolate=True,
                     )
 
                 # selection of outermost margin for u/v variables
-                for var in variable_info.keys():
-                    if var in bdry_data_vars:
-                        location = variable_info[var]["location"]
+                for var_name in variable_info.keys():
+                    if var_name in processed_fields:
+                        location = variable_info[var_name]["location"]
                         if location in ["u", "v"]:
-                            bdry_data_vars[var] = bdry_data_vars[var].isel(
-                                **bdry_coords[location][direction]
-                            )
+                            processed_fields[var_name] = processed_fields[
+                                var_name
+                            ].isel(**bdry_coords[location][direction])
+
+                if not self.apply_2d_horizontal_fill:
+                    processed_fields = apply_1d_horizontal_fill(processed_fields)
 
                 # vertical regridding
                 for location in ["rho", "u", "v"]:
@@ -177,53 +207,47 @@ class BoundaryForcing:
                         if info["location"] == location and info["is_3d"]
                     ]
                     if len(var_names) > 0:
-                        bdry_data_vars = _vertical_regrid(
-                            data,
+                        vertical_regrid = VerticalRegrid(
                             self.grid.ds[f"layer_depth_{location}"].isel(
-                                **bdry_coords[location][direction],
+                                **bdry_coords[location][direction]
                             ),
-                            bdry_data_vars,
-                            var_names,
+                            bdry_data.ds[bdry_data.dim_names["depth"]],
                         )
+                        for var_name in var_names:
+                            if var_name in processed_fields:
+                                processed_fields[var_name] = vertical_regrid.apply(
+                                    processed_fields[var_name]
+                                )
 
                 # compute barotropic velocities
                 if "u" in variable_info and "v" in variable_info:
-                    for var in ["u", "v"]:
-                        bdry_data_vars[f"{var}bar"] = compute_barotropic_velocity(
-                            bdry_data_vars[var],
-                            self.grid.ds[f"interface_depth_{var}"].isel(
-                                **bdry_coords[var][direction]
+                    for var_name in ["u", "v"]:
+                        processed_fields[
+                            f"{var_name}bar"
+                        ] = compute_barotropic_velocity(
+                            processed_fields[var_name],
+                            self.grid.ds[f"interface_depth_{var_name}"].isel(
+                                **bdry_coords[var_name][direction]
                             ),
                         )
 
                 # Reorder dimensions
-                for var in bdry_data_vars.keys():
-                    bdry_data_vars[var] = transpose_dimensions(bdry_data_vars[var])
+                for var_name in processed_fields.keys():
+                    processed_fields[var_name] = transpose_dimensions(
+                        processed_fields[var_name]
+                    )
 
                 # Write the boundary data into dataset
-                ds = self._write_into_dataset(direction, bdry_data_vars, ds)
+                ds = self._write_into_dataset(direction, processed_fields, ds)
 
         # Add global information
         ds = self._add_global_metadata(data, ds)
 
-        # NaN values at wet points indicate that the raw data did not cover the domain, and the following will raise a ValueError
-        # this check works only for 2D fields because for 3D I extrapolate to bottom which eliminates NaNs
-        for direction in ["south", "east", "north", "west"]:
-            if self.boundaries[direction]:
-                if type == "physics":
-                    nan_check(
-                        ds[f"zeta_{direction}"].isel(bry_time=0),
-                        self.grid.ds.mask_rho.isel(**bdry_coords["rho"][direction]),
-                    )
-                elif type == "bgc":
-                    nan_check(
-                        ds[f"ALK_{direction}"].isel(bry_time=0, s_rho=-1),
-                        self.grid.ds.mask_rho.isel(**bdry_coords["rho"][direction]),
-                    )
+        self._validate(ds, variable_info, bdry_coords)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
-        for var in ds.data_vars:
-            ds[var] = substitute_nans_by_fillvalue(ds[var])
+        for var_name in ds.data_vars:
+            ds[var_name] = substitute_nans_by_fillvalue(ds[var_name])
 
         object.__setattr__(self, "ds", ds)
 
@@ -335,25 +359,27 @@ class BoundaryForcing:
             }
         elif self.type == "bgc":
             variable_info = {}
-            for var in data.var_names.keys():
-                variable_info[var] = default_info
+            for var_name in data.var_names.keys():
+                variable_info[var_name] = default_info
 
         return variable_info
 
-    def _write_into_dataset(self, direction, data_vars, ds=None):
+    def _write_into_dataset(self, direction, processed_fields, ds=None):
         if ds is None:
             ds = xr.Dataset()
 
         d_meta = get_variable_metadata()
 
-        for var in data_vars.keys():
-            ds[f"{var}_{direction}"] = data_vars[var].astype(np.float32)
+        for var_name in processed_fields.keys():
+            ds[f"{var_name}_{direction}"] = processed_fields[var_name].astype(
+                np.float32
+            )
 
-            ds[f"{var}_{direction}"].attrs[
+            ds[f"{var_name}_{direction}"].attrs[
                 "long_name"
-            ] = f"{direction}ern boundary {d_meta[var]['long_name']}"
+            ] = f"{direction}ern boundary {d_meta[var_name]['long_name']}"
 
-            ds[f"{var}_{direction}"].attrs["units"] = d_meta[var]["units"]
+            ds[f"{var_name}_{direction}"].attrs["units"] = d_meta[var_name]["units"]
 
         # Gracefully handle dropping variables that might not be present
         variables_to_drop = [
@@ -371,7 +397,7 @@ class BoundaryForcing:
             "lat_v",
             "lon_v",
         ]
-        existing_vars = [var for var in variables_to_drop if var in ds]
+        existing_vars = [var_name for var_name in variables_to_drop if var_name in ds]
         ds = ds.drop_vars(existing_vars)
 
         return ds
@@ -480,9 +506,76 @@ class BoundaryForcing:
 
         return ds
 
+    def _validate(self, ds, variable_info, bdry_coords):
+        """Validate the dataset for NaN values at the first time step based on the fill
+        method used.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            The dataset to validate.
+
+        variable_info : dict
+            A dictionary containing metadata about the variables, including their locations (e.g., 'rho', 'u', 'v').
+
+        bdry_coords : dict
+            A dictionary containing the boundary coordinates for each variable location.
+
+        Raises
+        ------
+        ValueError
+            If NaN values are found in any of the specified variables at wet points,
+            indicating incomplete data coverage.
+
+        Notes
+        -----
+        Validation is performed on the initial boundary time step (`bry_time=0`) for each
+        variable in the dataset. If the `apply_2d_horizontal_fill` attribute is set to False,
+        a warning is issued instead of a strict NaN check, as the data may not be reliably validated.
+        Conversely, if `apply_2d_horizontal_fill` is True, a strict NaN check is performed, raising
+        a ValueError if any NaN values are detected.
+        """
+        if self.apply_2d_horizontal_fill:
+            # Strict NaN check with ValueError makes sense to be applied
+            for var_name in variable_info:
+                location = variable_info[var_name]["location"]
+
+                # Select the appropriate mask based on variable location
+                if location == "rho":
+                    mask = self.grid.ds.mask_rho
+                elif location == "u":
+                    mask = self.grid.ds.mask_u
+                elif location == "v":
+                    mask = self.grid.ds.mask_v
+                else:
+                    continue  # Skip if location is not recognized
+
+                for direction in ["south", "east", "north", "west"]:
+                    if self.boundaries[direction]:
+                        bdry_var_name = f"{var_name}_{direction}"
+
+                        # Check for NaN values at the first time step using the nan_check function
+                        nan_check(
+                            ds[bdry_var_name].isel(bry_time=0),
+                            mask.isel(**bdry_coords[location][direction]),
+                        )
+        else:
+            # Can't apply strict NaN check because land values haven't been filled before regridding step; instead warn user
+            for direction in ["south", "east", "north", "west"]:
+                if self.boundaries[direction]:
+                    for var_name in variable_info:
+                        bdry_var_name = f"{var_name}_{direction}"
+                        if ds[bdry_var_name].isel(bry_time=0).isnull().any().values:
+                            warnings.warn(
+                                f"NaN values detected in regridded variables along the {direction}ern boundary. This may indicate that the entire boundary is on land in the source data, or that the source data does not cover this boundary.",
+                                UserWarning,
+                            )
+                            # Break after the first warning for this direction to avoid duplicates
+                            break
+
     def plot(
         self,
-        varname,
+        var_name,
         time=0,
         layer_contours=False,
     ) -> None:
@@ -490,7 +583,7 @@ class BoundaryForcing:
 
         Parameters
         ----------
-        varname : str
+        var_name : str
             The name of the boundary forcing field to plot. Options include:
 
             - "temp_{direction}": Potential temperature,
@@ -550,37 +643,44 @@ class BoundaryForcing:
         Raises
         ------
         ValueError
-            If the specified varname is not one of the valid options.
+            If the specified var_name is not one of the valid options.
         """
 
-        if varname not in self.ds:
-            raise ValueError(f"Variable '{varname}' is not found in dataset.")
+        if var_name not in self.ds:
+            raise ValueError(f"Variable '{var_name}' is not found in dataset.")
+
+        field = self.ds[var_name].isel(bry_time=time)
+
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                field = field.load()
 
-        field = self.ds[varname].isel(bry_time=time).load()
         title = field.long_name
 
         if "s_rho" in field.dims:
-            if varname.startswith(("u_", "ubar_")):
+            if var_name.startswith(("u_", "ubar_")):
                 point = "u"
-            elif varname.startswith(("v_", "vbar_")):
+            elif var_name.startswith(("v_", "vbar_")):
                 point = "v"
             else:
                 point = "rho"
-            direction = varname.split("_")[-1]
+            direction = var_name.split("_")[-1]
 
             layer_depth, interface_depth = self._get_coordinates(direction, point)
 
             field = field.assign_coords({"layer_depth": layer_depth})
 
         # chose colorbar
-        if varname.startswith(("u", "v", "ubar", "vbar", "zeta")):
+        if var_name.startswith(("u", "v", "ubar", "vbar", "zeta")):
             vmax = max(field.max().values, -field.min().values)
             vmin = -vmax
             cmap = plt.colormaps.get_cmap("RdBu_r")
         else:
             vmax = field.max().values
             vmin = field.min().values
-            if varname.startswith(("temp", "salt")):
+            if var_name.startswith(("temp", "salt")):
                 cmap = plt.colormaps.get_cmap("YlOrRd")
             else:
                 cmap = plt.colormaps.get_cmap("YlGn")
@@ -606,24 +706,26 @@ class BoundaryForcing:
             _line_plot(field, title=title)
 
     def save(
-        self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
+        self,
+        filepath: Union[str, Path],
+        np_eta: int = None,
+        np_xi: int = None,
+        group: bool = False,
     ) -> None:
-        """Save the boundary forcing fields to netCDF4 files.
-
-        This method saves the dataset by grouping it into subsets based on the data frequency. The subsets are then written
-        to one or more netCDF4 files. The filenames of the output files reflect the temporal coverage of the data.
+        """Save the boundary forcing fields to one or more netCDF4 files.
 
-        There are two modes of saving the dataset:
+        This method saves the dataset either as a single file or as multiple files depending on the partitioning and grouping options.
+        The dataset can be saved in two modes:
 
-          1. **Single File Mode (default)**:
+        1. **Single File Mode (default)**:
+            - If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single netCDF4 file.
+            - The file is named based on the `filepath`, with `.nc` automatically appended.
 
-            If both `np_eta` and `np_xi` are `None`, the entire dataset, divided by temporal subsets, 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 partitioned into spatial tiles along the `eta` and `xi` axes.
+            - Each tile is saved as a separate netCDF4 file, and filenames are modified with an index (e.g., `"filepath_YYYYMM.0.nc"`, `"filepath_YYYYMM.1.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.
+        Additionally, if `group` is set to `True`, the dataset is first grouped into temporal subsets, resulting in multiple grouped files before partitioning and saving.
 
         Parameters
         ----------
@@ -635,6 +737,8 @@ class BoundaryForcing:
             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.
+        group: bool, optional
+            If `True`, groups the dataset into multiple files based on temporal data frequency. Defaults to `False`.
 
         Returns
         -------
@@ -649,7 +753,18 @@ class BoundaryForcing:
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
 
-        dataset_list, output_filenames = group_dataset(self.ds.load(), str(filepath))
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                self.ds.load()
+
+        if group:
+            dataset_list, output_filenames = group_dataset(self.ds, str(filepath))
+        else:
+            dataset_list = [self.ds]
+            output_filenames = [str(filepath)]
+
         saved_filenames = save_datasets(
             dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
         )
@@ -703,7 +818,7 @@ class BoundaryForcing:
             # Write header
             file.write(header)
             # Write YAML data
-            yaml.dump(yaml_data, file, default_flow_style=False)
+            yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
 
     @classmethod
     def from_yaml(
@@ -798,3 +913,51 @@ def get_boundary_info():
     }
 
     return bdry_coords
+
+
+def apply_1d_horizontal_fill(processed_fields: dict) -> dict:
+    """Forward and backward fill NaN values in horizontal direction for open boundaries.
+
+    Parameters
+    ----------
+    processed_fields : dict
+        A dictionary of variables to be updated, where each value is an
+        `xarray.DataArray`.
+
+    Returns
+    -------
+    dict of str : xarray.DataArray
+        The updated dictionary of variables, with NaN values filled.
+
+    Raises
+    ------
+    ValueError
+        If more than one horizontal dimension is found or none at all.
+    """
+
+    horizontal_dims = ["eta_rho", "eta_v", "xi_rho", "xi_u"]
+
+    for var_name in processed_fields.keys():
+        selected_horizontal_dim = None
+        # Determine the horizontal dimension to fill
+        for dim in horizontal_dims:
+            if dim in processed_fields[var_name].dims:
+                if selected_horizontal_dim is not None:
+                    raise ValueError(
+                        f"More than one horizontal dimension found in variable '{var_name}'."
+                    )
+                selected_horizontal_dim = dim
+
+        if selected_horizontal_dim is None:
+            raise ValueError(
+                f"No valid horizontal dimension found for variable '{var_name}'."
+            )
+        # Forward and backward fill in the horizontal direction
+        filled = one_dim_fill(
+            processed_fields[var_name], selected_horizontal_dim, direction="forward"
+        )
+        processed_fields[var_name] = one_dim_fill(
+            filled, selected_horizontal_dim, direction="backward"
+        )
+
+    return processed_fields