Creating custom accessors

Creating custom accessors#

Introduction#

An accessor is a way of attaching a custom function to xarray objects so that it can be called as if it were a method while retaining a clear separation between the “core” xarray API and custom API. It enables you to easily extend (which is why you’ll sometimes see it referred to as an extension) and customize xarray’s functionality while limiting naming conflicts and minimizing the chances of your code breaking with xarray upgrades.

If you’ve used rioxarray (e.g. da.rio.crs) or hvplot (e.g. ds.hvplot()), you may have already used an xarray accessor without knowing it!

The Xarray documentation has some more technical details, and this tutorial provides example custom accessors and their uses.

Why create a custom accessor#

You can easily create a custom suite of tools that work on Xarray objects
It keeps your workflows cleaner and simpler
Your project-specific code is easy to share
It’s easy to implement: you don’t need to integrate any code into Xarray
It makes it easier to perform checks and write code documentation because you only have to create them once!

Easy steps to create your own accessor#

Create your custom class, including the mandatory __init__ method
Add the xr.register_dataarray_accessor() or xr.register_dataset_accessor()
Use your custom functions

Example 3: creating your own workflows with locally stored corrections#

Consider someone who frequently converts their elevations to be relative to the geoid (rather than the ellipsoid) using a custom, local conversion (otherwise, we’d recommend using an established conversion library like pyproj to switch between datums).

An accessor provides an elegant way to build (once) and apply (as often as needed!) this custom conversion on top of the existing xarray ecosystem without the need to copy-paste the code into the start of each project. By standardizing our approach and adding a few sanity checks within the accessor, we also eliminate the risk of accidentally applying the correction multiple times.

import rasterio
import xarray as xr


@xr.register_dataset_accessor("geoidxr")
class GeoidXR:
    """
    An extension for an XArray dataset that will calculate geoidal elevations from a local source file.
    """

    # ----------------------------------------------------------------------
    # Constructors

    def __init__(
        self,
        xrds,
    ):
        self._xrds = xrds
        # Running this function on init will check that my dataset has all the needed dimensions and variables
        # as specific to my workflow, saving time and headache later if they were missing and the computation fails
        # partway through.
        self._validate(
            self, req_dim=['x', 'y', 'dtime'], req_vars={'elevation': ['x', 'y', 'dtime']}
        )

    # ----------------------------------------------------------------------
    # Methods

    @staticmethod
    def _validate(self, req_dim=None, req_vars=None):
        '''
        Make sure the xarray dataset has the correct dimensions and variables

        Parameters
        ----------
        req_dim : list of str
            List of all required dimension names
        req_vars : list of str
            List of all required variable  names
        '''

        if req_dim is not None:
            if all([dim not in list(self._xrds.dims) for dim in req_dim]):
                raise AttributeError("Required dimensions are missing")
        if req_vars is not None:
            if all([var not in self._xrds.variables for var in req_vars.keys()]):
                raise AttributeError("Required variables are missing")

    # Notice that 'geoid' has been added to the req_vars list
    def to_geoid(
        self,
        req_dim=['dtime', 'x', 'y'],
        req_vars={'elevation': ['x', 'y', 'dtime', 'geoid']},
        source=None,
    ):
        """
        Get geoid layer from your local file, which is provided to the function as "source",
        and apply the offset to all elevation values.
        Adds 'geoid_offset' keyword to "offsets" attribute so you know the geoid offset was applied.

        Parameters
        ----------
        req_dim : list of str
            List of all required dimension names.
        req_vars : list of str
            List of all required variable  names
        source : str
            Full path to your source file containing geoid offsets
        """

        # check to make sure you haven't already run this function (and are thus applying the offset twice)
        try:
            values = self._xrds.attrs['offset_names']
            assert 'geoid_offset' not in values, "You've already applied the geoid offset!"
            values = list([values]) + ['geoid_offset']
        except KeyError:
            values = ['geoid_offset']

        self._validate(self, req_dim, req_vars)

        # read in your geoid values
        # WARNING: this implementation assumes your geoid values are in the same CRS and grid as the data you are applying
        # them to. If not, you will need to reproject and/or resample them to match the data to which you are applying them.
        # That step is not included here to emphasize the accessor aspect of the workflow.
        with rasterio.open(source) as src:
            geoid = src['geoid_varname']

        # As noted above, this step will fail or produce unreliable results if your data is not properly gridded
        self._xrds['elevation'] = self._xrds.elevation - geoid

        self._xrds.attrs['offset_names'] = values

        return self._xrds

Now, each time we want to convert our ellipsoid data to the geoid, we only have to run one line of code, and it will also perform a multitude of checks for us to make sure we’re performing exactly the operation we expect. Imagine the possibilities (and decrease in frustration)!

ds = ds.geoidxr.to_geoid(source='/Path/to/Custom/source/file.nc')
ds

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
File ~/micromamba/envs/xarray-tutorial/lib/python3.11/site-packages/xarray/core/extensions.py:38, in _CachedAccessor.__get__(self, obj, cls)
     37 try:
---> 38     accessor_obj = self._accessor(obj)
     39 except AttributeError:
     40     # __getattr__ on data object will swallow any AttributeErrors
     41     # raised when initializing the accessor, so we need to raise as
     42     # something else (GH933):

Cell In[7], line 22, in GeoidXR.__init__(self, xrds)
     19 # Running this function on init will check that my dataset has all the needed dimensions and variables
     20 # as specific to my workflow, saving time and headache later if they were missing and the computation fails
     21 # partway through.
---> 22 self._validate(
     23     self, req_dim=['x', 'y', 'dtime'], req_vars={'elevation': ['x', 'y', 'dtime']}
     24 )

Cell In[7], line 47, in GeoidXR._validate(self, req_dim, req_vars)
     46 if all([var not in self._xrds.variables for var in req_vars.keys()]):
---> 47     raise AttributeError("Required variables are missing")

AttributeError: Required variables are missing

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
Cell In[8], line 1
----> 1 ds = ds.geoidxr.to_geoid(source='/Path/to/Custom/source/file.nc')
      2 ds

File ~/micromamba/envs/xarray-tutorial/lib/python3.11/site-packages/xarray/core/extensions.py:43, in _CachedAccessor.__get__(self, obj, cls)
     38     accessor_obj = self._accessor(obj)
     39 except AttributeError:
     40     # __getattr__ on data object will swallow any AttributeErrors
     41     # raised when initializing the accessor, so we need to raise as
     42     # something else (GH933):
---> 43     raise RuntimeError(f"error initializing {self._name!r} accessor.")
     45 cache[self._name] = accessor_obj
     46 return accessor_obj

RuntimeError: error initializing 'geoidxr' accessor.

Creating custom accessors

Contents

Creating custom accessors#

Introduction#

Why create a custom accessor#

Easy steps to create your own accessor#

Example 1: accessing scipy functionality#

Example 2: creating your own workflows#

Example 3: creating your own workflows with locally stored corrections#