Source code for xbout.calc.turbulence

import numpy as np
import xarray as xr


# TODO write a decorator to add functions as accessors to dataarrays?


def _rms_gufunc(x):
    """Generalized ufunc to calculate root mean squared value of an array."""
    squares = np.square(x)
    return np.sqrt(np.mean(squares, axis=-1))


[docs] def rms(da, dim=None, dask="parallelized", keep_attrs=True): """ Reduces a dataarray by calculating the root mean square along the dimension dim. """ # TODO If dim is None then take the root mean square along all dimensions? if dim is None: raise ValueError("Must supply a dimension along which to calculate rms") rms = xr.apply_ufunc( _rms_gufunc, da, input_core_dims=[[dim]], dask=dask, output_dtypes=[da.dtype], keep_attrs=keep_attrs, ) # Return the name of the da as variable_rms rms.name = str(da.name) + "_rms" return rms