bleh

Author: manishvenu

regional_mom6/validate.py

@@ -6,47 +6,56 @@
 
 from pathlib import Path
 import xarray as xr
+import numpy as np
+import re
 from .utils import setup_logger
 
 logger = setup_logger(__name__)
 
 
 def get_file(file: Path | xr.Dataset):
     """accept a filepath or xarray dataset and return the xarray dataset"""
-    if type(file) == xr.Dataset:
+    if isinstance(file, xr.Dataset):
         return file
     else:
         return xr.open_dataset(file)
 
 
 def check(condition, warning):
-    condition or logger.warn(warning)
+    if not condition:
+        logger.warning(warning)
+    return condition
 
 
 # Individual validation rule functions
 def _check_fill_value(da: xr.DataArray):
     """Check that fill values are set correctly"""
-    check("_FillValue" in da.attrs, f"{var_name} does not have a FillValue attribute")
-
-    check(
-        not np.isnan(da.attrs["_FillValue"]),
-        f"Fill Value for variable {var_name} is NaN (normally not wanted)",
+    condition = check(
+        "_FillValue" in da.attrs, f"{da.name} does not have a FillValue attribute"
     )
 
+    if condition:
+        check(
+            not np.isnan(da.attrs["_FillValue"]),
+            f"Fill Value for variable {da.name} is NaN (normally not wanted)",
+        )
+
 
 def _check_coordinates(ds: xr.Dataset, var_name: str):
-    """Check that missing values are set correctly"""
+    """Check that coordinates attribute exists and all listed coordinates are present in the dataset"""
 
     assert var_name in ds
-    check(
+    condition = check(
         "coordinates" in ds[var_name].attrs,
         f"{var_name} does not have a coordinates attribute",
     )
-
-    coordinates = ds[var_name].attrs["coordinates"]
-    coordinates = coordinates.strip(" ")
-    for coord in coordinates:
-        check(coord in ds, f"Coordinate {coord} for variable {var_name} does not exist")
+    if condition:
+        coordinates = ds[var_name].attrs["coordinates"].strip()
+        for coord in coordinates.split():
+            check(
+                coord in ds,
+                f"Coordinate {coord} for variable {var_name} does not exist",
+            )
 
 
 def _check_required_dimensions(da: xr.DataArray, surface=False):
@@ -60,12 +69,13 @@ def _check_required_dimensions(da: xr.DataArray, surface=False):
 
 
 def validate_obc_file(
-    file: Path | xr.Dataset, variable_names: list, encoding_dict={}, surface_var="eta"
+    file: Path | xr.Dataset, variable_names: list, encoding_dict=None, surface_var="eta"
 ):
     """Validate boundary condition file specifically (requires additional segment number validation)"""
+    if encoding_dict is None:
+        encoding_dict = {}
     ds = get_file(file)
 
-    # Check individual data variable specifications (nothing that starts with dz)
     print(
         "This function identifies variables by if they have the word 'segment' in the name and don't start with nz,dz,lon,lat."
     )
@@ -79,12 +89,12 @@ def validate_obc_file(
         )
         check(
             "segment" in var,
-            f"Variable {var} does not end with a 3 digit number. OBC file variables must end with a number",
+            f"Variable {var} does not contain 'segment'. OBC file variables must include 'segment'",
         )
 
         # Add encodings
         if var in encoding_dict:
-            for key, value in encoding_dict[var].item():
+            for key, value in encoding_dict[var].items():
                 ds[var].attrs[key] = value
 
         # Check if there is a non-NaN fill value
@@ -94,7 +104,7 @@ def validate_obc_file(
         _check_coordinates(ds, var_name=var)
 
         # Check the correct number of dimensions
-        _check_required_dimensions(ds[var], surface=(var == surface_var))  # just two
+        _check_required_dimensions(ds[var], surface=(var == surface_var))
 
         # Check for thickness variable
         if var != surface_var:
@@ -105,4 +115,4 @@ def validate_obc_file(
 
 
 def ends_with_3_digits(s: str) -> bool:
-    return bool(re.search(r"\d{3}$", s))
+    return bool(re.search(r"_\d{3}$", s))

tests/test_validate.py

@@ -0,0 +1,181 @@
+"""
+Test suite for the validate module
+"""
+
+import pytest
+import xarray as xr
+import numpy as np
+from pathlib import Path
+from regional_mom6.validate import (
+    _check_fill_value,
+    _check_coordinates,
+    _check_required_dimensions,
+    ends_with_3_digits,
+    validate_obc_file,
+)
+
+
+# _check_fill_value tests
+def test_check_fill_value_valid(caplog):
+    """DataArray with valid fill value logs no warnings"""
+    da = xr.DataArray(
+        [1, 2, 3], dims="x", name="temperature", attrs={"_FillValue": -999.0}
+    )
+    _check_fill_value(da)
+    assert "FillValue" not in caplog.text
+
+
+def test_check_fill_value_missing(caplog):
+    """DataArray without _FillValue attribute logs warning"""
+    da = xr.DataArray([1, 2, 3], dims="x", name="temperature")
+    _check_fill_value(da)
+    assert "FillValue" in caplog.text
+
+
+def test_check_fill_value_nan(caplog):
+    """DataArray with NaN fill value logs warning"""
+    da = xr.DataArray(
+        [1, 2, 3], dims="x", name="temperature", attrs={"_FillValue": np.nan}
+    )
+    _check_fill_value(da)
+    assert "NaN" in caplog.text
+
+
+# _check_coordinates tests
+def test_check_coordinates_valid(caplog):
+    """DataArray with valid coordinates attribute logs no warnings"""
+    ds = xr.Dataset(
+        {
+            "temperature": (["x", "y"], np.random.rand(3, 4)),
+            "lon": (["x", "y"], np.random.rand(3, 4)),
+            "lat": (["x", "y"], np.random.rand(3, 4)),
+        }
+    )
+    ds["temperature"].attrs["coordinates"] = "lon lat"
+    _check_coordinates(ds, "temperature")
+    assert "coordinate" not in caplog.text.lower()
+
+
+def test_check_coordinates_missing_attribute(caplog):
+    """DataArray without coordinates attribute logs warning"""
+    ds = xr.Dataset({"temperature": (["x", "y"], np.random.rand(3, 4))})
+    _check_coordinates(ds, "temperature")
+    assert "coordinates" in caplog.text.lower()
+
+
+def test_check_coordinates_missing_in_dataset(caplog):
+    """Missing coordinate variable logs warning"""
+    ds = xr.Dataset({"temperature": (["x", "y"], np.random.rand(3, 4))})
+    ds["temperature"].attrs["coordinates"] = "lon lat"
+    _check_coordinates(ds, "temperature")
+    assert "does not exist" in caplog.text
+
+
+# _check_required_dimensions tests
+def test_check_required_dimensions_valid_4d(caplog):
+    """4D variable passes check when surface=False"""
+    da = xr.DataArray(
+        np.random.rand(2, 3, 4, 5), dims=["time", "z", "x", "y"], name="temperature"
+    )
+    _check_required_dimensions(da, surface=False)
+    assert "dimension" not in caplog.text.lower()
+
+
+def test_check_required_dimensions_invalid_3d_for_4d(caplog):
+    """3D variable fails check when surface=False"""
+    da = xr.DataArray(np.random.rand(3, 4, 5), dims=["x", "y", "z"], name="temperature")
+    _check_required_dimensions(da, surface=False)
+    assert "dimension" in caplog.text.lower()
+
+
+def test_check_required_dimensions_valid_3d_surface(caplog):
+    """3D variable passes check when surface=True"""
+    da = xr.DataArray(np.random.rand(2, 3, 4), dims=["time", "x", "y"], name="eta")
+    _check_required_dimensions(da, surface=True)
+    assert "dimension" not in caplog.text.lower()
+
+
+def test_check_required_dimensions_invalid_4d_for_surface(caplog):
+    """4D variable fails check when surface=True"""
+    da = xr.DataArray(
+        np.random.rand(2, 3, 4, 5), dims=["time", "z", "x", "y"], name="eta"
+    )
+    _check_required_dimensions(da, surface=True)
+    assert "dimension" in caplog.text.lower()
+
+
+# ends_with_3_digits tests
+def test_ends_with_3_digits_valid_cases():
+    """String ending with 3 digits returns True"""
+    assert ends_with_3_digits("temp_001") is True
+    assert ends_with_3_digits("var_999") is True
+    assert ends_with_3_digits("_000") is True
+    assert ends_with_3_digits("temp_01") is False
+    assert ends_with_3_digits("temp_0001") is False
+    assert ends_with_3_digits("temp_abc") is False
+    assert ends_with_3_digits("temp") is False
+
+
+# validate_obc_file tests
+
+
+def test_validate_obc_file_valid(caplog):
+    """Valid OBC file with all required attributes passes"""
+    ds = xr.Dataset(
+        {
+            "temp_segment_001": (["time", "z", "x", "y"], np.random.rand(2, 3, 4, 5)),
+            "dz_temp_segment_001": (
+                ["time", "z", "x", "y"],
+                np.random.rand(2, 3, 4, 5),
+            ),
+            "eta_segment_001": (["time", "x", "y"], np.random.rand(2, 4, 5)),
+            "lon": (["x", "y"], np.random.rand(4, 5)),
+            "lat": (["x", "y"], np.random.rand(4, 5)),
+        }
+    )
+
+    for var in ds.data_vars:
+        ds[var].attrs["_FillValue"] = -999.0
+        ds[var].attrs["coordinates"] = "lon lat"
+
+    validate_obc_file(ds, ["temp_segment_001"], surface_var="eta_segment_001")
+
+
+def test_validate_obc_file_issues(caplog):
+    """OBC file with missing segment and thickness logs warnings"""
+    ds = xr.Dataset(
+        {
+            "temp_001": (["time", "z", "x", "y"], np.random.rand(2, 3, 4, 5)),
+            "lon": (["x", "y"], np.random.rand(4, 5)),
+            "lat": (["x", "y"], np.random.rand(4, 5)),
+        }
+    )
+    ds["temp_001"].attrs["_FillValue"] = -999.0
+    ds["temp_001"].attrs["coordinates"] = "lon lat"
+
+    validate_obc_file(ds, ["temp_001"])
+    assert "segment" in caplog.text
+    assert "thickness" in caplog.text or "dz_temp_001" in caplog.text
+
+
+def test_validate_obc_file_encoding_dict():
+    """Encoding dict is applied to variables"""
+    ds = xr.Dataset(
+        {
+            "temp_segment_001": (["time", "z", "x", "y"], np.random.rand(2, 3, 4, 5)),
+            "dz_temp_segment_001": (
+                ["time", "z", "x", "y"],
+                np.random.rand(2, 3, 4, 5),
+            ),
+            "lon": (["x", "y"], np.random.rand(4, 5)),
+            "lat": (["x", "y"], np.random.rand(4, 5)),
+        }
+    )
+    ds["temp_segment_001"].attrs["_FillValue"] = -999.0
+    ds["temp_segment_001"].attrs["coordinates"] = "lon lat"
+    ds["dz_temp_segment_001"].attrs["_FillValue"] = -999.0
+
+    encoding_dict = {"temp_segment_001": {"units": "celsius"}}
+    validate_obc_file(ds, ["temp_segment_001"], encoding_dict=encoding_dict)
+
+    assert ds["temp_segment_001"].attrs["units"] == "celsius"