"""Contains the TArray, a DataArray with automatic unit handling."""
from dataclasses import dataclass
from functools import partial
from typing import Optional
from numbers import Number
import operator
import scipy
import pint
import xarray as xr
import numpy as np
import dask.array as da
from torx.units import Quantity
from torx.decorators import AutoDocClass
def _is_plain(other) -> bool:
"""Check if other is an array without norm information."""
if isinstance(other, xr.DataArray):
return "norm" not in other.attrs
return isinstance(other, (np.ndarray, da.Array))
def _assert_scalar(quantity):
"""Reject quantities with array magnitude, since they have no dims."""
if np.ndim(quantity.m) != 0:
raise TypeError("Quantity with array magnitude is not supported. "
"Convert it to TArray with dims first.")
def _drop_attrs(data):
"""Return the DataArray without attributes, used if the unit changes."""
data = data.copy(deep=False)
data.attrs = {}
return data
def _unsupported(first, second):
"""Return the error for unsupported operand types."""
return TypeError(f"Unsupported operand type(s) for "
f"{type(first).__name__} and {type(second).__name__}.")
[docs]
@dataclass
class TArray(AutoDocClass):
"""
Wrapper for xarray.DataArray with automatic unit handling.
The physical values are the data times the norm, a scalar Quantity.
Operations accept TArray, scalar quantities, DataArray objects with a
norm attribute, and numbers and arrays without norm (numpy, dask or
DataArray). Numbers and arrays without norm are dimensionless, so they
can only be added to or compared with a dimensionless TArray. As an
exception, the numbers zero and NaN have any unit. Arrays must be the
right operand, and a 0-d TArray needs a DataArray with dims instead of
a numpy or dask array with dims. Attributes other than norm are kept,
unless the unit changes. Numpy functions are dispatched to the methods
of the same name.
"""
data: xr.DataArray
norm: Optional[Quantity] = None
def __post_init__(self):
"""Take the norm from the DataArray attributes, defaulting to 1."""
if not isinstance(self.data, xr.DataArray):
self.data = xr.DataArray(self.data)
if "norm" in self.data.attrs:
data_norm = self.data.attrs["norm"]
if (self.norm is not None
and Quantity(self.norm) != Quantity(data_norm)):
raise ValueError(f"The norm {self.norm} conflicts with the "
f"norm {data_norm} of the data attributes.")
self.norm = data_norm
# A shallow copy shares the values but not the attributes, such that
# the input array is not modified
data = self.data.copy(deep=False)
data.attrs.pop("norm", None)
data.name = None
self.data = data
if self.norm is None:
self.norm = Quantity(1)
elif not isinstance(self.norm, Quantity):
self.norm = Quantity(self.norm)
_assert_scalar(self.norm)
# Methods like real or max assume a real norm
if np.iscomplexobj(self.norm.m):
raise TypeError(f"The norm {self.norm} is complex, only real "
"norms are supported.")
# A positive norm keeps the order of the data, e.g. for max or <
if self.norm.m < 0:
self.data, self.norm = -self.data, -self.norm
@property
def shape(self):
"""Shape of the array."""
return self.data.shape
@property
def ndim(self):
"""Number of dimensions of the array."""
return self.data.ndim
@property
def size(self):
"""Number of elements of the array."""
return self.data.size
@property
def dims(self):
"""Dimensions of the array."""
return self.data.dims
@property
def coords(self):
"""Coordinates of the array."""
return self.data.coords
@property
def values(self):
"""Data values of the array."""
return self.data.values
@property
def denormalized_values(self):
"""Data values of the array multiplied by the norm."""
return self.__array__()
@property
def unit(self):
"""Unit of the data values."""
return self.norm.u
def to(self, unit):
"""Convert data to provided unit."""
norm = self.norm.to(unit)
data = self.data if norm.u == self.norm.u else _drop_attrs(self.data)
return TArray(data, norm)
def tobase(self):
"""Convert data to base unit, e.g. km to m."""
return self.to(self.norm.to_base_units().u)
def denormalize(self):
"""Return the TArray with the norm magnitude applied to the data."""
return TArray(self.data * self.norm.m, Quantity(1, self.norm.u))
def item(self):
"""Return the value of a TArray with one element as Quantity."""
return Quantity(self.data.values.item() * self.norm.m, self.norm.u)
def __array__(self, dtype=None, copy=None):
"""Return the array value times the norm."""
if self.norm.m == 1:
return np.asarray(self.data.values, dtype=dtype, copy=copy)
# Applying the norm creates a new array
if copy is False:
raise ValueError("Unable to avoid a copy, since the norm is "
"applied to the data.")
return np.asarray(self.data.values * self.norm.m, dtype=dtype)
def __float__(self):
"""Return the value of a dimensionless TArray with one element."""
return float(self.item().to("dimensionless").m)
def __bool__(self):
"""Return the truth value of a TArray with one element."""
return bool(self.data.values)
[docs]
def __len__(self):
"""Return the length of the first dimension, like numpy."""
return len(self.data)
def __iter__(self):
"""Return an iterator over the first dimension."""
for i in self.data:
yield TArray(i, self.norm)
def _detect_plot(self, key):
"""Detect if there is a plot operation from pyplot."""
if isinstance(key, (tuple, list)):
return any(k is np.newaxis for k in key)
return key is np.newaxis
[docs]
def __getitem__(self, key):
"""Return an element of the array, a TArray can be used as mask."""
if self._detect_plot(key):
return self.data.values.__getitem__(key) * self.norm.m
if isinstance(key, TArray):
key = key.data
return TArray(self.data.__getitem__(key), self.norm)
def _coerce(self, other):
"""Promote a DataArray with norm attribute to TArray."""
if isinstance(other, xr.DataArray) and "norm" in other.attrs:
return TArray(other)
return other
def _as_tarray(self, other):
"""Return other as TArray, if it carries unit information."""
other = self._coerce(other)
if isinstance(other, TArray):
return other
if isinstance(other, Number):
return TArray(xr.DataArray(other))
if isinstance(other, Quantity):
_assert_scalar(other)
return TArray(xr.DataArray(other.m), Quantity(1, other.u))
raise _unsupported(self, other)
def _scaled(self, other):
"""Return the data of other in the norm of self."""
if isinstance(other, Number) and (other == 0 or np.isnan(other)):
return other
if _is_plain(other):
# Arrays without norm are dimensionless numbers
self._assert_aligned(other)
factor = self.norm.to("dimensionless").m
return other if factor == 1 else other / factor
other = self._as_tarray(other)
if other.norm == self.norm:
return other.data
return other.data * (other.norm / self.norm).to("dimensionless").m
def _assert_aligned(self, other: np.ndarray | da.Array | xr.DataArray):
"""Reject arrays without dims for a 0-d TArray, which has no dims."""
if (self.ndim == 0 and not isinstance(other, xr.DataArray)
and np.ndim(other) > 0):
raise TypeError("A 0-d TArray has no dims to align with a numpy "
"or dask array. Use a DataArray with dims.")
def _assert_not_plain(self, other):
"""Reject arrays as left operand, since the norm would be lost."""
if _is_plain(other) or isinstance(other, xr.DataArray):
raise TypeError(f"Unsupported operand order {type(other).__name__}"
" and TArray. Use the TArray as left operand.")
def _combine(self, other, op):
"""Combine with other of compatible norm, keeping the norm of self."""
return TArray(op(self.data, self._scaled(other)), self.norm)
def _dimensionless_data(self):
"""Return the data of a dimensionless TArray with norm 1."""
factor = self.norm.to("dimensionless").m
return self.data if factor == 1 else self.data * factor
def _dimensionless(self, func):
"""Apply a function to a dimensionless TArray."""
return TArray(_drop_attrs(func(self._dimensionless_data())))
def _in_unit(self, func):
"""Apply a function to the values in the unit of the norm."""
if self.norm.dimensionless:
return self._dimensionless(func)
denormalized = self.denormalize()
return TArray(func(denormalized.data), denormalized.norm)
def __add__(self, other):
"""Add an object with compatible norm, keeping the norm of self."""
return self._combine(other, operator.add)
def __radd__(self, other):
"""Add from the right side."""
self._assert_not_plain(other)
return self._combine(other, operator.add)
def __sub__(self, other):
"""Subtract an object with compatible norm."""
return self._combine(other, operator.sub)
def __rsub__(self, other):
"""Subtract from the right side."""
self._assert_not_plain(other)
return self._combine(other, lambda x, y: y - x)
def _scale(self, other, op):
"""Multiply or divide, a Quantity only changes the norm."""
other = self._coerce(other)
if isinstance(other, TArray):
return TArray(_drop_attrs(op(self.data, other.data)),
op(self.norm, other.norm))
if isinstance(other, Number):
return TArray(op(self.data, other), self.norm)
if isinstance(other, Quantity):
_assert_scalar(other)
return TArray(_drop_attrs(self.data), op(self.norm, other))
if _is_plain(other):
# Arrays without norm are dimensionless numbers
self._assert_aligned(other)
return TArray(op(self.data, other), self.norm)
raise _unsupported(self, other)
def __mul__(self, other):
"""Multiply with TArray objects, numbers, quantities and arrays."""
return self._scale(other, operator.mul)
def __rmul__(self, other):
"""Multiply from the right side."""
self._assert_not_plain(other)
return self * other
def __truediv__(self, other):
"""Divide by TArray objects, numbers, quantities and arrays."""
return self._scale(other, operator.truediv)
def __rtruediv__(self, other):
"""Divide from the right side."""
self._assert_not_plain(other)
return self.reciprocal() * other
def __pow__(self, exponent):
"""Raise to the power of a number."""
return TArray(_drop_attrs(self.data**exponent), self.norm**exponent)
def __rpow__(self, base):
"""Raise a number to the power of a dimensionless TArray."""
self._assert_not_plain(base)
return self._dimensionless(lambda x: base**x)
def __neg__(self):
"""Multiply by negative one."""
return TArray(-self.data, self.norm)
def __pos__(self):
"""Multiply by positive one."""
return self
def __abs__(self):
"""Return the absolute value."""
return TArray(abs(self.data), self.norm)
add = __add__
subtract = __sub__
multiply = __mul__
divide = __truediv__
power = __pow__
negative = __neg__
positive = __pos__
absolute = abs = __abs__
def real(self):
"""Return the real part."""
return TArray(self.data.real, self.norm)
def imag(self):
"""Return the imaginary part."""
return TArray(self.data.imag, self.norm)
def conjugate(self):
"""Return the complex conjugate."""
return TArray(self.data.conjugate(), self.norm)
def sqrt(self):
"""Return the square root, the norm is square rooted."""
return TArray(_drop_attrs(np.sqrt(self.data)), self.norm**0.5)
def cbrt(self):
"""Return the cubic root, the norm is cubic rooted."""
return TArray(_drop_attrs(np.cbrt(self.data)), self.norm**(1 / 3))
def square(self):
"""Return the square, the norm is squared."""
return TArray(_drop_attrs(np.square(self.data)), self.norm**2)
def reciprocal(self):
"""Return the reciprocal, the norm is inverted."""
# 1/x instead of np.reciprocal, which truncates integers
return TArray(_drop_attrs(1 / self.data), 1 / self.norm)
def maximum(self, other):
"""Return the elementwise maximum with an object of compatible norm."""
return self._combine(other, np.maximum)
def minimum(self, other):
"""Return the elementwise minimum with an object of compatible norm."""
return self._combine(other, np.minimum)
def fmax(self, other):
"""Return the elementwise maximum, ignoring NaN."""
return self._combine(other, np.fmax)
def fmin(self, other):
"""Return the elementwise minimum, ignoring NaN."""
return self._combine(other, np.fmin)
def hypot(self, other):
"""Return sqrt(self**2 + other**2) for an object of compatible norm."""
return self._combine(other, np.hypot)
def arctan2(self, other):
"""Return the angle arctan(self/other) in the correct quadrant."""
return TArray(_drop_attrs(np.arctan2(self.data, self._scaled(other))))
def clip(self, a_min=None, a_max=None, **bounds):
"""Limit the values to bounds of compatible norm."""
# Like numpy, the bounds can be given as a_min, a_max or min, max
lower = bounds.pop("min", a_min)
upper = bounds.pop("max", a_max)
if bounds:
raise TypeError(f"Keywords {list(bounds)} are not implemented "
"for clip.")
lower = None if lower is None else self._scaled(lower)
upper = None if upper is None else self._scaled(upper)
return TArray(self.data.clip(lower, upper), self.norm)
def where(self, cond, other=np.nan):
"""Keep values where cond is True, otherwise use other."""
if isinstance(cond, TArray):
cond = cond.data
return TArray(self.data.where(cond, self._scaled(other)), self.norm)
def fillna(self, value):
"""Fill NaN values with a value of compatible norm."""
return TArray(self.data.fillna(self._scaled(value)), self.norm)
def round(self, decimals=0):
"""Round the values in the unit of the norm."""
return self._in_unit(lambda x: x.round(decimals))
def rint(self):
"""Round the values to the nearest integer in the unit of the norm."""
return self._in_unit(np.rint)
def floor(self):
"""Return the floor of the values in the unit of the norm."""
return self._in_unit(np.floor)
def ceil(self):
"""Return the ceiling of the values in the unit of the norm."""
return self._in_unit(np.ceil)
def trunc(self):
"""Truncate the values in the unit of the norm."""
return self._in_unit(np.trunc)
def sin(self):
"""Return the sine of a dimensionless TArray."""
return self._dimensionless(np.sin)
def cos(self):
"""Return the cosine of a dimensionless TArray."""
return self._dimensionless(np.cos)
def tan(self):
"""Return the tangent of a dimensionless TArray."""
return self._dimensionless(np.tan)
def arcsin(self):
"""Return the inverse sine of a dimensionless TArray."""
return self._dimensionless(np.arcsin)
def arccos(self):
"""Return the inverse cosine of a dimensionless TArray."""
return self._dimensionless(np.arccos)
def arctan(self):
"""Return the inverse tangent of a dimensionless TArray."""
return self._dimensionless(np.arctan)
def sinh(self):
"""Return the hyperbolic sine of a dimensionless TArray."""
return self._dimensionless(np.sinh)
def cosh(self):
"""Return the hyperbolic cosine of a dimensionless TArray."""
return self._dimensionless(np.cosh)
def tanh(self):
"""Return the hyperbolic tangent of a dimensionless TArray."""
return self._dimensionless(np.tanh)
def arcsinh(self):
"""Return the inverse hyperbolic sine of a dimensionless TArray."""
return self._dimensionless(np.arcsinh)
def arccosh(self):
"""Return the inverse hyperbolic cosine of a dimensionless TArray."""
return self._dimensionless(np.arccosh)
def arctanh(self):
"""Return the inverse hyperbolic tangent of a dimensionless TArray."""
return self._dimensionless(np.arctanh)
def exp(self):
"""Return the exponential of a dimensionless TArray."""
return self._dimensionless(np.exp)
def expm1(self):
"""Return exp(x) - 1 of a dimensionless TArray."""
return self._dimensionless(np.expm1)
def log(self):
"""Return the natural logarithm of a dimensionless TArray."""
return self._dimensionless(np.log)
def log10(self):
"""Return the base 10 logarithm of a dimensionless TArray."""
return self._dimensionless(np.log10)
def log2(self):
"""Return the base 2 logarithm of a dimensionless TArray."""
return self._dimensionless(np.log2)
def log1p(self):
"""Return log(1 + x) of a dimensionless TArray."""
return self._dimensionless(np.log1p)
def sign(self):
"""Return the sign of the values."""
return TArray(_drop_attrs(np.sign(self.data)))
def isfinite(self):
"""Test elementwise for finite values."""
return TArray(_drop_attrs(np.isfinite(self.data)))
def isinf(self):
"""Test elementwise for infinite values."""
return TArray(_drop_attrs(np.isinf(self.data)))
def isnan(self):
"""Test elementwise for NaN."""
return TArray(_drop_attrs(np.isnan(self.data)))
def _reduce(self, name, *args, power=1, **kwargs):
"""Apply a DataArray reduction, raising the norm to power."""
data = getattr(self.data, name)(*args, **kwargs)
data = data if power == 1 else _drop_attrs(data)
return TArray(data, self.norm**power)
def min(self, *args, **kwargs):
"""Return the minimum, see xarray.DataArray.min."""
return self._reduce("min", *args, **kwargs)
def max(self, *args, **kwargs):
"""Return the maximum, see xarray.DataArray.max."""
return self._reduce("max", *args, **kwargs)
def sum(self, *args, **kwargs):
"""Return the sum, see xarray.DataArray.sum."""
return self._reduce("sum", *args, **kwargs)
def cumsum(self, *args, **kwargs):
"""Return the cumulative sum, see xarray.DataArray.cumsum."""
return self._reduce("cumsum", *args, **kwargs)
def mean(self, *args, **kwargs):
"""Return the mean, see xarray.DataArray.mean."""
return self._reduce("mean", *args, **kwargs)
def median(self, *args, **kwargs):
"""Return the median, see xarray.DataArray.median."""
return self._reduce("median", *args, **kwargs)
def std(self, *args, **kwargs):
"""Return the standard deviation, see xarray.DataArray.std."""
return self._reduce("std", *args, **kwargs)
def var(self, *args, **kwargs):
"""Return the variance, the norm is squared."""
return self._reduce("var", *args, power=2, **kwargs)
def quantile(self, *args, **kwargs):
"""Return the quantile, see xarray.DataArray.quantile."""
return self._reduce("quantile", *args, **kwargs)
def all(self, *args, **kwargs):
"""Test whether all values are nonzero."""
return TArray(_drop_attrs(self.data.all(*args, **kwargs)))
def any(self, *args, **kwargs):
"""Test whether any value is nonzero."""
return TArray(_drop_attrs(self.data.any(*args, **kwargs)))
def prod(self, dim=None, **kwargs):
"""Return the product, the norm has the power of the factor count."""
# NaN values are not skipped, since they would change the count
data = self.data.prod(dim, skipna=False, **kwargs)
count = self.size // data.size if data.size else 0
return TArray(_drop_attrs(data), self.norm**count)
def nanprod(self, dim=None, **kwargs):
"""Return the product skipping NaN, along dims if dimensionless."""
if self.norm.dimensionless:
return self._dimensionless(
lambda x: x.prod(dim, skipna=True, **kwargs))
# With unit, the number of factors would vary along remaining dims
if dim is not None or kwargs:
raise ValueError("Only a dimensionless TArray supports nanprod "
"along dims, since the unit varies otherwise.")
count = int(self.data.notnull().sum())
return TArray(_drop_attrs(self.data.prod(skipna=True)),
self.norm**count)
def cumprod(self, *args, **kwargs):
"""Return the cumulative product of a dimensionless TArray."""
return self._dimensionless(lambda x: x.cumprod(*args, **kwargs))
def _get_axis_from_dims(self, axis=None, dims=None):
"""Return the axis from dimensions."""
if dims is None:
return axis
# The axis is kept for numpy calls like np.gradient(tarray, axis=0)
if axis is not None:
raise ValueError("Provide either axis or dims, not both.")
dims = (dims,) if isinstance(dims, str) else dims
return tuple(self.data.dims.index(dim) for dim in dims)
def _argfunc(self, func, axis=None, dims=None):
"""Perform function on TArray objects returning indices."""
axis = self._get_axis_from_dims(axis, dims)
if isinstance(axis, tuple):
(axis,) = axis
return TArray(func(self.values, axis=axis))
def argmin(self, axis=None, dims=None):
"""Return the index where the minimum is found."""
return self._argfunc(np.argmin, axis=axis, dims=dims)
def nanargmin(self, axis=None, dims=None):
"""Return the index where the minimum is found, ignoring NaN."""
return self._argfunc(np.nanargmin, axis=axis, dims=dims)
def argmax(self, axis=None, dims=None):
"""Return the index where the maximum is found."""
return self._argfunc(np.argmax, axis=axis, dims=dims)
def nanargmax(self, axis=None, dims=None):
"""Return the index where the maximum is found, ignoring NaN."""
return self._argfunc(np.nanargmax, axis=axis, dims=dims)
def argwhere(self):
"""Return indices where the values are nonzero."""
res = np.argwhere(self.values)
return TArray(xr.DataArray(res, dims=("matches", "ndim")))
def gradient(self, x=None, *, axis=None, dims=None, edge_order=1):
"""Return the derivative, x is the coordinates or the spacing."""
axis = self._get_axis_from_dims(axis, dims)
norm = self.norm
if isinstance(x, (TArray, Quantity)):
x = self._as_tarray(x)
norm = norm / x.norm
x = x.values
args = () if x is None else (x,)
grad = np.gradient(self.data, *args, axis=axis, edge_order=edge_order)
grad = xr.DataArray(grad, dims=self.dims, coords=self.coords)
return TArray(grad, norm)
def gradlog(self, x=None, *, axis=None, dims=None):
"""Compute the gradient of the log of the TArray."""
return TArray(self.data).log().gradient(x, axis=axis, dims=dims)
def trapezoid(self, *args, x=None, dx=1.0, axis=None, dims=None,
cumulative=False):
"""Return the integral along axes, args are TArray coordinates."""
if x is not None:
if args:
raise ValueError("Provide the coordinates either as args or "
"as x, not both.")
args = (x,)
# The cumulative integral starts with zero to keep the coordinates
func = (partial(scipy.integrate.cumulative_trapezoid, initial=0)
if cumulative else np.trapezoid)
dims = (dims,) if isinstance(dims, str) else dims
axis = self._get_axis_from_dims(axis, dims)
axis = np.lib.array_utils.normalize_axis_tuple(
-1 if axis is None else axis, self.ndim)
if args and len(args) != len(axis):
raise ValueError("Provide one coordinate per axis.")
# Integrate the last axis first, such that the other axes stay valid
coords = args if args else (None,) * len(axis)
res, norm = self.data, self.norm
for ax, coord in sorted(zip(axis, coords), key=lambda pair: pair[0],
reverse=True):
if coord is None:
res = func(res, dx=dx, axis=ax)
else:
res = func(res, x=coord.data, axis=ax)
norm = norm * coord.norm
if cumulative:
res = xr.DataArray(res, dims=self.dims, coords=self.coords)
elif dims is not None:
residual_dims = tuple(dim for dim in self.dims if dim not in dims)
residual_coords = [self.coords[dim] for dim in residual_dims]
res = xr.DataArray(res, dims=residual_dims, coords=residual_coords)
return TArray(res, norm)
def cumulative_trapezoid(self, *args, x=None, dx=1.0, axis=None,
dims=None):
"""Calculate the cumulative integral in the given axis."""
return self.trapezoid(*args, x=x, dx=dx, axis=axis, dims=dims,
cumulative=True)
def _compare(self, other, op):
"""Compare with an object of compatible norm or an array."""
return TArray(_drop_attrs(op(self.data, self._scaled(other))))
def __lt__(self, other):
"""Return if TArray is less than other (elementwise)."""
return self._compare(other, operator.lt)
def __gt__(self, other):
"""Return if TArray is greater than other (elementwise)."""
return self._compare(other, operator.gt)
def __le__(self, other):
"""Return if TArray is less than or equal to other (elementwise)."""
return self._compare(other, operator.le)
def __ge__(self, other):
"""Return if TArray is greater than or equal to other (elementwise)."""
return self._compare(other, operator.ge)
def __eq__(self, other):
"""Return if TArray is equal to other (elementwise)."""
return self._compare(other, operator.eq)
def __ne__(self, other):
"""Return if TArray is not equal to other (elementwise)."""
return self._compare(other, operator.ne)
less = __lt__
greater = __gt__
less_equal = __le__
greater_equal = __ge__
equal = __eq__
not_equal = __ne__
def isclose(self, other, rtol=1e-05, atol=0, equal_nan=False):
"""Return if close to other (elementwise), atol of compatible norm."""
atol = self._scaled(atol)
# A 0-d DataArray has no dims to align with numpy or dask arrays
if isinstance(atol, xr.DataArray) and atol.ndim == 0:
atol = atol.data
def close(x, y):
result = abs(x - y) <= atol + rtol * abs(y)
if equal_nan:
result = result | (np.isnan(x) & np.isnan(y))
return result
return self._compare(other, close)
def allclose(self, other, rtol=1e-05, atol=0, equal_nan=False):
"""Return if TArray is close to other (all values)."""
return self.isclose(other, rtol, atol, equal_nan).all()
def _logical(self, other, op):
"""Perform logical operation of dimensionless objects."""
if _is_plain(other):
self._assert_aligned(other)
else:
other = self._as_tarray(other)._dimensionless_data()
return TArray(op(self._dimensionless_data(), other))
def __and__(self, other):
"""Return logical and of TArray (elementwise)."""
return self._logical(other, np.logical_and)
def __or__(self, other):
"""Return logical or of TArray (elementwise)."""
return self._logical(other, np.logical_or)
def __xor__(self, other):
"""Return logical xor of TArray (elementwise)."""
return self._logical(other, np.logical_xor)
def __invert__(self):
"""Return logical not of TArray (elementwise)."""
return TArray(np.logical_not(self._dimensionless_data()))
logical_and = __rand__ = __and__
logical_or = __ror__ = __or__
logical_xor = __rxor__ = __xor__
logical_not = __invert__
def sel(self, *args, **kwargs):
"""Select by coordinate labels, see xarray.DataArray.sel."""
return TArray(self.data.sel(*args, **kwargs), self.norm)
def isel(self, *args, **kwargs):
"""Select by index, see xarray.DataArray.isel."""
return TArray(self.data.isel(*args, **kwargs), self.norm)
def compute(self, **kwargs):
"""Compute the chunks, adding full data to memory."""
return TArray(self.data.compute(**kwargs), self.norm)
def persist(self, **kwargs):
"""Compute the chunks, keeping the data as dask array."""
return TArray(self.data.persist(**kwargs), self.norm)
def load(self, **kwargs):
"""Load the data into memory."""
return TArray(self.data.load(**kwargs), self.norm)
def chunk(self, *args, **kwargs):
"""Chunk the data with dask."""
return TArray(self.data.chunk(*args, **kwargs), self.norm)
def copy(self, deep=True):
"""Return a copy of the TArray."""
return TArray(self.data.copy(deep=deep), self.norm)
def diff(self, *args, **kwargs):
"""Return the n-th order differences, see xarray.DataArray.diff."""
return TArray(self.data.diff(*args, **kwargs), self.norm)
def transpose(self, *dims, **kwargs):
"""Reorder the dimensions."""
return TArray(self.data.transpose(*dims, **kwargs), self.norm)
def roll(self, *args, **kwargs):
"""Roll the data along dims."""
return TArray(self.data.roll(*args, **kwargs), self.norm)
def shift(self, *args, **kwargs):
"""Shift the data along dims, filling with NaN."""
return TArray(self.data.shift(*args, **kwargs), self.norm)
def squeeze(self, *args, **kwargs):
"""Remove dims of length one."""
return TArray(self.data.squeeze(*args, **kwargs), self.norm)
def expand_dims(self, *args, **kwargs):
"""Add dims of length one."""
return TArray(self.data.expand_dims(*args, **kwargs), self.norm)
def assign_coords(self, *args, **kwargs):
"""Assign coordinates."""
return TArray(self.data.assign_coords(*args, **kwargs), self.norm)
def drop_vars(self, *args, **kwargs):
"""Drop coordinates."""
return TArray(self.data.drop_vars(*args, **kwargs), self.norm)
def interp(self, *args, **kwargs):
"""Interpolate the data to new coordinates."""
return TArray(self.data.interp(*args, **kwargs), self.norm)
def reindex(self, *args, **kwargs):
"""Conform the data to new coordinates."""
return TArray(self.data.reindex(*args, **kwargs), self.norm)
def broadcast_like(self, other, **kwargs):
"""Broadcast the data against other."""
if isinstance(other, TArray):
other = other.data
return TArray(self.data.broadcast_like(other, **kwargs), self.norm)
# Methods called if a TArray is the second argument of a binary ufunc
_REFLECTED = dict(
add="__radd__", subtract="__rsub__", multiply="__rmul__",
divide="__rtruediv__", power="__rpow__", less="__gt__",
less_equal="__ge__", greater="__lt__", greater_equal="__le__",
equal="__eq__", not_equal="__ne__", maximum="maximum",
minimum="minimum", fmax="fmax", fmin="fmin", hypot="hypot",
logical_and="__and__", logical_or="__or__", logical_xor="__xor__",
)
def __array_ufunc__(self, ufunc, method, *args, **kwargs):
"""Dispatch numpy ufuncs to the methods of the same name."""
name = ufunc.__name__
if method != "__call__" or not hasattr(TArray, name):
raise TypeError(f"Function '{name}' is not implemented for "
"TArray.")
if kwargs.get("out") is not None:
raise TypeError("The out keyword is not implemented for TArray.")
if isinstance(args[0], TArray):
return getattr(args[0], name)(*args[1:], **kwargs)
if len(args) == 2 and name in self._REFLECTED:
if kwargs:
raise TypeError(f"Keywords {list(kwargs)} are not "
f"implemented for '{name}' with TArray as "
"second argument.")
args[1]._assert_not_plain(args[0])
return getattr(args[1], self._REFLECTED[name])(args[0])
raise TypeError(f"Function '{name}' is not implemented for TArray.")
# numpy functions dispatched to methods or attributes of the same
# signature, None applies the function to the denormalized values
_NP_FUNCTIONS = dict(
shape="shape", ndim="ndim", size="size", real="real", imag="imag",
round="round", around="round", clip="clip", isclose="isclose",
allclose="allclose", gradient="gradient", trapezoid="trapezoid",
argmin="argmin", argmax="argmax", nanargmin="nanargmin",
nanargmax="nanargmax", argwhere="argwhere", atleast_1d=None,
atleast_2d=None, atleast_3d=None,
)
# numpy reductions dispatched to methods, with skipna following the numpy
# NaN handling, since xarray skips NaN by default
_NP_REDUCTIONS = dict(
all=("all", None), any=("any", None), prod=("prod", None),
nanprod=("nanprod", None), cumprod=("cumprod", False),
min=("min", False), amin=("min", False), nanmin=("min", True),
max=("max", False), amax=("max", False), nanmax=("max", True),
sum=("sum", False), nansum=("sum", True), cumsum=("cumsum", False),
nancumsum=("cumsum", True), mean=("mean", False),
nanmean=("mean", True), median=("median", False),
nanmedian=("median", True), std=("std", False), nanstd=("std", True),
var=("var", False), nanvar=("var", True),
)
def __array_function__(self, func, types, args, kwargs):
"""Dispatch numpy functions with a TArray as first argument."""
name = func.__name__
if (not all(issubclass(t, (TArray, Quantity)) for t in types)
or not args or not isinstance(args[0], TArray)):
return NotImplemented
if name in self._NP_REDUCTIONS:
return args[0]._np_reduce(name, *args[1:], **kwargs)
if name not in self._NP_FUNCTIONS:
return NotImplemented
if self._NP_FUNCTIONS[name] is None:
return func(args[0].denormalized_values, *args[1:], **kwargs)
attr = getattr(args[0], self._NP_FUNCTIONS[name])
return attr(*args[1:], **kwargs) if callable(attr) else attr
def _np_reduce(self, name, axis=None, **kwargs):
"""Apply a reduction with the numpy signature, like np.sum(a, axis)."""
method, skipna = self._NP_REDUCTIONS[name]
# numpy accumulates the flattened array, xarray along every dim
if method.startswith("cum") and axis is None and self.ndim > 1:
raise TypeError(f"An axis is required for np.{name} of a TArray "
"with more than one dimension.")
if skipna is not None:
kwargs["skipna"] = skipna
if axis is not None:
kwargs["axis"] = axis
return getattr(self, method)(**kwargs)
# Make pint return NotImplemented for a TArray operand, such that
# Quantity * TArray is handled by TArray.__rmul__
pint.compat.upcast_type_map[f"{TArray.__module__}.TArray"] = TArray
