roms-tools 1.4.0__py3-none-any.whl → 1.4.2__py3-none-any.whl

roms_tools/setup/mixins.py

@@ -12,22 +12,20 @@ import numpy as np
 
 @dataclass(frozen=True, kw_only=True)
 class ROMSToolsMixins:
-    """
-    Represents a mixin tool for ROMS-Tools with capabilities shared by the various
+    """Represents a mixin tool for ROMS-Tools with capabilities shared by the various
     ROMS-Tools dataclasses.
 
     Parameters
     ----------
     grid : Grid
         Object representing the grid information used for the model.
-
     """
 
     grid: Grid
 
-    def get_target_lon_lat(self, use_coarse_grid=False):
-        """
-        Retrieves the longitude and latitude arrays from the grid and adjusts them based on the grid's orientation.
+    def _get_target_lon_lat(self, use_coarse_grid=False):
+        """Retrieves the longitude and latitude arrays from the grid and adjusts them
+        based on the grid's orientation.
 
         This method provides longitude and latitude coordinates, with options for using a coarse grid
         if specified. It also handles longitudes to ensure they are between -180 and 180 degrees and adjusts
@@ -67,10 +65,9 @@ class ROMSToolsMixins:
 
         return lon, lat, angle, straddle
 
-    def regrid_data(self, data, vars_2d, vars_3d, lon, lat):
-
-        """
-        Interpolates data onto the desired grid and processes it for 2D and 3D variables.
+    def _regrid_data(self, data, vars_2d, vars_3d, lon, lat):
+        """Interpolates data onto the desired grid and processes it for 2D and 3D
+        variables.
 
         This method interpolates the specified 2D and 3D variables onto a new grid defined by the provided
         longitude and latitude coordinates. It handles both 2D and 3D data, performing extrapolation for 3D
@@ -164,10 +161,9 @@ class ROMSToolsMixins:
 
         return data_vars
 
-    def process_velocities(self, data_vars, angle, uname, vname, interpolate=True):
-        """
-        Process and rotate velocity components to align with the grid orientation and optionally interpolate
-        them to the appropriate grid points.
+    def _process_velocities(self, data_vars, angle, uname, vname, interpolate=True):
+        """Process and rotate velocity components to align with the grid orientation and
+        optionally interpolate them to the appropriate grid points.
 
         This method performs the following steps:
 

roms_tools/setup/surface_forcing.py

@@ -28,8 +28,7 @@ from pathlib import Path
 
 @dataclass(frozen=True, kw_only=True)
 class SurfaceForcing(ROMSToolsMixins):
-    """
-    Represents surface forcing input data for ROMS.
+    """Represents surface forcing input data for ROMS.
 
     Parameters
     ----------
@@ -40,14 +39,22 @@ class SurfaceForcing(ROMSToolsMixins):
     end_time : datetime
         End time of the desired surface forcing data.
     source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
-        Dictionary specifying the source of the surface forcing data:
-        - "name" (str): Name of the data source (e.g., "ERA5").
-        - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file(s). Can be a single string (with or without wildcards),
-          a single Path object, or a list of strings or Path objects containing multiple files.
-        - "climatology" (bool): Indicates if the data is climatology data. Defaults to False.
+        Dictionary specifying the source of the surface forcing data. Keys include:
+
+          - "name" (str): Name of the data source (e.g., "ERA5").
+          - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file(s). This can be:
+
+            - A single string (with or without wildcards).
+            - A single Path object.
+            - A list of strings or Path objects containing multiple files.
+          - "climatology" (bool): Indicates if the data is climatology data. Defaults to False.
+
     type : str
-        Specifies the type of forcing data, either "physics" for physical
-        atmospheric forcing or "bgc" for biogeochemical forcing.
+        Specifies the type of forcing data. Options are:
+
+          - "physics": for physical atmospheric forcing.
+          - "bgc": for biogeochemical forcing.
+
     correct_radiation : bool
         Whether to correct shortwave radiation. Default is False.
     use_coarse_grid: bool
@@ -57,12 +64,6 @@ class SurfaceForcing(ROMSToolsMixins):
     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.
 
-    Attributes
-    ----------
-    ds : xr.Dataset
-        Xarray Dataset containing the surface forcing data.
-
-
     Examples
     --------
     >>> surface_forcing = SurfaceForcing(
@@ -90,7 +91,7 @@ class SurfaceForcing(ROMSToolsMixins):
     def __post_init__(self):
 
         self._input_checks()
-        lon, lat, angle, straddle = super().get_target_lon_lat(self.use_coarse_grid)
+        lon, lat, angle, straddle = super()._get_target_lon_lat(self.use_coarse_grid)
         object.__setattr__(self, "target_lon", lon)
         object.__setattr__(self, "target_lat", lat)
 
@@ -107,10 +108,10 @@ class SurfaceForcing(ROMSToolsMixins):
             vars_2d = data.var_names.keys()
         vars_3d = []
 
-        data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
+        data_vars = super()._regrid_data(data, vars_2d, vars_3d, lon, lat)
 
         if self.type == "physics":
-            data_vars = super().process_velocities(
+            data_vars = super()._process_velocities(
                 data_vars, angle, "uwnd", "vwnd", interpolate=False
             )
             if self.correct_radiation:
@@ -136,7 +137,7 @@ class SurfaceForcing(ROMSToolsMixins):
                 vars_2d = ["swr_corr"]
                 vars_3d = []
                 # spatial interpolation
-                data_vars_corr = super().regrid_data(
+                data_vars_corr = super()._regrid_data(
                     correction_data, vars_2d, vars_3d, lon, lat
                 )
                 # temporal interpolation
@@ -333,13 +334,13 @@ class SurfaceForcing(ROMSToolsMixins):
         return ds
 
     def plot(self, varname, time=0) -> None:
-        """
-        Plot the specified surface forcing field for a given time slice.
+        """Plot the specified surface forcing field for a given time slice.
 
         Parameters
         ----------
         varname : str
             The name of the surface forcing field to plot. Options include:
+
             - "uwnd": 10 meter wind in x-direction.
             - "vwnd": 10 meter wind in y-direction.
             - "swrad": Downward short-wave (solar) radiation.
@@ -353,6 +354,7 @@ class SurfaceForcing(ROMSToolsMixins):
             - "dust": Dust decomposition.
             - "nox": NOx decomposition.
             - "nhy": NHy decomposition.
+
         time : int, optional
             The time index to plot. Default is 0, which corresponds to the first
             time slice.
@@ -416,21 +418,22 @@ class SurfaceForcing(ROMSToolsMixins):
     def save(
         self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
     ) -> None:
-        """
-        Save the surface forcing fields to netCDF4 files.
+        """Save the surface 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.
 
         There are two modes of saving the dataset:
 
-        1. **Single File Mode (default)**:
-           - 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`.
+          1. **Single File Mode (default)**:
 
-        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.
+            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 divided into spatial tiles along the eta-axis and xi-axis.
+            - Each spatial tile is saved as a separate netCDF4 file.
 
         Parameters
         ----------
@@ -464,8 +467,8 @@ class SurfaceForcing(ROMSToolsMixins):
         return saved_filenames
 
     def to_yaml(self, filepath: Union[str, Path]) -> None:
-        """
-        Export the parameters of the class to a YAML file, including the version of roms-tools.
+        """Export the parameters of the class to a YAML file, including the version of
+        roms-tools.
 
         Parameters
         ----------
@@ -520,8 +523,7 @@ class SurfaceForcing(ROMSToolsMixins):
     def from_yaml(
         cls, filepath: Union[str, Path], use_dask: bool = False
     ) -> "SurfaceForcing":
-        """
-        Create an instance of the SurfaceForcing class from a YAML file.
+        """Create an instance of the SurfaceForcing class from a YAML file.
 
         Parameters
         ----------

roms_tools/setup/tides.py

@@ -24,18 +24,22 @@ from pathlib import Path
 
 @dataclass(frozen=True, kw_only=True)
 class TidalForcing(ROMSToolsMixins):
-    """
-    Represents tidal forcing data used in ocean modeling.
+    """Represents tidal forcing data used in ocean modeling.
 
     Parameters
     ----------
     grid : Grid
         The grid object representing the ROMS grid associated with the tidal forcing data.
     source : Dict[str, Union[str, Path, List[Union[str, Path]]]]
-        Dictionary specifying the source of the tidal data:
-        - "name" (str): Name of the data source (e.g., "TPXO").
-        - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file(s). Can be a single string (with or without wildcards),
-          a single Path object, or a list of strings or Path objects containing multiple files.
+        Dictionary specifying the source of the tidal data. Keys include:
+
+          - "name" (str): Name of the data source (e.g., "TPXO").
+          - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file(s). This can be:
+
+            - A single string (with or without wildcards).
+            - A single Path object.
+            - A list of strings or Path objects containing multiple files.
+
     ntides : int, optional
         Number of constituents to consider. Maximum number is 14. Default is 10.
     allan_factor : float, optional
@@ -45,11 +49,6 @@ class TidalForcing(ROMSToolsMixins):
     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.
 
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the tidal forcing data.
-
     Examples
     --------
     >>> tidal_forcing = TidalForcing(
@@ -69,7 +68,7 @@ class TidalForcing(ROMSToolsMixins):
     def __post_init__(self):
 
         self._input_checks()
-        lon, lat, angle, straddle = super().get_target_lon_lat()
+        lon, lat, angle, straddle = super()._get_target_lon_lat()
 
         data = self._get_data()
 
@@ -98,12 +97,12 @@ class TidalForcing(ROMSToolsMixins):
         ]
         vars_3d = []
 
-        data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
+        data_vars = super()._regrid_data(data, vars_2d, vars_3d, lon, lat)
 
-        data_vars = super().process_velocities(
+        data_vars = super()._process_velocities(
             data_vars, angle, "u_Re", "v_Re", interpolate=False
         )
-        data_vars = super().process_velocities(
+        data_vars = super()._process_velocities(
             data_vars, angle, "u_Im", "v_Im", interpolate=False
         )
 
@@ -180,13 +179,13 @@ class TidalForcing(ROMSToolsMixins):
         return ds
 
     def plot(self, varname, ntides=0) -> None:
-        """
-        Plot the specified tidal forcing variable for a given tidal constituent.
+        """Plot the specified tidal forcing variable for a given tidal constituent.
 
         Parameters
         ----------
         varname : str
             The tidal forcing variable to plot. Options include:
+
             - "ssh_Re": Real part of tidal elevation.
             - "ssh_Im": Imaginary part of tidal elevation.
             - "pot_Re": Real part of tidal potential.
@@ -195,6 +194,7 @@ class TidalForcing(ROMSToolsMixins):
             - "u_Im": Imaginary part of tidal velocity in the x-direction.
             - "v_Re": Real part of tidal velocity in the y-direction.
             - "v_Im": Imaginary part of tidal velocity in the y-direction.
+
         ntides : int, optional
             The index of the tidal constituent to plot. Default is 0, which corresponds
             to the first constituent.
@@ -258,17 +258,19 @@ class TidalForcing(ROMSToolsMixins):
     def save(
         self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
     ) -> None:
-        """
-        Save the tidal forcing information to a netCDF4 file.
+        """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 file at the specified `filepath.nc`.
+          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.
-           - The files are saved as `filepath.0.nc`, `filepath.1.nc`, ..., where the numbering corresponds to the partition index.
+          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
         ----------
@@ -302,8 +304,8 @@ class TidalForcing(ROMSToolsMixins):
         return saved_filenames
 
     def to_yaml(self, filepath: Union[str, Path]) -> None:
-        """
-        Export the parameters of the class to a YAML file, including the version of roms-tools.
+        """Export the parameters of the class to a YAML file, including the version of
+        roms-tools.
 
         Parameters
         ----------
@@ -351,8 +353,7 @@ class TidalForcing(ROMSToolsMixins):
     def from_yaml(
         cls, filepath: Union[str, Path], use_dask: bool = False
     ) -> "TidalForcing":
-        """
-        Create an instance of the TidalForcing class from a YAML file.
+        """Create an instance of the TidalForcing class from a YAML file.
 
         Parameters
         ----------
@@ -400,10 +401,10 @@ class TidalForcing(ROMSToolsMixins):
         return cls(grid=grid, **tidal_forcing_params, use_dask=use_dask)
 
     def _correct_tides(self, data):
-        """
-        Apply tidal corrections to the dataset.
-        This method corrects the dataset for equilibrium tides, self-attraction and loading (SAL) effects, and
-        adjusts phases and amplitudes of tidal elevations and transports using Egbert's correction.
+        """Apply tidal corrections to the dataset. This method corrects the dataset for
+        equilibrium tides, self-attraction and loading (SAL) effects, and adjusts phases
+        and amplitudes of tidal elevations and transports using Egbert's correction.
+
         Parameters
         ----------
         data : Dataset
@@ -460,8 +461,7 @@ class TidalForcing(ROMSToolsMixins):
 
 
 def modified_julian_days(year, month, day, hour=0):
-    """
-    Calculate the Modified Julian Day (MJD) for a given date and time.
+    """Calculate the Modified Julian Day (MJD) for a given date and time.
 
     The Modified Julian Day (MJD) is a modified Julian day count starting from
     November 17, 1858 AD. It is commonly used in astronomy and geodesy.
@@ -519,9 +519,8 @@ def modified_julian_days(year, month, day, hour=0):
 
 
 def egbert_correction(date):
-    """
-    Correct phases and amplitudes for real-time runs using parts of the
-    post-processing code from Egbert's & Erofeeva's (OSU) TPXO model.
+    """Correct phases and amplitudes for real-time runs using parts of the post-
+    processing code from Egbert's & Erofeeva's (OSU) TPXO model.
 
     Parameters
     ----------
@@ -541,7 +540,6 @@ def egbert_correction(date):
     ----------
     - Egbert, G.D., and S.Y. Erofeeva. "Efficient inverse modeling of barotropic ocean
       tides." Journal of Atmospheric and Oceanic Technology 19, no. 2 (2002): 183-204.
-
     """
 
     year = date.year
