Grid Ufunc Examples#
This page contains examples of different Grid Ufuncs that you might find useful. They are intended to be more advanced or realistic cases than the simple differencing operations we showed in the introduction to Grid Ufuncs.
As each of these cases is a vector calculus operator, we will demonstrate their use on an example two-dimensional vector field.
Firstly we need a two-dimensional grid, and we use similar coordinate names to the 1D example we used when first introducing Simple Grids.
In [1]: import numpy as np
In [2]: import xarray as xr
In [3]: grid_ds = xr.Dataset(
...: coords={
...: "x_c": (
...: ["x_c"],
...: np.arange(1, 20),
...: {"axis": "X"},
...: ),
...: "x_g": (
...: ["x_g"],
...: np.arange(0.5, 19),
...: {"axis": "X", "c_grid_axis_shift": -0.5},
...: ),
...: "y_c": (
...: ["y_c"],
...: np.arange(1, 20),
...: {"axis": "Y"},
...: ),
...: "y_g": (
...: ["y_g"],
...: np.arange(0.5, 19),
...: {"axis": "Y", "c_grid_axis_shift": -0.5},
...: ),
...: }
...: )
...:
In [4]: grid_ds
Out[4]:
<xarray.Dataset>
Dimensions: (x_c: 19, x_g: 19, y_c: 19, y_g: 19)
Coordinates:
* x_c (x_c) int64 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
* x_g (x_g) float64 0.5 1.5 2.5 3.5 4.5 5.5 ... 14.5 15.5 16.5 17.5 18.5
* y_c (y_c) int64 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
* y_g (y_g) float64 0.5 1.5 2.5 3.5 4.5 5.5 ... 14.5 15.5 16.5 17.5 18.5
Data variables:
*empty*
In [5]: from xgcm import Grid
In [6]: grid = Grid(
...: grid_ds,
...: coords={
...: "X": {"center": "x_c", "left": "x_g"},
...: "Y": {"center": "y_c", "left": "y_g"},
...: },
...: autoparse_metadata=False,
...: )
...:
In [7]: grid
Out[7]:
<xgcm.Grid>
X Axis (periodic, boundary='periodic'):
* center x_c --> left
* left x_g --> center
Y Axis (periodic, boundary='periodic'):
* center y_c --> left
* left y_g --> center
Now we need some data.
We will create a 2D vector field, with components U
and V
.
We will treat these velocities as if they lie on a vector C-grid, so as velocities they will lie at the cell faces.
In [8]: U = np.sin(grid_ds.x_g * 2 * np.pi / 9.5) * np.cos(grid_ds.y_c * 2 * np.pi / 19)
In [9]: V = np.cos(grid_ds.x_c * 2 * np.pi / 19) * np.sin(grid_ds.y_g * 2 * np.pi / 19)
In [10]: ds = xr.Dataset({"V": V, "U": U})
In [11]: ds
Out[11]:
<xarray.Dataset>
Dimensions: (x_c: 19, y_g: 19, x_g: 19, y_c: 19)
Coordinates:
* x_c (x_c) int64 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
* y_g (y_g) float64 0.5 1.5 2.5 3.5 4.5 5.5 ... 14.5 15.5 16.5 17.5 18.5
* x_g (x_g) float64 0.5 1.5 2.5 3.5 4.5 5.5 ... 14.5 15.5 16.5 17.5 18.5
* y_c (y_c) int64 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
Data variables:
V (x_c, y_g) float64 0.1557 0.4502 0.6959 ... -0.7357 -0.4759 -0.1646
U (x_g, y_c) float64 0.3071 0.2562 0.1776 ... -0.2562 -0.3071 -0.3247
Interpolation#
It would be nice to see what our vector field looks like before we start doing calculus with it,
but the U
and V
velocities are defined on different points,
so plotting the vectors as arrows originating from a common point would technically be incorrect for our C-grid data.
We can fix this by interpolating the vectors onto co-located points.
In [12]: colocated = xr.Dataset()
In [13]: colocated["U"] = grid.interp(U, axis="X", to="center")
In [14]: colocated["V"] = grid.interp(V, axis="Y", to="center")
In [15]: colocated
Out[15]:
<xarray.Dataset>
Dimensions: (x_c: 19, y_c: 19)
Coordinates:
* x_c (x_c) int64 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
* y_c (y_c) int64 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
Data variables:
U (x_c, y_c) float64 0.5495 0.4584 0.3177 ... 1.665e-16 1.665e-16
V (x_c, y_c) float64 0.3029 0.573 0.781 ... -0.6058 -0.3203 8.327e-17
We can now show what this co-located vector field looks like
In [16]: colocated.plot.quiver("x_c", "y_c", u="U", v="V")
Out[16]: <matplotlib.quiver.Quiver at 0x7f1569c9a070>
Divergence#
Let’s first import the decorator.
In [17]: from xgcm import as_grid_ufunc
In two dimensions, the divergence operator accepts two vector components and returns one scalar result. A divergence is the sum of multiple partial derivatives, so first let’s define a derivative function like this
In [18]: def diff_forward_1d(a):
....: return a[..., 1:] - a[..., :-1]
....:
In [19]: def diff(arr, axis):
....: """First order forward difference along any axis"""
....: return np.apply_along_axis(diff_forward_1d, axis, arr)
....:
Each vector component will be differentiated along one axis, and doing so with a first order forward difference would
shift the data’s position along that axis.
Therefore our signature should look something like this "(X:left,Y:center),(X:center,Y:left)->(X:center,Y:center)"
.
We also need to pad the data to replace the elements that will be removed by the diff function, so our grid ufunc can be defined like this
In [20]: @as_grid_ufunc(
....: "(X:left,Y:center),(X:center,Y:left)->(X:center,Y:center)",
....: boundary_width={"X": (0, 1), "Y": (0, 1)},
....: )
....: def divergence(u, v):
....: u_diff_x = diff(u, axis=-2)
....: v_diff_y = diff(v, axis=-1)
....: div = u_diff_x[..., :-1] + v_diff_y[..., :-1, :]
....: return div
....:
Here we have treated the components of the (U, V)
vector as independent scalars.
Now we can compute the divergence of our example vector field
In [21]: div = divergence(grid, ds["U"], ds["V"], axis=[("X", "Y"), ("X", "Y")])
We can see the result lies on the expected coordinate positions
In [22]: div.coords
Out[22]:
Coordinates:
* x_c (x_c) int64 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
* y_c (y_c) int64 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
and the resulting divergence looks like it corresponds with the arrows of the vector field above
In [23]: import matplotlib.pyplot as plt
In [24]: div.plot(x="x_c", y="y_c")
Out[24]: <matplotlib.collections.QuadMesh at 0x7f156a016ca0>
In [25]: colocated.plot.quiver("x_c", "y_c", u="U", v="V")
Out[25]: <matplotlib.quiver.Quiver at 0x7f1569e11ac0>
In [26]: plt.gcf()
Out[26]: <Figure size 640x480 with 2 Axes>
Gradient#
The gradient is almost like the opposite of divergence in the sense that it accepts one scalar and returns multiple vectors.
Let’s first define a tracer field T
, which we imagine will start off localised near the center of the domain.
In [27]: def gaussian(x_coord, y_coord, x_pos, y_pos, A, w):
....: return A * np.exp(
....: -0.5 * ((x_coord - x_pos) ** 2 + (y_coord - y_pos) ** 2) / w**2
....: )
....:
In [28]: ds["T"] = gaussian(grid_ds.x_c, grid_ds.y_c, x_pos=7.5, y_pos=7.5, A=50, w=2)
In [29]: ds["T"].plot.contourf(x="x_c", vmax=60)
Out[29]: <matplotlib.contour.QuadContourSet at 0x7f156a001bb0>
Computing the first-order gradient will again move the data onto different grid positions, so the signature for a gradient ufunc will need to reflect this and our definition is similar to the derivative case.
In [30]: def gradient(a):
....: a_diff_x = diff(a, axis=-2)
....: a_diff_y = diff(a, axis=-1)
....: return a_diff_x[..., :-1], a_diff_y[..., :-1, :]
....:
Now we can compute the gradient of our example scalar field
In [31]: ds["grad_T_x"], ds["grad_T_y"] = grid.apply_as_grid_ufunc(
....: gradient,
....: ds["T"],
....: axis=[("X", "Y")],
....: signature="(X:center,Y:center)->(X:left,Y:center),(X:center,Y:left)",
....: boundary_width={"X": (1, 0), "Y": (1, 0)},
....: )
....:
Note
Notice we used the apply_as_grid_ufunc
syntax here instead of the as_grid_ufunc
decorator.
The result is the same.
Again in order to plot this as a vector field we should first interpolate it
In [32]: colocated["grad_T_x"] = grid.interp(ds["grad_T_x"], axis="X", to="center")
In [33]: colocated["grad_T_y"] = grid.interp(ds["grad_T_y"], axis="Y", to="center")
In [34]: colocated
Out[34]:
<xarray.Dataset>
Dimensions: (x_c: 19, y_c: 19)
Coordinates:
* x_c (x_c) int64 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
* y_c (y_c) int64 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
Data variables:
U (x_c, y_c) float64 0.5495 0.4584 0.3177 ... 1.665e-16 1.665e-16
V (x_c, y_c) float64 0.3029 0.573 0.781 ... -0.3203 8.327e-17
grad_T_x (x_c, y_c) float64 3.77e-08 0.002898 ... 1.603e-06 1.316e-07
grad_T_y (x_c, y_c) float64 3.77e-08 1.232e-07 ... -3.245e-10 1.316e-07
Now we can plot the gradient of the tracer field as a vector field
In [35]: ds["T"].plot.contourf(x="x_c", vmax=60)
Out[35]: <matplotlib.contour.QuadContourSet at 0x7f156a231400>
In [36]: colocated.plot.quiver(
....: "x_c", "y_c", u="grad_T_x", v="grad_T_y", color="0.5", scale=200
....: )
....:
Out[36]: <matplotlib.quiver.Quiver at 0x7f1569372ca0>
In [37]: plt.gcf()
Out[37]: <Figure size 640x480 with 2 Axes>
Advection#
We can also do “mixed” operations that involve both vectors and scalars, such as calculating the advective flux of a scalar tracer field due to a vector flow field.
Now we can define a simple flux operator (which internally calls our previous gradient function)
In [38]: def interp_forward_1d(a):
....: return (a[..., :-1] + a[..., 1:]) / 2.0
....:
In [39]: def interp_forward(arr, axis):
....: """First order forward interpolation along any axis"""
....: return np.apply_along_axis(interp_forward_1d, axis, arr)
....:
In [40]: @as_grid_ufunc(
....: "(X:left,Y:center),(X:center,Y:left),(X:center,Y:center)->(X:left,Y:center),(X:center,Y:left)",
....: boundary_width={"X": (1, 0), "Y": (1, 0)},
....: )
....: def flux(u, v, T):
....: """First order flux"""
....: T_at_U_position = interp_forward(T, axis=-2)
....: T_at_V_position = interp_forward(T, axis=-1)
....: T_flux_x = u[..., :-1, :-1] * T_at_U_position[..., :-1]
....: T_flux_y = v[..., :-1, :-1] * T_at_V_position[..., :-1, :]
....: return T_flux_x, T_flux_y
....:
We can use this operator in conjunction with our divergence operator in order to build an advection operator, with which we can solve the basic continuity equation
In [41]: def advect(T, U, V, delta_t):
....: """Simple solution to the continuity equation for a single timestep of length delta_t."""
....: T_flux_x, T_flux_y = flux(grid, U, V, T, axis=[("X", "Y")] * 3)
....: advected_T = T - delta_t * divergence(
....: grid, T_flux_x, T_flux_y, axis=[("X", "Y")] * 2
....: )
....: return advected_T
....:
Evaluating this function updates our tracer to what the tracer field might look like one (arbitrary-length) timestep later:
In [42]: new_T = advect(ds["T"], ds["U"], ds["V"], delta_t=3)
In [43]: new_T.plot.contourf(x="x_c", vmin=0, vmax=60)
Out[43]: <matplotlib.contour.QuadContourSet at 0x7f156a2f6fd0>
In [44]: colocated.plot.quiver("x_c", "y_c", u="U", v="V")
Out[44]: <matplotlib.quiver.Quiver at 0x7f15691accd0>
In [45]: plt.gcf()
Out[45]: <Figure size 640x480 with 2 Axes>
Vorticity#
We can compute vector fields from vector fields too, such as vorticity.
In [46]: @as_grid_ufunc(
....: "(X:left,Y:center),(X:center,Y:left)->(X:left,Y:left)",
....: boundary_width={"X": (1, 0), "Y": (1, 0)},
....: )
....: def vorticity(u, v):
....: v_diff_x = diff(v, axis=-2)
....: u_diff_y = diff(u, axis=-1)
....: return v_diff_x[..., 1:] - u_diff_y[..., 1:, :]
....:
In [47]: vort = vorticity(grid, ds["U"], ds["V"], axis=[("X", "Y"), ("X", "Y")])
In [48]: vort.plot(x="x_g", y="y_g")
Out[48]: <matplotlib.collections.QuadMesh at 0x7f1568fe9610>
In [49]: colocated.plot.quiver("x_c", "y_c", u="U", v="V")
Out[49]: <matplotlib.quiver.Quiver at 0x7f15690e3cd0>
In [50]: plt.gcf()
Out[50]: <Figure size 640x480 with 2 Axes>