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

roms_tools/setup/utils.py

@@ -1,10 +1,14 @@
 import xarray as xr
 import numpy as np
-from typing import Union
+from typing import Union, Any, Dict, Type
 import pandas as pd
 import cftime
 from roms_tools.utils import partition
 from pathlib import Path
+from datetime import datetime
+from dataclasses import fields, asdict
+import importlib.metadata
+import yaml
 
 
 def nan_check(field, mask, error_message=None) -> None:
@@ -724,23 +728,31 @@ def get_target_coords(grid, use_coarse_grid=False):
     Returns
     -------
     dict
-        Dictionary containing the longitude, latitude, and angle arrays, along with a boolean indicating
+        Dictionary containing the longitude, latitude, angle and mask arrays, along with a boolean indicating
         if the grid straddles the meridian.
     """
     # Select grid variables based on whether the coarse grid is used
     if use_coarse_grid:
-        lat, lon, angle = (
-            grid.ds.lat_coarse,
-            grid.ds.lon_coarse,
-            grid.ds.angle_coarse,
+        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_psi = grid.ds.get("lat_psi_coarse")
         lon_psi = grid.ds.get("lon_psi_coarse")
+
     else:
-        lat, lon, angle = (
+        lat, lon, angle, mask = (
             grid.ds.lat_rho,
             grid.ds.lon_rho,
             grid.ds.angle,
+            grid.ds.mask_rho,
         )
         lat_psi = grid.ds.get("lat_psi")
         lon_psi = grid.ds.get("lon_psi")
@@ -763,6 +775,7 @@ def get_target_coords(grid, use_coarse_grid=False):
         "lat_psi": lat_psi,
         "lon_psi": lon_psi,
         "angle": angle,
+        "mask": mask,
         "straddle": straddle,
     }
 
@@ -912,3 +925,257 @@ def get_vector_pairs(variable_info):
                 processed.update([var_name, vector_pair])
 
     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).
+
+    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.
+
+    Parameters
+    ----------
+    lon1, lat1 : float
+        Longitude and latitude of the first point in degrees.
+    lon2, lat2 : float
+        Longitude and latitude of the second point in degrees.
+
+    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
+
+    # Difference in latitudes and longitudes
+    dlat = lat2 - lat1
+    dlon = lon2 - lon1
+
+    # Haversine formula
+    dang = 2 * np.arcsin(
+        np.sqrt(
+            np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
+        )
+    )
+
+    # Radius of the Earth in meters
+    r_earth = 6371315.0
+
+    # Distance in meters
+    dis = r_earth * dang
+
+    return dis
+
+
+def convert_to_roms_time(ds, model_reference_date, climatology, time_name="time"):
+
+    if climatology:
+        ds.attrs["climatology"] = str(True)
+        month = xr.DataArray(range(1, 13), dims=time_name)
+        month.attrs["long_name"] = "Month index (1-12)"
+        ds = ds.assign_coords({"month": month})
+        # Preserve absolute time coordinate for readability
+        abs_time = np.datetime64(model_reference_date) + ds[time_name]
+        # Convert to pandas TimedeltaIndex
+        timedelta_index = pd.to_timedelta(ds[time_name].values)
+
+        # Determine the start of the year for the base_datetime
+        start_of_year = datetime(model_reference_date.year, 1, 1)
+
+        # Calculate the offset from midnight of the new year
+        offset = model_reference_date - start_of_year
+
+        # Convert the timedelta to nanoseconds first, then to days
+        time = xr.DataArray(
+            (timedelta_index - offset).view("int64") / 3600 / 24 * 1e-9,
+            dims=time_name,
+        )
+        time.attrs["cycle_length"] = 365.25
+
+    else:
+        # Preserve absolute time coordinate for readability
+        abs_time = ds[time_name]
+
+        time = (
+            (ds[time_name] - np.datetime64(model_reference_date)).astype("float64")
+            / 3600
+            / 24
+            * 1e-9
+        )
+
+    attrs = [key for key in abs_time.attrs]
+    for attr in attrs:
+        del abs_time.attrs[attr]
+    abs_time.attrs["long_name"] = "absolute time"
+    ds = ds.assign_coords({"abs_time": abs_time})
+
+    time.attrs["long_name"] = f"relative time: days since {str(model_reference_date)}"
+    time.encoding["units"] = "days"
+    time.attrs["units"] = "days"
+    ds.encoding["unlimited_dims"] = "time"
+
+    return ds, time
+
+
+def _to_yaml(forcing_object, filepath: Union[str, Path]) -> None:
+    """Serialize a forcing object (including its grid) into a YAML file.
+
+    This function serializes a dataclass object (forcing_object) and its associated
+    `grid` attribute into a YAML file. It includes additional metadata, such as
+    the version of the `roms-tools` package, and omits fields like `grid` and `ds`
+    that are not serializable or meant to be excluded.
+
+    The function also converts datetime fields to ISO format strings for proper
+    serialization.
+
+    Parameters
+    ----------
+    forcing_object : object
+        The object that contains the forcing data, typically a dataclass with attributes
+        such as `grid`, `start_time`, `end_time`, etc.
+    filepath : Union[str, Path]
+        The path where the serialized YAML file will be saved.
+
+    Returns
+    -------
+    None
+        The function writes the serialized data directly to a YAML file at the specified path.
+    """
+
+    # Convert the filepath to a Path object
+    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}
+
+    # Step 2: Get ROMS Tools version
+    # Fetch the version of the 'roms-tools' package for inclusion in the YAML header
+    try:
+        roms_tools_version = importlib.metadata.version("roms-tools")
+    except importlib.metadata.PackageNotFoundError:
+        roms_tools_version = "unknown"
+
+    # Create YAML header with version information
+    header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
+
+    # Step 3: Prepare Forcing Data
+    # Prepare the forcing object fields, excluding 'grid' and 'ds'
+    forcing_data = {}
+    field_names = [field.name for field in fields(forcing_object)]
+    filtered_field_names = [
+        param
+        for param in field_names
+        if param not in ("grid", "ds", "use_dask", "climatology")
+    ]
+
+    for field_name in filtered_field_names:
+        # Retrieve the value of each field using getattr
+        value = getattr(forcing_object, field_name)
+
+        # If the field is a datetime object, convert it to ISO format
+        if isinstance(value, datetime):
+            value = value.isoformat()
+
+        # Add the field and its value to the forcing_data dictionary
+        forcing_data[field_name] = value
+
+    # Step 4: Combine Grid and Forcing Data
+    # Combine grid and forcing data into a single dictionary for the final YAML content
+    yaml_data = {
+        **grid_yaml_data,  # Add the grid data to the final YAML structure
+        forcing_object.__class__.__name__: forcing_data,  # Include the serialized forcing object data
+    }
+
+    # Step 5: Write to YAML file
+    with filepath.open("w") as file:
+        # Write the header first
+        file.write(header)
+        # Write the serialized YAML data
+        yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
+
+
+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.
+
+    This function reads a YAML file, searches for the configuration data associated
+    with the class name of the forcing object, and returns the configuration data
+    as a dictionary. The dictionary contains the forcing parameters extracted from
+    the YAML file, with any date fields converted from ISO format.
+
+    Parameters
+    ----------
+    filepath : Union[str, Path]
+        The path to the YAML file from which the parameters will be read.
+    forcing_object : Type
+        The class type (e.g., TidalForcing) whose configuration data is to be loaded
+        from the YAML file. The class name is used to locate the relevant data in
+        the YAML structure.
+
+    Returns
+    -------
+    dict
+        A dictionary containing the forcing parameters extracted from the YAML file.
+        This dictionary contains key-value pairs where the keys are the parameter
+        names, and the values are the corresponding values from the YAML file.
+        Any date fields are converted from ISO format if necessary.
+
+    Raises
+    ------
+    ValueError
+        If no configuration for the specified class name is found in the YAML file.
+    """
+
+    # Read the entire file content
+    with filepath.open("r") as file:
+        file_content = file.read()
+
+    # Split the content into YAML documents
+    documents = list(yaml.safe_load_all(file_content))
+
+    forcing_data = None
+    forcing_object_name = forcing_object.__name__
+
+    # Process the YAML documents to find the forcing data for the given object
+    for doc in documents:
+        if doc is None:
+            continue
+        if forcing_object_name in doc:
+            forcing_data = doc[forcing_object_name]
+            break
+
+    if forcing_data is None:
+        raise ValueError(
+            f"No {forcing_object_name} configuration found in the YAML file."
+        )
+
+    # Convert any date fields from ISO format if necessary
+    for key, value in forcing_data.items():
+        forcing_data[key] = _convert_from_iso_format(value)
+
+    # Return the forcing data as a dictionary
+    return forcing_data
+
+
+def _convert_from_iso_format(value):
+    try:
+        # Return the parsed datetime object if successful
+        return datetime.fromisoformat(str(value))
+    except ValueError:
+        # Return None or raise an exception if parsing fails
+        return value

roms_tools/tests/test_setup/test_boundary_forcing.py

@@ -1,17 +1,19 @@
 import pytest
 from datetime import datetime
-from roms_tools import BoundaryForcing
+import xarray as xr
+from roms_tools import Grid, BoundaryForcing
 import textwrap
 from roms_tools.setup.download import download_test_data
 from conftest import calculate_file_hash
 from pathlib import Path
+import logging
 
 
 @pytest.mark.parametrize(
     "boundary_forcing_fixture",
     [
         "boundary_forcing",
-        "boundary_forcing_with_2d_fill",
+        # "boundary_forcing_with_2d_fill",
     ],
 )
 def test_boundary_forcing_creation(boundary_forcing_fixture, request):
@@ -51,7 +53,7 @@ def test_boundary_forcing_creation(boundary_forcing_fixture, request):
     "boundary_forcing_fixture",
     [
         "bgc_boundary_forcing_from_climatology",
-        "bgc_boundary_forcing_from_climatology_with_2d_fill",
+        # "bgc_boundary_forcing_from_climatology_with_2d_fill",
     ],
 )
 def test_boundary_forcing_creation_with_bgc(boundary_forcing_fixture, request):
@@ -85,6 +87,99 @@ def test_boundary_forcing_creation_with_bgc(boundary_forcing_fixture, request):
     assert hasattr(boundary_forcing.ds, "climatology")
 
 
