parcels/tests/test_convert.py at 435fd7aeeeb92582ef2ba95e20114439e5cccb74 · VeckoTheGecko/parcels · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import pytest
import xarray as xr

import parcels
import parcels.convert as convert
from parcels import FieldSet
from parcels._core.utils import sgrid
from parcels._datasets.structured.circulation_models import datasets as datasets_circulation_models
from parcels.interpolators._xinterpolators import _get_offsets_dictionary


def test_nemo_to_sgrid():
    data_folder = parcels.download_example_dataset("NemoCurvilinear_data")
    U = xr.open_mfdataset(data_folder.glob("*U.nc4"))
    V = xr.open_mfdataset(data_folder.glob("*V.nc4"))
    coords = xr.open_dataset(data_folder / "mesh_mask.nc4")

    ds = convert.nemo_to_sgrid(fields=dict(U=U, V=V), coords=coords)

    assert ds["grid"].attrs == {
        "cf_role": "grid_topology",
        "topology_dimension": 2,
        "node_dimensions": "x y",
        "face_dimensions": "x_center:x (padding:low) y_center:y (padding:low)",
        "node_coordinates": "lon lat",
        "vertical_dimensions": "depth_center:depth (padding:high)",
    }

    meta = sgrid.parse_grid_attrs(ds["grid"].attrs)

    # Assuming that node_dimension1 and node_dimension2 correspond to X and Y respectively
    # check that U and V are properly defined on the staggered grid
    assert {
        meta.get_value_by_id("node_dimension1"),  # X edge
        meta.get_value_by_id("face_dimension2"),  # Y center
    }.issubset(set(ds["U"].dims))
    assert {
        meta.get_value_by_id("face_dimension1"),  # X center
        meta.get_value_by_id("node_dimension2"),  # Y edge
    }.issubset(set(ds["V"].dims))


def test_convert_nemo_offsets():
    data_folder = parcels.download_example_dataset("NemoCurvilinear_data")
    U = xr.open_mfdataset(data_folder.glob("*U.nc4"))
    V = xr.open_mfdataset(data_folder.glob("*V.nc4"))
    coords = xr.open_dataset(data_folder / "mesh_mask.nc4")

    ds = convert.nemo_to_sgrid(fields=dict(U=U, V=V), coords=coords)
    fieldset = FieldSet.from_sgrid_conventions(ds)

    offsets = _get_offsets_dictionary(fieldset.UV.grid)
    assert offsets["X"] == 1
    assert offsets["Y"] == 1
    assert offsets["Z"] == 0


def test_convert_mitgcm_offsets():
    data_folder = parcels.download_example_dataset("MITgcm_example_data")
    ds_fields = xr.open_dataset(data_folder / "mitgcm_UV_surface_zonally_reentrant.nc")
    coords = ds_fields[["XG", "YG", "Zl", "time"]]
    ds_fset = convert.mitgcm_to_sgrid(fields={"U": ds_fields.UVEL, "V": ds_fields.VVEL}, coords=coords)
    fieldset = FieldSet.from_sgrid_conventions(ds_fset)
    offsets = _get_offsets_dictionary(fieldset.UV.grid)
    assert offsets["X"] == 0
    assert offsets["Y"] == 0
    assert offsets["Z"] == 0


def test_convert_croco_offsets():
    ds = datasets_circulation_models["ds_CROCO_idealized"]
    coords = ds[["x_rho", "y_rho", "s_w", "time"]]

    ds = convert.croco_to_sgrid(fields={"U": ds["u"], "V": ds["v"]}, coords=coords)
    fieldset = FieldSet.from_sgrid_conventions(ds)

    offsets = _get_offsets_dictionary(fieldset.UV.grid)
    assert offsets["X"] == 0
    assert offsets["Y"] == 0
    assert offsets["Z"] == 0


_COPERNICUS_DATASETS = [
    datasets_circulation_models["ds_copernicusmarine"],
    datasets_circulation_models["ds_copernicusmarine_waves"],
]


@pytest.mark.parametrize("ds", _COPERNICUS_DATASETS)
def test_convert_copernicusmarine(ds, caplog):
    if "uo" in ds:
        fields = {"U": ds["uo"], "V": ds["vo"]}
    elif "VSDX" in ds:
        fields = {"U": ds["VSDX"], "V": ds["VSDY"]}
    else:
        raise ValueError("Test dataset does not contain recognized current variables.")
    ds_fset = convert.copernicusmarine_to_sgrid(fields=fields)
    fieldset = FieldSet.from_sgrid_conventions(ds_fset)
    assert "U" in fieldset.fields
    assert "V" in fieldset.fields
    assert "UV" in fieldset.fields


def test_convert_copernicusmarine_no_currents(caplog):
    ds = datasets_circulation_models["ds_copernicusmarine"]
    ds_fset = convert.copernicusmarine_to_sgrid(fields={"do": ds["uo"]})
    fieldset = FieldSet.from_sgrid_conventions(ds_fset)
    assert "U" not in fieldset.fields
    assert "V" not in fieldset.fields
    assert "UV" not in fieldset.fields
    assert caplog.text == ""


@pytest.mark.parametrize("ds", _COPERNICUS_DATASETS)
def test_convert_copernicusmarine_no_logs(ds, caplog):
    ds = ds.copy()
    zeros = xr.zeros_like(list(ds.data_vars.values())[0])
    ds["U"] = zeros
    ds["V"] = zeros

    ds_fset = convert.copernicusmarine_to_sgrid(fields={"U": ds["U"], "V": ds["V"]})
    fieldset = FieldSet.from_sgrid_conventions(ds_fset)
    assert "U" in fieldset.fields
    assert "V" in fieldset.fields
    assert "UV" in fieldset.fields
    assert caplog.text == ""