<jupyter_start><jupyter_text>Transforming Vertical CoordinatesA common need in the analysis of ocean and atmospheric data is to transform the vertical coordinate from its original coordinate (e.g. depth) to a new coordinate (e.g. density).Xgcm supports this sort of one-dimensional coordinate transform on `Axis` and `Grid` objects using the `transform` method.Two algorithms are implemented:- _Linear interpolation:_ Linear interpolation is designed to interpolate intensive quantities (e.g. temperature) from one coordinate to another. This method is suitable when the target coordinate is monotonically increasing or decreasing and the data variable is intensive. For example, you want to visualize oxygen on density surfaces from a z-coordinate ocean model.- _Conservative remapping:_ This algorithm is designed to conserve extensive quantities (e.g. transport, heat content). It requires knowledge of cell bounds in both the source and target coordinate. It also handles non-monotonic target coordinates.On this page, we explain how to use these coordinate transformation capabilities.<jupyter_code>from xgcm import Grid
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt<jupyter_output><empty_output><jupyter_text>1D Toy Data ExampleFirst we will create a simple, one-dimensional dataset to illustrate how the `transform` function works.This dataset contains- a coordinate called `z`, representing the original depth coordinate- a data variable called `theta`, a function of `z`, which we want as our new vertical coordinate- a data variable called `phi`, akso a function of `z`, which represents the data we want to transform into this new coordinate spaceIn an oceanic context `theta` might be density and `phi` might be oxygen.In an atmospheric context, `theta` might be potential temperature and `phi` might be potential vorticity.<jupyter_code>z = np.arange(2, 12)
theta = xr.DataArray(np.log(z), dims=['z'], coords={'z': z})
phi = xr.DataArray(np.flip(np.log(z)*0.5+ np.random.rand(len(z))), dims=['z'], coords={'z':z})
ds = xr.Dataset({'phi': phi, 'theta': theta})
ds<jupyter_output><empty_output><jupyter_text>Let's plot this data. Note that, for a simple 1D profile, we can easily visualize `phi` in `theta` space by simply plotting `phi` vs. `theta`. This is essentially a form of linear interpolation, performed automatically by matplotlib when it draws lines between the discrete points of our data.<jupyter_code>def plot_profile():
    fig, (ax1, ax2, ax3) = plt.subplots(ncols=3, figsize=[8,5])
    ds.theta.plot(ax=ax1, y='z', marker='.', yincrease=False)
    ds.phi.plot(ax=ax2, y='z', marker='.', yincrease=False)
    ds.swap_dims({'z': 'theta'}).phi.plot(ax=ax3, y='theta', marker='.', yincrease=False)
    fig.subplots_adjust(wspace=0.5)
    return ax3

plot_profile();<jupyter_output><empty_output><jupyter_text>Linear transformationOk now lets transform `phi` to `theta` coordinates using linear interpolation.A key part of this is to define specific `theta` levels onto which we want to interpolate the data.<jupyter_code># First create an xgcm grid object
grid = Grid(ds, coords={'Z': {'center':'z'}}, periodic=False)

# define the target values in density, linearly spaced
theta_target = np.linspace(0, 3, 20)

# and transform
phi_transformed = grid.transform(ds.phi, 'Z', theta_target, target_data=ds.theta)
phi_transformed<jupyter_output><empty_output><jupyter_text>Now let's see what the result looks like.<jupyter_code>ax = plot_profile()
phi_transformed.plot(ax=ax, y='theta', marker='.')<jupyter_output><empty_output><jupyter_text>Not too bad. We can increase the number of interpolation levels to capture more of the small scale structure.<jupyter_code>target_theta = np.linspace(0,3, 100)
phi_transformed = grid.transform(ds.phi, 'Z', target_theta, target_data=ds.theta)
ax = plot_profile()
phi_transformed.plot(ax=ax, y='theta', marker='.')<jupyter_output><empty_output><jupyter_text>Note that by default, values of `theta_target` which lie outside the range of `theta` have been masked (set to `NaN`).To disable this behavior, you can pass `mask_edges=False`; values outside the range of `theta` will be filled with the nearest valid value.<jupyter_code>target_theta = np.linspace(0,3, 60)
phi_transformed = grid.transform(ds.phi, 'Z', target_theta, target_data=ds.theta, mask_edges=False)
ax = plot_profile()
phi_transformed.plot(ax=ax, y='theta', marker='.')<jupyter_output><empty_output><jupyter_text>Conservative transformationConservative transformation is designed to preseve the total sum of `phi` over the `Z` axis.It presumes that `phi` is an _extensive quantity_, i.e. a quantity that is already volume weighted, with respect to the Z axis: for example, units of `Kelvins * meters` for heat content, rather than just `Kelvins`.The conservative method requires more input data at the moment.You have to not only specify the coordinates of the cell centers, but also the cell faces (or bounds/boundaries). In xgcm we achieve this by defining the bounding coordinates as the `outer` axis position.The target `theta` values are likewise intepreted as cell boundaries in `theta`-space.In this way, conservative transformation is similar to calculating a histogram.<jupyter_code># define the cell bounds in depth
zc = np.arange(1,12)+0.5

