Uh oh!
There was an error while loading. Please reload this page.
XGrid localization - #2082
Conversation
Below is a small testing script along with output. # %%importnumpyasnpimportxarrayasxrfrompprintimportpprintfromparcelsimportxgcmfromparcels._datasets.structured.genericimportX, Y, Zfromparcels.xgridimportXGridT=2Z=Y=X=3TIME=xr.date_range("2000", "2001", T)
ds_mitgcm=xr.Dataset(
{
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
"U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
"V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
},
coords={
"XG": (
["XG"],
np.arange(0, X),
{"axis": "X", "c_grid_axis_shift": -0.5},
),
"XC": (["XC"], np.arange(0, X) +0.5, {"axis": "X"}),
"YG": (
["YG"],
np.arange(0, Y),
{"axis": "Y", "c_grid_axis_shift": -0.5},
),
"YC": (
["YC"],
np.arange(0, Y) +0.5,
{"axis": "Y"},
),
"ZG": (
["ZG"],
np.arange(Z),
{"axis": "Z", "c_grid_axis_shift": -0.5},
),
"ZC": (
["ZC"],
np.arange(Z) +0.5,
{"axis": "Z"},
),
"lon": (["XG"], np.arange(0, X)),
"lat": (["YG"], np.arange(0, Y)),
"depth": (["ZG"], np.arange(Z)),
"time": (["time"], TIME, {"axis": "T"}),
},
)
ds_nemo=xr.Dataset(
{
"data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
"U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
"U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
"V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
},
coords={
"XG": (
["XG"],
np.arange(0, X),
{"axis": "X", "c_grid_axis_shift": 0.5},
),
"XC": (["XC"], np.arange(0, X) -0.5, {"axis": "X"}),
"YG": (
["YG"],
np.arange(0, Y),
{"axis": "Y", "c_grid_axis_shift": 0.5},
),
"YC": (
["YC"],
np.arange(0, Y) -0.5,
{"axis": "Y"},
),
"ZG": (
["ZG"],
np.arange(Z),
{"axis": "Z", "c_grid_axis_shift": -0.5},
),
"ZC": (
["ZC"],
np.arange(Z) +0.5,
{"axis": "Z"},
),
"lon": (["XG"], np.arange(0, X)),
"lat": (["YG"], np.arange(0, Y)),
"depth": (["ZG"], np.arange(Z)),
"time": (["time"], TIME, {"axis": "T"}),
},
)
grid_mitgcm=XGrid(xgcm.Grid(ds_mitgcm, periodic=False))
grid_nemo=XGrid(xgcm.Grid(ds_nemo, periodic=False))
print("XGCM repr of MITgcm grid:")
print(grid_mitgcm.xgcm_grid)
print("\nXGCM repr of NEMO grid:")
print(grid_nemo.xgcm_grid)
# %%defshow_point_on_grid(grid, z, y, x):
"""Pretty printing of some info"""position=grid.search(z, y, x)
print(f"Position wrt. fpoints (lon/lat grid):")
pprint(position)
fordaingrid.xgcm_grid._ds.data_vars.values():
local_position=grid.localize(position, da.dims)
print(f"On {da.name=} with {da.dims=}, local position:")
pprint(local_position)
print("----Working with MITgcm grid----")
show_point_on_grid(grid_mitgcm, 0, 0.8, 0.8)
print("\n----Working with NEMO grid----")
show_point_on_grid(grid_nemo, 0, 0.8, 0.8)output: |
erikvansebille
commented
Jul 11, 2025
Note that some datasets have an attribute |
VeckoTheGecko
commented
Jul 11, 2025
Yes, this is already taken into account by xgcm during grid ingestion to determine the nature of the grid staggering. In fact, the only difference between |
if we're comfortable with this approach, I'll go and clear up some stuff (e.g., variable naming since it might be a bit confusing) |
f3949e4 to
0d0b0e5Compare
erikvansebille
left a comment
There was a problem hiding this comment.
Looks good. And how would a user use this in an interpolation method? Or will that come in a next PR?
Uh oh!
There was an error while loading. Please reload this page.
0d0b0e5 to
e2e549cCompareI realised that this localization step isn't actually needed.  If a point is defined at (in X) Still we need to build tooling around getting the dimension names that correspond to the axes (i.e, for |
VeckoTheGecko
commented
Jul 16, 2025
I realised that this is wrong actually 😅 . Localisation will be needed when writing curvilinear interpolators such as in the following being investigated in #2081 : defXTriCurviLinear(
field: Field,
ti: int,
position: dict[_XGRID_AXES, tuple[int, float|np.ndarray]],
tau: np.float32|np.float64,
t: np.float32|np.float64,
z: np.float32|np.float64,
y: np.float32|np.float64,
x: np.float32|np.float64,
):
"""Trilinear interpolation on a curvilinear grid."""xi, xsi=position["X"]
yi, eta=position["Y"]
zi, zeta=position["Z"]
data=field.dataaxis_dim=field.grid.get_axis_dim_mapping(field.data.dims)
return (
(
(1-xsi) * (1-eta) *data.isel({axis_dim["Y"]: yi, axis_dim["X"]: xi})
+xsi* (1-eta) *data.isel({axis_dim["Y"]: yi, axis_dim["X"]: xi+1})
+xsi*eta*data.isel({axis_dim["Y"]: yi+1, axis_dim["X"]: xi+1})
+ (1-xsi) *eta*data.isel({axis_dim["Y"]: yi+1, axis_dim["X"]: xi})
)
.interp(time=t, **{axis_dim["Z"]: zi+zeta})
.values
)Here, Perhaps there are other ways to achieve this without a localization step - but I need to get more familiar with Xarray internals and coordinate aware interpolation for that. I'll mark this as unstable API subject to change in the docstring - this is an isolated change we can remove later. @erikvansebille any additional thoughts? |
erikvansebille
commented
Jul 16, 2025
Yep agree; let's keep it in for now |
|
Currently we have a search method which returns a particle position relative to the F points. However, since our data can be defined on a staggered grid, it's important to "localize" this particle position to the grid for the array of interest. This mainly applies for C grids and when working with MITgcm and NEMO where their F points and C points are defined differently relative to each other (see diagram in docs or in #2037) .
This PR introduces this grid localization. This is really just the first draft to get feedback (code is a bit more messy than I would like - and there are no tests). This good localization will help with writing interpolators.
v4-devfor v4 changes)