Some handy functions to retrieve coordinate info

[lee1043]

Below is for the functions I have made for the PMP, which I found it is pretty handy to have. If any interest to host this capability in the xCDAT, I'd be happy to help! Any feedback also would be appreciated.

import xcdat as xc
from typing import Union


# Retrieve coordinate key names


def get_axis_list(ds: Union[xr.Dataset, xr.DataArray]) -> list[str]:
    axes = list(ds.coords.keys())
    return(axes)


def get_time_key(ds: Union[xr.Dataset, xr.DataArray]) -> str:
    return xc.get_dim_keys(ds, "T")


def get_lat_key(ds: Union[xr.Dataset, xr.DataArray]) -> str:
    return xc.get_dim_keys(ds, "Y")


def get_lon_key(ds: Union[xr.Dataset, xr.DataArray]) -> str:
    return xc.get_dim_keys(ds, "X")


# Retrieve bounds key names


def get_time_bounds_key(ds: Union[xr.Dataset, xr.DataArray]) -> str:
    lat_key = get_time_key(ds)
    return ds[lat_key].attrs["bounds"]        


def get_lat_bounds_key(ds: Union[xr.Dataset, xr.DataArray]) -> str:
    lat_key = get_lat_key(ds)
    return ds[lat_key].attrs["bounds"]        


def get_lon_bounds_key(ds: Union[xr.Dataset, xr.DataArray]) -> str:
    lon_key = get_lon_key(ds)
    return ds[lon_key].attrs["bounds"]  


# Extract coordinate data


def get_time(ds: Union[xr.Dataset, xr.DataArray]) -> xr.DataArray:
    time_key = get_time_key(ds)
    time = ds[time_key]
    return(time)


def get_longitude(ds: Union[xr.Dataset, xr.DataArray]) -> xr.DataArray:
    lon_key = get_lon_key(ds)
    lon = ds[lon_key]
    return(lon)


def get_latitude(ds: Union[xr.Dataset, xr.DataArray]) -> xr.DataArray:
    lat_key = get_lat_key(ds)
    lat = ds[lat_key]
    return(lat)


# Extract coordinate bounds data


def get_time_bounds(ds: Union[xr.Dataset, xr.DataArray]) -> xr.DataArray:
    time_bounds_key = get_time_bounds_key(ds)
    time_bounds = ds[time_bounds_key]
    return(time_bounds)


def get_longitude_bounds(ds: Union[xr.Dataset, xr.DataArray]) -> xr.DataArray:
    lon_bounds_key = get_lon_bounds_key(ds)
    lon_bounds = ds[lon_bounds_key]
    return(lon_bounds)


def get_latitude_bounds(ds: Union[xr.Dataset, xr.DataArray]) -> xr.DataArray:
    lat_bounds_key = get_lat_bounds_key(ds)
    lat_bounds = ds[lat_bounds_key]
    return(lat_bounds)

[lee1043]

@acordonez FYI

[tomvothecoder]

Thanks everyone for the input. It sounds like these functions will be useful because they are clearly named, simple to use, and return the exact coordinates/bounds from an Xarray object.

I also get the sense that the use cases for xc.get_dim_keys() and xc.get_dim_coords() aren't as explicitly clear to some users (e.g., CDAT users, those not used to "dims" or "axis"). Both functions use Xarray's terminology of "dimension" to represent axis/coordinates, which I admit often confused me when I was first learning Xarray.

How about combining Jiwoo's functions into general functions that accept a name arg? On the dev side it is cleaner and easier to maintain, and on the user side they only need to import a few functions rather than many functions.

StandardName = Literal["lat", "lon", "time", "level")

def get_coords(obj: xr.Dataset | xr.DataArray, name: StandardName) -> xr.DataArray:
    pass

def get_coords_key(obj: xr.Dataset | xr.DataArray, name: StandardName) -> str | List[str]:
    pass

def get_coords_bnds(obj: xr.Dataset, name: StandardName): -> xr.DataArray
    # We have to consider that bounds are only available on xr.Dataset objects.
    pass

def get_coords_bnds_key(obj: xr.Dataset, name: StandardName) -> str | List[str]:
    # We have to consider that bounds are only available on xr.Dataset objects.
    # It is possible to extract the "bounds" key from xr.DataArray coordinates metadata though.
    pass

[pochedls]

I frequently forgot or confuse the difference between dim, coords, and axis, but I think it might make sense to use the word axis. If I understand correctly, a dim is the name for each axis (e.g., lat). A coordinate (singular) defines a point along that axis (e.g., time=["2020-01-01"]). And coords are DataArrays that represent all the points along an axis (e.g., latitude=[-89, -87, ..., 87, 89]).

But in climate science we call "arrays (coordinates) that label each point" "axes".

I don't know what the base for these functions should be called, but there is a case for axis-based names (i.e., get_axis(...)).

For me, I will definitely remember get_latitude(). I'll likely remember get_axis(). get_coords starts to become less intuitive. get_dim would be slightly more intuitive, but I think xarray considers dim a name for a coordinate axis (and not the axis itself).

I think my preference is still for get_latitude, get_time, etc. (this seems like a reasonable set of relevant functions given the scope and goal of xcdat as geospatial climate analysis software), but get_axis would be pretty okay, too (this would mean replacing "coords" with "axis" in the functions @tomvothecoder outlined above).

[lee1043]

Thank you all for the discussion. I think this is great and very helpful.

I agree to @pochedls and others and believe having such capabilities could lower the learning curve for the users. I also have been confused between dim and coords, and had to spend quite some time to get myself more familiar with them. While I do understand the concerns from @tomvothecoder that this might increase number of codes to maintain and which could be duplicates of some of simple functionality that already exists. However, I would appreciate if the tool provides more intuitive and straightforward APIs that could be used without studying its details. That also could provide more flexibility on at which level API that users want to use (e.g., get_latitude or get_dim_keys).

In the meantime I have been implementing similar features to the PMP (latest code in progress), which includes:

get_data_list
get_grid
get_latitude
get_latitude_bounds
get_latitude_bounds_key
get_latitude_key
get_longitude
get_longitude_bounds
get_longitude_bounds_key
get_longitude_key
get_time
get_time_bounds
get_time_bounds_key
get_time_key

Maybe some of these functions could be useful for xcdat if we decide to implement.

[tomvothecoder]

Thanks @pochedls and @lee1043. I see how those functions can be valuable for xCDAT users because they are easy to use with clear naming conventions (no need to figure out dims/coords/axis). The implementations are well-defined and simple so we most likely won't need to touch the code in the future. The "lower-level" xCDAT APIs that use Xarray terminology (xc.get_dim_coords(), xc.get_dim_keys()) can also be re-used in the implementation.

This could be a good first issue for a somebody else to try tackling.

[pochedls]

One thought on this...some datasets have multiple latitude axes. I'm not sure what these functions would do in that case (or what CDAT did). Something to think about.

[tomvothecoder]

Hey Jiwoo, I thought more about your suggestion. The main benefit I see is that the explicit naming convention for the utility functions might be helpful for some users, specifically former CDAT users. The trade-off is that we are duplicating functionality by wrapping simple xCDAT APIs in other xCDAT APIs. This results in more code to maintain and test, and more APIs options for the same functionalities which might confuse users. Also the xCDAT APIs are purposely written with general configurations and to follow Xarray terminology.

Unless there is a big enough friction with the existing APIs, I currently lean towards users learning the existing xCDAT APIs. They can extend these APIs if they really want/need to (like in your case). Happy to hear thoughts from others! @xCDAT/core-developers

Examples of duplication:

1. The get_X_bounds functions are duplicates of these functions:
   • ds.bounds.get_bounds() -- return bounds for a specific dataset axis or data variable axis
   • xr.Dataset.bounds.map() -- returns a dictionary mapping coords/axes to bounds keys
2. Using get_dims_coords() and get_dim_keys() requires almost the same number of lines as those helper functions (xc.get_lat_key(ds) vs. xc.get_dim_keys(ds, "Y"))
3. cf-xarray provides the .cf repr, and provides .cf.axes, .cf.coords, .cf.bounds mapping tables.

# https://cf-xarray.readthedocs.io/en/latest/coord_axes.html#using-the-repr
airds.cf

Coordinates:
             CF Axes: * X: ['lon']
                      * Y: ['lat']
                      * T: ['time']
                        Z: n/a

      CF Coordinates: * longitude: ['lon']
                      * latitude: ['lat']
                      * time: ['time']
                        vertical: n/a

       Cell Measures:   area: ['cell_area']
                        volume: n/a

      Standard Names: * latitude: ['lat']
                      * longitude: ['lon']
                      * time: ['time']

              Bounds:   n/a

       Grid Mappings:   n/a

Data Variables:
       Cell Measures:   area, volume: n/a

      Standard Names:   air_temperature: ['air']

              Bounds:   n/a

       Grid Mappings:   n/a

An alternative that might be worth considering is providing general functions that return mapping tables (get_axes, get_coords()), which the user can parse for the metadata explicitly. Or we can cite the cf-xarray .cf methods in our API docs if they want to access those mapping tables.

[durack1]

I am just hitting similar issues defining a generic Z coord for CMIP6 ocean data, this includes lev, depth, deptht, olevel, rho. I don't recall seeing any other vertical coord identifiers in earlier CMIP phases, but it wouldn't be hard to check

[pochedls]

xcdat's underlying capabilities are largely (or entirely?) inherited from cf-xarray. I think if the axis is defined as Z in the data, it would not be hard to be able to retrieve the vertical coordinates (I'm not sure how much is inferred from axis names).

[tomvothecoder]

If you need that functionality now, xCDAT has xc.get_dim_coords() which supports getting Z coords generally. xCDAT uses a combination of cf-xarray's interpretation of CF metadata and xCDAT's own fall-back mapping table (used if CF metadata is not set and interprets axis names).

[durack1]

Yep, IF :axis = "Z" ; which miraculously has worked for EVERY SINGLE CMIP6 historical thetao file I tried, nice - @tomvothecoder xc.get_dim_coords() works a treat - thanks gents

[tomvothecoder]

@durack1 awesome to hear that xc.get_dim_coords() works great for your needs!!

[jypeter]

I have went through this thread again, and I like of course any utility function (or simple wrapper) that will help beginning or in-a-hurry user easily find the information/data they need. I also understand the need for making the code easy to maintain. I think it is nice if a user can start typing ds.get and then check what the auto-completion feature suggests

Now, once we have a time axis, is there an intuitive way of checking the absolute time values, or should there also be a handy function similar to some_cdms2_variable.getTime().asComponentTime(). This was discussed in #310 (comment) and around

[tomvothecoder]

Now, once we have a time axis, is there an intuitive way of checking the absolute time values, or should there also be a handy function similar to some_cdms2_variable.getTime().asComponentTime(). This was discussed in #310 (comment) and around

When using open_dataset/open_mfdataset, you specify decode_times=True to get datetime/cftime values, or decode_times=False for absolute integer time values. You can also specify decode_times=False for absolute time values and use xr.decode_cf() if you want to decode times values to datetime/cftime later.

As far as I'm aware, Xarray (and xCDAT) does not have a built-in method to convert datetime/cftime time values BACK to absolute integer time values.

[pochedls]

I think you can use this function for encoding times back to integer time values: xr.coding.times.encode_cf_datetime().

[tomvothecoder]

Nice. I wonder why Xarray doesn't list some of these public APIs in their docs.

[durack1]

@tomvothecoder I have been a little frustrated with the same, but reflecting, that was 400% true with cdms/cdat as well, so people in glass houses..