roms-tools 2.0.0__py3-none-any.whl → 2.2.0__py3-none-any.whl

roms_tools/setup/utils.py

@@ -745,6 +745,8 @@ def get_target_coords(grid, use_coarse_grid=False):
         mask = grid.ds.get("mask_coarse")
         if mask is not None:
             mask = mask.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
+            mask_u = interpolate_from_rho_to_u(mask, method="multiplicative")
+            mask_v = interpolate_from_rho_to_v(mask, method="multiplicative")
 
         lat_psi = grid.ds.get("lat_psi_coarse")
         lon_psi = grid.ds.get("lon_psi_coarse")
@@ -754,6 +756,8 @@ def get_target_coords(grid, use_coarse_grid=False):
         lon = grid.ds.lon_rho
         angle = grid.ds.angle
         mask = grid.ds.get("mask_rho")
+        mask_u = grid.ds.get("mask_u")
+        mask_v = grid.ds.get("mask_v")
         lat_psi = grid.ds.get("lat_psi")
         lon_psi = grid.ds.get("lon_psi")
 
@@ -776,6 +780,8 @@ def get_target_coords(grid, use_coarse_grid=False):
         "lon_psi": lon_psi,
         "angle": angle,
         "mask": mask,
+        "mask_u": mask_u,
+        "mask_v": mask_v,
         "straddle": straddle,
     }
 
@@ -927,37 +933,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,12 +1067,21 @@ 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)
-    grid_data.pop("verbose", None)
-    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}
+
+    # Ensure forcing_object.source.path is a string (convert if it's a pathlib object)
+    if hasattr(forcing_object, "source") and "path" in forcing_object.source:
+        forcing_object.source["path"] = str(forcing_object.source["path"])
 
     # Step 2: Get ROMS Tools version
     # Fetch the version of the 'roms-tools' package for inclusion in the YAML header
@@ -1082,7 +1100,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:
@@ -1111,6 +1138,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.
 
@@ -1203,3 +1238,84 @@ def handle_boundaries(field):
     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/tests/test_setup/test_boundary_forcing.py

@@ -19,7 +19,7 @@ import logging
 def test_boundary_forcing_creation(boundary_forcing_fixture, request):
     """Test the creation of the BoundaryForcing object."""
 
-    fname = download_test_data("GLORYS_coarse_test_data.nc")
+    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
     boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
     assert boundary_forcing.start_time == datetime(2021, 6, 29)
     assert boundary_forcing.end_time == datetime(2021, 6, 30)
@@ -59,7 +59,9 @@ def test_boundary_forcing_creation(boundary_forcing_fixture, request):
 def test_boundary_forcing_creation_with_bgc(boundary_forcing_fixture, request):
     """Test the creation of the BoundaryForcing object."""
 
-    fname_bgc = download_test_data("CESM_regional_coarse_test_data_climatology.nc")
+    fname_bgc = Path(
+        download_test_data("CESM_regional_coarse_test_data_climatology.nc")
+    )
     boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
 
     assert boundary_forcing.start_time == datetime(2021, 6, 29)
@@ -103,7 +105,7 @@ def test_unsuccessful_boundary_forcing_creation_with_1d_fill(use_dask):
         hc=250.0,  # critical depth
     )
 
-    fname = download_test_data("GLORYS_coarse_test_data.nc")
+    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
 
     with pytest.raises(ValueError, match="consists entirely of NaNs"):
 
@@ -138,7 +140,7 @@ def test_boundary_divided_by_land_warning(caplog, use_dask):
         nx=5, ny=5, size_x=500, size_y=500, center_lon=-10, center_lat=65, rot=0
     )
 
-    fname = download_test_data("GLORYS_coarse_test_data.nc")
+    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
 
     with caplog.at_level(logging.WARNING):
         BoundaryForcing(
@@ -158,7 +160,7 @@ def test_1d_and_2d_fill_coincide_if_no_land(use_dask):
     # this grid lies entirely over open ocean
     grid = Grid(nx=5, ny=5, size_x=300, size_y=300, center_lon=-5, center_lat=65, rot=0)
 
-    fname = download_test_data("GLORYS_coarse_test_data.nc")
+    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
 
     kwargs = {
         "grid": grid,

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

@@ -120,8 +120,7 @@
                 "nriver"
             ],
             "coordinates": "abs_time river_name tracer_name",
-            "long_name": "River tracer data",
-            "units": "degrees C [temperature]; psu [salinity]"
+            "long_name": "River tracer data"
         },
         "river_volume/.zarray": {
             "chunks": [
@@ -165,7 +164,7 @@
                 "id": "blosc",
                 "shuffle": 1
             },
-            "dtype": "<U11",
+            "dtype": "<U4",
             "fill_value": null,
             "filters": null,
             "order": "C",

roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_tracer/.zattrs

@@ -5,6 +5,5 @@
         "nriver"
     ],
     "coordinates": "abs_time river_name tracer_name",
-    "long_name": "River tracer data",
-    "units": "degrees C [temperature]; psu [salinity]"
+    "long_name": "River tracer data"
 }

roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_name/.zarray

@@ -9,7 +9,7 @@
         "id": "blosc",
         "shuffle": 1
     },