# add them to the existing dataset
ds = ds.assign_coords({'zc': zc})
ds
# Recreate the grid object with a staggered `center`/`outer` coordinate layout
grid = Grid(ds, coords={'Z':{'center':'z', 'outer':'zc'}},
            periodic=False)
grid<jupyter_output><empty_output><jupyter_text>Currently the `target_data`(`theta` in this case) has to be located on the `outer` coordinate for the conservative method (compared to the `center` for the linear method).We can easily interpolate `theta` on the outer coordinate with the grid object.<jupyter_code>ds['theta_outer'] = grid.interp(ds.theta, 'Z', boundary='fill')
ds['theta_outer']<jupyter_output><empty_output><jupyter_text>Now lets transform the data using the conservative method. Note that the target values will now be interpreted as cell bounds and not cell centers as before.<jupyter_code># define the target values in density
theta_target = np.linspace(0,3, 20)

# and transform
phi_transformed_cons = grid.transform(ds.phi,
                                      'Z',
                                      theta_target,
                                      method='conservative',
                                      target_data=ds.theta_outer)
phi_transformed_cons
phi_transformed_cons.plot(y='theta_outer', marker='.', yincrease=False)<jupyter_output><empty_output><jupyter_text>There is no point in comparing `phi_transformed_cons` directly to `phi` or the results of linear interoplation, since here we have reinterpreted `phi` as an extensive quantity.However, we can verify that the sum of the two quantities over the Z axis is exactly the same.<jupyter_code>ds.phi.sum().values
phi_transformed_cons.sum().values<jupyter_output><empty_output><jupyter_text>Realistic Data ExampleTo illustrate these features in a more realistic example, we use data from the [NCEP Global Ocean Data Assimilation](https://www.cpc.ncep.noaa.gov/products/GODAS/) (GODAS).This data are available from the [Pangeo Cloud Data Library](https://catalog.pangeo.io/browse/master/ocean/GODAS/).We can see that this is a full, global, 4D ocean dataset.<jupyter_code>from intake import open_catalog

cat = open_catalog("https://raw.githubusercontent.com/pangeo-data/pangeo-datastore/master/intake-catalogs/ocean.yaml")
ds  = cat["GODAS"].to_dask()
ds<jupyter_output><empty_output><jupyter_text>The grid is missing an `outer` coordinate for the Z axis, so we will construct one.This will be needed for conservative interpolation.<jupyter_code>level_outer_data = np.pad(ds.level_w.values, [1, 0])
level_outer = xr.DataArray(
    level_outer_data,
    dims=['level_outer'],
    coords={'level_outer': ('level_outer', level_outer_data)}
)
ds = ds.assign_coords({'level_outer': level_outer})<jupyter_output><empty_output><jupyter_text>Now we create a `Grid` object for this dataset.<jupyter_code>grid = Grid(ds, coords={'Z': {'center': 'level', 'outer': 'level_outer'},
                        'X': {'center': 'lon', 'right': 'lon_u'},
                        'Y': {'center': 'lat', 'right': 'lat_u'}},
            periodic=['X'])
grid<jupyter_output><empty_output><jupyter_text>Linear InterpolationTo illustrate linear interpolation, we will interpolate salinity onto temperature surface.<jupyter_code># convert to standard units
theta = ds.pottmp - 273.15 
salt = 1000 * ds.salt
target_theta_levels = np.arange(-2, 36)
salt_on_theta = grid.transform(salt, 'Z', target_theta_levels, target_data=theta, method='linear')
salt_on_theta<jupyter_output><empty_output><jupyter_text>Note that the computation is lazy. (No data has been downloaded or computed yet.)We can trigger computation by plotting something.<jupyter_code>salt_on_theta.isel(time=0).sel(pottmp=18).plot()
salt_on_theta.isel(time=0).mean(dim='lon').plot(x='lat')<jupyter_output>/home/jovyan/xgcm/xgcm/transform.py:60: RuntimeWarning: invalid value encountered in _interp_1d_linear
  return _interp_1d_linear(phi, theta, target_theta_levels, mask_edges)
/srv/conda/envs/notebook/lib/python3.7/site-packages/dask/array/numpy_compat.py:41: RuntimeWarning: invalid value encountered in true_divide
  x = np.divide(x1, x2, out)<jupyter_text>Conservative InterpolationTo do conservative interpolation, we will attempt to calculate the meridional overturning in temperature space.Note that this is not a perfectly precise calculation, since the GODAS data are not exactly volume conserving as provided.However, it's sufficient to illustrate the basic principles of the calculation.To use conservative interpolation, we have to go from an intensive quantity (velocity) to an extensive one (velocity times cell thickness).We fill any missing values with 0, since they don't contribute to the transport.<jupyter_code>thickness = grid.diff(ds.level_outer, 'Z')
v_transport =  ds.vcur * thickness
v_transport = v_transport.fillna(0.).rename('v_transport')
v_transport<jupyter_output><empty_output><jupyter_text>We also need to interpolate `theta`, our target data for interoplation, to the same horizontal position as `v_transport`. This means moving from cell center to cell corner.This step introduces some considerable errors, particularly near the boundaries of bathymetry.(Xgcm currently has no special treatment for internal boundary conditions--see issue [222](https://github.com/xgcm/xgcm/issues/240).)<jupyter_code>theta = grid.interp(theta,['X', 'Y'], boundary='extend').rename('theta')
theta
v_transport_theta = grid.transform(v_transport, 'Z', target_theta_levels,
                                   target_data=theta, method='conservative')
v_transport_theta<jupyter_output>/home/jovyan/xgcm/xgcm/grid.py:909: UserWarning: The `target data` input is not located on the cell bounds. This method will continue with linear interpolation with repeated boundary values. For most accurate results provide values on cell bounds.
  UserWarning,<jupyter_text>Notice that this produced a warning. The `conservative` transformation method natively needs `target_data` to be provided on the cell bounds (here `level_outer`).Since transforming onto tracer coordinates is a very common scenario, xgcm uses linear interpolation to infer the values on the `outer` axis position.To demonstrate how to provide provide `target_data` on the outer grid position, we reproduce the steps xgcm executes internally:<jupyter_code>theta_outer = grid.interp(theta,['Z'], boundary='extend')
# the data cannot be chunked along the transformation axis
theta_outer = theta_outer.chunk({'level_outer': -1}).rename('theta')
theta_outer<jupyter_output><empty_output><jupyter_text>When we apply the transformation we can see that the results in this case are equivalent:<jupyter_code>v_transport_theta_manual = grid.transform(v_transport, 'Z', target_theta_levels,
                                   target_data=theta_outer, method='conservative')

# Warning: this step takes a long time to compute. We will only compare the first time value
xr.testing.assert_allclose(v_transport_theta_manual.isel(time=0), v_transport_theta.isel(time=0))<jupyter_output><empty_output><jupyter_text>Now  we verify visually that the vertically integrated transport is conserved under this transformation.<jupyter_code>v_transport.isel(time=0).sum(dim='level').plot(robust=True)
v_transport_theta.isel(time=0).sum(dim='theta').plot(robust=True)<jupyter_output><empty_output><jupyter_text>Finally, we attempt to plot a crude meridional overturning streamfunction for a single timestep.<jupyter_code>dx = 110e3 * np.cos(np.deg2rad(ds.lat_u))
(v_transport_theta.isel(time=0) * dx).sum(dim='lon_u').cumsum(dim='theta').plot.contourf(x='lat_u', levels=31)<jupyter_output>/srv/conda/envs/notebook/lib/python3.7/site-packages/xarray/plot/utils.py:624: MatplotlibDeprecationWarning: The 'extend' parameter to Colorbar has no effect because it is overridden by the mappable; it is deprecated since 3.3 and will be removed two minor releases later.
  cbar = fig.colorbar(primitive, **cbar_kwargs)
{'fraction_non_alphanumeric': 0.09954648526077098, 'fraction_numerical': 0.00763416477702192, 'mean_word_length': 5.632080200501253, 'pattern_counts': {'":': 0, '<': 86, '<?xml version=': 0, '>': 86, 'https://': 4, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '22', 'dir': 'jupyter-structured-clean-dedup', 'id': '4503', 'license': 'MIT', 'n_tokens_mistral': 4016, 'n_tokens_neox': 3846, 'n_words': 1314, 'path': 'doc/transform.ipynb', 'repo_name': 'jbusecke/xgcm', 'size': '13230'}