def test_trim_and_sqeeze_by_yx_slice(y_slice, x_slice):
N = 102
ds = xr.Dataset(
coords={
"degen": (["degen"], [1]),
"t": (["t"], [np.datetime64('1970-01-01')]),
"z_c": (["z_c"], [1]),
"z_l": (["z_l"], [1]),
"y": ([
"y",
], range(N)),
"x": ([
"x",
], range(N))
})
ds_t = trim_and_squeeze(ds, y_slice=y_slice, x_slice=x_slice)
# To also cover partial overrides, we check for None's here
ds_trimmed_here = ds
if y_slice is None:
ds_trimmed_here = ds_trimmed_here.isel(y=slice(1, -1))
else:
ds_trimmed_here = ds_trimmed_here.isel(y=slice(*y_slice))
if x_slice is None:
ds_trimmed_here = ds_trimmed_here.isel(x=slice(1, -1))
else:
ds_trimmed_here = ds_trimmed_here.isel(x=slice(*x_slice))
assert "degen" not in ds_t.dims
assert "t" in ds_t.dims
assert "z_c" in ds_t.dims
assert "z_l" in ds_t.dims
assert ds_t.dims["y"] == ds_trimmed_here.dims["y"]
assert ds_t.dims["x"] == ds_trimmed_here.dims["x"]
def test_copy_coords(dims, variables, set_mm_coords, override_coords):
target_coords = [
"depth_c", "depth_l", "llat_cc", "llat_cr", "llat_rc", "llat_rr",
"llon_cc", "llon_cr", "llon_rc", "llon_rr"
]
_kwargs = {}
if override_coords:
_kwargs["update_orca_coords"] = {
"sigma_c": {
"dims": [
"z_c",
]
}
}
target_coords.append("sigma_c")
variables = variables.copy()
variables.update({"sigma_c": (
"t",
"z",
)})
mock_up_mm = _get_nan_filled_data_set(dims, variables)
mock_up_mm = trim_and_squeeze(mock_up_mm).squeeze()
if set_mm_coords:
mock_up_mm = mock_up_mm.set_coords(
[v for v in mock_up_mm.data_vars.keys()])
return_ds = create_minimal_coords_ds(mock_up_mm, **_kwargs)
return_ds = copy_coords(return_ds, mock_up_mm, **_kwargs)
assert all((tc in return_ds.coords) for tc in target_coords)
def test_trim_and_sqeeze_by_model_config(model_config_and_trimming):
N = 102
ds = xr.Dataset(
coords={
"degen": (["degen"], [1]),
"t": (["t"], [np.datetime64('1970-01-01')]),
"z_c": (["z_c"], [1]),
"z_l": (["z_l"], [1]),
"y": ([
"y",
], range(N)),
"x": ([
"x",
], range(N))
})
model_config = model_config_and_trimming["model_config"]
x_slice = model_config_and_trimming["x"]
y_slice = model_config_and_trimming["y"]
ds_t = trim_and_squeeze(ds, model_config=model_config)
ds_trimmed_here = ds.isel(x=slice(*x_slice), y=slice(*y_slice))
assert "degen" not in ds_t.dims
assert "t" in ds_t.dims
assert "z_c" in ds_t.dims
assert "z_l" in ds_t.dims
assert ds_t.dims["y"] == ds_trimmed_here.dims["y"]
assert ds_t.dims["x"] == ds_trimmed_here.dims["x"]
def test_copy_vars(dims, variables):
mock_up_mm = _get_nan_filled_data_set(dims, variables)
mock_up_mm = trim_and_squeeze(mock_up_mm).squeeze()
return_ds = create_minimal_coords_ds(mock_up_mm)
return_ds = copy_vars(return_ds, mock_up_mm)
target_vars = [
"e1t", "e2t", "e3t", "e1u", "e2u", "e3u", "e1v", "e2v", "e3v", "tmask",
"umask", "vmask", "fmask"
]
assert all((tv in return_ds.data_vars) for tv in target_vars)