roms-tools 1.7.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

roms_tools/setup/utils.py

@@ -41,7 +41,7 @@ def nan_check(field, mask, error_message=None) -> None:
     da = xr.where(mask == 1, field, 0)
     if error_message is None:
         error_message = (
-            "NaN values found in interpolated field. This likely occurs because the ROMS grid, including "
+            "NaN values found in regridded field. This likely occurs because the ROMS grid, including "
             "a small safety margin for interpolation, is not fully contained within the dataset's longitude/latitude range. Please ensure that the "
             "dataset covers the entire area required by the ROMS grid."
         )
@@ -107,10 +107,10 @@ def interpolate_from_rho_to_u(field, method="additive"):
     else:
         raise NotImplementedError(f"Unsupported method '{method}' specified.")
 
-    if "lat_rho" in field_interpolated.coords:
-        field_interpolated.drop_vars(["lat_rho"])
-    if "lon_rho" in field_interpolated.coords:
-        field_interpolated.drop_vars(["lon_rho"])
+    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"})
 
@@ -155,10 +155,10 @@ def interpolate_from_rho_to_v(field, method="additive"):
     else:
         raise NotImplementedError(f"Unsupported method '{method}' specified.")
 
-    if "lat_rho" in field_interpolated.coords:
-        field_interpolated.drop_vars(["lat_rho"])
-    if "lon_rho" in field_interpolated.coords:
-        field_interpolated.drop_vars(["lon_rho"])
+    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"})
 
@@ -733,27 +733,27 @@ def get_target_coords(grid, use_coarse_grid=False):
     """
     # Select grid variables based on whether the coarse grid is used
     if use_coarse_grid:
-        lat, lon, angle, mask = (
-            grid.ds.lat_coarse.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}),
-            grid.ds.lon_coarse.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}),
-            grid.ds.angle_coarse.rename(
-                {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
-            ),
-            grid.ds.mask_coarse.rename(
-                {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
-            ),
+        lat = grid.ds.lat_coarse.rename(
+            {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
         )
+        lon = grid.ds.lon_coarse.rename(
+            {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
+        )
+        angle = grid.ds.angle_coarse.rename(
+            {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
+        )
+        mask = grid.ds.get("mask_coarse")
+        if mask is not None:
+            mask = mask.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
 
         lat_psi = grid.ds.get("lat_psi_coarse")
         lon_psi = grid.ds.get("lon_psi_coarse")
 
     else:
-        lat, lon, angle, mask = (
-            grid.ds.lat_rho,
-            grid.ds.lon_rho,
-            grid.ds.angle,
-            grid.ds.mask_rho,
-        )
+        lat = grid.ds.lat_rho
+        lon = grid.ds.lon_rho
+        angle = grid.ds.angle
+        mask = grid.ds.get("mask_rho")
         lat_psi = grid.ds.get("lat_psi")
         lon_psi = grid.ds.get("lon_psi")
 
@@ -927,37 +927,40 @@ def get_vector_pairs(variable_info):
     return vector_pairs
 
 
-def gc_dist(lon1, lat1, lon2, lat2):
-    """Calculate the great circle distance between two points on the Earth's surface.
-    Latitude and longitude must be provided in degrees (they will be converted to
-    radians).
+def gc_dist(lon1, lat1, lon2, lat2, input_in_degrees=True):
+    """Calculate the great circle distance between two points on the Earth's surface
+    using the Haversine formula.
 
-    The function uses the Haversine formula to compute the shortest distance
-    along the surface of a sphere (Earth), assuming the Earth is a perfect sphere.
+    Latitude and longitude are assumed to be in degrees by default. If `input_in_degrees` is set to `False`,
+    the input is assumed to already be in radians.
 
     Parameters
     ----------
     lon1, lat1 : float
-        Longitude and latitude of the first point in degrees.
+        Longitude and latitude of the first point.
     lon2, lat2 : float
-        Longitude and latitude of the second point in degrees.
+        Longitude and latitude of the second point.
+    input_in_degrees : bool, optional
+        If True (default), the input coordinates are assumed to be in degrees and will be converted to radians.
+        If False, the input is assumed to be in radians and no conversion is applied.
 
     Returns
     -------
     dis : float
         The great circle distance between the two points in meters.
-        This is the shortest distance along the surface of a sphere (Earth).
 
     Notes
     -----
     The radius of the Earth is taken to be 6371315 meters.
     """
+
     # Convert degrees to radians
-    d2r = np.pi / 180
-    lon1 = lon1 * d2r
-    lat1 = lat1 * d2r
-    lon2 = lon2 * d2r
-    lat2 = lat2 * d2r
+    if input_in_degrees:
+        d2r = np.pi / 180
+        lon1 = lon1 * d2r
+        lat1 = lat1 * d2r
+        lon2 = lon2 * d2r
+        lat2 = lat2 * d2r
 
     # Difference in latitudes and longitudes
     dlat = lat2 - lat1
@@ -1058,11 +1061,17 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
     filepath = Path(filepath)
 
     # Step 1: Serialize Grid data
-    # Convert the grid attribute to a dictionary and remove non-serializable fields
-    grid_data = asdict(forcing_object.grid)
-    grid_data.pop("ds", None)  # Remove 'ds' attribute (non-serializable)
-    grid_data.pop("straddle", None)  # Remove 'straddle' if it's non-essential
-    grid_yaml_data = {"Grid": grid_data}
+    # Check if the forcing_object has a grid attribute
+    if hasattr(forcing_object, "grid") and forcing_object.grid is not None:
+        grid_data = asdict(forcing_object.grid)
+        grid_yaml_data = {"Grid": _pop_grid_data(grid_data)}
+    else:
+        parent_grid_data = asdict(forcing_object.parent_grid)
+        parent_grid_yaml_data = {"ParentGrid": _pop_grid_data(parent_grid_data)}
+        child_grid_data = asdict(forcing_object.child_grid)
+        child_grid_yaml_data = {"ChildGrid": _pop_grid_data(child_grid_data)}
+
+        grid_yaml_data = {**parent_grid_yaml_data, **child_grid_yaml_data}
 
     # Step 2: Get ROMS Tools version
     # Fetch the version of the 'roms-tools' package for inclusion in the YAML header
@@ -1081,7 +1090,16 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
     filtered_field_names = [
         param
         for param in field_names
-        if param not in ("grid", "ds", "use_dask", "climatology")
+        if param
+        not in (
+            "grid",
+            "parent_grid",
+            "child_grid",
+            "ds",
+            "use_dask",
+            "bypass_validation",
+            "climatology",
+        )
     ]
 
     for field_name in filtered_field_names:
@@ -1110,6 +1128,14 @@ def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
         yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
 
 
