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

roms_tools/regrid.py

@@ -1,7 +1,8 @@
-import xgcm
-import xarray as xr
 import warnings
 
+import xarray as xr
+import xgcm
+
 
 class LateralRegridToROMS:
     """Handles lateral regridding of data onto a new spatial grid."""
@@ -24,7 +25,6 @@ class LateralRegridToROMS:
             Maps the dimension names to the corresponding latitude and longitude
             DataArrays, providing easy access to target grid coordinates.
         """
-
         self.coords = {
             source_dim_names["latitude"]: target_coords["lat"],
             source_dim_names["longitude"]: target_coords["lon"],
@@ -96,7 +96,6 @@ class LateralRegridFromROMS:
         ImportError
             If xESMF is not installed.
         """
-
         try:
             import xesmf as xe
 
@@ -129,7 +128,6 @@ class LateralRegridFromROMS:
         xarray.DataArray
             The regridded data array.
         """
-
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=UserWarning, module="xesmf")
             regridded = self.regridder(da, keep_attrs=True)
@@ -166,7 +164,6 @@ class VerticalRegridToROMS:
             - `upper_mask`, `lower_mask` : Masks for valid interpolation bounds.
             - `factor` : Weight for blending values between levels.
         """
-
         self.depth_dim = source_depth_coords.dims[0]
         source_depth = source_depth_coords
         dims = {"dim": self.depth_dim}
@@ -214,7 +211,6 @@ class VerticalRegridToROMS:
             are replaced if `fill_nans=True`, with extrapolation allowed at the surface
             and bottom layers to minimize gaps.
         """
-
         dims = {"dim": self.depth_dim}
 
         da_below = da.where(self.coeff["is_below"]).sum(**dims)
@@ -283,7 +279,6 @@ class VerticalRegridFromROMS:
         xarray.DataArray
             A new `xarray.DataArray` containing the regridded data at the specified target depth levels.
         """
