roms-tools 0.1.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

roms_tools/setup/download.py

@@ -0,0 +1,118 @@
+import pooch
+import xarray as xr
+
+# Create a Pooch object to manage the global topography data
+topo_data = pooch.create(
+    # Use the default cache folder for the operating system
+    path=pooch.os_cache("roms-tools"),
+    base_url="https://github.com/CWorthy-ocean/roms-tools-data/raw/main/",
+    # The registry specifies the files that can be fetched
+    registry={
+        "etopo5.nc": "sha256:23600e422d59bbf7c3666090166a0d468c8ee16092f4f14e32c4e928fbcd627b",
+    },
+)
+
+# Create a Pooch object to manage the global SWR correction data
+correction_data = pooch.create(
+    # Use the default cache folder for the operating system
+    path=pooch.os_cache("roms-tools"),
+    base_url="https://github.com/CWorthy-ocean/roms-tools-data/raw/main/",
+    # The registry specifies the files that can be fetched
+    registry={
+        "etopo5.nc": "sha256:23600e422d59bbf7c3666090166a0d468c8ee16092f4f14e32c4e928fbcd627b",
+        "SSR_correction.nc": "sha256:a170c1698e6cc2765b3f0bb51a18c6a979bc796ac3a4c014585aeede1f1f8ea0",
+    },
+)
+
+# Create a Pooch object to manage the test data
+pup_test_data = pooch.create(
+    # Use the default cache folder for the operating system
+    path=pooch.os_cache("roms-tools"),
+    base_url="https://github.com/CWorthy-ocean/roms-tools-test-data/raw/main/",
+    # The registry specifies the files that can be fetched
+    registry={
+        "GLORYS_test_data.nc": "648f88ec29c433bcf65f257c1fb9497bd3d5d3880640186336b10ed54f7129d2",
+        "ERA5_regional_test_data.nc": "bd12ce3b562fbea2a80a3b79ba74c724294043c28dc98ae092ad816d74eac794",
+        "ERA5_global_test_data.nc": "8ed177ab64c02caf509b9fb121cf6713f286cc603b1f302f15f3f4eb0c21dc4f",
+        "TPXO_global_test_data.nc": "457bfe87a7b247ec6e04e3c7d3e741ccf223020c41593f8ae33a14f2b5255e60",
+        "TPXO_regional_test_data.nc": "11739245e2286d9c9d342dce5221e6435d2072b50028bef2e86a30287b3b4032",
+        "CESM_regional_test_data_one_time_slice.nc": "43b578ecc067c85f95d6b97ed7b9dc8da7846f07c95331c6ba7f4a3161036a17",
+        "CESM_regional_test_data_climatology.nc": "986a200029d9478fd43e6e4a8bc43e8a8f4407554893c59b5fcc2e86fd203272",
+        "CESM_surface_global_test_data_climatology.nc": "a072757110c6f7b716a98f867688ef4195a5966741d2f368201ac24617254e35",
+        "CESM_surface_global_test_data.nc": "874106ffbc8b1b220db09df1551bbb89d22439d795b4d1e5a24ee775e9a7bf6e",
+    },
+)
+
+
+def fetch_topo(topography_source: str) -> xr.Dataset:
+    """
+    Load the global topography data as an xarray Dataset.
+
+    Parameters
+    ----------
+    topography_source : str
+        The source of the topography data to be loaded. Available options:
+        - "ETOPO5"
+
+    Returns
+    -------
+    xr.Dataset
+        The global topography data as an xarray Dataset.
+    """
+    # Mapping from user-specified topography options to corresponding filenames in the registry
+    topo_dict = {"ETOPO5": "etopo5.nc"}
+
+    # Fetch the file using Pooch, downloading if necessary
+    fname = topo_data.fetch(topo_dict[topography_source])
+
+    # Load the dataset using xarray and return it
+    ds = xr.open_dataset(fname)
+    return ds
+
+
+def download_correction_data(filename: str) -> str:
+    """
+    Download the correction data file.
+
+    Parameters
+    ----------
+    filename : str
+        The name of the test data file to be downloaded. Available options:
+        - "SSR_correction.nc"
+
+    Returns
+    -------
+    str
+        The path to the downloaded test data file.
+    """
+    # Fetch the file using Pooch, downloading if necessary
+    fname = correction_data.fetch(filename)
+
+    return fname
+
+
+def download_test_data(filename: str) -> str:
+    """
+    Download the test data file.
+
+    Parameters
+    ----------
+    filename : str
+        The name of the test data file to be downloaded. Available options:
+        - "GLORYS_test_data.nc"
+        - "ERA5_regional_test_data.nc"
+        - "ERA5_global_test_data.nc"
+        - "TPXO_global_test_data.nc"
+        - "TPXO_regional_test_data.nc"
+        - "CESM_regional_test_data_one_time_slice.nc"
+        - "CESM_regional_test_data_climatology.nc"
+
+    Returns
+    -------
+    str
+        The path to the downloaded test data file.
+    """
+    # Fetch the file using Pooch, downloading if necessary
+    fname = pup_test_data.fetch(filename)
+
+    return fname