+def _pop_grid_data(grid_data):
+    grid_data.pop("ds", None)  # Remove 'ds' attribute (non-serializable)
+    grid_data.pop("straddle", None)
+    grid_data.pop("verbose", None)
+
+    return grid_data
+
+
 def _from_yaml(forcing_object: Type, filepath: Union[str, Path]) -> Dict[str, Any]:
     """Extract the configuration data for a given forcing object from a YAML file.
 
@@ -1179,3 +1205,107 @@ def _convert_from_iso_format(value):
     except ValueError:
         # Return None or raise an exception if parsing fails
         return value
+
+
+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 get_boundary_coords():
+    """This function determines the boundary points for the grid variables by specifying
+    the indices for the south, east, north, and west boundaries.
+
+    Returns
+    -------
+    dict
+        A dictionary containing the boundary coordinates for different variable types.
+        The dictionary has the following structure:
+        - Keys: Variable types ("rho", "u", "v", "vector").
+        - Values: Nested dictionaries that map each direction ("south", "east", "north", "west")
+          to another dictionary specifying the boundary coordinates, represented by grid indices
+          for the respective variable types. For example:
+          - "rho" variables (e.g., `eta_rho`, `xi_rho`)
+          - "u" variables (e.g., `xi_u`)
+          - "v" variables (e.g., `eta_v`)
+          - "vector" variables with lists of indices for multiple grid points (e.g., `eta_rho`, `xi_rho`).
+    """
+
+    bdry_coords = {
+        "rho": {
+            "south": {"eta_rho": 0},
+            "east": {"xi_rho": -1},
+            "north": {"eta_rho": -1},
+            "west": {"xi_rho": 0},
+        },
+        "u": {
+            "south": {"eta_rho": 0},
+            "east": {"xi_u": -1},
+            "north": {"eta_rho": -1},
+            "west": {"xi_u": 0},
+        },
+        "v": {
+            "south": {"eta_v": 0},
+            "east": {"xi_rho": -1},
+            "north": {"eta_v": -1},
+            "west": {"xi_rho": 0},
+        },
+        "vector": {
+            "south": {"eta_rho": [0, 1]},
+            "east": {"xi_rho": [-2, -1]},
+            "north": {"eta_rho": [-2, -1]},
+            "west": {"xi_rho": [0, 1]},
+        },
+    }
+
+    return bdry_coords
+
+
+def wrap_longitudes(grid_ds, straddle):
+    """Adjusts longitude values in a dataset to handle dateline crossing.
+
+    Parameters
+    ----------
+    grid_ds : xr.Dataset
+        The dataset containing longitude variables to adjust.
+    straddle : bool
+        If True, adjusts longitudes to the range [-180, 180] for datasets
+        that straddle the dateline. If False, adjusts longitudes to the
+        range [0, 360].
+
+    Returns
+    -------
+    xr.Dataset
+        The dataset with adjusted longitude values.
+    """
+    for lon_dim in ["lon_rho", "lon_u", "lon_v"]:
+        if straddle:
+            grid_ds[lon_dim] = xr.where(
+                grid_ds[lon_dim] > 180,
+                grid_ds[lon_dim] - 360,
+                grid_ds[lon_dim],
+            )
+        else:
+            grid_ds[lon_dim] = xr.where(
+                grid_ds[lon_dim] < 0, grid_ds[lon_dim] + 360, grid_ds[lon_dim]
+            )
+
+    return grid_ds

roms_tools/setup/vertical_coordinate.py

@@ -1,5 +1,6 @@
 import numpy as np
 import xarray as xr
+from roms_tools.setup.utils import transpose_dimensions
 
 
 def compute_cs(sigma, theta_s, theta_b):
@@ -83,7 +84,7 @@ def compute_depth(zeta, h, hc, cs, sigma):
 
     Parameters
     ----------
-    zeta : xr.DataArray
+    zeta : xr.DataArray or scalar
         The sea surface height.
     h : xr.DataArray
         The depth of the sea bottom.
@@ -98,23 +99,11 @@ def compute_depth(zeta, h, hc, cs, sigma):
     -------
     z : xr.DataArray
         The depth at different sigma levels.
-
-    Raises
-    ------
-    ValueError
-        If theta_s or theta_b are less than or equal to zero.
     """
 
-    # Expand dimensions
-    sigma = sigma.expand_dims(dim={"eta_rho": h.eta_rho, "xi_rho": h.xi_rho})
-    cs = cs.expand_dims(dim={"eta_rho": h.eta_rho, "xi_rho": h.xi_rho})
-
-    s = (hc * sigma + h * cs) / (hc + h)
-    z = zeta + (zeta + h) * s
+    z = (hc * sigma + h * cs) / (hc + h)
+    z = zeta + (zeta + h) * z
 
-    if "s_rho" in z.dims:
-        z = z.transpose("s_rho", "eta_rho", "xi_rho")
-    elif "s_w" in z.dims:
-        z = z.transpose("s_w", "eta_rho", "xi_rho")
+    z = -transpose_dimensions(z)
 
     return z

roms_tools/tests/test_setup/test_boundary_forcing.py

@@ -241,10 +241,14 @@ def test_boundary_forcing_save(boundary_forcing, tmp_path):
 def test_bgc_boundary_forcing_plot(bgc_boundary_forcing_from_climatology):
     """Test plot method."""
 
-    bgc_boundary_forcing_from_climatology.plot(var_name="ALK_south")
-    bgc_boundary_forcing_from_climatology.plot(var_name="ALK_east")
-    bgc_boundary_forcing_from_climatology.plot(var_name="ALK_north")
-    bgc_boundary_forcing_from_climatology.plot(var_name="ALK_west")
+    bgc_boundary_forcing_from_climatology.plot(
+        var_name="ALK_south", layer_contours=True
+    )
+    bgc_boundary_forcing_from_climatology.plot(var_name="ALK_east", layer_contours=True)
+    bgc_boundary_forcing_from_climatology.plot(
+        var_name="ALK_north", layer_contours=True
+    )
+    bgc_boundary_forcing_from_climatology.plot(var_name="ALK_west", layer_contours=True)
 
 
 def test_bgc_boundary_forcing_save(bgc_boundary_forcing_from_climatology, tmp_path):