@@ -676,8 +674,7 @@ def egbert_correction(date):
 
 
 def compute_equilibrium_tide(lon, lat):
-    """
-    Compute equilibrium tide for given longitudes and latitudes.
+    """Compute equilibrium tide for given longitudes and latitudes.
 
     Parameters
     ----------
@@ -699,7 +696,6 @@ def compute_equilibrium_tide(lon, lat):
         - 2: semidiurnal
         - 1: diurnal
         - 0: long-term
-
     """
 
     # Amplitudes and elasticity factors for 15 tidal constituents

roms_tools/setup/topography.py

@@ -12,8 +12,8 @@ from itertools import count
 def _add_topography_and_mask(
     ds, topography_source, hmin, smooth_factor=8.0, rmax=0.2
 ) -> xr.Dataset:
-    """
-    Adds topography and a land/water mask to the dataset based on the provided topography source.
+    """Adds topography and a land/water mask to the dataset based on the provided
+    topography source.
 
     This function performs the following operations:
     1. Interpolates topography data onto the desired grid.
@@ -61,6 +61,10 @@ def _add_topography_and_mask(
 
     # fill enclosed basins with land
     mask = _fill_enclosed_basins(mask.values)
+
+    # adjust mask boundaries by copying values from adjacent cells
+    mask = _handle_boundaries(mask)
+
     ds["mask_rho"] = xr.DataArray(mask.astype(np.int32), dims=("eta_rho", "xi_rho"))
     ds["mask_rho"].attrs = {
         "long_name": "Mask at rho-points",
@@ -82,9 +86,8 @@ def _add_topography_and_mask(
 
 
 def _make_raw_topography(lon, lat, topography_source) -> np.ndarray:
-    """
-    Given a grid of (lon, lat) points, fetch the topography file and interpolate height values onto the desired grid.
-    """
+    """Given a grid of (lon, lat) points, fetch the topography file and interpolate
+    height values onto the desired grid."""
 
     topo_ds = fetch_topo(topography_source)
 
@@ -146,9 +149,7 @@ def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
 
 
 def _fill_enclosed_basins(mask) -> np.ndarray:
-    """
-    Fills in enclosed basins with land
-    """
+    """Fills in enclosed basins with land."""
 
     # Label connected regions in the mask
     reg, nreg = label(mask)
@@ -170,9 +171,7 @@ def _fill_enclosed_basins(mask) -> np.ndarray:
 
 
 def _smooth_topography_locally(h, hmin=5, rmax=0.2):
-    """
-    Smoothes topography locally to satisfy r < rmax
-    """
+    """Smoothes topography locally to satisfy r < rmax."""
     # Compute rmax_log
     if rmax > 0.0:
         rmax_log = np.log((1.0 + rmax * 0.9) / (1.0 - rmax * 0.9))
@@ -233,10 +232,7 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
         )
 
         # No gradient at the domain boundaries
-        h_log[0, :] = h_log[1, :]
-        h_log[-1, :] = h_log[-2, :]
-        h_log[:, 0] = h_log[:, 1]
-        h_log[:, -1] = h_log[:, -2]
+        h_log = _handle_boundaries(h_log)
 
         # Update h
         h = hmin * np.exp(h_log)
@@ -252,10 +248,32 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
     return h
 
 
-def _compute_rfactor(h):
-    """
-    Computes slope parameter (or r-factor) r = |Delta h| / 2h in both horizontal grid directions.
+def _handle_boundaries(field):
+    """Adjust the boundaries of a 2D field by copying values from adjacent cells.
+
+    Parameters
+    ----------
+    field : numpy.ndarray or xarray.DataArray
+        A 2D array representing a field (e.g., topography or mask) whose boundary values
+        need to be adjusted.
+
+    Returns
+    -------
+    field : numpy.ndarray or xarray.DataArray
+        The input field with adjusted boundary values.
     """
+
+    field[0, :] = field[1, :]
+    field[-1, :] = field[-2, :]
+    field[:, 0] = field[:, 1]
+    field[:, -1] = field[:, -2]
+
+    return field
+
+
+def _compute_rfactor(h):
+    """Computes slope parameter (or r-factor) r = |Delta h| / 2h in both horizontal grid
+    directions."""
     # compute r_{i-1/2} = |h_i - h_{i-1}| / (h_i + h_{i+1})
     r_eta = np.abs(h.diff("eta_rho")) / (h + h.shift(eta_rho=1)).isel(
         eta_rho=slice(1, None)