+def test_unsuccessful_boundary_forcing_creation_with_1d_fill(use_dask):
+
+    grid = Grid(
+        nx=2,
+        ny=2,
+        size_x=500,
+        size_y=1000,
+        center_lon=0,
+        center_lat=55,
+        rot=10,
+        N=3,  # number of vertical levels
+        theta_s=5.0,  # surface control parameter
+        theta_b=2.0,  # bottom control parameter
+        hc=250.0,  # critical depth
+    )
+
+    fname = download_test_data("GLORYS_coarse_test_data.nc")
+
+    with pytest.raises(ValueError, match="consists entirely of NaNs"):
+
+        BoundaryForcing(
+            grid=grid,
+            start_time=datetime(2021, 6, 29),
+            end_time=datetime(2021, 6, 30),
+            source={"name": "GLORYS", "path": fname},
+            apply_2d_horizontal_fill=False,
+            use_dask=use_dask,
+        )
+
+    fname_bgc = download_test_data("CESM_regional_coarse_test_data_climatology.nc")
+
+    with pytest.raises(ValueError, match="consists entirely of NaNs"):
+
+        BoundaryForcing(
+            grid=grid,
+            start_time=datetime(2021, 6, 29),
+            end_time=datetime(2021, 6, 30),
+            source={"path": fname_bgc, "name": "CESM_REGRIDDED", "climatology": True},
+            type="bgc",
+            apply_2d_horizontal_fill=False,
+            use_dask=use_dask,
+        )
+
+
+def test_boundary_divided_by_land_warning(caplog, use_dask):
+
+    # Iceland intersects the western boundary of the following grid
+    grid = Grid(
+        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")
+
+    with caplog.at_level(logging.WARNING):
+        BoundaryForcing(
+            grid=grid,
+            start_time=datetime(2021, 6, 29),
+            end_time=datetime(2021, 6, 30),
+            source={"path": fname, "name": "GLORYS", "climatology": False},
+            apply_2d_horizontal_fill=False,
+            use_dask=use_dask,
+        )
+    # Verify the warning message in the log
+    assert "the western boundary is divided by land" in caplog.text
+
+
+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")
+
+    kwargs = {
+        "grid": grid,
+        "start_time": datetime(2021, 6, 29),
+        "end_time": datetime(2021, 6, 29),
+        "source": {"path": fname, "name": "GLORYS", "climatology": False},
+        "use_dask": use_dask,
+    }
+
+    bf_1d_fill = BoundaryForcing(
+        **kwargs,
+        apply_2d_horizontal_fill=False,
+    )
+    bf_2d_fill = BoundaryForcing(
+        **kwargs,
+        apply_2d_horizontal_fill=True,
+    )
+
+    xr.testing.assert_allclose(bf_1d_fill.ds, bf_2d_fill.ds, rtol=1.0e-4)
+
+
 def test_boundary_forcing_plot(boundary_forcing):
     """Test plot."""
 
@@ -296,9 +391,7 @@ def test_from_yaml_missing_boundary_forcing(tmp_path, request, use_dask):
       center_lat: 61
       rot: -20
       topography_source: ETOPO5
-      smooth_factor: 8
       hmin: 5.0
-      rmax: 0.2
     """
     )
     # Create a temporary filepath using the tmp_path fixture

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

@@ -3,7 +3,7 @@
     "end_time": "2021-06-30 00:00:00",
     "hc": 250.0,
     "model_reference_date": "2000-01-01 00:00:00",
-    "roms_tools_version": "0.1.dev157+dirty",
+    "roms_tools_version": "0.1.dev169+dirty",
     "source": "CESM_REGRIDDED",
     "start_time": "2021-06-29 00:00:00",
     "theta_b": 2.0,