@@ -390,7 +394,8 @@ def test_from_yaml_missing_boundary_forcing(tmp_path, request, use_dask):
       center_lon: -10
       center_lat: 61
       rot: -20
-      topography_source: ETOPO5
+      topography_source:
+        name: ETOPO5
       hmin: 5.0
     """
     )

roms_tools/tests/test_setup/test_data/grid.zarr/.zattrs

@@ -1,7 +1,6 @@
 {
     "center_lat": 0,
     "center_lon": -20,
-    "coordinates": "interface_depth_rho interface_depth_u interface_depth_v layer_depth_rho layer_depth_u layer_depth_v",
     "hc": 300.0,
     "hmin": 5.0,
     "roms_tools_version": "0.1.dev157+dirty",

roms_tools/tests/test_setup/test_data/grid.zarr/.zmetadata

@@ -3,7 +3,6 @@
         ".zattrs": {
             "center_lat": 0,
             "center_lon": -20,
-            "coordinates": "interface_depth_rho interface_depth_u interface_depth_v layer_depth_rho layer_depth_u layer_depth_v",
             "hc": 300.0,
             "hmin": 5.0,
             "roms_tools_version": "0.1.dev157+dirty",
@@ -42,7 +41,7 @@
             "_ARRAY_DIMENSIONS": [
                 "s_rho"
             ],
-            "long_name": "S-coordinate stretching curves at rho-points",
+            "long_name": "Vertical stretching function at rho-points",
             "units": "nondimensional"
         },
         "Cs_w/.zarray": {
@@ -69,7 +68,7 @@
             "_ARRAY_DIMENSIONS": [
                 "s_w"
             ],
-            "long_name": "S-coordinate stretching curves at w-points",
+            "long_name": "Vertical stretching function at w-points",
             "units": "nondimensional"
         },
         "angle/.zarray": {
@@ -196,105 +195,6 @@
             "long_name": "Final bathymetry at rho-points",
             "units": "meter"
         },
-        "interface_depth_rho/.zarray": {
-            "chunks": [
-                101,
-                3,
-                3
-            ],
-            "compressor": {
-                "blocksize": 0,
-                "clevel": 5,
-                "cname": "lz4",
-                "id": "blosc",
-                "shuffle": 1
-            },
-            "dtype": "<f4",
-            "fill_value": "NaN",
-            "filters": null,
-            "order": "C",
-            "shape": [
-                101,
-                3,
-                3
-            ],
-            "zarr_format": 2
-        },
-        "interface_depth_rho/.zattrs": {
-            "_ARRAY_DIMENSIONS": [
-                "s_w",
-                "eta_rho",
-                "xi_rho"
-            ],
-            "long_name": "Interface depth at rho-points",
-            "units": "m"
-        },
-        "interface_depth_u/.zarray": {
-            "chunks": [
-                101,
-                3,
-                2
-            ],
-            "compressor": {
-                "blocksize": 0,
-                "clevel": 5,
-                "cname": "lz4",
-                "id": "blosc",
-                "shuffle": 1
-            },
-            "dtype": "<f4",
-            "fill_value": "NaN",
-            "filters": null,
-            "order": "C",
-            "shape": [
-                101,
-                3,
-                2
-            ],
-            "zarr_format": 2
-        },
-        "interface_depth_u/.zattrs": {
-            "_ARRAY_DIMENSIONS": [
-                "s_w",
-                "eta_rho",
-                "xi_u"
-            ],
-            "long_name": "Interface depth at u-points",
-            "units": "m"
-        },
-        "interface_depth_v/.zarray": {
-            "chunks": [
-                101,
-                2,
-                3
-            ],
-            "compressor": {
-                "blocksize": 0,
-                "clevel": 5,
-                "cname": "lz4",
-                "id": "blosc",
-                "shuffle": 1
-            },
-            "dtype": "<f4",
-            "fill_value": "NaN",
-            "filters": null,
-            "order": "C",
-            "shape": [
-                101,
-                2,
-                3
-            ],
-            "zarr_format": 2
-        },
-        "interface_depth_v/.zattrs": {
-            "_ARRAY_DIMENSIONS": [
-                "s_w",
-                "eta_v",
-                "xi_rho"
-            ],
-            "long_name": "Interface depth at v-points",
-            "units": "m"
-        },
         "lat_coarse/.zarray": {
             "chunks": [
                 2,
@@ -415,105 +315,6 @@
             "long_name": "latitude of v-points",
             "units": "degrees North"
         },
-        "layer_depth_rho/.zarray": {
-            "chunks": [
-                100,
-                3,
-                3
-            ],
-            "compressor": {
-                "blocksize": 0,
-                "clevel": 5,
-                "cname": "lz4",
-                "id": "blosc",
-                "shuffle": 1
-            },
-            "dtype": "<f4",
-            "fill_value": "NaN",
-            "filters": null,
-            "order": "C",
-            "shape": [
-                100,
-                3,
-                3
-            ],
-            "zarr_format": 2
-        },
-        "layer_depth_rho/.zattrs": {
-            "_ARRAY_DIMENSIONS": [
-                "s_rho",
-                "eta_rho",
-                "xi_rho"
-            ],
-            "long_name": "Layer depth at rho-points",
-            "units": "m"
-        },
-        "layer_depth_u/.zarray": {
-            "chunks": [
-                100,
-                3,
-                2
-            ],
-            "compressor": {
-                "blocksize": 0,
-                "clevel": 5,
-                "cname": "lz4",
-                "id": "blosc",
-                "shuffle": 1
-            },
-            "dtype": "<f4",
-            "fill_value": "NaN",
-            "filters": null,
-            "order": "C",
-            "shape": [
-                100,
-                3,
-                2
-            ],
-            "zarr_format": 2
-        },
-        "layer_depth_u/.zattrs": {
-            "_ARRAY_DIMENSIONS": [
-                "s_rho",
-                "eta_rho",
-                "xi_u"
-            ],
-            "long_name": "Layer depth at u-points",
-            "units": "m"
-        },
-        "layer_depth_v/.zarray": {
-            "chunks": [
-                100,
-                2,
-                3
-            ],
-            "compressor": {
-                "blocksize": 0,
-                "clevel": 5,
-                "cname": "lz4",
-                "id": "blosc",
-                "shuffle": 1
-            },
-            "dtype": "<f4",
-            "fill_value": "NaN",
-            "filters": null,
-            "order": "C",
-            "shape": [
-                100,
-                2,
-                3
-            ],
-            "zarr_format": 2
-        },
-        "layer_depth_v/.zattrs": {
-            "_ARRAY_DIMENSIONS": [
-                "s_rho",
-                "eta_v",
-                "xi_rho"
-            ],
-            "long_name": "Layer depth at v-points",
-            "units": "m"
-        },
         "lon_coarse/.zarray": {
             "chunks": [
                 2,
@@ -820,6 +621,60 @@
             "long_name": "Curvilinear coordinate metric in eta-direction",
             "units": "meter-1"
         },
+        "sigma_r/.zarray": {
+            "chunks": [
+                100
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<f4",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [
+                100
+            ],
+            "zarr_format": 2
+        },
+        "sigma_r/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "s_rho"
+            ],
+            "long_name": "Fractional vertical stretching coordinate at rho-points",
+            "units": "nondimensional"
+        },
+        "sigma_w/.zarray": {
+            "chunks": [
+                101
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<f4",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [
+                101
+            ],
+            "zarr_format": 2
+        },
+        "sigma_w/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "s_w"
+            ],
+            "long_name": "Fractional vertical stretching coordinate at w-points",
+            "units": "nondimensional"
+        },
         "spherical/.zarray": {
             "chunks": [],
             "compressor": null,

roms_tools/tests/test_setup/test_data/grid.zarr/Cs_r/.zattrs

@@ -2,6 +2,6 @@
     "_ARRAY_DIMENSIONS": [
         "s_rho"
     ],
-    "long_name": "S-coordinate stretching curves at rho-points",
+    "long_name": "Vertical stretching function at rho-points",
     "units": "nondimensional"
 }