In oceanography it is often required to transform the vertical coordinate system of data. The large scale circulation for example, follows surfaces of constant density rather than constant depth levels, thus certain analyses require a coordinate transformation.
The general process of vertical coordinates transformation as discussed here, consists of two steps: Regridding and remapping.
Regridding: The process of definining the target depth coordinates (e.g. the depth defined by certain density levels or e.g. just a different depth spacing)
Remapping: Converting a data profile from the source depth profile to the target depth profile.
There are several methods to perform both regridding and remapping implemented in xarrayutils.vertical_coordinates. In the following we will see how to implement different combinations of method for both processes. For general applications the conservative conservative_remap function is recommended for remapping, due to the property of conserving total tracer mass (see examples below).
At the moment these functions require very explicit input (dimension names etc), which will be simplified with a wrapper function in the future.
In [1]:
%load_ext autoreload
%autoreload 2
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
We are going to investigate the coordinate transformations using data from an indealized high resolution model run produced by Dhruv Balwada(to read more about the data see his paper).
In [2]:
import fsspec
ds = xr.open_zarr(fsspec.get_mapper('gcs://pangeo-data/balwada/channel_ridge_resolutions/20km/tracer_10day_snap'), consolidated=True)
ds = ds.isel(time=slice(0, 20))
# ds.to_zarr('offline_backup.zarr')
# ds = xr.open_zarr('offline_backup.zarr')
All the following steps crucially assume that the depth values are increasing with the depth dimension. The values can be negative but then have progress towards less negative values as you follow the logical depth index. This dataset has negative depth values that decrease. The simplest fix is to flip the sign of all depth values.
In [3]:
for dim in ['Z', 'Zp1', 'Zl', 'Zu']:
ds.coords[dim] = -ds[dim]
ds.Z
Out[3]:
Another issue with this dataset is that in the surface, the profile is sometimes unstable (density does not strictly increase with depth). This will cause issues when interpolating, and thus for now we use a time mean of the dataset between ~200-1500m depth.
In [4]:
# check monotonicity for all profiles
ds = ds.mean('time')
ds = ds.isel(YC=slice(10,-10), Z=slice(28,60), Zp1=slice(28,61))
# since we operate on the depth dimension we convert it to a single chunk to avoid doing this all the time later
ds = ds.chunk({'Z':-1})
assert (ds.T.diff('Z') < 0).all(['Z']).all()
In this case we will transform the dataset to different depth coordinates (we leave out the regridding step and manually provide a new depth grid). This example might be useful when you want to convert different depth grids (e.g. from several observational products) into a uniform grid.
For the simplest case (using linear interpolation for both the regridding and remapping), xarray offers all the necessary tools built in.
In [5]:
# define a new depth array
z_new = np.arange(10,2000, 20)
ds_z_new = ds.interp(Z=z_new)
ds_z_new.Z
Out[5]:
The depth coordinate Z is now regularly spaced instead of the surface refined resolution of the original dataset, we 'remapped' all data onto a new vertical grid using linear interpolation. Lets compare the actual data.
In [6]:
ds.PTRACER01.isel(XC=40, YC=40).plot(y='Z', yincrease=False)
ds_z_new.PTRACER01.isel(XC=40, YC=40).plot(ls='--', y='Z', yincrease=False)
Out[6]:
Visually that looks pretty good and these results might be sufficient for certain applications. The biggest downside of this approach is that the total amount of tracer is not conserved.
In [7]:
dz_original = ds.drF #vertical cell thickness of the model grid
tracer_intz_original = (ds.PTRACER01 * dz_original).sum('Z')
dz_new = 20 #This is easy to infer since the grid is uniformly spaced
tracer_intz_new = (ds_z_new.PTRACER01 * dz_new).sum('Z')
print(tracer_intz_original.isel(XC=40, YC=40).load())
print(tracer_intz_new.isel(XC=40, YC=40).load())
The difference might seem small but for certain applications (e.g. budget reconstruction), this is not acceptable.
However, we can do better by using the conservative_remap function:
For this method we need not only the cell centers, but instead need to provide the depths of the vertical bounding surfaces.
In [8]:
from xarrayutils.vertical_coordinates import conservative_remap
# the conservative remapping needs information about the upper and lower bounds of the source and target cells.
bounds_original = ds.Zp1 # depth position of vertical bounding surface
bounds_new = xr.DataArray(np.arange(0,2020, 20), dims=['new_bounds'])
# now we can remap the tracer data again (at the moment the depth dimensions have to be explicitly defined).
ds_z_cons_new = conservative_remap(ds.PTRACER01,bounds_original, bounds_new,
z_dim='Z', z_bnd_dim='Zp1', z_bnd_dim_target='new_bounds', mask=True) # the associated depth dimensions for each array
# replace the new depth dimension values with the appropriate depth
ds_z_cons_new.coords['remapped'] = xr.DataArray(z_new, coords=[('remapped', z_new)])
In [9]:
ds_z_cons_new
Out[9]:
In [10]:
ds.PTRACER01.isel(XC=40, YC=40).plot(y='Z', yincrease=False)
ds_z_cons_new.isel(XC=40, YC=40).plot(ls='--', y='remapped', yincrease=False)
Out[10]:
conservative_remap takes into account every overlap between source and target cells. See for instance the uppermost value of ds_z_cons_new, the low value is due to the fact that there is a traceramount in the upper half of the first source cell, which is then distributed over the much larger targer cell. This ensures that the full tracer amount is conserved to floating point precision.
In [11]:
tracer_intz_cons_new = (ds_z_cons_new * dz_new).sum('remapped')
np.isclose(tracer_intz_original.isel(XC=40, YC=40).load(), tracer_intz_cons_new.isel(XC=40, YC=40).load())
Out[11]:
This is in fact true for every grid position:
In [12]:
np.isclose(tracer_intz_original, tracer_intz_cons_new).all()
Out[12]:
Ok this is nice, but it really gets interesting when we define our new depth coordinates
The xarray internals can only help us if we want to interpolate on values of a dimension (1D), e.g. if we want to know the temperature as a function of a new depth. For the case of temperature coordinates, we aim to do the opposite: Find the depth for a given temperature value. This is not possible with xarray at the moment, but can be achieved using linear_interpolation_regrid.
The values used to find the new cell depths have to be monotonic! Currently there is no check implemented for this but non-monotonic fields can lead to undesired behaviour. See above for how to check for monotonicity.
In [13]:
from xarrayutils.vertical_coordinates import linear_interpolation_regrid
In [14]:
t_vals = np.arange(0.6,3, 0.01)
temperature_values = xr.DataArray(t_vals, coords=[('t', t_vals)]) # define the new temperature grid
z_temp_coord = linear_interpolation_regrid(ds.Z, ds.T, temperature_values, target_value_dim='t')
plt.subplot(2,1,1)
ds.T.isel(XC=40).plot(x='YC', yincrease=False)
ds.T.isel(XC=40).plot.contour(levels=[2], colors='w',x='YC', yincrease=False)
plt.subplot(2,1,2)
z_temp_coord.isel(XC=40).plot(x='YC')
plt.axhline(2, color='w')
Out[14]:
As you can see in this example, the line of constant temperature is moving deeper with increasing y, and that is reflected in the depth values along a constant depth coordinate in the regridded values. Now we can remap other data values on the corresponding depths. The simplest method is again linear interpolation:
In [15]:
from xarrayutils.vertical_coordinates import linear_interpolation_remap
In [16]:
# we cant have nans, so just fill them with an out of bounds value
# z_temp_coord = z_temp_coord.fillna(1e8) # this should be fixed now?
ds_temp_linear = linear_interpolation_remap(ds.Z, ds.T, z_temp_coord)
# this requires me to downgrade dask, should be solved soon (see here https://github.com/pydata/xarray/pull/3660)
ds_temp_linear
Out[16]:
In [17]:
ds_temp_linear.isel(XC=40).plot(x='YC', robust=True)
Out[17]:
As expected, when we remap the temperature field itself, we get a matching horizontal stratification. Now lets do something more interesting and look at the tracer field on a constant temperature surface of 2 deg:
In [18]:
ds_temp_linear = linear_interpolation_remap(ds.Z, ds.PTRACER01, z_temp_coord)
ds_temp_linear.sel(remapped=2, method='nearest').plot(robust=True)
Out[18]:
In [19]:
# we need to capture all the tracer cells, so if we dont specify our temperature range covering the Tracer min and max, the total tracer amount will not be conserved.
t_vals = np.hstack([ds.T.min().load().data[np.newaxis], np.arange(0.6,3, 0.01), ds.T.max().load().data[np.newaxis]])
t_vals
# for now the results are ordered by the bin values. Since the temperature is decreasing, we need to flip the logical axis.
# We could also do that after the regridding, but this seems more elegant.
t_vals = np.flip(t_vals)
temperature_values = xr.DataArray(t_vals, coords=[('t', t_vals)]) # define the new temperature grid
temperature_values
# # additionally to covering the full range of `ds.T`, we need to provide the bounds of `z` to the function.
z_temp_bounds = linear_interpolation_regrid(ds.Z, ds.T, temperature_values, z_bounds=ds.Zp1 ,target_value_dim='t', z_bounds_dim='Zp1')
z_temp_bounds.isel(XC=40, YC=40).load()
Out[19]:
And now we can use these cell bounding values exactly like we did before.
In [20]:
# now we can remap the tracer data again (at the moment the depth dimensions have to be explicitly defined).
ds_temp_cons = conservative_remap(ds.PTRACER01,bounds_original, z_temp_bounds,
z_dim='Z', z_bnd_dim='Zp1', z_bnd_dim_target='regridded', mask=True) # the associated depth dimensions for each array
# # replace the new depth dimension values with the appropriate depth (here the middle of the temperature cell bounds)
t_vals = z_temp_bounds.coords['regridded'].data
t_vals = 0.5 * (t_vals[1:] + t_vals[0:-1])
ds_temp_cons.coords['remapped'] = xr.DataArray(t_vals, coords=[('remapped', t_vals)])
In [21]:
plt.figure(figsize=[20,5])
plt.subplot(1,3,1)
ds_temp_linear.sel(remapped=2, method='nearest').plot(robust=True)
plt.subplot(1,3,2)
ds_temp_cons.sel(remapped=2, method='nearest').plot(robust=True)
plt.subplot(1,3,3)
(ds_temp_cons.sel(remapped=2, method='nearest')-ds_temp_linear.sel(remapped=2, method='nearest')).plot(robust=True)
Out[21]:
And most importantly, the vertical tracer content is again conserved:
In [22]:
dz_remapped = z_temp_bounds.diff('regridded').rename({'regridded':'remapped'})
dz_remapped.coords['remapped'] = ds_temp_cons.coords['remapped']
tracer_intz_remapped_temp = (ds_temp_cons*dz_remapped).sum('remapped')
xr.testing.assert_allclose(tracer_intz_original, tracer_intz_remapped_temp)
In [ ]: