[Add feature] Add cf compliant attributes

[rcaneill]

Hello, thanks for the great gsw package!
While using it with xarray, I thought that it would be very nice that gsw adds all the cf compliant (or at least the standard name + unit) attributes to the returned DataArray. I quickly went through the code, but everything is wrapped / ufuncs so I did not find any easy solution to add the attributes.

Does any of the core developer has any idea on this? If yes, I'll be happy to implement / test it (at least for the most common variables, as e.g. conservative and potential temperatures, sigma0, absolute salinity). If that was done, it would make it very easy to use the cf-xarray package with the outputs of gsw.

Another solution would be to create another package, e.g. cf-xarray-gsw, that would wrap all the gsw function, and add the proper attributes. Such a package could then work on datasets by using the cf attributes for the standard names to automatically get the proper input variables for the function calls (maybe a discussion on this should be started elsewhere than here).

[efiring]

I'm not sure how hard this would be. We are already using a decorator to do a little processing of input arguments and corresponding modification of output arguments to match: see match_args_return in gsw/_utilities.py. I'm guessing that a similar but much more complicated decorator would be needed for the xarray attribute handling you are proposing.

[rcaneill]

Ok thanks I'll try to give it a try when I have time.

[ocefpaf]

I have to confess that I believe this would be a nice extra package. Putting this inside gsw could make it harder to maintain. As a new package, maybe a plugin(?), this package could evolve faster and easier than adding the code here.

With that said, I'm not against it, I just believe an external package is a better solution in the long term.

[efiring]

I agree; I would not want a bunch of extra logic for xarray in the execution path for cases where the inputs are simple ndarrays or masked arrays. The xarray logic needs to be kept quite separate in order to keep things manageable. My first question is, can everything be done by substituting a different decorator for match_args_return? Or at least almost everything?

[DocOtak]

This has significant potential... I did a hack up with a decorator and some argo data.

import gsw
import xarray as xr
from functools import wraps

def cf_attrs(standard_name, units, **extra):
    def cf_attrs_decorator(func):
        @wraps(func)
        def cf_attrs_wrapper(*args, **kwargs):
            rv = func(*args, **kwargs)
            if isinstance(rv, xr.DataArray):
                rv.attrs["standard_name"] = standard_name
                rv.attrs["units"] = units
                rv.attrs.update(extra)
            
            return rv
        return cf_attrs_wrapper
    return cf_attrs_decorator

z_from_p = cf_attrs("height_above_mean_sea_level", "m")(gsw.z_from_p)

# load some argo file
ds = xr.load_dataset("3900198_prof.nc")

# try the wrapped
z_from_p(ds.PRES, ds.LATITUDE)

The output looked like

<xarray.DataArray 'PRES' (N_PROF: 60, N_LEVELS: 49)>
array([[   -4.65666253,    -9.31321878,   -19.02229371, ...,
        -1775.1811473 , -1872.70786343, -1970.8769359 ],
       [   -4.26028283,    -9.51120668,   -19.2201051 , ...,
        -1774.85619597, -1872.87245914, -1971.43237369],
       [   -4.26039255,    -9.41237592,   -19.12152908, ...,
        -1775.19665624, -1873.11701823, -1971.67937895],
       ...,
       [   -4.85707271,    -9.61490906,   -19.52671354, ...,
        -1770.60043627, -1868.46971246, -1966.39211757],
       [   -4.85722746,    -9.71433934,   -19.52733564, ...,
        -1770.36192252, -1868.0380133 ,            nan],
       [   -5.05538638,    -9.41679577,   -18.83315719, ...,
        -1770.42776011, -1868.29857141,            nan]])
Dimensions without coordinates: N_PROF, N_LEVELS
Attributes:
    long_name:       Sea water pressure, equals 0 at sea-level
    standard_name:   height_above_mean_sea_level
    units:           m
    valid_min:       0.0
    valid_max:       12000.0
    C_format:        %7.1f
    FORTRAN_format:  F7.1
    resolution:      1.0
    axis:            Z

If I were to expand on this (and I think I do want to), I'd probably make a xarray accessor that relies on cf-xarray to tell me what is available, and then figure out what gsw properties I could have it make. There are some cases about what to do when multiple variable have the same standard name in a dataset (for argo e.g. PRES and PRES_ADJUSTED both have sea_water_pressure as standard name). I also think in my toy example, that having additional arguments to z_form_p (the geopotential ones) might invalidate the standard name I chose, so that would need some handling.

[rcaneill]

Very nice!
Probably the most annoying work would be to write a dict containing all the standard names of the gsw functions outputs. @DocOtak inspired from your example, I made a quick tests that uses the name of the function to determine the attributes. It relies on cf-xarray to add the cf attributes:

import gsw
import xarray as xr
from functools import wraps
import numpy as np
import cf_xarray

# Create a dict containing the standard names of the outputs of the gsw functions
cf_name_gsw = {
    'z_from_p': 'height_above_mean_sea_level',
    'CT_from_t':'sea_water_conservative_temperature'
}

def cf_attrs(func):
    @wraps(func)
    def cf_attrs_wrapper(*args, **kwargs):
        rv = func(*args, **kwargs)
        if isinstance(rv, xr.DataArray):
            rv.attrs["standard_name"] = cf_name_gsw[func.__name__]
            rv = rv.cf.add_canonical_attributes(verbose=False)
        return rv
    return cf_attrs_wrapper

z_from_p = cf_attrs(gsw.z_from_p)
CT_from_t = cf_attrs(gsw.CT_from_t)

temp = 15 + 8 * np.random.randn(2, 2, 3)
SA = 35 + np.random.rand(2, 2, 3)
lon = [[-99.83, -99.32], [-99.79, -99.23]]
lat = [[42.25, 42.21], [42.63, 42.59]]
pres = [0, 100, 200]

ds = xr.Dataset(
    {
        "temp": (["x", "y", "time"], CT),
        "SA": (["x", "y", "time"], SA),
    },
        coords={
        "lon": (["x", "y"], lon),
        "lat": (["x", "y"], lat),
        "pres": pres,
    },
)
ds.pres.attrs['standard_name'] = 'sea_water_pressure_due_to_sea_water'
ds.lat.attrs['standard_name'] = 'latitude'
ds.lon.attrs['standard_name'] = 'longitude'

# try the wrapped
print(z_from_p(ds.pres, ds.lat))

The output is:

<xarray.DataArray 'pres' (pres: 3, x: 2, y: 2)>
array([[[  -0.        ,   -0.        ],
        [  -0.        ,   -0.        ]],

       [[ -99.189482  ,  -99.18984644],
        [ -99.18601757,  -99.18638243]],

       [[-198.33092659, -198.33165528],
        [-198.32399948, -198.32472902]]])
Coordinates:
  * pres     (pres) int64 0 100 200
    lon      (x, y) float64 -99.83 -99.32 -99.79 -99.23
    lat      (x, y) float64 42.25 42.21 42.63 42.59