roms_tools/setup/fill.py

@@ -3,7 +3,103 @@ import xarray as xr
 from numba import jit
 
 
-def lateral_fill(var, land_mask, dims=["latitude", "longitude"]):
+def fill_and_interpolate(
+    field,
+    mask,
+    fill_dims,
+    coords,
+    method="linear",
+    fillvalue_fill=0.0,
+    fillvalue_interp=np.nan,
+):
+    """
+    Propagates ocean values into land areas and interpolates the data to specified coordinates using a given method.
+
+    Parameters
+    ----------
+    field : xr.DataArray
+        The data array to be interpolated, typically containing oceanographic or atmospheric data
+        with dimensions such as latitude and longitude.
+
+    mask : xr.DataArray
+        A data array with the same spatial dimensions as `field`, where `1` indicates ocean points
+        and `0` indicates land points. This mask is used to identify land and ocean areas in the dataset.
+
+    fill_dims : list of str
+        List specifying the dimensions along which to perform the lateral fill, typically the horizontal
+        dimensions such as latitude and longitude, e.g., ["latitude", "longitude"].
+
+    coords : dict
+        Dictionary specifying the target coordinates for interpolation. The keys should match the dimensions
+        of `field` (e.g., {"longitude": lon_values, "latitude": lat_values, "depth": depth_values}).
+        This dictionary provides the new coordinates onto which the data array will be interpolated.
+
+    method : str, optional, default='linear'
+        The interpolation method to use. Valid options are those supported by `xarray.DataArray.interp`,
+        such as 'linear' or 'nearest'.
+
+    fillvalue_fill : float, optional, default=0.0
+        Value to use in the fill step if an entire data slice along the fill dimensions contains only NaNs.
+
+    fillvalue_interp : float, optional, default=np.nan
+        Value to use in the interpolation step. `np.nan` means that no extrapolation is applied.
+        `None` means that extrapolation is applied, which often makes sense when interpolating in the
+        vertical direction to avoid NaNs at the surface if the lowest depth is greater than zero.
+
+    Returns
+    -------
+    xr.DataArray
+        The interpolated data array. This array has the same dimensions as the input `field` but with values
+        interpolated to the new coordinates specified in `coords`.
+
+    Notes
+    -----
+    This method performs the following steps:
+    1. Sets land values to NaN based on the provided mask to ensure that interpolation does not cross
+       the land-ocean boundary.
+    2. Uses the `lateral_fill` function to propagate ocean values into the land interior, helping to fill
+       gaps in the dataset.
+    3. Interpolates the filled data array over the specified coordinates using the selected interpolation method.
+
+    Example
+    -------
+    >>> import xarray as xr
+    >>> field = xr.DataArray(...)
+    >>> mask = xr.DataArray(...)
+    >>> fill_dims = ["latitude", "longitude"]
+    >>> coords = {"latitude": new_lat_values, "longitude": new_lon_values}
+    >>> interpolated_field = fill_and_interpolate(
+    ...     field, mask, fill_dims, coords, method="linear"
+    ... )
+    >>> print(interpolated_field)
+    """
+    if not isinstance(field, xr.DataArray):
+        raise TypeError("field must be an xarray.DataArray")
+    if not isinstance(mask, xr.DataArray):
+        raise TypeError("mask must be an xarray.DataArray")
+    if not isinstance(coords, dict):
+        raise TypeError("coords must be a dictionary")
+    if not all(dim in field.dims for dim in coords.keys()):
+        raise ValueError("All keys in coords must match dimensions of field")
+    if method not in ["linear", "nearest"]:
+        raise ValueError(
+            "Unsupported interpolation method. Choose from 'linear', 'nearest'"
+        )
+
+    # Set land values to NaN
+    field = field.where(mask)
+
+    # Propagate ocean values into land interior before interpolation
+    field = lateral_fill(field, 1 - mask, fill_dims, fillvalue_fill)
+
+    field_interpolated = field.interp(
+        coords, method=method, kwargs={"fill_value": fillvalue_interp}
+    ).drop_vars(list(coords.keys()))
+
+    return field_interpolated
+
+
+def lateral_fill(var, land_mask, dims=["latitude", "longitude"], fillvalue=0.0):
     """
     Perform lateral fill on an xarray DataArray using a land mask.
 
@@ -20,6 +116,9 @@ def lateral_fill(var, land_mask, dims=["latitude", "longitude"]):
     dims : list of str, optional, default=['latitude', 'longitude']
         Dimensions along which to perform the fill. The default is ['latitude', 'longitude'].
 
+    fillvalue : float, optional, default=0.0
+        Value to use if an entire data slice along the dims contains only NaNs.
+
     Returns
     -------
     var_filled : xarray.DataArray
@@ -27,21 +126,25 @@ def lateral_fill(var, land_mask, dims=["latitude", "longitude"]):
         specified by `land_mask` where NaNs are preserved.
 
     """