-
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=FutureWarning, module="xgcm")
             transformed = self.grid.transform(

roms_tools/setup/boundary_forcing.py

@@ -1,39 +1,41 @@
-import xarray as xr
-import numpy as np
-from scipy.ndimage import label
-import logging
 import importlib.metadata
-from typing import Dict, Union, List, Optional
+import logging
 from dataclasses import dataclass, field
 from datetime import datetime
-import matplotlib.pyplot as plt
 from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
+from scipy.ndimage import label
+
 from roms_tools import Grid
+from roms_tools.plot import line_plot, section_plot
 from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
-from roms_tools.utils import save_datasets
-from roms_tools.vertical_coordinate import compute_depth
-from roms_tools.plot import _section_plot, _line_plot
-from roms_tools.utils import (
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
-    transpose_dimensions,
-)
-from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset, UnifiedBGCDataset
+from roms_tools.setup.datasets import CESMBGCDataset, GLORYSDataset, UnifiedBGCDataset
 from roms_tools.setup.utils import (
-    get_variable_metadata,
-    group_dataset,
-    get_target_coords,
-    rotate_velocities,
+    add_time_info_to_ds,
     compute_barotropic_velocity,
     compute_missing_bgc_variables,
-    one_dim_fill,
+    from_yaml,
+    get_boundary_coords,
+    get_target_coords,
+    get_variable_metadata,
+    group_dataset,
     nan_check,
+    one_dim_fill,
+    rotate_velocities,
     substitute_nans_by_fillvalue,
-    convert_to_roms_time,
-    get_boundary_coords,
-    _to_yaml,
-    _from_yaml,
+    to_dict,
+    write_to_yaml,
 )
+from roms_tools.utils import (
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+    save_datasets,
+    transpose_dimensions,
+)
+from roms_tools.vertical_coordinate import compute_depth
 
 
 @dataclass(kw_only=True)
@@ -101,9 +103,12 @@ class BoundaryForcing:
     """
 
     grid: Grid
-    start_time: Optional[datetime] = None
-    end_time: Optional[datetime] = None
-    boundaries: Dict[str, bool] = field(
+    """Object representing the grid information."""
+    start_time: datetime | None = None
+    """The start time of the desired surface forcing data."""
+    end_time: datetime | None = None
+    """The end time of the desired surface forcing data."""
+    boundaries: dict[str, bool] = field(
         default_factory=lambda: {
             "south": True,
             "east": True,
@@ -111,18 +116,29 @@ class BoundaryForcing:
             "west": True,
         }
     )
-    source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
+    """Dictionary specifying which boundaries are forced (south, east, north, west)."""
+    source: dict[str, str | Path | list[str | Path]]
+    """Dictionary specifying the source of the boundary forcing data."""
     type: str = "physics"
+    """Specifies the type of forcing data ("physics", "bgc")."""
     apply_2d_horizontal_fill: bool = False
+    """Whether to perform a two-dimensional horizontal fill on the source data prior to
+    regridding to boundaries."""
     adjust_depth_for_sea_surface_height: bool = False
+    """Whether to account for sea surface height (`zeta`) variations when computing
+    depth coordinates."""
     model_reference_date: datetime = datetime(2000, 1, 1)
+    """Reference date for the model."""
     use_dask: bool = False
+    """Whether to use dask for processing."""
     bypass_validation: bool = False
+    """Whether to skip validation checks in the processed data."""
 
     ds: xr.Dataset = field(init=False, repr=False)
+    """An xarray Dataset containing post-processed variables ready for input into
+    ROMS."""
 
     def __post_init__(self):
-
         self._input_checks()
         # Dataset for depth coordinates
         self.ds_depth_coords = xr.Dataset()
@@ -132,7 +148,6 @@ class BoundaryForcing:
         data = self._get_data()
 
         if self.apply_2d_horizontal_fill:
-
             data.choose_subdomain(
                 target_coords,
                 buffer_points=20,  # lateral fill needs good buffer from data margin
@@ -168,7 +183,6 @@ class BoundaryForcing:
 
         for direction in ["south", "east", "north", "west"]:
             if self.boundaries[direction]:
-
                 bdry_target_coords = {
                     "lat": target_coords["lat"].isel(
                         **self.bdry_coords["vector"][direction]
@@ -250,7 +264,10 @@ class BoundaryForcing:
                     angle = target_coords["angle"].isel(
                         **self.bdry_coords["vector"][direction]
                     )
-                    (processed_fields["u"], processed_fields["v"],) = rotate_velocities(
+                    (
+                        processed_fields["u"],
+                        processed_fields["v"],
+                    ) = rotate_velocities(
                         processed_fields["u"],
                         processed_fields["v"],
                         angle,
@@ -332,13 +349,13 @@ class BoundaryForcing:
                     self._get_depth_coordinates(zeta_u, direction, "u", "interface")
                     self._get_depth_coordinates(zeta_v, direction, "v", "interface")
                     for location in ["u", "v"]:
-                        processed_fields[
-                            f"{location}bar"
-                        ] = compute_barotropic_velocity(
-                            processed_fields[location],
-                            self.ds_depth_coords[
-                                f"interface_depth_{location}_{direction}"
-                            ],
+                        processed_fields[f"{location}bar"] = (
+                            compute_barotropic_velocity(
+                                processed_fields[location],
+                                self.ds_depth_coords[
+                                    f"interface_depth_{location}_{direction}"
+                                ],
+                            )
                         )
 
                 # Reorder dimensions
@@ -409,7 +426,6 @@ class BoundaryForcing:
             )
 
     def _get_data(self):
-
         data_dict = {
             "filename": self.source["path"],
             "start_time": self.start_time,
@@ -428,7 +444,6 @@ class BoundaryForcing:
 
         elif self.type == "bgc":
             if self.source["name"] == "CESM_REGRIDDED":
-
                 data = CESMBGCDataset(**data_dict)
             elif self.source["name"] == "UNIFIED":
                 data = UnifiedBGCDataset(**data_dict)
@@ -530,9 +545,9 @@ class BoundaryForcing:
                 np.float32
             )
 
-            ds[f"{var_name}_{direction}"].attrs[
-                "long_name"
-            ] = f"{direction}ern boundary {d_meta[var_name]['long_name']}"
+            ds[f"{var_name}_{direction}"].attrs["long_name"] = (
+                f"{direction}ern boundary {d_meta[var_name]['long_name']}"
+            )
 
             ds[f"{var_name}_{direction}"].attrs["units"] = d_meta[var_name]["units"]
 
@@ -582,7 +597,6 @@ class BoundaryForcing:
               to their corresponding boundary coordinates. The coordinates are specified in terms of
               grid indices for the respective variable types.
         """
-
         bdry_coords = get_boundary_coords()
 
         self.bdry_coords = bdry_coords
@@ -663,7 +677,6 @@ class BoundaryForcing:
             self.ds_depth_coords[key] = depth
 
     def _add_global_metadata(self, data, ds=None):
-
         if ds is None:
             ds = xr.Dataset()
         ds.attrs["title"] = "ROMS boundary forcing file created by ROMS-Tools"
@@ -686,8 +699,7 @@ class BoundaryForcing:
         ds.attrs["theta_b"] = self.grid.ds.attrs["theta_b"]
         ds.attrs["hc"] = self.grid.ds.attrs["hc"]
 
-        # Convert the time coordinate to the format expected by ROMS
-        ds, bry_time = convert_to_roms_time(
+        ds, bry_time = add_time_info_to_ds(
             ds, self.model_reference_date, data.climatology
         )
 
@@ -719,7 +731,6 @@ class BoundaryForcing:
         None
             If a boundary is divided by land, a warning is issued. No return value is provided.
         """
-
         for var_name in processed_fields.keys():
             if self.variable_info[var_name]["validate"]:
                 location = self.variable_info[var_name]["location"]
@@ -876,7 +887,6 @@ class BoundaryForcing:
         ValueError
             If the specified var_name is not one of the valid options.
         """
-
         if var_name not in self.ds:
             raise ValueError(f"Variable '{var_name}' is not found in dataset.")
 
@@ -939,7 +949,7 @@ class BoundaryForcing:
             else:
                 interface_depth = None
 
-            _section_plot(
+            section_plot(
                 field,
                 interface_depth=interface_depth,
                 title=title,
@@ -947,11 +957,11 @@ class BoundaryForcing:
                 ax=ax,
             )
         else:
-            _line_plot(field.where(mask), title=title, ax=ax)
+            line_plot(field.where(mask), title=title, ax=ax)
 
     def save(
         self,
-        filepath: Union[str, Path],
+        filepath: str | Path,
         group: bool = True,
     ) -> None:
         """Save the boundary forcing fields to one or more netCDF4 files.
@@ -973,7 +983,6 @@ class BoundaryForcing:
         List[Path]
             A list of Path objects for the filenames that were saved.
         """
-
         # Ensure filepath is a Path object
         filepath = Path(filepath)
 
@@ -993,7 +1002,7 @@ class BoundaryForcing:
 
         return saved_filenames
 
-    def to_yaml(self, filepath: Union[str, Path]) -> None:
+    def to_yaml(self, filepath: str | Path) -> None:
         """Export the parameters of the class to a YAML file, including the version of
         roms-tools.
 
@@ -1002,13 +1011,13 @@ class BoundaryForcing:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-
-        _to_yaml(self, filepath)
+        forcing_dict = to_dict(self, exclude=["use_dask"])
+        write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(
         cls,
-        filepath: Union[str, Path],
+        filepath: str | Path,
         use_dask: bool = False,
     ) -> "BoundaryForcing":
         """Create an instance of the BoundaryForcing class from a YAML file.
@@ -1028,7 +1037,7 @@ class BoundaryForcing:
         filepath = Path(filepath)
 
         grid = Grid.from_yaml(filepath)
-        params = _from_yaml(cls, filepath)
+        params = from_yaml(cls, filepath)
 
         # Create and return an instance of InitialConditions
         return cls(grid=grid, **params, use_dask=use_dask)
@@ -1053,7 +1062,6 @@ def apply_1d_horizontal_fill(data_array: xr.DataArray) -> xr.DataArray:
     ValueError
         If more than one horizontal dimension is found or none at all.
     """
-
     horizontal_dims = ["eta_rho", "eta_v", "xi_rho", "xi_u"]
     selected_horizontal_dim = None