Dimensions without coordinates: x, y
Attributes:
    standard_name:  height_above_mean_sea_level
    units:          m
    description:    "Height_above_X" means the vertical distance above the na...
    history:        Wed Dec 15 10:11:43 2021: cf.add_canonical_attributes(ove...

And for CT:

print(CT_from_t(ds.SA, ds.temp, ds.pres))

Output:

<xarray.DataArray 'SA' (x: 2, y: 2, time: 3, pres: 3)>
array([[[[ 9.27778514,  9.26676965,  9.25563131],
         [29.41616824, 29.39159422, 29.36698863],
         [ 8.81176841,  8.80097534,  8.79005785]],

        [[22.48839017, 22.46820436, 22.44796052],
         [16.30911126, 16.29305312, 16.27690924],
         [24.04467169, 24.02352564, 24.00232751]]],


       [[[ 5.29091601,  5.28293494,  5.27480682],
         [27.95646065, 27.93282581, 27.90915392],
         [27.39322922, 27.36996217, 27.34665588]],

        [[20.66688995, 20.64789477, 20.62883384],
         [ 1.09462962,  1.09006399,  1.08532215],
         [-0.15600342, -0.15952862, -0.16323929]]]])
Coordinates:
    lon      (x, y) float64 -99.83 -99.32 -99.79 -99.23
    lat      (x, y) float64 42.25 42.21 42.63 42.59
  * pres     (pres) int64 0 100 200
Dimensions without coordinates: x, y, time
Attributes:
    standard_name:  sea_water_conservative_temperature
    units:          K
    description:    Conservative Temperature is defined as part of the Thermo...
    history:        Wed Dec 15 10:06:31 2021: cf.add_canonical_attributes(ove...

[DocOtak]

Well I went a little nuts: https://github.com/DocOtak/gsw-xarray/blob/master/gsw_xarray/__init__.py

Clearly it's a fresh project and needs more docs and some... more work. However, it's installable right now with pip instal gsw-xarray and it's a drop in replacement/wrapper for the normal gsw library.

Right now it only adds/modified the standard_name and units attributes and leaves everything else alone, this can cause some... discrepancies.

>>> import gsw_xarray as gsw
>>> import xarray as xr
>>> data = xr.load_dataset("3900198_prof.nc")
>>> SA = gsw.SA_from_SP(data.PSAL, data.PRES, data.LONGITUDE, data.LATITUDE)
>>> SA
<xarray.DataArray 'PSAL' (N_PROF: 60, N_LEVELS: 49)>
array([[34.07422795, 34.07323092, 34.07219222, ...,  0.        ,
         0.        ,  0.        ],
       [33.97201818, 33.97102936, 33.97499554, ...,  0.        ,
         0.        ,  0.        ],
       [33.91766935, 33.91969222, 33.92065451, ...,  0.        ,
         0.        ,  0.        ],
       ...,
       [34.27673275, 34.27672559, 34.27667937, ..., 34.83418206,
        34.83428946, 34.83540519],
       [34.25456125, 34.25455565, 34.25552293, ..., 34.8340494 ,
        34.83315774,         nan],
       [34.25963474, 34.2596309 , 34.25958462, ..., 34.83512316,
        34.83422725,         nan]])
Dimensions without coordinates: N_PROF, N_LEVELS
Attributes:
    long_name:       Practical salinity
    standard_name:   sea_water_absolute_salinity
    units:           g kg-1
    valid_min:       2.0
    valid_max:       41.0
    C_format:        %9.3f
    FORTRAN_format:  F9.3
    resolution:      0.001

As you can see, the long_name is incorrect now, but the standard name and units attributes are correct. This becomes even "worse" as the xarray array_func wrapper will take the attributes of the first xarray data structure argument:

>>> gsw.CT_from_t(SA, data.TEMP, data.PRES)
<xarray.DataArray 'PSAL' (N_PROF: 60, N_LEVELS: 49)>
array([[ 3.37232957,  3.37704811,  3.34644193, ...,  1.93462053,
         1.87903843,  1.82540169],
       [ 2.06071968,  2.05745565,  1.98290691, ...,  1.75328904,
         1.71236196,  1.63094742],
       [ 1.89502887,  1.82472764,  1.73521056, ...,  1.32840633,
         1.24957374,  1.16962728],
       ...,
       [ 0.1346279 ,  0.13346063,  0.13511035, ...,  0.28749637,
         0.24708016,  0.22620097],
       [-0.23026039, -0.23341496, -0.23473247, ...,  0.26388665,
         0.21465588,         nan],
       [-0.27928389, -0.27942102, -0.27771855, ...,  0.29341063,
         0.24315915,         nan]])
Dimensions without coordinates: N_PROF, N_LEVELS
Attributes:
    long_name:       Practical salinity
    standard_name:   sea_water_conservative_temperature
    units:           degC
    valid_min:       2.0
    valid_max:       41.0
    C_format:        %9.3f
    FORTRAN_format:  F9.3
    resolution:      0.001

"unwrapped" functions just fallback to the gsw implementation, i.e. the attrs will be the same as the first xr.DataArray argument of the function:

>>> gsw.latentheat_evap_t(SA, data.TEMP)
<xarray.DataArray 'PSAL' (N_PROF: 60, N_LEVELS: 49)>
array([[2490534.85708187, 2490522.99775165, 2490594.33716134, ...,
        2494680.71568463, 2494805.55376088, 2494925.59402496],
       [2493656.90259427, 2493664.06380862, 2493839.93384263, ...,
        2495100.85974986, 2495192.09725111, 2495376.98026767],
       [2494053.16401431, 2494219.58838256, 2494431.23488775, ...,
        2496085.38585371, 2496265.51140843, 2496448.0483373 ],
       ...,
       [2498224.94795744, 2498227.32778714, 2498222.56993685, ...,
        2497608.73075887, 2497689.60083074, 2497722.86847858],
       [2499094.26112732, 2499101.40161442, 2499103.7523934 , ...,
        2497665.82164466, 2497768.1308122 ,              nan],
       [2499210.73340175, 2499210.73351825, 2499205.97460375, ...,
        2497594.43089151, 2497699.11731516,              nan]])
Dimensions without coordinates: N_PROF, N_LEVELS
Attributes:
    long_name:       Practical salinity
    standard_name:   sea_water_absolute_salinity
    units:           g kg-1
    valid_min:       2.0
    valid_max:       41.0
    C_format:        %9.3f
    FORTRAN_format:  F9.3
    resolution:      0.001

It's also safe to use non xarray structures:

>>> gsw.z_from_p(1000, 0)
-992.1173890198451

[rcaneill]

@DocOtak seems like we did the same thing :)

[rcaneill]

Should we move this discussion to the new project and close this issue?