xarray.py

import numpy as np
import xarray as xr
import pandas as pd
from functools import reduce

from .cartopy import transform_points

def transform_dataset(ds, coords, from_crs, to_crs):
 """
 Transform coordinates of a Dataset or DataArray.

 Output can be directly assigned into a dataset.

 Arguments
 ---------
 ds: xarray dataset
 coords: iterable of dimension, coordinates, or data variable names
 that hold the coordinates to be transformed
 from_crs, to_crs: cartopy CRS instances
 """
 x = ds[coords[0]]
 y = ds[coords[1]]

 x, y = xr.broadcast(x, y) # coerce to at least 2D
 dims = x.dims

 xp, yp = transform_points(from_crs, to_crs, x.values, y.values)

 return ((dims, xp), (dims, yp))

def crop(ds, **constraints):
 """
 Crop an xarray dataset.

 Discards all coordinates where a given coordinate or data variable is outside
 the bounds given by the respective tuple.

 Arguments
 ---------
 ds: Dataset
 constraints: dict, where keys are the name of a coordinate or data_var,
 and values are tuples that will bound that coordinate or data_var.
 """
 ds = ds.set_coords(list(constraints))
 # set nan
 for var, (varmin, varmax) in constraints.items():
 ds = ds.where((ds[var]>varmin) & (ds[var]<varmax))
 # crop array
 for dim in ds.dims:
 other_dims = [d for d in ds.dims if d != dim]
 all_nans = [ds[var].isnull().all(other_dims) for var in ds.data_vars]
 all_nan = reduce(lambda x, y: x&y, all_nans)
 ds = ds.isel({dim: ~all_nan})

 return ds

def apply_1d(over_da, func, dim, **kwargs):
 """
 For those occasions where you'd think that ds.reduce() should do the trick,
 but you somehow don't have a function that already handles ndarrays.

 Parameters
 ----------
 over_da : xarray DataArray or list thereof
 func : Function that can handle a 1-dimensional ndarray
 dim : Dimension over which to apply func

 Usage
 -----
 e.g.: apply_1d(da, ols, dim='Depth', param='slope')

 License
 -------
 GNU-GPLv3, (C) A. Randelhoff
 (https://github.com/poplarShift/python-data-science-utils)
 """
 if not isinstance(over_da, list):
 over_da = [over_da]

 da_dropped = over_da[0].isel({dim: 0}).drop(dim)
 dims = da_dropped.coords.dims

 results = np.nan * da_dropped
 for idx, _ in np.ndenumerate(da_dropped):
 sel_dict = {c: i for c, i in zip(dims, idx)}
 res = func(
 *(da[sel_dict] for da in over_da),
 **kwargs
 )
 results[sel_dict] = res
 return results


def critical_index_value(vals, crit_val, dim, smaller_than):
 """
 Find the first value of an index (`dim`: str) where some value (`vals`: xr.DataArray) becomes
 larger or smaller (`smaller_than`: bool) than `crit_val` (float).

 Example
 -----
 To find Zeu, where iPAR drops below 1% of surface iPAR0-:
 ds['Zeu'] = critical_index_value(ds['iPAR']/ds['iPAR0minus'], 0.01, 'Depth', True)
 """
 if smaller_than:
 criterion = vals<crit_val
 else:
 criterion = vals>crit_val
 idx = criterion.reduce(np.argmax, dim=dim) #True>False
 return xr.where(idx>0, vals[dim].isel(**{dim:idx}), np.nan)
 

def ols(da, param='slope'):
 """
 Apply statsmodel's OLS regression to a 1d DataArray.
 Handles datetimes (in that case, regression is against days).

 Parameters
 ----------
 data : xarray DataArray
 str, one of ['slope', 'intercept', 'slope_pvalue',
 'intercept_pvalue', 'slope_se', 'intercept_se',]
 sought-after regression parameter

 Notes
 -----
 https://www.statsmodels.org/dev/generated/statsmodels.regression.linear_model.OLS.html

 License
 -------
 GNU-GPLv3, (C) A. Randelhoff
 (https://github.com/poplarShift/python-data-science-utils)
 """
 import statsmodels.api as sm

 data = da.dropna(dim=da.dims[0])

 # specify function with which to retrieve sought-after
 # parameter from statsmodels RegressionResultsWrapper
 res_fn = {
 'intercept': lambda res: res.params[0],
 'slope': lambda res: res.params[1],
 'intercept_pvalue': lambda res: res.pvalues[0],
 'slope_pvalue': lambda res: res.pvalues[1],
 'intercept_se': lambda res: res.bse[0],
 'slope_se': lambda res: res.bse[1],
 }

 if len(data)>=2:
 y = data.values
 xdata = data[data.dims[0]]
 if xdata.dtype.kind in ['M']:
 x = xdata.astype(float).values/1e9/86400.
 else:
 x = xdata.values
 ols = sm.OLS(y, sm.add_constant(x))
 res = ols.fit()
 return res_fn[param](res)
 else:
 return np.nan


# implement precision for uniqueness?
# np.round(12.3456789, decimals=4)

# np.unique(ds.longitude.isel(profile_id=0))[2]

def get_unique(x, axis=0):
 """
 Return unique non-nan value along specified xarray axis.
 Use this to squeeze out dimensions with length>1.

 Raises
 ------
 ValueError: if there are more than one non-nan values

 Usage
 -----
 ds.reduce(get_unique, dim=some_dim)

 License
 -------
 GNU-GPLv3, (C) A. Randelhoff
 (https://github.com/poplarShift/python-data-science-utils)
 """
 is_obj = x.dtype.char == 'O'
 is_dt = np.issubdtype(x.dtype, np.datetime64)
 x_ = np.moveaxis(x, source=axis, destination=-1)
 iteridx = x_.shape[:-1]
 if is_dt:
 u = np.zeros(iteridx, dtype=x.dtype)
 elif is_obj:
 u = np.zeros(iteridx, dtype=x.dtype)
 else: # numeric
 u = np.nan*np.zeros(iteridx)

 if not isinstance(u, np.ndarray):
 # if iteridx was empty tuple
 u = np.array(u)

 isnull = lambda x: pd.isnull(x) | (x == '')

 for i in np.ndindex(iteridx):
 u_1dim = np.unique(x_[i])
 u_non_null = u_1dim[~isnull(u_1dim)]
 if len(u_non_null)==1:
 u[i] = u_non_null[0]
 elif isnull(u_1dim[0]):
 if is_dt:
 u[i] = np.datetime64('NaT')
 elif is_obj:
 u[i] = ''
 else:
 u[i] = np.nan
 else:
 raise ValueError(f'Non-unique slices encountered at {i}!')
 return u