+
     var_filled = xr.apply_ufunc(
         _lateral_fill_np_array,
         var,
         land_mask,
         input_core_dims=[dims, dims],
         output_core_dims=[dims],
-        dask="parallelized",
         output_dtypes=[var.dtype],
+        dask="parallelized",
         vectorize=True,
+        kwargs={"fillvalue": fillvalue},
     )
 
     return var_filled
 
 
-def _lateral_fill_np_array(var, isvalid_mask, tol=1.0e-4, rc=1.8, max_iter=10000):
+def _lateral_fill_np_array(
+    var, isvalid_mask, fillvalue=0.0, tol=1.0e-4, rc=1.8, max_iter=10000
+):
     """
     Perform lateral fill on a numpy array.
 
@@ -55,6 +158,9 @@ def _lateral_fill_np_array(var, isvalid_mask, tol=1.0e-4, rc=1.8, max_iter=10000
         Valid values mask: `True` where data should be filled. Must have same shape
         as `var`.
 
+    fillvalue: float
+        Value to use if the full field `var` contains only  NaNs. Default is 0.0.
+
     tol : float, optional, default=1.0e-4
         Convergence criteria: stop filling when the value change is less than
         or equal to `tol * var`, i.e., `delta <= tol * np.abs(var[j, i])`.
@@ -90,14 +196,14 @@ def _lateral_fill_np_array(var, isvalid_mask, tol=1.0e-4, rc=1.8, max_iter=10000
 
     fillmask = np.isnan(var)  # Fill all NaNs
     keepNaNs = ~isvalid_mask & np.isnan(var)
-    var = _iterative_fill_sor(nlat, nlon, var, fillmask, tol, rc, max_iter)
+    var = _iterative_fill_sor(nlat, nlon, var, fillmask, tol, rc, max_iter, fillvalue)
     var[keepNaNs] = np.nan  # Replace NaNs in areas not designated for filling
 
     return var
 
 
 @jit(nopython=True, parallel=True)
-def _iterative_fill_sor(nlat, nlon, var, fillmask, tol, rc, max_iter):
+def _iterative_fill_sor(nlat, nlon, var, fillmask, tol, rc, max_iter, fillvalue=0.0):
     """
     Perform an iterative land fill algorithm using the Successive Over-Relaxation (SOR)
     solution of the Laplace Equation.
@@ -126,6 +232,9 @@ def _iterative_fill_sor(nlat, nlon, var, fillmask, tol, rc, max_iter):
     max_iter : int
         Maximum number of iterations allowed before the process is terminated.
 
+    fillvalue: float
+        Value to use if the full field is NaNs. Default is 0.0.
+
     Returns
     -------
     None
@@ -155,6 +264,10 @@ def _iterative_fill_sor(nlat, nlon, var, fillmask, tol, rc, max_iter):
     if np.max(np.fabs(var)) == 0.0:
         var = np.zeros_like(var)
         return var
+    # If field consists only of NaNs, fill NaNs with fill value
+    if np.isnan(var).all():
+        var = fillvalue * np.ones_like(var)
+        return var
 
     # Compute a zonal mean to use as a first guess
     zoncnt = np.zeros(nlat)

roms_tools/setup/grid.py

@@ -1,12 +1,16 @@
 import copy
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, asdict
 
 import numpy as np
 import xarray as xr
+import yaml
+import importlib.metadata
 