-    "dtype": "<U11",
+    "dtype": "<U4",
     "fill_value": null,
     "filters": null,
     "order": "C",

roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/.zmetadata

@@ -121,7 +121,7 @@
         "river_tracer/.zarray": {
             "chunks": [
                 12,
-                2,
+                34,
                 6
             ],
             "compressor": {
@@ -137,7 +137,7 @@
             "order": "C",
             "shape": [
                 12,
-                2,
+                34,
                 6
             ],
             "zarr_format": 2
@@ -149,8 +149,7 @@
                 "nriver"
             ],
             "coordinates": "abs_time month river_name tracer_name",
-            "long_name": "River tracer data",
-            "units": "degrees C [temperature]; psu [salinity]"
+            "long_name": "River tracer data"
         },
         "river_volume/.zarray": {
             "chunks": [
@@ -185,7 +184,7 @@
         },
         "tracer_name/.zarray": {
             "chunks": [
-                2
+                34
             ],
             "compressor": {
                 "blocksize": 0,
@@ -199,7 +198,7 @@
             "filters": null,
             "order": "C",
             "shape": [
-                2
+                34
             ],
             "zarr_format": 2
         },

roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_tracer/.zarray

@@ -1,7 +1,7 @@
 {
     "chunks": [
         12,
-        2,
+        34,
         6
     ],
     "compressor": {
@@ -17,7 +17,7 @@
     "order": "C",
     "shape": [
         12,
-        2,
+        34,
         6
     ],
     "zarr_format": 2

roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_tracer/.zattrs

@@ -5,6 +5,5 @@
         "nriver"
     ],
     "coordinates": "abs_time month river_name tracer_name",
-    "long_name": "River tracer data",
-    "units": "degrees C [temperature]; psu [salinity]"
+    "long_name": "River tracer data"
 }

roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/tracer_name/.zarray

@@ -1,6 +1,6 @@
 {
     "chunks": [
-        2
+        34
     ],
     "compressor": {
         "blocksize": 0,
@@ -14,7 +14,7 @@
     "filters": null,
     "order": "C",
     "shape": [
-        2
+        34
     ],
     "zarr_format": 2
 }

roms_tools/tests/test_setup/test_datasets.py

@@ -431,8 +431,8 @@ def test_era5_correction_choose_subdomain(use_dask):
     data = ERA5Correction(use_dask=use_dask)
     lats = data.ds.latitude[10:20]
     lons = data.ds.longitude[10:20]
-    coords = {"latitude": lats, "longitude": lons}
-    data.choose_subdomain(coords, straddle=False)
+    target_coords = {"lat": lats, "lon": lons}
+    data.choose_subdomain(target_coords, straddle=False)
     assert (data.ds["latitude"] == lats).all()
     assert (data.ds["longitude"] == lons).all()
 

roms_tools/tests/test_setup/test_initial_conditions.py

@@ -26,7 +26,7 @@ def test_initial_conditions_creation(ic_fixture, request):
     assert ic.ini_time == datetime(2021, 6, 29)
     assert ic.source == {
         "name": "GLORYS",
-        "path": download_test_data("GLORYS_coarse_test_data.nc"),
+        "path": Path(download_test_data("GLORYS_coarse_test_data.nc")),
         "climatology": False,
     }
     assert isinstance(ic.ds, xr.Dataset)
@@ -83,7 +83,7 @@ def test_initial_conditions_missing_physics_path(example_grid, use_dask):
 # Test initialization with missing 'name' in bgc_source
 def test_initial_conditions_missing_bgc_name(example_grid, use_dask):
 
-    fname = download_test_data("GLORYS_coarse_test_data.nc")
+    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
     with pytest.raises(
         ValueError, match="`bgc_source` must include a 'name' if it is provided."
     ):
@@ -99,7 +99,7 @@ def test_initial_conditions_missing_bgc_name(example_grid, use_dask):
 # Test initialization with missing 'path' in bgc_source
 def test_initial_conditions_missing_bgc_path(example_grid, use_dask):
 
-    fname = download_test_data("GLORYS_coarse_test_data.nc")
+    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
     with pytest.raises(
         ValueError, match="`bgc_source` must include a 'path' if it is provided."
     ):
@@ -115,7 +115,7 @@ def test_initial_conditions_missing_bgc_path(example_grid, use_dask):
 # Test default climatology value
 def test_initial_conditions_default_climatology(example_grid, use_dask):
 
-    fname = download_test_data("GLORYS_coarse_test_data.nc")
+    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
 
     initial_conditions = InitialConditions(
         grid=example_grid,
@@ -130,8 +130,8 @@ def test_initial_conditions_default_climatology(example_grid, use_dask):
 
 def test_initial_conditions_default_bgc_climatology(example_grid, use_dask):
 
-    fname = download_test_data("GLORYS_coarse_test_data.nc")
-    fname_bgc = download_test_data("CESM_regional_test_data_one_time_slice.nc")
+    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
+    fname_bgc = Path(download_test_data("CESM_regional_test_data_one_time_slice.nc"))
 
     initial_conditions = InitialConditions(
         grid=example_grid,