mirror of
https://github.com/simonw/datasette.git
synced 2025-12-10 16:51:24 +01:00
306 lines
9.8 KiB
Python
306 lines
9.8 KiB
Python
"""
|
|
pint.facets.numpy.quantity
|
|
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
|
|
:copyright: 2022 by Pint Authors, see AUTHORS for more details.
|
|
:license: BSD, see LICENSE for more details.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import math
|
|
import warnings
|
|
from typing import Any, Generic
|
|
|
|
from ..._typing import Shape
|
|
from ...compat import HAS_NUMPY, _to_magnitude, np
|
|
from ...errors import DimensionalityError, PintTypeError, UnitStrippedWarning
|
|
from ..plain import MagnitudeT, PlainQuantity
|
|
from .numpy_func import (
|
|
HANDLED_UFUNCS,
|
|
copy_units_output_ufuncs,
|
|
get_op_output_unit,
|
|
matching_input_copy_units_output_ufuncs,
|
|
matching_input_set_units_output_ufuncs,
|
|
numpy_wrap,
|
|
op_units_output_ufuncs,
|
|
set_units_ufuncs,
|
|
)
|
|
|
|
try:
|
|
import uncertainties.unumpy as unp
|
|
from uncertainties import UFloat, ufloat
|
|
|
|
HAS_UNCERTAINTIES = True
|
|
except ImportError:
|
|
unp = np
|
|
ufloat = Ufloat = None
|
|
HAS_UNCERTAINTIES = False
|
|
|
|
|
|
def method_wraps(numpy_func):
|
|
if isinstance(numpy_func, str):
|
|
numpy_func = getattr(np, numpy_func, None)
|
|
|
|
def wrapper(func):
|
|
func.__wrapped__ = numpy_func
|
|
|
|
return func
|
|
|
|
return wrapper
|
|
|
|
|
|
class NumpyQuantity(Generic[MagnitudeT], PlainQuantity[MagnitudeT]):
|
|
""" """
|
|
|
|
# NumPy function/ufunc support
|
|
__array_priority__ = 17
|
|
|
|
def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
|
|
if method != "__call__":
|
|
# Only handle ufuncs as callables
|
|
return NotImplemented
|
|
|
|
# Replicate types from __array_function__
|
|
types = {
|
|
type(arg)
|
|
for arg in list(inputs) + list(kwargs.values())
|
|
if hasattr(arg, "__array_ufunc__")
|
|
}
|
|
|
|
return numpy_wrap("ufunc", ufunc, inputs, kwargs, types)
|
|
|
|
def __array_function__(self, func, types, args, kwargs):
|
|
return numpy_wrap("function", func, args, kwargs, types)
|
|
|
|
_wrapped_numpy_methods = ["flatten", "astype", "item"]
|
|
|
|
def _numpy_method_wrap(self, func, *args, **kwargs):
|
|
"""Convenience method to wrap on the fly NumPy ndarray methods taking
|
|
care of the units.
|
|
"""
|
|
|
|
# Set input units if needed
|
|
if func.__name__ in set_units_ufuncs:
|
|
self.__ito_if_needed(set_units_ufuncs[func.__name__][0])
|
|
|
|
value = func(*args, **kwargs)
|
|
|
|
# Set output units as needed
|
|
if func.__name__ in (
|
|
matching_input_copy_units_output_ufuncs
|
|
+ copy_units_output_ufuncs
|
|
+ self._wrapped_numpy_methods
|
|
):
|
|
output_unit = self._units
|
|
elif func.__name__ in set_units_ufuncs:
|
|
output_unit = set_units_ufuncs[func.__name__][1]
|
|
elif func.__name__ in matching_input_set_units_output_ufuncs:
|
|
output_unit = matching_input_set_units_output_ufuncs[func.__name__]
|
|
elif func.__name__ in op_units_output_ufuncs:
|
|
output_unit = get_op_output_unit(
|
|
op_units_output_ufuncs[func.__name__],
|
|
self.units,
|
|
list(args) + list(kwargs.values()),
|
|
self._magnitude.size,
|
|
)
|
|
else:
|
|
output_unit = None
|
|
|
|
if output_unit is not None:
|
|
return self.__class__(value, output_unit)
|
|
|
|
return value
|
|
|
|
def __array__(self, t=None) -> np.ndarray:
|
|
if HAS_NUMPY and isinstance(self._magnitude, np.ndarray):
|
|
warnings.warn(
|
|
"The unit of the quantity is stripped when downcasting to ndarray.",
|
|
UnitStrippedWarning,
|
|
stacklevel=2,
|
|
)
|
|
return _to_magnitude(self._magnitude, force_ndarray=True)
|
|
|
|
def clip(self, min=None, max=None, out=None, **kwargs):
|
|
if min is not None:
|
|
if isinstance(min, self.__class__):
|
|
min = min.to(self).magnitude
|
|
elif self.dimensionless:
|
|
pass
|
|
else:
|
|
raise DimensionalityError("dimensionless", self._units)
|
|
|
|
if max is not None:
|
|
if isinstance(max, self.__class__):
|
|
max = max.to(self).magnitude
|
|
elif self.dimensionless:
|
|
pass
|
|
else:
|
|
raise DimensionalityError("dimensionless", self._units)
|
|
return self.__class__(self.magnitude.clip(min, max, out, **kwargs), self._units)
|
|
|
|
def fill(self: NumpyQuantity, value) -> None:
|
|
self._units = value._units
|
|
return self.magnitude.fill(value.magnitude)
|
|
|
|
def put(self: NumpyQuantity, indices, values, mode="raise") -> None:
|
|
if isinstance(values, self.__class__):
|
|
values = values.to(self).magnitude
|
|
elif self.dimensionless:
|
|
values = self.__class__(values, "").to(self)
|
|
else:
|
|
raise DimensionalityError("dimensionless", self._units)
|
|
self.magnitude.put(indices, values, mode)
|
|
|
|
@property
|
|
def real(self) -> NumpyQuantity:
|
|
return self.__class__(self._magnitude.real, self._units)
|
|
|
|
@property
|
|
def imag(self) -> NumpyQuantity:
|
|
return self.__class__(self._magnitude.imag, self._units)
|
|
|
|
@property
|
|
def T(self):
|
|
return self.__class__(self._magnitude.T, self._units)
|
|
|
|
@property
|
|
def flat(self):
|
|
for v in self._magnitude.flat:
|
|
yield self.__class__(v, self._units)
|
|
|
|
@property
|
|
def shape(self) -> Shape:
|
|
return self._magnitude.shape
|
|
|
|
@property
|
|
def dtype(self):
|
|
return self._magnitude.dtype
|
|
|
|
@shape.setter
|
|
def shape(self, value):
|
|
self._magnitude.shape = value
|
|
|
|
def searchsorted(self, v, side="left", sorter=None):
|
|
if isinstance(v, self.__class__):
|
|
v = v.to(self).magnitude
|
|
elif self.dimensionless:
|
|
v = self.__class__(v, "").to(self)
|
|
else:
|
|
raise DimensionalityError("dimensionless", self._units)
|
|
return self.magnitude.searchsorted(v, side)
|
|
|
|
def dot(self, b):
|
|
"""Dot product of two arrays.
|
|
|
|
Wraps np.dot().
|
|
"""
|
|
|
|
return np.dot(self, b)
|
|
|
|
@method_wraps("prod")
|
|
def prod(self, *args, **kwargs):
|
|
"""Return the product of quantity elements over a given axis
|
|
|
|
Wraps np.prod().
|
|
"""
|
|
return np.prod(self, *args, **kwargs)
|
|
|
|
def __ito_if_needed(self, to_units):
|
|
if self.unitless and to_units == "radian":
|
|
return
|
|
|
|
self.ito(to_units)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._magnitude)
|
|
|
|
def __getattr__(self, item) -> Any:
|
|
if item.startswith("__array_"):
|
|
# Handle array protocol attributes other than `__array__`
|
|
raise AttributeError(f"Array protocol attribute {item} not available.")
|
|
elif item in HANDLED_UFUNCS or item in self._wrapped_numpy_methods:
|
|
magnitude_as_duck_array = _to_magnitude(
|
|
self._magnitude, force_ndarray_like=True
|
|
)
|
|
try:
|
|
attr = getattr(magnitude_as_duck_array, item)
|
|
return functools.partial(self._numpy_method_wrap, attr)
|
|
except AttributeError:
|
|
raise AttributeError(
|
|
f"NumPy method {item} not available on {type(magnitude_as_duck_array)}"
|
|
)
|
|
except TypeError as exc:
|
|
if "not callable" in str(exc):
|
|
raise AttributeError(
|
|
f"NumPy method {item} not callable on {type(magnitude_as_duck_array)}"
|
|
)
|
|
else:
|
|
raise exc
|
|
elif (
|
|
HAS_UNCERTAINTIES and item == "ndim" and isinstance(self._magnitude, UFloat)
|
|
):
|
|
# Dimensionality of a single UFloat is 0, like any other scalar
|
|
return 0
|
|
|
|
try:
|
|
return getattr(self._magnitude, item)
|
|
except AttributeError:
|
|
raise AttributeError(
|
|
"Neither Quantity object nor its magnitude ({}) "
|
|
"has attribute '{}'".format(self._magnitude, item)
|
|
)
|
|
|
|
def __getitem__(self, key):
|
|
try:
|
|
return type(self)(self._magnitude[key], self._units)
|
|
except PintTypeError:
|
|
raise
|
|
except TypeError:
|
|
raise TypeError(
|
|
"Neither Quantity object nor its magnitude ({})"
|
|
"supports indexing".format(self._magnitude)
|
|
)
|
|
|
|
def __setitem__(self, key, value):
|
|
try:
|
|
# If we're dealing with a masked single value or a nan, set it
|
|
if (
|
|
isinstance(self._magnitude, np.ma.MaskedArray)
|
|
and np.ma.is_masked(value)
|
|
and getattr(value, "size", 0) == 1
|
|
) or (getattr(value, "ndim", 0) == 0 and math.isnan(value)):
|
|
self._magnitude[key] = value
|
|
return
|
|
except TypeError:
|
|
pass
|
|
|
|
try:
|
|
if isinstance(value, self.__class__):
|
|
factor = self.__class__(
|
|
value.magnitude, value._units / self._units
|
|
).to_root_units()
|
|
else:
|
|
factor = self.__class__(value, self._units ** (-1)).to_root_units()
|
|
|
|
if isinstance(factor, self.__class__):
|
|
if not factor.dimensionless:
|
|
raise DimensionalityError(
|
|
value,
|
|
self.units,
|
|
extra_msg=". Assign a quantity with the same dimensionality "
|
|
"or access the magnitude directly as "
|
|
f"`obj.magnitude[{key}] = {value}`.",
|
|
)
|
|
self._magnitude[key] = factor.magnitude
|
|
else:
|
|
self._magnitude[key] = factor
|
|
|
|
except PintTypeError:
|
|
raise
|
|
except TypeError as exc:
|
|
raise TypeError(
|
|
f"Neither Quantity object nor its magnitude ({self._magnitude}) "
|
|
"supports indexing"
|
|
) from exc
|