Adding new data variables

Adding new data variables#

or more verbosely:

Scenario 1: Adding new data variables without altering the existing data structure#

If you want to write your own function or method for an xsnowDataset, simply code your function in a standard python module and decide whether you want to keep this module private to yourself or host it in a public repository.

Here is a cheat sheet for the steps you have to take. You will find more detailed explanations and a demonstration further below:

Recipe

  1. Import: from xsnow import xsnowDataset, register_xsnowDataset_method, registry

  2. Definition of function/method: @register_xsnowDataset_method def my_method(xs: xsnowDataset) -> Any:

  3. (optional) If returning an xsnowDataset, ._rewrap() the Dataset back into the correct class and make sure you created metadata for the newly added DataArray.

Regarding 1.)

  • If you use or return an xsnowDataset you have to import the class; alternatively, you can import xsnow and then use the class as xsnow.xsnowDataset.

  • (optional) If you want to create a method (and not only a function), import register_xsnowDataset_method.

  • (optional) If you add a variable that may exist in the xsnow registry of known variables with their names, units, etc, then import it.

Regarding 2.)

  • Define your function.

  • (optional) If you want to create a method, add @register_xsnowDataset_method in the line above your function signature and ensure that the first argument is an xsnowDataset.

Regarding 3.)

  • (optional) If you want to return an xsnowDataset, ensure that the class of the returned object is the same class as the input object. To be compatible with Scenario-2-extensions, please do that with the pattern
    xs_out = xs_in._rewrap(xs_modified)

  • (optional) If you want to return an xsnowDataset and added a new variable, please ensure that the new variable contains metadata, such as long_name, standard_name, and units.

The following code block shows an example from the Critical crack length extension, which applies all of these three steps. Feel free to use it as a template!

import numpy as np
import xarray as xr

from xsnow import xsnowDataset, register_xsnowDataset_method, registry
from xsnow.extensions.critical_crack_length \
    import compute_rho_slab, compute_rc_flat_richter2019

@register_xsnowDataset_method
def compute_critical_crack_length(xs: xsnowDataset) -> xsnowDataset:
    """
    Demo for customizing a new xsnowDataset method

    Parameters
    ----------
    xs : xsnowDataset
        The first object always has to be an xsnowDataset if this function is
        supposed to become a method!

    Returns
    -------
    new_xs : xsnowDataset
        A new xsnowDataset with a new data variable `rc_flat` [m].

    Examples
    --------
    import xsnow
    import mymodule  # wherever your custom method lives
    xs = xsnow.single_profile_timeseries()
    xs_new = xs.compute_critical_crack_length()
    """

    if xs.data is None:
        raise ValueError("Cannot compute on an empty dataset.")

    required = ['density', 'grain_size', 'shear_strength', 'height']
    for v in required:
        if v not in xs.data:
            raise ValueError(
                f"Missing required variable for rc_flat calculation: '{v}'"
            )

    density_arr = xs.data['density'].to_numpy()
    gsize_arr = xs.data['grain_size'].to_numpy()
    shearstrength_arr = xs.data['shear_strength'].to_numpy()
    height_arr = xs.data['height'].to_numpy()

    # Derive per-layer thickness from cumulative height along 'layer'.
    # layer=0 is bottom (height ~ 0), layer increases upward.
    height_bottom = height_arr.shift(layer=1, fill_value=0)
    thick_arr = height_arr - height_bottom  # [m]

    # Compute density of slab above each layer
    rho_slab_arr = compute_rho_slab(
        density=density_arr,
        thickness=thick_arr,
        xp=np,
    )

    # Physics core: Richter et al. (2019) flat-terrain rc
    rc_flat_arr = compute_rc_flat_richter2019(
        rho_layer=density_arr,
        grain_size_mm=gsize_arr,
        shear_strength_kpa=shearstrength_arr,
        rho_slab=rho_slab_arr,
        xp=np,
    )

    # Wrap result back into xarray using same dims/coords as height
    rc_flat_da = xr.DataArray(
        rc_flat_arr,
        dims=xs.data['height'].dims,
        coords=xs.data['height'].coords,
    )

    # Attach registry metadata if available
    if 'rc_flat' in registry:
        reg_entry = registry['rc_flat']
        rc_flat_da.attrs = {
            'long_name': reg_entry.get('long_name', 
                                       'critical crack length (flat terrain)'),
            'units': reg_entry.get('si_units', 'm'),
            'standard_name': reg_entry.get('standard_name', 'rc_flat'),
        }
    else:
        rc_flat_da.attrs = {
            'long_name': 'critical crack length on flat terrain',
            'units': 'm',
            'standard_name': 'rc_flat',
        }

    ## create new xarray.Dataset and wrap back into correct class:
    new_xs = xs.assign(rc_flat=rc_flat_da)
    new_xs = xs._rewrap(new_xs)
    return new_xs
Contents