python4oceanographers

Turning ripples into waves

Using get_variables_by_attributes with xray

In my last post I showed the new get_variables_by_attributes method that will be available in netCDF4-python 1.2.0. But what if you want to use that now? You can just copy the method as a local function and change the self argument to take the the netCDF4.Dataset (nc) object.

In fact we can even expand that to xray or iris. Well... If your dataset is CF-compliant then it makes no sense to use this in iris. Just use iris methods instead ;-)

Below is an example with xray.

In [3]:
def get_variables_by_attributes(ds, **kwargs):
    vs = []
    
    has_value_flag  = False
    for vname, var in ds.items():
        for k, v in kwargs.items():
            if callable(v):
                has_value_flag = v(getattr(var, k, None))
                if has_value_flag is False:
                    break
            elif hasattr(var, k) and getattr(var, k) == v:
                has_value_flag = True
            else:
                has_value_flag = False
                break
        if has_value_flag is True:
            vs.append(ds[vname])

    return vs

And now let's do the same thing we did in the last post: construct the dimensioned vertical coordinate.

In [4]:
import xray


url = ("http://omgsrv1.meas.ncsu.edu:8080/thredds/dodsC/fmrc/sabgom/"
       "SABGOM_Forecast_Model_Run_Collection_best.ncd")


ds = xray.open_dataset(url)
In [5]:
formula_terms = lambda v: v is not None

var = get_variables_by_attributes(ds, formula_terms=formula_terms)[0]
In [6]:
var.attrs.get('formula_terms')
Out[6]:
u's: s_rho C: Cs_r eta: zeta depth: h depth_c: hc'
In [7]:
var.attrs.get('standard_name')
Out[7]:
u'ocean_s_coordinate_g1'

Now that we know the variable holding the formula_terms and all the variables in the formula_terms itself we can get them from the dataset.

Note that I am getting only the last time-step for eta. That is because I was unable to find a way to lazily re-shape arrays in xray.

In [8]:
s = ds['s_rho']
C = ds['Cs_r']
eta = ds['zeta'][-1, ...]  # Last time step.
depth = ds['h']
depth_c = ds['hc']

Remember that we need to match all the shapes before computing the vertical coordinate. Here I will take a different approach from the last post, that soon will be the default behavior in odvc, I will align the shapes based on a final desired shape.

Right now this approach is incomplete because it does not guess the final shape, but in odvc this will be possible. The current shapes are,

In [9]:
s.shape, C.shape, eta.shape, depth.shape
Out[9]:
((36,), (36,), (320, 440), (320, 440))

We want the vertical coordinate to be the leftmost dimension. That means the final shape should be (36, 320, 440).

In [10]:
final = (36, 320, 440)

def align(shape_in, shape_out):
    dims = [k for k in shape_in]
    def select(dim):
        return dim if dim in dims else 1
    return tuple(select(k) for k in shape_out)


def to_shape(arr, shape):
    reshape = align(shape_in=arr.shape, shape_out=final)
    return arr.data.reshape(reshape)
In [11]:
s = to_shape(s, shape=final)
C = to_shape(C, shape=final)
eta = to_shape(eta, shape=final)
depth = to_shape(depth, shape=final)
In [12]:
s.shape, C.shape, eta.shape, depth.shape
Out[12]:
((36, 1, 1), (36, 1, 1), (1, 320, 440), (1, 320, 440))

Nice! It works!! Now let's end the post with a fancy figure.

In [13]:
from odvc import ocean_s_coordinate_g1

z = ocean_s_coordinate_g1(s, C, eta, depth, depth_c.data)
In [14]:
salt = get_variables_by_attributes(ds, standard_name='sea_water_salinity')[0]

temp = get_variables_by_attributes(ds, standard_name='sea_water_potential_temperature')[0]
In [15]:
import matplotlib.pyplot as plt
import mpl_toolkits.axisartist as AA
from mpl_toolkits.axes_grid1 import host_subplot

import seaborn
seaborn.set(style='ticks')


def adjust_xlim(ax, offset):
    x1, x2 = ax.get_xlim()
    ax.set_xlim(x1 - offset, x2 + offset)


fig = plt.figure(figsize=(6, 9))
fig.subplots_adjust(bottom=0.1, top=0.85)
ax0 = host_subplot(111, axes_class=AA.Axes)

ax0.invert_yaxis()
ax1 = ax0.twiny()

new_axis0 = ax0.get_grid_helper().new_fixed_axis
new_axis1 = ax1.get_grid_helper().new_fixed_axis

ax0.axis["top"] = new_axis0(loc="top", axes=ax0, offset=(0, 0))
ax1.axis["top"] = new_axis1(loc="top", axes=ax1, offset=(0, 40))

ax0.axis["bottom"].toggle(all=False)
ax1.axis["bottom"].toggle(all=False)

ax0.set_ylabel("Depth (m)")
ax0.set_xlabel(u"Temperature (\xb0C)")
ax1.set_xlabel(r"Salinity (g kg$^{-1}$)")

kw = dict(linewidth=2.5)
l0, = ax0.plot(temp[-1, :, 5, 364], -z[:, 5, 364], **kw)
l1, = ax1.plot(salt[-1, :, 5, 364], -z[:, 5, 364], **kw)

adjust_xlim(ax0, offset=0.05)
adjust_xlim(ax1, offset=0.05)

ax0.axis["top"].label.set_color(l0.get_color())
ax1.axis["top"].label.set_color(l1.get_color())
In [16]:
HTML(html)
Out[16]:

This post was written as an IPython notebook. It is available for download or as a static html.

Creative Commons License
python4oceanographers by Filipe Fernandes is licensed under a Creative Commons Attribution-ShareAlike 4.0 International License.
Based on a work at https://ocefpaf.github.io/.

Comments