Regridding OBC Velocity & Tracers Test 1 (#276)

* First Try

* New test

* Rm tides test

* This

---------

Co-authored-by: Ashley Barnes <53282288+ashjbarnes@users.noreply.github.com>

Author: manishvenu

tests/conftest.py

@@ -3,6 +3,7 @@
 import xarray as xr
 import numpy as np
 import regional_mom6 as rmom6
+from regional_mom6 import experiment
 
 # Define the path where the curvilinear hgrid file is expected in the Docker container
 DOCKER_FILE_PATH = "/data/small_curvilinear_hgrid.nc"
@@ -12,6 +13,56 @@
 LOCAL_FILE_PATH = str(os.getenv("local_curvilinear_hgrid"))
 
 
+@pytest.fixture
+def toy_glorys_ds():
+    lat = np.linspace(2, 4, 2)
+    lon = np.linspace(2, 4, 2)
+    depth = np.linspace(0, 1, 2)
+    time = np.arange(1, dtype=float)
+    s4 = (len(time), len(depth), len(lat), len(lon))
+    s3 = (len(time), len(lat), len(lon))
+    c4 = {
+        "time": time,
+        "depth": depth,
+        "lat": lat,
+        "lon": lon,
+    }
+    c3 = {"time": time, "lat": lat, "lon": lon}
+    d4 = ["time", "depth", "lat", "lon"]
+    d3 = ["time", "lat", "lon"]
+    # Define 2D spatial pattern
+    temp_2d = np.array([[20.0, 22.0], [24.0, 26.0]])  # shape (lat, lon)
+    temp_4d = np.broadcast_to(temp_2d[np.newaxis, np.newaxis, :, :], s4).copy()
+    ds = xr.Dataset(
+        {
+            "temp": xr.DataArray(temp_4d, dims=d4, coords=c4),
+            "salt": xr.DataArray(np.full(s4, 35.0), dims=d4, coords=c4),
+            "u": xr.DataArray(np.zeros(s4), dims=d4, coords=c4),
+            "v": xr.DataArray(np.zeros(s4), dims=d4, coords=c4),
+            "eta": xr.DataArray(np.zeros(s3), dims=d3, coords=c3),
+        }
+    )
+    ds.time.attrs = {"units": "days"}
+    return ds
+
+
+@pytest.fixture
+def simple_experiment(tmp_path):
+    return experiment(
+        longitude_extent=[-5, 5],
+        latitude_extent=[0, 10],
+        date_range=["2003-01-01 00:00:00", "2003-01-01 00:00:00"],
+        resolution=0.1,
+        number_vertical_layers=5,
+        layer_thickness_ratio=1,
+        depth=1000,
+        mom_run_dir=tmp_path / "rundir",
+        mom_input_dir=tmp_path / "inputdir",
+        fre_tools_dir="toolpath",
+        hgrid_type="even_spacing",
+    )
+
+
 @pytest.fixture
 def get_curvilinear_hgrid():
     # Check if the file exists in the Docker-specific location

tests/test_regridding.py

@@ -4,6 +4,9 @@
 import pytest
 import xarray as xr
 import numpy as np
+import pandas as pd
+from regional_mom6.regional_mom6 import segment
+import shutil
 
 
 # Not testing get_arakawa_c_points, coords, & create_regridder
@@ -264,3 +267,101 @@ def test_mask_dataset(get_curvilinear_hgrid):
     ).all() == True and (
         (ds["temp"][0, 0, 0][(((end_ind * 2) + 1) + 1 + 1) : -1] == fill_value)
     ).all() == True  # Ensure data is not in land area
+
+
+def test_regrid_velocity_tracers(toy_glorys_ds, tmp_path):
+    """
+    Correctness test for segment.regrid_velocity_tracers.
+
+    Checks:
+    - Output OBC file is written
+    - Variables follow the {var}_{segment_name} naming convention
+    - Temperature is in Celsius (< 100)
+    - 3D fields have companion dz_* variables
+    - Vertical coordinate is re-encoded as incremental integers
+    - Perpendicular dimension has size 1
+    """
+    lat = np.linspace(2, 4, 2)
+    lon = np.linspace(2, 4, 2)
+    hgrid = rmom6.generate_rectangular_hgrid(lat, lon)
+
+    seg_name = "segment_001"
+    outfolder = tmp_path / "inputdir"
+    outfolder.mkdir()
+
+    # Minimal synthetic boundary dataset covering the east edge of the hgrid (lon ≈ 10, lat 0-10)
+
+    infile = tmp_path / "east_raw.nc"
+    ds = toy_glorys_ds
+    ds.to_netcdf(infile)
+    ds.close()
+
+    varnames = {
+        "xh": "lon",
+        "yh": "lat",
+        "time": "time",
+        "eta": "eta",
+        "zl": "depth",
+        "u": "u",
+        "v": "v",
+        "tracers": {"temp": "temp", "salt": "salt"},
+    }
+
+    seg = segment(
+        hgrid=hgrid,
+        outfolder=outfolder,
+        segment_name=seg_name,
+        orientation="east",
+        startdate="2003-01-01 00:00:00",
+    )
+    segment_out, _ = seg.regrid_velocity_tracers(infile, varnames, arakawa_grid="A")
+
+    # Salt is spatially constant, so all ocean points must match exactly.
+    salt_vals = segment_out[f"salt_{seg_name}"].values
+    np.testing.assert_allclose(salt_vals, 35.0, rtol=1e-4)
+
+    # Temp is spatially varying (20–26 °C), so just check values are exact (hgrid overlap with seg exactly)
+    temp_vals = segment_out[f"temp_{seg_name}"].values
+    assert temp_vals[0, 0, 0, 0] == 22
+    assert temp_vals[0, 0, 1, 0] == 26
+
+    seg_north = segment(
+        hgrid=hgrid,
+        outfolder=outfolder,
+        segment_name=seg_name,
+        orientation="north",
+        startdate="2003-01-01 00:00:00",
+    )
+    segment_out_north, _ = seg_north.regrid_velocity_tracers(
+        infile, varnames, arakawa_grid="A"
+    )
+    temp_vals = segment_out_north[f"temp_{seg_name}"].values
+    assert temp_vals[0, 0, 0, 0] == 24
+    assert temp_vals[0, 0, 0, 1] == 26
+
+    # Mess with hgrid, subtract 1
+
+    folder = outfolder / "weights"
+    shutil.rmtree(
+        folder
+    )  # removes weights so they aren't saved with the old hgrid, and forces them to be recomputed with the new hgrid
+    folder.mkdir()  # recreate the empty folder
+    hgrid["x"] = hgrid.x + 1
+    hgrid["y"] = hgrid.y + 1
+    seg_regrid = segment(
+        hgrid=hgrid,
+        outfolder=outfolder,
+        segment_name=seg_name,
+        orientation="west",
+        startdate="2003-01-01 00:00:00",
+    )
+    seg_regridded, _ = seg_regrid.regrid_velocity_tracers(
+        infile, varnames, arakawa_grid="A", regridding_method="bilinear"
+    )
+    temp_vals = seg_regridded[f"temp_{seg_name}"].values
+    assert (
+        np.abs(temp_vals[0, 0, 0, 0] - 23) < 0.01
+    )  # bilinear at this point is nearly the average of the toy_glory_ds values (22 and 24 and 26 and 20)
+    assert (
+        temp_vals[0, 0, 1, 0] == 0
+    )  # The bilinear regridding would be zero here because there isn't 4 points