-
-from roms_tools.setup.topography import _add_topography_and_mask
+from roms_tools.setup.topography import _add_topography_and_mask, _add_velocity_masks
 from roms_tools.setup.plot import _plot
+from roms_tools.setup.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
+
+import warnings
 
 RADIUS_OF_EARTH = 6371315.0  # in m
 
@@ -41,7 +45,7 @@ class Grid:
          The default is 0, which means that the x-direction of the grid is aligned with lines of constant latitude.
     topography_source : str, optional
         Specifies the data source to use for the topography. Options are
-        "etopo5". The default is "etopo5".
+        "ETOPO5". The default is "ETOPO5".
     smooth_factor : float, optional
         The smoothing factor used in the domain-wide Gaussian smoothing of the
         topography. Smaller values result in less smoothing, while larger
@@ -97,7 +101,7 @@ class Grid:
     center_lon: float
     center_lat: float
     rot: float = 0
-    topography_source: str = "etopo5"
+    topography_source: str = "ETOPO5"
     smooth_factor: int = 8
     hmin: float = 5.0
     rmax: float = 0.2
@@ -130,7 +134,7 @@ class Grid:
         self._straddle()
 
     def add_topography_and_mask(
-        self, topography_source="etopo5", smooth_factor=8, hmin=5.0, rmax=0.2
+        self, topography_source="ETOPO5", smooth_factor=8, hmin=5.0, rmax=0.2
     ) -> None:
         """
         Add topography and mask to the grid dataset.
@@ -144,7 +148,7 @@ class Grid:
         ----------
         topography_source : str, optional
             Specifies the data source to use for the topography. Options are
-            "etopo5". The default is "etopo5".
+            "ETOPO5". The default is "ETOPO5".
         smooth_factor : float, optional
             The smoothing factor used in the domain-wide Gaussian smoothing of the
             topography. Smaller values result in less smoothing, while larger
@@ -204,6 +208,37 @@ class Grid:
         """
         self.ds.to_netcdf(filepath)
 
+    def to_yaml(self, filepath: str) -> None:
+        """
+        Export the parameters of the class to a YAML file, including the version of roms-tools.
+
+        Parameters
+        ----------
+        filepath : str
+            The path to the YAML file where the parameters will be saved.
+        """
+        data = asdict(self)
+        data.pop("ds", None)
+        data.pop("straddle", None)
+
+        # Include the version of roms-tools
+        try:
+            roms_tools_version = importlib.metadata.version("roms-tools")
+        except importlib.metadata.PackageNotFoundError:
+            roms_tools_version = "unknown"
+
+        # Create header
+        header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
+
+        # Use the class name as the top-level key
+        yaml_data = {self.__class__.__name__: data}
+
+        with open(filepath, "w") as file:
+            # Write header
+            file.write(header)
+            # Write YAML data
+            yaml.dump(yaml_data, file, default_flow_style=False)
+
     @classmethod
     def from_file(cls, filepath: str) -> "Grid":
         """
@@ -222,6 +257,11 @@ class Grid:
         # Load the dataset from the file
         ds = xr.open_dataset(filepath)
 
+        if not all(mask in ds for mask in ["mask_u", "mask_v"]):
+            ds = _add_velocity_masks(ds)
+        if not all(coord in ds for coord in ["lat_u", "lon_u", "lat_v", "lon_v"]):
+            ds = _add_lat_lon_at_velocity_points(ds)
+
         # Create a new Grid instance without calling __init__ and __post_init__
         grid = cls.__new__(cls)
 
@@ -251,6 +291,62 @@ class Grid:
 
         return grid
 
+    @classmethod
+    def from_yaml(cls, filepath: str) -> "Grid":
+        """
+        Create an instance of the class from a YAML file.
+
+        Parameters
+        ----------
+        filepath : str
+            The path to the YAML file from which the parameters will be read.
+
+        Returns
+        -------
+        Grid
+            An instance of the Grid class.
+        """
+        # Read the entire file content
+        with open(filepath, "r") as file:
+            file_content = file.read()
+
+        # Split the content into YAML documents
+        documents = list(yaml.safe_load_all(file_content))
+
+        header_data = None
+        grid_data = None
+
+        # Iterate over documents to find the header and grid configuration
+        for doc in documents:
+            if doc is None:
+                continue
+            if "roms_tools_version" in doc:
+                header_data = doc
+            elif "Grid" in doc:
+                grid_data = doc["Grid"]
+
+        if header_data is None:
+            raise ValueError("Version of ROMS-Tools not found in the YAML file.")
+        else:
+            # Check the roms_tools_version
+            roms_tools_version_header = header_data.get("roms_tools_version")
+            # Get current version of roms-tools
+            try:
+                roms_tools_version_current = importlib.metadata.version("roms-tools")
+            except importlib.metadata.PackageNotFoundError:
+                roms_tools_version_current = "unknown"
+
+            if roms_tools_version_header != roms_tools_version_current:
+                warnings.warn(
+                    f"Current roms-tools version ({roms_tools_version_current}) does not match the version in the YAML header ({roms_tools_version_header}).",
+                    UserWarning,
+                )
+
+        if grid_data is None:
+            raise ValueError("No Grid configuration found in the YAML file.")
+
+        return cls(**grid_data)
+
     # override __repr__ method to only print attributes that are actually set
     def __repr__(self) -> str:
         cls = self.__class__
@@ -304,6 +400,7 @@ class Grid:
 
         if bathymetry:
             kwargs = {"cmap": "YlGnBu"}
+
             _plot(
                 self.ds,
                 field=self.ds.h.where(self.ds.mask_rho),
@@ -719,6 +816,18 @@ def _create_grid_ds(
         },
     )
 
+    ds["tra_lon"] = center_lon
+    ds["tra_lon"].attrs["long_name"] = "Longitudinal translation of base grid"
+    ds["tra_lon"].attrs["units"] = "degrees East"
+
+    ds["tra_lat"] = center_lat
+    ds["tra_lat"].attrs["long_name"] = "Latitudinal translation of base grid"
+    ds["tra_lat"].attrs["units"] = "degrees North"
+
+    ds["rotate"] = rot
+    ds["rotate"].attrs["long_name"] = "Rotation of base grid"
+    ds["rotate"].attrs["units"] = "degrees"
+
     ds["lon_rho"] = xr.Variable(
         data=lon * 180 / np.pi,
         dims=["eta_rho", "xi_rho"],
@@ -731,23 +840,21 @@ def _create_grid_ds(
         attrs={"long_name": "latitude of rho-points", "units": "degrees North"},
     )
 
-    ds["tra_lon"] = center_lon
-    ds["tra_lon"].attrs["long_name"] = "Longitudinal translation of base grid"
-    ds["tra_lon"].attrs["units"] = "degrees East"
-
-    ds["tra_lat"] = center_lat
-    ds["tra_lat"].attrs["long_name"] = "Latitudinal translation of base grid"
-    ds["tra_lat"].attrs["units"] = "degrees North"
-
-    ds["rotate"] = rot
-    ds["rotate"].attrs["long_name"] = "Rotation of base grid"
-    ds["rotate"].attrs["units"] = "degrees"
+    ds = _add_lat_lon_at_velocity_points(ds)
 
     return ds
 
 
 def _add_global_metadata(ds, size_x, size_y):
-    ds.attrs["Type"] = "ROMS grid produced by roms-tools"
+    ds.attrs["title"] = "ROMS grid created by ROMS-Tools"
+
+    # Include the version of roms-tools
+    try:
+        roms_tools_version = importlib.metadata.version("roms-tools")
+    except importlib.metadata.PackageNotFoundError:
+        roms_tools_version = "unknown"
+
+    ds.attrs["roms_tools_version"] = roms_tools_version
     ds.attrs["size_x"] = size_x
     ds.attrs["size_y"] = size_y
 
@@ -807,3 +914,22 @@ def _f2c_xdir(f):
     fc[-1, :] = f[-1, :] + 0.5 * (f[-1, :] - f[-2, :])
 
     return fc
+
+
+def _add_lat_lon_at_velocity_points(ds):
+
+    lat_u = interpolate_from_rho_to_u(ds["lat_rho"])
+    lon_u = interpolate_from_rho_to_u(ds["lon_rho"])
+    lat_v = interpolate_from_rho_to_v(ds["lat_rho"])
+    lon_v = interpolate_from_rho_to_v(ds["lon_rho"])
+
+    lat_u.attrs = {"long_name": "latitude of u-points", "units": "degrees North"}
+    lon_u.attrs = {"long_name": "longitude of u-points", "units": "degrees East"}
+    lat_v.attrs = {"long_name": "latitude of v-points", "units": "degrees North"}
+    lon_v.attrs = {"long_name": "longitude of v-points", "units": "degrees East"}
+
+    ds = ds.assign_coords(
+        {"lat_u": lat_u, "lon_u": lon_u, "lat_v": lat_v, "lon_v": lon_v}
+    )
+
+    return ds